andygrove opened a new issue, #5710:
URL: https://github.com/apache/datafusion-comet/issues/5710
Part of #5572.
`ds.map(f)` drops a three-operator island into the middle of an
otherwise-native plan, and every operator in it falls back. Reproduced on
`81d637b9b` against a Parquet table with a `case class ScratchRec(a: Int, b:
String)`:
```
*(1) SerializeFromObject [invoke(knownnotnull(assertnotnull(input[0,
ScratchRec, true])).a()) AS a#21, \
static_invoke(UTF8String.fromString(invoke(...).b())) AS b#22]
+- *(1) MapElements <lambda>, obj#18: ScratchRec
+- *(1) DeserializeToObject newInstance(class ScratchRec), obj#15:
ScratchRec
+- *(1) CometColumnarToRow
+- CometProject [a#4, b#5], [_1#2 AS a#4, _2#3 AS b#5]
+- CometNativeScan parquet [_1#2,_2#3]
```
Only `DeserializeToObject` carries a fallback reason, because it is the
lowest node whose children are all native; the two above it just inherit the
island. Neither operator has an entry in `CometExecRule.nativeExecs`
(`spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:75`), so both
take the generic "`X` is not supported" path at `:507`.
It cascades past the island. `ds.map(f).groupBy("b").count()` loses the
aggregate and the exchange too:
```
HashAggregate [COMET: Comet aggregate that merges intermediate buffers
requires a Comet child aggregate ...]
+- AQEShuffleRead
+- Exchange
+- HashAggregate
+- SerializeFromObject
+- MapElements
+- DeserializeToObject [COMET: DeserializeToObject is not
supported]
+- CometColumnarToRow
+- CometProject
+- CometNativeScan parquet
Comet accelerated 2 out of 8 eligible operators (25%).
```
## Why the mixin is not the answer here
`DeserializeToObjectExec.output` is exactly one attribute of
`ObjectType(cls)` (Spark's `ObjectProducerExec.output`), and
`SerializeFromObjectExec`'s input is the same. `ObjectType` is absent from
`QueryPlanSerde.supportedDataType` (`QueryPlanSerde.scala:555`) and
`serializeDataType` returns `None` for it (`:612`), and it is outside
`CometBatchKernelCodegen.isSupportedDataType`
(`CometBatchKernelCodegen.scala:85`) as the EPIC already notes. A JVM object
reference cannot live in an Arrow `FieldVector`. Each of these operators has an
object on one side of it, so neither can be a kernel in isolation — the kernel
contract is Arrow-in, Arrow-out.
`CodegenDispatchFallback`'s self-type is also `CometExpressionSerde[_]`
(`CometExpressionSerde.scala:125`), so there is no operator-level dispatch hook
to mix it into. This is the same structural wall the EPIC records for the
aggregate serdes.
## Why the fused sandwich is
Two things make the fused form legal, and neither needs new native code.
`canHandle` only inspects the **root** `dataType`
(`CometBatchKernelCodegen.scala:120`) and every **`BoundReference`** `dataType`
(`:171`). Intermediate nodes are never type-checked. So an `ObjectType` that
exists strictly *inside* the tree passes the gate — it is only the boundary
that has to be Arrow-representable, and for this sandwich the outer boundary is
ordinary SQL data on both sides.
Spark also already builds the fused expression itself.
`MapElementsExec.doConsume` constructs
```scala
val funcObj = Literal.create(func, ObjectType(funcClass))
val callFunc = Invoke(funcObj, funcName, outputObjectType, child.output,
propagateNull = false)
```
so the user closure has a first-class Catalyst representation and we do not
have to invent one. `DeserializeToObjectExec.doConsume` is
`bindReference(deserializer, child.output).genCode(ctx)` and
`SerializeFromObjectExec.doConsume` is the same over each serializer. The fused
tree is therefore just: substitute `Invoke(func, deserializer)` into each
serializer's `BoundReference(0)` — exactly what whole-stage codegen produces by
chaining the three `doConsume`s today.
That means the shape of the change is a **plan rewrite into `ProjectExec`**,
not a new operator serde. Recognize the sandwich, emit a `ProjectExec` whose
expressions are the fused trees, and let the existing `CometProjectExec` plus
`CometScalaUDF.emitJvmCodegenDispatch` path take it from there. No proto
change, no native change.
#5692 landed `Invoke` as a `CometCodegenDispatch[Invoke]`, which is the
piece that makes the middle of the sandwich dispatchable. Note that #5575
explicitly scoped *out* the encoder and deserializer trees on the grounds that
their `ObjectType` boundary makes `canHandle` reject them — correct for a bare
`Invoke`, and this issue is the follow-up that removes the boundary by fusing.
## The one design problem worth deciding up front
`emitJvmCodegenDispatch` emits one `JvmScalarUdf` and gets one output vector
back. `SerializeFromObject` has N serializers that all share the *same*
`Invoke`. Dispatching them as N separate expressions means N closure
invocations per row: N times the cost, and observably wrong for a closure with
side effects, since Spark's fused loop calls it once.
The fix is to make the kernel's single root a
`CreateNamedStruct(serializer)` and project the N fields off it natively with
`GetStructField`. Struct output is already supported —
`CometBatchKernelCodegenOutput.outputVectorClass` maps `StructType` to
`StructVector` (`CometBatchKernelCodegenOutput.scala:180`) — and both
`CometCreateNamedStruct` and `CometGetStructField` already exist
(`serde/structs.scala:34`, `:73`). `Invoke` is deterministic in Catalyst
(`InvokeLike` does not override `deterministic`, and both children are
deterministic), so the `doSubexpressionElimination = true` path in
`generateSource` (`CometBatchKernelCodegen.scala:261`) hoists the shared
`Invoke` into a single per-row call.
## Scope
| Shape | Middle operator | Fusable |
| --- | --- | --- |
| `ds.map(f)`, `ds.map(MapFunction)` | `MapElementsExec` | yes, per-row |
| `ds.mapPartitions`, `ds.flatMap` | `MapPartitionsExec(Iterator[Any] =>
Iterator[Any])` | no — arbitrary row count, no per-row expression exists |
| `groupByKey(...).mapGroups` / `cogroup` | `FlatMapGroupsExec`,
`CoGroupExec` | no — group-at-a-time |
| `groupByKey` key extraction | `AppendColumnsExec` | per-row but widens the
schema to `child.output ++ serializer`; a Project with passthrough, worth a
follow-up rather than this issue |
| `ds.filter(func)` | none — `EliminateSerialization` and `TypedFilter`
collapse to a plain `FilterExec` with an `Invoke` condition | already covered
by #5692 |
So the target is `MapElementsExec`, i.e. `ds.map`. `ds.mapPartitions` was
confirmed in the same run to produce the identical sandwich with
`MapPartitions` in the middle, and it is not reachable this way.
## Payoff, and where it could regress
The closure does not get faster. It is the same `Invoke` on the same JVM
lambda, once per row, in a Janino-compiled kernel loop instead of a
whole-stage-codegen loop. The win is deleting the island, and it is worth
having only when the typed operation sits *between* native operators — the
`groupBy` case above, or a typed map feeding a join.
When the typed operation is at the top of the plan, which
`ds.map(f).collect()` is, the gain is roughly nil: we would trade a
`CometColumnarToRow` before the map for one after it, and it could come out
slightly behind, since the kernel writes into Arrow only for something to
immediately read rows back out. Worth checking whether
`RevertNativeForTransitionHeavyStages` already covers that shape before
assuming a new heuristic is needed.
## Risk to check before writing code
The lesson from #5575 applies directly: the mixin's contract covers whether
`doGenCode` *compiles*, not limits that live at the **Arrow output boundary**.
That is where the Iceberg `truncate(w, decimal)` regression came from.
Encoder serializers for decimals are the case I would check first. A
`BigDecimal` encoder declares `DecimalType(38, 18)`, and the
`StaticInvoke(Decimal.fromDecimal)` in the serializer can produce a value wider
than the declared precision — Spark nulls it at row materialization, but the
Arrow `Decimal128(p, s)` writer the kernel targets may not. Same failure shape
as the Iceberg bug, so it needs a fixture before this lands.
Two smaller ones: `NewInstance` for a nested (non-top-level) case class
carries an `outerPointer` closure over the enclosing instance, which
closure-serialization would drag along — `ScalaUDF` has the same hazard, but it
is worth a test. And any encoder that touches a type outside
`isSupportedDataType` (`NullType`, UDTs, Variant) has to fall back cleanly at
plan time rather than fail the Janino compile at execute time, which is what
the `canHandle` gate is for.
## Verification note
The plan shapes and fallback reasons above were reproduced on `81d637b9b`
with a throwaway suite; the fused-kernel design has not been prototyped. The
`DecimalType` concern in particular is a hypothesis drawn from the #5575
postmortem, not a reproduced bug.
--
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]
