fix: decode dictionary input for PyArrow UDFs

[sunchao]

Which issue does this PR close?

Extracted from #5557 while addressing #5555.

Rationale for this change

JVM Comet shuffle can dictionary-encode repeated string and binary values. The accelerated mapInArrow / mapInPandas runner previously sent their indices to an Arrow writer without a dictionary provider, failing before Python received the batch. Decoding an entire compact batch can also expand it beyond Arrow's offset limits or accumulate excessive transport-buffer memory.

What changes are included in this PR?

• Decode top-level dictionaries into temporary logical vectors before constructing the IPC schema.
• Split all columns at matching row boundaries. Apply Spark's record limit to every input and estimate decoded dictionary bytes before allocating.
• Emit one slice per writer call so Spark can drain its transport buffer. Keep the source borrowed, without accessing upstream iterators, until the last slice is written.
• Retain the dictionary-size fast path, Spark's soft byte-limit semantics, and the preventive 32-bit size check. Reject unsupported nested dictionaries with their field path.
• Share decoding cleanup, consolidate test fixtures and parameterized cases, and update documentation and CI path coverage.

The byte estimate excludes plain columns, leaves a single oversized row intact, and is not an allocation ceiling. Preserving dictionaries across IPC is tracked separately in #5906.

How are these changes tested?

The simplification reduces this PR from 2,082 to 1,583 added lines while retaining the regression scenarios.

Current local validation on Spark 4.1.3 / Scala 2.13 / JDK 17:

• Native debug build: passed.
• 45/45 tests passed across the three modified JVM suites, recompiling current sources against verified cached dependencies; Arrow C-data JNI and allocator checks were exercised.
• Both transport regressions still fail in a control using the pre-fix writer and pass with the current writer.
• 6/6 dictionary-shuffle and 133/133 general PyArrow tests passed with real Python workers (PySpark 4.1.3, PyArrow 25.0.1, pandas 3.0.5).
• Java/Scala formatting, Python syntax, suite registration, and git diff --check: passed. Standalone Scalastyle reports no new findings compared with the previous head.

The full Maven reactor build is blocked by the inherited GCS Maven repository being unavailable. Local worker validation uses an unshaded test package containing the newly compiled classes and fresh native library; it does not replace CI validation of the shaded release package. CometMapInBatchSuite and Spark 4.0/4.2 were not rerun locally.

[andygrove]

Thanks for splitting this out of #5557. I checked it out locally, built against Spark 4.1 / Scala 2.13, and ran CometArrowPythonRunnerSuite (13/13) and CometMapInBatchSuite (5/5). I could not run the pytest module here because PyPI is blocked on my machine, but CI covers it on 4.0/4.1/4.2 and is green.

I wrote a handful of extra probes against the new code. Five of them pass, which is good news: a split batch mixing a dictionary column with plain fixed-width, plain var-width and a nested struct stays correctly row-aligned; an all-null dictionary column works; a zero-row dictionary batch works; non-positive limits behave as unlimited; and over 200 randomized configs inputBatchRanges always returns contiguous ranges that start at 0 and sum to numRows. So I have no correctness concern about the main path.

The one probe that fails is a dictionary nested inside a struct, which still hits the same NPE this PR fixes. I left a comment on that, plus a performance measurement on inputBatchRanges that I think is worth acting on, and a few smaller things.

No blockers from me.

[viirya]

I reviewed this independently and reached the same overall conclusion as @andygrove: the core fix is correct and I have no blockers. Recording what I verified, plus one place where I think his line-441 suggestion should only be partly taken (left as a reply on that thread).

What I verified. The crash mechanism is precise: the old code handed CometDecodedVector.getValueVector to serializeBatch, which for a dictionary column is the indices vector whose Field carries a DictionaryEncoding, while the ArrowStreamWriter is constructed with a null provider (line 160). Decoding first means batchFields — and therefore streamFields — now come from the decoded vector and advertise Utf8/Binary. That ordering is the fix and it's right.

