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 the dictionary indices and metadata to ArrowStreamWriter without a dictionary provider, so Arrow failed before the Python worker received the batch.

Simply decoding an entire compact shuffle batch would fix that crash but could expand a small index vector into a very large string or binary vector. For example, thousands of references to one 300 KiB value occupy little space as dictionary indices but exceed Arrow's regular 2 GiB variable-width buffer limit when decoded together. The runner therefore has to bound dictionary materialization before decoding while keeping every column on the same row boundaries.

What changes are included in this PR?

The runner now recognizes top-level CometDictionaryVector columns and decodes their logical values into temporary vectors owned by the runner allocator. Before decoding, it computes row-aligned ranges from the dictionary indices and logical string or binary lengths. It applies Spark's Arrow record threshold and decoded-dictionary byte threshold, preserves Spark's one-row soft-limit behavior, and adds a hard guard for regular Arrow variable-width buffers.

For each range, every input column is sliced at the same boundaries. Dictionary slices are decoded, all columns are serialized synchronously into the existing Arrow stream, and both decoded vectors and slice references are released before the next range. A batch that does not need splitting avoids the slice copies, and plain vectors keep their existing borrowed-buffer path.

This batching is intentionally scoped to dictionary materialization. The byte estimate covers decoded dictionary vectors rather than every plain vector in the record batch, and inputs without dictionaries preserve their upstream Comet batch boundaries. This PR also does not add spark.sql.execution.arrow.useLargeVarTypes=true support.

The regression coverage includes string and binary dictionaries, null, empty, repeated, and Unicode values; record and byte limits; an allocator-capped proof that slicing happens before decoding; exact IPC batch boundaries; source-buffer reference counts; and cleanup when serialization fails on a later slice. Real-worker tests exercise both mapInArrow and mapInPandas after an actual JVM Comet shuffle. Negative controls fail both on the original missing dictionary provider and on a dictionary-capable runner without limit wiring.

The documentation now states that both JVM and native Comet columnar shuffle modes can feed CometMapInBatch, and the dictionary-ratio documentation covers binary as well as string columns. The PyArrow workflow watches the relevant shuffle, vector, shared runner, and version-specific wiring files and runs the dictionary module separately for Spark 4.0 and 4.1 workers.

The benchmark retains its low-cardinality JVM-shuffle workload and optional workload selector. It compares vanilla and accelerated Python execution on the same shuffled input.

How are these changes tested?

• Native library debug build: passed.
• Spark 4.0 / Scala 2.13 root-reactor runner suite: 13/13 passed, with BUILD SUCCESS.
• Spark 4.0 / Scala 2.13 root-reactor package build: passed, with BUILD SUCCESS.
• PySpark 4.0.4 / Python 3.11.15 / PyArrow 25.0.1 / pandas 3.0.5: the existing worker module passed 117/117, and the dictionary-shuffle module passed 6/6. The limit cases produced exactly five 2-row batches and ten 1-row batches for both APIs.
• CI on this head: the Spark 4.0 and Spark 4.1 real-worker jobs both reported Maven BUILD SUCCESS, 117/117 general PyArrow tests, and 6/6 dictionary-shuffle tests. The dedicated Spark 4.1 build also passed.
• Benchmark smoke: the dictionary JVM-shuffle workload completed all four API/mode cases with 4,096 rows against the bounded implementation.
• Scala formatting/style, Python formatting, workflow YAML parsing, and git diff --check: passed.

Local Spark 4.1 and 4.2 package attempts were blocked before compilation because the configured Maven mirror timed out and the cache lacks jackson-bom:2.21.2. The version-specific call sites and the Spark 4.0.4, 4.1.3, and 4.2.0 SQLConf getter signatures were checked directly. CI has now covered Spark 4.1; the broader Spark 4.2 matrix remains gated on its shared native-build job.

Diff size

