perf: avoid FFI import/export between native subtree and ShuffleWriter

[mbutrovich]

Which issue does this PR close?

Closes #3925. Replaces #3930.

Rationale for this change

When CometShuffleExchangeExec has a CometNativeExec child, the current code runs two CometExecIterators per partition: one for the upstream subtree, one for Scan("ShuffleWriterInput") -> ShuffleWriter consuming the FFI'd batches back. The JVM doesn't read the data, so the round-trip is pure overhead.

This PR collapses the two iterators into one CometExecIterator rooted at ShuffleWriter(child = childNativeOp): no JVM intermediary, no Arrow FFI export/import, one fewer createPlan/releasePlan pair per partition. The proto already supported this nested form.

Although the Arrow C Data Interface is zero-copy, this native-to-JVM-and-back round-trip was not. The old Scan("ShuffleWriterInput") left arrow_ffi_safe at its proto default of false, so on import ScanExec ran copy_array on every batch: an unconditional deep copy into freshly-sized buffers. Collapsing the iterators removes that per-batch copy.

Why this replaces #3930

#3930 also removed the FFI but kept the two-iterator structure, passing batches between them through a JVM-side u64-keyed Mutex<HashMap<u64, RecordBatch>>. This PR achieves the same result without the global registry, the batch-handle protocol, or the second iterator: the child plan is just a child operator in the same DataFusion plan as the writer.

What changes are included in this PR?

• CometNativeExec exposes NativeExecContext (buildNativeContext() / executeColumnarWithContext(ctx)) so the shuffle path reuses doExecuteColumnar's per-partition setup (broadcast alignment, plan-data injection, subqueries, encryption).
• CometNativeShuffleInputRDD is a thin scheduling-anchor RDD that resolves leaf partitions on the driver and carries per-partition leaf/shuffle-block iterators.
• NativeShuffleSpec (child native op, metric node, exec context) is carried by CometShuffleDependency and only populated on native-child shuffles.
• CometShuffleExchangeExec dispatches: native-child shuffles use prepareNativeShuffleDependency and the thin RDD; everything else uses the existing 5-arg prepareShuffleDependency with a synthesized Scan("ShuffleWriterInput") placeholder. Behavior unchanged for CometSparkToColumnarExec, CometCollectLimitExec, CometTakeOrderedAndProjectExec.
• CometNativeShuffleWriter rewritten to one code path over ShuffleWriter(child = spec.childNativeOp).
• CometNativeShuffleWriter calls reportScanInputMetrics so a native scan now executing inside the writer's plan still populates the task's inputMetrics (bytesRead/recordsRead). The split-driver flow reported this from the scan's own CometExecRDD.compute, which no longer runs once the scan is inlined.

Removing the FFI boundary surfaces four issues the old deep-copy + ScanExec cast was masking:

• Shuffle-writer memory accounting. copy_array compacted every batch into tightly-sized buffers, so the writer's data_size metric and spill reservation (both from RecordBatch::get_array_memory_size) were accurate. With the child inlined, the writer sees the native subtree's buffers directly: a partial HashAggregate emits one group-values buffer sliced into batch_size chunks, and get_array_memory_size charges that buffer's capacity once per slice, overstating resident memory by the chunk count (TPC-H q17 reported a 91.2 TiB data size vs 177.9 GiB on main, with tens of thousands of spurious spills). The writer now charges each distinct backing allocation once via count_new_buffers. The sum of get_slice_memory_size is not a substitute: a slice pins its whole backing buffer, which the group-values Vec rounds up to the next power of two, so live-row size under-counts resident memory.
• DataFusion-vs-catalyst Arrow type drift (width_bucket, date_trunc, collect_set, ...). New SchemaAlignExec between the inlined child and the writer casts per column and widens nullability to actual.nullable || expected.nullable. Each distinct drift logs one process-deduped warning pointing at DataFusion / DataFusion-Spark functions whose Arrow return type drifts from Spark catalyst's declared type #4515. New ShuffleWriter.expected_output_schema proto field.
• Catalyst-pruned HashAggregateExec.resultExpressions (EXISTS / row-existence-only subqueries) where the JVM-declared output=[] disagrees with the native aggregate's natural shape. Fixed in CometBaseAggregate.doConvert: when natural shape differs from declared shape, emit an explicit Projection proto op above the HashAgg. Native OpStruct::HashAgg always returns its natural shape. HashAggregate.result_exprs (field 3) and apply_result_projection (field 8) reserved.
• Typed exception loss through ShuffleWriterExec. With the writer now the root operator for any stage that ends in a shuffle, an upstream DataFusionError::External(SparkError) (e.g. decimal overflow in ANSI mode) was being re-wrapped by external_shuffle's ArrowError::ExternalError adapter, hiding the typed exception from the JNI bridge's single-level downcast. Drop the wrap; try_flatten already handles the matching DataFusionError types directly.

