[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 plan node with per-expression AST evaluation inside GpuProjectExec. Each top-level Project output is independently wrapped in GpuProjectAstExpression if it is fixed-width, AST-compatible, and not a top-level null literal; sibling outputs that fail these checks fall back to the regular GPU expression path. A shared cuDF Table is built once per tier/batch and reused across all AST outputs in that tier, avoiding redundant host→device copies.

• GpuProjectAstExpression — a new AutoCloseable expression marker that lazily compiles and caches the native CompiledExpression, registers a task-completion callback for cleanup, and exposes computeColumn(table) for direct AST evaluation against a shared input table.
• buildExprTiers / rewrapAstTiers — CSE now unwraps AST markers before calling GpuEquivalentExpressions, then re-applies the markers to intermediate tier aliases reachable from AST outputs, preserving CSE across mixed AST/non-AST expressions.
• GpuArrayHofFusion.project — updated to accept an evalColumn function, routing HOF non-fused expressions through the same AST-aware dispatcher so the shared input table is respected inside fused HOF projections.

Confidence Score: 5/5

This PR is safe to merge. The refactor correctly maintains result semantics for all expression types and execution paths.

The change is a well-scoped refactor with thorough test coverage across retry, CSE multi-tier propagation, HOF fusion, and task-completion lifecycle. Resource management follows established withResource/closeOnExcept patterns, the OOM retry flow is unchanged (GpuTieredProject still drives it), and the canThisBeAst refactoring is a pure identity transform. No regressions were found in the logic for mixed AST/non-AST evaluation, null-literal routing, or explain output generation.

Files Needing Attention: No files require special attention.

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala	New file: defines GpuProjectAstExpression (lazy compile, task-completion close, synchronized access) and its companion object (wrap/unwrap/buildExprTiers/tableFromBatch). Resource management and CSE logic are correct.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala	GpuProjectAstExec removed; GpuProjectExecMeta.convertToGpu now wraps eligible per-expression AST; GpuProjectExec.project dispatches to a shared-table AST path when any expression is AST-wrapped. OOM retry flows through GpuTieredProject unchanged.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala	GpuArrayHofFusion.project and projectWithFusedGroups updated to accept an evalColumn callback, routing non-fused slots through the same AST-aware dispatcher. Logic is correct and validated by the new HOF fusion test.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala	canThisBeAst refactored as canSelfBeAst && childExprs.forall(canThisBeAst) (semantically equivalent). printAst now uses canSelfBeAst to print only node-local blockers. willWorkInAstInfo improved to report all three blocking conditions correctly.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala	Inline CSE replaced by delegation to GpuProjectAstExpression.buildExprTiers, which handles both AST-marker-aware tier building and CSE in one place.
integration_tests/src/main/python/ast_test.py	Test assertions updated from GpuProjectAstExec plan nodes to GpuProjectAstExpression. Null literal test now verifies the regular-path routing. New mixed-expression and folded-null-literal tests added. non_exist_classes='' is safe (guarded by Python truthiness check in asserts.py).
tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala	AST retry test updated to use GpuTieredProject directly. New tests for multi-level CSE preservation and task-completion lifecycle. Resource cleanup (withResource/finally) is correct.
tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala	New test verifies that fused HOF projection shares the same cuDF Table instance across sibling AST expressions. Uses Mockito ArgumentCaptor to capture and compare table instances.
tests/src/test/scala/com/nvidia/spark/rapids/AstUtilSuite.scala	Two new tests validate that explainAst prints only node-local AST blockers (canSelfBeAst) rather than whole-subtree failures, matching the changed printAst logic.
tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala	Test updated to verify YearMonthInterval unary-positive stays on GpuProjectExec with no GpuProjectAstExpression, since the expression is optimized to a pass-through.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[GpuProjectExecMeta.convertToGpu] --> B{conf.isProjectAstEnabled?}
    B -- No --> C[GpuProjectExec with all GPU exprs]
    B -- Yes --> D{Per-expression check:\nisFixedWidth AND canThisBeAst\nAND NOT topLevelNullLiteral?}
    D -- Yes --> E[wrap in GpuProjectAstExpression]
    D -- No --> F[keep as regular GPU expr]
    E --> G[GpuProjectExec with mixed expr list]
    F --> G

    G --> H[GpuTieredProject.buildExprTiers]
    H --> I[unwrap AST markers]
    I --> J[CSE: GpuEquivalentExpressions]
    J --> K[GpuEquivalentExpressions.getExprTiers]
    K --> L[rewrapAstTiers: trace AST outputs\nback through CSE tiers]
    L --> M[Bound + metric-injected tier expressions]

    M --> N[GpuProjectExec.project per tier/batch]
    N --> O{hasAstExpressions?}
    O -- Yes --> P[tableFromBatch: build shared cuDF Table]
    P --> Q[GpuArrayHofFusion.project with evalColumn lambda]
    Q --> R{Expression type?}
    R -- AST --> S[astExpression.computeColumn table]
    R -- Regular --> T[expression.columnarEval cb]
    O -- No --> U[GpuArrayHofFusion.project with columnarEval]

    S --> V[GpuColumnVector output]
    T --> V
    U --> V
    V --> W[ColumnarBatch result]

Loading

_{Reviews (11): 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

[igorpeshansky]

build

[thirtiseven]

build

[thirtiseven]

build

[thirtiseven]

build

[thirtiseven]

build

[igorpeshansky]

build