[FEA] Enable per-expression legacy AST projection [databricks] [fast-ut]

[thirtiseven]

Related to #8954.

Description

The motivation for this PR is to prepare for AST JIT work.

Legacy AST Project execution is currently all-or-nothing: if any top-level Project output cannot be represented by the legacy AST backend, eligible sibling outputs cannot use AST either.

This change selects the legacy AST backend independently for each top-level Project expression. Supported fixed-width outputs use AST while unsupported outputs in the same Project continue through the regular GPU expression path. Top-level null literals also remain on the regular projection path so their outputs can reuse its cached null vector. It uses the existing spark.rapids.sql.projectAstEnabled configuration and does not change result semantics or add a user-facing configuration.

This is independent of the experimental AST JIT work in #15312. It refactors the existing legacy AST path and does not depend on AST JIT, LTO, or precompiled fragments.

This also changes physical-plan rendering: AST projections now appear under GpuProject as AST(...) AS x instead of using a dedicated GpuProjectAstExec node. This does not change result semantics, APIs, or configuration.

Implementation details:

• Represent each eligible output with a lightweight GpuProjectAstExpression.
• Keep top-level null literals on the regular projection path so multiple null outputs can reuse its cached null vector; non-null literals remain AST-eligible.
• Share one cuDF input Table across all AST outputs evaluated for the same tier and batch.
• Preserve TieredProject common-subexpression elimination for duplicate outputs and shared subtrees.
• Route fused higher-order-function projections through the same expression evaluator so AST outputs retain the shared input Table.
• Close compiled legacy AST expressions at task completion, including retry execution paths.

Performance testing:

• Benchmark script: project_ast_per_expression_perf.scala
• Environment: Spark 3.5.2 with local[4], plugin 26.08.0-SNAPSHOT, and 2 NVIDIA RTX 5880 Ada Generation GPUs.
• Workload: 20 million Parquet rows in 8 partitions, 84 projected outputs, one warmup, and five measured iterations. A global GPU aggregate materializes every projected output.
• LIBCUDF_JIT_ENABLED=0 isolates the legacy AST backend. AST off/on order and case order are reversed on alternating iterations.
• Speedup is AST off / AST on; values above 1 mean AST is faster. The Project metric is opTimeLegacy, accumulated across tasks, so it is not directly comparable to E2E wall-clock time.

Case	Expressions	Coverage / CSE shape	AST tiers (off/on)	AST outputs (off/on)	AST off/on E2E median (ms)	E2E speedup	AST off/on Project op median (ms)	Project op speedup
broad_unique_ast	84 unique AST-compatible	42 AST operator families	[0] / [84]	0 / 84	670.249 / 589.848	1.136x	684.796 / 356.251	1.922x
whole_output_duplicates	42 AST expressions, each projected twice	whole-output CSE	[0,0] / [42,0]	0 / 84	601.863 / 526.478	1.143x	355.906 / 155.668	2.286x
cheap_partial_cse	84 AST-compatible	one shared add	[0,0] / [1,84]	0 / 84	516.149 / 467.530	1.104x	214.336 / 189.046	1.134x
expensive_partial_cse	84 AST-compatible	one shared transcendental subtree	[0,0] / [1,84]	0 / 84	476.142 / 440.932	1.080x	239.745 / 175.069	1.369x
mixed_half	42 AST-compatible + 42 regular GPU	mixed execution	[0] / [42]	0 / 42	793.670 / 712.257	1.114x	1167.327 / 868.695	1.344x

No median regression was observed in the five main benchmark cases in either E2E or Project op time.

Top-level literal routing follow-up

A targeted follow-up compared top-level literals on the regular and AST paths. It used 20 million rows, 84 Project outputs, two warmups, and 15 alternating measured iterations. The benchmark directly executed and consumed the columnar GpuProject, avoiding Catalyst constant folding and unrelated aggregate, shuffle, or ColumnarToRow work. Speedup is regular / AST; values above 1 mean AST is faster.

Case	Expressions	Regular/AST E2E median (ms)	E2E speedup	Regular/AST Project op median (ms)	Project op speedup
literal_only	84 unique non-null long literals	61.418 / 60.952	1.008x	117.137 / 104.202	1.124x
mixed_ast_literals	42 AST-compatible expressions + 42 unique non-null long literals	231.422 / 230.816	1.003x	241.066 / 239.654	1.006x
null_literal_only	84 duplicate double null literals	33.257 / 137.528	0.242x	8.224 / 407.906	0.020x

Non-null literals were neutral in E2E time in both the literal-only and mixed workloads, so there is no evidence for excluding all top-level literals from AST. Null literals were different: the regular projection was 4.1x faster E2E and 49.6x faster in Project op time because it can reuse its cached null vector across outputs, whereas per-expression AST evaluates the null outputs independently. Based on these results, this PR keeps only top-level null literals on the regular projection path and leaves non-null literals AST-eligible.

TPC-H/TPC-DS coverage

