[FEA] Add Java bindings for AST JIT operations

cpp/src/ast/expression_parser.cpp

@@ -265,6 +265,9 @@ cudf::size_type expression_parser::visit(operation const& expr)
   auto const output = [&]() {
     if (expression_index == 0) {
       // This expression is the root. Output should be directed to the output column.
+      CUDF_EXPECTS(data_type.id() != cudf::type_id::DECIMAL128,
+                   "decimal128 is not supported as an AST expression output.",
+                   cudf::data_type_error);
       return detail::device_data_reference(
         detail::device_data_reference_type::COLUMN, data_type, 0, table_reference::OUTPUT);
     } else {

cpp/tests/ast/transform_tests.cpp

@@ -1514,9 +1514,8 @@ TYPED_TEST(TransformTest, Decimal128Unsupported)
 
   // column = {2000.00, 2000.00}  (rep 200000 @ scale -2)
   auto const scale = numeric::scale_type{-2};
-  auto const col   = cudf::test::fixed_point_column_wrapper<__int128>{
-    {__int128_t{200000}, __int128_t{200000}}, scale};
-  auto table = cudf::table_view{{col}};
+  auto const col   = cudf::test::fixed_point_column_wrapper<__int128>{{200000, 200000}, scale};
+  auto table       = cudf::table_view{{col}};
 
   // literal = 0.5 @ scale -2  (rep 50)
   auto half = numeric::decimal128{numeric::scaled_integer<numeric::decimal128::rep>{50, scale}};
@@ -1526,20 +1525,37 @@ TYPED_TEST(TransformTest, Decimal128Unsupported)
   auto cr = cudf::ast::column_reference(0);
 
   auto ast = cudf::ast::operation(cudf::ast::ast_operator::MUL, lr, cr);
-  EXPECT_THROW(cudf::compute_column(table, ast), cudf::data_type_error);
-
   if constexpr (std::is_same_v<Executor, executor_ast>) {
-    auto ast2 = cudf::ast::operation(cudf::ast::ast_operator::MUL, lr, cr);
-    EXPECT_THROW(cudf::compute_column(table, ast), cudf::data_type_error);
+    EXPECT_THROW(Executor::compute_column(table, ast), cudf::data_type_error);
   } else {
     auto result = Executor::compute_column(table, ast);
     // Expected: 0.5 * 2000.00 = 1000.00  => rep 10000000 @ scale -4
     EXPECT_EQ(result->type().id(), cudf::type_id::DECIMAL128);
     EXPECT_EQ(result->type().scale(), numeric::scale_type{-4});
-    auto expected = cudf::test::fixed_point_column_wrapper<__int128>{
-      {__int128_t{10000000}, __int128_t{10000000}}, numeric::scale_type{-4}};
+    auto expected = cudf::test::fixed_point_column_wrapper<__int128>{{10000000, 10000000},
+                                                                     numeric::scale_type{-4}};
     CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, expected);
   }
 }
 
+TYPED_TEST(TransformTest, Decimal128IdentityOutput)
+{
+  using Executor = TypeParam;
+
+  auto const input   = column_wrapper<int32_t>{0, 0};
+  auto const table   = cudf::table_view{{input}};
+  auto const scale   = numeric::scale_type{-2};
+  auto literal_value = cudf::fixed_point_scalar<numeric::decimal128>(12345, scale, true);
+  auto literal       = cudf::ast::literal(literal_value);
+  auto expression    = cudf::ast::operation(cudf::ast::ast_operator::IDENTITY, literal);
+
+  if constexpr (std::is_same_v<Executor, executor_ast>) {
+    EXPECT_THROW(Executor::compute_column(table, expression), cudf::data_type_error);
+  } else {
+    auto result   = Executor::compute_column(table, expression);
+    auto expected = cudf::test::fixed_point_column_wrapper<__int128>({12345, 12345}, scale);
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity);
+  }
+}
+
 CUDF_TEST_PROGRAM_MAIN()

java/src/main/java/ai/rapids/cudf/Cudf.java

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 

java/src/main/java/ai/rapids/cudf/ast/AstExpression.java

