andygrove opened a new pull request, #5714:
URL: https://github.com/apache/datafusion-comet/pull/5714
## Which issue does this PR close?
Closes #5710. Part of #5572.
## Rationale for this change
`ds.map(f)` drops a three-operator island into the middle of an otherwise
native plan, and every operator in it falls back:
```
*(1) SerializeFromObject [invoke(knownnotnull(assertnotnull(input[0,
TypedRec, true])).a()) AS a#21, ...]
+- *(1) MapElements <lambda>, obj#18: TypedRec
+- *(1) DeserializeToObject newInstance(class TypedRec), obj#15: TypedRec
+- *(1) CometColumnarToRow
+- CometProject [a#4, b#5], [_1#2 AS a#4, _2#3 AS b#5]
+- CometNativeScan parquet [_1#2,_2#3]
```
It cascades past the island. `ds.map(f).groupBy("b").count()` loses the
aggregate and the exchange too — 2 of 8 eligible operators accelerated.
Neither object operator can be handled on its own. `DeserializeToObjectExec`
outputs a single `ObjectType` attribute and `SerializeFromObjectExec` consumes
one; `ObjectType` is outside `QueryPlanSerde.supportedDataType` and outside
`CometBatchKernelCodegen.isSupportedDataType`, because a JVM object reference
cannot live in an Arrow vector. `CodegenDispatchFallback`'s self-type is
`CometExpressionSerde[_]`, so there is no operator-level dispatch hook to mix
in either.
Fusing the sandwich works, though, for two reasons. `canHandle` only
type-checks the **root** `dataType` and every **`BoundReference`** —
intermediate nodes are never inspected — so the object may exist strictly
_inside_ the tree while the outer boundary stays ordinary SQL data. And Spark
already builds the fused expression: `MapElementsExec.doConsume` constructs
`Invoke(Literal.create(func, ObjectType(funcClass)), funcName,
outputObjectType, child.output, propagateNull = false)`, so the closure has a
first-class Catalyst representation. The rule just does statically what
whole-stage codegen does by chaining the three `doConsume`s.
## What changes are included in this PR?
`RewriteTypedDatasetMap` (new), run as a pre-pass in `CometExecRule._apply`
before `transform`, since the bottom-up conversion would otherwise reach
`DeserializeToObjectExec` and fall the island back before anything could fuse
it. It rewrites the sandwich into a projection over the deserializer's child.
The projection then converts through the ordinary `CometProjectExec` path and
the fused tree routes through the JVM codegen dispatcher — **no proto change
and no native change**.
**Multiple output columns.** `emitJvmCodegenDispatch` emits one
`JvmScalarUdf` per expression and gets one Arrow vector back, but
`SerializeFromObject` has N serializers that all share one `Invoke`. Converting
them separately would produce N kernels and call the user closure N times per
row where Spark calls it once — a real behavioural difference for a
side-effecting closure, not just a slowdown. So for N > 1 the rule emits two
stacked projections: an inner one producing a single `CreateNamedStruct`, and
an outer one extracting the fields with `GetStructField`. Both node types
already have serdes, and `CometBatchKernelCodegenOutput` already maps
`StructType` to `StructVector`.
**`CometScalaUDF.FORCE_DISPATCH`.** The struct needs to compile into _one_
kernel, but `CreateNamedStruct` has a perfectly good native serde that would
convert each field independently and put us back at N kernels. A `TreeNodeTag`
checked at the top of `exprToProtoInternal` forces whole-subtree dispatch. I
used a tag rather than a Comet-specific `Expression` subclass deliberately: the
rewritten plan stays built entirely from stock Spark expressions, so it still
executes correctly if the enclosing operator falls back to Spark for an
unrelated reason.
**Scope.** Only `MapElementsExec`, including a chain of them
(`ds.map(f).map(g)` leaves two adjacent `MapElements` under one
Serialize/Deserialize pair). `MapPartitionsExec`, `FlatMapGroupsExec` and
`CoGroupExec` consume iterators or groups, so no per-row expression exists.
`AppendColumnsExec` is per-row but widens the schema; left for a follow-up.
Typed filters never produce this sandwich — Catalyst lowers them to a
`FilterExec` over an `Invoke`, which #5692 already dispatches.
**Declines rather than guesses** when the serializer has a non-`Alias`
element, reads an unexpected bound reference, produces a dangling attribute
reference, `canHandle` refuses the tree, the dispatcher is disabled, or (for N
> 1) subexpression elimination is off. Each records a fallback reason so
`EXPLAIN` says why instead of showing the bare "not supported".
**Off by default** (`spark.comet.exec.typedDatasetMap.enabled`). The rewrite
pays off when the typed operation sits between native operators; when it is at
the top of the plan (`ds.map(f).collect()`) the gain is roughly nil and could
be slightly negative, since the kernel writes into Arrow only for something to
read rows straight back out. Flipping the default needs benchmarks, and it is
worth checking whether `RevertNativeForTransitionHeavyStages` already covers
that shape.
## How are these changes tested?
New `CometTypedDatasetSuite`, 16 tests, registered in both
`pr_build_linux.yml` and `pr_build_macos.yml`. Green on the default profile and
on `-Pspark-3.5`; main and test code compile on `spark-3.4`, `spark-3.5`,
`spark-4.0` and the default. No regressions in `CometCodegenSuite` (86),
`CometCodegenSourceSuite` (60), `CometExpressionSuite` (141),
`CometExecRuleSuite` (29) or `CometCoverageStatsSuite`.
Beyond `checkSparkAnswerAndOperator` on single-column, multi-column, wide
(decimal / string / `Option`), nested-struct-and-array, chained-map and Java
`MapFunction` shapes, the tests that carry the most weight:
- **`closure runs exactly once per row with multiple output columns`** — a
JVM-static counter asserts 50 calls for 50 rows. This is the test that would
catch the N-kernel regression the struct wrapper exists to prevent.
- **`fusion unblocks the aggregate and shuffle above it`** — asserts the
`groupBy().count()` case is fully native, i.e. the cascade is actually fixed.
- **`output schema is unchanged by the rewrite`** — compares the schema tree
with the flag on and off, guarding the nullability reasoning behind
`GetStructField` over a non-nullable `CreateNamedStruct`.
- **`AssertNotNull inside the fused kernel still raises like Spark`** —
returning null for a non-nullable product is an error in Spark; the
serializer's `assertnotnull` has to survive the fuse or Comet would silently
emit a null row. The stack trace confirms it fires from
`SpecificCometBatchKernel.subExpr_0$`, which incidentally also confirms CSE
hoisted the shared `Invoke`.
- **`decimal overflow in the serializer is caught before the Arrow write`**
— I raised this on #5710 as the #5575-shaped risk: an encoder-declared
`decimal(38,18)` receiving a wider value that Spark nulls at row
materialization but the kernel's `DecimalVector` write might not. **It does not
apply.** The encoder's serializer already wraps the value in `CheckOverflow`,
so the fused tree raises under ANSI and nulls under non-ANSI exactly where
Spark does, ahead of the write. The test pins both directions. I verified the
same is true on the pre-existing plain-`ScalaUDF` decimal path.
- Negative tests for each decline path: off by default, dispatcher disabled,
CSE disabled, and `mapPartitions` left alone.
Not covered, and worth a reviewer's attention: no benchmark numbers yet,
which is why the flag is off. Nested (non-top-level) case classes carry a
`NewInstance` `outerPointer` closure over the enclosing instance that closure
serialization would drag along — `ScalaUDF` has the same exposure, but I have
not written a test for it.
--
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]