We also performed physical-plan sweeps over stock TPC-H/NDS-H and TPC-DS workloads.

• TPC-H/NDS-H with the standard Decimal schema produced no legacy Project AST expressions across 24 physical plans.
• TPC-DS produced only three AST expressions across 102 physical plans. Two operate on a small filtered date dimension, and the remaining expression is in a post-aggregation Project.

These workloads therefore spend nearly all of their time in unrelated scans, joins, and aggregations and do not provide a meaningful performance signal for this change. The performance results above use purpose-built AST-only and mixed-expression Projects so that the affected execution path is exercised directly, while still materializing every projected expression.

Checklists

Documentation

• Updated for new or modified user-facing features or behaviors
• No user-facing change

Testing

• Added or modified tests to cover new code paths
• Covered by existing tests
  (Please provide the names of the existing tests in the PR description.)
• Not required

Performance

• Tests ran and results are added in the PR description
• Issue filed with a link in the PR description
• Not required

[thirtiseven]

@greptile full review

[greptile-apps]

Greptile Summary

This PR replaces the all-or-nothing GpuProjectAstExec node with per-expression AST eligibility inside the existing GpuProjectExec. Each top-level Project output is independently wrapped in the new GpuProjectAstExpression marker when eligible (fixed-width output, subtree AST-compatible, non-null literal). The shared input Table is created once per batch and passed to all AST outputs, preserving the shared-table optimisation. CSE (tiered project) is extended to propagate AST markers backwards through intermediate tiers.

• New expression wrapper: GpuProjectAstExpression is a lazy-compiled, AutoCloseable expression that registers a task-completion callback to release the CompiledExpression; compilation is guarded by synchronized and deferred to the first batch.
• Routing logic: GpuProjectExecMeta.convertToGpu() individually wraps each eligible output; null-literal outputs are kept on the regular GPU path to exploit its cached-null-vector optimisation. The GpuArrayHofFusion project path is updated to accept an evalColumn callback so HOF-fused and AST outputs share the same Table.
• Plan rendering change: AST outputs now appear as AST(…) AS x inside GpuProjectExec instead of the separate GpuProjectAstExec node; integration-test plan assertions are updated from checking for GpuProjectAstExec to GpuProjectAstExpression.

Confidence Score: 5/5

Safe to merge. Resource management (compiled expressions, intermediate Tables) is correct, OOM split-retry is preserved via GpuTieredProject, and per-expression routing logic is well-tested at both unit and integration level.

The core change — replacing the all-or-nothing GpuProjectAstExec with per-expression GpuProjectAstExpression markers — is architecturally clean and the implementation handles the key correctness concerns: lazy compilation with synchronized access, AutoCloseable + task-completion cleanup, shared Table creation guarded by withResource, and correct CSE backward-propagation through rewrapAstTiers. The two observations are minor explain verbosity and a test coverage gap in the Python integration tests; neither affects runtime correctness.

Files Needing Attention: No files require special attention. The integration test null-literal routing case could use a mixed-type companion test for stronger isolation, but this is not a blocker.

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala	New file: core marker expression that wraps a GPU expression for AST evaluation. Lazy compilation with synchronized, AutoCloseable + task-completion cleanup, and tableFromBatch helpers are correct. The buildExprTiers/rewrapAstTiers logic properly propagates AST markers backwards through CSE intermediate tiers while checking isFixedWidth.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala	Per-expression AST routing replaces the old all-or-nothing GpuProjectAstExec. The shared-table path in GpuProjectExec.project is guarded by withResource and correctly falls back to columnarEval for non-AST outputs. GpuProjectAstExec and its retry iterator are fully removed.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala	canThisBeAst refactored to canSelfBeAst && childExprs.forall(_.canThisBeAst) which is semantically equivalent. printAst now uses canSelfBeAst so only the specific failing node is printed rather than all ancestors, correctly matching per-expression eligibility semantics.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala	GpuArrayHofFusion.project and projectWithFusedGroups accept a new evalColumn callback so AST expressions outside fused HOF groups receive the shared input Table from the caller.
tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala	Retry test migrated from GpuProjectAstExec.buildRetryableAstIterator to GpuTieredProject.projectAndCloseStreamingWithSplitRetry; three new tests cover AST lifecycle (task-completion close), multi-level CSE tier preservation, and the HOF shared-table invariant.
integration_tests/src/main/python/ast_test.py	Plan assertions updated from GpuProjectAstExec to GpuProjectAstExpression; test_null_literal correctly switched to assert_gpu_project_without_ast; new test_project_ast_mixed_expressions covers the mixed AST + GPU case.

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into legacy-ast-per-..." | Re-trigger Greptile}

[nvauto]

NOTE: release/26.08 has been created from main. Please retarget your PR to release/26.08 if it should be included in the release.

[igorpeshansky]

LGTM modulo one remaining missing comment.

[thirtiseven]

build

[igorpeshansky]

build

[igorpeshansky]

build