@@ -12,21 +12,22 @@
 public abstract class AstExpression {
   /**
    * Enumeration for the types of AST nodes that can appear in a serialized AST.
-   * NOTE: This must be kept in sync with the `jni_serialized_expression_type` in CompiledExpression.cpp!
+   * NOTE: This must be kept in sync with `jni_serialized_expression_type` in
+   * CompiledExpression.cpp!
    */
   protected enum ExpressionType {
     VALID_LITERAL(0),
     NULL_LITERAL(1),
     COLUMN_REFERENCE(2),
     UNARY_EXPRESSION(3),
     BINARY_EXPRESSION(4),
-    COLUMN_NAME_REFERENCE(5);
+    COLUMN_NAME_REFERENCE(5),
+    JIT_EXPRESSION(6);
 
     private final byte nativeId;
 
     ExpressionType(int nativeId) {
-      this.nativeId = (byte) nativeId;
-      assert this.nativeId == nativeId;
+      this.nativeId = AstUtils.checkByte(nativeId);
     }
 
     /** Get the size in bytes to serialize this node type */

java/src/main/java/ai/rapids/cudf/ast/AstUtils.java

@@ -0,0 +1,19 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package ai.rapids.cudf.ast;
+
+final class AstUtils {
+  private AstUtils() {
+  }
+
+  static byte checkByte(int value) {
+    byte result = (byte) value;
+    if (result != value) {
+      throw new IllegalArgumentException("value does not fit in a byte: " + value);
+    }
+    return result;
+  }
+}

java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -40,8 +40,7 @@ public enum BinaryOperator {
   private final byte nativeId;
 
   BinaryOperator(int nativeId) {
-    this.nativeId = (byte) nativeId;
-    assert this.nativeId == nativeId;
+    this.nativeId = AstUtils.checkByte(nativeId);
   }
 
   /** Get the size in bytes to serialize this operator */

java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -77,6 +77,19 @@ public ColumnVector computeColumn(Table table) {
     return new ColumnVector(computeColumn(cleaner.nativeHandle, table.getNativeView()));
   }
 
+  /**
+   * Compute a new column by applying this expression with the libcudf JIT executor, independent
+   * of the process-level backend selected for {@link #computeColumn}.
+   *
+   * @param table input table for this expression
+   * @return new column computed from this expression applied to the input table
+   * @throws ai.rapids.cudf.CudfException if the expression refers to
+   *         {@link TableReference#RIGHT}, or if JIT compilation or evaluation fails
+   */
+  public ColumnVector computeColumnJit(Table table) {
+    return new ColumnVector(computeColumnJit(cleaner.nativeHandle, table.getNativeView()));
+  }
+
   @Override
   public synchronized void close() {
     cleaner.delRef();
@@ -95,5 +108,6 @@ public long getNativeHandle() {
 
   private static native long compile(byte[] serializedExpression);
   private static native long computeColumn(long astHandle, long tableHandle);
+  private static native long computeColumnJit(long astHandle, long tableHandle);
   private static native void destroy(long handle);
 }

java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java

@@ -0,0 +1,34 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package ai.rapids.cudf.ast;
+
+import java.nio.ByteBuffer;
+
+/**
+ * Error handling policy for fallible JIT AST operations.
+ *
+ * NOTE: This must be kept in sync with `jni_to_jit_error_policy` in CompiledExpression.cpp!
+ */
+public enum JitErrorPolicy {
+  /** Propagate an evaluation error to the caller. */
+  PROPAGATE(0),
+  /** Produce null for a row where evaluation fails. */
+  NULLIFY(1);
+
+  private final byte nativeId;
+
+  JitErrorPolicy(int nativeId) {
+    this.nativeId = AstUtils.checkByte(nativeId);
+  }
+
+  int getSerializedSize() {
+    return Byte.BYTES;
+  }
+
+  void serialize(ByteBuffer bb) {
+    bb.put(nativeId);
+  }
+}

java/src/main/java/ai/rapids/cudf/ast/JitOperation.java

@@ -0,0 +1,125 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package ai.rapids.cudf.ast;
+
+import java.nio.ByteBuffer;
+import java.util.Objects;
+
+/**
+ * A libcudf JIT operation. Expressions containing a JIT operation must be evaluated with
+ * {@link CompiledExpression#computeColumnJit}.
+ * Operator arity, error policy, and target-scale constraints are validated when the expression
+ * is compiled.
+ */
+public final class JitOperation extends AstExpression {
+  private final JitOperator op;
+  private final JitErrorPolicy errorPolicy;
+  private final AstExpression[] inputs;
+  private final Integer targetScale;
+
+  /**
+   * Construct an operation that propagates evaluation errors.
+   *
+   * @param op operator to apply
+   * @param inputs operator inputs
+   * @throws NullPointerException if {@code op}, {@code inputs}, or an input is null
+   */
+  public JitOperation(JitOperator op, AstExpression... inputs) {
+    this(op, JitErrorPolicy.PROPAGATE, null, inputs);
+  }
+
+  /**
+   * Construct an operation with an explicit error policy.
+   *
+   * @param op operator to apply
+   * @param errorPolicy error handling policy
+   * @param inputs operator inputs
+   * @throws NullPointerException if any argument or input is null
+   */
+  public JitOperation(JitOperator op, JitErrorPolicy errorPolicy, AstExpression... inputs) {
+    this(op, errorPolicy, null, inputs);
+  }
+
+  /**
+   * Construct a target-scale operation that propagates evaluation errors.
+   * The target scale is valid only for {@link JitOperator#RESCALE}.
+   *
+   * @param op operator to apply
+   * @param targetScale target fixed-point scale
+   * @param inputs operator inputs
+   * @throws NullPointerException if {@code op}, {@code inputs}, or an input is null
+   */
+  public JitOperation(JitOperator op, int targetScale, AstExpression... inputs) {
+    this(op, JitErrorPolicy.PROPAGATE, Integer.valueOf(targetScale), inputs);
+  }
+
+  private JitOperation(
+      JitOperator op,
+      JitErrorPolicy errorPolicy,
+      Integer targetScale,
+      AstExpression... inputs) {
+    this.op = Objects.requireNonNull(op, "op is null");
+    this.errorPolicy = Objects.requireNonNull(errorPolicy, "errorPolicy is null");
+    this.inputs = Objects.requireNonNull(inputs, "inputs is null").clone();
+    this.targetScale = targetScale;
+    for (int i = 0; i < this.inputs.length; i++) {
+      Objects.requireNonNull(this.inputs[i], "input " + i + " is null");
+    }
+  }
+
+  @Override
+  int getSerializedSize() {
+    int size = ExpressionType.JIT_EXPRESSION.getSerializedSize() +
+        op.getSerializedSize() +
+        errorPolicy.getSerializedSize() +
+        Byte.BYTES +  // targetScale present
+        Byte.BYTES;   // inputs.length
+    if (targetScale != null) {
+      size += Integer.BYTES;
+    }
+    for (AstExpression input : inputs) {
+      size += input.getSerializedSize();
+    }
+    return size;
+  }
+
+  @Override
+  void serialize(ByteBuffer bb) {
+    ExpressionType.JIT_EXPRESSION.serialize(bb);
+    op.serialize(bb);
+    errorPolicy.serialize(bb);
+    bb.put((byte) (targetScale == null ? 0 : 1));
+    if (targetScale != null) {
+      bb.putInt(targetScale);
+    }
+    bb.put((byte) inputs.length);
+    for (AstExpression input : inputs) {
+      input.serialize(bb);
+    }
+  }
+
+  @Override
+  public String toString() {
+    StringBuilder ret = new StringBuilder(op.toString());
+    if (errorPolicy != JitErrorPolicy.PROPAGATE) {
+      ret.append("[").append(errorPolicy).append("]");
+    }
+    ret.append("(");
+    for (int i = 0; i < inputs.length; i++) {
+      if (i > 0) {
+        ret.append(", ");
+      }
+      ret.append(inputs[i]);
+    }
+    if (targetScale != null) {
+      if (inputs.length > 0) {
+        ret.append(", ");
+      }
+      ret.append("scale=").append(targetScale);
+    }
+    return ret.append(")").toString();
+  }
+}