apache / apache/datafusion-comet

Preserve dictionary encoding in Arrow IPC inputs to PyArrow UDFs

Open
#5,906 1 comment 0 reactions 1 assignee Claimed by @peterxcli View on GitHub
enhancement requires-triage
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 6h
Merged PRs (30d)
190

Description

## What is the problem the feature request solves?

PR #5560 fixes dictionary-encoded shuffle inputs to accelerated `mapInArrow` / `mapInPandas` by decoding row-aligned slices on the JVM before Arrow IPC serialization. Returning to Spark between slices lets the transport buffer drain, but repeated dictionary values still expand before transport. For example, thousands of rows referencing one long string still send that string repeatedly to Python.

Preserving dictionary encoding across IPC could reduce transport bytes and JVM decoding allocations. The throughput and overall memory benefits need measurement against the implementation in #5560.

## Describe the potential solution

Investigate keeping dictionaries encoded until they reach the Python worker:

- Retain the direct `MessageSerializer.serialize(...)` path, supply a dictionary provider for schema construction, and emit `ArrowDictionaryBatch` messages before record batches that reference them.
- Define dictionary IDs and replacement/update handling across source batches, including changes between dictionary and plain columns under a stable IPC schema.
- Establish how to integrate a Python normalization layer with the supported Spark workers. Decode before `mapInArrow` user code or pandas conversion so user-visible values and input types remain compatible with vanilla Spark.
- Split before decoding to bound newly decoded buffers before allocation, preserving row alignment, record limits, soft byte-limit semantics, and Arrow offset safety. The encoded source dictionary may remain resident; this is not a bound on total worker memory.
- Preserve buffer ownership, transport draining, and cleanup on completion, cancellation, and serialization/decoding failure.

### Validation

- Cover mixed dictionary/plain columns, nulls, empty batches, changing dictionaries, encoding changes, large repeated values, and failure/lifetime cases.
- Run real-worker regressions for both Python APIs across supported Spark versions, including Spark 4.0, 4.1, and 4.2.
- Compare IPC bytes, JVM and Python peak memory, and end-to-end throughput with the JVM-decoding implementation on the same inputs. Include both highly repetitive data and cases with little dictionary benefit.

## Additional context

Raised in [peterxcli's review of #5560](https://github.com/apache/datafusion-comet/pull/5560#pullrequestreview-5189828303).

This tracks a separate performance improvement. The JVM-decoding fix in #5560 keeps compatibility with existing Spark Python workers; dictionary-preserving transport adds worker integration and dictionary-stream state management.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.