[FEA] Add Java bindings for multi-output AST JIT

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

@@ -41,12 +41,31 @@ void serialize(ByteBuffer bb) {
     }
   }
 
+  /**
+   * Compile this expression for execution with the process-level backend selection.
+   *
+   * @return expression compatible with default AST consumers
+   */
   public CompiledExpression compile() {
+    return compile(CompiledExpression.CompilationMode.DEFAULT);
+  }
+
+  /**
+   * Compile this expression for explicit execution with the libcudf JIT backend.
+   * The returned expression cannot be used as a join or scan predicate.
+   *
+   * @return expression specialized for JIT execution
+   */
+  public CompiledExpression compileJit() {
+    return compile(CompiledExpression.CompilationMode.JIT);
+  }
+
+  private CompiledExpression compile(CompiledExpression.CompilationMode mode) {
     int size = getSerializedSize();
     ByteBuffer bb = ByteBuffer.allocate(size);
     bb.order(ByteOrder.nativeOrder());
     serialize(bb);
-    return new CompiledExpression(bb.array());
+    return new CompiledExpression(bb.array(), mode);
   }
 
   /** Get the size in bytes of the serialized form of this node and all child nodes */

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

@@ -12,8 +12,15 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Objects;
+
 /** This class wraps a native compiled AST and must be closed to avoid native memory leaks. */
 public class CompiledExpression implements AutoCloseable {
+  enum CompilationMode {
+    DEFAULT,
+    JIT
+  }
+
   static {
     NativeDepsLoader.loadNativeDeps();
   }
@@ -52,42 +59,101 @@ public boolean isClean() {
   }
 
   private final CompiledExpressionCleaner cleaner;
+  private final CompilationMode mode;
   private boolean isClosed = false;
 
   /** Construct a compiled expression from a serialized AST */
-  CompiledExpression(byte[] serializedExpression) {
-    this(compile(serializedExpression));
+  CompiledExpression(byte[] serializedExpression, CompilationMode mode) {
+    this(mode == CompilationMode.JIT ? compileJit(serializedExpression) :
+        compile(serializedExpression), mode);
   }
 
   /** Construct a compiled expression from a native compiled AST pointer */
-  CompiledExpression(long nativeHandle) {
+  private CompiledExpression(long nativeHandle, CompilationMode mode) {
     this.cleaner = new CompiledExpressionCleaner(nativeHandle);
+    this.mode = mode;
     MemoryCleaner.register(this, cleaner);
     cleaner.addRef();
   }
 
   /**
-   * Compute a new column by applying this AST expression to the specified table. All
-   * {@link ColumnReference} instances within the expression will use the sole input table,
-   * even if they try to specify a non-existent table, e.g.: {@link TableReference#RIGHT}.
+   * Compute a new column by applying this AST expression to the specified table. All column
+   * references must use {@link TableReference#LEFT}; references to {@link TableReference#RIGHT}
+   * are rejected because this operation has only one input table.
+   * An expression produced by {@link AstExpression#compileJit()} always uses the JIT backend.
+   * Otherwise, execution uses the process-level backend selection.
+   *
    * @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 compilation or evaluation fails
    */
   public ColumnVector computeColumn(Table table) {
-    return new ColumnVector(computeColumn(cleaner.nativeHandle, table.getNativeView()));
+    long result;
+    try {
+      result = computeColumn(cleaner.nativeHandle, table.getNativeView());
+    } finally {
+      reachabilityFence(this);
+      reachabilityFence(table);
+    }
+    return new ColumnVector(result);
   }
 
   /**
-   * Compute a new column by applying this expression with the libcudf JIT executor, independent
-   * of the process-level backend selected for {@link #computeColumn}.
+   * Compute a new table by applying expressions to the input table in one multi-output JIT
+   * transform. Output column {@code i} contains the result of {@code expressions[i]}.
    *
-   * @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
+   * @param table input table for the expressions
+   * @param expressions non-empty JIT-compiled expressions to evaluate in output order
+   * @return table containing one output column per expression
+   * @throws NullPointerException if the table, expression array, or an expression is null
+   * @throws IllegalArgumentException if no expressions are provided or an expression was not
+   *         produced by {@link AstExpression#compileJit()}
+   * @throws IllegalStateException if the table or an expression is closed
+   * @throws ai.rapids.cudf.CudfException if JIT compilation or evaluation fails
    */
-  public ColumnVector computeColumnJit(Table table) {
-    return new ColumnVector(computeColumnJit(cleaner.nativeHandle, table.getNativeView()));
+  public static Table computeTableJit(Table table, CompiledExpression... expressions) {
+    Objects.requireNonNull(table, "table");
+    Objects.requireNonNull(expressions, "expressions");
+    if (expressions.length == 0) {
+      throw new IllegalArgumentException("At least one expression is required");
+    }
+
+    long tableHandle = table.getNativeView();
+    if (tableHandle == 0) {
+      throw new IllegalStateException("Table is closed");
+    }
+
+    CompiledExpression[] expressionRefs = expressions.clone();
+    long[] nativeHandles = new long[expressionRefs.length];
+    for (int i = 0; i < expressionRefs.length; i++) {
+      CompiledExpression expression = Objects.requireNonNull(
+          expressionRefs[i], "expression " + i + " is null");
+      if (expression.mode != CompilationMode.JIT) {
+        throw new IllegalArgumentException("Expression " + i + " was not compiled for JIT");
+      }
+      nativeHandles[i] = expression.cleaner.nativeHandle;
+      if (nativeHandles[i] == 0) {
+        throw new IllegalStateException("Expression " + i + " is closed");
+      }
+    }
+
+    long[] result;
+    try {
+      result = computeTableJitNative(nativeHandles, tableHandle);
+    } finally {
+      reachabilityFence(table);
+      reachabilityFence(expressionRefs);
+    }
+    return new Table(result);
+  }
+
+  private static void reachabilityFence(Object object) {
+    if (object != null) {
+      synchronized (object) {
+        // The monitor operation is a Java 8 reachability fence.
+      }
+    }
   }
 
   @Override
@@ -101,13 +167,24 @@ public synchronized void close() {
     isClosed = true;
   }
 
-  /** Returns the native address of a compiled expression. Intended for internal cudf use only. */
+  /**
+   * Returns the native address of a default-compatible compiled expression.
+   * Intended for internal cudf use only.
+   *
+   * @throws IllegalStateException if this expression was produced by
+   *         {@link AstExpression#compileJit()}
+   */
   public long getNativeHandle() {
+    if (mode == CompilationMode.JIT) {
+      throw new IllegalStateException(
+          "JIT-compiled expressions cannot be used by a default AST consumer");
+    }
     return cleaner.nativeHandle;
   }
 
   private static native long compile(byte[] serializedExpression);
+  private static native long compileJit(byte[] serializedExpression);
   private static native long computeColumn(long astHandle, long tableHandle);
-  private static native long computeColumnJit(long astHandle, long tableHandle);
+  private static native long[] computeTableJitNative(long[] astHandles, long tableHandle);
   private static native void destroy(long handle);
 }

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

@@ -9,8 +9,8 @@
 import java.util.Objects;
 
 /**
- * A libcudf JIT operation. Expressions containing a JIT operation must be evaluated with
- * {@link CompiledExpression#computeColumnJit}.
+ * A libcudf JIT operation. Expressions containing a JIT operation must be compiled with
+ * {@link AstExpression#compileJit()}.
  * Operator arity, error policy, and target-scale constraints are validated when the expression
  * is compiled.
  */

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

@@ -127,10 +127,9 @@ public static Literal ofDouble(Double value) {
   /**
    * Construct a decimal literal with the specified type and unscaled value.
    * A null {@code unscaledValue} produces a null literal of the requested type.
-   * Root literals of type {@code DECIMAL32} or {@code DECIMAL64} can be evaluated with either
-   * {@link CompiledExpression#computeColumn} or {@link CompiledExpression#computeColumnJit}.
-   * A {@code DECIMAL128} root literal must use {@code computeColumnJit}; the legacy executor
-   * cannot materialize it correctly.
+   * Root literals of type {@code DECIMAL32} or {@code DECIMAL64} can use either compilation mode.
+   * A {@code DECIMAL128} root literal must use {@link AstExpression#compileJit()}; the default AST
+   * executor cannot materialize it correctly.
    *
    * @param type decimal storage type and scale
    * @param unscaledValue unscaled decimal value, or null

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

@@ -452,6 +452,10 @@ cudf::ast::operation const& compile_binary_expression(cudf::jni::ast::compiled_e
 cudf::ast::expression const& compile_jit_expression(cudf::jni::ast::compiled_expr& compiled_expr,
                                                     jni_serialized_ast& jni_ast)
 {
+  if (!compiled_expr.is_jit()) {
+    throw std::invalid_argument("JIT operations require an expression compiled for JIT");
+  }
+
   auto const jni_op_value     = jni_ast.read_byte();
   auto const op_info          = jni_to_jit_operator(jni_op_value);
   auto const jni_policy_value = jni_ast.read_byte();
@@ -530,53 +534,70 @@ cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& c
 }
 
 /** Decode a serialized AST into a native libcudf AST and associated resources */
-std::unique_ptr<cudf::jni::ast::compiled_expr> compile_serialized_ast(jni_serialized_ast& jni_ast)
+std::unique_ptr<cudf::jni::ast::compiled_expr> compile_serialized_ast(
+  jni_serialized_ast& jni_ast, cudf::jni::ast::compilation_mode mode)
 {
-  auto jni_expr_ptr = std::make_unique<cudf::jni::ast::compiled_expr>();
+  auto jni_expr_ptr = std::make_unique<cudf::jni::ast::compiled_expr>(mode);
   (void)compile_expression(*jni_expr_ptr, jni_ast);
 
   if (!jni_ast.at_eof()) { throw std::invalid_argument("Extra bytes at end of serialized AST"); }
 
   // The expression may be handed to a thread with a different default stream.
-  if (jni_expr_ptr->has_literals()) { cudf::get_default_stream().synchronize(); }
+  if (jni_expr_ptr->has_literals()) {
+    cudf::get_default_stream().synchronize();
+    // JIT literals retain only the copied one-row columns after construction completes.
+    jni_expr_ptr->release_jit_staging_scalars();
+  }
 
   return jni_expr_ptr;
 }
 
-enum class execution_backend { DEFAULT, JIT };
-
-jlong execute_compiled_expression(jlong j_ast, jlong j_table, execution_backend backend)
+jlong execute_compiled_expression(jlong j_ast, jlong j_table)
 {
   auto compiled_expr_ptr = reinterpret_cast<cudf::jni::ast::compiled_expr const*>(j_ast);
   auto tview_ptr         = reinterpret_cast<cudf::table_view const*>(j_table);
-  auto const& expression = compiled_expr_ptr->get_top_expression();
-  std::unique_ptr<cudf::column> result = backend == execution_backend::JIT
-                                           ? cudf::compute_column_jit(*tview_ptr, expression)
-                                           : cudf::compute_column(*tview_ptr, expression);
+  std::unique_ptr<cudf::column> result =
+    compiled_expr_ptr->is_jit()
+      ? cudf::compute_column_jit(*tview_ptr, compiled_expr_ptr->get_jit_top_expression())
+      : cudf::compute_column(*tview_ptr, compiled_expr_ptr->get_top_expression());
   return reinterpret_cast<jlong>(result.release());
 }
 
-}  // anonymous namespace
-
-extern "C" {
-
-JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_compile(JNIEnv* env,
-                                                                           jclass,
-                                                                           jbyteArray jni_data)
+jlong compile_serialized_expression(JNIEnv* env,
+                                    jbyteArray jni_data,
+                                    cudf::jni::ast::compilation_mode mode)
 {
   JNI_NULL_CHECK(env, jni_data, "Serialized AST data is null", 0);
   JNI_TRY
   {
     cudf::jni::auto_set_device(env);
     cudf::jni::native_jbyteArray jbytes(env, jni_data);
     jni_serialized_ast jni_ast(jbytes);
-    auto compiled_expr_ptr = compile_serialized_ast(jni_ast);
+    auto compiled_expr_ptr = compile_serialized_ast(jni_ast, mode);
     jbytes.cancel();
     return reinterpret_cast<jlong>(compiled_expr_ptr.release());
   }
   JNI_CATCH(env, 0);
 }
 
+}  // anonymous namespace
+
+extern "C" {
+
+JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_compile(JNIEnv* env,
+                                                                           jclass,
+                                                                           jbyteArray jni_data)
+{
+  return compile_serialized_expression(env, jni_data, cudf::jni::ast::compilation_mode::DEFAULT);
+}
+
+JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_compileJit(JNIEnv* env,
+                                                                              jclass,
+                                                                              jbyteArray jni_data)
+{
+  return compile_serialized_expression(env, jni_data, cudf::jni::ast::compilation_mode::JIT);
+}
+
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumn(JNIEnv* env,
                                                                                  jclass,
                                                                                  jlong j_ast,
@@ -587,24 +608,37 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumn
   JNI_TRY
   {
     cudf::jni::auto_set_device(env);
-    return execute_compiled_expression(j_ast, j_table, execution_backend::DEFAULT);
+    return execute_compiled_expression(j_ast, j_table);
   }
   JNI_CATCH(env, 0);
 }
 
-JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumnJit(JNIEnv* env,
-                                                                                    jclass,
-                                                                                    jlong j_ast,
-                                                                                    jlong j_table)
+JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeTableJitNative(
+  JNIEnv* env, jclass, jlongArray j_asts, jlong j_table)
 {
-  JNI_NULL_CHECK(env, j_ast, "Compiled AST pointer is null", 0);
-  JNI_NULL_CHECK(env, j_table, "Table view pointer is null", 0);
+  JNI_NULL_CHECK(env, j_asts, "Compiled AST pointer array is null", nullptr);
+  JNI_NULL_CHECK(env, j_table, "Table view pointer is null", nullptr);
   JNI_TRY
   {
     cudf::jni::auto_set_device(env);
-    return execute_compiled_expression(j_ast, j_table, execution_backend::JIT);
+    cudf::jni::native_jlongArray ast_handles(env, j_asts);
+    if (ast_handles.size() == 0) { throw std::invalid_argument("At least one AST is required"); }
+
+    std::vector<std::reference_wrapper<cudf::ast::expression const>> expressions;
+    expressions.reserve(ast_handles.size());
+    for (auto const handle : ast_handles) {
+      if (handle == 0) { throw std::invalid_argument("Compiled AST pointer is null"); }
+      auto const* compiled_expr_ptr =
+        reinterpret_cast<cudf::jni::ast::compiled_expr const*>(handle);
+      expressions.emplace_back(compiled_expr_ptr->get_jit_top_expression());
+    }
+    ast_handles.cancel();
+
+    auto const* tview_ptr = reinterpret_cast<cudf::table_view const*>(j_table);
+    return cudf::jni::convert_table_for_return(env,
+                                               cudf::compute_table_jit(*tview_ptr, expressions));
   }
-  JNI_CATCH(env, 0);
+  JNI_CATCH(env, nullptr);
 }
 
 JNIEXPORT void JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_destroy(JNIEnv* env,