sunchao commented on code in PR #5452:
URL: https://github.com/apache/datafusion-comet/pull/5452#discussion_r3849285190
##########
spark/src/main/scala/org/apache/comet/serde/literals.scala:
##########
@@ -215,4 +225,104 @@ 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
+ *
+ * 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.
+ * - `StructType` at any depth. 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. Its proto message also carries no
type, so Spark's
+ * declared field nullability cannot survive the wire either way.
+ * - Map key types whose Spark equality semantics native lookup cannot
honor, see
+ * [[hasUnsafeMapKeyType]].
+ * - Folded maps with duplicate keys. The rebuilt `CreateMap` evaluates
through
+ * `ArrayBasedMapBuilder`, which throws under
`MAP_KEY_DEDUP_POLICY=EXCEPTION`, where the
+ * original literal had already folded cleanly.
+ */
+ private def expandComplexLiteral(expr: Literal): Option[Expression] = {
+ if (expr.value == null) return None
+ expr.dataType match {
+ case ArrayType(et, _) if needsExpansion(et) =>
+ val arr = expr.value.asInstanceOf[ArrayData]
+ if (arr.numElements() == 0) {
+ None
+ } else {
+ val elements = (0 until arr.numElements()).map(i =>
asNullable(literalAt(arr, i, et)))
+ Some(CreateArray(elements, useStringTypeWhenEmpty = false))
+ }
+ case MapType(kt, vt, _) =>
+ val mapData = expr.value.asInstanceOf[MapData]
+ val keys = mapData.keyArray()
+ if (mapData.numElements() == 0 || hasUnsafeMapKeyType(kt) ||
+ hasDuplicateMapKeys(keys, kt)) {
+ None
+ } else {
+ val values = mapData.valueArray()
+ val children = (0 until keys.numElements()).flatMap(i =>
+ Seq(literalAt(keys, i, kt), asNullable(literalAt(values, i, vt))))
+ 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: a
`StructType` is not
+ * expandable at all (see [[expandComplexLiteral]]), so the walk stops there.
+ */
+ 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 map key type has Spark equality semantics that native map
lookup (Arrow
+ * bytewise comparison) cannot honor, in which case declining expansion
keeps the projection on
+ * Spark. `NormalizeFloatingNumbers` makes Spark treat `+0.0` and `-0.0` as
the same key and
+ * canonicalises NaN, and a non-default collation compares under rules
native applies as
+ * `UTF8_BINARY`. Both are checked at every nesting level of the key type.
+ */
+ private def hasUnsafeMapKeyType(kt: DataType): Boolean =
+ hasNonDefaultStringCollation(kt) ||
Review Comment:
[P2] Preserve nullable components of complex map lookup keys
This admits folded maps whose key type is `ArrayType(IntegerType, false)`,
but Spark permits a dynamic lookup key containing a null element. With normal
folding and Parquet `id` values `1, 2, 3`:
```sql
SELECT id, element_at(
map(array(1), 7),
array(IF(id = 2, CAST(NULL AS INT), id))) AS v
FROM t
```
Spark returns `7, NULL, NULL`. Native `map_extract` instead coerces the
lookup to the literal key's exact Arrow type; the inserted cast aborts with
`Non-nullable field of ListArray "item" cannot contain nulls`. The null is
inside the lookup array, not a null map key. This needs a
nullability-compatible lookup type or fallback for these literals.
Reproduced on Spark 3.5.9 and 4.1.3 with native `CometProject` assertions,
in both ANSI modes. The non-null lookup control and dispatcher-off path pass;
replacing only the literal serializer with the exact base version also restores
the correct rows. No floating-point keys, collations, or duplicate keys are
involved.
##########
spark/src/main/scala/org/apache/comet/serde/literals.scala:
##########
@@ -215,4 +225,104 @@ 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
+ *
+ * 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.
+ * - `StructType` at any depth. 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. Its proto message also carries no
type, so Spark's
+ * declared field nullability cannot survive the wire either way.
+ * - Map key types whose Spark equality semantics native lookup cannot
honor, see
+ * [[hasUnsafeMapKeyType]].
+ * - Folded maps with duplicate keys. The rebuilt `CreateMap` evaluates
through
+ * `ArrayBasedMapBuilder`, which throws under
`MAP_KEY_DEDUP_POLICY=EXCEPTION`, where the
+ * original literal had already folded cleanly.
+ */
+ private def expandComplexLiteral(expr: Literal): Option[Expression] = {
+ if (expr.value == null) return None
+ expr.dataType match {
+ case ArrayType(et, _) if needsExpansion(et) =>
+ val arr = expr.value.asInstanceOf[ArrayData]
+ if (arr.numElements() == 0) {
+ None
+ } else {
+ val elements = (0 until arr.numElements()).map(i =>
asNullable(literalAt(arr, i, et)))
+ Some(CreateArray(elements, useStringTypeWhenEmpty = false))
+ }
+ case MapType(kt, vt, _) =>
+ val mapData = expr.value.asInstanceOf[MapData]
+ val keys = mapData.keyArray()
+ if (mapData.numElements() == 0 || hasUnsafeMapKeyType(kt) ||
+ hasDuplicateMapKeys(keys, kt)) {
+ None
+ } else {
+ val values = mapData.valueArray()
+ val children = (0 until keys.numElements()).flatMap(i =>
Review Comment:
[P2] Guard nested map literals before native map_entries
A map stored inside a map value is handed to the JVM dispatcher as a
literal, so its original `valueContainsNull=false` survives. This enables
another previously falling-back query with normal folding over Parquet `id`
values `1, 2, 3`:
```sql
SELECT id, map_entries(
element_at(map(1, map(1, 2)), id)) AS e
FROM t
```
The inner map reaches native `map_entries`, which declares a nullable
`value` field but reuses entry arrays whose field is non-nullable. Arrow's
`ListArray::new` then panics with a child-type mismatch instead of returning
Spark's `[{1, 2}], NULL, NULL`. There are no mixed map siblings here: the inner
type has not been changed; the incompatible schema is created by its consumer.
Reproduced on Spark 3.5.9 and 4.1.3 with native `CometProject` assertions,
in both ANSI modes. Nullable-inner-map, `map_keys`, and `map_values` controls
pass; the exact-base literal serializer and dispatcher-off path restore Spark's
result. Please retain fallback for unsupported nested-map shapes until their
native consumers preserve matching entry types.
--
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]
