sunchao commented on code in PR #5452:
URL: https://github.com/apache/datafusion-comet/pull/5452#discussion_r3865543786
##########
spark/src/main/scala/org/apache/comet/serde/literals.scala:
##########
@@ -215,4 +225,168 @@ 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.
`getSupportLevel`
+ * probes with this same function and reports `Compatible` when it returns a
`Some`, so
+ * admission and expansion cannot diverge. 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. A non-nullable-value map into
`map_entries` is handled by
+ * the native planner widening its argument rather than declined here.
+ *
+ * Map key semantics, however, are declined here at admission (see
[[mapKeyTypesExpandable]])
+ * rather than left to consumers. A map lookup gates itself on
[[MapKeySupport]], but not every
+ * native map consumer does: `map_contains_key` lowers to
`array_contains(map_keys(...), key)`,
+ * and neither kernel checks the key type. A map that is only the value of
another map is also
+ * opaque once rebuilt, because `CometCreateMap` hands the whole `CreateMap`
to the JVM codegen
+ * dispatcher, so a nested unsupported key type never revisits this serde.
Declining any folded
+ * literal that (recursively) contains an unsupported or non-orderable map
key type keeps all of
+ * those paths on Spark.
+ *
+ * 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.
+ * - Any map key type native map kernels cannot reproduce Spark equality
for (floating-point,
+ * collated string, complex), or whose interpreted ordering is undefined
+ * (`CalendarIntervalType`), at any nesting level. See
[[mapKeyTypesExpandable]].
+ * - 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 || !mapKeyTypesExpandable(expr.dataType)) 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] Retain a safe execution path for large folded map literals
Rebuilding every stored entry as constructor children newly sends large
folded maps through the dispatcher's nested generated classes. With normal
folding and a Parquet `id INT` column containing `1, 2, 3`:
```sql
SELECT id, element_at(
map_from_arrays(sequence(1, 100000), sequence(1, 100000)),
id) AS v
FROM t
```
Spark returns `(1,1), (2,2), (3,3)`, but this head selects `CometProject`
and fails with `IllegalAccessError`:
`SpecificCometBatchKernel$NestedClass_8.CreateMap_0$` accesses the protected
`CometBatchKernel.references` field. The generated helper uses that field for
its map builder. The schema-field-count admission check does not constrain map
entry count, and there is no execution-time fallback.
This reproduced twice on Spark 4.0.4/JDK 17. The exact-base literal
serializer and dispatcher-off controls both return the correct rows; a
10,000-entry native control passes. Please retain fallback for expansions the
dispatcher cannot execute, preserve the compact literal representation, or fix
generated-helper reference access before admitting them.
--
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]
