diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 2ba71e93c4d..1d18e5f0aa5 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -624,8 +624,8 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | Function | Status | Implementation | Notes | | --- | --- | --- | --- | -| `named_struct` | ✅ | Native | Duplicate field names fall back | -| `struct` | ✅ | Native | | +| `named_struct` | ✅ | Hybrid | Duplicate field names route through the JVM codegen dispatcher | +| `struct` | ✅ | Hybrid | Duplicate field names route through the JVM codegen dispatcher | --- diff --git a/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java b/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java index 94c4916f9f1..158cc4cbe1d 100644 --- a/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java +++ b/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java @@ -19,6 +19,8 @@ package org.apache.arrow.c; +import java.util.function.Function; + import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.FieldVector; @@ -53,11 +55,19 @@ Field importField(ArrowSchema schema, CDataDictionaryProvider provider) { public FieldVector importVector( ArrowArray array, ArrowSchema schema, CDataDictionaryProvider provider) { + return importVector(array, schema, provider, field -> field.createVector(allocator)); + } + + public FieldVector importVector( + ArrowArray array, + ArrowSchema schema, + CDataDictionaryProvider provider, + Function vectorFactory) { Field field = null; FieldVector vector = null; try { field = importField(schema, provider); - vector = field.createVector(allocator); + vector = vectorFactory.apply(field); ArrayImporter importer = new ArrayImporter(allocator, vector, provider); importer.importArray(array); return vector; diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index d8dea731353..fb0500c1c9f 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -20,18 +20,22 @@ package org.apache.comet.udf; import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Function; import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowImporter; import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.spark.TaskContext; import org.apache.spark.comet.CometTaskContextShim; import org.apache.spark.util.TaskCompletionListener; import org.apache.comet.util.ClassLoaders; +import org.apache.comet.vector.NativeUtil$; /** * JNI entry point for native execution to invoke a {@link CometUDF}. Matches the static-method @@ -210,6 +214,9 @@ private static void evaluateInternal( assert udf != null : "reflective instantiation returned null for " + udfClassName; BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator(); + ArrowImporter importer = new ArrowImporter(allocator); + Function vectorFactory = + field -> NativeUtil$.MODULE$.createVectorForImport(field, allocator); ValueVector[] inputs = new ValueVector[inputArrayPtrs.length]; ValueVector result = null; @@ -217,7 +224,7 @@ private static void evaluateInternal( for (int i = 0; i < inputArrayPtrs.length; i++) { ArrowArray inArr = ArrowArray.wrap(inputArrayPtrs[i]); ArrowSchema inSch = ArrowSchema.wrap(inputSchemaPtrs[i]); - inputs[i] = Data.importVector(allocator, inArr, inSch, null); + inputs[i] = importer.importVector(inArr, inSch, null, vectorFactory); } result = udf.evaluate(inputs, numRows); diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 33e6c0c0355..371c7116d84 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -22,7 +22,6 @@ package org.apache.comet.codegen import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.types.pojo.{ArrowType, Field} @@ -32,6 +31,7 @@ import org.apache.spark.sql.types._ import org.apache.comet.CometArrowAllocator import org.apache.comet.shims.CometTypeShim +import org.apache.comet.vector.NativeUtil /** * Output-side emitters for the codegen kernel: [[allocateOutput]], [[emitOutputWriter]] @@ -43,8 +43,8 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { /** * Spark `DataType` to an Arrow `Field` with names Comet expects on FFI export. Spark's * `Utils.toArrowField` names list children `"element"`; this rewrites them to `"item"`. Pair - * with the [[RenamedListVector]] / [[RenamedMapVector]] / [[RenamedStructVector]] subclasses in - * [[allocateOutput]], which pin `getField()` so the cached Field actually reaches export. + * with [[NativeUtil.createVector]], whose complex-vector wrappers pin `getField()` so the + * cached Field actually reaches export. */ def toFfiArrowField(name: String, dataType: DataType, nullable: Boolean): Field = renameForArrowRustFfi(Utils.toArrowField(name, dataType, nullable, "UTC")) @@ -71,8 +71,8 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { * Allocate an Arrow output vector from a pre-built `Field`. Callers cache the Field per * `(expression, schema)` and pass it on every batch. * - * Complex top-level types route through a [[RenamedListVector]] / [[RenamedMapVector]] / - * [[RenamedStructVector]] (see those for the runtime-vs-export naming gap). + * Complex top-level types route through [[NativeUtil.createVector]] to bridge runtime child + * names and the exported schema. * * `estimatedBytes` pre-sizes the data buffer for variable-length scalar outputs. Ignored for * other root types, and not propagated into nested var-width children (their `allocateNew` runs @@ -87,22 +87,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { * Closes the vector on any failure so a partially-initialized tree doesn't leak buffers. */ def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int): FieldVector = { - val vec: FieldVector = field.getType match { - case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList => - val v = new RenamedListVector(field, CometArrowAllocator) - v.initializeChildrenFromFields(field.getChildren) - v - case _: ArrowType.Map => - val v = new RenamedMapVector(field, CometArrowAllocator) - v.initializeChildrenFromFields(field.getChildren) - v - case _: ArrowType.Struct => - val v = new RenamedStructVector(field, CometArrowAllocator) - v.initializeChildrenFromFields(field.getChildren) - v - case _ => - field.createVector(CometArrowAllocator).asInstanceOf[FieldVector] - } + val vec = NativeUtil.createVector(field, CometArrowAllocator) try { vec.setInitialCapacity(numRows) vec match { @@ -122,27 +107,6 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { } } - /** - * Pin `getField()` to the cached Field so FFI export carries the names Comet expects. - * `ListVector.getField` rebuilds child labels from the runtime data vector, which - * `addOrGetVector` hardcodes to `"$data$"`. Applied to `MapVector` and `StructVector` too - * because their `getField` recurses and can pick up a buried `ListVector`'s `"$data$"`. - */ - private final class RenamedListVector(exportField: Field, allocator: BufferAllocator) - extends ListVector(exportField, allocator, null) { - override def getField: Field = exportField - } - - private final class RenamedMapVector(exportField: Field, allocator: BufferAllocator) - extends MapVector(exportField, allocator, null) { - override def getField: Field = exportField - } - - private final class RenamedStructVector(exportField: Field, allocator: BufferAllocator) - extends StructVector(exportField, allocator, null) { - override def getField: Field = exportField - } - /** * Returns `(concreteVectorClassName, batchSetup, perRowSnippet)`. `output` is cast to the * concrete class in `process`'s prelude so `emitWrite`'s complex-type branches can hoist child diff --git a/spark/src/main/scala/org/apache/comet/serde/structs.scala b/spark/src/main/scala/org/apache/comet/serde/structs.scala index 2f9619d491f..b5dc8da44c9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/structs.scala +++ b/spark/src/main/scala/org/apache/comet/serde/structs.scala @@ -31,12 +31,16 @@ import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.DataTypeSupport import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} -object CometCreateNamedStruct extends CometExpressionSerde[CreateNamedStruct] { +object CometCreateNamedStruct + extends CometExpressionSerde[CreateNamedStruct] + with CodegenDispatchFallback { private val duplicateNamesReason = - "`CreateNamedStruct` with duplicate field names is not supported" + "`CreateNamedStruct` with duplicate field names cannot use native execution" - override def getUnsupportedReasons(): Seq[String] = Seq(duplicateNamesReason) + override def getUnsupportedReasons(): Seq[String] = Seq( + "Duplicate field names are routed through the JVM codegen dispatcher " + + "(Spark's own `doGenCode`).") override def getSupportLevel(expr: CreateNamedStruct): SupportLevel = { if (expr.names.length != expr.names.distinct.length) { diff --git a/spark/src/main/scala/org/apache/comet/vector/CometArrowStreamReader.scala b/spark/src/main/scala/org/apache/comet/vector/CometArrowStreamReader.scala new file mode 100644 index 00000000000..0aef3d62c2f --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/vector/CometArrowStreamReader.scala @@ -0,0 +1,186 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.vector + +import java.nio.channels.ReadableByteChannel +import java.util + +import scala.collection.JavaConverters._ + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.util.AutoCloseables +import org.apache.arrow.vector.{FieldVector, VectorLoader, VectorSchemaRoot} +import org.apache.arrow.vector.compression.CompressionCodec +import org.apache.arrow.vector.dictionary.Dictionary +import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel} +import org.apache.arrow.vector.ipc.message.{ArrowDictionaryBatch, ArrowRecordBatch, MessageChannelReader} +import org.apache.arrow.vector.types.pojo.{Field, Schema} +import org.apache.arrow.vector.util.{DictionaryUtility, VectorBatchAppender} + +/** + * Arrow IPC reader that keeps struct children positional when a schema contains duplicate names. + * + * ArrowReader normally allocates each field with `Field.createVector`, which indexes direct + * struct children by name and collapses duplicates. Reuse NativeUtil's import factory here so IPC + * and C Data imports have the same physical layout and the ordinary no-duplicate path stays + * unchanged. + */ +final class CometArrowStreamReader( + messageReader: MessageChannelReader, + allocator: BufferAllocator, + compressionFactory: CompressionCodec.Factory) + extends ArrowStreamReader(messageReader, allocator, compressionFactory) { + + def this(messageReader: MessageChannelReader, allocator: BufferAllocator) = + this(messageReader, allocator, CompressionCodec.Factory.INSTANCE) + + def this(channel: ReadableByteChannel, allocator: BufferAllocator) = + this( + new MessageChannelReader(new ReadChannel(channel), allocator), + allocator, + CompressionCodec.Factory.INSTANCE) + + private var cometInitialized = false + private var cometResourcesClosed = false + private var cometSourceClosed = false + private var cometRoot: VectorSchemaRoot = _ + private var cometLoader: VectorLoader = _ + + override protected def initialize(): Unit = { + val originalSchema = readSchema() + val fields = new util.ArrayList[Field](originalSchema.getFields.size()) + val vectors = new util.ArrayList[FieldVector](originalSchema.getFields.size()) + val importedDictionaries = new util.HashMap[java.lang.Long, Dictionary]() + + try { + originalSchema.getFields.asScala.foreach { field => + val updated = DictionaryUtility.toMemoryFormat(field, allocator, importedDictionaries) + fields.add(updated) + vectors.add(NativeUtil.createVectorForImport(updated, allocator)) + } + cometRoot = + new VectorSchemaRoot(new Schema(fields, originalSchema.getCustomMetadata), vectors, 0) + cometLoader = new VectorLoader(cometRoot, compressionFactory) + dictionaries = util.Collections.unmodifiableMap(importedDictionaries) + cometInitialized = true + } catch { + case failure: Throwable => + AutoCloseables.close(failure, vectors) + AutoCloseables.close( + failure, + importedDictionaries.values().asScala.map(_.getVector).asJava) + cometRoot = null + cometLoader = null + throw failure + } + } + + override protected def ensureInitialized(): Unit = { + if (!cometInitialized) initialize() + } + + override def getVectorSchemaRoot: VectorSchemaRoot = { + ensureInitialized() + cometRoot + } + + override def getDictionaryVectors: util.Map[java.lang.Long, Dictionary] = { + ensureInitialized() + dictionaries + } + + override def lookup(id: Long): Dictionary = { + if (!cometInitialized) { + throw new IllegalStateException("Unable to lookup until reader has been initialized") + } + dictionaries.get(id) + } + + override def getDictionaryIds: util.Set[java.lang.Long] = { + if (!cometInitialized) { + throw new IllegalStateException( + "Unable to list dictionaries until reader has been initialized") + } + dictionaries.keySet() + } + + override protected def prepareLoadNextBatch(): Unit = { + ensureInitialized() + cometRoot.setRowCount(0) + } + + override protected def loadRecordBatch(batch: ArrowRecordBatch): Unit = { + try cometLoader.load(batch) + finally batch.close() + } + + override protected def loadDictionary(dictionaryBatch: ArrowDictionaryBatch): Unit = { + val dictionary = dictionaries.get(dictionaryBatch.getDictionaryId) + if (dictionary == null) { + throw new IllegalArgumentException( + s"Dictionary ID ${dictionaryBatch.getDictionaryId} not defined in schema") + } + + val vector = dictionary.getVector + if (dictionaryBatch.isDelta) { + val deltaVector = NativeUtil.createVectorForImport(vector.getField, allocator) + try { + loadDictionaryBatch(dictionaryBatch, deltaVector) + VectorBatchAppender.batchAppend(vector, deltaVector) + } finally { + deltaVector.close() + } + } else { + loadDictionaryBatch(dictionaryBatch, vector) + } + } + + private def loadDictionaryBatch( + dictionaryBatch: ArrowDictionaryBatch, + vector: FieldVector): Unit = { + val root = new VectorSchemaRoot( + util.Collections.singletonList(vector.getField), + util.Collections.singletonList(vector), + 0) + val loader = new VectorLoader(root, compressionFactory) + try loader.load(dictionaryBatch.getDictionary) + finally dictionaryBatch.close() + } + + override def close(): Unit = close(closeReadSource = true) + + override def close(closeReadSource: Boolean): Unit = { + val resources = new util.ArrayList[AutoCloseable]() + if (!cometResourcesClosed) { + cometResourcesClosed = true + if (cometRoot != null) resources.add(cometRoot) + if (dictionaries != null) { + dictionaries.values().asScala.foreach(dictionary => resources.add(dictionary.getVector)) + } + } + if (closeReadSource && !cometSourceClosed) { + cometSourceClosed = true + resources.add(new AutoCloseable { + override def close(): Unit = CometArrowStreamReader.super.closeReadSource() + }) + } + AutoCloseables.close(resources) + } +} diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 173086d2fd7..dcb1ac7a162 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -19,12 +19,18 @@ package org.apache.comet.vector +import java.util.{ArrayList, HashSet} +import java.util.function.Function + import scala.collection.mutable import org.apache.arrow.c.{ArrowArray, ArrowImporter, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.util.AutoCloseables -import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.{AbstractStructVector, ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.DictionaryProvider +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} import org.apache.spark.SparkException import org.apache.spark.sql.comet.execution.arrow.ConstantColumnVectors import org.apache.spark.sql.comet.util.Utils @@ -53,6 +59,10 @@ class NativeUtil extends AutoCloseable { /** ArrowImporter does not hold any state and does not need to be closed */ private val importer = new ArrowImporter(allocator) + /** Reuse one factory on the per-batch import hot path. */ + private val importVectorFactory: Function[Field, FieldVector] = + field => NativeUtil.createVectorForImport(field, allocator) + /** * Dictionary provider to use for the lifetime of this instance of NativeUtil. The dictionary * provider is closed when NativeUtil is closed. @@ -268,7 +278,8 @@ class NativeUtil extends AutoCloseable { // importField's finally consumes the schema. ArrayImporter takes the array normally, while // ArrowImporter's catch releases it if the import fails before that transfer. firstUnconsumed = i + 1 - val arrowVector = importer.importVector(arrowArray, arrowSchema, dictionaryProvider) + val arrowVector = + importer.importVector(arrowArray, arrowSchema, dictionaryProvider, importVectorFactory) val cometVector = try CometVector.getVector(arrowVector, dictionaryProvider) catch { @@ -322,6 +333,192 @@ class NativeUtil extends AutoCloseable { } object NativeUtil { + + /** + * Create a vector whose physical struct children remain positional when the exported Arrow + * schema contains duplicate names. Arrow's default struct factory indexes children by name and + * collapses such fields. + */ + private[comet] def createVector(field: Field, allocator: BufferAllocator): FieldVector = { + val runtimeField = fieldForAllocation(field) + createPinnedVector(runtimeField, field, allocator) + } + + /** + * Preserve Arrow's default allocation path unless a duplicate-name struct needs positional + * runtime children. This is called for every imported column of every native batch. + */ + private[comet] def createVectorForImport( + field: Field, + allocator: BufferAllocator): FieldVector = { + val runtimeField = fieldForAllocation(field) + if (runtimeField eq field) { + field.createVector(allocator).asInstanceOf[FieldVector] + } else { + createPinnedVector(runtimeField, field, allocator) + } + } + + /** Build an IPC root with the same duplicate-safe allocation used by C Data imports. */ + def createVectorSchemaRootForImport( + schema: Schema, + allocator: BufferAllocator): VectorSchemaRoot = { + val fields = schema.getFields + val vectors = new ArrayList[FieldVector](fields.size()) + try { + var ordinal = 0 + while (ordinal < fields.size()) { + vectors.add(createVectorForImport(fields.get(ordinal), allocator)) + ordinal += 1 + } + new VectorSchemaRoot(schema, vectors, 0) + } catch { + case failure: Throwable => + AutoCloseables.close(failure, vectors) + throw failure + } + } + + /** + * Build a C Stream root whose physical and advertised schemas use the same duplicate-safe field + * names. Arrow's C Data exporter reconstructs nested vectors from the advertised schema and + * otherwise collapses duplicate struct children before loading the record batch. + */ + def createVectorSchemaRootForExport( + schema: Schema, + allocator: BufferAllocator): VectorSchemaRoot = { + val fields = schema.getFields + val runtimeFields = new ArrayList[Field](fields.size()) + val vectors = new ArrayList[FieldVector](fields.size()) + try { + var ordinal = 0 + while (ordinal < fields.size()) { + val runtimeField = fieldForAllocation(fields.get(ordinal)) + runtimeFields.add(runtimeField) + vectors.add(runtimeField.createVector(allocator).asInstanceOf[FieldVector]) + ordinal += 1 + } + new VectorSchemaRoot(new Schema(runtimeFields), vectors, 0) + } catch { + case failure: Throwable => + AutoCloseables.close(failure, vectors) + throw failure + } + } + + private def createPinnedVector( + runtimeField: Field, + exportField: Field, + allocator: BufferAllocator): FieldVector = { + exportField.getType match { + case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList => + val vector = new RenamedListVector(runtimeField, exportField, allocator) + vector.initializeChildrenFromFields(runtimeField.getChildren) + vector + case _: ArrowType.Map => + val vector = new RenamedMapVector(runtimeField, exportField, allocator) + vector.initializeChildrenFromFields(runtimeField.getChildren) + vector + case _: ArrowType.Struct => + val vector = new RenamedStructVector(runtimeField, exportField, allocator) + // Arrow's Field-based StructVector constructor creates children through a writer whose + // cache lower-cases field names. Build the direct children positionally instead so case- + // distinct names such as `a` and `A` remain separate physical vectors. + vector.initializeChildrenFromFields(runtimeField.getChildren) + vector + case _ => exportField.createVector(allocator).asInstanceOf[FieldVector] + } + } + + private def fieldForAllocation(field: Field): Field = { + val children = field.getChildren + if (children.isEmpty) return field + + val names = field.getType match { + case _: ArrowType.Struct if children.size() > 1 => new HashSet[String](children.size()) + case _ => null + } + + var hasDuplicateNames = false + var runtimeChildren: ArrayList[Field] = null + var ordinal = 0 + while (ordinal < children.size()) { + val child = children.get(ordinal) + val runtimeChild = fieldForAllocation(child) + + if ((runtimeChild ne child) && runtimeChildren == null) { + runtimeChildren = new ArrayList[Field](children.size()) + var priorOrdinal = 0 + while (priorOrdinal < ordinal) { + runtimeChildren.add(children.get(priorOrdinal)) + priorOrdinal += 1 + } + } + if (runtimeChildren != null) runtimeChildren.add(runtimeChild) + + if (names != null && !names.add(child.getName)) hasDuplicateNames = true + ordinal += 1 + } + + val childrenForAllocation = if (runtimeChildren == null) children else runtimeChildren + if (hasDuplicateNames) { + val renamedChildren = new ArrayList[Field](children.size()) + ordinal = 0 + while (ordinal < childrenForAllocation.size()) { + val child = childrenForAllocation.get(ordinal) + var runtimeName = s"__comet_runtime_field_$ordinal" + while (names.contains(runtimeName)) runtimeName = s"_$runtimeName" + names.add(runtimeName) + renamedChildren.add(new Field(runtimeName, child.getFieldType, child.getChildren)) + ordinal += 1 + } + new Field(field.getName, field.getFieldType, renamedChildren) + } else if (runtimeChildren != null) { + new Field(field.getName, field.getFieldType, runtimeChildren) + } else { + field + } + } + + /** + * Pin `getField()` to the imported Field so FFI keeps the original child labels. ListVector's + * runtime data-vector label is `"$data$"`; struct runtime names may be private and unique. + */ + private final class RenamedListVector( + runtimeField: Field, + exportField: Field, + allocator: BufferAllocator) + extends ListVector(runtimeField, allocator, null) { + override def getField: Field = exportField + } + + private final class RenamedMapVector( + runtimeField: Field, + exportField: Field, + allocator: BufferAllocator) + extends MapVector(runtimeField, allocator, null) { + override def getField: Field = exportField + } + + private final class RenamedStructVector( + runtimeField: Field, + exportField: Field, + allocator: BufferAllocator) + extends StructVector( + runtimeField.getName, + allocator, + runtimeField.getFieldType, + null, + AbstractStructVector.ConflictPolicy.CONFLICT_ERROR, + true) { + override def getField: Field = { + // StructVector's writer calls getField during construction. Keep the superclass's in-progress + // field visible until every positional child exists, then publish the original metadata. The + // child count avoids a separate construction-state flag. + if (size() == exportField.getChildren.size()) exportField else super.getField + } + } + def rootAsBatch(arrowRoot: VectorSchemaRoot): ColumnarBatch = { rootAsBatch(arrowRoot, null) } diff --git a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala index 805eae988e8..4f83338cf23 100644 --- a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala +++ b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala @@ -24,7 +24,7 @@ import java.nio.channels.ReadableByteChannel import scala.util.control.NonFatal import org.apache.arrow.vector.VectorSchemaRoot -import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel} +import org.apache.arrow.vector.ipc.ReadChannel import org.apache.arrow.vector.ipc.message.MessageChannelReader import org.apache.spark.sql.vectorized.ColumnarBatch @@ -36,7 +36,7 @@ import org.apache.comet.CometArrowAllocator case class StreamReader(channel: ReadableByteChannel, source: String) extends AutoCloseable { private val channelReader = new MessageChannelReader(new ReadChannel(channel), CometArrowAllocator) - private var arrowReader = new ArrowStreamReader(channelReader, CometArrowAllocator) + private var arrowReader = new CometArrowStreamReader(channelReader, CometArrowAllocator) // Reading the schema allocates the root's vectors, so it can fail with buffers already taken. // No caller holds this reader until its constructor returns, so close it here or nothing will. @@ -64,7 +64,6 @@ case class StreamReader(channel: ReadableByteChannel, source: String) extends Au override def close(): Unit = { if (root != null) { arrowReader.close() - root.close() arrowReader = null root = null diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala index 379b2cfb3ab..43e7456128d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala @@ -19,18 +19,20 @@ package org.apache.spark.sql.comet.execution.arrow +import java.util import java.util.{ArrayList => JArrayList} import scala.collection.mutable.ListBuffer import org.apache.arrow.memory.BufferAllocator -import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot, VectorUnloader} -import org.apache.arrow.vector.dictionary.DictionaryEncoder +import org.apache.arrow.vector.{FieldVector, VectorLoader, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryEncoder} import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch import org.apache.arrow.vector.types.pojo.Schema import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.vector.{CometDictionaryVector, CometVector} +import org.apache.comet.vector.{CometDictionaryVector, CometVector, NativeUtil} /** * `ArrowReader` over an iterator of Arrow-backed `ColumnarBatch`es. The unload/load step @@ -44,12 +46,68 @@ private[comet] class ColumnarBatchArrowReader( source: Iterator[ColumnarBatch]) extends ArrowReader(allocator) { + private var cometInitialized = false + private var cometClosed = false + private var cometRoot: VectorSchemaRoot = _ + private var cometLoader: VectorLoader = _ + override protected def readSchema(): Schema = arrowSchema + override protected def initialize(): Unit = { + cometRoot = NativeUtil.createVectorSchemaRootForExport(readSchema(), allocator) + cometLoader = new VectorLoader(cometRoot) + cometInitialized = true + } + + override protected def ensureInitialized(): Unit = { + if (!cometInitialized) initialize() + } + + override def getVectorSchemaRoot: VectorSchemaRoot = { + ensureInitialized() + cometRoot + } + + override def getDictionaryVectors: util.Map[java.lang.Long, Dictionary] = { + ensureInitialized() + util.Collections.emptyMap() + } + + override def lookup(id: Long): Dictionary = { + if (!cometInitialized) { + throw new IllegalStateException("Unable to lookup until reader has been initialized") + } + null + } + + override def getDictionaryIds: util.Set[java.lang.Long] = { + ensureInitialized() + util.Collections.emptySet() + } + + override protected def prepareLoadNextBatch(): Unit = { + ensureInitialized() + cometRoot.setRowCount(0) + } + + override protected def loadRecordBatch(batch: ArrowRecordBatch): Unit = { + try cometLoader.load(batch) + finally batch.close() + } + override def bytesRead(): Long = 0L override protected def closeReadSource(): Unit = () + override def close(): Unit = close(closeReadSource = true) + + override def close(closeReadSource: Boolean): Unit = { + if (!cometClosed) { + cometClosed = true + if (cometRoot != null) cometRoot.close() + } + } + override def loadNextBatch(): Boolean = { prepareLoadNextBatch() diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index f07848b2b38..7131844bef9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -30,7 +30,7 @@ import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider -import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} +import org.apache.arrow.vector.ipc.ArrowStreamWriter import org.apache.arrow.vector.types._ import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.arrow.vector.util.VectorSchemaRootAppender @@ -46,7 +46,7 @@ import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStrea import org.apache.comet.Constants.COMET_CONF_DIR_ENV import org.apache.comet.shims.CometTypeShim -import org.apache.comet.vector.CometVector +import org.apache.comet.vector.{CometArrowStreamReader, CometVector, NativeUtil} object Utils extends CometTypeShim with Logging { def getConfPath(confFileName: String): String = { @@ -377,7 +377,7 @@ object Utils extends CometTypeShim with Logging { val compressedInputStream = new DataInputStream(codec.compressedInputStream(bytes.toInputStream())) val reader = - new ArrowStreamReader(Channels.newChannel(compressedInputStream), allocator) + new CometArrowStreamReader(Channels.newChannel(compressedInputStream), allocator) try { // Comet decodes dictionaries during execution, so this shouldn't happen. // If it does, fall back to the original uncoalesced buffers because each @@ -397,7 +397,8 @@ object Utils extends CometTypeShim with Logging { while (reader.loadNextBatch()) { val sourceRoot = reader.getVectorSchemaRoot if (targetRoot == null) { - targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator) + targetRoot = + NativeUtil.createVectorSchemaRootForImport(sourceRoot.getSchema, allocator) targetRoot.allocateNew() } try { diff --git a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql index a3ad834c3ff..565148a588a 100644 --- a/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql +++ b/spark/src/test/resources/sql-tests/expressions/struct/create_named_struct.sql @@ -33,3 +33,33 @@ SELECT named_struct('x', 1, 'y', 'hello', 'z', 3.14) query SELECT named_struct('x', a, 'y', 'fixed_val', 'z', c) FROM test_named_struct + +-- duplicate names dispatch through Spark codegen while preserving ordinal values +query +SELECT named_struct('x', a, 'x', b) FROM test_named_struct + +-- struct() lowers to CreateNamedStruct and derives duplicate names from repeated children +query +SELECT struct(a, a) FROM test_named_struct + +-- nested duplicate-name structs exercise list and map roots during Arrow import +query +SELECT array(named_struct('x', a, 'x', b)) FROM test_named_struct + +query +SELECT map('row', named_struct('x', a, 'x', b)) FROM test_named_struct + +-- nested structs, three duplicates, and an all-null row +query +SELECT named_struct('outer', named_struct('x', a, 'x', b)) FROM test_named_struct + +query +SELECT named_struct('x', a, 'x', b, 'x', c) FROM test_named_struct + +query +SELECT named_struct('x', a, 'x', b, 'x', c) FROM test_named_struct WHERE a IS NULL + +-- construct the duplicate-name struct after a supported primitive-key shuffle boundary +query +SELECT named_struct('x', a, 'x', b) +FROM (SELECT /*+ REPARTITION(2, a) */ a, b FROM test_named_struct) shuffled diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5806cb35015..e196f9e8adb 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -111,6 +111,28 @@ class CometCodegenSuite } } + test("codegen kernel preserves duplicate struct field names positionally") { + val expr = + CreateNamedStruct(Seq(Literal("a"), Literal(true), Literal("a"), Literal(7.toByte))) + + val field = CometBatchKernelCodegen.toFfiArrowField("out", expr.dataType, nullable = true) + val output = CometBatchKernelCodegen.allocateOutput(field, 1, 0) + try { + assert(output.getChildrenFromFields.get(0).isInstanceOf[BitVector]) + assert(output.getChildrenFromFields.get(1).isInstanceOf[TinyIntVector]) + assert(output.getField.getChildren.get(0).getName === "a") + assert(output.getField.getChildren.get(1).getName === "a") + } finally { + output.close() + } + + val actual = runKernel(expr, 1) { vector => + val row = vector.getStruct(0) + row.getBoolean(0) -> row.getByte(1) + } + assert(actual === (true -> 7.toByte)) + } + test("ScalaUDF over concat(c1, c2) suppresses the null short-circuit") { // Concat is not NullIntolerant. The dispatcher's short-circuit guard inspects every node in // the bound tree and must skip the whole-tree null short-circuit because one child is diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 25b43d9fe6e..d0185a28b49 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -2474,26 +2474,6 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("named_struct with duplicate field names") { - Seq(true, false).foreach { dictionaryEnabled => - withTempDir { dir => - val path = new Path(dir.toURI.toString, "test.parquet") - makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = dictionaryEnabled, 10000) - withParquetTable(path.toString, "tbl") { - checkSparkAnswerAndOperator( - "SELECT named_struct('a', _1, 'a', _2) FROM tbl", - classOf[ProjectExec]) - checkSparkAnswerAndOperator( - "SELECT named_struct('a', _1, 'a', 2) FROM tbl", - classOf[ProjectExec]) - checkSparkAnswerAndOperator( - "SELECT named_struct('a', named_struct('b', _1, 'b', _2)) FROM tbl", - classOf[ProjectExec]) - } - } - } - } - test("to_json") { withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToJson]) -> "true") { Seq(true, false).foreach { dictionaryEnabled => diff --git a/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala b/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala index 4a4df2e9b1a..6b7bab1969e 100644 --- a/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala @@ -62,9 +62,29 @@ class CometJsonJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("from_json preserves case-distinct fields via JVM engine") { + withJsonTable { + checkSparkAnswerAndOperator(sql("SELECT from_json(j, 'a INT, A INT') FROM t")) + } + } + test("to_json round-trip via JVM engine") { withJsonTable { checkSparkAnswerAndOperator(sql("SELECT to_json(from_json(j, 'a INT, b STRING')) FROM t")) } } + + test("to_json preserves a retained duplicate-name producer") { + withTable("t") { + sql("CREATE TABLE t (a INT, b INT) USING parquet") + sql("INSERT INTO t VALUES (1, 10), (2, 20), (3, 30)") + withSQLConf("spark.sql.optimizer.collapseProjectAlwaysInline" -> "false") { + checkSparkAnswerAndOperator(sql("""SELECT to_json(s), s + |FROM ( + | SELECT named_struct('x', a, 'x', b) AS s + | FROM t + |) q""".stripMargin)) + } + } + } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 7f114aa160e..9eb0a215d6c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -732,6 +732,32 @@ class CometJoinSuite extends CometTestBase { } } + test("Broadcast hash join preserves duplicate struct fields") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.PREFER_SORTMERGEJOIN.key -> "false", + CometConf.COMET_EXEC_BROADCAST_FORCE_ENABLED.key -> "true") { + withParquetTable(Seq((1, 10), (2, 20)), "build") { + withParquetTable(Seq(1, 2, 3).map(Tuple1(_)), "probe") { + val query = + """SELECT /*+ BROADCAST(b) */ p._1, b.s + |FROM probe p + |JOIN ( + | SELECT _1 AS k, + | named_struct( + | '__comet_runtime_field_0', _1, + | '__comet_runtime_field_0', _2) AS s + | FROM build + |) b ON p._1 = b.k""".stripMargin + + checkSparkAnswerAndOperator( + sql(query), + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + } + } + } + } + // Reproducer for SPARK-43113: full outer SMJ with a join filter that references // a nullable column should not match when the filter evaluates to NULL. test("SPARK-43113: Full outer SMJ with NULL in join filter") { diff --git a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala index e60b26ff61a..359a49e2339 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala @@ -21,13 +21,14 @@ package org.apache.comet.vector import java.io.IOException import java.nio.charset.StandardCharsets +import java.util.{Arrays, Collections} import scala.util.Using import org.apache.arrow.c.{ArrowArray, ArrowSchema, Data} import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.{IntVector, UInt4Vector, VarCharVector} -import org.apache.arrow.vector.complex.StructVector +import org.apache.arrow.vector.{BitVector, IntVector, TinyIntVector, UInt4Vector, VarCharVector} +import org.apache.arrow.vector.complex.{ListVector, StructVector} import org.apache.arrow.vector.dictionary.Dictionary import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType} @@ -365,4 +366,77 @@ class NativeUtilSuite extends CometTestBase { nativeUtil.close() } } + + test("importVector preserves duplicate struct fields positionally") { + Using.Manager { use => + val allocator = use(new RootAllocator(Long.MaxValue)) + val nativeUtil = use(new NativeUtil) + val children = Arrays.asList( + new Field( + "a", + FieldType.nullable(ArrowType.Bool.INSTANCE), + Collections.emptyList[Field]()), + new Field( + "a", + FieldType.nullable(new ArrowType.Int(8, true)), + Collections.emptyList[Field]())) + val field = new Field("value", FieldType.nullable(ArrowType.Struct.INSTANCE), children) + val source = use(NativeUtil.createVector(field, allocator).asInstanceOf[StructVector]) + + source.allocateNew() + val bool = source.getChildByOrdinal(0).asInstanceOf[BitVector] + val byte = source.getChildByOrdinal(1).asInstanceOf[TinyIntVector] + bool.setSafe(0, 1) + byte.setSafe(0, 7) + bool.setValueCount(1) + byte.setValueCount(1) + source.setIndexDefined(0) + source.setValueCount(1) + + val array = ArrowArray.allocateNew(allocator) + val schema = ArrowSchema.allocateNew(allocator) + var handedToImport = false + try { + Data.exportVector(allocator, source, null, array, schema) + + // importVector owns both structs once called, including on its failure path. + handedToImport = true + val imported = use(nativeUtil.importVector(Array(array), Array(schema)).head) + val row = imported.getStruct(0) + assert(row.getBoolean(0)) + assert(row.getByte(1) === 7.toByte) + assert(imported.getValueVector.getField.getChildren.get(0).getName === "a") + assert(imported.getValueVector.getField.getChildren.get(1).getName === "a") + } finally { + if (!handedToImport) { + try array.release() + finally { + try array.close() + finally { + try schema.release() + finally schema.close() + } + } + } + } + }.get + } + + test("import vector factory retains Arrow's default path without duplicate struct names") { + Using.resource(new RootAllocator(Long.MaxValue)) { allocator => + val child = new Field( + "item", + FieldType.nullable(new ArrowType.Int(32, true)), + Collections.emptyList[Field]()) + val field = new Field( + "values", + FieldType.nullable(ArrowType.List.INSTANCE), + Collections.singletonList(child)) + + Using.resource(NativeUtil.createVectorForImport(field, allocator)) { vector => + assert(vector.isInstanceOf[ListVector]) + assert(vector.getField.getChildren.get(0).getName === "$data$") + } + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 4510f9d0ac1..5c9cd80d581 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,14 +19,21 @@ package org.apache.spark.sql.comet.util +import java.util.{Arrays, Collections} + +import scala.collection.JavaConverters._ + import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.vector.IntVector +import org.apache.arrow.vector.complex.StructVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.CometArrowAllocator -import org.apache.comet.vector.CometVector +import org.apache.comet.vector.{CometVector, NativeUtil} class UtilsSuite extends CometTestBase { @@ -58,6 +65,49 @@ class UtilsSuite extends CometTestBase { assert(decoded.map(_.numRows()).sum == expected) } + test("coalesceBroadcastBatches preserves duplicate struct fields") { + val allocator = + CometArrowAllocator.newChildAllocator("duplicate-struct-broadcast", 0, Long.MaxValue) + try { + val intType = FieldType.nullable(new ArrowType.Int(32, true)) + val children = Arrays.asList( + new Field("x", intType, Collections.emptyList[Field]()), + new Field("x", intType, Collections.emptyList[Field]())) + val field = new Field("s", FieldType.nullable(ArrowType.Struct.INSTANCE), children) + val root = + NativeUtil.createVectorSchemaRootForImport(new Schema(Seq(field).asJava), allocator) + val source = root.getVector(0).asInstanceOf[StructVector] + source.allocateNew() + source.getChildByOrdinal(0).asInstanceOf[IntVector].setSafe(0, 1) + source.getChildByOrdinal(1).asInstanceOf[IntVector].setSafe(0, 10) + source.setIndexDefined(0) + source.setValueCount(1) + + val batch = new ColumnarBatch(Array(CometVector.getVector(source, null)), 1) + val serialized = + try Utils.serializeBatches(Iterator(batch)).map(_._2).toArray + finally batch.close() + + val (coalesced, batchCount, totalRows) = + Utils.coalesceBroadcastBatches(serialized.iterator) + assert(batchCount == 1) + assert(totalRows == 1) + + val decoded = coalesced.iterator.flatMap(Utils.decodeBatches(_, "test")) + assert(decoded.hasNext) + val output = decoded.next() + val value = output.column(0).asInstanceOf[CometVector] + val row = value.getStruct(0) + assert(row.getInt(0) == 1) + assert(row.getInt(1) == 10) + assert(value.getValueVector.getField.getChildren.get(0).getName == "x") + assert(value.getValueVector.getField.getChildren.get(1).getName == "x") + assert(!decoded.hasNext) + } finally { + allocator.close() + } + } + test("serializeBatches materializes ConstantColumnVector columns") { // Spark wraps file-source partition columns and other per-batch constants in // ConstantColumnVector. When such a batch reaches Comet's serialization/export path