Add experimental per-expression AST JIT for integral add and multiply

[thirtiseven]

Related to #15069.

Description

This draft PR proposes a minimal plugin-side integration framework for using cuDF AST JIT within GpuProjectExec. It is intended as an architectural proposal and an experimental foundation for future AST JIT work, rather than as a generally beneficial performance feature in its current form.

The implementation intentionally supports only non-ANSI IntegerType and LongType addition and multiplication. Other operators and data types are not enabled for AST JIT by this PR.

It adds the internal spark.rapids.sql.projectAstJitEnabled configuration, which is disabled by default. When enabled, the plugin recursively identifies maximal subtrees composed of the supported add and multiply expressions and wraps them in GpuAstJitExpression.

Expressions outside those supported subtrees continue to use the normal GPU Project evaluation path. This allows, for example, an add or multiply subtree to use AST JIT while an enclosing subtraction continues to use the existing GPU expression implementation. JIT and non-JIT expressions can also coexist as separate outputs of the same GpuProjectExec.

The main purpose of this PR is to establish and review the plugin/JNI integration boundary, per-expression selection and fallback behavior, retry handling, resource ownership, configuration, and plan visibility. It does not implement the following planned cuDF optimizations:

• Multi-output AST JIT and common subexpression elimination across Project outputs. Using those capabilities will require grouping eligible Spark expressions and submitting them through a future multi-output API.
• Linked or precompiled LTO fragments that reduce cold compilation cost. The current integration provides a place to consume that backend capability when it becomes available, but this PR does not reduce the initial compilation cost.

Performance

The primary benchmark reads 100 million rows from 16 Parquet partitions, evaluates eight distinct INT/BIGINT Project expressions with 16 alternating add and multiply operators each, and consumes every output with an aggregate. Steady-state modes use one discarded warmup and five measured iterations. Cold-state modes use three fresh Spark processes.

Environment: Spark 3.5.2 with local[8], NVIDIA RTX 5880 Ada Generation, CUDA/NVRTC 12.8, and LIBCUDF_JIT_ENABLED=0. AST JIT is enabled only through spark.rapids.sql.projectAstJitEnabled=true.

Benchmark scripts: Project AST JIT benchmark gist

Mode	Cache state	Runs	E2E avg (ms)	Speedup vs GPU Project	Speedup vs legacy AST
GPU_PROJECT	none	5	875.9	1.00x	0.49x
GPU_AST_LEGACY	none	5	425.5	2.06x	1.00x
GPU_AST_JIT_COLD	cold	3	6399.4	0.14x	0.07x
GPU_AST_JIT_DISK_WARM	disk warm	3	2263.5	0.39x	0.19x
GPU_AST_JIT_PCH_WARM_KERNEL_COLD	PCH warm, kernel cold	3	2874.4	0.30x	0.15x
GPU_AST_JIT_HOT	hot	5	430.8	2.03x	0.99x

All modes produced the same row count and aggregate checksum. Cold-state results intentionally include the time-to-first-result penalty. For this workload, hot AST JIT is effectively at parity with legacy AST, while non-warm JIT modes are substantially slower because of compilation cost.

An additional exploratory benchmark used one expression with 64 alternating add and multiply operators, five discarded warmups, and two independent applications with 20 measured iterations each:

Mode	Measured iterations	E2E avg (ms)	Speedup vs legacy AST
GPU_AST_LEGACY	40	213.9	1.00x
GPU_AST_JIT_HOT	40	191.4	1.12x

This indicates that expression-specific JIT code can help sufficiently deep expressions, but the current results do not demonstrate a broad performance benefit for typical expressions. Spark Project operation metrics are not included because they measure host-side operator time and do not reliably attribute asynchronous CUDA kernel execution; E2E time is used for the comparison.

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
• 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 introduces an experimental per-expression AST JIT path for GpuProjectExec, limited intentionally to non-ANSI IntegerType/LongType add and multiply. It is disabled by default and requires an internal config flag plus a cuDF runtime environment variable.

• New GpuAstJitExpression wraps maximal qualifying subtrees during planning, compiles them lazily via cuDF JIT, and integrates with the existing OOM retry (Retryable) and task-completion cleanup (onTaskCompletion) infrastructure.
• Planning hook in GpuProjectExecMeta.convertToGpu() runs the JIT wrapping pass before the existing AST and tiered-project paths, returning a standard GpuProjectExec so mixed JIT/non-JIT expressions coexist naturally.
• New traits/overrides on GpuExpression, GpuBoundReference, GpuLiteral, GpuAddBase, and GpuMultiply propagate JIT eligibility bottom-up through the expression tree.

