[FEA] Add Java bindings for AST JIT operations by thirtiseven · Pull Request #23117 · NVIDIA/cudf

thirtiseven · 2026-07-06T10:25:16Z

Description

Adds cuDF Java/JNI bindings for the libcudf AST JIT operations introduced by #22602.

Why

cudf-spark needs to lower fused ANSI and TRY expressions to libcudf row IR while preserving row-level overflow and error semantics. JIT-only AST operations also need an explicit Java execution path: routing them through the existing CompiledExpression.computeColumn makes execution depend on the process-level LIBCUDF_JIT_ENABLED backend selection.

What changed

Add JitOperation, JitOperator, and JitErrorPolicy, with Java and JNI validation for operator arity, error-policy use, and decimal target scales.
Add CompiledExpression.computeColumnJit(Table) to call cudf::compute_column_jit directly.
Add Cudf.initializeJitRuntime() as an idempotent JIT runtime and cache preflight. It does not change the backend selected by computeColumn.
Add decimal AST literals for DECIMAL32, DECIMAL64, and DECIMAL128, including null and range validation.
Cover regular and overflow-checking arithmetic, PROPAGATE/NULLIFY, shifts, casts, decimal precision and rescaling, conditionals, and nested legacy/JIT expression trees.

API impact

This adds the public JitOperation, JitOperator, and JitErrorPolicy types, CompiledExpression.computeColumnJit(Table), Cudf.initializeJitRuntime(), and Literal.ofDecimal(DType, BigInteger). Existing APIs remain source-compatible, and CompiledExpression.computeColumn retains its current backend-selection behavior. Expressions containing JitOperation must be evaluated with computeColumnJit unless the process-level JIT backend is enabled.

Validation

All 134 CompiledExpressionTest cases pass on a local GPU against matching locally built shared libcudf and JNI libraries with LIBCUDF_JIT_ENABLED=0.
A clean cudf-spark Spark 3.5.2 sql-plugin compile passes against the resulting JNI jar.
Relevant clang-format, codespell, whitespace, and EOF pre-commit checks pass.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

copy-pr-bot · 2026-07-06T10:25:19Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

coderabbitai · 2026-07-07T09:21:50Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR adds Java and JNI support for libcudf JIT AST expressions, decimal literal serialization, native AST construction, separate default and JIT execution paths, and validation and execution tests.

Changes

JIT AST Execution

Layer / File(s)	Summary
JIT AST and serialization contracts `java/src/main/java/ai/rapids/cudf/ast/AstExpression.java`, `java/src/main/java/ai/rapids/cudf/ast/AstUtils.java`, `java/src/main/java/ai/rapids/cudf/ast/Jit.java`, `java/src/main/java/ai/rapids/cudf/ast/Operator.java`, `java/src/main/java/ai/rapids/cudf/ast/TableReference.java`	Adds JIT expression types, operators, error policies, checked native-id conversion, and serialized JIT operation support.
Decimal literal representation `java/src/main/java/ai/rapids/cudf/ast/Literal.java`	Adds BigInteger decimal literal construction, four-byte decimal scale serialization, and decimal128 byte-order conversion.
Native AST tree construction `java/src/main/native/src/jni_compiled_expr.hpp`	Refactors native expression storage to use an AST tree and adds reference-based construction APIs.
Native JIT compilation and execution `java/src/main/native/src/CompiledExpression.cpp`, `java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java`	Decodes and validates JIT expressions, preserves decimal scale, and exposes separate default and JIT execution paths.
JIT and decimal validation `java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java`	Adds tests for decimal literals, JIT operators, overflow policies, conditionals, casts, rescaling, precision handling, and execution errors.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related issues

rapidsai/cudf#18023: Covers expanding Java bindings with libcudf JIT expression support.

Suggested labels: libcudf, improvement

Suggested reviewers: revans2, ttnghia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 8.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding Java bindings for AST JIT operations.
Description check	✅ Passed	The description is directly related to the changeset and describes the new JIT bindings, decimal literals, and tests.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

thirtiseven · 2026-07-08T03:00:12Z

@coderabbitai full review

coderabbitai · 2026-07-08T03:00:17Z

✅ Action performed

Full review finished.

coderabbitai

🧹 Nitpick comments (1)

java/src/main/java/ai/rapids/cudf/ast/Literal.java (1)
130-132: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Unverified doc claim: legacy computeColumn support for DECIMAL32/64 literals.

The javadoc states only DECIMAL128 literals require computeColumnJit, implying DECIMAL32/DECIMAL64 decimal literals work with legacy computeColumn. No test in CompiledExpressionTest exercises a decimal literal via computeColumn (all testJitDecimalLiteralTransform cases use computeColumnJit), so this claim is currently untested here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@java/src/main/java/ai/rapids/cudf/ast/Literal.java` around lines 130 - 132,
The javadoc in Literal is making an unverified claim about decimal literal
support with computeColumn versus computeColumnJit. Update the documentation to
match the actual behavior for DECIMAL32/DECIMAL64/DECIMAL128 literals, and add
or adjust coverage in CompiledExpressionTest (for example around
testJitDecimalLiteralTransform and computeColumn) so the supported execution
path is explicitly verified.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@java/src/main/java/ai/rapids/cudf/ast/Literal.java`:
- Around line 130-132: The javadoc in Literal is making an unverified claim
about decimal literal support with computeColumn versus computeColumnJit. Update
the documentation to match the actual behavior for
DECIMAL32/DECIMAL64/DECIMAL128 literals, and add or adjust coverage in
CompiledExpressionTest (for example around testJitDecimalLiteralTransform and
computeColumn) so the supported execution path is explicitly verified.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 61abb19d-7bb0-44a3-991b-d21254a99a70

📥 Commits

Reviewing files that changed from the base of the PR and between ba255bb and 83f500a.

📒 Files selected for processing (11)

java/src/main/java/ai/rapids/cudf/Cudf.java
java/src/main/java/ai/rapids/cudf/ast/AstExpression.java
java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java
java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java
java/src/main/java/ai/rapids/cudf/ast/JitOperation.java
java/src/main/java/ai/rapids/cudf/ast/JitOperator.java
java/src/main/java/ai/rapids/cudf/ast/Literal.java
java/src/main/native/src/CompiledExpression.cpp
java/src/main/native/src/CudfJni.cpp
java/src/main/native/src/jni_compiled_expr.hpp
java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java

…indings-main-upmerge

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

lamarrr

The C++ changes look good to me.
I took a look at the Java changes too, but I believe we'll need more eyes on it.

lamarrr · 2026-07-14T21:08:02Z

+  auto const op_info      = jni_to_jit_operator(jni_ast.read_byte());
+  auto const error_policy = jni_to_jit_error_policy(jni_ast.read_byte());
+  auto const arity        = static_cast<int32_t>(jni_ast.read_byte());
+  if (arity < 0) { throw std::invalid_argument("unexpected JNI AST JIT operator arity"); }


i think checking the arity comparison below is sufficient

Agreed. I removed the separate negative-arity check. The single expected-arity comparison rejects both negative and otherwise mismatched values.

lamarrr · 2026-07-14T21:12:54Z

/ok to test 30cceaf

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@java/src/main/native/src/CompiledExpression.cpp`:
- Around line 560-562: Add benchmark coverage for the JIT branch in the result
computation around compute_column_jit, using representative GPU expressions that
exercise standard AST execution, decimal handling, and overflow-policy behavior.
Follow existing benchmark conventions and retain equivalent coverage for the
non-JIT path where applicable.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: fc60ee7f-d81f-4b70-a8af-4d97ba53b4cc

📥 Commits

Reviewing files that changed from the base of the PR and between 30cceaf and 4f2dd1b.

📒 Files selected for processing (13)

java/src/main/java/ai/rapids/cudf/Cudf.java
java/src/main/java/ai/rapids/cudf/ast/AstExpression.java
java/src/main/java/ai/rapids/cudf/ast/AstUtils.java
java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java
java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java
java/src/main/java/ai/rapids/cudf/ast/JitOperation.java
java/src/main/java/ai/rapids/cudf/ast/JitOperator.java
java/src/main/java/ai/rapids/cudf/ast/Literal.java
java/src/main/java/ai/rapids/cudf/ast/TableReference.java
java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java
java/src/main/native/src/CompiledExpression.cpp
java/src/main/native/src/CudfJni.cpp
java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java

💤 Files with no reviewable changes (2)

java/src/main/native/src/CudfJni.cpp
java/src/main/java/ai/rapids/cudf/Cudf.java

🚧 Files skipped from review as they are similar to previous changes (5)

java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java
java/src/main/java/ai/rapids/cudf/ast/AstExpression.java
java/src/main/java/ai/rapids/cudf/ast/Literal.java
java/src/main/java/ai/rapids/cudf/ast/JitOperator.java
java/src/main/java/ai/rapids/cudf/ast/JitOperation.java

igorpeshansky · 2026-07-15T18:15:12Z

+  auto const& expression = compiled_expr_ptr->get_top_expression();
+  if (backend == execution_backend::DEFAULT) {
+    auto const* literal = dynamic_cast<cudf::ast::literal const*>(&expression);
+    // The legacy evaluator silently corrupts decimal128 literal outputs in release builds.


Would a DECIMAL128 literal nested within an expression also be corrupted, or does that only apply to root literals?

Good catch. The limitation is not specific to a bare root literal: the legacy evaluator cannot materialize any DECIMAL128 root result. I moved the validation into expression_parser so it uses the inferred root output type, while nested comparisons producing BOOL8 remain supported. This follows the fail-fast behavior introduced for unsupported DECIMAL128 AST operations in #22982.

@lamarrr @davidwendt, could you sanity-check whether the root-output routing is the right place for this validation? The alternative would be to keep this as a separate libcudf follow-up, but the JNI-only literal check was too narrow and allowed explicit IDENTITY expressions to reach the same silent-corruption path.

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

igorpeshansky

LGTM modulo a tiny consistency nit.

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

thirtiseven · 2026-07-20T02:11:40Z

/ok to test f869b70

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

thirtiseven · 2026-07-20T02:30:20Z

/ok to test 3aeba52

thirtiseven · 2026-07-20T05:59:42Z

/merge

Adds cuDF Java/JNI bindings for the libcudf AST JIT operations introduced by NVIDIA#22602. ### Why cudf-spark needs to lower fused ANSI and TRY expressions to libcudf row IR while preserving row-level overflow and error semantics. JIT-only AST operations also need an explicit Java execution path: routing them through the existing `CompiledExpression.computeColumn` makes execution depend on the process-level `LIBCUDF_JIT_ENABLED` backend selection. ### What changed - Add `JitOperation`, `JitOperator`, and `JitErrorPolicy`, with Java and JNI validation for operator arity, error-policy use, and decimal target scales. - Add `CompiledExpression.computeColumnJit(Table)` to call `cudf::compute_column_jit` directly. - Add `Cudf.initializeJitRuntime()` as an idempotent JIT runtime and cache preflight. It does not change the backend selected by `computeColumn`. - Add decimal AST literals for `DECIMAL32`, `DECIMAL64`, and `DECIMAL128`, including null and range validation. - Cover regular and overflow-checking arithmetic, `PROPAGATE`/`NULLIFY`, shifts, casts, decimal precision and rescaling, conditionals, and nested legacy/JIT expression trees. ### API impact This adds the public `JitOperation`, `JitOperator`, and `JitErrorPolicy` types, `CompiledExpression.computeColumnJit(Table)`, `Cudf.initializeJitRuntime()`, and `Literal.ofDecimal(DType, BigInteger)`. Existing APIs remain source-compatible, and `CompiledExpression.computeColumn` retains its current backend-selection behavior. Expressions containing `JitOperation` must be evaluated with `computeColumnJit` unless the process-level JIT backend is enabled. ### Validation - All 134 `CompiledExpressionTest` cases pass on a local GPU against matching locally built shared libcudf and JNI libraries with `LIBCUDF_JIT_ENABLED=0`. - A clean cudf-spark Spark 3.5.2 `sql-plugin` compile passes against the resulting JNI jar. - Relevant clang-format, codespell, whitespace, and EOF pre-commit checks pass. Authors: - Haoyang Li (https://github.com/thirtiseven) Approvers: - Basit Ayantunde (https://github.com/lamarrr) - Igor Peshansky (https://github.com/igorpeshansky) URL: NVIDIA#23117

thirtiseven added 2 commits July 4, 2026 09:04

Expose JIT AST operations in Java

fdde92b

Align Java AST JIT bindings with libcudf APIs

db1e27a

github-actions Bot assigned thirtiseven Jul 6, 2026

github-actions Bot added the Java Affects Java cuDF API. label Jul 6, 2026

thirtiseven added feature request New feature or request non-breaking Non-breaking change labels Jul 6, 2026

thirtiseven added 2 commits July 7, 2026 15:37

local review

4680b9a

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

local review comments address

6f79a23

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

thirtiseven marked this pull request as ready for review July 7, 2026 09:18

thirtiseven requested a review from a team as a code owner July 7, 2026 09:18

Merge branch 'main' into codex-ansi-jit-java-bindings-main-upmerge

83f500a

coderabbitai Bot reviewed Jul 8, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into codex-ansi-jit-java-b…

500eae5

…indings-main-upmerge

thirtiseven requested review from igorpeshansky and ttnghia July 8, 2026 06:36

address code rabbit comments

19f974b

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

thirtiseven requested review from lamarrr and revans2 July 9, 2026 08:25

lamarrr reviewed Jul 9, 2026

View reviewed changes

Comment thread java/src/main/java/ai/rapids/cudf/ast/JitOperation.java Outdated

Comment thread java/src/main/java/ai/rapids/cudf/ast/JitOperator.java

address comments

30cceaf

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

GregoryKimball added this to libcudf Jul 13, 2026

GregoryKimball moved this to Burndown in libcudf Jul 13, 2026

lamarrr approved these changes Jul 14, 2026

View reviewed changes

igorpeshansky reviewed Jul 15, 2026

View reviewed changes

address comments

4f2dd1b

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

coderabbitai Bot reviewed Jul 15, 2026

View reviewed changes

Comment thread java/src/main/native/src/CompiledExpression.cpp

igorpeshansky reviewed Jul 15, 2026

View reviewed changes

address comments

122192f

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

thirtiseven requested a review from a team as a code owner July 16, 2026 04:01

thirtiseven requested a review from vuule July 16, 2026 04:01

github-actions Bot added the libcudf Affects libcudf (C++/CUDA) code. label Jul 16, 2026

igorpeshansky reviewed Jul 16, 2026

View reviewed changes

Comment thread cpp/tests/ast/transform_tests.cpp Outdated

Comment thread cpp/tests/ast/transform_tests.cpp Outdated

address comments

27d63d8

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

igorpeshansky approved these changes Jul 17, 2026

View reviewed changes

Comment thread cpp/tests/ast/transform_tests.cpp Outdated

Comment thread cpp/tests/ast/transform_tests.cpp

thirtiseven added 2 commits July 20, 2026 10:10

address comments

842c4a0

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

Merge branch 'main' into codex-ansi-jit-java-bindings-main-upmerge

f869b70

style

3aeba52

Signed-off-by: Haoyang Li <haoyangl@nvidia.com>

rapids-bot Bot merged commit 4206875 into NVIDIA:main Jul 20, 2026
255 of 261 checks passed

GregoryKimball moved this from Burndown to Landed in libcudf Jul 20, 2026

This was referenced Aug 10, 2026

[FEA] Add support for scalar column views in AST expressions and JIT execution #23615

Merged

[FEA] Implement Multi-output AST JIT & IR CSE #23621

Merged

josephine-wolf-oberholtzer added this to libcudf Aug 12, 2026

josephine-wolf-oberholtzer moved this to Landed in libcudf Aug 12, 2026

GregoryKimball removed this from libcudf Aug 21, 2026

Uh oh!

Conversation

thirtiseven commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Why

What changed

API impact

Validation

Checklist

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

coderabbitai Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

thirtiseven commented Jul 8, 2026

Uh oh!

coderabbitai Bot commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

lamarrr left a comment

Choose a reason for hiding this comment

Uh oh!

lamarrr Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

thirtiseven Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

lamarrr commented Jul 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

igorpeshansky Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

thirtiseven Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

igorpeshansky left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

thirtiseven commented Jul 20, 2026

Uh oh!

thirtiseven commented Jul 20, 2026

Uh oh!

thirtiseven commented Jul 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

thirtiseven commented Jul 6, 2026 •

edited

Loading

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading

coderabbitai Bot commented Jul 8, 2026 •

edited

Loading