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/01 04:33:58 UTC
(datafusion-comet) 01/01: feat: admit string maps in Spark-to-Comet conversion (#6036)
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-6036-cb0b584bb9c603e960e5dba394caf36798d47f92
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 92b43cfaaea476ad8c146f9a252c640594921766
Author: Chao Sun <su...@apache.org>
AuthorDate: Thu Oct 1 04:33:48 2026 +0000
feat: admit string maps in Spark-to-Comet conversion (#6036)
Co-authored-by: Jeff Wang <je...@openai.com>
---
docs/source/user-guide/latest/datasources.md | 16 +-
.../spark/sql/comet/CometSparkToColumnarExec.scala | 1 +
.../org/apache/comet/exec/CometExecSuite.scala | 101 ++++++++++-
.../execution/arrow/CometArrowStreamSuite.scala | 184 ++++++++++++++++++++-
4 files changed, 285 insertions(+), 17 deletions(-)
diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md
index 197b31f7df..c6d9debaf0 100644
--- a/docs/source/user-guide/latest/datasources.md
+++ b/docs/source/user-guide/latest/datasources.md
@@ -32,7 +32,7 @@ Arrow format, allowing the Comet pipeline to take over after that, but the proce
Comet accelerates Iceberg scans of Parquet files and has an experimental, opt-in native Iceberg writer.
See the [Iceberg Guide] and [Iceberg Writes](iceberg-writes.md) for more information.
-[Iceberg Guide]: iceberg.md
+[iceberg guide]: iceberg.md
### CSV
@@ -51,11 +51,17 @@ converted into Arrow format, allowing the Comet pipeline to take over after that
### Spark-to-Comet conversion types
-Spark-to-Comet conversion supports `ARRAY<STRING>` with binary string semantics, including
-nullable arrays and nullable elements, both as top-level fields and inside supported structs.
+Spark-to-Comet conversion supports `ARRAY<STRING>` and `MAP<STRING,STRING>` with binary
+string semantics, both as top-level fields and inside supported structs. Arrays and maps may
+be null; array elements and map values may also be null. Map keys must be non-null.
This applies to Spark row and columnar inputs when conversion is enabled for the source.
-Other array element types, nested arrays, arrays of structs, maps, and non-binary string
-collations remain unsupported at this conversion boundary. Source defaults are unchanged.
+Other array element types, other map key/value types, nested collections, and non-binary
+string collations remain unsupported at this conversion boundary. Source defaults are unchanged.
+
+This includes row-backed `ExistingRDD` inputs when
+`spark.comet.sparkToColumnar.enabled=true` and
+`spark.comet.sparkToColumnar.supportedOperatorList` includes `RDDScan`. Spark still produces
+the RDD rows; conversion lets eligible downstream operators execute in Comet.
## Data Catalogs
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
index 6f91248e78..c13a595078 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
@@ -142,6 +142,7 @@ object CometSparkToColumnarExec extends CometSink[SparkPlan] with DataTypeSuppor
name: String,
fallbackReasons: ListBuffer[String]): Boolean = dt match {
case ArrayType(StringType, _) => true
+ case MapType(StringType, StringType, _) => true
case _: ArrayType | _: MapType => false
case _ => super.isTypeSupported(dt, name, fallbackReasons)
}
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
index 5bdecfd3a9..3e7cf698f5 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -3844,21 +3844,34 @@ class CometExecSuite extends CometTestBase {
})
}
- test("SparkToColumnar admits only binary string arrays through the array gate") {
+ test("SparkToColumnar admits only binary string arrays and maps through the collection gate") {
for (nullable <- Seq(false, true); containsNull <- Seq(false, true)) {
- val field = StructField("tags", ArrayType(StringType, containsNull), nullable)
- Seq(StructType(Seq(field)), StructType(Seq(StructField("nested", StructType(Seq(field))))))
- .foreach { schema =>
- assert(CometSparkToColumnarExec.isSchemaSupported(schema, ListBuffer.empty))
- }
+ for (dataType <- Seq(
+ ArrayType(StringType, containsNull),
+ MapType(StringType, StringType, containsNull))) {
+ val field = StructField("tags", dataType, nullable)
+ Seq(
+ StructType(Seq(field)),
+ StructType(Seq(StructField("nested", StructType(Seq(field))))))
+ .foreach { schema =>
+ assert(CometSparkToColumnarExec.isSchemaSupported(schema, ListBuffer.empty))
+ }
+ }
}
val unsupported = Seq(
ArrayType(IntegerType),
ArrayType(BinaryType),
ArrayType(ArrayType(StringType)),
ArrayType(StructType(Seq(StructField("s", StringType)))),
- MapType(StringType, StringType)) ++
- (if (isSpark40Plus) Seq(DataType.fromDDL("ARRAY<STRING COLLATE UTF8_LCASE>"))
+ MapType(StringType, IntegerType),
+ MapType(IntegerType, StringType),
+ MapType(StringType, ArrayType(StringType)),
+ MapType(StringType, StructType(Seq(StructField("s", StringType))))) ++
+ (if (isSpark40Plus)
+ Seq(
+ DataType.fromDDL("ARRAY<STRING COLLATE UTF8_LCASE>"),
+ DataType.fromDDL("MAP<STRING COLLATE UTF8_LCASE,STRING>"),
+ DataType.fromDDL("MAP<STRING,STRING COLLATE UTF8_LCASE>"))
else Seq.empty)
unsupported.foreach { dataType =>
val schema = StructType(Seq(StructField("value", dataType)))
@@ -3936,6 +3949,78 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("SparkToColumnar string maps cross RDD and Parquet native boundaries") {
+ val schema = new StructType()
+ .add("id", IntegerType)
+ .add("partitionValues", MapType(StringType, StringType))
+ .add("nested", new StructType().add("tags", MapType(StringType, StringType)))
+ val rows = Seq(
+ Row(1, null, null),
+ Row(2, Map.empty[String, String], Row(null)),
+ Row(3, Map("hour" -> null), Row(Map.empty[String, String])),
+ Row(
+ 4,
+ Map("hour" -> "2026-09-02T22", "" -> "é", "東京" -> "a\u0000b"),
+ Row(Map("nullable" -> null, "x" -> "value"))),
+ Row(5, Map("hour" -> "2026-09-03T00", "long" -> ("東京" * 32768)), null),
+ Row(6, Map("hour" -> "2026-09-03T01"), Row(Map("x" -> "value"))))
+ for (sourceType <- Seq("rdd", "parquet-row", "parquet-columnar")) {
+ val vectorized = sourceType == "parquet-columnar"
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> vectorized.toString,
+ "spark.sql.parquet.enableNestedColumnVectorizedReader" -> "true",
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false",
+ CometConf.COMET_BATCH_SIZE.key -> "2",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true",
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true",
+ CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "RDDScan") {
+ withTempPath { dir =>
+ def rdd = spark.createDataFrame(spark.sparkContext.parallelize(rows, 1), schema)
+ if (sourceType != "rdd") {
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ rdd.write.parquet(dir.toString)
+ }
+ }
+ def source =
+ if (sourceType == "rdd") rdd else spark.read.schema(schema).parquet(dir.toString)
+ def query = source
+ .filter("id > 0")
+ .selectExpr("id", "partitionValues", "nested", "partitionValues['hour'] AS hour")
+ val (_, plan) = checkSparkAnswerAndOperator(
+ query,
+ includeClasses = Seq(
+ classOf[CometSparkToColumnarExec],
+ classOf[CometProjectExec],
+ classOf[CometFilterExec]))
+ val conversions = collect(plan) { case c: CometSparkToColumnarExec => c }
+ assert(conversions.size == 1)
+ assert(conversions.head.child.supportsColumnar == vectorized)
+ if (sourceType == "rdd") {
+ assert(conversions.head.child.collect { case scan: RDDScanExec => scan }.nonEmpty)
+ }
+ checkSparkSchema(query)
+ checkSparkAnswerAndOperator(
+ source.filter("partitionValues['hour'] BETWEEN '2026-09-02T22' AND '2026-09-03T00'"),
+ includeClasses = Seq(classOf[CometSparkToColumnarExec], classOf[CometFilterExec]))
+ val (_, shuffled) = checkSparkAnswerAndOperator(
+ query.repartition(2, col("id")),
+ includeClasses = Seq(classOf[CometShuffleExchangeExec]))
+ val exchanges = collect(shuffled) { case s: CometShuffleExchangeExec => s }
+ assert(exchanges.nonEmpty && exchanges.forall(_.shuffleType == CometNativeShuffle))
+ checkSparkAnswer(query.limit(1))
+ withSQLConf(
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false",
+ CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "false") {
+ val (_, disabled) = checkSparkAnswer(query)
+ assert(collect(disabled) { case c: CometSparkToColumnarExec => c }.isEmpty)
+ }
+ }
+ }
+ }
+ }
+
test("SparkToColumnar over BatchScan (Spark Parquet reader)") {
Seq("", "parquet").foreach { v1List =>
Seq(true, false).foreach { parquetVectorized =>
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala
index f2ff9e9790..21c6c4150e 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala
@@ -31,17 +31,17 @@ import org.scalatest.matchers.should.Matchers
import org.apache.arrow.memory.{AllocationListener, RootAllocator}
import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVector, BitVector, DecimalVector, IntervalMonthDayNanoVector, IntVector, VarCharVector, VectorLoader, VectorSchemaRoot, VectorUnloader}
-import org.apache.arrow.vector.complex.ListVector
+import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
import org.apache.arrow.vector.dictionary.{Dictionary => ArrowDictionary}
import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider
import org.apache.arrow.vector.ipc.ArrowReader
import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType, Schema}
import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters}
-import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData}
import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector}
-import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType}
-import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch, ColumnVector}
+import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType}
+import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch, ColumnarMap, ColumnVector}
import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil}
@@ -599,6 +599,182 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers {
}
}
+ for (nullable <- Seq(false, true); valueContainsNull <- Seq(false, true);
+ columnar <- Seq(false, true)) {
+ test(s"string maps preserve slices and ownership: $nullable/$valueContainsNull/$columnar") {
+ val mapType = MapType(StringType, StringType, valueContainsNull)
+ val schema = StructType(Seq(StructField("tags", mapType, nullable)))
+ val arrowSchema = Utils.toArrowSchema(schema, "UTC")
+ val values: Seq[Seq[(String, String)]] = Seq(
+ Seq("" -> "", "é" -> "東京", "a\u0000b" -> "duplicate", "b" -> "duplicate"),
+ Seq.empty,
+ null,
+ null,
+ Seq.empty,
+ Seq("nullable" -> (if (valueContainsNull) null else "value")),
+ Seq.tabulate(4097)(i => s"key-$i" -> (if (i == 1) "東京" * 32768 else s"value-$i")),
+ Seq.empty,
+ Seq("last" -> "value")) ++ (if (nullable) Seq.fill(5)(null) else Seq.empty)
+ val expected = if (nullable) values else values.filter(_ != null)
+ val allocator = new RootAllocator(Long.MaxValue)
+ val maps = new OnHeapColumnVector(expected.size, mapType)
+ val keys = maps.getChild(0)
+ val vals = maps.getChild(1)
+ val capacity = expected.filter(_ != null).map(_.size).sum + 2
+ keys.reserve(capacity)
+ vals.reserve(capacity)
+ var elementOffset = 2 // Exercise a nonzero offset in both child vectors.
+ expected.zipWithIndex.foreach { case (entries, row) =>
+ if (entries == null) maps.putNull(row)
+ else {
+ maps.putArray(row, elementOffset, entries.size)
+ entries.foreach { case (key, value) =>
+ keys.putByteArray(elementOffset, key.getBytes(StandardCharsets.UTF_8))
+ if (value == null) vals.putNull(elementOffset)
+ else vals.putByteArray(elementOffset, value.getBytes(StandardCharsets.UTF_8))
+ elementOffset += 1
+ }
+ }
+ }
+ val input = new ColumnarBatch(Array[ColumnVector](maps), expected.size)
+ val empty = new ColumnarBatch(Array[ColumnVector](maps), 0)
+ val reusedRow = new GenericInternalRow(1)
+ val backing = ArrayBuffer.empty[Array[Byte]]
+ def borrowedString(text: String): UTF8String = {
+ if (text == null) null
+ else {
+ val bytes = ("!" + text + "!").getBytes(StandardCharsets.UTF_8)
+ backing += bytes
+ UTF8String.fromBytes(bytes, 1, bytes.length - 2)
+ }
+ }
+ val rows = expected.iterator.map { entries =>
+ backing.foreach(bytes => java.util.Arrays.fill(bytes, 0.toByte))
+ backing.clear()
+ reusedRow.update(
+ 0,
+ if (entries == null) null
+ else
+ new ArrayBasedMapData(
+ new GenericArrayData(entries.map(e => borrowedString(e._1)).toArray[Any]),
+ new GenericArrayData(entries.map(e => borrowedString(e._2)).toArray[Any])))
+ reusedRow
+ }
+ val reader: ArrowReader = if (columnar) {
+ new SparkColumnarArrowReader(allocator, arrowSchema, Iterator(empty, input, empty), 2)
+ } else {
+ new RowArrowReader(allocator, arrowSchema, rows, 2)
+ }
+ val retained = ArrayBuffer.empty[VectorSchemaRoot]
+ try {
+ while (reader.loadNextBatch()) {
+ val batch = new VectorUnloader(reader.getVectorSchemaRoot).getRecordBatch
+ val output = VectorSchemaRoot.create(arrowSchema, allocator)
+ retained += output
+ try new VectorLoader(output).load(batch)
+ finally batch.close()
+ }
+ reader.close()
+ backing.foreach(bytes => java.util.Arrays.fill(bytes, 0.toByte))
+ Seq(keys, vals).foreach { child =>
+ val bytes = child.getChild(0)
+ (0 until bytes.getElementsAppended).foreach(bytes.putByte(_, 0.toByte))
+ }
+ var start = 0
+ retained.foreach { root =>
+ val chunk = expected.slice(start, start + 2)
+ root.getRowCount shouldBe chunk.size
+ root.getSchema shouldBe arrowSchema
+ val field = root.getSchema.getFields.get(0)
+ field.isNullable shouldBe nullable
+ Utils.fromArrowField(field) shouldBe mapType
+ val entriesField = field.getChildren.get(0)
+ entriesField.isNullable shouldBe false
+ entriesField.getChildren.get(0).isNullable shouldBe false
+ entriesField.getChildren.get(1).isNullable shouldBe valueContainsNull
+ val output = root.getVector(0).asInstanceOf[MapVector]
+ val entries = output.getDataVector.asInstanceOf[StructVector]
+ val outKeys = entries.getChild(MapVector.KEY_NAME).asInstanceOf[VarCharVector]
+ val outVals = entries.getChild(MapVector.VALUE_NAME).asInstanceOf[VarCharVector]
+ val childCount = chunk.filter(_ != null).map(_.size).sum
+ entries.getValueCount shouldBe childCount
+ outKeys.getValueCount shouldBe childCount
+ outVals.getValueCount shouldBe childCount
+ var offset = 0
+ chunk.zipWithIndex.foreach { case (value, row) =>
+ output.isNull(row) shouldBe (value == null)
+ output.getOffsetBuffer.getInt(row * 4L) shouldBe offset
+ if (value != null) value.foreach { case (key, text) =>
+ entries.isNull(offset) shouldBe false
+ outKeys.isNull(offset) shouldBe false
+ new String(outKeys.get(offset), StandardCharsets.UTF_8) shouldBe key
+ outVals.isNull(offset) shouldBe (text == null)
+ if (text != null)
+ new String(outVals.get(offset), StandardCharsets.UTF_8) shouldBe text
+ offset += 1
+ }
+ output.getOffsetBuffer.getInt((row + 1) * 4L) shouldBe offset
+ }
+ start += chunk.size
+ }
+ start shouldBe expected.size
+ retained.foreach(_.close())
+ retained.clear()
+ maps.getMap(0).numElements() shouldBe 4 // Closing output must not close Spark input.
+ } finally {
+ reader.close()
+ retained.foreach(_.close())
+ input.close()
+ allocator.close()
+ }
+ }
+ }
+
+ test("string maps release reader allocations when a key or value cannot be encoded") {
+ for (failInKey <- Seq(false, true); columnar <- Seq(false, true)) {
+ val allocator = new RootAllocator(Long.MaxValue)
+ val mapType = MapType(StringType, StringType, valueContainsNull = true)
+ val schema = StructType(Seq(StructField("tags", mapType)))
+ val failure = new IllegalStateException("map entry encoding failed")
+ def strings(failHere: Boolean) = new ConstantColumnVector(2, StringType) {
+ override def getUTF8String(row: Int): UTF8String = {
+ if (failHere && row == 1) throw failure
+ UTF8String.fromString(s"value-$row")
+ }
+ }
+ val keys = strings(failInKey)
+ val values = strings(!failInKey)
+ val map = new ColumnarMap(keys, values, 0, 2)
+ val maps = new ConstantColumnVector(1, mapType) {
+ override def getMap(row: Int): ColumnarMap = map
+ }
+ val input = new ColumnarBatch(Array[ColumnVector](maps), 1)
+ val arrowSchema = Utils.toArrowSchema(schema, "UTC")
+ val reader: ArrowReader = if (columnar) {
+ new SparkColumnarArrowReader(allocator, arrowSchema, Iterator.single(input), 2)
+ } else {
+ new RowArrowReader(
+ allocator,
+ arrowSchema,
+ Iterator(new GenericInternalRow(Array[Any](map))),
+ 2)
+ }
+ try {
+ try {
+ intercept[IllegalStateException](
+ reader.loadNextBatch()) should be theSameInstanceAs failure
+ } finally reader.close()
+ allocator.getAllocatedMemory shouldBe 0L
+ } finally {
+ reader.close()
+ input.close()
+ keys.close()
+ values.close()
+ allocator.close()
+ }
+ }
+ }
+
test("Spark columnar reader preserves split input slice offsets") {
val allocator = new RootAllocator(Long.MaxValue)
val numRows = 12
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org