sunchao opened a new pull request, #5560: URL: https://github.com/apache/datafusion-comet/pull/5560
## Which issue does this PR close? Extracted from #5557 while addressing #5555. ## Rationale for this change A JVM Comet shuffle can encode repeated string and binary columns as dictionaries. The accelerated `mapInArrow` / `mapInPandas` runner currently takes the dictionary indices vector, advertises its dictionary metadata, and starts `ArrowStreamWriter` without a dictionary provider. Arrow then fails before the Python worker can receive the batch. This affects the existing accelerated path with `spark.sql.execution.arrow.useLargeVarTypes=false`, so it is an independent crash fix rather than part of large-offset support. For example, a Parquet input with repeated string and binary values followed by `repartition(2, "id")`, using `spark.comet.shuffle.mode=jvm` and `spark.comet.shuffle.jvm.preferDictionary.ratio=1.01`, produces `CometDictionaryVector` input and fails on current `main`. ## What changes are included in this PR? The runner now identifies top-level `CometDictionaryVector` columns and decodes their logical values into temporary vectors owned by the runner allocator. It derives the stream schema and performs the synchronous IPC write from those decoded vectors, then closes them on both success and failure. Plain Comet vectors keep the existing borrowed-buffer path and ownership model. The regression coverage includes string and binary dictionaries, repeated, empty, Unicode, and null values, injected write failure, buffer reference counts, and allocator cleanup. A real-worker test exercises both `mapInArrow` and `mapInPandas` after an actual JVM Comet shuffle. The PyArrow workflow runs that module separately on each supported Spark worker so its shuffle-manager Spark session cannot inherit the main module's shuffle-disabled session. The benchmark gains a low-cardinality JVM-shuffle workload and an optional `BENCHMARK_WORKLOAD` selector. The workload requires `CometColumnarExchange` in both modes and compares vanilla and accelerated Python execution on the same shuffled input. ## How are these changes tested? Native code was built before the JVM checks; this PR does not change native code. - Spark 4.0 / Scala 2.13 root-reactor package build: **11/11 focused JVM tests passed**, including dictionary success and injected-failure ownership cases. - PySpark 4.0.4 / JDK 17 / PyArrow 25.0.1 / pandas 3.0.5: the existing **117 tests passed**, and the new dictionary-shuffle module passed **2/2 tests** in the separate process used by CI. - Benchmark smoke: the new workload completed all four API/mode cases with 4,096 rows. - A local `local[2]` debug-build timing with 2,000,000 rows, two warmups, and five measured iterations produced median end-to-end times of **2.453 s vanilla / 1.794 s accelerated (1.37x)** for `mapInArrow`, and **2.415 s / 1.797 s (1.34x)** for `mapInPandas`. These numbers show that the required dictionary materialization still preserves an end-to-end gain in this setup; they are not a general performance claim. - Scala formatting/style, Python formatting/syntax, workflow YAML parsing, and `git diff --check` passed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
