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

cpp/src/transform/transform.cu

@@ -1521,8 +1521,12 @@ transform_program::transform_program(
     impl_->ast_input_types_.push_back(std::visit([](auto& view) { return view.type(); }, input));
     impl_->ast_input_nullable_.push_back(
       std::visit([](auto& view) { return view.nullable(); }, input));
+    if (auto const* scalar = std::get_if<scalar_column_view>(&input)) {
+      // The program must outlive non-owning scalar-column literals in the source AST.
+      impl_->ast_scalar_columns_.push_back(
+        std::make_unique<column>(scalar->as_column_view(), stream, mr));
+    }
   }
-  impl_->ast_scalar_columns_       = std::move(args.scalar_columns);
   impl_->ast_input_column_indices_ = std::move(args.input_column_indices);
   impl_->ast_outputs_              = std::move(args.outputs);
 }

cpp/tests/ast/transform_tests.cpp

@@ -110,6 +110,52 @@ TEST_F(TransformProgramTest, ReusesAstWithCompatibleTable)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity);
 }
 
+TEST_F(TransformProgramTest, OwnsScalarColumnViewLiterals)
+{
+  std::unique_ptr<cudf::transform_program> program;
+  {
+    auto construction_input = column_wrapper<int32_t>{3, 20, 1, 50};
+    auto construction_table = cudf::table_view{{construction_input}};
+    auto literal_column     = column_wrapper<int32_t>{2};
+    auto column_ref         = cudf::ast::column_reference{0};
+    auto literal            = cudf::ast::literal{cudf::scalar_column_view{literal_column}};
+    auto expression = cudf::ast::operation{cudf::ast::ast_operator::ADD, column_ref, literal};
+    std::reference_wrapper<cudf::ast::expression const> expressions[] = {expression};
+
+    program = std::make_unique<cudf::transform_program>(construction_table, expressions);
+  }
+
+  auto input    = column_wrapper<int32_t>{10, 20, 30};
+  auto table    = cudf::table_view{{input}};
+  auto expected = column_wrapper<int32_t>{12, 22, 32};
+  auto result   = std::move(program->run(table)->release().front());
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity);
+}
+
+TEST_F(TransformProgramTest, OwnsStringScalarColumnViewLiterals)
+{
+  std::unique_ptr<cudf::transform_program> program;
+  {
+    auto construction_input = cudf::test::strings_column_wrapper{"a", "ccc"};
+    auto construction_table = cudf::table_view{{construction_input}};
+    auto literal_column     = cudf::test::strings_column_wrapper{"ccc"};
+    auto column_ref         = cudf::ast::column_reference{0};
+    auto literal            = cudf::ast::literal{cudf::scalar_column_view{literal_column}};
+    auto expression = cudf::ast::operation{cudf::ast::ast_operator::LESS, column_ref, literal};
+    std::reference_wrapper<cudf::ast::expression const> expressions[] = {expression};
+
+    program = std::make_unique<cudf::transform_program>(construction_table, expressions);
+  }
+
+  auto input    = cudf::test::strings_column_wrapper{"a", "ccc", "dddd"};
+  auto table    = cudf::table_view{{input}};
+  auto expected = column_wrapper<bool>{true, false, false};
+  auto result   = std::move(program->run(table)->release().front());
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity);
+}
+
 TEST_F(TransformProgramTest, RejectsIncompatibleTable)
 {
   auto construction_input = column_wrapper<int32_t>{3, 20, 1, 50};

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

@@ -7,6 +7,7 @@
 
 package ai.rapids.cudf;
 
+import ai.rapids.cudf.ast.AstJitProgram;
 import ai.rapids.cudf.ast.CompiledExpression;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -367,7 +368,13 @@ static void register(CuFileHandle handle, Cleaner cleaner) {
   }
 
   public static void register(CompiledExpression expr, Cleaner cleaner) {
-    all.put(cleaner.id, new CleanerWeakReference(expr, cleaner, collected, false));
+    // Both compilation modes retain device-backed literal values.
+    all.put(cleaner.id, new CleanerWeakReference(expr, cleaner, collected, true));
+  }
+
+  public static void register(AstJitProgram program, Cleaner cleaner) {
+    // AST programs retain copied literal columns across evaluations.
+    all.put(cleaner.id, new CleanerWeakReference(program, cleaner, collected, true));
   }
 
   static void register(HybridScanReader reader, Cleaner cleaner) {

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

@@ -41,14 +41,39 @@ void serialize(ByteBuffer bb) {
     }
   }
 
+  /**
+   * Compile this expression for execution with the process-level backend selection.
+   *
+   * @return expression compatible with default AST consumers
+   * @throws IllegalArgumentException if a root literal requires JIT compilation
+   * @throws ai.rapids.cudf.CudfException if compilation fails
+   */
   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
+   * @throws ai.rapids.cudf.CudfException if compilation fails
+   */
+  public CompiledExpression compileJit() {
+    return compile(CompiledExpression.CompilationMode.JIT);
+  }
+
+  private CompiledExpression compile(CompiledExpression.CompilationMode mode) {
+    validateCompilationMode(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);
   }
 
+  void validateCompilationMode(CompiledExpression.CompilationMode mode) {}
+
   /** Get the size in bytes of the serialized form of this node and all child nodes */
   abstract int getSerializedSize();
 

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

@@ -0,0 +1,159 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package ai.rapids.cudf.ast;
+
+import ai.rapids.cudf.MemoryCleaner;
+import ai.rapids.cudf.NativeDepsLoader;
+import ai.rapids.cudf.Table;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Objects;
+
+/**
+ * A reusable AST JIT program specialized to an input schema.
+ * Construction lowers the expressions and retrieves their JIT kernel. Subsequent calls reuse that
+ * kernel with tables whose referenced columns have compatible types and nullability.
+ * Callers must ensure that {@link #close()} does not overlap with {@link #computeTable(Table)}.
+ */
+public final class AstJitProgram implements AutoCloseable {
+  static {
+    NativeDepsLoader.loadNativeDeps();
+  }
+
+  private static final Logger log = LoggerFactory.getLogger(AstJitProgram.class);
+
+  private static final class AstJitProgramCleaner extends MemoryCleaner.Cleaner {
+    private long nativeHandle;
+
+    AstJitProgramCleaner(long nativeHandle) {
+      this.nativeHandle = nativeHandle;
+    }
+
+    @Override
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
+      boolean alreadyClean = nativeHandle == 0;
+      if (alreadyClean) {
+        return false;
+      }
+      long origAddress = nativeHandle;
+      try {
+        destroy(nativeHandle);
+      } finally {
+        nativeHandle = 0;
+      }
+      if (logErrorIfNotClean) {
+        log.error("AN AST JIT PROGRAM WAS LEAKED (ID: {} {})", id,
+            Long.toHexString(origAddress));
+      }
+      return true;
+    }
+
+    @Override
+    public boolean isClean() {
+      return nativeHandle == 0;
+    }
+  }
+
+  private final AstJitProgramCleaner cleaner;
+  private boolean isClosed = false;
+
+  private AstJitProgram(long nativeHandle) {
+    AstJitProgramCleaner newCleaner = null;
+    try {
+      newCleaner = new AstJitProgramCleaner(nativeHandle);
+      cleaner = newCleaner;
+      MemoryCleaner.register(this, cleaner);
+      cleaner.addRef();
+    } catch (Throwable t) {
+      try {
+        if (newCleaner == null) {
+          destroy(nativeHandle);
+        } else {
+          newCleaner.clean(false);
+        }
+      } catch (Throwable cleanupFailure) {
+        t.addSuppressed(cleanupFailure);
+      }
+      throw t;
+    }
+  }
+
+  /**
+   * Compile a reusable program from one or more JIT-compiled expressions.
+   * The schema table and expressions are inspected during construction but are not retained. The
+   * returned program owns any literal values required by later evaluations.
+   *
+   * @param schemaTable table whose referenced column schema is used to compile the program
+   * @param expressions non-empty JIT-compiled expressions in output order
+   * @return a reusable AST JIT program
+   * @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 fails
+   */
+  public static AstJitProgram compile(Table schemaTable, CompiledExpression... expressions) {
+    long tableHandle = Objects.requireNonNull(schemaTable, "schemaTable").getNativeView();
+    CompiledExpression.JitExpressionArgs expressionArgs =
+        CompiledExpression.getJitExpressionArgs(expressions);
+    if (tableHandle == 0) {
+      throw new IllegalStateException("Table is closed");
+    }
+
+    try {
+      return new AstJitProgram(create(expressionArgs.nativeHandles, tableHandle));
+    } finally {
+      CompiledExpression.reachabilityFence(schemaTable);
+      CompiledExpression.reachabilityFence(expressionArgs.expressionRefs);
+    }
+  }
+
+  /**
+   * Evaluate this program on a table with a compatible referenced-column schema.
+   * The row count and unreferenced columns may differ from the schema table used at compilation.
+   * Calling {@link #close()} while an evaluation is in progress is unsupported.
+   *
+   * @param table input table for expression evaluation
+   * @return table containing the program outputs in expression order
+   * @throws NullPointerException if the table is null
+   * @throws IllegalStateException if the program or table is closed
+   * @throws ai.rapids.cudf.CudfException if the referenced-column schema is incompatible or
+   *         evaluation fails
+   */
+  public Table computeTable(Table table) {
+    long tableHandle = Objects.requireNonNull(table, "table").getNativeView();
+    long programHandle = cleaner.nativeHandle;
+    if (programHandle == 0) {
+      throw new IllegalStateException("AST JIT program is closed");
+    }
+    if (tableHandle == 0) {
+      throw new IllegalStateException("Table is closed");
+    }
+
+    try {
+      return new Table(computeTableNative(programHandle, tableHandle));
+    } finally {
+      CompiledExpression.reachabilityFence(this);
+      CompiledExpression.reachabilityFence(table);
+    }
+  }
+
+  @Override
+  public synchronized void close() {
+    cleaner.delRef();
+    if (isClosed) {
+      cleaner.logRefCountDebug("double free " + this);
+      throw new IllegalStateException("Close called too many times " + this);
+    }
+    cleaner.clean(false);
+    isClosed = true;
+  }
+
+  private static native long create(long[] astHandles, long tableHandle);
+  private static native long[] computeTableNative(long programHandle, long tableHandle);
+  private static native void destroy(long handle);
+}