You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/03 11:58:36 UTC
(datafusion-comet) 01/01: fix: check accelerated mapInArrow output against the declared schema (#6365)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6365-77aa7f32f4db607193ae16574122419f568da31d
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 4bde6895a75b3738a9709afb948dcabb7f10315e
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sat Oct 3 11:58:17 2026 +0000
fix: check accelerated mapInArrow output against the declared schema (#6365)
PySpark does not cast mapInArrow output to the declared schema, and the
accelerated runner read each returned column at its declared type. So
int64 data under a declared int read as wrong values, int32 data under
bigint read past the end of its buffer, a decimal with another scale
read at the declared scale, and dictionary-encoded output failed with a
NullPointerException.
Check the output stream's schema once, when the runner reads it, and
raise ARROW_TYPE_MISMATCH (a SparkException with the same message on
Spark 4.0) when a column cannot be read as its declared type.
---
docs/source/user-guide/latest/pyarrow-udfs.md | 11 ++
.../sql/comet/shims/ShimCometMapInBatch.scala | 1 +
.../execution/python/CometArrowPythonRunner.scala | 11 +-
.../sql/comet/shims/ShimCometMapInBatch.scala | 1 +
.../execution/python/CometArrowPythonRunner.scala | 14 +-
.../sql/comet/shims/ShimCometMapInBatch.scala | 1 +
.../execution/python/CometArrowPythonRunner.scala | 14 +-
.../sql/comet/shims/Spark4xMapInBatchSupport.scala | 3 +
.../python/CometArrowPythonRunnerBase.scala | 101 ++++++++++-
.../src/test/resources/pyspark/test_pyarrow_udf.py | 144 +++++++++++++++
.../python/CometArrowPythonRunnerSuite.scala | 195 ++++++++++++++++++++-
11 files changed, 484 insertions(+), 12 deletions(-)
diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md
index eb704deb6f..d3f405fc90 100644
--- a/docs/source/user-guide/latest/pyarrow-udfs.md
+++ b/docs/source/user-guide/latest/pyarrow-udfs.md
@@ -207,6 +207,17 @@ 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.
+- PySpark does not convert the batches a `mapInArrow` UDF returns to the declared output schema,
+ so Comet checks their Arrow schema before reading them. When a column cannot be read as its
+ declared type, Comet raises Spark's `ARROW_TYPE_MISMATCH` error (on Spark 4.0, a
+ `SparkException` with the same message). Column names, nullability, `large_string` and
+ `large_binary` offsets, the time zone of a timestamp, and an all-null column of Arrow type
+ `null` may differ from the declared schema. Vanilla Spark reads some mismatches that Comet
+ rejects, unless `spark.sql.execution.arrow.pyspark.validateSchema.enabled` is set: it rescales
+ a decimal to the declared precision and scale, reads a timestamp without a time zone as
+ `TimestampType`, reads other types of the same width (for example `int64` as a timestamp), and
+ ignores extra columns. Return the declared types, or set
+ `spark.comet.exec.pyarrowUDF.enabled=false` for such UDFs.
- Comet applies `spark.sql.execution.arrow.maxRecordsPerBatch` to every input batch, including
batches with only plain columns. Before decoding dictionary-encoded shuffle columns, Comet also
compares their estimated decoded size with `spark.sql.execution.arrow.maxBytesPerBatch`.
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
index b0cd245586..d249fd9711 100644
--- a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
+++ b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
@@ -41,6 +41,7 @@ trait ShimCometMapInBatch extends Spark4xMapInBatchSupport {
evalType,
argOffsets,
schema,
+ runnerInputs.outputSchema,
runnerInputs.pythonRunnerConf,
pythonMetrics,
runnerInputs.jobArtifactUUID,
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala b/spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
index 663de88200..78c0c47a92 100644
--- a/spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
+++ b/spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
@@ -21,9 +21,10 @@ package org.apache.spark.sql.execution.python
import java.io.DataOutputStream
+import org.apache.spark.SparkException
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions}
import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.sql.vectorized.ColumnarBatch
/**
@@ -36,6 +37,7 @@ class CometArrowPythonRunner(
evalType: Int,
argOffsets: Array[Array[Int]],
override val schema: StructType,
+ override val outputSchema: DataType,
override val workerConf: Map[String, String],
override val pythonMetrics: Map[String, SQLMetric],
jobArtifactUUID: Option[String],
@@ -51,4 +53,11 @@ class CometArrowPythonRunner(
override protected def writeUDF(dataOut: DataOutputStream): Unit =
PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, jobArtifactUUID)
+
+ // Spark 4.0 has no error condition for this; raise the message Spark 4.1 uses.
+ override protected def outputSchemaMismatchError(
+ operation: String,
+ expected: String,
+ actual: String): Throwable =
+ new SparkException(s"Invalid schema from $operation: expected $expected, got $actual.")
}
diff --git a/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
index 32aec97685..2cf8372647 100644
--- a/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
+++ b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
@@ -41,6 +41,7 @@ trait ShimCometMapInBatch extends Spark4xMapInBatchSupport {
evalType,
argOffsets,
schema,
+ runnerInputs.outputSchema,
runnerInputs.pythonRunnerConf,
pythonMetrics,
runnerInputs.jobArtifactUUID,
diff --git a/spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
index 658e559adc..8e7855632d 100644
--- a/spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
+++ b/spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
@@ -21,9 +21,10 @@ package org.apache.spark.sql.execution.python
import java.io.DataOutputStream
+import org.apache.spark.SparkException
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions}
import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.sql.vectorized.ColumnarBatch
/**
@@ -37,6 +38,7 @@ class CometArrowPythonRunner(
evalType: Int,
argOffsets: Array[Array[Int]],
override val schema: StructType,
+ override val outputSchema: DataType,
override val workerConf: Map[String, String],
override val pythonMetrics: Map[String, SQLMetric],
jobArtifactUUID: Option[String],
@@ -53,4 +55,14 @@ class CometArrowPythonRunner(
override protected def writeUDF(dataOut: DataOutputStream): Unit =
PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+
+ override protected def outputSchemaMismatchError(
+ operation: String,
+ expected: String,
+ actual: String): Throwable =
+ new SparkException(
+ errorClass = "ARROW_TYPE_MISMATCH",
+ messageParameters =
+ Map("operation" -> operation, "outputTypes" -> expected, "actualDataTypes" -> actual),
+ cause = null)
}
diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
index 32aec97685..2cf8372647 100644
--- a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
+++ b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala
@@ -41,6 +41,7 @@ trait ShimCometMapInBatch extends Spark4xMapInBatchSupport {
evalType,
argOffsets,
schema,
+ runnerInputs.outputSchema,
runnerInputs.pythonRunnerConf,
pythonMetrics,
runnerInputs.jobArtifactUUID,
diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
index 2ec17fd4cd..78cb53257e 100644
--- a/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
+++ b/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala
@@ -21,9 +21,10 @@ package org.apache.spark.sql.execution.python
import java.io.DataOutputStream
+import org.apache.spark.SparkException
import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions}
import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.sql.vectorized.ColumnarBatch
/**
@@ -36,6 +37,7 @@ class CometArrowPythonRunner(
evalType: Int,
argOffsets: Array[Array[Int]],
override val schema: StructType,
+ override val outputSchema: DataType,
pythonRunnerConf: Map[String, String],
override val pythonMetrics: Map[String, SQLMetric],
jobArtifactUUID: Option[String],
@@ -61,4 +63,14 @@ class CometArrowPythonRunner(
override protected def writeUDF(dataOut: DataOutputStream): Unit =
PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets)
+
+ override protected def outputSchemaMismatchError(
+ operation: String,
+ expected: String,
+ actual: String): Throwable =
+ new SparkException(
+ errorClass = "ARROW_TYPE_MISMATCH",
+ messageParameters =
+ Map("operation" -> operation, "outputTypes" -> expected, "actualDataTypes" -> actual),
+ cause = null)
}
diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala
index e096f87cd3..88931aedd7 100644
--- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala
@@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.PythonUDF
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.python.{ArrowPythonRunner, MapInArrowExec, MapInPandasExec}
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.DataType
/**
* Shared 4.x bits for `ShimCometMapInBatch`. The matchers and `runnerInputs` helper are identical
@@ -62,6 +63,7 @@ trait Spark4xMapInBatchSupport {
/** Inputs every 4.x `ArrowPythonRunner` constructor needs in the same shape. */
protected case class RunnerInputs(
chainedFunc: Seq[(ChainedPythonFunctions, Long)],
+ outputSchema: DataType,
pythonRunnerConf: Map[String, String],
jobArtifactUUID: Option[String],
arrowMaxRecordsPerBatch: Int,
@@ -75,6 +77,7 @@ trait Spark4xMapInBatchSupport {
protected def runnerInputs(pythonUDF: PythonUDF, conf: SQLConf): RunnerInputs =
RunnerInputs(
chainedFunc = Seq((ChainedPythonFunctions(Seq(pythonUDF.func)), pythonUDF.resultId.id)),
+ outputSchema = pythonUDF.dataType,
pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf),
jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid),
arrowMaxRecordsPerBatch = conf.arrowMaxRecordsPerBatch,
diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
index fecf781baa..cf2ac1ab0c 100644
--- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
+++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala
@@ -25,19 +25,21 @@ import java.util.ArrayList
import java.util.concurrent.atomic.AtomicBoolean
import scala.jdk.CollectionConverters._
+import scala.util.Try
import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, NullVector, VectorSchemaRoot, VectorUnloader}
import org.apache.arrow.vector.complex.StructVector
import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel}
import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch, MessageSerializer}
-import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.arrow.vector.types.TimeUnit
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
import org.apache.spark.{SparkEnv, TaskContext}
-import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.api.python.{BasePythonRunner, PythonEvalType, PythonRDD, PythonWorker, SpecialLengths}
import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, MapType, NullType, StringType, StructType, TimestampType, UserDefinedType}
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
import org.apache.comet.CometArrowAllocator
@@ -97,6 +99,21 @@ private[python] trait CometArrowPythonRunnerBase
*/
protected def schema: StructType
+ /**
+ * Declared output type of the UDF: a struct of the output columns, which the Python worker
+ * returns as a single struct column.
+ */
+ protected def outputSchema: DataType
+
+ /**
+ * Error for Python output whose Arrow schema cannot be read as [[outputSchema]]. Spark 4.1
+ * added the `ARROW_TYPE_MISMATCH` condition for this; Spark 4.0 has no equivalent.
+ */
+ protected def outputSchemaMismatchError(
+ operation: String,
+ expected: String,
+ actual: String): Throwable
+
/** Arrow input limits captured from the driver-side SQLConf. */
protected def arrowMaxRecordsPerBatch: Int
protected def arrowMaxBytesPerBatch: Long
@@ -307,6 +324,7 @@ private[python] trait CometArrowPythonRunnerBase
case SpecialLengths.START_ARROW_STREAM =>
reader = new ArrowStreamReader(stream, allocator)
root = reader.getVectorSchemaRoot()
+ checkOutputSchema(root.getSchema)
read()
case SpecialLengths.TIMING_DATA =>
handleTimingData()
@@ -323,6 +341,22 @@ private[python] trait CometArrowPythonRunnerBase
}
}
+ /**
+ * PySpark does not cast mapInArrow output to the declared schema, and Comet's vectors read
+ * buffers at the declared types, so check the output schema before any batch is read. Every
+ * batch in the Arrow stream has the schema read at its start.
+ */
+ private def checkOutputSchema(streamSchema: Schema): Unit = {
+ CometArrowPythonRunnerBase
+ .outputSchemaMismatch(outputSchema, streamSchema.getFields.asScala.toSeq)
+ .foreach { actual =>
+ throw outputSchemaMismatchError(
+ PythonEvalType.toString(evalType),
+ outputSchema.toString,
+ actual)
+ }
+ }
+
/**
* Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
* whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
@@ -619,6 +653,67 @@ private[python] object CometArrowPythonRunnerBase {
}
}
+ /**
+ * Checks the schema of a Python output stream, whose single struct column holds the output
+ * columns, against the declared output type. Returns a description of the stream schema when
+ * Comet's vectors cannot read it as the declared type, or None when they can.
+ */
+ private[python] def outputSchemaMismatch(
+ expected: DataType,
+ streamFields: Seq[Field]): Option[String] =
+ streamFields match {
+ case Seq(struct) if isReadableAs(expected, struct, nullable = false) => None
+ case _ => Some(describeOutput(expected, streamFields))
+ }
+
+ /**
+ * Whether Comet's vectors read `actual` as `expected`. The Arrow type must be the one Comet
+ * uses for the Spark type (see `Utils.toArrowType`), except for differences that do not change
+ * the values read, and that Spark's `ArrowColumnVector` reads the same way: field names and
+ * nullability, 64-bit string and binary offsets, the time zone of a timestamp, and an all-null
+ * `null` column where nulls are allowed. A dictionary-encoded column never matches; Spark
+ * cannot read one as a string either.
+ */
+ private def isReadableAs(expected: DataType, actual: Field, nullable: Boolean): Boolean = {
+ val children = actual.getChildren
+ actual.getDictionary == null && ((expected, actual.getType) match {
+ case (udt: UserDefinedType[_], _) => isReadableAs(udt.sqlType, actual, nullable)
+ case (_, ArrowType.Null.INSTANCE) => nullable || expected == NullType
+ case (StructType(fields), ArrowType.Struct.INSTANCE) =>
+ fields.length == children.size &&
+ fields.zip(children.asScala).forall { case (field, child) =>
+ isReadableAs(field.dataType, child, field.nullable)
+ }
+ case (ArrayType(elementType, containsNull), ArrowType.List.INSTANCE) =>
+ isReadableAs(elementType, children.get(0), containsNull)
+ case (MapType(keyType, valueType, valueContainsNull), _: ArrowType.Map) =>
+ val entry = children.get(0).getChildren
+ isReadableAs(keyType, entry.get(0), nullable = false) &&
+ isReadableAs(valueType, entry.get(1), valueContainsNull)
+ case (_: StringType, ArrowType.LargeUtf8.INSTANCE) => true
+ case (BinaryType, ArrowType.LargeBinary.INSTANCE) => true
+ case (TimestampType, timestamp: ArrowType.Timestamp) =>
+ timestamp.getUnit == TimeUnit.MICROSECOND && timestamp.getTimezone != null
+ case (dataType, arrowType) =>
+ Try(Utils.toArrowType(dataType, "UTC")).toOption.contains(arrowType)
+ })
+ }
+
+ /**
+ * Describes the stream schema with Spark type names, as Spark's error does, unless they cannot
+ * show the difference: a dictionary encoding, an Arrow type with no Spark equivalent, or a
+ * layout Comet does not read, such as a 256-bit decimal. Those show the Arrow schema instead.
+ */
+ private def describeOutput(expected: DataType, streamFields: Seq[Field]): String = {
+ val sparkTypes =
+ if (streamFields.exists(hasDictionary)) None
+ else Try(streamFields.map(Utils.fromArrowField).mkString(", ")).toOption
+ sparkTypes.filter(_ != expected.toString).getOrElse(streamFields.mkString(", "))
+ }
+
+ private def hasDictionary(field: Field): Boolean =
+ field.getDictionary != null || field.getChildren.asScala.exists(hasDictionary)
+
/**
* Serialize source vectors directly beneath the non-null struct advertised in the IPC stream.
*
diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf.py b/spark/src/test/resources/pyspark/test_pyarrow_udf.py
index 9a54300eed..683045743e 100644
--- a/spark/src/test/resources/pyspark/test_pyarrow_udf.py
+++ b/spark/src/test/resources/pyspark/test_pyarrow_udf.py
@@ -315,6 +315,150 @@ def test_map_in_arrow_python_exception_propagates(spark, tmp_path, accelerated):
)
+# PySpark does not cast mapInArrow output to the declared schema. The accelerated path checks
+# the returned Arrow schema before reading it and raises Spark's ARROW_TYPE_MISMATCH error
+# (a plain SparkException with the same message on Spark 4.0).
+_OUTPUT_TYPE_MISMATCH = "Invalid schema from SQL_MAP_ARROW_ITER_UDF"
+
+
+def _collect_error(df) -> str:
+ with pytest.raises(Exception) as exc_info:
+ df.collect()
+ return str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+ "returned_type,declared_type",
+ [(pa.int64(), T.IntegerType()), (pa.int32(), T.LongType())],
+ ids=["int64_as_int", "int32_as_bigint"],
+)
+def test_map_in_arrow_integer_width_mismatch_raises(
+ spark, tmp_path, accelerated, returned_type, declared_type
+):
+ """
+ Spark raises when it reads an integer column whose Arrow width differs from the declared
+ type. Reading the buffer at the declared width instead returns wrong values (int64 read as
+ int) or reads past the end of the buffer (int32 read as bigint).
+ """
+ src = str(tmp_path / "src.parquet")
+ spark.range(64).write.parquet(src)
+
+ def cast_id(iterator):
+ for batch in iterator:
+ yield pa.RecordBatch.from_arrays(
+ [batch.column(0).cast(returned_type)], ["id"]
+ )
+
+ schema = T.StructType([T.StructField("id", declared_type)])
+ result_df = spark.read.parquet(src).mapInArrow(cast_id, schema)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ message = _collect_error(result_df)
+ if accelerated:
+ assert _OUTPUT_TYPE_MISMATCH in message, message
+
+
+def test_map_in_arrow_decimal_scale_mismatch(spark, tmp_path, accelerated):
+ """
+ Spark rescales a returned decimal(4, 3) value to a declared decimal(10, 2), so 1.234 reads
+ as 1.23. The accelerated path does not rescale; it raises rather than reading the unscaled
+ value at the declared scale (12.34).
+ """
+ src = str(tmp_path / "src.parquet")
+ spark.range(1).write.parquet(src)
+
+ def to_decimal(iterator):
+ for batch in iterator:
+ amounts = pa.array([Decimal("1.234")] * batch.num_rows, pa.decimal128(4, 3))
+ yield pa.RecordBatch.from_arrays([amounts], ["amount"])
+
+ schema = T.StructType([T.StructField("amount", T.DecimalType(10, 2))])
+ result_df = spark.read.parquet(src).mapInArrow(to_decimal, schema)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ if accelerated:
+ message = _collect_error(result_df)
+ assert _OUTPUT_TYPE_MISMATCH in message, message
+ else:
+ assert [r["amount"] for r in result_df.collect()] == [Decimal("1.23")]
+
+
+def test_map_in_arrow_dictionary_encoded_output_raises(spark, tmp_path, accelerated):
+ """
+ Spark cannot read a dictionary-encoded column returned for a declared string column. The
+ accelerated path used to fail with a NullPointerException; it now reports the mismatch.
+ """
+ src = str(tmp_path / "src.parquet")
+ spark.range(10).write.parquet(src)
+
+ def encode(iterator):
+ for batch in iterator:
+ names = [f"name_{i % 3}" for i in batch.column(0).to_pylist()]
+ yield pa.RecordBatch.from_arrays(
+ [pa.array(names).dictionary_encode()], ["name"]
+ )
+
+ schema = T.StructType([T.StructField("name", T.StringType())])
+ result_df = spark.read.parquet(src).mapInArrow(encode, schema)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ message = _collect_error(result_df)
+ if accelerated:
+ assert _OUTPUT_TYPE_MISMATCH in message, message
+ assert "dictionary" in message, message
+
+
+def test_map_in_arrow_output_differences_that_read_the_same(
+ spark, tmp_path, accelerated
+):
+ """
+ Output that differs from the declared schema only in ways that do not change the values
+ read is accepted on both paths: column names, nullability, 64-bit string and binary offsets,
+ the time zone label of a timestamp, and an all-null column that PyArrow types as `null`.
+ """
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("ts", T.TimestampType()),
+ ]
+ )
+ rows = [(1, dt.datetime(2024, 1, 1, 12, 30, 45)), (2, None)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(rows, schema_in).write.parquet(src)
+
+ def relabel(iterator):
+ for batch in iterator:
+ ids = batch.column(0).to_pylist()
+ yield pa.RecordBatch.from_arrays(
+ [
+ batch.column(0),
+ batch.column(1).cast(pa.timestamp("us", tz="America/Los_Angeles")),
+ pa.array([f"s{i}" for i in ids], pa.large_string()),
+ pa.array([str(i).encode() for i in ids], pa.large_binary()),
+ pa.nulls(batch.num_rows),
+ ],
+ ["key", "ts", "s", "b", "missing"],
+ )
+
+ schema_out = T.StructType(
+ [
+ T.StructField("id", T.LongType(), nullable=False),
+ T.StructField("ts", T.TimestampType()),
+ T.StructField("s", T.StringType()),
+ T.StructField("b", T.BinaryType()),
+ T.StructField("missing", T.StringType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInArrow(relabel, schema_out)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = {
+ (r["id"], r["ts"], r["s"], bytes(r["b"]), r["missing"])
+ for r in result_df.collect()
+ }
+ assert out == {(i, ts, f"s{i}", str(i).encode(), None) for i, ts in rows}
+
+
def test_map_in_arrow_decimal_type(spark, tmp_path, accelerated):
schema_in = T.StructType(
[
diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
index f6974f02b2..e2f4c9e788 100644
--- a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
+++ b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala
@@ -19,11 +19,12 @@
package org.apache.spark.sql.execution.python
-import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataOutputStream, IOException}
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream, IOException}
import java.nio.ByteBuffer
import java.nio.channels.{Channels, WritableByteChannel}
import java.nio.charset.StandardCharsets
import java.nio.file.{Files, Paths}
+import java.util.concurrent.atomic.AtomicBoolean
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
@@ -38,13 +39,14 @@ import org.apache.arrow.vector.{BaseVariableWidthVector, BigIntVector, FieldVect
import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider}
import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel}
-import org.apache.arrow.vector.types.TimeUnit
+import org.apache.arrow.vector.types.{IntervalUnit, TimeUnit}
import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType, Schema}
-import org.apache.spark.{SparkConf, SparkEnv, TaskContext, TaskContextImpl}
-import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions, PythonEvalType, SimplePythonFunction}
+import org.apache.spark.{SparkConf, SparkEnv, SparkException, TaskContext, TaskContextImpl}
+import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions, PythonEvalType, SimplePythonFunction, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.{hasCompatibleSchema, inputBatchRanges, serializeBatch, withInputBatchRange, withMaterializedInputVectors}
-import org.apache.spark.sql.types.{StringType, StructField, StructType}
+import org.apache.spark.sql.execution.python.CometArrowPythonRunnerBase.{hasCompatibleSchema, inputBatchRanges, outputSchemaMismatch, serializeBatch, withInputBatchRange, withMaterializedInputVectors}
+import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
import org.apache.spark.util.DirectByteBufferOutputStream
@@ -71,14 +73,28 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
Map("pythonDataSent" -> new SQLMetric("size", 0L))
override protected val schema: StructType = StructType(
Seq(StructField("struct", StructType(Seq(StructField("text", StringType))))))
+ override protected val outputSchema: DataType =
+ StructType(Seq(StructField("id", IntegerType)))
override protected val arrowMaxRecordsPerBatch: Int = 1
override protected val arrowMaxBytesPerBatch: Long = Long.MaxValue
override protected def writeUDF(dataOut: DataOutputStream): Unit = ()
+ override protected def outputSchemaMismatchError(
+ operation: String,
+ expected: String,
+ actual: String): Throwable =
+ new SparkException(s"Invalid schema from $operation: expected $expected, got $actual.")
+
def inputWriter(input: Iterator[Iterator[ColumnarBatch]], context: TaskContext): Writer =
newWriter(SparkEnv.get, null, input, 0, context)
+ def outputReader(
+ stream: DataInputStream,
+ writer: Writer,
+ context: TaskContext): Iterator[ColumnarBatch] =
+ newReaderIterator(stream, writer, 0L, SparkEnv.get, null, None, new AtomicBoolean, context)
+
def bytesSent: Long = pythonMetrics("pythonDataSent").value
def transportBufferSize: Int = bufferSize
@@ -601,6 +617,173 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers {
}
}
+ private def arrowField(name: String, arrowType: ArrowType, children: Field*): Field =
+ new Field(name, FieldType.nullable(arrowType), children.asJava)
+
+ /** The Python worker returns the output columns in a single struct column. */
+ private def outputStream(columns: Field*): Seq[Field] =
+ Seq(arrowField("_0", ArrowType.Struct.INSTANCE, columns: _*))
+
+ private def declared(columns: (String, DataType)*): StructType =
+ StructType(columns.map { case (name, dataType) => StructField(name, dataType) })
+
+ private def arrowMap(key: Field, value: Field): Field =
+ arrowField(
+ "c",
+ new ArrowType.Map(false),
+ new Field(
+ "entries",
+ FieldType.notNullable(ArrowType.Struct.INSTANCE),
+ Seq(key, value).asJava))
+
+ private val int32 = new ArrowType.Int(32, true)
+ private val int64 = new ArrowType.Int(64, true)
+
+ test("output schema check rejects integer widths and decimal scales other than declared") {
+ outputSchemaMismatch(declared("id" -> IntegerType), outputStream(arrowField("id", int32)))
+ .shouldBe(None)
+ outputSchemaMismatch(declared("id" -> IntegerType), outputStream(arrowField("id", int64)))
+ .shouldBe(Some("StructType(StructField(id,LongType,true))"))
+ outputSchemaMismatch(declared("id" -> LongType), outputStream(arrowField("id", int32)))
+ .shouldBe(Some("StructType(StructField(id,IntegerType,true))"))
+
+ val amount = declared("amount" -> DecimalType(10, 2))
+ def decimal(precision: Int, scale: Int, bitWidth: Int): Seq[Field] =
+ outputStream(arrowField("amount", new ArrowType.Decimal(precision, scale, bitWidth)))
+ outputSchemaMismatch(amount, decimal(10, 2, 128)) shouldBe None
+ outputSchemaMismatch(amount, decimal(4, 3, 128)) shouldBe
+ Some("StructType(StructField(amount,DecimalType(4,3),true))")
+ // Both map to the declared Spark type, so the description shows the Arrow layout instead.
+ outputSchemaMismatch(amount, decimal(10, 2, 256)).get should include("Decimal(10, 2, 256)")
+ }
+
+ test("output schema check rejects layouts that Comet vectors would misread") {
+ val encoding = new DictionaryEncoding(7L, false, int32)
+ val dictionary = new Field("name", new FieldType(true, int32, encoding), Seq.empty.asJava)
+ outputSchemaMismatch(declared("name" -> StringType), outputStream(dictionary)).get should
+ include("name: Int(32, true)[dictionary: 7]")
+
+ val struct = StructType(Seq(StructField("a", IntegerType)))
+ Seq(
+ BinaryType -> arrowField("c", new ArrowType.FixedSizeBinary(4)),
+ DayTimeIntervalType() -> arrowField("c", new ArrowType.Interval(IntervalUnit.DAY_TIME)),
+ TimestampType -> arrowField("c", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")),
+ TimestampType -> arrowField("c", new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)),
+ TimestampNTZType -> arrowField("c", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")),
+ IntegerType -> arrowField("c", new ArrowType.Int(32, false)),
+ ArrayType(IntegerType) -> arrowField(
+ "c",
+ ArrowType.LargeList.INSTANCE,
+ arrowField("e", int32)),
+ ArrayType(IntegerType) -> arrowField("c", ArrowType.List.INSTANCE, arrowField("e", int64)),
+ MapType(StringType, IntegerType) ->
+ arrowMap(arrowField("k", ArrowType.Utf8.INSTANCE), arrowField("v", int64)),
+ struct -> arrowField("c", ArrowType.Struct.INSTANCE, arrowField("a", int64)),
+ struct -> arrowField(
+ "c",
+ ArrowType.Struct.INSTANCE,
+ arrowField("a", int32),
+ arrowField("b", int32)),
+ // An all-null column is only readable where the declared type allows nulls.
+ MapType(IntegerType, StringType) ->
+ arrowMap(
+ arrowField("k", ArrowType.Null.INSTANCE),
+ arrowField("v", ArrowType.Utf8.INSTANCE)))
+ .foreach { case (dataType, field) =>
+ withClue(s"$dataType read from $field: ") {
+ outputSchemaMismatch(declared("c" -> dataType), outputStream(field)) shouldBe defined
+ }
+ }
+
+ val nonNullable = StructType(Seq(StructField("c", StringType, nullable = false)))
+ outputSchemaMismatch(nonNullable, outputStream(arrowField("c", ArrowType.Null.INSTANCE)))
+ .shouldBe(defined)
+ // Missing or extra columns, or a stream that is not a single struct column.
+ val column = arrowField("a", int32)
+ outputSchemaMismatch(declared("a" -> IntegerType, "b" -> IntegerType), outputStream(column))
+ .shouldBe(defined)
+ outputSchemaMismatch(declared("a" -> IntegerType), outputStream(column, column))
+ .shouldBe(defined)
+ outputSchemaMismatch(declared("a" -> IntegerType), Seq(column)) shouldBe defined
+ outputSchemaMismatch(
+ declared("a" -> IntegerType),
+ outputStream(column) ++ outputStream(column))
+ .shouldBe(defined)
+ }
+
+ test("output schema check accepts differences that do not change the values read") {
+ // Comet's own Arrow layout for each type is accepted, whatever the field names and
+ // nullability.
+ val schema = StructType(
+ Seq(
+ StructField("bool", BooleanType),
+ StructField("byte", ByteType),
+ StructField("short", ShortType),
+ StructField("int", IntegerType),
+ StructField("long", LongType),
+ StructField("float", FloatType),
+ StructField("double", DoubleType),
+ StructField("string", StringType),
+ StructField("binary", BinaryType),
+ StructField("date", DateType),
+ StructField("timestamp", TimestampType),
+ StructField("timestamp_ntz", TimestampNTZType),
+ StructField("decimal", DecimalType(38, 18)),
+ StructField("year_month", YearMonthIntervalType()),
+ StructField("day_time", DayTimeIntervalType()),
+ StructField("null", NullType),
+ StructField("array", ArrayType(StructType(Seq(StructField("x", IntegerType))))),
+ StructField("map", MapType(StringType, ArrayType(LongType)))))
+ val renamed =
+ StructType(schema.fields.map(f => f.copy(name = s"${f.name}_out", nullable = false)))
+ val stream = Utils.toArrowField("_0", renamed, nullable = true, "UTC")
+ outputSchemaMismatch(schema, Seq(stream)) shouldBe None
+
+ val longBacked = new UserDefinedType[java.lang.Long] {
+ override def sqlType: DataType = LongType
+ override def serialize(value: java.lang.Long): Any = value
+ override def deserialize(datum: Any): java.lang.Long = datum.asInstanceOf[java.lang.Long]
+ override def userClass: Class[java.lang.Long] = classOf[java.lang.Long]
+ }
+ Seq(
+ StringType -> arrowField("c", ArrowType.LargeUtf8.INSTANCE),
+ StringType("UTF8_LCASE") -> arrowField("c", ArrowType.Utf8.INSTANCE),
+ BinaryType -> arrowField("c", ArrowType.LargeBinary.INSTANCE),
+ TimestampType ->
+ arrowField("c", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "America/Los_Angeles")),
+ DayTimeIntervalType(DayTimeIntervalType.HOUR, DayTimeIntervalType.MINUTE) ->
+ arrowField("c", new ArrowType.Duration(TimeUnit.MICROSECOND)),
+ YearMonthIntervalType(YearMonthIntervalType.YEAR) ->
+ arrowField("c", new ArrowType.Interval(IntervalUnit.YEAR_MONTH)),
+ longBacked -> arrowField("c", int64),
+ StringType -> arrowField("c", ArrowType.Null.INSTANCE),
+ ArrayType(StringType) ->
+ arrowField("c", ArrowType.List.INSTANCE, arrowField("e", ArrowType.Null.INSTANCE)))
+ .foreach { case (dataType, field) =>
+ withClue(s"$dataType read from $field: ") {
+ outputSchemaMismatch(declared("c" -> dataType), outputStream(field)) shouldBe None
+ }
+ }
+ }
+
+ test("output reader checks the stream schema before reading any batch") {
+ withInputWriterRunner { (runner, context) =>
+ val bytes = new ByteArrayOutputStream()
+ new DataOutputStream(bytes).writeInt(SpecialLengths.START_ARROW_STREAM)
+ Using.resource(new RootAllocator(Long.MaxValue)) { allocator =>
+ withWriter(Seq(arrowField("id", int64)), allocator, Channels.newChannel(bytes))(_ => ())
+ }
+ val output = runner.outputReader(
+ new DataInputStream(new ByteArrayInputStream(bytes.toByteArray)),
+ runner.inputWriter(Iterator.empty, context),
+ context)
+ intercept[SparkException](output.hasNext).getMessage shouldBe
+ "Invalid schema from SQL_MAP_ARROW_ITER_UDF: " +
+ "expected StructType(StructField(id,IntegerType,true)), " +
+ "got StructType(StructField(id,LongType,true))."
+ }
+ }
+
test("direct batches retain borrowed buffers without copying them into the writer allocator") {
val sourceAllocator = new RootAllocator(Long.MaxValue)
val writerAllocator = new RootAllocator(1024)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org