Skip to content

Preserve dictionary encoding in Arrow IPC inputs to PyArrow UDFs #5906

Description

@sunchao

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.

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions