You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/09/05 16:26:08 UTC
[PR] feat: fuse the typed Dataset map sandwich into a Comet projection [datafusion-comet]
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: fuse the typed Dataset map sandwich into a Comet projection [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove closed pull request #5714: feat: fuse the typed Dataset map sandwich into a Comet projection
URL: https://github.com/apache/datafusion-comet/pull/5714
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: fuse the typed Dataset map sandwich into a Comet projection [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #5714:
URL: https://github.com/apache/datafusion-comet/pull/5714#issuecomment-5553599223
Pushed a review pass. One finding was a real limitation rather than a cleanup, so flagging it here:
**`canHandle`'s `spark.sql.codegen.maxFields` gate counted `BoundReference` occurrences, not distinct ordinals.** The fused struct puts the shared deserializer in every field deliberately, so a 12-column typed map counted 12 + 12*12 = 156 against the default cap of 100 and silently declined to fuse. The feature capped out around 9 columns and nothing caught it, since the widest record in the suite had 4 fields. Fixed at the root: WSCG gates on the operator's schema, where a column read twice contributes once, and the kernel likewise emits one typed field and one getter per ordinal. Now pinned by a 12-column test that I confirmed fails without the fix. Note this touches shared code, so it also slightly relaxes the gate for `ScalaUDF` — in the correct direction.
Also worth a reviewer's attention, and deliberately **not** changed here:
**This rewrite could live in `convertNode` instead of as an `_apply` pre-pass, and probably should.** My original comment justified the pre-pass by claiming the bottom-up walk would have already fallen the island back — that is wrong, and I have corrected it. `convertNode` only tags `DeserializeToObjectExec` with a fallback reason; it never replaces it, so the sandwich is structurally intact when the walk reaches `SerializeFromObjectExec`. Doing the fusion there would let the rule see whether `deserialize.child` actually converted to a `CometNativeExec` and decline when it did not — which is exactly the profitability signal this rewrite currently lacks and the reason the config is off by default. That is a ~55 line change and a behavior change (it would decline more often), so I left it out of a cleanup pass, but it is the thing most likely to let the flag default to `true`.
The rest is cleanup: extracted `CometScalaUDF.canDispatch` (the rule was re-deriving the bind-then-`canHandle` gate and had already drifted — it omitted the `RuntimeReplaceable` unwrap); moved `FORCE_DISPATCH` onto `QueryPlanSerde` which reads it, matching what every other behavior tag in the codebase does, and documented that the tag bypasses the per-expression policy layer; moved the tag check ahead of the version shim so the override is unconditional; moved the pre-pass above `normalizePlan` so the synthesized projections get the same NaN / -0.0 treatment as every other `ProjectExec`.
On tests: one assertion was vacuous — `Dataset.schema` comes from the analyzed plan, which a physical rule cannot touch — so it now compares the executed plan's output attributes instead. Both decline paths now pin their message via `checkSparkAnswerAndFallbackReason` rather than just asserting the sandwich survived.
17 tests, green on the default profile and `-Pspark-3.5`; compiles on 3.4/3.5/4.0/default. No regressions in `CometCodegenSuite` (86), `CometCodegenSourceSuite` (60), `CometCodegenFuzzSuite` (28), `CometExpressionSuite` (141), `CometExecRuleSuite` (29).
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: fuse the typed Dataset map sandwich into a Comet projection [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #5714:
URL: https://github.com/apache/datafusion-comet/pull/5714#issuecomment-5970550332
Closing this in favor of #6564.
Since #6542, the codegen dispatcher only calls into Spark's own classes, so the fused tree this PR builds now falls back to Spark. I confirmed it by merging main locally: 10 of the 18 tests in `CometTypedDatasetSuite` fail with `invoke calls scala.Function1, which is not part of Spark`. Getting it working again would take an exemption to #6542's policy plus special cases for decimal serializers.
#6564 takes a simpler route. It leaves the typed operation on Spark and converts the output of `SerializeFromObjectExec` to Arrow, so the operators above it can run natively. Measured with this PR's benchmark, it got 93-100% of the fused speedup wherever either approach paid off, and it covers `flatMap`, `mapPartitions`, `mapGroups` and `cogroup` as well as `map`. A trimmed version of the benchmark moved over with 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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org