How are these changes tested?

Existing tests, plus:

• New CometAggregateSuite regression test for the catalyst-pruned aggregate path. The original Spark SQL test that surfaced the aggregate bug (subquery/exists-subquery/exists-orderby-limit.sql query 19) is also covered by the SQL-tests harness CI.
• New multi_partition Rust test (shuffle_partitioner_charges_shared_buffer_once): a batch sliced into chunks that share one backing buffer is charged once, so it stays under the memory limit and reports a data_size that tracks the buffer rather than the chunk count.
• New CometTaskMetricsSuite test (native shuffle reports task input metrics for its scan child): a native shuffle over a CometNativeScanExec populates the ShuffleMapTask's inputMetrics. It fails (zero recordsRead) without the reportScanInputMetrics call.

[mbutrovich]

Ran TPC-H SF1000 with Comet on Spark 3.5.8:

Query	main (s)	#4507 (s)
TPCH-01	10.495	10.151
TPCH-02	23.455	22.177
TPCH-03	12.366	14.064
TPCH-04	14.067	13.360
TPCH-05	30.023	29.119
TPCH-06	1.128	0.710
TPCH-07	16.368	15.073
TPCH-08	36.566	35.680
TPCH-09	44.722	41.845
TPCH-10	29.848	29.106
TPCH-11	15.373	14.837
TPCH-12	7.919	7.486
TPCH-13	10.194	10.103
TPCH-14	2.631	2.672
TPCH-15	11.613	12.524
TPCH-16	11.207	9.778
TPCH-17	31.002	28.755
TPCH-18	62.856	56.601
TPCH-19	9.253	9.319
TPCH-20	9.015	8.610
TPCH-21	74.530	69.337
TPCH-22	10.640	9.574
Total	475.272	450.879

Looks to be a decent improvement.

[mbutrovich]

TPC-DS SF1000 is a little over 6% faster as well.

Query	main (s)	#4507 (s)
TPCDS-01	4.303	4.239
TPCDS-02	5.111	5.413
TPCDS-03	1.571	1.538
TPCDS-04	22.027	21.094
TPCDS-05	5.552	5.723
TPCDS-06	1.838	1.693
TPCDS-07	2.503	2.379
TPCDS-08	1.601	1.454
TPCDS-09	9.673	9.771
TPCDS-10	2.920	2.973
TPCDS-11	11.558	11.298
TPCDS-12	1.101	1.083
TPCDS-13	2.528	2.521
TPCDS-14a	31.533	28.518
TPCDS-14b	31.786	27.090
TPCDS-15	1.639	1.815
TPCDS-16	6.272	6.292
TPCDS-17	3.577	3.671
TPCDS-18	3.092	3.716
TPCDS-19	1.648	1.585
TPCDS-20	1.087	1.109
TPCDS-21	0.779	0.721
TPCDS-22	2.982	2.902
TPCDS-23a	35.984	34.227
TPCDS-23b	46.507	45.944
TPCDS-24a	22.057	22.422
TPCDS-24b	23.459	20.861
TPCDS-25	2.376	2.361
TPCDS-26	1.312	1.491
TPCDS-27	2.193	2.259
TPCDS-28	12.102	11.772
TPCDS-29	6.224	6.129
TPCDS-30	2.216	2.431
TPCDS-31	4.839	4.737
TPCDS-32	1.094	0.890
TPCDS-33	1.404	1.351
TPCDS-34	2.167	1.932
TPCDS-35	4.462	4.379
TPCDS-36	2.782	2.720
TPCDS-37	2.912	2.412
TPCDS-38	6.699	6.255
TPCDS-39a	5.590	2.514
TPCDS-39b	4.184	2.883
TPCDS-40	3.036	2.812
TPCDS-41	0.459	0.421
TPCDS-42	0.877	0.898
TPCDS-43	2.322	2.032
TPCDS-44	7.295	6.833
TPCDS-45	1.722	1.749
TPCDS-46	2.902	2.633
TPCDS-47	4.631	4.533
TPCDS-48	2.403	2.208
TPCDS-49	3.760	3.921
TPCDS-50	16.671	16.843
TPCDS-51	6.923	6.748
TPCDS-52	0.886	0.766
TPCDS-53	2.473	2.014
TPCDS-54	5.476	5.625
TPCDS-55	0.964	0.835
TPCDS-56	1.543	1.148
TPCDS-57	3.646	3.509
TPCDS-58	1.456	1.267
TPCDS-59	7.807	7.210
TPCDS-60	1.775	1.516
TPCDS-61	1.972	1.735
TPCDS-62	3.638	3.172
TPCDS-63	2.009	1.786
TPCDS-64	11.576	11.327
TPCDS-65	5.588	5.338
TPCDS-66	2.639	2.474
TPCDS-67	32.790	30.212
TPCDS-68	5.592	1.607
TPCDS-69	2.273	2.247
TPCDS-70	6.047	4.519
TPCDS-71	2.236	1.356
TPCDS-72	9.072	8.243
TPCDS-73	1.262	1.112
TPCDS-74	8.803	8.449
TPCDS-75	15.381	14.322
TPCDS-76	6.932	6.637
TPCDS-77	1.489	1.469
TPCDS-78	15.289	14.661
TPCDS-79	3.082	2.071
TPCDS-80	5.332	5.350
TPCDS-81	3.647	3.526
TPCDS-82	4.227	3.861
TPCDS-83	0.983	0.962
TPCDS-84	2.682	2.223
TPCDS-85	4.294	3.429
TPCDS-86	1.504	1.396
TPCDS-87	6.849	6.398
TPCDS-88	21.062	20.543
TPCDS-89	2.041	2.069
TPCDS-90	3.313	3.734
TPCDS-91	0.782	0.846
TPCDS-92	1.077	0.975
TPCDS-93	20.542	19.973
TPCDS-94	5.443	5.285
TPCDS-95	19.409	18.508
TPCDS-96	4.313	3.873
TPCDS-97	4.889	4.643
TPCDS-98	1.496	1.433
TPCDS-99	4.055	3.276
Total	681.881	639.129

[andygrove]

@mbutrovich WDYT about adding some unit tests for SchemaAlignExec?

[mbutrovich]

I would love for this class to be short-lived and eventually be gone, but yeah I can add some tests in the meantime.

[mbutrovich]

Added Rust unit tests for SchemaAlignExec in schema_align.rs, one per concrete drift tracked in #4515:

• aligns_width_bucket_int32_to_int64 - top-level Int32 -> Int64 cast, values preserved.
• aligns_date_trunc_timestamp_timezone - Timestamp(us) -> Timestamp(us, "UTC"), instants unchanged.
• aligns_collect_set_list_element_nullability - the interesting one: List(nullable Int32) -> List(non-null Int32). This narrows the inner element nullability, and I wasn't sure arrow's list cast would do that rather than fail the re-stamp on DataType inequality. It does. The test confirms the cast reshapes the inner field and the batch re-stamps cleanly, so the operator genuinely covers the nested collect_set case, not just the top-level ones.

All three pass. The module doc-comment also records the intent: each test is anchored to a drift pair, so when a function's upstream return type is corrected the pair collapses, try_new_or_passthrough returns the child unwrapped, and the matching test flips to a passthrough assertion. That is the signal the workaround for that function can be removed, which lines up with wanting this class to eventually go away.

[andygrove]

Thanks @mbutrovich. I've been wanting to see this optimization for a while now. Left a comment about adding unit tests but will leave that up to you.

[andygrove]

Thanks for adding the tests!