I also checked inputBatchRanges against Spark's BatchedPythonArrowInput.writeSizedBatch semantics across 8 boundary configurations (record limit, byte limit, single oversized row, 1 row, limit=1) and the range output is identical, including the subtlety that the row crossing the byte soft limit stays in the current batch. That's easy to get wrong; nice.

Resource handling holds up: foreachInputBatch closes slices in reverse, withMaterializedInputVectors closes decoded vectors in a finally, and sliced CometDictionaryVectors carry isAlias=true so the shared dictionary isn't closed early. The allocator-capped test is my favourite one here — asserting getPeakMemoryAllocation < fullDecodedDataBytes actually proves slicing precedes decoding rather than just checking the output.

Empty batches are unchanged (numRows == 0 → Seq(0 -> 0) → the fast path calls the body once), and metrics still aggregate correctly since startData is captured outside the loop.

The CI path additions are a substantive fix rather than housekeeping, incidentally: this feature genuinely depends on row.rs (the dictionary-encoding decision) and comet/vector/** (CometDictionaryVector), and neither was watched before, so the most relevant changes wouldn't have triggered the test.

[viirya]

Thanks for addressing the earlier feedback. I reviewed the latest head (a75538b75), including the previous discussions and the updated implementation.

The dictionary decoding and schema ordering look correct. The dictionary-size fast path, uniform record limits, shared decoding helper, and mixed-column alignment tests address the earlier concerns well.

I found one remaining issue at the boundary with Spark's Python transport: all slices are serialized within a single writeNextInputToStream call, so their IPC bytes accumulate in Spark's transport buffer before anything is sent to the worker. This keeps the Arrow decoding allocations small but leaves transport memory proportional to the entire expanded source batch.

I reproduced this using the PR's batching/serialization code and Spark 4.1.3's DirectByteBufferOutputStream; details are inline. I think we should fix this before merging, since large dictionary expansion is explicitly part of the problem this PR addresses.

My validation was a focused JVM reproducer, not a rerun of the full JVM or Python suites.

[viirya]

Thanks for the update. I re-reviewed 279bbe303, including the interaction with upstream batch ownership and Spark's transport buffering. Both of my previous concerns are addressed.

The writer now emits one slice per call and leaves the upstream iterators untouched while slices remain. This is important because CometExecIterator.hasNext can close the previous batch and reuse its buffers. Temporary vectors remain scoped to each write, while source cleanup stays with the upstream owner.

I compiled the updated runner and tests against Spark 4.1.3 / JDK 17 and ran the four new transport/lifetime tests: all four passed. As a negative control, both transport tests fail against the previous runner, confirming that they catch the original accumulation issue.

No remaining blockers from me. The PyArrow jobs for Spark 4.0, 4.1, and 4.2 have passed on this head. My local validation was focused; I did not rerun the full build or Python worker suites.

[peterxcli]

Did you consider preserving dictionaries in IPC and decoding them on the Python side?

One option would be to keep the current direct MessageSerializer.serialize(...) path, supply a dictionary provider for schema construction, and emit ArrowDictionaryBatch messages before the record batches that reference them, handling dictionary changes across batches. This could avoid expanding repeated values before transport without switching to ArrowWriter.writeBatch().

We would still need bounded decoding on the Python side before invoking user code, preserving vanilla Spark’s input types. I’d like to understand that tradeoff—does integrating this into the Python worker make JVM-side decoding simpler overall?

[sunchao]

Thanks @peterxcli, this looks like a useful follow-up. Keeping dictionaries in IPC is feasible with the existing direct MessageSerializer path and could reduce transport bytes and JVM decoding allocations.

For this PR, I'd keep the JVM-side decoding so the fix works with existing Spark Python workers. Moving decoding to Python would also require worker integration to split before decoding and preserve vanilla Spark's input types before user code or pandas conversion, plus dictionary ID/update handling across batches. It moves the bounded-decoding work rather than eliminating it.

I opened #5906 to track that improvement, including compatibility and lifetime tests and benchmarks of transport volume, JVM/Python peak memory, and throughput against this implementation.