Add experimental Project AST JIT for integral add and multiply [databricks]

docs/dev/nvtx_profiling.md

@@ -69,3 +69,87 @@ sudo sh -c 'echo [level] >/proc/sys/kernel/perf_event_paranoid'
 where valid values are { 1, 2 }. Refer to
 [NVIDIA Nsight Systems documentation](https://docs.nvidia.com/nsight-systems/)
 for further details.
+
+## Project AST backend selection
+
+`spark.rapids.sql.projectAstEnabled` and `spark.rapids.sql.projectAstJitEnabled`
+independently enable interpreted AST and AST JIT. Both remain disabled by default.
+Project and extracted broadcast build-side projections use a shared tier planner;
+generic tiered binders continue to require an explicit backend opt-in and drop
+existing AST wrappers otherwise.
+
+The planner first materializes deterministic AST subexpressions shared across
+backends, so their consumers reuse one computed column. It then preserves a complete
+expression when an enabled AST backend supports it. JIT wins when both support the
+complete expression; a complete interpreted AST
+is preferred over splitting an expression only to use JIT for part of it. Otherwise,
+eligible subexpressions can produce columns for a later backend tier, including
+GPU inputs to CPU bridges. Only deterministic values without side effects move to an
+earlier tier, where they are computed for every row, including rows that a conditional
+such as `IF` or `AND` would skip. Lambda bodies are not split. With tiered projection disabled, only complete output expressions are wrapped.
+Legacy eligibility comes from the existing expression metadata and is carried by
+driver-side tree tags. Tiering replaces those decisions with executable AST wrappers
+before the bound project is serialized to executors.
+
+Interpreted AST uses plugin-side CSE and can execute several expressions separately
+in one tier. JIT roots in the same multi-output group keep shared subtrees for cuDF
+CSE. Cross-tier dependencies are materialized and managed by the existing tiered
+projection lifecycle. Cross-backend sharing takes precedence over complete expression
+matching; shared descendants are not separately materialized when sharing their parent
+already eliminates the repeated work. A value that would otherwise stay inline is not
+materialized only for sharing when that would leave an interpreted consumer with a
+single simple arithmetic operation. This is a fixed policy, not a cost model:
+materializing a cheap operation can cost more than recomputing it, and enabling JIT
+does not imply that cold compilation will be amortized.
+
+Standalone add, subtract, and multiply expressions with only column or literal
+inputs stay on regular GPU execution instead of interpreted AST. The planner
+checks this before boundary selection, treating children read from earlier tiers as
+columns, and again after CSE, when a fused expression may have become a single
+operation. These operations remain eligible inside compound AST expressions, and
+this rule does not change JIT selection.
+
+## Project AST JIT diagnostics
+
+With either Project AST backend enabled, `spark.rapids.sql.explain=ALL` logs the
+final bound tiers under `FINAL PROJECT AST SELECTION`, including any JIT execution
+groups. Input ordinals refer to the current tier's
+input batch; output positions refer to its output batch. The log identifies
+forwarded outputs and computed outputs materialized for the next tier.
+`plugin_unique_ops` and `plugin_shared_ops` count operations using plugin expression
+equivalence, not the final cuDF Row IR or GPU instructions. Grouping honors the
+multi-output execution setting. These are final physical tiers, not the original
+planner's dependency-wave identifiers.
+
+The final backend labels distinguish the plugin's explicit `AST JIT` path,
+its legacy `AST Interpreted` path, and the regular GPU projection. These labels
+describe wrapper selection, not a trace of native execution. The legacy wrapper
+uses `compile()` with cuDF's default backend selection; `LIBCUDF_JIT_ENABLED=1`
+can route that path to native JIT even when the plugin's Project JIT switch is off.
+Set `LIBCUDF_JIT_ENABLED=0` when comparing interpreted AST with explicit AST JIT,
+and use native tracing to confirm the executed backend. The Spark plan string
+shows a backend for an output only when the output's final root runs on it; it does
+not show AST subexpressions extracted to earlier tiers.
+
+With `spark.rapids.sql.metrics.level=DEBUG`, `GpuProjectExec` exposes:
+
+| Metric key | Meaning |
+|---|---|
+| `astJitProgramBuildAttempts` | Calls to construct a schema-bound native program, including failed attempts. |
+| `astJitProgramCacheHits` | Reuse of the owning expression's existing schema-compatible program. |
+| `astJitProgramBuildTime` | Host elapsed time constructing programs, including failures, native cache lookup, and any compilation/waiting inside that call. |
+| `astJitEvalAttempts` | Group evaluation calls, including failed attempts. A multi-output group counts once. |
+| `astJitEvalRows` | Input rows summed over group evaluation attempts, including retries and repeated processing by different groups. |
+| `astJitEvalTime` | Host elapsed time of group evaluation, including failed attempts and output wrapping; excludes explicit program construction. |
+
+These counters measure attempted work, not committed output. In particular, a
+batch processed by two groups contributes twice its row count. Program reuse is
+local to the bound expression; a program build can still hit the native kernel
+cache and perform no NVRTC compilation. Neither counter measures native cache
+hits or misses. Projections evaluated within joins only receive these counters
+if their owning operator supplies them.
+
+The timing metrics are task-accumulated host times, not query wall time or isolated
+GPU kernel time. Their work also contributes to the existing operator timing;
+do not add them to operator time. Use native NVTX/compiler tracing to distinguish
+actual NVRTC compilation from program construction and waiting.

docs/dev/nvtx_ranges.md

@@ -115,6 +115,7 @@ ProjectExec|Executing projection operation on columnar batch
 Columnar batch serialize row only|Serializing row-only batch (no GPU data)
 disk spill|Spilling data from host memory to disk
 Async Shuffle Buffer|Asynchronous shuffle buffering operation
+Compile AST JIT|Compiling an AST expression for JIT evaluation
 CSV decode|Decoding CSV data
 get batch|Getting join batch
 Spark Task|Spark task execution range for stage and task tracking
@@ -187,6 +188,7 @@ GpuGenerateExec|Executing generate operation on GPU
 Join gather|Gathering join results
 parquet get blocks with filter|Retrieving Parquet blocks after applying filters
 dynamic sort heuristic|Applying dynamic sort heuristic for aggregation
+Project AST JIT|Applying JIT-compiled AST projection to batch
 shuffle concat load batch|Concatenating and loading batch in shuffle operation
 parquet read footer bytes|Reading raw footer bytes from Parquet file
 spill batch|Spilling join batch

integration_tests/src/main/python/ast_test.py

@@ -69,6 +69,12 @@
          is_spark_412_or_later()))]
 
 _project_ast_enabled_conf = {"spark.rapids.sql.projectAstEnabled": "true"}
