From a69107524ebd15b158c3faffe667b5cbe9be9830 Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Tue, 1 Sep 2026 05:53:36 +0800 Subject: [PATCH 1/7] fix: preserve duplicate named_struct fields --- docs/source/user-guide/latest/expressions.md | 2 +- .../org/apache/arrow/c/ArrowImporter.java | 12 ++- .../CometBatchKernelCodegenOutput.scala | 48 ++--------- .../org/apache/comet/serde/structs.scala | 4 +- .../org/apache/comet/vector/NativeUtil.scala | 86 ++++++++++++++++++- .../org/apache/comet/CometCodegenSuite.scala | 22 +++++ .../apache/comet/CometExpressionSuite.scala | 11 +-- .../apache/comet/vector/NativeUtilSuite.scala | 46 +++++++++- 8 files changed, 175 insertions(+), 56 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 2ba71e93c4d..013ababdb03 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -624,7 +624,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | Function | Status | Implementation | Notes | | --- | --- | --- | --- | -| `named_struct` | ✅ | Native | Duplicate field names fall back | +| `named_struct` | ✅ | Hybrid | Duplicate field names route through the JVM codegen dispatcher | | `struct` | ✅ | Native | | --- 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/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..ecac127a5c9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/structs.scala +++ b/spark/src/main/scala/org/apache/comet/serde/structs.scala @@ -31,7 +31,9 @@ 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" 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..70322a2d24c 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -20,11 +20,15 @@ package org.apache.comet.vector import scala.collection.mutable +import scala.jdk.CollectionConverters._ 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} import org.apache.spark.SparkException import org.apache.spark.sql.comet.execution.arrow.ConstantColumnVectors import org.apache.spark.sql.comet.util.Utils @@ -268,7 +272,11 @@ 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, + field => NativeUtil.createVector(field, allocator)) val cometVector = try CometVector.getVector(arrowVector, dictionaryProvider) catch { @@ -322,6 +330,80 @@ 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) + field.getType match { + case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList => + val vector = new RenamedListVector(runtimeField, field, allocator) + vector.initializeChildrenFromFields(runtimeField.getChildren) + vector + case _: ArrowType.Map => + val vector = new RenamedMapVector(runtimeField, field, allocator) + vector.initializeChildrenFromFields(runtimeField.getChildren) + vector + case _: ArrowType.Struct => + val vector = new RenamedStructVector(runtimeField, field, allocator) + vector.initializeChildrenFromFields(runtimeField.getChildren) + vector + case _ => field.createVector(allocator).asInstanceOf[FieldVector] + } + } + + private def fieldForAllocation(field: Field): Field = { + val children = field.getChildren.asScala.map(fieldForAllocation).toIndexedSeq + val runtimeChildren = field.getType match { + case _: ArrowType.Struct if children.map(_.getName).distinct.size != children.size => + children.zipWithIndex.map { case (child, ordinal) => + new Field(s"__comet_runtime_field_$ordinal", child.getFieldType, child.getChildren) + } + case _ => children + } + new Field(field.getName, field.getFieldType, runtimeChildren.asJava) + } + + /** + * 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) { + private var constructed = false + constructed = true + + override def getField: Field = if (constructed) exportField else super.getField + } + def rootAsBatch(arrowRoot: VectorSchemaRoot): ColumnarBatch = { rootAsBatch(arrowRoot, null) } 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..dfdc6f51ab9 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -2480,15 +2480,10 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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") + checkSparkAnswerAndOperator("SELECT named_struct('a', _1, 'a', 2) FROM 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]) + "SELECT named_struct('a', named_struct('b', _1, 'b', _2)) FROM tbl") } } } 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..a9e6a5fa72f 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala @@ -21,12 +21,13 @@ 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.{BitVector, IntVector, TinyIntVector, UInt4Vector, VarCharVector} import org.apache.arrow.vector.complex.StructVector import org.apache.arrow.vector.dictionary.Dictionary import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider @@ -365,4 +366,47 @@ class NativeUtilSuite extends CometTestBase { nativeUtil.close() } } + + test("importVector preserves duplicate struct fields positionally") { + val allocator = new RootAllocator(Long.MaxValue) + val nativeUtil = 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 = NativeUtil.createVector(field, allocator).asInstanceOf[StructVector] + var imported: CometVector = null + + try { + 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) + Data.exportVector(allocator, source, null, array, schema) + source.close() + + imported = 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 { + source.close() + if (imported != null) imported.close() + nativeUtil.close() + allocator.close() + } + } } From 1646528e78bfdb76d012f7789d5607ef89fce2e0 Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Tue, 1 Sep 2026 10:22:13 +0800 Subject: [PATCH 2/7] fix: preserve nested duplicate struct imports --- docs/source/user-guide/latest/expressions.md | 2 +- .../org/apache/comet/serde/structs.scala | 6 +- .../org/apache/comet/vector/NativeUtil.scala | 122 ++++++++++++++---- .../struct/create_named_struct.sql | 30 +++++ .../apache/comet/CometExpressionSuite.scala | 15 --- .../apache/comet/vector/NativeUtilSuite.scala | 84 ++++++++---- 6 files changed, 189 insertions(+), 70 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 013ababdb03..1d18e5f0aa5 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -625,7 +625,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `named_struct` | ✅ | Hybrid | Duplicate field names route through the JVM codegen dispatcher | -| `struct` | ✅ | Native | | +| `struct` | ✅ | Hybrid | Duplicate field names route through the JVM codegen dispatcher | --- 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 ecac127a5c9..b5dc8da44c9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/structs.scala +++ b/spark/src/main/scala/org/apache/comet/serde/structs.scala @@ -36,9 +36,11 @@ object CometCreateNamedStruct 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/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 70322a2d24c..20c2882b96c 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -19,8 +19,10 @@ package org.apache.comet.vector +import java.util.{ArrayList, HashSet} +import java.util.function.Function + import scala.collection.mutable -import scala.jdk.CollectionConverters._ import org.apache.arrow.c.{ArrowArray, ArrowImporter, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.memory.BufferAllocator @@ -57,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. @@ -272,11 +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, - field => NativeUtil.createVector(field, allocator)) + val arrowVector = + importer.importVector(arrowArray, arrowSchema, dictionaryProvider, importVectorFactory) val cometVector = try CometVector.getVector(arrowVector, dictionaryProvider) catch { @@ -338,33 +341,101 @@ object NativeUtil { */ private[comet] def createVector(field: Field, allocator: BufferAllocator): FieldVector = { val runtimeField = fieldForAllocation(field) - field.getType match { + 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) + } + } + + 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, field, allocator) + val vector = new RenamedListVector(runtimeField, exportField, allocator) vector.initializeChildrenFromFields(runtimeField.getChildren) vector case _: ArrowType.Map => - val vector = new RenamedMapVector(runtimeField, field, allocator) + val vector = new RenamedMapVector(runtimeField, exportField, allocator) vector.initializeChildrenFromFields(runtimeField.getChildren) vector case _: ArrowType.Struct => - val vector = new RenamedStructVector(runtimeField, field, allocator) - vector.initializeChildrenFromFields(runtimeField.getChildren) + val vector = new RenamedStructVector(runtimeField, exportField, allocator) + // The Field-based StructVector constructor creates the direct children. Initialize each + // child's descendants from the runtime schema without adding the direct children twice. + val runtimeChildren = runtimeField.getChildren + var ordinal = 0 + while (ordinal < runtimeChildren.size()) { + vector + .getChildByOrdinal(ordinal) + .asInstanceOf[FieldVector] + .initializeChildrenFromFields(runtimeChildren.get(ordinal).getChildren) + ordinal += 1 + } vector - case _ => field.createVector(allocator).asInstanceOf[FieldVector] + case _ => exportField.createVector(allocator).asInstanceOf[FieldVector] } } private def fieldForAllocation(field: Field): Field = { - val children = field.getChildren.asScala.map(fieldForAllocation).toIndexedSeq - val runtimeChildren = field.getType match { - case _: ArrowType.Struct if children.map(_.getName).distinct.size != children.size => - children.zipWithIndex.map { case (child, ordinal) => - new Field(s"__comet_runtime_field_$ordinal", child.getFieldType, child.getChildren) + 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 } - case _ => children + } + 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) + renamedChildren.add( + new Field(s"__comet_runtime_field_$ordinal", 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 } - new Field(field.getName, field.getFieldType, runtimeChildren.asJava) } /** @@ -392,16 +463,17 @@ object NativeUtil { exportField: Field, allocator: BufferAllocator) extends StructVector( - runtimeField.getName, + runtimeField, allocator, - runtimeField.getFieldType, null, AbstractStructVector.ConflictPolicy.CONFLICT_ERROR, true) { - private var constructed = false - constructed = true - - override def getField: Field = if (constructed) exportField else super.getField + override def getField: Field = { + // StructVector's constructor calls getField before creating its children. Keep the unique + // runtime field visible for that call, then publish the original metadata once all children + // exist. The child count avoids a separate construction-state flag. + if (size() == exportField.getChildren.size()) exportField else super.getField + } } def rootAsBatch(arrowRoot: VectorSchemaRoot): ColumnarBatch = { 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/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index dfdc6f51ab9..d0185a28b49 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -2474,21 +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") - checkSparkAnswerAndOperator("SELECT named_struct('a', _1, 'a', 2) FROM tbl") - checkSparkAnswerAndOperator( - "SELECT named_struct('a', named_struct('b', _1, 'b', _2)) FROM tbl") - } - } - } - } - test("to_json") { withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToJson]) -> "true") { Seq(true, false).foreach { dictionaryEnabled => 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 a9e6a5fa72f..359a49e2339 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala @@ -28,7 +28,7 @@ import scala.util.Using import org.apache.arrow.c.{ArrowArray, ArrowSchema, Data} import org.apache.arrow.memory.RootAllocator import org.apache.arrow.vector.{BitVector, IntVector, TinyIntVector, UInt4Vector, VarCharVector} -import org.apache.arrow.vector.complex.StructVector +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} @@ -368,19 +368,21 @@ class NativeUtilSuite extends CometTestBase { } test("importVector preserves duplicate struct fields positionally") { - val allocator = new RootAllocator(Long.MaxValue) - val nativeUtil = 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 = NativeUtil.createVector(field, allocator).asInstanceOf[StructVector] - var imported: CometVector = null + 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]) - try { source.allocateNew() val bool = source.getChildByOrdinal(0).asInstanceOf[BitVector] val byte = source.getChildByOrdinal(1).asInstanceOf[TinyIntVector] @@ -393,20 +395,48 @@ class NativeUtilSuite extends CometTestBase { val array = ArrowArray.allocateNew(allocator) val schema = ArrowSchema.allocateNew(allocator) - Data.exportVector(allocator, source, null, array, schema) - source.close() - - imported = 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 { - source.close() - if (imported != null) imported.close() - nativeUtil.close() - allocator.close() + 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$") + } } } } From d966c4951f6aa1035e0ec86a3b86204425f6a1a5 Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Tue, 1 Sep 2026 15:23:11 +0800 Subject: [PATCH 3/7] fix: preserve case-distinct struct fields --- .../org/apache/comet/vector/NativeUtil.scala | 24 +++++++------------ .../org/apache/comet/CometJsonJvmSuite.scala | 6 +++++ 2 files changed, 15 insertions(+), 15 deletions(-) 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 20c2882b96c..110d81e80b1 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -374,17 +374,10 @@ object NativeUtil { vector case _: ArrowType.Struct => val vector = new RenamedStructVector(runtimeField, exportField, allocator) - // The Field-based StructVector constructor creates the direct children. Initialize each - // child's descendants from the runtime schema without adding the direct children twice. - val runtimeChildren = runtimeField.getChildren - var ordinal = 0 - while (ordinal < runtimeChildren.size()) { - vector - .getChildByOrdinal(ordinal) - .asInstanceOf[FieldVector] - .initializeChildrenFromFields(runtimeChildren.get(ordinal).getChildren) - ordinal += 1 - } + // 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] } @@ -463,15 +456,16 @@ object NativeUtil { exportField: Field, allocator: BufferAllocator) extends StructVector( - runtimeField, + runtimeField.getName, allocator, + runtimeField.getFieldType, null, AbstractStructVector.ConflictPolicy.CONFLICT_ERROR, true) { override def getField: Field = { - // StructVector's constructor calls getField before creating its children. Keep the unique - // runtime field visible for that call, then publish the original metadata once all children - // exist. The child count avoids a separate construction-state flag. + // 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 } } diff --git a/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala b/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala index 4a4df2e9b1a..25d0b0be4a3 100644 --- a/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala @@ -62,6 +62,12 @@ 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")) From 91bba712752eb0191efab376805a22a52a1ed785 Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Tue, 1 Sep 2026 23:22:11 +0800 Subject: [PATCH 4/7] fix: preserve duplicate fields at Arrow readers --- .../org/apache/comet/udf/CometUdfBridge.java | 9 +- .../comet/vector/CometArrowStreamReader.scala | 186 ++++++++++++++++++ .../org/apache/comet/vector/NativeUtil.scala | 22 ++- .../apache/comet/vector/StreamReader.scala | 5 +- .../apache/spark/sql/comet/util/Utils.scala | 9 +- .../org/apache/comet/CometJsonJvmSuite.scala | 14 ++ .../apache/comet/exec/CometJoinSuite.scala | 23 +++ .../spark/sql/comet/util/UtilsSuite.scala | 52 ++++- 8 files changed, 310 insertions(+), 10 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/vector/CometArrowStreamReader.scala 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/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 110d81e80b1..75091fb7356 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -30,7 +30,7 @@ import org.apache.arrow.util.AutoCloseables 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} +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 @@ -359,6 +359,26 @@ object NativeUtil { } } + /** 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 + } + } + private def createPinnedVector( runtimeField: Field, exportField: Field, 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/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/scala/org/apache/comet/CometJsonJvmSuite.scala b/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala index 25d0b0be4a3..6b7bab1969e 100644 --- a/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometJsonJvmSuite.scala @@ -73,4 +73,18 @@ class CometJsonJvmSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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..92c36a55b56 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,29 @@ 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('x', _1, 'x', _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/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 From 08d43fe6da747dfe137f56b569e7ac2bfdc0ec47 Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Thu, 3 Sep 2026 04:02:14 +0800 Subject: [PATCH 5/7] fix: preserve duplicate fields in columnar batch streams --- .../org/apache/comet/vector/NativeUtil.scala | 27 ++++++++ .../arrow/ColumnarBatchArrowReader.scala | 64 ++++++++++++++++++- 2 files changed, 88 insertions(+), 3 deletions(-) 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 75091fb7356..41468d8dfee 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -379,6 +379,33 @@ object NativeUtil { } } + /** + * 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, 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() From ba8ca2ccaaf8508b21562a938fa5394525d4ed3b Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Thu, 3 Sep 2026 06:42:41 +0800 Subject: [PATCH 6/7] fix: avoid runtime field name collisions --- .../src/main/scala/org/apache/comet/vector/NativeUtil.scala | 5 ++++- .../test/scala/org/apache/comet/exec/CometJoinSuite.scala | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) 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 41468d8dfee..749ced23ae8 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -466,8 +466,11 @@ object NativeUtil { 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(s"__comet_runtime_field_$ordinal", child.getFieldType, child.getChildren)) + new Field(runtimeName, child.getFieldType, child.getChildren)) ordinal += 1 } new Field(field.getName, field.getFieldType, renamedChildren) 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 92c36a55b56..9eb0a215d6c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -743,7 +743,10 @@ class CometJoinSuite extends CometTestBase { """SELECT /*+ BROADCAST(b) */ p._1, b.s |FROM probe p |JOIN ( - | SELECT _1 AS k, named_struct('x', _1, 'x', _2) AS s + | 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 From e1ecb7fb54dec61fc1cd59068d9c0437e7a8e872 Mon Sep 17 00:00:00 2001 From: RRXXZZYY Date: Thu, 3 Sep 2026 12:39:34 +0800 Subject: [PATCH 7/7] style: apply scalafmt to runtime field allocation --- spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 749ced23ae8..dcb1ac7a162 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -469,8 +469,7 @@ object NativeUtil { 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)) + renamedChildren.add(new Field(runtimeName, child.getFieldType, child.getChildren)) ordinal += 1 } new Field(field.getName, field.getFieldType, renamedChildren)