apache / apache/datafusion-comet

Fuse the SerializeFromObject / MapElements / DeserializeToObject sandwich into a Comet projection instead of falling back

Open
#5,710 0 comments 0 reactions 1 assignee Claimed by @andygrove View on GitHub
area:expressions enhancement performance requires-triage
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 4h
Merged PRs (30d)
198

Description

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 , 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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.