sunchao commented on code in PR #5452:
URL: https://github.com/apache/datafusion-comet/pull/5452#discussion_r3856509162
##########
spark/src/main/scala/org/apache/comet/serde/literals.scala:
##########
@@ -215,4 +225,129 @@ object CometLiteral extends CometExpressionSerde[Literal]
with Logging {
}
listLiteralBuilder
}
+
+ /**
+ * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` /
`CreateMap` over
+ * primitive-typed Literals, or `None` when the shape cannot be rebuilt. The
native `Literal`
+ * proto carries scalars and nested `ListLiteral`s but no map values, so a
Literal whose type
+ * contains a `MapType` has to be expanded before serialization. Teaching
the proto to transport
+ * maps directly would remove the need for this rewrite and for the declines
below:
+ * https://github.com/apache/datafusion-comet/issues/1937
+ *
+ * The rebuilt tree is the tree Spark itself had before `ConstantFolding`
collapsed it, down to
+ * every container's declared nullability (see [[withNullability]]). That
equivalence is the
+ * safety property this rewrite rests on: whatever the rebuilt expression
does natively is what
+ * the same query already does with `ConstantFolding` disabled, so expansion
cannot introduce a
+ * folding-only behaviour difference. Shapes the native map kernels cannot
handle are therefore
+ * declined by their own serdes rather than here, see [[MapKeySupport]] and
[[CometMapEntries]].
+ *
+ * Declined shapes:
+ * - Null values and empty top-level containers: a synthesized `Create*`
with no children
+ * cannot recover the original element type. Empty `ArrayType` literals
still serialize via
+ * `makeListLiteral`, which keeps the type.
+ * - Arrays whose elements are structs, because [[needsExpansion]] does
not walk into a
+ * `StructType`. Native `CreateNamedStruct` builds a 1-row `StructArray`
whenever all of its
+ * children are scalars (`values_to_arrays`), which collides with the
surrounding batch's
+ * row count, and its proto message carries no type, so Spark's declared
field nullability
+ * cannot survive the wire either way. A struct that is only a map value
is safe:
+ * `CometCreateMap` hands the whole rebuilt `CreateMap` to the JVM
codegen dispatcher, so
+ * Spark's own code builds the struct.
+ * - Folded maps with duplicate keys, see [[hasDuplicateMapKeys]].
+ */
+ private def expandComplexLiteral(expr: Literal): Option[Expression] = {
+ if (expr.value == null) return None
+ expr.dataType match {
+ case ArrayType(et, containsNull) if needsExpansion(et) =>
+ val arr = expr.value.asInstanceOf[ArrayData]
+ if (arr.numElements() == 0) {
+ None
+ } else {
+ val elements = (0 until arr.numElements())
+ .map(i => withNullability(literalAt(arr, i, et), containsNull))
+ Some(CreateArray(elements, useStringTypeWhenEmpty = false))
+ }
+ case MapType(kt, vt, valueContainsNull) =>
+ val mapData = expr.value.asInstanceOf[MapData]
+ val keys = mapData.keyArray()
+ if (mapData.numElements() == 0 || hasDuplicateMapKeys(keys, kt)) {
+ None
+ } else {
+ val values = mapData.valueArray()
+ val children = (0 until keys.numElements()).flatMap(i =>
+ Seq(
+ literalAt(keys, i, kt),
+ withNullability(literalAt(values, i, vt), valueContainsNull)))
+ Some(CreateMap(children, useStringTypeWhenEmpty = false))
+ }
+ case _ => None
+ }
+ }
+
+ /**
+ * True when a Literal of this type has to be expanded rather than
serialized, because the
+ * native `Literal` proto carries no map values. Walks array nesting only:
an array of structs
+ * is not expandable (see [[expandComplexLiteral]]), so the walk stops at a
`StructType`.
+ */
+ private def needsExpansion(dataType: DataType): Boolean = dataType match {
+ case _: MapType => true
+ case ArrayType(et, _) => needsExpansion(et)
+ case _ => false
+ }
+
+ /** Element `i` of `arr`, or `null` for a null slot. */
+ private def valueAt(arr: ArrayData, i: Int, dt: DataType): Any =
+ if (arr.isNullAt(i)) null else arr.get(i, dt)
+
+ /** Element `i` of `arr` as a Literal of type `dt`. */
+ private def literalAt(arr: ArrayData, i: Int, dt: DataType): Literal =
+ Literal(valueAt(arr, i, dt), dt)
+
+ /**
+ * True when the folded map's key array holds duplicates under Spark's own
key equality, in
+ * which case rebuilding the value as a `CreateMap` would change semantics
the folded `MapData`
+ * had already settled: `ArrayBasedMapBuilder` throws `DUPLICATED_MAP_KEY`
under
+ * `MAP_KEY_DEDUP_POLICY=EXCEPTION` and silently drops the earlier entry
under `LAST_WIN`, where
+ * the original literal (from `from_json`, or a cast of one) kept both
entries.
+ *
+ * `ArrayBasedMapBuilder` hashes most atomic keys but compares `BinaryType`,
collated
+ * `StringType` and complex keys through `TypeUtils.getInterpretedOrdering`,
and it normalizes
+ * `-0.0` and `NaN` before either. The interpreted ordering is the one
comparison that agrees
+ * with all of those: `Array[Byte]` keys compare by content rather than by
identity, a collated
+ * `StringType` compares under its collation, and `nanSafeCompareDoubles`
already treats `-0.0`
+ * and `+0.0`, and any two `NaN`s, as equal. A null key cannot occur in a
map Spark built, so
+ * treat one as a duplicate and keep the projection on Spark instead of
asking the ordering to
+ * compare it.
+ */
+ private def hasDuplicateMapKeys(keys: ArrayData, keyType: DataType): Boolean
= {
+ val n = keys.numElements()
+ if (n < 2) {
+ false
+ } else if ((0 until n).exists(keys.isNullAt)) {
+ true
+ } else {
+ val ordering = TypeUtils.getInterpretedOrdering(keyType)
+ val sorted = (0 until n).map(i => keys.get(i, keyType)).sorted(ordering)
+ sorted.sliding(2).exists(pair => ordering.compare(pair.head, pair.last)
== 0)
Review Comment:
[P2] Decline non-orderable map keys before requesting an ordering
`TypeUtils.getInterpretedOrdering` is not defined for every valid Spark
map-key type. Spark supports `CalendarIntervalType` keys using hash equality,
so normal folding produces a valid two-entry literal for:
```sql
SELECT id, map(make_interval(1), 1, make_interval(2), 2) AS m FROM t
```
With Parquet `id` values `1, 2, 3`, this new call throws `Type
PhysicalCalendarIntervalType does not support ordered operations` from
`CometLiteral.getSupportLevel`, before the dispatcher can decline the
unsupported type. The keys are distinct and non-null. I reproduced the planning
failure with the exact current serde on Spark 3.5.9 and 4.1.3; Spark and the
exact-base serde return all three rows correctly. Please retain fallback for
non-orderable key types, or use the matching Spark hash/equality path instead
of unconditionally requesting an ordering.
##########
spark/src/main/scala/org/apache/comet/serde/literals.scala:
##########
@@ -215,4 +225,129 @@ object CometLiteral extends CometExpressionSerde[Literal]
with Logging {
}
listLiteralBuilder
}
+
+ /**
+ * Rebuild a folded complex Literal as an equivalent tree of `CreateArray` /
`CreateMap` over
+ * primitive-typed Literals, or `None` when the shape cannot be rebuilt. The
native `Literal`
+ * proto carries scalars and nested `ListLiteral`s but no map values, so a
Literal whose type
+ * contains a `MapType` has to be expanded before serialization. Teaching
the proto to transport
+ * maps directly would remove the need for this rewrite and for the declines
below:
+ * https://github.com/apache/datafusion-comet/issues/1937
+ *
+ * The rebuilt tree is the tree Spark itself had before `ConstantFolding`
collapsed it, down to
+ * every container's declared nullability (see [[withNullability]]). That
equivalence is the
+ * safety property this rewrite rests on: whatever the rebuilt expression
does natively is what
+ * the same query already does with `ConstantFolding` disabled, so expansion
cannot introduce a
+ * folding-only behaviour difference. Shapes the native map kernels cannot
handle are therefore
+ * declined by their own serdes rather than here, see [[MapKeySupport]] and
[[CometMapEntries]].
+ *
+ * Declined shapes:
+ * - Null values and empty top-level containers: a synthesized `Create*`
with no children
+ * cannot recover the original element type. Empty `ArrayType` literals
still serialize via
+ * `makeListLiteral`, which keeps the type.
+ * - Arrays whose elements are structs, because [[needsExpansion]] does
not walk into a
+ * `StructType`. Native `CreateNamedStruct` builds a 1-row `StructArray`
whenever all of its
+ * children are scalars (`values_to_arrays`), which collides with the
surrounding batch's
+ * row count, and its proto message carries no type, so Spark's declared
field nullability
+ * cannot survive the wire either way. A struct that is only a map value
is safe:
+ * `CometCreateMap` hands the whole rebuilt `CreateMap` to the JVM
codegen dispatcher, so
+ * Spark's own code builds the struct.
+ * - Folded maps with duplicate keys, see [[hasDuplicateMapKeys]].
+ */
+ private def expandComplexLiteral(expr: Literal): Option[Expression] = {
+ if (expr.value == null) return None
+ expr.dataType match {
+ case ArrayType(et, containsNull) if needsExpansion(et) =>
+ val arr = expr.value.asInstanceOf[ArrayData]
+ if (arr.numElements() == 0) {
+ None
+ } else {
+ val elements = (0 until arr.numElements())
+ .map(i => withNullability(literalAt(arr, i, et), containsNull))
+ Some(CreateArray(elements, useStringTypeWhenEmpty = false))
+ }
+ case MapType(kt, vt, valueContainsNull) =>
+ val mapData = expr.value.asInstanceOf[MapData]
+ val keys = mapData.keyArray()
+ if (mapData.numElements() == 0 || hasDuplicateMapKeys(keys, kt)) {
+ None
+ } else {
+ val values = mapData.valueArray()
+ val children = (0 until keys.numElements()).flatMap(i =>
+ Seq(
+ literalAt(keys, i, kt),
+ withNullability(literalAt(values, i, vt), valueContainsNull)))
+ Some(CreateMap(children, useStringTypeWhenEmpty = false))
Review Comment:
[P2] Preserve NULL short-circuiting for newly admitted map lookups
Admitting the nested map literal exposes eager evaluation of a lookup key
that Spark skips after a NULL map. With ANSI enabled and Parquet `id` values
`1, 2, 3`:
```sql
SELECT id, element_at(
element_at(map(1, map(0, 7)), id),
id % (id - 2)) AS v
FROM t
```
Spark returns `7, NULL, NULL`: for `id=2`, the inner lookup returns NULL, so
the outer `ElementAt` never evaluates the remainder. The native `map_extract`
scalar expression evaluates every argument over the batch first and raises a
divide/remainder-by-zero error. All map keys are integers, so the new key-type
guards do not prevent it. I reproduced this with the exact current serde on
Spark 3.5.9 and 4.1.3 with `CometProject` asserted; the exact-base serde and a
nonzero-divisor control pass. Please preserve per-row lazy key evaluation, or
retain fallback for these newly admitted expressions until the native consumer
has Spark's NULL behavior.
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -502,20 +503,25 @@ object CometCreateArray extends
CometExpressionSerde[CreateArray] {
}
/**
- * Rewrites a type so that container nullability (`ArrayType.containsNull`,
- * `MapType.valueContainsNull`) is forced to `true` everywhere, while struct
field nullability
- * is left intact. Two CreateArray children whose types differ ONLY in
container nullability are
- * tolerated by DataFusion's `make_array` (coerced), so they normalize equal
here; a difference
- * in a struct field's nullability survives normalization and triggers the
decline above.
+ * Rewrites a type so that `ArrayType.containsNull` is forced to `true`
everywhere, while struct
+ * field nullability and `MapType.valueContainsNull` are left intact. Two
CreateArray children
+ * whose types differ ONLY in `containsNull` are tolerated by DataFusion's
`make_array`, because
+ * `type_union_resolution_coercion` routes `List` arguments through
`list_coercion`, which
+ * merges the element field's nullability and lets Comet's planner insert
the unifying cast. Its
+ * fallback chain has no `map_coercion` arm in DataFusion 54.1, so two `Map`
arguments that
+ * differ in value nullability fail to unify, no cast is inserted, and
`MutableArrayData`
+ * panics. Keep `valueContainsNull` significant here so those children are
declined instead. The
+ * same is true of a struct field's nullability, which `coerce_struct_by_*`
would merge but
+ * which `MutableArrayData` still rejects for the array's own element type.
*/
private def normalizeContainerNullability(dt: DataType): DataType = dt match
{
case ArrayType(elementType, _) =>
ArrayType(normalizeContainerNullability(elementType), containsNull =
true)
- case MapType(keyType, valueType, _) =>
+ case MapType(keyType, valueType, valueContainsNull) =>
MapType(
normalizeContainerNullability(keyType),
normalizeContainerNullability(valueType),
- valueContainsNull = true)
+ valueContainsNull)
Review Comment:
[P2] Keep array nullability inside map types significant
The recursion still erases `ArrayType.containsNull` inside a `MapType`, even
though DataFusion 54.1 cannot coerce two differing map types. With normal
folding and a Parquet `id INT` column containing `1, 2, 3`:
```sql
SELECT id, array(map(1, array(1)), map(2, array(id))) AS a FROM t
```
The folded map has value type `ArrayType(IntegerType,false)` and the dynamic
sibling has `ArrayType(IntegerType,true)`. Both maps retain
`valueContainsNull=false`. Normalization admits them, but native `make_array`
receives unequal Arrow map types and panics with `Arrays with inconsistent
types passed to MutableArrayData`. This reproduces with the exact current serde
on Spark 3.5.9 and 4.1.3 with `CometProject` asserted; replacing `id` with
`coalesce(id,0)` passes, as does the exact-base serde retaining Spark
projection. Unlike the earlier `valueContainsNull` finding, reconstruction
preserves both original map flags here. Please preserve nested array
nullability across the map boundary, or unify the complete map types before
calling `make_array`.
--
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]
