fix: decode dictionary input for PyArrow UDFs

.github/workflows/pyarrow_udf_test.yml

@@ -28,23 +28,29 @@ on:
     paths: &feature-paths
       - "pom.xml"
       - "common/pom.xml"
-      - "common/src/main/scala/org/apache/comet/CometConf.scala"
+      - "native/shuffle/src/spark_unsafe/row.rs"
       - "spark/pom.xml"
+      - "spark/src/main/java/org/apache/comet/vector/**"
+      - "spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java"
+      - "spark/src/main/scala/org/apache/comet/CometConf.scala"
       - "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+      - "spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala"
+      - "spark/src/main/scala/org/apache/comet/vector/StreamReader.scala"
       - "spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala"
       - "spark/src/main/scala/org/apache/spark/sql/comet/shims/MapInBatchInfo.scala"
       - "spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
-      - "spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
-      - "spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala"
+      - "spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala"
+      - "spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
       - "spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala"
       - "spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala"
       - "spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala"
       - "spark/src/test/resources/pyspark/conftest.py"
       - "spark/src/test/resources/pyspark/test_pyarrow_udf.py"
+      - "spark/src/test/resources/pyspark/test_pyarrow_udf_dictionary_shuffle.py"
       - "spark/src/test/spark-3.5/org/apache/spark/sql/comet/CometMapInBatchSuite.scala"
       - "spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala"
       - "spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala"
@@ -130,3 +136,5 @@ jobs:
         run: |
           /tmp/venv/bin/python -m pytest -v \
             spark/src/test/resources/pyspark/test_pyarrow_udf.py
+          /tmp/venv/bin/python -m pytest -v \
+            spark/src/test/resources/pyspark/test_pyarrow_udf_dictionary_shuffle.py

docs/source/user-guide/latest/pyarrow-udfs.md

@@ -183,8 +183,9 @@ on the unoptimized path.
 - The optimization currently applies only to `mapInArrow` and `mapInPandas`. Scalar pandas UDFs
   (`@pandas_udf`) and grouped operations (`applyInPandas`) are not yet supported.
 - The optimization requires Arrow data on the input side. If a shuffle sits between the upstream
-  Comet operator and the Python UDF, you need Comet's native shuffle for the optimization to
-  apply. Set `spark.shuffle.manager` to
+  Comet operator and the Python UDF, use Comet's columnar shuffle for the optimization to apply.
+  Both the `jvm` and `native` shuffle modes can feed `CometMapInBatch`. Set
+  `spark.shuffle.manager` to
   `org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager` and enable
   `spark.comet.shuffle.enabled=true` at session startup. With a vanilla Spark `Exchange`
   in the plan the data leaves the shuffle as rows and the optimization cannot fire.
@@ -209,8 +210,12 @@ on the unoptimized path.
   requested by the configuration. `EliminateRedundantTransitions` therefore skips the rewrite
   and vanilla Spark handles the operation. Comet can read `large_string` and `large_binary`
   columns returned by a Python worker; that output support does not widen the input vectors.
-- Comet writes input Arrow IPC record batches directly from its existing vector buffers. The
-  only additional Arrow buffer is the validity bitmap for the non-null struct that wraps the
-  input columns. Writing the IPC bytes to the Python worker's pipe still requires one copy;
-  that copy is inherent to Spark's process-based Python transport. This path does not transfer
-  buffers between Arrow allocators or change their ownership.
+- Comet writes input Arrow IPC record batches directly from existing plain vector buffers. The
+  only additional Arrow buffer for plain inputs is the validity bitmap for the non-null struct
+  that wraps the input columns. Before decoding dictionary-encoded shuffle columns, Comet uses
+  Spark's Arrow record threshold and the decoded dictionary size against Spark's byte threshold to
+  split the compact batch. Each temporary logical slice is released after its synchronous write.
+  Plain-only inputs continue to preserve their upstream Comet batch boundaries. Writing the IPC
+  bytes to the Python worker's pipe still requires one copy; that copy is inherent to Spark's
+  process-based Python transport. Borrowed buffers are not transferred between Arrow allocators or
+  given new ownership.

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -596,9 +596,10 @@ object CometConf extends ShimCometConf {
     .withAlternative("spark.comet.shuffle.preferDictionary.ratio")
     .category(CATEGORY_SHUFFLE)
     .doc(
-      "The ratio of total values to distinct values in a string column to decide whether to " +
+      "The ratio of total values to distinct values in a string or binary column to decide " +
+        "whether to " +
         "prefer dictionary encoding when shuffling the column. If the ratio is higher than " +
-        "this config, dictionary encoding will be used on shuffling string column. This config " +
+        "this config, dictionary encoding will be used when shuffling the column. This config " +
         "is effective if it is higher than 1.0. Note that this " +
         "config is only used when `spark.comet.shuffle.mode` is `jvm`.")
     .doubleConf

spark/src/main/scala/org/apache/comet/Native.scala

@@ -136,9 +136,9 @@ class Native extends NativeBase {
    * @param file
    *   the file path to write to.
    * @param preferDictionaryRatio
-   *   the ratio of total values to distinct values in a string column that makes the writer to
-   *   prefer dictionary encoding. If it is larger than the specified ratio, dictionary encoding
-   *   will be used when writing columns of string type.
+   *   the ratio of total values to distinct values in a string or binary column that makes the
+   *   writer prefer dictionary encoding. If it is larger than the specified ratio, dictionary
+   *   encoding will be used when writing columns of either type.
    * @param batchSize
    *   the batch size on the native side to buffer outputs during the row to columnar conversion
    *   before writing them out to disk.

spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala

@@ -43,5 +43,7 @@ trait ShimCometMapInBatch extends Spark4xMapInBatchSupport {
       schema,
       runnerInputs.pythonRunnerConf,
       pythonMetrics,
-      runnerInputs.jobArtifactUUID).compute(batchIter, partitionId, context)
+      runnerInputs.jobArtifactUUID,
+      runnerInputs.arrowMaxRecordsPerBatch,
+      runnerInputs.arrowMaxBytesPerBatch).compute(batchIter, partitionId, context)
 }

spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala

@@ -38,7 +38,9 @@ class CometArrowPythonRunner(
     override val schema: StructType,
     override val workerConf: Map[String, String],
     override val pythonMetrics: Map[String, SQLMetric],
-    jobArtifactUUID: Option[String])
+    jobArtifactUUID: Option[String],
+    override val arrowMaxRecordsPerBatch: Int,
+    override val arrowMaxBytesPerBatch: Long)
     extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch](
       funcs.map(_._1),
       evalType,

spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala

@@ -44,5 +44,7 @@ trait ShimCometMapInBatch extends Spark4xMapInBatchSupport {
       runnerInputs.pythonRunnerConf,
       pythonMetrics,
       runnerInputs.jobArtifactUUID,
-      None).compute(batchIter, partitionId, context)
+      None,
+      runnerInputs.arrowMaxRecordsPerBatch,
+      runnerInputs.arrowMaxBytesPerBatch).compute(batchIter, partitionId, context)
 }

spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala

@@ -40,7 +40,9 @@ class CometArrowPythonRunner(
     override val workerConf: Map[String, String],
     override val pythonMetrics: Map[String, SQLMetric],
     jobArtifactUUID: Option[String],
-    sessionUUID: Option[String])
+    sessionUUID: Option[String],
+    override val arrowMaxRecordsPerBatch: Int,
+    override val arrowMaxBytesPerBatch: Long)
     extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch](
       funcs.map(_._1),
       evalType,

spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala

@@ -44,5 +44,7 @@ trait ShimCometMapInBatch extends Spark4xMapInBatchSupport {
       runnerInputs.pythonRunnerConf,
       pythonMetrics,
       runnerInputs.jobArtifactUUID,
-      None).compute(batchIter, partitionId, context)
+      None,
+      runnerInputs.arrowMaxRecordsPerBatch,
+      runnerInputs.arrowMaxBytesPerBatch).compute(batchIter, partitionId, context)
 }

spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala

@@ -39,7 +39,9 @@ class CometArrowPythonRunner(
     pythonRunnerConf: Map[String, String],
     override val pythonMetrics: Map[String, SQLMetric],
     jobArtifactUUID: Option[String],
-    sessionUUID: Option[String])
+    sessionUUID: Option[String],
+    override val arrowMaxRecordsPerBatch: Int,
+    override val arrowMaxBytesPerBatch: Long)
     extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch](
       funcs.map(_._1),
       evalType,

spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala

@@ -63,7 +63,9 @@ trait Spark4xMapInBatchSupport {
   protected case class RunnerInputs(
       chainedFunc: Seq[(ChainedPythonFunctions, Long)],
       pythonRunnerConf: Map[String, String],
-      jobArtifactUUID: Option[String])
+      jobArtifactUUID: Option[String],
+      arrowMaxRecordsPerBatch: Int,
+      arrowMaxBytesPerBatch: Long)
 
   /**
    * Resolves the `SQLConf`-derived inputs the `ArrowPythonRunner` needs. Must be called on the
@@ -74,5 +76,7 @@ trait Spark4xMapInBatchSupport {
     RunnerInputs(
       chainedFunc = Seq((ChainedPythonFunctions(Seq(pythonUDF.func)), pythonUDF.resultId.id)),
       pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf),
-      jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid))
+      jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid),
+      arrowMaxRecordsPerBatch = conf.arrowMaxRecordsPerBatch,
+      arrowMaxBytesPerBatch = conf.arrowMaxBytesPerBatch)
 }