[FEA] Add Java bindings for AST JIT operations by thirtiseven · Pull Request #23117 · rapidsai/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

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a17ddb75-f37f-45a5-a61e-88e16de37de0

📥 Commits

Reviewing files that changed from the base of the PR and between 19f974b and 30cceaf.

📒 Files selected for processing (3)

java/src/main/java/ai/rapids/cudf/ast/JitOperation.java
java/src/main/java/ai/rapids/cudf/ast/JitOperator.java
java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java

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

java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java

📝 Walkthrough

Summary by CodeRabbit

New Features
- Added initializeJitRuntime (idempotent) and a new computeColumnJit execution option for compiled expressions.
- Introduced JIT AST components: JitOperator, JitErrorPolicy, and JitOperation (including RESCALE and nullification/error-handling behavior).
- Added Literal.ofDecimal(DType, BigInteger) for constructing large decimal literals.
Bug Fixes
- Fixed decimal scale serialization/deserialization to use full integer precision.
Tests
- Expanded JIT tests for initialization repeatability, constructor/validation failures, casting/precision behavior, conditionals, overflow handling, and end-to-end correctness.
Documentation
- Updated public Javadoc/enum notes to keep JIT and serialized definitions aligned.

Walkthrough

This PR adds Java and JNI support for libcudf JIT AST execution, including runtime initialization, JIT operators and expressions, decimal literal serialization, native compilation and execution paths, and comprehensive validation and execution tests.

Changes

JIT AST Execution

Layer / File(s)	Summary
JIT runtime initialization API `java/src/main/java/ai/rapids/cudf/Cudf.java`, `java/src/main/native/src/CudfJni.cpp`, `java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java`	Adds `initializeJitRuntime`, its JNI implementation, and repeat-initialization coverage.
JIT AST contracts `java/src/main/java/ai/rapids/cudf/ast/AstExpression.java`, `java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java`, `java/src/main/java/ai/rapids/cudf/ast/JitOperator.java`, `java/src/main/java/ai/rapids/cudf/ast/JitOperation.java`	Adds serialized JIT expression types, operators, error policies, and JIT operation serialization.
Decimal literal serialization `java/src/main/java/ai/rapids/cudf/ast/Literal.java`	Adds BigInteger decimal literal construction and updates decimal scale and byte-layout serialization.
Native AST tree construction `java/src/main/native/src/jni_compiled_expr.hpp`	Refactors native expression ownership 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: Directly covers expanding Java bindings with libcudf JIT expression support.

Possibly related PRs

rapidsai/cudf#22602: Modifies the underlying compute_column_jit error-policy and execution plumbing used by this JNI path.
rapidsai/cudf#22654: Refactors the JIT runtime/cache initialization used by initializeJitRuntime.
rapidsai/cudf#22456: Also changes serialized AST node-type synchronization around COLUMN_NAME_REFERENCE.

Suggested labels: libcudf

Suggested reviewers: vyasr, bdice, mhaseeb123, igorpeshansky, mythrocks

🚥 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 detailed and directly matches the changeset and PR objectives.
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

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

/ok to test 30cceaf

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[FEA] Add Java bindings for AST JIT operations#23117

[FEA] Add Java bindings for AST JIT operations#23117
thirtiseven wants to merge 8 commits into
rapidsai:mainfrom
thirtiseven:codex-ansi-jit-java-bindings-main-upmerge

thirtiseven commented Jul 6, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 6, 2026

Uh oh!

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading

Summary by CodeRabbit

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!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

lamarrr left a comment

Uh oh!

lamarrr Jul 14, 2026

Uh oh!

lamarrr commented Jul 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

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!

Summary by CodeRabbit

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!

lamarrr commented Jul 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

thirtiseven commented Jul 6, 2026 •

edited

Loading

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading

coderabbitai Bot commented Jul 8, 2026 •

edited

Loading