fix: decode dictionary input for PyArrow UDFs

[sunchao]

Which issue does this PR close?

Extracted from #5557 while addressing #5555.

Rationale for this change

A JVM Comet shuffle can dictionary-encode repeated string and binary columns. The accelerated mapInArrow / mapInPandas runner previously passed their dictionary indices to an Arrow writer without a dictionary provider, failing before Python received the batch. Decoding an entire compact batch can also expand repeated large values beyond regular Arrow's 32-bit offset range. Serializing every decoded slice in a single writer call accumulates all of their bytes in Spark's transport buffer before Spark can drain it.

What changes are included in this PR?

• Materialize top-level dictionary columns into temporary logical vectors before advertising the IPC schema. Share dictionary lookup, decoding, and failure cleanup with the existing Arrow stream reader.
• Apply Spark's Arrow record limit to every input, including plain-only batches. Bound dictionary materialization using decoded logical bytes before allocating, and slice every column at the same row boundaries.
• Write one slice per writeNextInputToStream call so Spark can drain its direct transport buffer. Keep the source batch borrowed across calls and avoid all upstream iterator access until its last slice is written; release temporary vectors on each call and preserve upstream ownership on failure or cancellation.
• Skip the row scan when a dictionary-cardinality upper bound proves that the batch fits. Otherwise use an unboxed loop that reads each selected dictionary index once, including rows that start a new slice.
• Preserve Spark's soft byte-limit semantics: the crossing row stays in the current batch. Retain the separate preventive 32-bit size check when combining large rows. A single oversized row stays intact; this estimate does not promise an actual allocation ceiling, and it excludes plain vectors.
• Reject nested dictionaries with a named field path before writing. The current Comet shuffle does not produce them; recursive dictionary decoding remains unsupported.
• Add mixed dictionary/plain/nested IPC alignment coverage, empty and null cases, disabled limits, seeded range-property checks, index-read counts, synthetic large-size boundaries, and failure/lifetime tests. Keep real-worker regressions for both Python APIs after a JVM Comet shuffle.
• Update the user guide and string/binary dictionary configuration description. The PyArrow workflow watches the relevant shuffle, Java and Scala vector directories, runner, and version-wiring files; the existing benchmark compares vanilla and accelerated execution on the same dictionary-shuffled input.

How are these changes tested?

Local validation on Spark 4.1.3 / Scala 2.13 / JDK 17:

• Root-reactor package build: passed, using the native debug library freshly built earlier in this PR work; this follow-up changes no native code.
• Four focused JVM suites: 48/48 passed (CometArrowPythonRunnerSuite, CometVectorUtilsSuite, CometArrowStreamSuite, and CometMapInBatchSuite).
• The new transport regressions invoke the production writer with Spark's DirectByteBufferOutputStream, emulate Spark's fill/drain threshold, and round-trip the complete IPC stream. They cover 64 KiB and 16-byte dictionary values, small slices sharing a buffer, empty groups, consecutive sources, metrics, upstream iterator access, and temporary/source lifetimes on interruption and serialization failure.
• Regression check against the original runner classes (whose packaged source matches the pre-fix PR head): both transport tests fail. Pending bytes reach 16,832,728 and 139,416, exceeding their respective 132,096 and 66,576 bounds. Both pass with this fix; the checks measure pending transport bytes without depending on the JVM's direct-memory-limit enforcement.
• Packaged-JAR Python workers with PySpark 4.1.3, Python 3.12.13, PyArrow 25.0.1, and pandas 3.0.5: 133/133 general PyArrow tests and 6/6 dictionary-shuffle tests passed.
• Java/Scala formatting, Scala style, Markdown/YAML formatting, workflow Actionlint, suite-registration checks, and git diff --check: passed.

Hosted checks at the preceding head a75538b75 completed with 74 passed and 9 skipped, including the repaired Preflight suite registration and all three PyArrow worker versions. CI for this follow-up is pending; Spark 4.0 and 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.