sunchao commented on code in PR #5526:
URL: https://github.com/apache/datafusion-comet/pull/5526#discussion_r3898273998
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:
##########
@@ -113,13 +142,19 @@ object CometBatchKernelCodegen extends Logging with
CometExprTraitShim with Come
* back cleanly rather than crashing the Janino compile at execute time.
*
* Checks every `BoundReference`'s data type and the root `expr.dataType`
against
- * [[isSupportedDataType]], rejects aggregates / generators / `Unevaluable`,
and gates total
- * nested-field count on `spark.sql.codegen.maxFields`.
+ * [[isSupportedDataType]] and [[duplicateStructFieldNames]], rejects
aggregates / generators /
+ * `Unevaluable`, and gates total nested-field count on
`spark.sql.codegen.maxFields`.
*/
def canHandle(boundExpr: Expression): Option[String] = {
- if (!isSupportedDataType(boundExpr.dataType)) {
+ if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) {
Review Comment:
[P2] Preserve the item field when repeating map entries
The [NullType-leaf
normalization](https://github.com/apache/datafusion-comet/pull/5526#issuecomment-5478130993)
still leaves `array_repeat(map_entries(map(id, NULL)), 2)` failing over a
non-null native LONG input with codegen dispatch enabled. `map_entries`
produces a list whose struct item is non-nullable. The registered
SparkArrayRepeat wrapper delegates to native repeat, which rebuilds the inner
list with a nullable item while its outer field still expects the original
non-nullable item type. Arrow's ListArray constructor compares those nested
fields, including nullability, and rejects the mismatch. The item here is a
struct, so normalizing direct NullType leaves does not repair it. Please
preserve that item field or keep this newly admitted composition in Spark. BASE
rejects the Null-bearing map. The same admission already existed at the
previous reviewed head. Source-derived, not executed.
##########
spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala:
##########
@@ -197,4 +212,270 @@ class UtilsSuite extends CometTestBase {
}
}
}
+
+ /**
+ * One map column of `numRows` rows. With an `IntegerType` key every row is
a single entry `i ->
+ * NULL` (a `NullVector` map value); with a `NullType` key every row is an
empty map, as `map()`
+ * produces (a `NullVector` map key). Both nest a `NullVector` inside the
entries struct.
+ */
+ private def nullTypeMapBatch(numRows: Int, keyType: DataType): ColumnarBatch
= {
+ val field = Utils.toArrowField("m", MapType(keyType, NullType), nullable =
true, "UTC")
+ val vector =
field.createVector(CometArrowAllocator).asInstanceOf[MapVector]
+ vector.allocateNew()
+ val entries = vector.getDataVector.asInstanceOf[StructVector]
+ (0 until numRows).foreach { i =>
+ vector.startNewValue(i)
+ keyType match {
+ case NullType =>
+ vector.endValue(i, 0)
+ case _ =>
+ entries.setIndexDefined(i)
+
entries.getChild(MapVector.KEY_NAME).asInstanceOf[IntVector].setSafe(i, i)
+ vector.endValue(i, 1)
+ }
+ }
+ entries.setValueCount(if (keyType == NullType) 0 else numRows)
+ vector.setValueCount(numRows)
+ new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector,
null)), numRows)
+ }
+
+ private def mapKeyField(field: Field): Field =
field.getChildren.get(0).getChildren.get(0)
+
+ /** One `array<null>` column; row `i` holds `i` nulls. */
+ private def nullListBatch(numRows: Int): ColumnarBatch = {
+ val field = Utils.toArrowField("l", ArrayType(NullType), nullable = true,
"UTC")
+ val vector =
field.createVector(CometArrowAllocator).asInstanceOf[ListVector]
+ vector.allocateNew()
+ (0 until numRows).foreach { i =>
+ vector.startNewValue(i)
+ vector.endValue(i, i)
+ }
+ vector.setValueCount(numRows)
+ new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector,
null)), numRows)
+ }
+
+ /** One `array<struct<a: null>>` column; every row holds one struct. */
+ private def nullStructListBatch(numRows: Int): ColumnarBatch = {
+ val elementType = StructType(Seq(StructField("a", NullType)))
+ val field = Utils.toArrowField("l", ArrayType(elementType), nullable =
true, "UTC")
+ val vector =
field.createVector(CometArrowAllocator).asInstanceOf[ListVector]
+ vector.allocateNew()
+ val elements = vector.getDataVector.asInstanceOf[StructVector]
+ (0 until numRows).foreach { i =>
+ vector.startNewValue(i)
+ elements.setIndexDefined(i)
+ vector.endValue(i, 1)
+ }
+ elements.setValueCount(numRows)
+ vector.setValueCount(numRows)
+ new ColumnarBatch(Array[ColumnVector](CometVector.getVector(vector,
null)), numRows)
+ }
+
+ test("withNonNullableMapKeys restores the non-nullable key flag a NullVector
drops") {
+ val batch = nullTypeMapBatch(2, NullType)
+ val field =
batch.column(0).asInstanceOf[CometVector].getValueVector.getField
+ // `toArrowField` declares the key non-nullable, but Arrow's
`MinorType.NULL` factory builds the
+ // key `NullVector` from the name alone, so the vector reports a nullable
key. If this assertion
+ // starts failing, Arrow fixed that and `withNonNullableMapKeys` can go.
+ assert(mapKeyField(field).isNullable)
+
+ val repaired = Utils.withNonNullableMapKeys(field)
+ assert(!mapKeyField(repaired).isNullable)
+ assert(mapKeyField(repaired).getType.isInstanceOf[ArrowType.Null])
+ assert(repaired.getName == field.getName)
+ assert(repaired.getFieldType == field.getFieldType)
+ assert(repaired.getChildren.get(0).getFieldType ==
field.getChildren.get(0).getFieldType)
+ // Idempotent, and a no-op on fields that already satisfy the invariant.
+ assert(Utils.withNonNullableMapKeys(repaired) eq repaired)
+ batch.close()
+ }
+
+ test("newArrowStreamWriter keeps a root whose declared schema is already
valid") {
+ val batch = nullTypeMapBatch(2, NullType)
+ val vector =
+
batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector]
+ val declared = Utils.withNonNullableMapKeys(vector.getField)
+ // The live vector still reports a nullable key, so a root declared from
it would be swapped
+ // for a repaired copy. One declared from `declared` must be kept as-is:
the row count is set
+ // only after the writer exists, and a swapped root would not see it.
+ val root = new VectorSchemaRoot(Seq(declared).asJava, Seq(vector).asJava,
0)
+ val out = new ByteArrayOutputStream()
+ val (bound, writer) = Utils.newArrowStreamWriter(root, null,
Channels.newChannel(out))
+ assert(bound eq root)
+ root.setRowCount(2)
+ writer.start()
+ writer.writeBatch()
+ writer.end()
+
+ val reader =
+ new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray),
CometArrowAllocator)
+ assert(reader.loadNextBatch())
+ assert(reader.getVectorSchemaRoot.getRowCount == 2)
+
assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable)
+ reader.close()
+ batch.close()
+ }
+
+ test("newArrowStreamWriter returns the root a later row count must be set
on") {
+ val batch = nullTypeMapBatch(2, NullType)
+ val vector =
+
batch.column(0).asInstanceOf[CometVector].getValueVector.asInstanceOf[FieldVector]
+ // Declared from the live vector, so the key is nullable and the root must
be swapped. Setting
+ // the row count on the returned root has to reach the writer; setting it
on the original one
+ // would ship an empty batch ("Array length did not match record batch
length" downstream).
+ val root = new VectorSchemaRoot(Seq(vector).asJava)
+ val out = new ByteArrayOutputStream()
+ val (bound, writer) = Utils.newArrowStreamWriter(root, null,
Channels.newChannel(out))
+ assert(bound ne root)
+ bound.setRowCount(2)
+ writer.start()
+ writer.writeBatch()
+ writer.end()
+
+ val reader =
+ new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray),
CometArrowAllocator)
+ assert(reader.loadNextBatch())
+ assert(reader.getVectorSchemaRoot.getRowCount == 2)
+
assert(!mapKeyField(reader.getVectorSchemaRoot.getSchema.getFields.get(0)).isNullable)
+ reader.close()
+ batch.close()
+ }
+
+ test("serializeBatches round-trips a NullType map key through Arrow IPC") {
+ // The IPC reader rebuilds a MapVector from the stream's schema and
rejects a nullable key
+ // ("Map data key type should be a non-nullable"), which is exactly what a
NullVector key
+ // reports unless the written schema is repaired.
+ val numRows = 3
+ val batch = nullTypeMapBatch(numRows, NullType)
+ val (rowCount, buf) = Utils.serializeBatches(Iterator(batch)).next()
+ assert(rowCount == numRows)
+
+ val decoded = Utils.decodeBatches(buf, "test").toSeq
+ assert(decoded.map(_.numRows()).sum == numRows)
+ decoded.foreach(_.close())
+ }
+
+ test("coalesceBroadcastBatches ships struct-nested NullType uncoalesced") {
+ // VectorSchemaRootAppender cannot grow a NullVector nested in a struct,
including the map
+ // entries struct (NullVector.reAlloc is a no-op), so such buffers must be
passed through,
+ // not appended. The list case pins that a struct below a list is still a
struct.
+ val cases: Seq[(String, Int => ColumnarBatch)] = Seq(
+ "map<int, null>" -> (nullTypeMapBatch(_, IntegerType)),
+ "map<null, null>" -> (nullTypeMapBatch(_, NullType)),
+ "array<struct<a: null>>" -> nullStructListBatch)
+ cases.foreach { case (name, batch) =>
+ val numRows = 4
+ val numBatches = 3
+ val batches = (0 until numBatches).map(_ => batch(numRows))
+ val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq
+
+ val (result, batchCount, totalRows) =
Utils.coalesceBroadcastBatches(bufs.iterator)
+ // The pass-through signature: original buffers, nothing coalesced.
+ assert(batchCount == 0 && totalRows == 0, name)
+ assert(result.length == numBatches, name)
+
+ val decoded = result.iterator.flatMap(b => Utils.decodeBatches(b,
"test")).toSeq
+ assert(decoded.map(_.numRows()).sum == numRows.toLong * numBatches, name)
+ decoded.foreach(_.close())
+ }
+ }
+
+ test("coalesceBroadcastBatches bypasses exactly the schemas with a NullType
under a struct") {
+ // Exhaustive over the shape space of the bypass rule: VectorAppender
hangs only when a
+ // NullVector is a *direct* child of a struct (see
`Utils.hasNullDirectlyUnderStruct` for
+ // the Arrow mechanics). Each shape runs the real appender under a
timeout, so a rule that
+ // is too narrow shows up as a timeout on the hanging shapes instead of a
hung build, and
+ // one that is too wide shows up as a needless bypass.
+ val nullStruct = StructType(Seq(StructField("a", NullType)))
+ val shapes: Seq[(DataType, Any)] = Seq(
+ NullType -> null,
+ ArrayType(NullType) -> new GenericArrayData(Array[Any](null)),
+ ArrayType(ArrayType(NullType)) ->
+ new GenericArrayData(Array[Any](new
GenericArrayData(Array[Any](null)))),
+ nullStruct -> InternalRow(null),
+ ArrayType(nullStruct) -> new
GenericArrayData(Array[Any](InternalRow(null))),
+ StructType(Seq(StructField("l", ArrayType(NullType)))) ->
+ InternalRow(new GenericArrayData(Array[Any](null))),
+ MapType(IntegerType, NullType) -> ArrayBasedMapData(Array[Any](1),
Array[Any](null)),
+ // `map(k, array(NULL))`: the entry struct's direct child is a list, not
a NullVector, so
+ // this still coalesces.
+ MapType(IntegerType, ArrayType(NullType)) ->
+ ArrayBasedMapData(Array[Any](1), Array[Any](new
GenericArrayData(Array[Any](null)))),
+ MapType(NullType, NullType) -> ArrayBasedMapData(Array.empty[Any],
Array.empty[Any]))
+ // A list insulates whatever is below it, so `inStruct` resets when
descending into one.
+ def nullUnderStruct(dt: DataType, inStruct: Boolean): Boolean = dt match {
+ case NullType => inStruct
+ case ArrayType(element, _) => nullUnderStruct(element, inStruct = false)
+ case StructType(fields) => fields.exists(f =>
nullUnderStruct(f.dataType, inStruct = true))
+ case MapType(k, v, _) =>
+ nullUnderStruct(k, inStruct = true) || nullUnderStruct(v, inStruct =
true)
+ case _ => false
+ }
+
+ val numRows = 4
+ val numBatches = 3
+ shapes.foreach { case (dataType, value) =>
+ val name = dataType.simpleString
+ val schema = StructType(Seq(StructField("c", dataType)))
+ val batches = (0 until numBatches).map { _ =>
+ CometArrowConverters
+ .rowToArrowBatchIter(
+ Iterator.fill(numRows)(InternalRow(value)),
+ schema,
+ numRows,
+ "UTC",
+ CometArrowAllocator)
+ .next()
+ }
+ val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq
Review Comment:
[P2] Materialize serialization before closing the fixture batches
On Scala 2.12, `Iterator.toSeq` produces a lazy `Stream`: the first batch is
serialized here, but the tail is not. The next line closes all three input
batches. Closing their ListVectors clears the later value counts, so the
remaining serializations construct zero-row roots and this fixture supplies 4 +
0 + 0 rows instead of 12. This matches the `array<void>` assertion in the
[Spark 3.4 / Scala 2.12 CI
job](https://github.com/apache/datafusion-comet/actions/runs/33394656333/job/99563322366).
Please consume the iterator eagerly, for example with `toVector`, before
closing the inputs. This is a fixture lifetime defect. That CI failure does not
establish production coalescer row loss. I did not rerun the test.
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:
##########
@@ -113,13 +142,19 @@ object CometBatchKernelCodegen extends Logging with
CometExprTraitShim with Come
* back cleanly rather than crashing the Janino compile at execute time.
*
* Checks every `BoundReference`'s data type and the root `expr.dataType`
against
- * [[isSupportedDataType]], rejects aggregates / generators / `Unevaluable`,
and gates total
- * nested-field count on `spark.sql.codegen.maxFields`.
+ * [[isSupportedDataType]] and [[duplicateStructFieldNames]], rejects
aggregates / generators /
+ * `Unevaluable`, and gates total nested-field count on
`spark.sql.codegen.maxFields`.
*/
def canHandle(boundExpr: Expression): Option[String] = {
- if (!isSupportedDataType(boundExpr.dataType)) {
+ if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) {
Review Comment:
[P2] Preserve single evaluation of nullable stateful outputs
With codegen dispatch and ANSI enabled, consider
`element_at(transform(IF(monotonically_increasing_id() % 2 = 0, array(id),
CAST(NULL AS array<bigint>)), x -> named_struct('id', x, 'n', NULL)), 1)` over
a native LONG batch containing `id=0,1,2,3`. Spark evaluates the transform once
per row, retaining the struct for `id=2`. This gate now admits that
Null-bearing output, but the existing ElementAt ANSI conversion serializes its
left subtree into both the CASE predicate and the lookup. Native CASE tests all
four rows, then reevaluates the transform on the two selected rows. The second
selected row receives an odd counter and becomes NULL. Please materialize the
left value once or retain Spark fallback for this composition. Current BASE
rejects this output. The previous reviewed head admitted it but lacked the
duplicated-left CASE, so the newly failing combination is in this increment.
Source-derived witness, not executed.
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:
##########
@@ -81,21 +81,50 @@ object CometBatchKernelCodegen extends Logging with
CometExprTraitShim with Come
/**
* Type surface the kernel covers on both input and output sides. Recursive:
complex types are
* supported when their children are.
+ *
+ * `NullType` is output-only: [[CometBatchKernelCodegenOutput]] can write an
all-null Arrow
+ * `NullVector`, but `CometScalaUDFCodegen.specFor` cannot build an
[[ArrowColumnSpec]] for one,
+ * so a `NullType` input (nested or not) has to keep falling back to Spark.
*/
- def isSupportedDataType(dt: DataType): Boolean = dt match {
+ def isSupportedDataType(dt: DataType): Boolean = isSupportedDataType(dt,
allowNullType = false)
+
+ private def isSupportedDataType(dt: DataType, allowNullType: Boolean):
Boolean = dt match {
+ case NullType => allowNullType
Review Comment:
[P2] Preserve row count for newly admitted all-null array producers
With codegen dispatch enabled, `array(aggregate(array(id), NULL, (acc, x) ->
NULL))` over a live LONG input now sends a NullArray of the batch length to
native `make_array`. The locked implementation's all-Null branch builds one
list containing those nulls, rather than one list per row. With N > 1 input
rows, the physical scalar-function row-count guard therefore raises an error
for length 1 versus N. The ordinary nonfoldable ArrayAggregate survives the
inspected Spark 3.5/4.0 optimizer rules. A folded `array(NULL)` or a one-row
batch does not exercise this case. Please preserve batch cardinality in the
consumer or retain fallback for this composition. BASE rejects this producer
and stays in Spark. This is a previously unreported PR admission bug already
present at the prior reviewed head, not a new dependency change.
Source-derived, not executed.
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:
##########
@@ -113,13 +142,19 @@ object CometBatchKernelCodegen extends Logging with
CometExprTraitShim with Come
* back cleanly rather than crashing the Janino compile at execute time.
*
* Checks every `BoundReference`'s data type and the root `expr.dataType`
against
- * [[isSupportedDataType]], rejects aggregates / generators / `Unevaluable`,
and gates total
- * nested-field count on `spark.sql.codegen.maxFields`.
+ * [[isSupportedDataType]] and [[duplicateStructFieldNames]], rejects
aggregates / generators /
+ * `Unevaluable`, and gates total nested-field count on
`spark.sql.codegen.maxFields`.
*/
def canHandle(boundExpr: Expression): Option[String] = {
- if (!isSupportedDataType(boundExpr.dataType)) {
+ if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) {
Review Comment:
[P2] Retain real NULL entries in newly admitted array unions
With codegen dispatch enabled and
`spark.sql.legacy.createEmptyCollectionUsingStringType=false`,
`array_union(transform(array(id), x -> NULL), array())` over a live LONG input
should return `[NULL]` for each row. The new gate admits the transform's
List<Null> output, but native union treats `left.value_type().is_null()` as a
reason to return `distinct(right)`, discarding the left list's actual entries.
The empty right list is accepted by the locked Arrow row converter, so the
result is `[]`, not `[NULL]`. The inspected Spark 3.5/4.0 type guards
explicitly allow NullType here and their union implementation retains one NULL.
Please handle Null elements using the list offsets or retain fallback for this
composition. BASE refuses the producer. This previously unreported admission
bug was already possible at the prior reviewed head. Source-derived, not
executed.
--
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]
