From aa0963c8c2e7fcb96517f319c4db6f1130a7184b Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 26 Aug 2026 16:17:39 +0800 Subject: [PATCH 1/4] Add Java bindings for multi-output AST JIT Expose computeTableJit to evaluate compiled AST roots in one libcudf call. Keep scalar-column-backed JIT trees alongside regular trees so compiled literals can be reused across evaluations. Signed-off-by: Haoyang Li --- .../rapids/cudf/ast/CompiledExpression.java | 74 ++++++- .../main/native/src/CompiledExpression.cpp | 127 +++++++----- .../src/main/native/src/jni_compiled_expr.hpp | 80 ++++++-- .../cudf/ast/CompiledExpressionTest.java | 187 ++++++++++++++++++ 4 files changed, 397 insertions(+), 71 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java index 52c3c3ac03d7..031699aa229e 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java @@ -12,6 +12,8 @@ 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 { static { @@ -74,7 +76,14 @@ public boolean isClean() { * @return new column computed from this expression applied to the input table */ 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); } /** @@ -87,7 +96,67 @@ public ColumnVector computeColumn(Table table) { * {@link TableReference#RIGHT}, or if JIT compilation or evaluation fails */ public ColumnVector computeColumnJit(Table table) { - return new ColumnVector(computeColumnJit(cleaner.nativeHandle, table.getNativeView())); + long result; + try { + result = computeColumnJit(cleaner.nativeHandle, table.getNativeView()); + } finally { + reachabilityFence(this); + reachabilityFence(table); + } + return new ColumnVector(result); + } + + /** + * 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 the expressions + * @param expressions non-empty 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 + * @throws IllegalStateException if the table or an expression is closed + * @throws ai.rapids.cudf.CudfException if JIT compilation or evaluation fails + */ + 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"); + 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 @@ -109,5 +178,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 long[] computeTableJitNative(long[] astHandles, long tableHandle); private static native void destroy(long handle); } diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index 44748ef6b9b9..bd8a41863a79 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -293,10 +293,10 @@ cudf::ast::table_reference jni_to_table_reference(jbyte jni_value) struct make_literal { /** Construct an AST literal from a numeric value */ template ()>* = nullptr> - cudf::ast::literal const& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_numeric_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -312,10 +312,10 @@ struct make_literal { /** Construct an AST literal from a timestamp value */ template ()>* = nullptr> - cudf::ast::literal const& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_timestamp_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -331,10 +331,10 @@ struct make_literal { /** Construct an AST literal from a duration value */ template ()>* = nullptr> - cudf::ast::literal const& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_duration_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -350,10 +350,10 @@ struct make_literal { /** Construct an AST literal from a string value */ template >* = nullptr> - cudf::ast::literal const& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = [&]() { if (is_valid) { @@ -370,10 +370,10 @@ struct make_literal { /** Construct an AST literal from a fixed-point value */ template ()>* = nullptr> - cudf::ast::literal const& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { using rep_type = typename T::rep; auto const val = is_valid ? jni_ast.read() : rep_type{}; @@ -390,26 +390,26 @@ struct make_literal { std::enable_if_t() && !cudf::is_timestamp() && !cudf::is_duration() && !cudf::is_fixed_point() && !std::is_same_v>* = nullptr> - cudf::ast::literal const& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { throw std::logic_error("Unsupported AST literal type"); } }; /** Decode a serialized AST literal */ -cudf::ast::literal const& compile_literal(bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::jni::ast::expression_pair compile_literal(bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { auto const dtype = jni_ast.read_cudf_type(); return cudf::type_dispatcher(dtype, make_literal{}, dtype, is_valid, compiled_expr, jni_ast); } /** Decode a serialized AST column reference */ -cudf::ast::column_reference const& compile_column_reference( +cudf::jni::ast::expression_pair compile_column_reference( cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { auto const table_ref = jni_to_table_reference(jni_ast.read_byte()); @@ -418,7 +418,7 @@ cudf::ast::column_reference const& compile_column_reference( } /** Decode a serialized AST column name reference */ -cudf::ast::column_name_reference const& compile_column_name_reference( +cudf::jni::ast::expression_pair compile_column_name_reference( cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { std::string column_name = jni_ast.read(); @@ -426,31 +426,31 @@ cudf::ast::column_name_reference const& compile_column_name_reference( } // forward declaration -cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast); +cudf::jni::ast::expression_pair compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast); /** Decode a serialized AST unary expression */ -cudf::ast::operation const& compile_unary_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::jni::ast::expression_pair compile_unary_expression( + cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { - auto const ast_op = jni_to_unary_operator(jni_ast.read_byte()); - cudf::ast::expression const& child_expression = compile_expression(compiled_expr, jni_ast); + auto const ast_op = jni_to_unary_operator(jni_ast.read_byte()); + auto const child_expression = compile_expression(compiled_expr, jni_ast); return compiled_expr.add_operation(ast_op, child_expression); } /** Decode a serialized AST binary expression */ -cudf::ast::operation const& compile_binary_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::jni::ast::expression_pair compile_binary_expression( + cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { - auto const ast_op = jni_to_binary_operator(jni_ast.read_byte()); - cudf::ast::expression const& left_child = compile_expression(compiled_expr, jni_ast); - cudf::ast::expression const& right_child = compile_expression(compiled_expr, jni_ast); + auto const ast_op = jni_to_binary_operator(jni_ast.read_byte()); + auto const left_child = compile_expression(compiled_expr, jni_ast); + auto const right_child = compile_expression(compiled_expr, jni_ast); return compiled_expr.add_operation(ast_op, left_child, right_child); } /** Decode a serialized JIT AST expression */ -cudf::ast::expression const& compile_jit_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::jni::ast::expression_pair compile_jit_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { auto const jni_op_value = jni_ast.read_byte(); auto const op_info = jni_to_jit_operator(jni_op_value); @@ -493,21 +493,18 @@ cudf::ast::expression const& compile_jit_expression(cudf::jni::ast::compiled_exp op_info.arity)); } - std::vector> args; + std::vector args; args.reserve(arity); for (int32_t index = 0; index < arity; ++index) { args.emplace_back(compile_expression(compiled_expr, jni_ast)); } - return compiled_expr.add_jit_expression( - [&](cudf::ast::tree& tree) -> cudf::ast::expression const& { - return cudf::ast::jit::operation(tree, op_info.op, args, error_policy, target_scale); - }); + return compiled_expr.add_jit_operation(op_info.op, args, error_policy, target_scale); } /** Decode a serialized AST expression by reading the expression type and dispatching */ -cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::jni::ast::expression_pair compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { auto const expression_type = static_cast(jni_ast.read_byte()); switch (expression_type) { @@ -549,7 +546,9 @@ jlong execute_compiled_expression(jlong j_ast, jlong j_table, execution_backend { auto compiled_expr_ptr = reinterpret_cast(j_ast); auto tview_ptr = reinterpret_cast(j_table); - auto const& expression = compiled_expr_ptr->get_top_expression(); + auto const& expression = backend == execution_backend::JIT + ? compiled_expr_ptr->get_jit_top_expression() + : compiled_expr_ptr->get_top_expression(); std::unique_ptr result = backend == execution_backend::JIT ? cudf::compute_column_jit(*tview_ptr, expression) : cudf::compute_column(*tview_ptr, expression); @@ -607,6 +606,34 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumn JNI_CATCH(env, 0); } +JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeTableJitNative( + JNIEnv* env, jclass, jlongArray j_asts, jlong j_table) +{ + 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); + 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> 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(handle); + expressions.emplace_back(compiled_expr_ptr->get_jit_top_expression()); + } + ast_handles.cancel(); + + auto const* tview_ptr = reinterpret_cast(j_table); + return cudf::jni::convert_table_for_return(env, + cudf::compute_table_jit(*tview_ptr, expressions)); + } + JNI_CATCH(env, nullptr); +} + JNIEXPORT void JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_destroy(JNIEnv* env, jclass, jlong jni_handle) diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index f51457bd0606..493c9055ae81 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -6,9 +6,15 @@ #pragma once #include +#include +#include +#include #include +#include +#include #include +#include #include #include #include @@ -17,57 +23,93 @@ namespace cudf { namespace jni { namespace ast { +struct expression_pair { + std::reference_wrapper regular; + std::reference_wrapper jit; + + expression_pair(cudf::ast::expression const& regular, cudf::ast::expression const& jit) + : regular{regular}, jit{jit} + { + } +}; + /** A class to capture all resources associated with a compiled AST expression. */ class compiled_expr { + // Keep literal owners before both trees so their non-owning nodes are destroyed first. /** GPU scalar instances that correspond to literal nodes */ std::vector> scalars; - /** All expression nodes within the expression tree */ + /** One-row columns backing literals in the JIT expression tree */ + std::vector> scalar_columns; + + /** All expression nodes within the regular expression tree */ cudf::ast::tree expressions; + /** All expression nodes within the JIT expression tree */ + cudf::ast::tree jit_expressions; + public: template - cudf::ast::literal const& add_literal(ScalarType& scalar, - std::unique_ptr scalar_ptr) + expression_pair add_literal(ScalarType& scalar, std::unique_ptr scalar_ptr) { + auto scalar_column = cudf::make_column_from_scalar(scalar, 1); scalars.push_back(std::move(scalar_ptr)); - return expressions.emplace(scalar); + scalar_columns.push_back(std::move(scalar_column)); + return {expressions.emplace(scalar), + jit_expressions.emplace( + cudf::scalar_column_view{scalar_columns.back()->view()})}; } - cudf::ast::column_reference const& add_column_ref(cudf::size_type column_index, - cudf::ast::table_reference table_ref) + expression_pair add_column_ref(cudf::size_type column_index, cudf::ast::table_reference table_ref) { - return expressions.emplace(column_index, table_ref); + return {expressions.emplace(column_index, table_ref), + jit_expressions.emplace(column_index, table_ref)}; } - cudf::ast::column_name_reference const& add_column_name_ref(std::string column_name) + expression_pair add_column_name_ref(std::string column_name) { - return expressions.emplace(std::move(column_name)); + return {expressions.emplace(column_name), + jit_expressions.emplace(std::move(column_name))}; } - cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, - cudf::ast::expression const& child) + expression_pair add_operation(cudf::ast::ast_operator op, expression_pair const& child) { - return expressions.emplace(op, child); + return {expressions.emplace(op, child.regular.get()), + jit_expressions.emplace(op, child.jit.get())}; } - cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, - cudf::ast::expression const& left, - cudf::ast::expression const& right) + expression_pair add_operation(cudf::ast::ast_operator op, + expression_pair const& left, + expression_pair const& right) { - return expressions.emplace(op, left, right); + return {expressions.emplace(op, left.regular.get(), right.regular.get()), + jit_expressions.emplace(op, left.jit.get(), right.jit.get())}; } - template - cudf::ast::expression const& add_jit_expression(F&& factory) + expression_pair add_jit_operation(cudf::ast::jit::op op, + std::vector const& args, + cudf::error_policy error_policy, + std::optional target_scale) { - return factory(expressions); + std::vector> regular_args; + std::vector> jit_args; + regular_args.reserve(args.size()); + jit_args.reserve(args.size()); + for (auto const& arg : args) { + regular_args.emplace_back(arg.regular); + jit_args.emplace_back(arg.jit); + } + return {cudf::ast::jit::operation(expressions, op, regular_args, error_policy, target_scale), + cudf::ast::jit::operation(jit_expressions, op, jit_args, error_policy, target_scale)}; } [[nodiscard]] bool has_literals() const { return !scalars.empty(); } /** Return the expression node at the top of the tree */ cudf::ast::expression const& get_top_expression() const { return expressions.back(); } + + /** Return the expression node at the top of the JIT tree */ + cudf::ast::expression const& get_jit_top_expression() const { return jit_expressions.back(); } }; } // namespace ast diff --git a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java index 575e899bf32d..71b9dc57e3cc 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -487,6 +487,193 @@ void testJitMismatchedOperandTypes() { } } + @Test + void testJitLiteralWorksWithBothExecutorsAndRepeatedInputs() { + AstExpression expr = new BinaryOperation(BinaryOperator.ADD, + new ColumnReference(0), Literal.ofInt(7)); + try (Table firstInput = new Table.TestBuilder().column(1, 2, 3).build(); + Table secondInput = new Table.TestBuilder().column(10, 20).build(); + CompiledExpression compiled = expr.compile(); + ColumnVector regularFirst = compiled.computeColumn(firstInput); + ColumnVector jitFirst = compiled.computeColumnJit(firstInput); + ColumnVector jitSecond = compiled.computeColumnJit(secondInput); + ColumnVector regularSecond = compiled.computeColumn(secondInput); + ColumnVector expectedFirst = ColumnVector.fromInts(8, 9, 10); + ColumnVector expectedSecond = ColumnVector.fromInts(17, 27)) { + assertColumnsAreEqual(expectedFirst, regularFirst); + assertColumnsAreEqual(expectedFirst, jitFirst); + assertColumnsAreEqual(expectedSecond, jitSecond); + assertColumnsAreEqual(expectedSecond, regularSecond); + } + } + + @Test + void testJitStringLiterals() { + AstExpression lessThan = new BinaryOperation(BinaryOperator.LESS, + new ColumnReference(0), Literal.ofString("ccc")); + AstExpression nullEqual = new BinaryOperation(BinaryOperator.NULL_EQUAL, + new ColumnReference(0), Literal.ofString(null)); + try (Table input = new Table.TestBuilder().column("a", null, "ccc", "dddd").build(); + CompiledExpression lessThanCompiled = lessThan.compile(); + CompiledExpression nullEqualCompiled = nullEqual.compile(); + ColumnVector actualLessThan = lessThanCompiled.computeColumnJit(input); + ColumnVector actualNullEqual = nullEqualCompiled.computeColumnJit(input); + ColumnVector expectedLessThan = + ColumnVector.fromBoxedBooleans(true, null, false, false); + ColumnVector expectedNullEqual = + ColumnVector.fromBoxedBooleans(false, true, false, false)) { + assertColumnsAreEqual(expectedLessThan, actualLessThan); + assertColumnsAreEqual(expectedNullEqual, actualNullEqual); + } + } + + @Test + void testJitMultipleOutputTransform() { + AstExpression firstSum = new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1)); + AstExpression multiply = new JitOperation(JitOperator.MUL, + firstSum, Literal.ofInt(2)); + AstExpression secondSum = new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1)); + AstExpression subtract = new JitOperation(JitOperator.SUB, + secondSum, new ColumnReference(2)); + + Table actual; + try (Table input = new Table.TestBuilder() + .column(1, 2, 3, 4) + .column(10, 20, 30, 40) + .column(2, 3, 4, 5) + .build(); + CompiledExpression multiplyCompiled = multiply.compile(); + CompiledExpression subtractCompiled = subtract.compile(); + CompiledExpression sumCompiled = secondSum.compile()) { + actual = CompiledExpression.computeTableJit( + input, multiplyCompiled, subtractCompiled, sumCompiled); + } + + try (Table result = actual; + ColumnVector expectedMultiply = ColumnVector.fromInts(22, 44, 66, 88); + ColumnVector expectedSubtract = ColumnVector.fromInts(9, 19, 29, 39); + ColumnVector expectedSum = ColumnVector.fromInts(11, 22, 33, 44)) { + Assertions.assertEquals(3, result.getNumberOfColumns()); + assertColumnsAreEqual(expectedMultiply, result.getColumn(0)); + assertColumnsAreEqual(expectedSubtract, result.getColumn(1)); + assertColumnsAreEqual(expectedSum, result.getColumn(2)); + } + } + + @Test + void testJitMultipleOutputPerOutputNullability() { + AstExpression isNull = new UnaryOperation( + UnaryOperator.IS_NULL, new ColumnReference(0)); + AstExpression sum = new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1)); + try (Table input = new Table.TestBuilder() + .column(1, null, 3, null) + .column(10, 20, 30, 40) + .build(); + CompiledExpression isNullCompiled = isNull.compile(); + CompiledExpression sumCompiled = sum.compile(); + Table actual = CompiledExpression.computeTableJit( + input, isNullCompiled, sumCompiled); + ColumnVector expectedIsNull = + ColumnVector.fromBoxedBooleans(false, true, false, true); + ColumnVector expectedSum = ColumnVector.fromBoxedInts(11, null, 33, null)) { + assertColumnsAreEqual(expectedIsNull, actual.getColumn(0)); + assertColumnsAreEqual(expectedSum, actual.getColumn(1)); + } + } + + @Test + void testJitMultipleOutputIndependentNullMasks() { + AstExpression first = new UnaryOperation( + UnaryOperator.IDENTITY, new ColumnReference(0)); + AstExpression second = new UnaryOperation( + UnaryOperator.IDENTITY, new ColumnReference(1)); + try (Table input = new Table.TestBuilder() + .column(1, null, 3, null) + .column(10, 20, null, null) + .build(); + CompiledExpression firstCompiled = first.compile(); + CompiledExpression secondCompiled = second.compile(); + Table actual = CompiledExpression.computeTableJit( + input, firstCompiled, secondCompiled); + ColumnVector expectedFirst = ColumnVector.fromBoxedInts(1, null, 3, null); + ColumnVector expectedSecond = ColumnVector.fromBoxedInts(10, 20, null, null)) { + assertColumnsAreEqual(expectedFirst, actual.getColumn(0)); + assertColumnsAreEqual(expectedSecond, actual.getColumn(1)); + } + } + + @Test + void testJitMultipleOutputEmptyInput() { + AstExpression identity = new UnaryOperation( + UnaryOperator.IDENTITY, new ColumnReference(0)); + AstExpression sum = new JitOperation(JitOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (ColumnVector empty = ColumnVector.fromInts(); + Table input = new Table(empty); + CompiledExpression identityCompiled = identity.compile(); + CompiledExpression sumCompiled = sum.compile(); + Table actual = CompiledExpression.computeTableJit( + input, identityCompiled, sumCompiled); + ColumnVector expected = ColumnVector.fromInts()) { + Assertions.assertEquals(2, actual.getNumberOfColumns()); + assertColumnsAreEqual(expected, actual.getColumn(0)); + assertColumnsAreEqual(expected, actual.getColumn(1)); + } + } + + @Test + void testJitMultipleOutputValidation() { + AstExpression expr = new JitOperation(JitOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (Table input = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expr.compile()) { + Assertions.assertThrows(NullPointerException.class, + () -> CompiledExpression.computeTableJit(null, compiled)); + Assertions.assertThrows(NullPointerException.class, + () -> CompiledExpression.computeTableJit(input, (CompiledExpression[]) null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> CompiledExpression.computeTableJit(input)); + Assertions.assertThrows(NullPointerException.class, + () -> CompiledExpression.computeTableJit(input, compiled, null)); + } + + CompiledExpression closedExpression = expr.compile(); + closedExpression.close(); + try (Table input = new Table.TestBuilder().column(1, 2, 3).build()) { + Assertions.assertThrows(IllegalStateException.class, + () -> CompiledExpression.computeTableJit(input, closedExpression)); + } + + try (Table closedTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expr.compile()) { + closedTable.close(); + Assertions.assertThrows(IllegalStateException.class, + () -> CompiledExpression.computeTableJit(closedTable, compiled)); + } + } + + @Test + void testJitMultipleOutputFailureDoesNotConsumeInputs() { + AstExpression valid = new JitOperation(JitOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + AstExpression invalid = new JitOperation(JitOperator.ADD, + new ColumnReference(1), Literal.ofInt(1)); + try (Table input = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression validCompiled = valid.compile(); + CompiledExpression invalidCompiled = invalid.compile()) { + Assertions.assertThrows(CudfException.class, + () -> CompiledExpression.computeTableJit( + input, validCompiled, invalidCompiled).close()); + try (ColumnVector actual = validCompiled.computeColumnJit(input); + ColumnVector expected = ColumnVector.fromInts(2, 3, 4)) { + assertColumnsAreEqual(expected, actual); + } + } + } + @Test void testJitEmptyInputTransform() { JitOperation expr = new JitOperation(JitOperator.ADD, From e37a82971e67f96aad2b254f1ff2f379903e3d7c Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 26 Aug 2026 18:05:23 +0800 Subject: [PATCH 2/4] refator to new api Signed-off-by: Haoyang Li --- .../ai/rapids/cudf/ast/AstExpression.java | 21 ++- .../rapids/cudf/ast/CompiledExpression.java | 67 ++++--- .../java/ai/rapids/cudf/ast/JitOperation.java | 4 +- .../main/java/ai/rapids/cudf/ast/Literal.java | 7 +- .../main/native/src/CompiledExpression.cpp | 171 +++++++++--------- .../src/main/native/src/jni_compiled_expr.hpp | 111 ++++++------ .../cudf/ast/CompiledExpressionTest.java | 165 ++++++++++------- 7 files changed, 304 insertions(+), 242 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java b/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java index 503ba613d0ff..c133c140a456 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java +++ b/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 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java index 031699aa229e..616b93656da7 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java @@ -16,6 +16,11 @@ /** 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(); } @@ -54,51 +59,39 @@ 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}. - * @param table input table for this expression - * @return new column computed from this expression applied to the input table - */ - public ColumnVector computeColumn(Table table) { - 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 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 JIT compilation or evaluation fails + * {@link TableReference#RIGHT}, or if compilation or evaluation fails */ - public ColumnVector computeColumnJit(Table table) { + public ColumnVector computeColumn(Table table) { long result; try { - result = computeColumnJit(cleaner.nativeHandle, table.getNativeView()); + result = computeColumn(cleaner.nativeHandle, table.getNativeView()); } finally { reachabilityFence(this); reachabilityFence(table); @@ -111,10 +104,11 @@ public ColumnVector computeColumnJit(Table table) { * transform. Output column {@code i} contains the result of {@code expressions[i]}. * * @param table input table for the expressions - * @param expressions non-empty expressions to evaluate in output order + * @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 + * @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 */ @@ -135,6 +129,9 @@ public static Table computeTableJit(Table table, CompiledExpression... expressio 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"); @@ -170,14 +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); } diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java index 7418413702af..95baf1dc83b7 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java +++ b/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. */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/Literal.java b/java/src/main/java/ai/rapids/cudf/ast/Literal.java index 715fc1cce42e..0aa227507511 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/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 diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index bd8a41863a79..a5f7f86d45e6 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -293,10 +293,10 @@ cudf::ast::table_reference jni_to_table_reference(jbyte jni_value) struct make_literal { /** Construct an AST literal from a numeric value */ template ()>* = nullptr> - cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_numeric_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -312,10 +312,10 @@ struct make_literal { /** Construct an AST literal from a timestamp value */ template ()>* = nullptr> - cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_timestamp_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -331,10 +331,10 @@ struct make_literal { /** Construct an AST literal from a duration value */ template ()>* = nullptr> - cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_duration_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -350,10 +350,10 @@ struct make_literal { /** Construct an AST literal from a string value */ template >* = nullptr> - cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = [&]() { if (is_valid) { @@ -370,10 +370,10 @@ struct make_literal { /** Construct an AST literal from a fixed-point value */ template ()>* = nullptr> - cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { using rep_type = typename T::rep; auto const val = is_valid ? jni_ast.read() : rep_type{}; @@ -390,26 +390,26 @@ struct make_literal { std::enable_if_t() && !cudf::is_timestamp() && !cudf::is_duration() && !cudf::is_fixed_point() && !std::is_same_v>* = nullptr> - cudf::jni::ast::expression_pair operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) const + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const { throw std::logic_error("Unsupported AST literal type"); } }; /** Decode a serialized AST literal */ -cudf::jni::ast::expression_pair compile_literal(bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::ast::literal const& compile_literal(bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { auto const dtype = jni_ast.read_cudf_type(); return cudf::type_dispatcher(dtype, make_literal{}, dtype, is_valid, compiled_expr, jni_ast); } /** Decode a serialized AST column reference */ -cudf::jni::ast::expression_pair compile_column_reference( +cudf::ast::column_reference const& compile_column_reference( cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { auto const table_ref = jni_to_table_reference(jni_ast.read_byte()); @@ -418,7 +418,7 @@ cudf::jni::ast::expression_pair compile_column_reference( } /** Decode a serialized AST column name reference */ -cudf::jni::ast::expression_pair compile_column_name_reference( +cudf::ast::column_name_reference const& compile_column_name_reference( cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { std::string column_name = jni_ast.read(); @@ -426,32 +426,36 @@ cudf::jni::ast::expression_pair compile_column_name_reference( } // forward declaration -cudf::jni::ast::expression_pair compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast); +cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast); /** Decode a serialized AST unary expression */ -cudf::jni::ast::expression_pair compile_unary_expression( - cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) +cudf::ast::operation const& compile_unary_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { - auto const ast_op = jni_to_unary_operator(jni_ast.read_byte()); - auto const child_expression = compile_expression(compiled_expr, jni_ast); + auto const ast_op = jni_to_unary_operator(jni_ast.read_byte()); + cudf::ast::expression const& child_expression = compile_expression(compiled_expr, jni_ast); return compiled_expr.add_operation(ast_op, child_expression); } /** Decode a serialized AST binary expression */ -cudf::jni::ast::expression_pair compile_binary_expression( - cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) +cudf::ast::operation const& compile_binary_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { - auto const ast_op = jni_to_binary_operator(jni_ast.read_byte()); - auto const left_child = compile_expression(compiled_expr, jni_ast); - auto const right_child = compile_expression(compiled_expr, jni_ast); + auto const ast_op = jni_to_binary_operator(jni_ast.read_byte()); + cudf::ast::expression const& left_child = compile_expression(compiled_expr, jni_ast); + cudf::ast::expression const& right_child = compile_expression(compiled_expr, jni_ast); return compiled_expr.add_operation(ast_op, left_child, right_child); } /** Decode a serialized JIT AST expression */ -cudf::jni::ast::expression_pair compile_jit_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +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(); @@ -493,18 +497,21 @@ cudf::jni::ast::expression_pair compile_jit_expression(cudf::jni::ast::compiled_ op_info.arity)); } - std::vector args; + std::vector> args; args.reserve(arity); for (int32_t index = 0; index < arity; ++index) { args.emplace_back(compile_expression(compiled_expr, jni_ast)); } - return compiled_expr.add_jit_operation(op_info.op, args, error_policy, target_scale); + return compiled_expr.add_jit_expression( + [&](cudf::ast::tree& tree) -> cudf::ast::expression const& { + return cudf::ast::jit::operation(tree, op_info.op, args, error_policy, target_scale); + }); } /** Decode a serialized AST expression by reading the expression type and dispatching */ -cudf::jni::ast::expression_pair compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { auto const expression_type = static_cast(jni_ast.read_byte()); switch (expression_type) { @@ -527,41 +534,38 @@ cudf::jni::ast::expression_pair compile_expression(cudf::jni::ast::compiled_expr } /** Decode a serialized AST into a native libcudf AST and associated resources */ -std::unique_ptr compile_serialized_ast(jni_serialized_ast& jni_ast) +std::unique_ptr compile_serialized_ast( + jni_serialized_ast& jni_ast, cudf::jni::ast::compilation_mode mode) { - auto jni_expr_ptr = std::make_unique(); + auto jni_expr_ptr = std::make_unique(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(j_ast); auto tview_ptr = reinterpret_cast(j_table); - auto const& expression = backend == execution_backend::JIT - ? compiled_expr_ptr->get_jit_top_expression() - : compiled_expr_ptr->get_top_expression(); - std::unique_ptr result = backend == execution_backend::JIT - ? cudf::compute_column_jit(*tview_ptr, expression) - : cudf::compute_column(*tview_ptr, expression); + std::unique_ptr 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(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 @@ -569,13 +573,31 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_compile(JNIEn 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(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, @@ -586,22 +608,7 @@ 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); - } - JNI_CATCH(env, 0); -} - -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumnJit(JNIEnv* env, - jclass, - jlong j_ast, - 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_TRY - { - cudf::jni::auto_set_device(env); - return execute_compiled_expression(j_ast, j_table, execution_backend::JIT); + return execute_compiled_expression(j_ast, j_table); } JNI_CATCH(env, 0); } diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index 493c9055ae81..8e85f2a5bd7c 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -11,10 +11,8 @@ #include #include -#include -#include #include -#include +#include #include #include #include @@ -23,93 +21,98 @@ namespace cudf { namespace jni { namespace ast { -struct expression_pair { - std::reference_wrapper regular; - std::reference_wrapper jit; - - expression_pair(cudf::ast::expression const& regular, cudf::ast::expression const& jit) - : regular{regular}, jit{jit} - { - } -}; +enum class compilation_mode { DEFAULT, JIT }; /** A class to capture all resources associated with a compiled AST expression. */ class compiled_expr { - // Keep literal owners before both trees so their non-owning nodes are destroyed first. + compilation_mode const mode; + + // Keep literal owners before the tree so its non-owning nodes are destroyed first. /** GPU scalar instances that correspond to literal nodes */ std::vector> scalars; - /** One-row columns backing literals in the JIT expression tree */ + /** One-row columns backing literals in a JIT expression tree */ std::vector> scalar_columns; - /** All expression nodes within the regular expression tree */ + /** All expression nodes within the expression tree */ cudf::ast::tree expressions; - /** All expression nodes within the JIT expression tree */ - cudf::ast::tree jit_expressions; - public: + explicit compiled_expr(compilation_mode mode) : mode{mode} {} + template - expression_pair add_literal(ScalarType& scalar, std::unique_ptr scalar_ptr) + cudf::ast::literal const& add_literal(ScalarType& scalar, + std::unique_ptr scalar_ptr) { - auto scalar_column = cudf::make_column_from_scalar(scalar, 1); + if (is_jit()) { + auto scalar_column = cudf::make_column_from_scalar(scalar, 1); + scalars.push_back(std::move(scalar_ptr)); + scalar_columns.push_back(std::move(scalar_column)); + return expressions.emplace( + cudf::scalar_column_view{scalar_columns.back()->view()}); + } + scalars.push_back(std::move(scalar_ptr)); - scalar_columns.push_back(std::move(scalar_column)); - return {expressions.emplace(scalar), - jit_expressions.emplace( - cudf::scalar_column_view{scalar_columns.back()->view()})}; + return expressions.emplace(scalar); } - expression_pair add_column_ref(cudf::size_type column_index, cudf::ast::table_reference table_ref) + cudf::ast::column_reference const& add_column_ref(cudf::size_type column_index, + cudf::ast::table_reference table_ref) { - return {expressions.emplace(column_index, table_ref), - jit_expressions.emplace(column_index, table_ref)}; + return expressions.emplace(column_index, table_ref); } - expression_pair add_column_name_ref(std::string column_name) + cudf::ast::column_name_reference const& add_column_name_ref(std::string column_name) { - return {expressions.emplace(column_name), - jit_expressions.emplace(std::move(column_name))}; + return expressions.emplace(std::move(column_name)); } - expression_pair add_operation(cudf::ast::ast_operator op, expression_pair const& child) + cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, + cudf::ast::expression const& child) { - return {expressions.emplace(op, child.regular.get()), - jit_expressions.emplace(op, child.jit.get())}; + return expressions.emplace(op, child); } - expression_pair add_operation(cudf::ast::ast_operator op, - expression_pair const& left, - expression_pair const& right) + cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, + cudf::ast::expression const& left, + cudf::ast::expression const& right) { - return {expressions.emplace(op, left.regular.get(), right.regular.get()), - jit_expressions.emplace(op, left.jit.get(), right.jit.get())}; + return expressions.emplace(op, left, right); } - expression_pair add_jit_operation(cudf::ast::jit::op op, - std::vector const& args, - cudf::error_policy error_policy, - std::optional target_scale) + template + cudf::ast::expression const& add_jit_expression(F&& factory) { - std::vector> regular_args; - std::vector> jit_args; - regular_args.reserve(args.size()); - jit_args.reserve(args.size()); - for (auto const& arg : args) { - regular_args.emplace_back(arg.regular); - jit_args.emplace_back(arg.jit); + if (!is_jit()) { + throw std::invalid_argument("JIT operations require an expression compiled for JIT"); } - return {cudf::ast::jit::operation(expressions, op, regular_args, error_policy, target_scale), - cudf::ast::jit::operation(jit_expressions, op, jit_args, error_policy, target_scale)}; + return factory(expressions); } [[nodiscard]] bool has_literals() const { return !scalars.empty(); } - /** Return the expression node at the top of the tree */ - cudf::ast::expression const& get_top_expression() const { return expressions.back(); } + [[nodiscard]] bool is_jit() const { return mode == compilation_mode::JIT; } + + void release_jit_staging_scalars() + { + if (is_jit()) { scalars.clear(); } + } + + /** Return the expression node at the top of a default-compatible tree */ + cudf::ast::expression const& get_top_expression() const + { + if (is_jit()) { + throw std::logic_error("JIT-compiled expressions cannot be used by a default AST consumer"); + } + return expressions.back(); + } /** Return the expression node at the top of the JIT tree */ - cudf::ast::expression const& get_jit_top_expression() const { return jit_expressions.back(); } + cudf::ast::expression const& get_jit_top_expression() const + { + if (!is_jit()) { throw std::logic_error("Expression was not compiled for JIT"); } + return expressions.back(); + } }; } // namespace ast diff --git a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java index 71b9dc57e3cc..1f14786e5068 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -57,9 +57,10 @@ public void testInvalidColumnReferenceTransform() { // Verify that computeColumn throws when passed an expression operating on TableReference.RIGHT. ColumnReference expr = new ColumnReference(1, TableReference.RIGHT); try (Table t = new Table.TestBuilder().column(5, 4, 3, 2, 1).column(6, 7, 8, null, 10).build(); - CompiledExpression compiledExpr = expr.compile()) { + CompiledExpression compiledExpr = expr.compile(); + CompiledExpression compiledJitExpr = expr.compileJit()) { Assertions.assertThrows(CudfException.class, () -> compiledExpr.computeColumn(t).close()); - Assertions.assertThrows(CudfException.class, () -> compiledExpr.computeColumnJit(t).close()); + Assertions.assertThrows(CudfException.class, () -> compiledJitExpr.computeColumn(t).close()); } } @@ -384,8 +385,8 @@ public void testDecimal128LiteralComparisonLegacyTransform() { public void testJitDecimalLiteralTransform(DType type, BigInteger value) { Literal expr = Literal.ofDecimal(type, value); try (Table t = new Table.TestBuilder().column(1, 2, 3).build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); Scalar expectedScalar = value == null ? Scalar.fromNull(type) : Scalar.fromDecimal(value, type); ColumnVector expected = ColumnVector.fromScalar(expectedScalar, 3)) { @@ -470,40 +471,55 @@ void testJitOperationValidation() { } private static void assertJitCompileThrows(JitOperation expr) { + Assertions.assertThrows(CudfException.class, () -> { + try (CompiledExpression ignored = expr.compileJit()) { + } + }); + } + + @Test + void testJitOperationRequiresJitCompilation() { + JitOperation expr = new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1)); Assertions.assertThrows(CudfException.class, () -> { try (CompiledExpression ignored = expr.compile()) { } }); } + @Test + void testJitCompiledExpressionDoesNotExposeRegularHandle() { + AstExpression expr = new BinaryOperation(BinaryOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (CompiledExpression compiled = expr.compileJit()) { + Assertions.assertThrows(IllegalStateException.class, compiled::getNativeHandle); + } + } + @Test void testJitMismatchedOperandTypes() { JitOperation expr = new JitOperation(JitOperator.ADD, new ColumnReference(0), new ColumnReference(1)); try (Table t = new Table.TestBuilder().column(1).column(2L).build(); - CompiledExpression compiledExpr = expr.compile()) { + CompiledExpression compiledExpr = expr.compileJit()) { Assertions.assertThrows(CudfException.class, - () -> compiledExpr.computeColumnJit(t).close()); + () -> compiledExpr.computeColumn(t).close()); } } @Test - void testJitLiteralWorksWithBothExecutorsAndRepeatedInputs() { + void testJitLiteralWorksWithRepeatedInputs() { AstExpression expr = new BinaryOperation(BinaryOperator.ADD, new ColumnReference(0), Literal.ofInt(7)); try (Table firstInput = new Table.TestBuilder().column(1, 2, 3).build(); Table secondInput = new Table.TestBuilder().column(10, 20).build(); - CompiledExpression compiled = expr.compile(); - ColumnVector regularFirst = compiled.computeColumn(firstInput); - ColumnVector jitFirst = compiled.computeColumnJit(firstInput); - ColumnVector jitSecond = compiled.computeColumnJit(secondInput); - ColumnVector regularSecond = compiled.computeColumn(secondInput); + CompiledExpression compiled = expr.compileJit(); + ColumnVector jitFirst = compiled.computeColumn(firstInput); + ColumnVector jitSecond = compiled.computeColumn(secondInput); ColumnVector expectedFirst = ColumnVector.fromInts(8, 9, 10); ColumnVector expectedSecond = ColumnVector.fromInts(17, 27)) { - assertColumnsAreEqual(expectedFirst, regularFirst); assertColumnsAreEqual(expectedFirst, jitFirst); assertColumnsAreEqual(expectedSecond, jitSecond); - assertColumnsAreEqual(expectedSecond, regularSecond); } } @@ -514,10 +530,10 @@ void testJitStringLiterals() { AstExpression nullEqual = new BinaryOperation(BinaryOperator.NULL_EQUAL, new ColumnReference(0), Literal.ofString(null)); try (Table input = new Table.TestBuilder().column("a", null, "ccc", "dddd").build(); - CompiledExpression lessThanCompiled = lessThan.compile(); - CompiledExpression nullEqualCompiled = nullEqual.compile(); - ColumnVector actualLessThan = lessThanCompiled.computeColumnJit(input); - ColumnVector actualNullEqual = nullEqualCompiled.computeColumnJit(input); + CompiledExpression lessThanCompiled = lessThan.compileJit(); + CompiledExpression nullEqualCompiled = nullEqual.compileJit(); + ColumnVector actualLessThan = lessThanCompiled.computeColumn(input); + ColumnVector actualNullEqual = nullEqualCompiled.computeColumn(input); ColumnVector expectedLessThan = ColumnVector.fromBoxedBooleans(true, null, false, false); ColumnVector expectedNullEqual = @@ -544,9 +560,9 @@ void testJitMultipleOutputTransform() { .column(10, 20, 30, 40) .column(2, 3, 4, 5) .build(); - CompiledExpression multiplyCompiled = multiply.compile(); - CompiledExpression subtractCompiled = subtract.compile(); - CompiledExpression sumCompiled = secondSum.compile()) { + CompiledExpression multiplyCompiled = multiply.compileJit(); + CompiledExpression subtractCompiled = subtract.compileJit(); + CompiledExpression sumCompiled = secondSum.compileJit()) { actual = CompiledExpression.computeTableJit( input, multiplyCompiled, subtractCompiled, sumCompiled); } @@ -572,8 +588,8 @@ void testJitMultipleOutputPerOutputNullability() { .column(1, null, 3, null) .column(10, 20, 30, 40) .build(); - CompiledExpression isNullCompiled = isNull.compile(); - CompiledExpression sumCompiled = sum.compile(); + CompiledExpression isNullCompiled = isNull.compileJit(); + CompiledExpression sumCompiled = sum.compileJit(); Table actual = CompiledExpression.computeTableJit( input, isNullCompiled, sumCompiled); ColumnVector expectedIsNull = @@ -594,8 +610,8 @@ void testJitMultipleOutputIndependentNullMasks() { .column(1, null, 3, null) .column(10, 20, null, null) .build(); - CompiledExpression firstCompiled = first.compile(); - CompiledExpression secondCompiled = second.compile(); + CompiledExpression firstCompiled = first.compileJit(); + CompiledExpression secondCompiled = second.compileJit(); Table actual = CompiledExpression.computeTableJit( input, firstCompiled, secondCompiled); ColumnVector expectedFirst = ColumnVector.fromBoxedInts(1, null, 3, null); @@ -613,8 +629,8 @@ void testJitMultipleOutputEmptyInput() { new ColumnReference(0), Literal.ofInt(1)); try (ColumnVector empty = ColumnVector.fromInts(); Table input = new Table(empty); - CompiledExpression identityCompiled = identity.compile(); - CompiledExpression sumCompiled = sum.compile(); + CompiledExpression identityCompiled = identity.compileJit(); + CompiledExpression sumCompiled = sum.compileJit(); Table actual = CompiledExpression.computeTableJit( input, identityCompiled, sumCompiled); ColumnVector expected = ColumnVector.fromInts()) { @@ -629,7 +645,7 @@ void testJitMultipleOutputValidation() { AstExpression expr = new JitOperation(JitOperator.ADD, new ColumnReference(0), Literal.ofInt(1)); try (Table input = new Table.TestBuilder().column(1, 2, 3).build(); - CompiledExpression compiled = expr.compile()) { + CompiledExpression compiled = expr.compileJit()) { Assertions.assertThrows(NullPointerException.class, () -> CompiledExpression.computeTableJit(null, compiled)); Assertions.assertThrows(NullPointerException.class, @@ -640,7 +656,7 @@ void testJitMultipleOutputValidation() { () -> CompiledExpression.computeTableJit(input, compiled, null)); } - CompiledExpression closedExpression = expr.compile(); + CompiledExpression closedExpression = expr.compileJit(); closedExpression.close(); try (Table input = new Table.TestBuilder().column(1, 2, 3).build()) { Assertions.assertThrows(IllegalStateException.class, @@ -648,11 +664,22 @@ void testJitMultipleOutputValidation() { } try (Table closedTable = new Table.TestBuilder().column(1, 2, 3).build(); - CompiledExpression compiled = expr.compile()) { + CompiledExpression compiled = expr.compileJit()) { closedTable.close(); Assertions.assertThrows(IllegalStateException.class, () -> CompiledExpression.computeTableJit(closedTable, compiled)); } + + AstExpression defaultExpr = new BinaryOperation(BinaryOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (Table input = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression defaultCompiled = defaultExpr.compile(); + CompiledExpression jitCompiled = defaultExpr.compileJit()) { + Assertions.assertThrows(IllegalArgumentException.class, + () -> CompiledExpression.computeTableJit(input, defaultCompiled)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> CompiledExpression.computeTableJit(input, jitCompiled, defaultCompiled)); + } } @Test @@ -662,12 +689,12 @@ void testJitMultipleOutputFailureDoesNotConsumeInputs() { AstExpression invalid = new JitOperation(JitOperator.ADD, new ColumnReference(1), Literal.ofInt(1)); try (Table input = new Table.TestBuilder().column(1, 2, 3).build(); - CompiledExpression validCompiled = valid.compile(); - CompiledExpression invalidCompiled = invalid.compile()) { + CompiledExpression validCompiled = valid.compileJit(); + CompiledExpression invalidCompiled = invalid.compileJit()) { Assertions.assertThrows(CudfException.class, () -> CompiledExpression.computeTableJit( input, validCompiled, invalidCompiled).close()); - try (ColumnVector actual = validCompiled.computeColumnJit(input); + try (ColumnVector actual = validCompiled.computeColumn(input); ColumnVector expected = ColumnVector.fromInts(2, 3, 4)) { assertColumnsAreEqual(expected, actual); } @@ -679,8 +706,8 @@ void testJitEmptyInputTransform() { JitOperation expr = new JitOperation(JitOperator.ADD, new ColumnReference(0), Literal.ofInt(1)); try (Table t = new Table.TestBuilder().column(new Integer[0]).build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromInts()) { assertColumnsAreEqual(expected, actual); } @@ -699,8 +726,8 @@ void testJitNestedArithmeticTransform() { expr = new JitOperation(JitOperator.BITWISE_SHIFT_RIGHT, expr, Literal.ofInt(1)); try (Table t = new Table.TestBuilder().column(1, 2, 3, 4).build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromInts(6, 8, 2, 4)) { assertColumnsAreEqual(expected, actual); } @@ -710,8 +737,8 @@ void testJitNestedArithmeticTransform() { void testJitNegTransform() { JitOperation expr = new JitOperation(JitOperator.NEG, new ColumnReference(0)); try (Table t = new Table.TestBuilder().column(-5, 0, 7).build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromInts(5, 0, -7)) { assertColumnsAreEqual(expected, actual); } @@ -726,24 +753,24 @@ void testJitOverflowPolicies() { .build()) { JitOperation successExpr = new JitOperation(JitOperator.ADD_OVERFLOW, new ColumnReference(0), new ColumnReference(1)); - try (CompiledExpression compiledExpr = successExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = successExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromInts(11, 10)) { assertColumnsAreEqual(expected, actual); } JitOperation propagateExpr = new JitOperation(JitOperator.ADD_OVERFLOW, new ColumnReference(0), new ColumnReference(2)); - try (CompiledExpression compiledExpr = propagateExpr.compile()) { + try (CompiledExpression compiledExpr = propagateExpr.compileJit()) { Assertions.assertThrows(CudfException.class, - () -> compiledExpr.computeColumnJit(t).close()); + () -> compiledExpr.computeColumn(t).close()); } JitOperation nullifyExpr = new JitOperation(JitOperator.ADD_OVERFLOW, JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(2)); - try (CompiledExpression compiledExpr = nullifyExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = nullifyExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(11, null)) { assertColumnsAreEqual(expected, actual); } @@ -764,8 +791,8 @@ void testJitFusedNullifyingOverflowTransform() { expr, new ColumnReference(2)); expr = new JitOperation(JitOperator.DIV_OVERFLOW, JitErrorPolicy.NULLIFY, expr, new ColumnReference(3)); - try (CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(null, 65, null, null, null, 12)) { assertColumnsAreEqual(expected, actual); } @@ -780,24 +807,24 @@ void testJitUnaryAndSubtractOverflowTransform() { .build()) { JitOperation subExpr = new JitOperation(JitOperator.SUB_OVERFLOW, JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); - try (CompiledExpression compiledExpr = subExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = subExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(7, null, 1)) { assertColumnsAreEqual(expected, actual); } JitOperation negExpr = new JitOperation(JitOperator.NEG_OVERFLOW, JitErrorPolicy.NULLIFY, new ColumnReference(0)); - try (CompiledExpression compiledExpr = negExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = negExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(-10, null, -1)) { assertColumnsAreEqual(expected, actual); } JitOperation absExpr = new JitOperation(JitOperator.ABS_OVERFLOW, JitErrorPolicy.NULLIFY, new ColumnReference(0)); - try (CompiledExpression compiledExpr = absExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = absExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(10, null, 1)) { assertColumnsAreEqual(expected, actual); } @@ -812,16 +839,16 @@ void testJitTryDivModTransform() { .build()) { JitOperation divExpr = new JitOperation(JitOperator.DIV_OVERFLOW, JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); - try (CompiledExpression compiledExpr = divExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = divExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(5, null, null, null, null)) { assertColumnsAreEqual(expected, actual); } JitOperation modExpr = new JitOperation(JitOperator.MOD_OVERFLOW, JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); - try (CompiledExpression compiledExpr = modExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = modExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(0, null, null, null, 0)) { assertColumnsAreEqual(expected, actual); } @@ -841,8 +868,8 @@ void testJitMixedConditionalTransform() { new ColumnReference(0), Literal.ofInt(99)); JitOperation expr = new JitOperation(JitOperator.IF_ELSE, coalesced, new ColumnReference(1), predicate); - try (CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.fromBoxedInts(10, 20, 3, 40)) { assertColumnsAreEqual(expected, actual); } @@ -886,8 +913,8 @@ void testJitNumericCastTransform( JitOperator op, Supplier expectedFactory) { JitOperation expr = new JitOperation(op, new ColumnReference(0)); try (Table t = new Table.TestBuilder().column(0, 1, 2, 3).build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = expectedFactory.get()) { assertColumnsAreEqual(expected, actual); } @@ -910,8 +937,8 @@ void testJitDecimalCastTransform( JitOperator op, Supplier expectedFactory) { JitOperation expr = new JitOperation(op, new ColumnReference(0)); try (Table t = new Table.TestBuilder().decimal64Column(0, 0L, 1L, -2L, 3L).build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = expectedFactory.get()) { assertColumnsAreEqual(expected, actual); } @@ -923,8 +950,8 @@ void testJitDecimalRescaleTransform() { try (Table t = new Table.TestBuilder() .decimal32Column(0, 123, 1234, 12345, 123456, 1234567) .build(); - CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + CompiledExpression compiledExpr = expr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.decimalFromInts( -2, 12300, 123400, 1234500, 12345600, 123456700)) { assertColumnsAreEqual(expected, actual); @@ -936,15 +963,15 @@ void testJitDecimalPrecisionPolicies() { try (Table t = new Table.TestBuilder().decimal32Column(0, 3, 200, 250, 20000).build()) { JitOperation propagateExpr = new JitOperation(JitOperator.CHECK_PRECISION, new ColumnReference(0), Literal.ofInt(3)); - try (CompiledExpression compiledExpr = propagateExpr.compile()) { + try (CompiledExpression compiledExpr = propagateExpr.compileJit()) { Assertions.assertThrows(CudfException.class, - () -> compiledExpr.computeColumnJit(t).close()); + () -> compiledExpr.computeColumn(t).close()); } JitOperation nullifyExpr = new JitOperation(JitOperator.CHECK_PRECISION, JitErrorPolicy.NULLIFY, new ColumnReference(0), Literal.ofInt(3)); - try (CompiledExpression compiledExpr = nullifyExpr.compile(); - ColumnVector actual = compiledExpr.computeColumnJit(t); + try (CompiledExpression compiledExpr = nullifyExpr.compileJit(); + ColumnVector actual = compiledExpr.computeColumn(t); ColumnVector expected = ColumnVector.decimalFromBoxedInts(0, 3, 200, 250, null)) { assertColumnsAreEqual(expected, actual); } From 6ac19c854da125fb3b79bb4bdd70e159b6b7a4c0 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 1 Sep 2026 17:05:38 +0800 Subject: [PATCH 3/4] Fix transform program scalar literal ownership Copy prepared scalar column views into program-owned columns so reusable AST programs remain valid after their source expressions are destroyed. Signed-off-by: Haoyang Li --- cpp/src/transform/transform.cu | 6 +++++- cpp/tests/ast/transform_tests.cpp | 23 +++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index 0cdfdf74bd02..10e8bc92eaa2 100644 --- a/cpp/src/transform/transform.cu +++ b/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(&input)) { + // The program must outlive non-owning scalar-column literals in the source AST. + impl_->ast_scalar_columns_.push_back( + std::make_unique(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); } diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index 1e1385d15875..a8dd06b6c409 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -110,6 +110,29 @@ TEST_F(TransformProgramTest, ReusesAstWithCompatibleTable) CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity); } +TEST_F(TransformProgramTest, OwnsScalarColumnViewLiterals) +{ + std::unique_ptr program; + { + auto construction_input = column_wrapper{3, 20, 1, 50}; + auto construction_table = cudf::table_view{{construction_input}}; + auto literal_column = column_wrapper{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 expressions[] = {expression}; + + program = std::make_unique(construction_table, expressions); + } + + auto input = column_wrapper{10, 20, 30}; + auto table = cudf::table_view{{input}}; + auto expected = column_wrapper{12, 22, 32}; + 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{3, 20, 1, 50}; From 983fc16bac85f996dd02dd639a7a09e65566bc86 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 1 Sep 2026 17:06:09 +0800 Subject: [PATCH 4/4] Add Java bindings for reusable AST JIT programs Expose schema-specialized transform_program construction and reuse from Java. Preserve literal ownership across program reuse, validate compilation modes and schemas, and cover multi-output execution. Signed-off-by: Haoyang Li --- .../java/ai/rapids/cudf/MemoryCleaner.java | 9 +- .../ai/rapids/cudf/ast/AstJitProgram.java | 152 ++++++++++++++++++ .../rapids/cudf/ast/CompiledExpression.java | 30 ++-- java/src/main/native/CMakeLists.txt | 1 + java/src/main/native/src/AstJitProgram.cpp | 80 +++++++++ .../src/main/native/src/jni_compiled_expr.hpp | 2 + .../ai/rapids/cudf/ast/AstJitProgramTest.java | 136 ++++++++++++++++ 7 files changed, 396 insertions(+), 14 deletions(-) create mode 100644 java/src/main/java/ai/rapids/cudf/ast/AstJitProgram.java create mode 100644 java/src/main/native/src/AstJitProgram.cpp create mode 100644 java/src/test/java/ai/rapids/cudf/ast/AstJitProgramTest.java diff --git a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java b/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java index 7d48dfa34b74..e8a5fd6cc087 100644 --- a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java +++ b/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)); + // JIT expressions can own one-row literal columns. + 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) { diff --git a/java/src/main/java/ai/rapids/cudf/ast/AstJitProgram.java b/java/src/main/java/ai/rapids/cudf/ast/AstJitProgram.java new file mode 100644 index 000000000000..f9012760fc3f --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/AstJitProgram.java @@ -0,0 +1,152 @@ +/* + * 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. + */ +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) { + long origAddress = nativeHandle; + boolean neededCleanup = nativeHandle != 0; + if (neededCleanup) { + try { + destroy(nativeHandle); + } finally { + nativeHandle = 0; + } + if (logErrorIfNotClean) { + log.error("AN AST JIT PROGRAM WAS LEAKED (ID: " + + id + " " + Long.toHexString(origAddress)); + } + } + return neededCleanup; + } + + @Override + public boolean isClean() { + return nativeHandle == 0; + } + } + + private final AstJitProgramCleaner cleaner; + private boolean isClosed = false; + + private AstJitProgram(long nativeHandle) { + cleaner = new AstJitProgramCleaner(nativeHandle); + MemoryCleaner.register(this, cleaner); + cleaner.addRef(); + } + + /** + * 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) { + Objects.requireNonNull(schemaTable, "schemaTable"); + Objects.requireNonNull(expressions, "expressions"); + if (expressions.length == 0) { + throw new IllegalArgumentException("At least one expression is required"); + } + + long tableHandle = schemaTable.getNativeView(); + if (tableHandle == 0) { + throw new IllegalStateException("Table is closed"); + } + + CompiledExpression[] expressionRefs = expressions.clone(); + long[] nativeHandles = CompiledExpression.getJitNativeHandles(expressionRefs); + long programHandle; + try { + programHandle = create(nativeHandles, tableHandle); + } finally { + CompiledExpression.reachabilityFence(schemaTable); + CompiledExpression.reachabilityFence(expressionRefs); + } + return new AstJitProgram(programHandle); + } + + /** + * 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. + * + * @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) { + Objects.requireNonNull(table, "table"); + long programHandle = cleaner.nativeHandle; + if (programHandle == 0) { + throw new IllegalStateException("AST JIT program is closed"); + } + long tableHandle = table.getNativeView(); + if (tableHandle == 0) { + throw new IllegalStateException("Table is closed"); + } + + long[] result; + try { + result = computeTableNative(programHandle, tableHandle); + } finally { + CompiledExpression.reachabilityFence(this); + CompiledExpression.reachabilityFence(table); + } + return new Table(result); + } + + @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); +} diff --git a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java index 616b93656da7..cc16d3af1394 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java @@ -125,10 +125,22 @@ public static Table computeTableJit(Table table, CompiledExpression... expressio } CompiledExpression[] expressionRefs = expressions.clone(); - long[] nativeHandles = new long[expressionRefs.length]; - for (int i = 0; i < expressionRefs.length; i++) { + long[] nativeHandles = getJitNativeHandles(expressionRefs); + long[] result; + try { + result = computeTableJitNative(nativeHandles, tableHandle); + } finally { + reachabilityFence(table); + reachabilityFence(expressionRefs); + } + return new Table(result); + } + + static long[] getJitNativeHandles(CompiledExpression[] expressions) { + long[] nativeHandles = new long[expressions.length]; + for (int i = 0; i < expressions.length; i++) { CompiledExpression expression = Objects.requireNonNull( - expressionRefs[i], "expression " + i + " is null"); + expressions[i], "expression " + i + " is null"); if (expression.mode != CompilationMode.JIT) { throw new IllegalArgumentException("Expression " + i + " was not compiled for JIT"); } @@ -137,18 +149,10 @@ public static Table computeTableJit(Table table, CompiledExpression... expressio throw new IllegalStateException("Expression " + i + " is closed"); } } - - long[] result; - try { - result = computeTableJitNative(nativeHandles, tableHandle); - } finally { - reachabilityFence(table); - reachabilityFence(expressionRefs); - } - return new Table(result); + return nativeHandles; } - private static void reachabilityFence(Object object) { + static void reachabilityFence(Object object) { if (object != null) { synchronized (object) { // The monitor operation is a Java 8 reachability fence. diff --git a/java/src/main/native/CMakeLists.txt b/java/src/main/native/CMakeLists.txt index 1784631ddcad..8b56e2aa8fd8 100644 --- a/java/src/main/native/CMakeLists.txt +++ b/java/src/main/native/CMakeLists.txt @@ -156,6 +156,7 @@ add_library( cudfjni src/Aggregation128UtilsJni.cpp src/AggregationJni.cpp + src/AstJitProgram.cpp src/ChunkedPackJni.cpp src/ChunkedReaderJni.cpp src/CudfJni.cpp diff --git a/java/src/main/native/src/AstJitProgram.cpp b/java/src/main/native/src/AstJitProgram.cpp new file mode 100644 index 000000000000..03fff5a7e9d1 --- /dev/null +++ b/java/src/main/native/src/AstJitProgram.cpp @@ -0,0 +1,80 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cudf_jni_apis.hpp" +#include "jni_compiled_expr.hpp" + +#include + +#include +#include +#include +#include + +extern "C" { + +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_AstJitProgram_create(JNIEnv* env, + jclass, + jlongArray j_asts, + jlong j_table) +{ + JNI_NULL_CHECK(env, j_asts, "Compiled AST pointer array is null", 0); + JNI_NULL_CHECK(env, j_table, "Table view pointer is null", 0); + JNI_TRY + { + cudf::jni::auto_set_device(env); + 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> expressions; + expressions.reserve(ast_handles.size()); + auto has_literals = false; + 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(handle); + expressions.emplace_back(compiled_expr_ptr->get_jit_top_expression()); + has_literals |= compiled_expr_ptr->has_jit_literals(); + } + ast_handles.cancel(); + + auto const* table = reinterpret_cast(j_table); + auto const stream = cudf::get_default_stream(); + auto program = std::make_unique(*table, expressions, stream); + // Construction inputs may be released by a thread with a different default stream. + if (has_literals) { stream.synchronize(); } + return reinterpret_cast(program.release()); + } + JNI_CATCH(env, 0); +} + +JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_ast_AstJitProgram_computeTableNative( + JNIEnv* env, jclass, jlong j_program, jlong j_table) +{ + JNI_NULL_CHECK(env, j_program, "AST JIT program pointer is null", nullptr); + JNI_NULL_CHECK(env, j_table, "Table view pointer is null", nullptr); + JNI_TRY + { + cudf::jni::auto_set_device(env); + auto* program = reinterpret_cast(j_program); + auto const* table = reinterpret_cast(j_table); + return cudf::jni::convert_table_for_return(env, program->run(*table)); + } + JNI_CATCH(env, nullptr); +} + +JNIEXPORT void JNICALL Java_ai_rapids_cudf_ast_AstJitProgram_destroy(JNIEnv* env, + jclass, + jlong j_program) +{ + JNI_TRY + { + cudf::jni::auto_set_device(env); + delete reinterpret_cast(j_program); + } + JNI_CATCH(env, ); +} + +} // extern "C" diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index 8e85f2a5bd7c..4b7fc734fb1b 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -91,6 +91,8 @@ class compiled_expr { [[nodiscard]] bool has_literals() const { return !scalars.empty(); } + [[nodiscard]] bool has_jit_literals() const { return !scalar_columns.empty(); } + [[nodiscard]] bool is_jit() const { return mode == compilation_mode::JIT; } void release_jit_staging_scalars() diff --git a/java/src/test/java/ai/rapids/cudf/ast/AstJitProgramTest.java b/java/src/test/java/ai/rapids/cudf/ast/AstJitProgramTest.java new file mode 100644 index 000000000000..e5f469227171 --- /dev/null +++ b/java/src/test/java/ai/rapids/cudf/ast/AstJitProgramTest.java @@ -0,0 +1,136 @@ +/* + * 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.ColumnVector; +import ai.rapids.cudf.CudfException; +import ai.rapids.cudf.CudfTestBase; +import ai.rapids.cudf.Table; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual; + +public class AstJitProgramTest extends CudfTestBase { + @Test + void testReusesMultiOutputProgramAndOwnsLiterals() { + AstExpression shared = new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1)); + AstExpression multiply = new JitOperation(JitOperator.MUL, shared, Literal.ofInt(2)); + AstExpression sum = new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1)); + + AstJitProgram program; + try (Table schemaTable = new Table.TestBuilder() + .column(1, 2, 3) + .column(10, 20, 30) + .build(); + CompiledExpression multiplyCompiled = multiply.compileJit(); + CompiledExpression sumCompiled = sum.compileJit()) { + program = AstJitProgram.compile(schemaTable, multiplyCompiled, sumCompiled); + } + + try (AstJitProgram closeableProgram = program; + Table firstInput = new Table.TestBuilder() + .column(4, 5) + .column(40, 50) + .build(); + Table firstResult = closeableProgram.computeTable(firstInput); + ColumnVector firstMultiply = ColumnVector.fromInts(88, 110); + ColumnVector firstSum = ColumnVector.fromInts(44, 55); + Table secondInput = new Table.TestBuilder() + .column(6, 7, 8, 9) + .column(60, 70, 80, 90) + .build(); + Table secondResult = closeableProgram.computeTable(secondInput); + ColumnVector secondMultiply = ColumnVector.fromInts(132, 154, 176, 198); + ColumnVector secondSum = ColumnVector.fromInts(66, 77, 88, 99)) { + Assertions.assertEquals(2, firstResult.getNumberOfColumns()); + assertColumnsAreEqual(firstMultiply, firstResult.getColumn(0)); + assertColumnsAreEqual(firstSum, firstResult.getColumn(1)); + Assertions.assertEquals(2, secondResult.getNumberOfColumns()); + assertColumnsAreEqual(secondMultiply, secondResult.getColumn(0)); + assertColumnsAreEqual(secondSum, secondResult.getColumn(1)); + } + } + + @Test + void testRejectsIncompatibleReferencedColumnSchema() { + AstExpression expression = new JitOperation(JitOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (Table schemaTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expression.compileJit(); + AstJitProgram program = AstJitProgram.compile(schemaTable, compiled); + Table wrongType = new Table.TestBuilder().column(1L, 2L, 3L).build(); + Table nullable = new Table.TestBuilder().column(1, null, 3).build()) { + Assertions.assertThrows(CudfException.class, () -> program.computeTable(wrongType).close()); + Assertions.assertThrows(CudfException.class, () -> program.computeTable(nullable).close()); + } + } + + @Test + void testValidation() { + AstExpression expression = new JitOperation(JitOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (Table schemaTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expression.compileJit()) { + Assertions.assertThrows(NullPointerException.class, + () -> AstJitProgram.compile(null, compiled)); + Assertions.assertThrows(NullPointerException.class, + () -> AstJitProgram.compile(schemaTable, (CompiledExpression[]) null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> AstJitProgram.compile(schemaTable)); + Assertions.assertThrows(NullPointerException.class, + () -> AstJitProgram.compile(schemaTable, compiled, null)); + } + + try (Table schemaTable = new Table.TestBuilder().column(1, 2, 3).build()) { + CompiledExpression closedExpression = expression.compileJit(); + closedExpression.close(); + Assertions.assertThrows(IllegalStateException.class, + () -> AstJitProgram.compile(schemaTable, closedExpression)); + } + + AstExpression defaultExpression = new BinaryOperation(BinaryOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (Table schemaTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = defaultExpression.compile()) { + Assertions.assertThrows(IllegalArgumentException.class, + () -> AstJitProgram.compile(schemaTable, compiled)); + } + + try (Table closedTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expression.compileJit()) { + closedTable.close(); + Assertions.assertThrows(IllegalStateException.class, + () -> AstJitProgram.compile(closedTable, compiled)); + } + } + + @Test + void testClosedInputs() { + AstExpression expression = new JitOperation(JitOperator.ADD, + new ColumnReference(0), Literal.ofInt(1)); + try (Table schemaTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expression.compileJit()) { + AstJitProgram program = AstJitProgram.compile(schemaTable, compiled); + Assertions.assertThrows(NullPointerException.class, () -> program.computeTable(null)); + program.close(); + Assertions.assertThrows(IllegalStateException.class, + () -> program.computeTable(schemaTable)); + Assertions.assertThrows(IllegalStateException.class, program::close); + } + + try (Table schemaTable = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiled = expression.compileJit(); + AstJitProgram program = AstJitProgram.compile(schemaTable, compiled); + Table closedTable = new Table.TestBuilder().column(4, 5, 6).build()) { + closedTable.close(); + Assertions.assertThrows(IllegalStateException.class, + () -> program.computeTable(closedTable)); + } + } +}