grorge123 commented on code in PR #5526:
URL: https://github.com/apache/datafusion-comet/pull/5526#discussion_r3920830465
##########
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:
Fixed: the serialization is forced with `toVector` before the inputs are
closed, so the Scala 2.12 Stream laziness no longer drops the tail batches.
Verified with UtilsSuite under `-Pspark-3.5` (Scala 2.12), which reproduces the
CI environment.
##########
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:
Fixed: `CometElementAt` keeps a non-deterministic nullable collection in
Spark under ANSI, since the null guard's two copies each hold their own kernel
state and the THEN copy only sees the CASE-selected rows. Reproduced the
witness first (Comet returned `[null]` where Spark keeps `[[2,null]]`), and it
now matches. The same guard is built by `CometSize` (non-legacy mode),
`CometArrayAppend`, `CometMapFromArrays` and `CometCoalesce`, so they share the
gate (`NullGuard`); `CometSize` drops the guard in legacy mode, where native
already answers -1. The coalesce case is reachable on main without NullType
(`coalesce(IF(monotonically_increasing_id() % 2 = 0, array(id), NULL),
array(id))` NPEs in columnar-to-row because the result is declared
non-nullable); it is included since the sweep found it and the fix is one line
on the shared guard.
Tests: expect_fallback queries in element_at_ansi.sql, map_from_arrays.sql
and coalesce.sql, plus the non-deterministic dimension of
CometNullTypeCompositionSuite.
##########
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:
Fixed: `CometCreateArray` keeps non-foldable NullType arguments in Spark -
`make_array`'s all-Null branch (`SingleRowListArrayBuilder`) collapses the
batch to one row, reproduced as "Expected: 4, Got: 1". All-literal NULLs arrive
as scalars and broadcast correctly, so `array(NULL)` literals stay native.
Tests: expect_fallback query in create_array.sql.
##########
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:
Fixed: a dedicated `CometArrayRepeat` keeps `containsNull = false`
non-NullType items in Spark, since DataFusion's repeat rebuilds the item field
as nullable (reproduced the ListArray nested-field rejection). A NullType item
is declared nullable on the FFI boundary and matches the rebuild, so plain
`array<null>` stays native.
Tests: expect_fallback query in array_repeat.sql.
##########
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:
Fixed: `CometArrayUnion` keeps NullType-element unions in Spark - the set-op
kernel's `value_type().is_null()` branch returns `distinct(other side)`,
reproduced as `[]` where Spark keeps `[NULL]`. The existing
`array_union(array(), array())` / `array_union(array(), array(NULL))` queries
move to expect_fallback with it (their operands are Null-typed; the empty cases
are benign natively but the gate cannot see emptiness statically).
Tests: expect_fallback queries in array_union.sql, including the non-empty
witness.
--
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]
