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?

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

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