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

[andygrove]

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

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