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/09/18 19:03:57 UTC
(datafusion-comet) 01/01: feat: support string arrays in Spark-to-Comet conversion (#5954)
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-5954-c56cef77a973af89255553b0ac08a98385409e67
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit d528af7c2f52e121cfe79e0d32e583739dfd171c
Author: KUAN-HAO HUANG <10...@users.noreply.github.com>
AuthorDate: Fri Sep 18 19:03:33 2026 +0000
feat: support string arrays in Spark-to-Comet conversion (#5954)
---
docs/source/user-guide/latest/datasources.md | 8 +
.../spark/sql/comet/CometSparkToColumnarExec.scala | 1 +
.../org/apache/comet/exec/CometExecSuite.scala | 96 +++++++++++-
.../execution/arrow/CometArrowStreamSuite.scala | 166 ++++++++++++++++++++-
4 files changed, 268 insertions(+), 3 deletions(-)
diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md
index 97d3470fac..eff37b0a1e 100644
--- a/docs/source/user-guide/latest/datasources.md
+++ b/docs/source/user-guide/latest/datasources.md
@@ -47,6 +47,14 @@ converted into Arrow format, allowing the Comet pipeline to take over after that
Comet does not provide a Rust-based JSON scan, but when `spark.comet.convert.json.enabled` is enabled, data is immediately
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.
+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.
+
## Data Catalogs
### Apache Iceberg
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 ed4dc84c9c..a96069b7dc 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
@@ -141,6 +141,7 @@ object CometSparkToColumnarExec extends CometSink[SparkPlan] with DataTypeSuppor
dt: DataType,
name: String,
fallbackReasons: ListBuffer[String]): Boolean = dt match {
+ case ArrayType(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 5b5d43dbe8..387785dfc5 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -22,6 +22,7 @@ package org.apache.comet.exec
import java.sql.Date
import java.time.{Duration, Period}
+import scala.collection.mutable.ListBuffer
import scala.util.Random
import org.scalactic.source.Position
@@ -34,7 +35,7 @@ import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, Cat
import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, ExpressionInfo, Hex, Literal}
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate}
import org.apache.spark.sql.comet._
-import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec}
@@ -47,6 +48,7 @@ import org.apache.spark.sql.execution.window.WindowExec
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE
+import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
import org.apache.comet.{CometConf, CometExecIterator, ExtendedExplainInfo}
@@ -3727,6 +3729,98 @@ class CometExecSuite extends CometTestBase {
})
}
+ test("SparkToColumnar admits only binary string arrays through the array 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))
+ }
+ }
+ 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>"))
+ else Seq.empty)
+ unsupported.foreach { dataType =>
+ val schema = StructType(Seq(StructField("value", dataType)))
+ assert(!CometSparkToColumnarExec.isSchemaSupported(schema, ListBuffer.empty), dataType)
+ assert(
+ !CometSparkToColumnarExec.isSchemaSupported(
+ StructType(Seq(StructField("nested", schema))),
+ ListBuffer.empty),
+ dataType)
+ }
+ }
+
+ test("SparkToColumnar string arrays cross JSON and Parquet native boundaries") {
+ val schema = new StructType()
+ .add("id", IntegerType)
+ .add("tags", ArrayType(StringType))
+ .add("nested", new StructType().add("tags", ArrayType(StringType)))
+ val rows = Seq(
+ Row(1, null, null),
+ Row(2, Seq.empty[String], Row(null)),
+ Row(3, Seq(null), Row(Seq.empty[String])),
+ Row(4, Seq("", "é", "東京", "a\u0000b", "dup", "dup"), Row(Seq(null, "x"))),
+ Row(5, Seq("東京" * 32768), Row(Seq("long"))),
+ Row(6, null, Row(null)),
+ Row(7, Seq.empty[String], null))
+ for (format <- Seq("json", "parquet"); v1 <- Seq("", format);
+ vectorized <- (if (format == "parquet") Seq(false, true) else Seq(false))) {
+ val convertKey =
+ if (format == "json") CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key
+ else CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> v1,
+ 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",
+ convertKey -> "true") {
+ withTempPath { dir =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .createDataFrame(spark.sparkContext.parallelize(rows, 1), schema)
+ .write
+ .format(format)
+ .save(dir.toString)
+ }
+ def source = spark.read.schema(schema).format(format).load(dir.toString)
+ def query = source
+ .filter("id > 0")
+ .selectExpr("id", "tags", "nested", "size(tags) AS tag_count")
+ 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)
+ checkSparkSchema(query)
+ 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))
+ // Stop before draining the input to exercise normal task-completion cleanup.
+ checkSparkAnswer(query.limit(1))
+ withSQLConf(convertKey -> "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 a355d143d1..cfdefedd6e 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
@@ -22,6 +22,7 @@ package org.apache.spark.sql.comet.execution.arrow
import java.nio.ByteOrder
import java.nio.charset.StandardCharsets
+import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import scala.util.Using
@@ -29,16 +30,19 @@ import org.scalatest.funsuite.AnyFunSuite
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, VectorSchemaRoot}
+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.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.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.unsafe.types.CalendarInterval
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil}
@@ -347,6 +351,164 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers {
}
}
+ for (nullable <- Seq(false, true); containsNull <- Seq(false, true);
+ columnar <- Seq(false, true)) {
+ test(
+ "string arrays preserve slices and ownership: " +
+ s"nullable=$nullable, containsNull=$containsNull, columnar=$columnar") {
+ val arrayType = ArrayType(StringType, containsNull)
+ val schema = StructType(Seq(StructField("tags", arrayType, nullable)))
+ val arrowSchema = Utils.toArrowSchema(schema, "UTC")
+ val longString = "東京" * 32768
+ val many = Seq.tabulate(4097)(i => if (i == 1) longString else s"value-$i")
+ val values = Seq(
+ Seq("", "é", "a\u0000b", "dup", "dup"),
+ Seq.empty,
+ null,
+ null,
+ Seq.empty,
+ if (containsNull) Seq(null, "x", null) else Seq("x"),
+ many,
+ Seq.empty,
+ Seq("last")) ++ (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 arrays = new OnHeapColumnVector(expected.size, arrayType)
+ val elements = arrays.getChild(0)
+ elements.reserve(expected.filter(_ != null).map(_.size).sum + 2)
+ var elementOffset = 2 // The source child vector need not start at offset zero.
+ expected.zipWithIndex.foreach { case (value, row) =>
+ if (value == null) arrays.putNull(row)
+ else {
+ arrays.putArray(row, elementOffset, value.size)
+ value.foreach { text =>
+ if (text == null) elements.putNull(elementOffset)
+ else elements.putByteArray(elementOffset, text.getBytes(StandardCharsets.UTF_8))
+ elementOffset += 1
+ }
+ }
+ }
+ val input = new ColumnarBatch(Array[ColumnVector](arrays), expected.size)
+ val empty = new ColumnarBatch(Array[ColumnVector](arrays), 0)
+ val reusedRow = new GenericInternalRow(1)
+ val backing = ArrayBuffer.empty[Array[Byte]]
+ val rows = expected.iterator.map { value =>
+ backing.foreach(bytes => java.util.Arrays.fill(bytes, 0.toByte))
+ backing.clear()
+ reusedRow.update(
+ 0,
+ if (value == null) null
+ else
+ new GenericArrayData(
+ value
+ .map {
+ case null => null
+ case text =>
+ val bytes = ("!" + text + "!").getBytes(StandardCharsets.UTF_8)
+ backing += bytes
+ UTF8String.fromBytes(bytes, 1, bytes.length - 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()) {
+ // Retain the actual buffers, as an exported batch does, rather than copying values.
+ 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))
+ val sourceBytes = elements.getChild(0)
+ (0 until sourceBytes.getElementsAppended).foreach(sourceBytes.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
+ field.getChildren.get(0).isNullable shouldBe containsNull
+ Utils.fromArrowField(field) shouldBe arrayType
+ val lists = root.getVector(0).asInstanceOf[ListVector]
+ val strings = lists.getDataVector.asInstanceOf[VarCharVector]
+ val childCount = chunk.filter(_ != null).map(_.size).sum
+ lists.getValueCount shouldBe chunk.size
+ strings.getValueCount shouldBe childCount
+ var offset = 0
+ chunk.zipWithIndex.foreach { case (value, row) =>
+ lists.isNull(row) shouldBe (value == null)
+ lists.getOffsetBuffer.getInt(row * 4L) shouldBe offset
+ if (value != null) value.foreach { text =>
+ strings.isNull(offset) shouldBe (text == null)
+ if (text != null)
+ new String(strings.get(offset), StandardCharsets.UTF_8) shouldBe text
+ offset += 1
+ }
+ lists.getOffsetBuffer.getInt((row + 1) * 4L) shouldBe offset
+ }
+ start += chunk.size
+ }
+ start shouldBe expected.size
+ retained.foreach(_.close())
+ retained.clear()
+ // Reader/output cleanup must not close borrowed Spark columns.
+ arrays
+ .getArray(expected.indexWhere(v => v != null && v.nonEmpty))
+ .numElements() shouldBe 5
+ } finally {
+ reader.close()
+ retained.foreach(_.close())
+ input.close()
+ allocator.close()
+ }
+ }
+ }
+
+ test("string arrays release allocations when an element cannot be encoded") {
+ val allocator = new RootAllocator(Long.MaxValue)
+ val arrayType = ArrayType(StringType, containsNull = true)
+ val schema = StructType(Seq(StructField("tags", arrayType)))
+ val failure = new IllegalStateException("string element encoding failed")
+ val child = new ConstantColumnVector(2, StringType) {
+ override def getUTF8String(row: Int): UTF8String = {
+ if (row == 1) throw failure
+ UTF8String.fromString("first")
+ }
+ }
+ val arrays = new ConstantColumnVector(1, arrayType) {
+ override def getArray(row: Int): ColumnarArray = new ColumnarArray(child, 0, 2)
+ }
+ val input = new ColumnarBatch(Array[ColumnVector](arrays), 1)
+ try {
+ intercept[IllegalStateException] {
+ CometArrowConverters.columnarBatchToArrowBatch(
+ input,
+ Utils.toArrowSchema(schema, "UTC"),
+ allocator)
+ } should be theSameInstanceAs failure
+ allocator.getAllocatedMemory shouldBe 0L
+ val rows = Iterator(new GenericInternalRow(Array[Any](new ColumnarArray(child, 0, 2))))
+ intercept[IllegalStateException] {
+ CometArrowConverters.rowToArrowBatchIter(rows, schema, 2, "UTC", allocator).next()
+ } should be theSameInstanceAs failure
+ allocator.getAllocatedMemory shouldBe 0L
+ } finally {
+ input.close()
+ child.close()
+ allocator.close()
+ }
+ }
+
test("nested and foreign-vector fallback preserves slices and independently owned batches") {
val allocator = new RootAllocator(Long.MaxValue)
val arrayType = ArrayType(IntegerType, containsNull = true)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org