+_project_ast_jit_enabled_conf = {"spark.rapids.sql.projectAstJitEnabled": "true"}
+_project_ast_jit_and_legacy_enabled_conf = {
+    "spark.rapids.sql.projectAstEnabled": "true",
+    "spark.rapids.sql.projectAstJitEnabled": "true"
+}
+
 
 def assert_gpu_ast(is_supported, func, conf={}):
     ast_expression = "GpuProjectAstExpression"
@@ -364,7 +370,7 @@ def test_bitwise_xor(data_descr):
 @disable_ansi_mode
 def test_addition(data_descr):
     data_type = data_descr[0].data_type
-    assert_binary_ast(data_descr,
+    assert_binary_ast((data_descr[0], False),
         lambda df: df.select(
             f.col('a') + f.lit(100).cast(data_type),
             f.lit(-12).cast(data_type) + f.col('b'),
@@ -374,7 +380,7 @@ def test_addition(data_descr):
 @disable_ansi_mode
 def test_subtraction(data_descr):
     data_type = data_descr[0].data_type
-    assert_binary_ast(data_descr,
+    assert_binary_ast((data_descr[0], False),
         lambda df: df.select(
             f.col('a') - f.lit(100).cast(data_type),
             f.lit(-12).cast(data_type) - f.col('b'),
@@ -384,12 +390,160 @@ def test_subtraction(data_descr):
 @disable_ansi_mode
 def test_multiplication(data_descr):
     data_type = data_descr[0].data_type
-    assert_binary_ast(data_descr,
+    assert_binary_ast((data_descr[0], False),
         lambda df: df.select(
             f.col('a') * f.lit(100).cast(data_type),
             f.lit(-12).cast(data_type) * f.col('b'),
             f.col('a') * f.col('b')))
 
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_add_multiply(data_gen):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            f.col('a') + f.col('b'),
+            f.col('a') * f.col('b')),
+        exist_classes=r"GpuProject.*AST_JIT",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_partial_project_with_unique_unsupported_expression(data_gen):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            (f.col('a') + f.col('b')) - (f.col('a') * f.col('b'))),
+        exist_classes="GpuProject",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@pytest.mark.parametrize('backend', ['legacy', 'jit', 'both'])
+@disable_ansi_mode
+def test_ast_shared_subtree(data_gen, backend):
+    def project_shared_expression(spark):
+        df = binary_op_df(spark, data_gen)
+        shared = f.col('a') + f.col('b')
+        return df.select(
+            (shared * f.col('a')).alias('left'),
+            (shared * f.col('b')).alias('right'))
+
+    conf = {
+        'spark.rapids.sql.projectAstEnabled': str(backend != 'jit').lower(),
+        'spark.rapids.sql.projectAstJitEnabled': str(backend != 'legacy').lower(),
+    }
+    if backend == 'legacy':
+        # Plugin CSE extracts a + b and leaves simple arithmetic on the regular GPU.
+        exist_classes = 'GpuProject'
+        non_exist_classes = 'GpuProjectAst,AST_JIT'
+    else:
+        exist_classes = r"GpuProject.*AST_JIT.*AS left.*AST_JIT.*AS right"
+        non_exist_classes = 'GpuProjectAst'
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        project_shared_expression,
+        exist_classes=exist_classes,
+        non_exist_classes=non_exist_classes,
+        conf=conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_regular_jit_waves(data_gen):
+    def project_waves(spark):
+        df = binary_op_df(spark, data_gen)
+        shared = f.col('a') + f.col('b')
+        regular = f.greatest(shared, f.col('a'))
+        return df.select(
+            (shared * f.col('a')).alias('early'),
+            regular.alias('regular'),
+            (regular * f.col('b')).alias('late'))
+
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        project_waves,
+        exist_classes=r"GpuProject.*AST_JIT.*AS early.*AS regular.*AS late",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_cse_shared_subexpression(data_gen):
+    def project_shared_expression(spark):
+        df = binary_op_df(spark, data_gen)
+        shared = f.col('a') + f.col('b')
+        return df.select(
+            shared.alias('shared'),
+            (shared - f.col('a')).alias('left'),
+            (shared - f.col('b')).alias('right'))
+
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        project_shared_expression,
+        exist_classes=r"GpuProject.*AST_JIT.*AS shared.*AS left.*AS right",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@pytest.mark.parametrize(
+    'tiered_project_enabled', ['true', 'false'], ids=['tiered', 'single_tier'])
+@disable_ansi_mode
+def test_jit_mixed_project_expressions(data_gen, tiered_project_enabled):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            (f.col('a') + f.col('b')).alias('jit'),
+            (f.col('a') - f.col('b')).alias('gpu'),
+            ((f.col('a') * f.col('b')) - f.col('a')).alias('mixed')),
+        exist_classes=r"GpuProject.*AST_JIT.*AS jit.*AS gpu.*AS mixed",
+        non_exist_classes=r"GpuProjectAst,AS gpu.*AST_JIT",
+        conf=copy_and_update(_project_ast_jit_enabled_conf, {
+            'spark.rapids.sql.tiered.project.enabled': tiered_project_enabled
+        }))
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_and_legacy_ast_mixed_project_expressions(data_gen):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            (f.col('a') + f.col('b')).alias('jit'),
+            ((f.col('a') * f.col('a')) - f.col('b')).alias('legacy'),
+            ((f.col('a') + f.col('b')) -
+                (f.col('a') * f.col('b'))).alias('mixed')),
+        exist_classes=(
+            r"GpuProject.*AST_JIT.*AS jit.*AST\(.*AS legacy.*"
+            r"AST\(.*AS mixed,GpuProjectAstExpression"),
+        non_exist_classes=r"AS legacy.*AST_JIT",
+        conf=_project_ast_jit_and_legacy_enabled_conf)
+
+
+@pytest.mark.parametrize('backend', ['legacy', 'jit', 'both'])
+@pytest.mark.parametrize('expr', [
+    '(a < 0) AND ((b + 1L) > 0)',
+    '(a >= 0) OR ((b + 1L) > 0)',
+    'IF(a < 0, b + 1L, 0L)',
+    'CASE WHEN a < 0 THEN b + 1L ELSE 0L END'], ids=idfn)
+def test_project_ast_keeps_skipped_side_effects(expr, backend):
+    # The ANSI addition overflows only on the first row, where the conditional skips it.
+    data = [(9, LONG_MAX), (-1, 1), (None, None)]
+    def project(spark):
+        df = spark.createDataFrame(spark.sparkContext.parallelize(data), 'a long, b long')
+        return df.selectExpr(f'{expr} AS result')
+
+    conf = copy_and_update(ansi_enabled_conf, {
+        'spark.rapids.sql.projectAstEnabled': str(backend != 'jit').lower(),
+        'spark.rapids.sql.projectAstJitEnabled': str(backend != 'legacy').lower(),
+    })
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        project, exist_classes='GpuProjectExec', conf=conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+# The first remainder keeps a + b inline in both outputs; the second materializes it for both.
+@pytest.mark.parametrize('legacy_expr', ['(a + b) - a', '(a + b) - (a * b)'], ids=idfn)
+@disable_ansi_mode
+def test_cross_backend_shared_subtree(data_gen, legacy_expr):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).selectExpr(
+            f'{legacy_expr} as legacy', '(a + b) * b as jit'),
+        exist_classes=r'GpuProject.*AST\(.*AS legacy.*AST_JIT.*AS jit',
+        conf=_project_ast_jit_and_legacy_enabled_conf)
+
+
 # Each descriptor contains a list of data generators and a corresponding boolean
 # indicating whether that data type is supported by the AST
 # all the below desc are not supported by the AST because ANSI mode is on
@@ -459,9 +613,10 @@ def project(spark):
             ('col_string', string_gen)
         ])
         shared = f.col('a') + f.col('b')
+        # Compound outputs keep AST after tiered CSE extracts the shared addition.
         return df.select(
-            (shared * f.col('c')).alias('ast_first'),
-            (shared * f.col('d')).alias('ast_second'),
+            ((shared * f.col('c')) - f.col('d')).alias('ast_first'),
+            ((shared * f.col('d')) - f.col('c')).alias('ast_second'),
             f.greatest(shared, f.col('c')).alias('gpu_shared'),
             f.length(f.col('col_string')).alias('gpu_string'),
             f.col('col_string').alias('raw_string'))

integration_tests/src/main/python/cpu_bridge_test.py

@@ -60,6 +60,47 @@ def test_func(spark):
     assert_gpu_and_cpu_are_equal_collect(test_func, conf=conf)
 
 
+@allow_non_gpu('Subtract')
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@pytest.mark.parametrize('tiered', ['true', 'false'])
+@pytest.mark.parametrize('multi_output', ['true', 'false'])
+@pytest.mark.parametrize('backend', ['legacy', 'jit', 'both'])
+def test_cpu_bridge_with_ast_backends(data_gen, tiered, multi_output, backend):
+    def test_func(spark):
+        return binary_op_df(spark, data_gen).selectExpr(
+            "(a + b) * b as jit", "((a + b) * b) - b as bridged",
+            "(((a + b) * b) - b) * a + b as result")
+
+    conf = copy_and_update(create_cpu_bridge_fallback_conf(
+        ['Subtract'], ['org.apache.spark.sql.catalyst.expressions.Add']), {
+        'spark.rapids.sql.projectAstJitEnabled': str(backend != 'legacy').lower(),
+        'spark.rapids.sql.projectAstEnabled': str(backend != 'jit').lower(),
+        'spark.rapids.sql.tiered.project.enabled': tiered,
+        'spark.rapids.sql.projectAstJitMultiOutputEnabled': multi_output,
+    })
+    ast_pattern = r'AST\(' if backend == 'legacy' else 'AST_JIT'
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        test_func, exist_classes=rf'GpuProject.*{ast_pattern},GpuCpuBridgeExpression',
+        non_exist_classes='AST_JIT' if backend == 'legacy' else '', conf=conf)
+
+
+@allow_non_gpu('Cast')
+@pytest.mark.parametrize('backend', ['legacy', 'jit', 'both'])
+def test_cpu_bridge_short_circuit_with_ast_backends(backend):
+    # The bridged ANSI cast fails only on the first row, where AND skips it.
+    data = [(9, 'abc'), (-1, '5')]
+    def test_func(spark):
+        df = spark.createDataFrame(spark.sparkContext.parallelize(data), 'a long, s string')
+        return df.selectExpr('(a < 0) AND (CAST(s AS INT) > 0) AS result')
+
+    conf = copy_and_update(create_cpu_bridge_fallback_conf(['Cast']), ansi_enabled_conf, {
+        'spark.rapids.sql.projectAstJitEnabled': str(backend != 'legacy').lower(),
+        'spark.rapids.sql.projectAstEnabled': str(backend != 'jit').lower(),
+    })
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        test_func, exist_classes='GpuCpuBridgeExpression', conf=conf)
+
+
 # MIN_ROWS_PER_SUBBATCH is 500,000: these values exercise one and two sub-batch results.
 @allow_non_gpu('Add')
 @pytest.mark.parametrize('row_count', [500000, 500001], ids=idfn)