feat: fuse the typed Dataset map sandwich into a Comet projection

[andygrove]

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

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 Fuse the SerializeFromObject / MapElements / DeserializeToObject sandwich into a Comet projection instead of falling back #5710 as the Route unrecognized StaticInvoke and Invoke through the codegen dispatcher instead of falling back #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.

[andygrove]

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