andygrove commented on code in PR #5603:
URL: https://github.com/apache/datafusion-comet/pull/5603#discussion_r3942430315
##########
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
Review Comment:
This case does what the comment above it says, which is the problem: the
`REPARTITION(2, a)` hint is inside the subquery, so `a` and `b` cross the
exchange as plain `int` and `string` and the struct is built afterwards.
Nothing in the PR carries a duplicate-name struct *through* a shuffle, which is
the boundary the `StreamReader` change is for.
Moving the hint outside would cover it:
```sql
SELECT /*+ REPARTITION(3) */ s
FROM (SELECT named_struct('x', a, 'x', b) AS s FROM test_named_struct)
shuffled
```
Worth knowing before you try it that this version currently fails the
fixture's coverage assertion, for the shuffle-guard reason in my summary. It
passes once that guard is gone.
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -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())
Review Comment:
The identity fast path works as intended. A scalar column now costs one
`getChildren().isEmpty()` and goes straight to `field.createVector`, and the
new `$data$` regression pins that.
One allocation is left on the per-batch path. `names` is a fresh
`HashSet[String]` for every struct node with more than one child, created on
every `importVector` call for every complex column, and in the overwhelmingly
common no-duplicate case it gets filled and discarded having proved nothing.
For the child counts structs usually have, a linear scan over the
already-materialised `children` list would avoid it, or the set could be
created only on the first repeat.
This is second-order next to the `Field` tree that `importField` already
rebuilds on the same call, so only worth doing if it stays a one-liner.
##########
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") {
Review Comment:
Removing this makes sense, since it asserted `ProjectExec` and that is
precisely the behaviour the PR changes. The input diversity went with it
though: ten thousand rows, `dictionaryEnabled` both ways, and a literal child
in `named_struct('a', _1, 'a', 2)`. The new fixture has three rows, no
dictionary variation, and no duplicate-plus-literal case.
I reran these queries against the branch and they all match, so nothing is
broken. But a dictionary-encoded string child inside a duplicate-name struct is
cheap to keep and awkward to notice losing. Would you add a `dictionaryEnabled`
loop over a duplicate-name struct somewhere, or a `--CONFIG` matrix line on the
fixture?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala:
##########
@@ -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)
Review Comment:
This override drops three things `ArrowReader.initialize` did: the per-field
`DictionaryUtility.toMemoryFormat` call, `originalSchema.getCustomMetadata()`,
and populating `dictionaries`.
I traced all three and they look safe.
`CometArrowStream.reconcileStreamSchema` already decodes
`CometDictionaryVector` columns down to the dictionary's value type, so the
schema reaching here never carries a `DictionaryEncoding`, and neither
`Utils.toArrowSchema` nor `reconcileStreamSchema` ever sets schema-level
metadata, so there is none to lose.
That took a while to establish though, and the commit directly under this
branch (5552) was specifically about Arrow metadata surviving C Data exports.
Could a sentence here record why skipping both is correct, so the next person
does not have to re-derive it?
##########
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 = _
Review Comment:
This class and `ColumnarBatchArrowReader` now carry the same forty-odd lines
of shadow state: `cometInitialized`, `cometRoot` and `cometLoader`, plus
overrides of `ensureInitialized`, `getVectorSchemaRoot`,
`prepareLoadNextBatch`, `loadRecordBatch`, `lookup`, `getDictionaryIds`,
`getDictionaryVectors` and `close`.
The reason is good and not at all obvious from reading either file.
`ArrowReader.root`, `loader` and `initialized` are all private in Arrow 18.3,
so there is no way to reuse them, and every base method that reads them has to
be overridden or it silently sees `null` or `false`. I enumerated those methods
and the override set is complete in both classes, which is the part I most
expected to be wrong. The class doc here explains the allocation motive but not
the shadowing one.
Two smaller things suggest the second copy did not get the same pass as the
first. `getDictionaryIds` throws `IllegalStateException` when uninitialized
here but calls `ensureInitialized()` in `ColumnarBatchArrowReader`, and Arrow's
base does neither. And `ColumnarBatchArrowReader.close(closeReadSource:
Boolean)` ignores its parameter, which is harmless because `closeReadSource()`
is `()` in that class but reads as an oversight.
Would a shared `private[comet]` base in `org.apache.comet.vector` holding
the shadow root and loader and the mechanical overrides be worth it? Mostly so
the private-field constraint has one place to be written down, since that is
the thing most likely to get lost when someone next touches either class.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