Confidence Score: 4/5

Safe to merge with the feature disabled by default; the JIT path is only reachable with an explicit internal config, so production workloads are unaffected.

The core resource management is correctly implemented and consistent with the rest of the codebase. Two non-blocking issues exist: restore() leaves completionRegistered unreset (functionally correct but subtle), and there is no code-level guard against the config being enabled without LIBCUDF_JIT_ENABLED=1, which would produce a hard runtime failure rather than a graceful fallback. Mixed-expression integration tests also only cover IntegerType.

GpuAstJitExpression.scala — the restore()/completionRegistered interaction and the missing cuDF JIT availability guard are both worth addressing before the feature graduates from experimental.

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala	New expression wrapper that compiles a GPU subtree via cuDF JIT. OOM retry, task-completion cleanup, and zero-column dummy-table handling are present. Two subtle issues: restore() leaves completionRegistered unreset, and there is no guard against the config being enabled while LIBCUDF_JIT_ENABLED is absent.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala	JIT early-return injected in GpuProjectExecMeta.convertToGpu(). Logic is correct; JIT takes priority over the whole-project AST path when both configs are enabled.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala	Adds selfSupportsAstJit, selfIsAstJitOperator, supportsAstJit, and containsAstJitOperator to GpuExpression. Traversal correctly handles pre-binding AttributeReference children.
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala	Adds selfSupportsAstJit and selfIsAstJitOperator overrides to GpuAddBase and GpuMultiply, gated on !failOnError && (dataType == IntegerType
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala	New ENABLE_PROJECT_AST_JIT config key is correctly marked .internal(), defaults to false, and has a clear doc string.
integration_tests/src/main/python/ast_test.py	Four new JIT integration tests. Missing LongType parametrization for the two mixed-expression tests.
tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala	Unit tests for tree-wrapping logic covering default-disabled config, maximal-subtree wrapping, nested independent subtrees, and ANSI/float exclusion.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Meta as GpuProjectExecMeta
    participant Wrap as GpuAstJitExpression (object)
    participant JIT as GpuAstJitExpression (instance)
    participant Exec as GpuProjectExec
    participant cuDF as cuDF JIT API

    Meta->>Wrap: wrapProjectExpressions(gpuExprs)
    Wrap-->>Meta: jitProjectList
    Meta->>Exec: GpuProjectExec(jitProjectList)

    Note over Exec: Task starts
    Exec->>JIT: checkpoint()
    JIT->>cuDF: convertToAst().compile()
    cuDF-->>JIT: CompiledExpression
    JIT->>JIT: register onTaskCompletion(close)

    loop per batch
        Exec->>JIT: columnarEval(batch)
        JIT->>cuDF: computeColumnJit(table)
        cuDF-->>JIT: ColumnVector
        JIT-->>Exec: GpuColumnVector
    end

    alt GPU OOM
        JIT->>JIT: restore() close expression
        Exec->>JIT: retry getCompiledExpression recompile
    end

    Note over JIT: Task ends
    JIT->>JIT: onTaskCompletion close()

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Meta as GpuProjectExecMeta
    participant Wrap as GpuAstJitExpression (object)
    participant JIT as GpuAstJitExpression (instance)
    participant Exec as GpuProjectExec
    participant cuDF as cuDF JIT API

    Meta->>Wrap: wrapProjectExpressions(gpuExprs)
    Wrap-->>Meta: jitProjectList
    Meta->>Exec: GpuProjectExec(jitProjectList)

    Note over Exec: Task starts
    Exec->>JIT: checkpoint()
    JIT->>cuDF: convertToAst().compile()
    cuDF-->>JIT: CompiledExpression
    JIT->>JIT: register onTaskCompletion(close)

    loop per batch
        Exec->>JIT: columnarEval(batch)
        JIT->>cuDF: computeColumnJit(table)
        cuDF-->>JIT: ColumnVector
        JIT-->>Exec: GpuColumnVector
    end

    alt GPU OOM
        JIT->>JIT: restore() close expression
        Exec->>JIT: retry getCompiledExpression recompile
    end

    Note over JIT: Task ends
    JIT->>JIT: onTaskCompletion close()

Loading

_{Reviews (1): Last reviewed commit: "Update copyright years" | Re-trigger Greptile}

[thirtiseven]

Java binding is in 26.10, waiting for main branch to switch...

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