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}; 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/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/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 52c3c3ac03d7..cc16d3af1394 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java @@ -12,8 +12,15 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Objects; + /** This class wraps a native compiled AST and must be closed to avoid native memory leaks. */ public class CompiledExpression implements AutoCloseable { + enum CompilationMode { + DEFAULT, + JIT + } + static { NativeDepsLoader.loadNativeDeps(); } @@ -52,42 +59,105 @@ public boolean isClean() { } private final CompiledExpressionCleaner cleaner; + private final CompilationMode mode; private boolean isClosed = false; /** Construct a compiled expression from a serialized AST */ - CompiledExpression(byte[] serializedExpression) { - this(compile(serializedExpression)); + CompiledExpression(byte[] serializedExpression, CompilationMode mode) { + this(mode == CompilationMode.JIT ? compileJit(serializedExpression) : + compile(serializedExpression), mode); } /** Construct a compiled expression from a native compiled AST pointer */ - CompiledExpression(long nativeHandle) { + private CompiledExpression(long nativeHandle, CompilationMode mode) { this.cleaner = new CompiledExpressionCleaner(nativeHandle); + this.mode = mode; MemoryCleaner.register(this, cleaner); cleaner.addRef(); } /** - * Compute a new column by applying this AST expression to the specified table. All - * {@link ColumnReference} instances within the expression will use the sole input table, - * even if they try to specify a non-existent table, e.g.: {@link TableReference#RIGHT}. + * Compute a new column by applying this AST expression to the specified table. All column + * references must use {@link TableReference#LEFT}; references to {@link TableReference#RIGHT} + * are rejected because this operation has only one input table. + * An expression produced by {@link AstExpression#compileJit()} always uses the JIT backend. + * Otherwise, execution uses the process-level backend selection. + * * @param table input table for this expression * @return new column computed from this expression applied to the input table + * @throws ai.rapids.cudf.CudfException if the expression refers to + * {@link TableReference#RIGHT}, or if compilation or evaluation fails */ public ColumnVector computeColumn(Table table) { - return new ColumnVector(computeColumn(cleaner.nativeHandle, table.getNativeView())); + long result; + try { + result = computeColumn(cleaner.nativeHandle, table.getNativeView()); + } finally { + reachabilityFence(this); + reachabilityFence(table); + } + return new ColumnVector(result); } /** - * Compute a new column by applying this expression with the libcudf JIT executor, independent - * of the process-level backend selected for {@link #computeColumn}. + * Compute a new table by applying expressions to the input table in one multi-output JIT + * transform. Output column {@code i} contains the result of {@code expressions[i]}. * - * @param table input table for this expression - * @return new column computed from this expression applied to the input table - * @throws ai.rapids.cudf.CudfException if the expression refers to - * {@link TableReference#RIGHT}, or if JIT compilation or evaluation fails + * @param table input table for the expressions + * @param expressions non-empty JIT-compiled expressions to evaluate in output order + * @return table containing one output column per expression + * @throws NullPointerException if the table, expression array, or an expression is null + * @throws IllegalArgumentException if no expressions are provided or an expression was not + * produced by {@link AstExpression#compileJit()} + * @throws IllegalStateException if the table or an expression is closed + * @throws ai.rapids.cudf.CudfException if JIT compilation or evaluation fails */ - public ColumnVector computeColumnJit(Table table) { - return new ColumnVector(computeColumnJit(cleaner.nativeHandle, table.getNativeView())); + public static Table computeTableJit(Table table, CompiledExpression... expressions) { + Objects.requireNonNull(table, "table"); + Objects.requireNonNull(expressions, "expressions"); + if (expressions.length == 0) { + throw new IllegalArgumentException("At least one expression is required"); + } + + long tableHandle = table.getNativeView(); + if (tableHandle == 0) { + throw new IllegalStateException("Table is closed"); + } + + CompiledExpression[] expressionRefs = expressions.clone(); + long[] nativeHandles = 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( + expressions[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"); + } + } + return nativeHandles; + } + + static void reachabilityFence(Object object) { + if (object != null) { + synchronized (object) { + // The monitor operation is a Java 8 reachability fence. + } + } } @Override @@ -101,13 +171,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/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/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index 44748ef6b9b9..a5f7f86d45e6 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -452,6 +452,10 @@ cudf::ast::operation const& compile_binary_expression(cudf::jni::ast::compiled_e cudf::ast::expression const& compile_jit_expression(cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { + if (!compiled_expr.is_jit()) { + throw std::invalid_argument("JIT operations require an expression compiled for JIT"); + } + auto const jni_op_value = jni_ast.read_byte(); auto const op_info = jni_to_jit_operator(jni_op_value); auto const jni_policy_value = jni_ast.read_byte(); @@ -530,39 +534,38 @@ cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& c } /** Decode a serialized AST into a native libcudf AST and associated resources */ -std::unique_ptr 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 = 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 @@ -570,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, @@ -587,24 +608,37 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumn JNI_TRY { cudf::jni::auto_set_device(env); - return execute_compiled_expression(j_ast, j_table, execution_backend::DEFAULT); + return execute_compiled_expression(j_ast, j_table); } JNI_CATCH(env, 0); } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumnJit(JNIEnv* env, - jclass, - jlong j_ast, - jlong j_table) +JNIEXPORT jlongArray JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeTableJitNative( + JNIEnv* env, jclass, jlongArray j_asts, jlong j_table) { - JNI_NULL_CHECK(env, j_ast, "Compiled AST pointer is null", 0); - JNI_NULL_CHECK(env, j_table, "Table view pointer is null", 0); + JNI_NULL_CHECK(env, j_asts, "Compiled AST pointer array is null", nullptr); + JNI_NULL_CHECK(env, j_table, "Table view pointer is null", nullptr); JNI_TRY { cudf::jni::auto_set_device(env); - return execute_compiled_expression(j_ast, j_table, execution_backend::JIT); + cudf::jni::native_jlongArray ast_handles(env, j_asts); + if (ast_handles.size() == 0) { throw std::invalid_argument("At least one AST is required"); } + + std::vector> 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, 0); + JNI_CATCH(env, nullptr); } JNIEXPORT void JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_destroy(JNIEnv* env, diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index f51457bd0606..4b7fc734fb1b 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -6,9 +6,13 @@ #pragma once #include +#include +#include +#include #include #include +#include #include #include #include @@ -17,19 +21,37 @@ namespace cudf { namespace jni { namespace ast { +enum class compilation_mode { DEFAULT, JIT }; + /** A class to capture all resources associated with a compiled AST expression. */ class compiled_expr { + 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 a JIT expression tree */ + std::vector> scalar_columns; + /** All expression nodes within the expression tree */ cudf::ast::tree expressions; public: + explicit compiled_expr(compilation_mode mode) : mode{mode} {} + template cudf::ast::literal const& add_literal(ScalarType& scalar, std::unique_ptr scalar_ptr) { + 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)); return expressions.emplace(scalar); } @@ -61,13 +83,38 @@ class compiled_expr { template cudf::ast::expression const& add_jit_expression(F&& factory) { + if (!is_jit()) { + throw std::invalid_argument("JIT operations require an expression compiled for JIT"); + } 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 has_jit_literals() const { return !scalar_columns.empty(); } + + [[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 + { + 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/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)); + } + } +} 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..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,20 +471,233 @@ 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.computeColumn(t).close()); + } + } + + @Test + 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.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, jitFirst); + assertColumnsAreEqual(expectedSecond, jitSecond); + } + } + + @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.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 = + 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.compileJit(); + CompiledExpression subtractCompiled = subtract.compileJit(); + CompiledExpression sumCompiled = secondSum.compileJit()) { + 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.compileJit(); + CompiledExpression sumCompiled = sum.compileJit(); + 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.compileJit(); + CompiledExpression secondCompiled = second.compileJit(); + 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.compileJit(); + CompiledExpression sumCompiled = sum.compileJit(); + 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.compileJit()) { + 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.compileJit(); + 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.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 + 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.compileJit(); + CompiledExpression invalidCompiled = invalid.compileJit()) { Assertions.assertThrows(CudfException.class, - () -> compiledExpr.computeColumnJit(t).close()); + () -> CompiledExpression.computeTableJit( + input, validCompiled, invalidCompiled).close()); + try (ColumnVector actual = validCompiled.computeColumn(input); + ColumnVector expected = ColumnVector.fromInts(2, 3, 4)) { + assertColumnsAreEqual(expected, actual); + } } } @@ -492,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); } @@ -512,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); } @@ -523,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); } @@ -539,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); } @@ -577,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); } @@ -593,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); } @@ -625,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); } @@ -654,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); } @@ -699,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); } @@ -723,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); } @@ -736,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); @@ -749,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); }