diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index 97d3470facb..eff37b0a1ee 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` 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 ed4dc84c9cc..a96069b7dc3 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 5b5d43dbe8a..387785dfc5b 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")) + 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 a355d143d1b..cfdefedd6e6 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)