The combined PR changes 15 files with 819 insertions and 71 deletions. Of the additions, 507 are automated regression tests, 47 are benchmark support, 235 are runner and version-wiring code, and 30 are documentation, workflow, or configuration text.

[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.

[andygrove]

inputBatchRanges walks every row before any decoding happens, and with Comet's defaults it can never split: spark.comet.batchSize and spark.comet.shuffle.jvm.batchSize are both 8192, under spark.sql.execution.arrow.maxRecordsPerBatch (10000), and a decoded 8192-row batch is nowhere near spark.sql.execution.arrow.maxBytesPerBatch (64MB). I measured it on an 8192-row batch with one dictionary column (64 distinct ~28-byte values): the scan takes 227-320 us against 630-731 us for the whole ranges+decode+serialize path, so 34-43% of the work is thrown away.

Two things would help. First, a cheap upper bound that skips the scan entirely when the batch provably fits: one pass over the dictionary for max(getValueLength) is O(distinct values) rather than O(rows), and initialBytes + numRows * (maxLen + OFFSET_WIDTH + 1) bounds the decoded size from above. If that is under byteLimit and numRows <= recordLimit, return Seq(0 -> numRows) without touching the indices. That took the same batch from 320 us to 1.0 us for me, and I checked it against the exact scan over 300 randomized configs (row counts, dictionary sizes, null rates, both limits) with no disagreement on the 116 where it fired.

Second, the fallback scan itself: dictionaries.foldLeft(0L) boxes the accumulator and destructures a tuple once per row per column, and the split row's bytes are computed twice. A while loop over parallel arrays took it from 320 us to 99 us on the same input.

[andygrove]

columns.collect { case column: CometDictionaryVector => ... } only sees top-level columns, so a struct, list or map whose child is dictionary-encoded goes through the case column => arm in withMaterializedInputVectors unchanged and reaches startWriter with a DictionaryEncoding on the child field and a null provider. I built a Struct<child: Dictionary<Int32, Utf8>> and it fails with the same NPE this PR fixes:

java.lang.NullPointerException: Cannot invoke
  "org.apache.arrow.vector.dictionary.DictionaryProvider.lookup(long)" because "provider" is null

I do not think it is reachable today. builder_to_array in native/shuffle/src/spark_unsafe/row.rs dictionary-encodes only top-level Utf8 and Binary, so the JVM shuffle cannot produce it. But CometStructVector builds its children through CometVector.getVector, which does create nested CometDictionaryVectors, so an FFI-imported batch could. Would you either recurse into children or add an explicit check with a message that names the column, so this surfaces as something diagnosable rather than an NPE from inside Arrow?

[viirya]

I verified this reachability analysis independently and it holds, with one addition that I think argues for the cheaper of your two options.

row.rs:1317,1329 explicitly disables dictionary encoding for array elements and struct fields, so JVM shuffle can't produce it. I also checked the stacked-CometMapInBatchExec route, since that feeds one runner's flattened output into another: the output side calls CometVector.getVector(vector, null) (CometArrowPythonRunnerBase.scala:274), and getVector does dictionaryProvider.lookup(...) for any dictionary-encoded vector (CometVector.java:249), so it would NPE there first — a pre-existing issue outside this PR, but it means stacking can't quietly introduce a nested dictionary either.

Since it's genuinely unreachable today, I'd favour the explicit check naming the column over full recursion: recursion adds untestable complexity for an unreachable path, while a named check converts a future FFI-path failure from an NPE inside Arrow into something diagnosable.

[andygrove]

With this early return the worker's batch boundaries depend on whether the shuffle chose dictionary encoding for that column, which comes down to spark.comet.shuffle.jvm.preferDictionary.ratio and the data's cardinality. Vanilla Spark applies both limits to every mapInArrow / mapInPandas input (BatchedPythonArrowInput.writeSizedBatch), and Comet's plain path already exceeds maxRecordsPerBatch whenever spark.comet.batchSize is set above 10000.

I read the PR description and I see this is deliberate, and the plain path serializing existing buffers is a real argument for not slicing it. Would it be worth at least applying the record limit uniformly, since that one costs nothing to check, and saying in pyarrow-udfs.md that the byte limit is bounded only for dictionary inputs? Right now the doc says Comet "uses Spark's Arrow record threshold and the decoded dictionary size against Spark's byte threshold to split the compact batch", which reads as though both thresholds are always honoured.

[andygrove]

byteLimit is clamped to MaxDecodedBatchBytes a few lines above, so decodedBytes >= MaxDecodedBatchBytes is strictly implied by the decodedBytes >= byteLimit disjunct right next to it, and rowBytes > MaxDecodedBatchBytes - decodedBytes can only fire if someone sets spark.sql.execution.arrow.maxBytesPerBatch to within one row of 2GB, where Arrow's own OversizedAllocationException gets there first. saturatedAdd is similar: decodedBytes never exceeds byteLimit plus one row's worth, so a Long cannot overflow.

Dropping exceedsArrowLimit and saturatedAdd and keeping just the clamp would make the loop condition read as the one rule it actually implements, and this is also the hot loop from my other comment.

[viirya]

Partly disagreeing with this one. I think the three pieces should be split rather than removed together.

The first disjunct is dead, as you say. byteLimit = min(cfg, MaxDecodedBatchBytes) <= MaxDecodedBatchBytes, so decodedBytes >= MaxDecodedBatchBytes implies decodedBytes >= byteLimit sitting right next to it in the same OR. I removed only that disjunct and swept 4000 randomized legal configs (row counts, per-value sizes up to 2e9, both limits): 0 disagreements.

The second disjunct is not dead. The soft-limit semantics are what make it reachable: decodedBytes >= byteLimit is checked before the row is added, so the crossing row stays in and decodedBytes can reach byteLimit - 1 + maxSingleRowBytes. Concretely, with byteLimit == 2147483647, decodedBytes = 1073741828 and rowBytes = 1073741824: decodedBytes >= byteLimit is false while rowBytes > MaxDecodedBatchBytes - decodedBytes is 1073741824 > 1073741819 → true. Four 1GiB rows give [(0,1),(1,1),(2,1),(3,1)] with the guard and [(0,2),(2,2)] without it — a 2.0 GiB first batch. Over the same 4000 configs, dropping this disjunct changed 1552 of them, with a worst observed batch of 3.68 GiB.

Where the estimate of the trigger condition needs widening: I swept the threshold and at byteLimit = 64MiB / 256MiB / 512MiB / 1GiB the guard and no-guard results are identical (the soft limit splits each oversized row off anyway). Divergence starts above roughly 1GiB. So it isn't "within one row of 2GB" — it's the top half of the legal range. spark.sql.execution.arrow.maxBytesPerBatch is checkValue(x => x > 0 && x <= Int.MaxValue) (SQLConf.scala:4040), so anything up to 2147483647 is accepted; the default is 256MB on Spark 4.0 and 64MB on 4.1.

I'd also not rely on Arrow throwing first: a 2GiB Utf8 batch overflows 32-bit offsets, which is silent corruption rather than a reliable OversizedAllocationException.

saturatedAdd I agree is removable. getValueLength and getBufferSize both return int, so overflowing a Long would need ~4.3 billion dictionary columns in one batch.

Suggested resolution: drop the first disjunct and saturatedAdd, keep rowBytes > MaxDecodedBatchBytes - decodedBytes, and add a comment naming the reachable case (maxBytesPerBatch above ~1GiB with multi-hundred-MiB dictionary values) plus a test pinning it. Nothing in the suite currently exercises this branch, which is probably why it reads as dead.

[andygrove]

ColumnarBatchArrowReader.loadNextBatch has the same block: match CometDictionaryVector, look up the dictionary through the provider, DictionaryEncoder.decode into the caller's allocator, close the temporaries in a finally. The two have already drifted a little (d.provider there vs d.getDictionaryProvider here, and the reader swallows exceptions from close() while this one propagates them), and CometNativeArrowSource.actualFieldOf has a third copy of the lookup half. Would a shared helper next to CometVector be worth it, so the next person who touches dictionary materialization only has one place to look?

[andygrove]

Both new slicing tests use a column set made only of dictionary columns, so the part of the change that I would most want pinned down is untested: that every column is sliced at the same boundaries. Plain, nested and dictionary columns go through three different slice implementations, and nothing here would fail if one of them stopped being sliced.

I added a case locally with 10 rows, maxRecordsPerBatch=3, and columns [dictionary VarChar, plain BigInt with a null, plain VarChar, Struct<k: bigint>]. It passes on this head, batches come back [3,3,3,1] with every column on the right row, so this is regression coverage rather than a bug. Would you add it?

Three smaller ones in the same spirit, all passing today: an all-null dictionary column (with every index null decodedValueBytes returns 0 for every row, so the byte limit can never fire and only the record limit splits), a zero-row dictionary batch, and maxRecordsPerBatch / maxBytesPerBatch at 0 and -1. A randomized property over inputBatchRanges asserting the ranges are contiguous, start at 0, sum to numRows, and never exceed the record limit is cheap and covers a lot of future off-by-one ground. I ran 200 configs and it held.

[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]

Independent of the exceedsArrowLimit discussion: rowsInBatch > 0 means a single row is never split, which is the right Spark contract, but it also means the 2GiB protection is one order weaker than the PR description claims ("adds a hard guard for regular Arrow variable-width buffers"). A single logical value above 2GiB still reaches DictionaryEncoder.decode and overflows. Extreme under Spark's string limits, but worth one sentence in the comment or the doc so the boundary is stated rather than left to be inferred.

Relatedly, decodedValueBytes is a logical estimate — offsets plus value plus amortized validity bit. I checked the two branches are mutually consistent (the getBufferSizeFor(batchRow + 1) - getBufferSizeFor(batchRow) delta amortizes the validity byte the same way the var-width branch does), and it errs conservative, which is the right direction. But Arrow rounds real allocations up to powers of two, so actual memory can approach 2x the estimate and maxBytesPerBatch isn't an actual memory ceiling for users. Worth noting in pyarrow-udfs.md.

[viirya]

Seconding @andygrove's performance point, and I'd treat it as more than a nice-to-have. With Comet's defaults this scan can never split — spark.comet.batchSize and spark.comet.shuffle.jvm.batchSize are both 8192, below maxRecordsPerBatch (10000), and an 8192-row batch is far under maxBytesPerBatch (64MB on Spark 4.1, 256MB on 4.0). So the common case walks every row and discards the result. The O(distinct) upper bound he describes is a sound conservative estimate and short-circuits exactly that case.

Also worth noting the rowBytes fold appears twice (once before the split decision, once recomputed after), so the split row's bytes are computed twice. Rewriting as a while loop over parallel arrays addresses the cost and the duplication together.

[viirya]

Agreed with @andygrove, and I'd rank this the most valuable follow-up in the PR. The central guarantee of the change is that every column is sliced at the same boundaries, but plain / nested / dictionary go through three different slice implementations (CometPlainVector:213, CometStructVector:62, CometDictionaryVector:137) and both slicing tests use all-dictionary column sets — so if one of those stopped being sliced, nothing here goes red.

Separately: inputBatchRanges is deliberately private[python] but has no direct unit test — it's only reached through foreachInputBatch, and it's the subtlest arithmetic in the change. The randomized property he describes (contiguous, starts at 0, sums to numRows, never exceeds the record limit) is worth pointing straight at it.

[viirya]

Nit: the reflow left "whether to " + as an orphan line, which reads oddly. Rebreaking the string manually would be cleaner.

The wording change itself is correct — native/shuffle/src/spark_unsafe/row.rs:1451,1470 confirms both Utf8 and Binary are dictionary-encoded.