feat: support string arrays in Spark-to-Comet conversion

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

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)
   }

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 =>

spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala

@@ -22,23 +22,27 @@ 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
 
 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)