From fdde92b2fcbf0bc76bab54da8ff3675ceb568712 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 3 Jul 2026 11:30:53 +0800 Subject: [PATCH 01/11] Expose JIT AST operations in Java --- .../ai/rapids/cudf/ast/AstExpression.java | 3 +- .../ai/rapids/cudf/ast/JitComplianceMode.java | 33 ++ .../java/ai/rapids/cudf/ast/JitOperation.java | 109 +++++++ .../java/ai/rapids/cudf/ast/JitOperator.java | 76 +++++ .../main/java/ai/rapids/cudf/ast/Literal.java | 30 ++ .../main/native/src/CompiledExpression.cpp | 300 +++++++++++++++++- .../src/main/native/src/jni_compiled_expr.hpp | 18 +- .../cudf/ast/CompiledExpressionTest.java | 95 ++++++ 8 files changed, 646 insertions(+), 18 deletions(-) create mode 100644 java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java create mode 100644 java/src/main/java/ai/rapids/cudf/ast/JitOperation.java create mode 100644 java/src/main/java/ai/rapids/cudf/ast/JitOperator.java 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 27406402f84c..33974955f4ba 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java @@ -19,7 +19,8 @@ protected enum ExpressionType { NULL_LITERAL(1), COLUMN_REFERENCE(2), UNARY_EXPRESSION(3), - BINARY_EXPRESSION(4); + BINARY_EXPRESSION(4), + JIT_EXPRESSION(5); private final byte nativeId; diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java b/java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java new file mode 100644 index 000000000000..f3dbb6f3bca5 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +import java.nio.ByteBuffer; + +/** + * Arithmetic compliance mode for JIT AST operations. + * NOTE: This must be kept in sync with `jni_to_jit_compliance_mode` in CompiledExpression.cpp! + */ +public enum JitComplianceMode { + DEFAULT(0), + ANSI(1), + ANSI_TRY(2); + + private final byte nativeId; + + JitComplianceMode(int nativeId) { + this.nativeId = (byte) nativeId; + assert this.nativeId == nativeId; + } + + int getSerializedSize() { + return Byte.BYTES; + } + + void serialize(ByteBuffer bb) { + bb.put(nativeId); + } +} diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java new file mode 100644 index 000000000000..8f9cd6bea95f --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +import java.nio.ByteBuffer; +import java.util.Objects; + +/** A JIT operation consisting of a row IR opcode and operands. */ +public final class JitOperation extends AstExpression { + private final JitOperator op; + private final JitComplianceMode complianceMode; + private final AstExpression[] inputs; + private final Integer targetScale; + + public JitOperation(JitOperator op, AstExpression... inputs) { + this(op, JitComplianceMode.DEFAULT, null, inputs); + } + + public JitOperation(JitOperator op, JitComplianceMode complianceMode, AstExpression... inputs) { + this(op, complianceMode, null, inputs); + } + + public JitOperation(JitOperator op, int targetScale, AstExpression... inputs) { + this(op, JitComplianceMode.DEFAULT, Integer.valueOf(targetScale), inputs); + } + + private JitOperation( + JitOperator op, + JitComplianceMode complianceMode, + Integer targetScale, + AstExpression... inputs) { + this.op = Objects.requireNonNull(op, "op is null"); + this.complianceMode = Objects.requireNonNull(complianceMode, "complianceMode is null"); + this.inputs = Objects.requireNonNull(inputs, "inputs is null").clone(); + this.targetScale = targetScale; + if (this.inputs.length != op.getArity()) { + throw new IllegalArgumentException( + op + " requires " + op.getArity() + " inputs, found " + this.inputs.length); + } + if (!op.supportsComplianceMode() && complianceMode != JitComplianceMode.DEFAULT) { + throw new IllegalArgumentException(op + " does not support compliance mode " + complianceMode); + } + if (op == JitOperator.PRECISION_CHECK && complianceMode == JitComplianceMode.DEFAULT) { + throw new IllegalArgumentException("PRECISION_CHECK requires ANSI or ANSI_TRY compliance mode"); + } + if (op.requiresTargetScale() != (targetScale != null)) { + throw new IllegalArgumentException(op + " target scale usage is invalid"); + } + for (AstExpression input : this.inputs) { + Objects.requireNonNull(input, "input is null"); + } + } + + @Override + int getSerializedSize() { + int size = ExpressionType.JIT_EXPRESSION.getSerializedSize() + + op.getSerializedSize() + + complianceMode.getSerializedSize() + + Byte.BYTES + + Byte.BYTES; + if (targetScale != null) { + size += Integer.BYTES; + } + for (AstExpression input : inputs) { + size += input.getSerializedSize(); + } + return size; + } + + @Override + void serialize(ByteBuffer bb) { + ExpressionType.JIT_EXPRESSION.serialize(bb); + op.serialize(bb); + complianceMode.serialize(bb); + bb.put((byte) inputs.length); + bb.put((byte) (targetScale == null ? 0 : 1)); + if (targetScale != null) { + bb.putInt(targetScale); + } + for (AstExpression input : inputs) { + input.serialize(bb); + } + } + + @Override + public String toString() { + StringBuilder ret = new StringBuilder(op.toString()); + if (complianceMode != JitComplianceMode.DEFAULT) { + ret.append("[").append(complianceMode).append("]"); + } + ret.append("("); + for (int i = 0; i < inputs.length; i++) { + if (i > 0) { + ret.append(", "); + } + ret.append(inputs[i]); + } + if (targetScale != null) { + if (inputs.length > 0) { + ret.append(", "); + } + ret.append("scale=").append(targetScale); + } + return ret.append(")").toString(); + } +} diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java new file mode 100644 index 000000000000..e99720bbfe96 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +import java.nio.ByteBuffer; + +/** + * Enumeration of AST JIT operators backed by libcudf row IR opcodes. + * NOTE: This must be kept in sync with `compile_jit_expression` in CompiledExpression.cpp! + */ +public enum JitOperator { + ADD(0, 2, true, false), + SUB(1, 2, true, false), + MUL(2, 2, true, false), + DIV(3, 2, true, false), + MOD(4, 2, true, false), + ABS(5, 1, true, false), + NEG(6, 1, true, false), + PRECISION_CHECK(7, 2, true, false), + BITWISE_SHIFT_LEFT(8, 2, false, false), + BITWISE_SHIFT_RIGHT(9, 2, false, false), + COALESCE(10, 2, false, false), + PREDICATE(11, 1, false, false), + CAST_TO_BOOL8(12, 1, false, false), + CAST_TO_INT8(13, 1, false, false), + CAST_TO_INT16(14, 1, false, false), + CAST_TO_INT32(15, 1, false, false), + CAST_TO_INT64(16, 1, false, false), + CAST_TO_UINT8(17, 1, false, false), + CAST_TO_UINT16(18, 1, false, false), + CAST_TO_UINT32(19, 1, false, false), + CAST_TO_UINT64(20, 1, false, false), + CAST_TO_FLOAT32(21, 1, false, false), + CAST_TO_FLOAT64(22, 1, false, false), + CAST_TO_DECIMAL32(23, 1, false, false), + CAST_TO_DECIMAL64(24, 1, false, false), + CAST_TO_DECIMAL128(25, 1, false, false), + RESCALE(26, 1, false, true), + IF_ELSE(27, 3, false, false); + + private final byte nativeId; + private final int arity; + private final boolean supportsComplianceMode; + private final boolean requiresTargetScale; + + JitOperator(int nativeId, int arity, boolean supportsComplianceMode, boolean requiresTargetScale) { + this.nativeId = (byte) nativeId; + this.arity = arity; + this.supportsComplianceMode = supportsComplianceMode; + this.requiresTargetScale = requiresTargetScale; + assert this.nativeId == nativeId; + } + + int getArity() { + return arity; + } + + boolean supportsComplianceMode() { + return supportsComplianceMode; + } + + boolean requiresTargetScale() { + return requiresTargetScale; + } + + int getSerializedSize() { + return Byte.BYTES; + } + + void serialize(ByteBuffer bb) { + bb.put(nativeId); + } +} 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 0fcbeaf317da..719707a0a1c8 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -7,6 +7,7 @@ import ai.rapids.cudf.DType; +import java.math.BigInteger; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; @@ -123,6 +124,23 @@ public static Literal ofDouble(Double value) { return ofDouble(value.doubleValue()); } + /** Construct a decimal literal with the specified type and unscaled value. */ + public static Literal ofDecimal(DType type, BigInteger unscaledValue) { + if (!type.isDecimalType()) { + throw new IllegalArgumentException("type is not a decimal: " + type); + } + if (unscaledValue == null) { + return ofNull(type); + } + if (type.getTypeId() == DType.DTypeEnum.DECIMAL32) { + return ofIntBasedType(type, unscaledValue.intValueExact()); + } else if (type.getTypeId() == DType.DTypeEnum.DECIMAL64) { + return ofLongBasedType(type, unscaledValue.longValueExact()); + } else { + return new Literal(type, convertDecimal128FromJavaToCudf(unscaledValue.toByteArray(), type)); + } + } + /** Construct a timestamp days literal with the specified value. */ public static Literal ofTimestampDaysFromInt(int value) { return ofIntBasedType(DType.TIMESTAMP_DAYS, value); @@ -275,4 +293,16 @@ private static Literal ofLongBasedType(DType type, long value) { ByteBuffer.wrap(serializedValue).order(ByteOrder.nativeOrder()).putLong(value); return new Literal(type, serializedValue); } + + private static byte[] convertDecimal128FromJavaToCudf(byte[] bytes, DType type) { + byte[] finalBytes = new byte[type.getSizeInBytes()]; + byte signByte = (bytes[0] & 0x80) > 0 ? (byte) 0xff : (byte) 0x00; + for (int i = bytes.length; i < finalBytes.length; i++) { + finalBytes[i] = signByte; + } + for (int i = 0; i < bytes.length; i++) { + finalBytes[i] = bytes[bytes.length - i - 1]; + } + return finalBytes; + } } diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index ed4eb409138c..bdfa9b95441a 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,13 +7,16 @@ #include "jni_compiled_expr.hpp" #include +#include #include #include #include #include #include +#include #include +#include #include #include @@ -98,7 +101,8 @@ class jni_serialized_ast { return cudf::data_type(dtype_id); } case cudf::type_id::DECIMAL32: - case cudf::type_id::DECIMAL64: { + case cudf::type_id::DECIMAL64: + case cudf::type_id::DECIMAL128: { int32_t const scale = read_byte(); return cudf::data_type(dtype_id, scale); } @@ -116,7 +120,8 @@ enum class jni_serialized_expression_type : int8_t { NULL_LITERAL = 1, COLUMN_REFERENCE = 2, UNARY_OPERATION = 3, - BINARY_OPERATION = 4 + BINARY_OPERATION = 4, + JIT_OPERATION = 5 }; /** @@ -193,6 +198,80 @@ cudf::ast::ast_operator jni_to_binary_operator(jbyte jni_op_value) } } +enum class jni_jit_compliance_mode : int8_t { DEFAULT, ANSI, ANSI_TRY }; + +/** + * Convert a Java AST serialized byte representing a JIT compliance mode. + * NOTE: This must be kept in sync with the enumeration in JitComplianceMode.java! + */ +jni_jit_compliance_mode jni_to_jit_compliance_mode(jbyte jni_mode_value) +{ + switch (jni_mode_value) { + case 0: return jni_jit_compliance_mode::DEFAULT; + case 1: return jni_jit_compliance_mode::ANSI; + case 2: return jni_jit_compliance_mode::ANSI_TRY; + default: throw std::invalid_argument("unexpected JNI AST JIT compliance mode value"); + } +} + +void expect_default_jit_compliance_mode(jni_jit_compliance_mode mode) +{ + if (mode != jni_jit_compliance_mode::DEFAULT) { + throw std::invalid_argument("unexpected compliance mode for JNI AST JIT operator"); + } +} + +void expect_non_default_jit_compliance_mode(jni_jit_compliance_mode mode) +{ + if (mode == jni_jit_compliance_mode::DEFAULT) { + throw std::invalid_argument( + "expected ANSI or ANSI_TRY compliance mode for JNI AST JIT operator"); + } +} + +cudf::error_policy jni_to_jit_error_policy(jni_jit_compliance_mode mode) +{ + expect_non_default_jit_compliance_mode(mode); + return mode == jni_jit_compliance_mode::ANSI_TRY ? cudf::error_policy::NULLIFY + : cudf::error_policy::PROPAGATE; +} + +cudf::ast::expression const& dispatch_jit_expression( + cudf::ast::tree& tree, + std::vector> const& args, + jni_jit_compliance_mode mode, + cudf::ast::jit::op default_op, + cudf::ast::jit::op overflow_op) +{ + if (mode == jni_jit_compliance_mode::DEFAULT) { + return cudf::ast::jit::operation(tree, default_op, args); + } + return cudf::ast::jit::operation(tree, overflow_op, args, jni_to_jit_error_policy(mode)); +} + +void expect_jit_arity(std::vector> const& args, + std::size_t expected) +{ + if (args.size() != expected) { + throw std::invalid_argument("unexpected JNI AST JIT operator arity"); + } +} + +void expect_no_target_scale(std::optional const& target_scale) +{ + if (target_scale.has_value()) { + throw std::invalid_argument("unexpected target scale for JNI AST JIT operator"); + } +} + +int32_t require_target_scale(std::optional const& target_scale) +{ + if (!target_scale.has_value()) { + throw std::invalid_argument("expected target scale for JNI AST JIT operator"); + } + return target_scale.value(); +} + /** * Convert a Java AST serialized byte representing an AST table reference into the * corresponding libcudf AST table reference. @@ -291,10 +370,10 @@ struct make_literal { } /** Default functor implementation to catch type dispatch errors */ - template < - typename T, - std::enable_if_t() && !cudf::is_timestamp() && - !cudf::is_duration() && !std::is_same_v>* = nullptr> + template () && !cudf::is_timestamp() && + !cudf::is_duration() && !cudf::is_fixed_point() && + !std::is_same_v>* = nullptr> cudf::ast::literal& operator()(cudf::data_type dtype, bool is_valid, cudf::jni::ast::compiled_expr& compiled_expr, @@ -302,6 +381,24 @@ struct make_literal { { throw std::logic_error("Unsupported AST literal type"); } + + /** Construct an AST literal from a fixed-point value */ + template ()>* = nullptr> + cudf::ast::literal& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) + { + using rep_type = typename T::rep; + auto const val = is_valid ? jni_ast.read() : rep_type{}; + std::unique_ptr scalar_ptr = + cudf::make_fixed_point_scalar(val, numeric::scale_type{dtype.scale()}); + scalar_ptr->set_valid_async(is_valid); + + auto& fixed_point_scalar = static_cast&>(*scalar_ptr); + return compiled_expr.add_literal(std::make_unique(fixed_point_scalar), + std::move(scalar_ptr)); + } }; /** Decode a serialized AST literal */ @@ -324,15 +421,15 @@ cudf::ast::column_reference& compile_column_reference(cudf::jni::ast::compiled_e } // forward declaration -cudf::ast::expression& 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::ast::operation& 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& 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( std::make_unique(ast_op, child_expression)); } @@ -341,16 +438,185 @@ cudf::ast::operation& compile_unary_expression(cudf::jni::ast::compiled_expr& co cudf::ast::operation& 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& left_child = compile_expression(compiled_expr, jni_ast); - cudf::ast::expression& 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( std::make_unique(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) +{ + auto const op_byte = jni_ast.read_byte(); + auto const mode = jni_to_jit_compliance_mode(jni_ast.read_byte()); + auto const arity = static_cast(jni_ast.read_byte()); + if (arity < 0) { throw std::invalid_argument("unexpected JNI AST JIT operator arity"); } + auto const has_target_scale = jni_ast.read_byte(); + std::optional target_scale; + if (has_target_scale != 0) { target_scale = jni_ast.read(); } + + 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& { + switch (op_byte) { + case 0: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::ADD, cudf::ast::jit::op::ADD_OVERFLOW); + case 1: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::SUB, cudf::ast::jit::op::SUB_OVERFLOW); + case 2: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::MUL, cudf::ast::jit::op::MUL_OVERFLOW); + case 3: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::DIV, cudf::ast::jit::op::DIV_OVERFLOW); + case 4: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::MOD, cudf::ast::jit::op::MOD_OVERFLOW); + case 5: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::ABS, cudf::ast::jit::op::ABS_OVERFLOW); + case 6: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + return dispatch_jit_expression( + tree, args, mode, cudf::ast::jit::op::NEG, cudf::ast::jit::op::NEG_OVERFLOW); + case 7: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + return cudf::ast::jit::operation( + tree, cudf::ast::jit::op::CHECK_PRECISION, args, jni_to_jit_error_policy(mode)); + case 8: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::BITWISE_SHIFT_LEFT, args); + case 9: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::BITWISE_SHIFT_RIGHT, args); + case 10: + expect_jit_arity(args, 2); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::COALESCE, args); + case 11: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::PREDICATE, args); + case 12: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_BOOL8, args); + case 13: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT8, args); + case 14: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT16, args); + case 15: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT32, args); + case 16: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT64, args); + case 17: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT8, args); + case 18: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT16, args); + case 19: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT32, args); + case 20: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT64, args); + case 21: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_FLOAT32, args); + case 22: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_FLOAT64, args); + case 23: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_DECIMAL32, args); + case 24: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_DECIMAL64, args); + case 25: + expect_jit_arity(args, 1); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_DECIMAL128, args); + case 26: + expect_jit_arity(args, 1); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, + cudf::ast::jit::op::RESCALE, + args, + cudf::error_policy::PROPAGATE, + require_target_scale(target_scale)); + case 27: + expect_jit_arity(args, 3); + expect_no_target_scale(target_scale); + expect_default_jit_compliance_mode(mode); + return cudf::ast::jit::operation(tree, cudf::ast::jit::op::IF_ELSE, args); + default: throw std::invalid_argument("unexpected JNI AST JIT operator value"); + } + }); +} + /** Decode a serialized AST expression by reading the expression type and dispatching */ -cudf::ast::expression& 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) { @@ -364,6 +630,8 @@ cudf::ast::expression& compile_expression(cudf::jni::ast::compiled_expr& compile return compile_unary_expression(compiled_expr, jni_ast); case jni_serialized_expression_type::BINARY_OPERATION: return compile_binary_expression(compiled_expr, jni_ast); + case jni_serialized_expression_type::JIT_OPERATION: + return compile_jit_expression(compiled_expr, jni_ast); default: throw std::invalid_argument("data is not a serialized AST expression"); } } diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index 1754c4d41218..7d0e7f6214f2 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -29,32 +29,48 @@ class compiled_expr { /** All expression nodes within the expression tree */ std::vector> expressions; + /** Expression tree for JIT helper-created nodes */ + cudf::ast::tree jit_expressions; + /** GPU scalar instances that correspond to literal nodes */ std::vector> scalars; + cudf::ast::expression const* top_expression = nullptr; + public: cudf::ast::literal& add_literal(std::unique_ptr literal_ptr, std::unique_ptr scalar_ptr) { expressions.push_back(std::move(literal_ptr)); scalars.push_back(std::move(scalar_ptr)); + top_expression = expressions.back().get(); return static_cast(*expressions.back()); } cudf::ast::column_reference& add_column_ref(std::unique_ptr ref_ptr) { expressions.push_back(std::move(ref_ptr)); + top_expression = expressions.back().get(); return static_cast(*expressions.back()); } cudf::ast::operation& add_operation(std::unique_ptr expr_ptr) { expressions.push_back(std::move(expr_ptr)); + top_expression = expressions.back().get(); return static_cast(*expressions.back()); } + template + cudf::ast::expression const& add_jit_expression(F&& factory) + { + auto& expr = factory(jit_expressions); + top_expression = &expr; + return expr; + } + /** Return the expression node at the top of the tree */ - cudf::ast::expression& get_top_expression() const { return *expressions.back(); } + cudf::ast::expression const& get_top_expression() const { return *top_expression; } }; } // 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 cdfcbe596d07..7a0b03fb8576 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -18,6 +18,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.junit.jupiter.params.provider.NullSource; +import java.math.BigInteger; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -26,6 +27,7 @@ import java.util.stream.Stream; import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual; +import static org.junit.jupiter.api.Assumptions.assumeTrue; public class CompiledExpressionTest extends CudfTestBase { @Test @@ -295,6 +297,99 @@ public void testDurationNanoSecondsLiteralTransform(Long value) { } } + @Test + public void testDecimal128LiteralTransform() { + int scale = -4; + BigInteger value = new BigInteger("-123456789012345678901234567890"); + Literal expr = Literal.ofDecimal(DType.create(DType.DTypeEnum.DECIMAL128, scale), value); + try (Table t = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiledExpr = expr.compile(); + ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector expected = ColumnVector.decimalFromBigInt(scale, value, value, value)) { + assertColumnsAreEqual(expected, actual); + } + } + + @Test + void testJitOperationValidation() { + AstExpression[] inputs = new AstExpression[] { + new ColumnReference(0), + new ColumnReference(1), + new ColumnReference(2) + }; + for (JitOperator op : JitOperator.values()) { + if (op == JitOperator.RESCALE) { + Assertions.assertDoesNotThrow( + () -> new JitOperation(op, -2, Arrays.copyOf(inputs, op.getArity()))); + } else { + JitComplianceMode mode = op == JitOperator.PRECISION_CHECK + ? JitComplianceMode.ANSI : JitComplianceMode.DEFAULT; + Assertions.assertDoesNotThrow( + () -> new JitOperation(op, mode, Arrays.copyOf(inputs, op.getArity()))); + } + if (op.getArity() > 0) { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new JitOperation(op, Arrays.copyOf(inputs, op.getArity() - 1))); + } + } + + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new JitOperation(JitOperator.PRECISION_CHECK, + new ColumnReference(0), new ColumnReference(1))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new JitOperation(JitOperator.IF_ELSE, JitComplianceMode.ANSI, + new ColumnReference(0), new ColumnReference(1), new ColumnReference(2))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new JitOperation(JitOperator.RESCALE, new ColumnReference(0))); + } + + @Test + void testJitTryDivModTransform() { + assumeTrue("1".equals(System.getenv("LIBCUDF_JIT_ENABLED"))); + try (Table t = new Table.TestBuilder() + .column(10, 7, null, 6) + .column(2, 0, 3, null) + .build()) { + JitOperation divExpr = new JitOperation(JitOperator.DIV, JitComplianceMode.ANSI_TRY, + new ColumnReference(0), new ColumnReference(1)); + try (CompiledExpression compiledExpr = divExpr.compile(); + ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector expected = ColumnVector.fromBoxedInts(5, null, null, null)) { + assertColumnsAreEqual(expected, actual); + } + + JitOperation modExpr = new JitOperation(JitOperator.MOD, JitComplianceMode.ANSI_TRY, + new ColumnReference(0), new ColumnReference(1)); + try (CompiledExpression compiledExpr = modExpr.compile(); + ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector expected = ColumnVector.fromBoxedInts(0, null, null, null)) { + assertColumnsAreEqual(expected, actual); + } + } + } + + @Test + void testJitIfElseTransform() { + assumeTrue("1".equals(System.getenv("LIBCUDF_JIT_ENABLED"))); + try (Table t = new Table.TestBuilder() + .column(1, 2, 3, 4) + .column(10, 20, 30, 40) + .column(true, false, true, false) + .build()) { + JitOperation expr = new JitOperation(JitOperator.IF_ELSE, + new ColumnReference(0), new ColumnReference(1), new ColumnReference(2)); + try (CompiledExpression compiledExpr = expr.compile(); + ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector expected = ColumnVector.fromBoxedInts(1, 20, 3, 40)) { + assertColumnsAreEqual(expected, actual); + } + } + } + private static ArrayList mapArray(T[] input, Function func) { ArrayList result = new ArrayList<>(input.length); for (T t : input) { From db1e27a1805a40120248c6d864f4573e2be7e871 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 6 Jul 2026 18:24:55 +0800 Subject: [PATCH 02/11] Align Java AST JIT bindings with libcudf APIs --- java/src/main/java/ai/rapids/cudf/Cudf.java | 12 +- .../ai/rapids/cudf/ast/AstExpression.java | 5 +- .../rapids/cudf/ast/CompiledExpression.java | 16 +- ...omplianceMode.java => JitErrorPolicy.java} | 16 +- .../java/ai/rapids/cudf/ast/JitOperation.java | 62 ++- .../java/ai/rapids/cudf/ast/JitOperator.java | 90 +++-- .../main/java/ai/rapids/cudf/ast/Literal.java | 19 +- .../main/native/src/CompiledExpression.cpp | 321 +++++---------- java/src/main/native/src/CudfJni.cpp | 13 +- .../src/main/native/src/jni_compiled_expr.hpp | 2 +- .../cudf/ast/CompiledExpressionTest.java | 365 ++++++++++++++++-- 11 files changed, 605 insertions(+), 316 deletions(-) rename java/src/main/java/ai/rapids/cudf/ast/{JitComplianceMode.java => JitErrorPolicy.java} (53%) diff --git a/java/src/main/java/ai/rapids/cudf/Cudf.java b/java/src/main/java/ai/rapids/cudf/Cudf.java index 2d8cf6c88ff4..07cb67d1a519 100644 --- a/java/src/main/java/ai/rapids/cudf/Cudf.java +++ b/java/src/main/java/ai/rapids/cudf/Cudf.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,6 +11,16 @@ public class Cudf { NativeDepsLoader.loadNativeDeps(); } + /** + * Initialize the libcudf JIT runtime and program cache. + * This method may be called repeatedly. It validates runtime dependencies but does not actively + * enable JIT evaluation; {@code CompiledExpression.computeColumn} continues to use the + * process-level libcudf configuration. + * + * @throws CudfException if the JIT runtime cannot be initialized + */ + public static native void initializeJitRuntime(); + /** * cuDF copies that are smaller than the threshold will use a kernel to copy, instead * of cudaMemcpyAsync. 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 33974955f4ba..8175e24d6c85 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -12,7 +12,8 @@ public abstract class AstExpression { /** * Enumeration for the types of AST nodes that can appear in a serialized AST. - * NOTE: This must be kept in sync with the `jni_serialized_node_type` in CompiledExpression.cpp! + * NOTE: This must be kept in sync with `jni_serialized_expression_type` in + * CompiledExpression.cpp! */ protected enum ExpressionType { VALID_LITERAL(0), 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 5d92d623b38b..80b7a1217392 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -77,6 +77,19 @@ public ColumnVector computeColumn(Table table) { return new ColumnVector(computeColumn(cleaner.nativeHandle, table.getNativeView())); } + /** + * Compute a new column by applying this expression with the libcudf JIT executor, independent + * of the process-level backend selected for {@link #computeColumn}. + * + * @param table input table for this expression + * @return new column computed from this expression applied to the input table + * @throws ai.rapids.cudf.CudfException if the expression refers to + * {@link TableReference#RIGHT}, or if JIT compilation or evaluation fails + */ + public ColumnVector computeColumnJit(Table table) { + return new ColumnVector(computeColumnJit(cleaner.nativeHandle, table.getNativeView())); + } + @Override public synchronized void close() { cleaner.delRef(); @@ -95,5 +108,6 @@ public long getNativeHandle() { private static native long compile(byte[] serializedExpression); private static native long computeColumn(long astHandle, long tableHandle); + private static native long computeColumnJit(long astHandle, long tableHandle); private static native void destroy(long handle); } diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java similarity index 53% rename from java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java rename to java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java index f3dbb6f3bca5..e5dfc0069980 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitComplianceMode.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java @@ -8,17 +8,19 @@ import java.nio.ByteBuffer; /** - * Arithmetic compliance mode for JIT AST operations. - * NOTE: This must be kept in sync with `jni_to_jit_compliance_mode` in CompiledExpression.cpp! + * Error handling policy for fallible JIT AST operations. + * + * NOTE: This must be kept in sync with `jni_to_jit_error_policy` in CompiledExpression.cpp! */ -public enum JitComplianceMode { - DEFAULT(0), - ANSI(1), - ANSI_TRY(2); +public enum JitErrorPolicy { + /** Propagate an evaluation error to the caller. */ + PROPAGATE(0), + /** Produce null for a row where evaluation fails. */ + NULLIFY(1); private final byte nativeId; - JitComplianceMode(int nativeId) { + JitErrorPolicy(int nativeId) { this.nativeId = (byte) nativeId; assert this.nativeId == nativeId; } 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 8f9cd6bea95f..2d2f99dbb0ce 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java @@ -8,43 +8,71 @@ import java.nio.ByteBuffer; import java.util.Objects; -/** A JIT operation consisting of a row IR opcode and operands. */ +/** + * A libcudf row-IR operation. Expressions containing a JIT operation must be evaluated with + * {@link CompiledExpression#computeColumnJit}. + */ public final class JitOperation extends AstExpression { private final JitOperator op; - private final JitComplianceMode complianceMode; + private final JitErrorPolicy errorPolicy; private final AstExpression[] inputs; private final Integer targetScale; + /** + * Construct an operation that propagates evaluation errors. + * + * @param op operator to apply + * @param inputs operator inputs + * @throws NullPointerException if {@code op}, {@code inputs}, or an input is null + * @throws IllegalArgumentException if the operator arity or target-scale usage is invalid + */ public JitOperation(JitOperator op, AstExpression... inputs) { - this(op, JitComplianceMode.DEFAULT, null, inputs); + this(op, JitErrorPolicy.PROPAGATE, null, inputs); } - public JitOperation(JitOperator op, JitComplianceMode complianceMode, AstExpression... inputs) { - this(op, complianceMode, null, inputs); + /** + * Construct an operation with an explicit error policy. + * + * @param op operator to apply + * @param errorPolicy error handling policy + * @param inputs operator inputs + * @throws NullPointerException if any argument or input is null + * @throws IllegalArgumentException if the arity, policy, or target-scale usage is invalid + */ + public JitOperation(JitOperator op, JitErrorPolicy errorPolicy, AstExpression... inputs) { + this(op, errorPolicy, null, inputs); } + /** + * Construct a target-scale operation that propagates evaluation errors. + * The target scale is valid only for {@link JitOperator#RESCALE}. + * + * @param op operator to apply + * @param targetScale target fixed-point scale + * @param inputs operator inputs + * @throws NullPointerException if {@code op}, {@code inputs}, or an input is null + * @throws IllegalArgumentException if the operator is not {@link JitOperator#RESCALE} or its + * arity is invalid + */ public JitOperation(JitOperator op, int targetScale, AstExpression... inputs) { - this(op, JitComplianceMode.DEFAULT, Integer.valueOf(targetScale), inputs); + this(op, JitErrorPolicy.PROPAGATE, Integer.valueOf(targetScale), inputs); } private JitOperation( JitOperator op, - JitComplianceMode complianceMode, + JitErrorPolicy errorPolicy, Integer targetScale, AstExpression... inputs) { this.op = Objects.requireNonNull(op, "op is null"); - this.complianceMode = Objects.requireNonNull(complianceMode, "complianceMode is null"); + this.errorPolicy = Objects.requireNonNull(errorPolicy, "errorPolicy is null"); this.inputs = Objects.requireNonNull(inputs, "inputs is null").clone(); this.targetScale = targetScale; if (this.inputs.length != op.getArity()) { throw new IllegalArgumentException( op + " requires " + op.getArity() + " inputs, found " + this.inputs.length); } - if (!op.supportsComplianceMode() && complianceMode != JitComplianceMode.DEFAULT) { - throw new IllegalArgumentException(op + " does not support compliance mode " + complianceMode); - } - if (op == JitOperator.PRECISION_CHECK && complianceMode == JitComplianceMode.DEFAULT) { - throw new IllegalArgumentException("PRECISION_CHECK requires ANSI or ANSI_TRY compliance mode"); + if (!op.isFallible() && errorPolicy == JitErrorPolicy.NULLIFY) { + throw new IllegalArgumentException(op + " cannot nullify errors"); } if (op.requiresTargetScale() != (targetScale != null)) { throw new IllegalArgumentException(op + " target scale usage is invalid"); @@ -58,7 +86,7 @@ private JitOperation( int getSerializedSize() { int size = ExpressionType.JIT_EXPRESSION.getSerializedSize() + op.getSerializedSize() + - complianceMode.getSerializedSize() + + errorPolicy.getSerializedSize() + Byte.BYTES + Byte.BYTES; if (targetScale != null) { @@ -74,7 +102,7 @@ int getSerializedSize() { void serialize(ByteBuffer bb) { ExpressionType.JIT_EXPRESSION.serialize(bb); op.serialize(bb); - complianceMode.serialize(bb); + errorPolicy.serialize(bb); bb.put((byte) inputs.length); bb.put((byte) (targetScale == null ? 0 : 1)); if (targetScale != null) { @@ -88,8 +116,8 @@ void serialize(ByteBuffer bb) { @Override public String toString() { StringBuilder ret = new StringBuilder(op.toString()); - if (complianceMode != JitComplianceMode.DEFAULT) { - ret.append("[").append(complianceMode).append("]"); + if (errorPolicy != JitErrorPolicy.PROPAGATE) { + ret.append("[").append(errorPolicy).append("]"); } ret.append("("); for (int i = 0; i < inputs.length; i++) { diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java index e99720bbfe96..661a75689031 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java @@ -8,48 +8,68 @@ import java.nio.ByteBuffer; /** - * Enumeration of AST JIT operators backed by libcudf row IR opcodes. - * NOTE: This must be kept in sync with `compile_jit_expression` in CompiledExpression.cpp! + * The subset of libcudf row-IR operators exposed through the Java AST API. Standard comparison, + * logical, and mathematical operations remain available through {@link BinaryOperation} and + * {@link UnaryOperation}. + * Operators ending in {@code _OVERFLOW} and {@link #CHECK_PRECISION} are fallible and support + * {@link JitErrorPolicy#NULLIFY}. Other operators require the default + * {@link JitErrorPolicy#PROPAGATE} policy. + * + * NOTE: This must be kept in sync with `jni_to_jit_operator` in CompiledExpression.cpp! */ public enum JitOperator { - ADD(0, 2, true, false), - SUB(1, 2, true, false), - MUL(2, 2, true, false), - DIV(3, 2, true, false), - MOD(4, 2, true, false), - ABS(5, 1, true, false), - NEG(6, 1, true, false), - PRECISION_CHECK(7, 2, true, false), - BITWISE_SHIFT_LEFT(8, 2, false, false), - BITWISE_SHIFT_RIGHT(9, 2, false, false), - COALESCE(10, 2, false, false), - PREDICATE(11, 1, false, false), - CAST_TO_BOOL8(12, 1, false, false), - CAST_TO_INT8(13, 1, false, false), - CAST_TO_INT16(14, 1, false, false), - CAST_TO_INT32(15, 1, false, false), - CAST_TO_INT64(16, 1, false, false), - CAST_TO_UINT8(17, 1, false, false), - CAST_TO_UINT16(18, 1, false, false), - CAST_TO_UINT32(19, 1, false, false), - CAST_TO_UINT64(20, 1, false, false), - CAST_TO_FLOAT32(21, 1, false, false), - CAST_TO_FLOAT64(22, 1, false, false), - CAST_TO_DECIMAL32(23, 1, false, false), - CAST_TO_DECIMAL64(24, 1, false, false), - CAST_TO_DECIMAL128(25, 1, false, false), - RESCALE(26, 1, false, true), - IF_ELSE(27, 3, false, false); + /** Return the first non-null input. */ + COALESCE(0, 2, false, false), + /** Convert a nullable boolean input into an always-valid predicate. */ + PREDICATE(1, 1, false, false), + ADD(2, 2, false, false), + SUB(3, 2, false, false), + MUL(4, 2, false, false), + /** Divide without reporting arithmetic errors. */ + DIV(5, 2, false, false), + NEG(6, 1, false, false), + ABS(7, 1, false, false), + MOD(8, 2, false, false), + ADD_OVERFLOW(9, 2, true, false), + SUB_OVERFLOW(10, 2, true, false), + MUL_OVERFLOW(11, 2, true, false), + /** Divide and report division-by-zero and signed-overflow errors. */ + DIV_OVERFLOW(12, 2, true, false), + NEG_OVERFLOW(13, 1, true, false), + ABS_OVERFLOW(14, 1, true, false), + MOD_OVERFLOW(15, 2, true, false), + /** Verify that decimal input 0 fits the INT32 precision supplied by input 1. */ + CHECK_PRECISION(16, 2, true, false), + BITWISE_SHIFT_LEFT(17, 2, false, false), + BITWISE_SHIFT_RIGHT(18, 2, false, false), + CAST_TO_BOOL8(19, 1, false, false), + CAST_TO_INT8(20, 1, false, false), + CAST_TO_INT16(21, 1, false, false), + CAST_TO_INT32(22, 1, false, false), + CAST_TO_INT64(23, 1, false, false), + CAST_TO_UINT8(24, 1, false, false), + CAST_TO_UINT16(25, 1, false, false), + CAST_TO_UINT32(26, 1, false, false), + CAST_TO_UINT64(27, 1, false, false), + CAST_TO_FLOAT32(28, 1, false, false), + CAST_TO_FLOAT64(29, 1, false, false), + CAST_TO_DECIMAL32(30, 1, false, false), + CAST_TO_DECIMAL64(31, 1, false, false), + CAST_TO_DECIMAL128(32, 1, false, false), + /** Rescale a decimal input to the target scale supplied to {@link JitOperation}. */ + RESCALE(33, 1, false, true), + /** Select input 0 when input 2 is true, otherwise input 1. */ + IF_ELSE(34, 3, false, false); private final byte nativeId; private final int arity; - private final boolean supportsComplianceMode; + private final boolean fallible; private final boolean requiresTargetScale; - JitOperator(int nativeId, int arity, boolean supportsComplianceMode, boolean requiresTargetScale) { + JitOperator(int nativeId, int arity, boolean fallible, boolean requiresTargetScale) { this.nativeId = (byte) nativeId; this.arity = arity; - this.supportsComplianceMode = supportsComplianceMode; + this.fallible = fallible; this.requiresTargetScale = requiresTargetScale; assert this.nativeId == nativeId; } @@ -58,8 +78,8 @@ int getArity() { return arity; } - boolean supportsComplianceMode() { - return supportsComplianceMode; + boolean isFallible() { + return fallible; } boolean requiresTargetScale() { 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 719707a0a1c8..b8178fc4c306 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -124,7 +124,19 @@ public static Literal ofDouble(Double value) { return ofDouble(value.doubleValue()); } - /** Construct a decimal literal with the specified type and unscaled value. */ + /** + * Construct a decimal literal with the specified type and unscaled value. + * A null {@code unscaledValue} produces a null literal of the requested type. + * A {@code DECIMAL128} literal used as the root expression must be evaluated with + * {@link CompiledExpression#computeColumnJit}; the legacy executor cannot materialize it + * directly. + * + * @param type decimal storage type and scale + * @param unscaledValue unscaled decimal value, or null + * @return decimal literal + * @throws IllegalArgumentException if {@code type} is not a decimal type + * @throws ArithmeticException if {@code unscaledValue} does not fit in {@code type} + */ public static Literal ofDecimal(DType type, BigInteger unscaledValue) { if (!type.isDecimalType()) { throw new IllegalArgumentException("type is not a decimal: " + type); @@ -137,6 +149,9 @@ public static Literal ofDecimal(DType type, BigInteger unscaledValue) { } else if (type.getTypeId() == DType.DTypeEnum.DECIMAL64) { return ofLongBasedType(type, unscaledValue.longValueExact()); } else { + if (unscaledValue.bitLength() > type.getSizeInBytes() * Byte.SIZE - 1) { + throw new ArithmeticException("BigInteger out of range for " + type); + } return new Literal(type, convertDecimal128FromJavaToCudf(unscaledValue.toByteArray(), type)); } } diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index bdfa9b95441a..3794f1567899 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -198,78 +198,71 @@ cudf::ast::ast_operator jni_to_binary_operator(jbyte jni_op_value) } } -enum class jni_jit_compliance_mode : int8_t { DEFAULT, ANSI, ANSI_TRY }; +struct jni_jit_operator_info { + cudf::ast::jit::op op; + std::size_t arity; + bool is_fallible; + bool requires_target_scale; +}; /** - * Convert a Java AST serialized byte representing a JIT compliance mode. - * NOTE: This must be kept in sync with the enumeration in JitComplianceMode.java! + * Convert a serialized Java JIT operator into its libcudf definition. + * NOTE: This must be kept in sync with JitOperator.java! */ -jni_jit_compliance_mode jni_to_jit_compliance_mode(jbyte jni_mode_value) -{ - switch (jni_mode_value) { - case 0: return jni_jit_compliance_mode::DEFAULT; - case 1: return jni_jit_compliance_mode::ANSI; - case 2: return jni_jit_compliance_mode::ANSI_TRY; - default: throw std::invalid_argument("unexpected JNI AST JIT compliance mode value"); - } -} - -void expect_default_jit_compliance_mode(jni_jit_compliance_mode mode) -{ - if (mode != jni_jit_compliance_mode::DEFAULT) { - throw std::invalid_argument("unexpected compliance mode for JNI AST JIT operator"); - } -} - -void expect_non_default_jit_compliance_mode(jni_jit_compliance_mode mode) -{ - if (mode == jni_jit_compliance_mode::DEFAULT) { - throw std::invalid_argument( - "expected ANSI or ANSI_TRY compliance mode for JNI AST JIT operator"); - } -} - -cudf::error_policy jni_to_jit_error_policy(jni_jit_compliance_mode mode) +jni_jit_operator_info jni_to_jit_operator(jbyte jni_op_value) { - expect_non_default_jit_compliance_mode(mode); - return mode == jni_jit_compliance_mode::ANSI_TRY ? cudf::error_policy::NULLIFY - : cudf::error_policy::PROPAGATE; -} - -cudf::ast::expression const& dispatch_jit_expression( - cudf::ast::tree& tree, - std::vector> const& args, - jni_jit_compliance_mode mode, - cudf::ast::jit::op default_op, - cudf::ast::jit::op overflow_op) -{ - if (mode == jni_jit_compliance_mode::DEFAULT) { - return cudf::ast::jit::operation(tree, default_op, args); - } - return cudf::ast::jit::operation(tree, overflow_op, args, jni_to_jit_error_policy(mode)); -} - -void expect_jit_arity(std::vector> const& args, - std::size_t expected) -{ - if (args.size() != expected) { - throw std::invalid_argument("unexpected JNI AST JIT operator arity"); - } -} - -void expect_no_target_scale(std::optional const& target_scale) -{ - if (target_scale.has_value()) { - throw std::invalid_argument("unexpected target scale for JNI AST JIT operator"); + using enum cudf::ast::jit::op; + switch (jni_op_value) { + case 0: return {COALESCE, 2, false, false}; + case 1: return {PREDICATE, 1, false, false}; + case 2: return {ADD, 2, false, false}; + case 3: return {SUB, 2, false, false}; + case 4: return {MUL, 2, false, false}; + case 5: return {DIV, 2, false, false}; + case 6: return {NEG, 1, false, false}; + case 7: return {ABS, 1, false, false}; + case 8: return {MOD, 2, false, false}; + case 9: return {ADD_OVERFLOW, 2, true, false}; + case 10: return {SUB_OVERFLOW, 2, true, false}; + case 11: return {MUL_OVERFLOW, 2, true, false}; + case 12: return {DIV_OVERFLOW, 2, true, false}; + case 13: return {NEG_OVERFLOW, 1, true, false}; + case 14: return {ABS_OVERFLOW, 1, true, false}; + case 15: return {MOD_OVERFLOW, 2, true, false}; + case 16: return {CHECK_PRECISION, 2, true, false}; + case 17: return {BITWISE_SHIFT_LEFT, 2, false, false}; + case 18: return {BITWISE_SHIFT_RIGHT, 2, false, false}; + case 19: return {CAST_TO_BOOL8, 1, false, false}; + case 20: return {CAST_TO_INT8, 1, false, false}; + case 21: return {CAST_TO_INT16, 1, false, false}; + case 22: return {CAST_TO_INT32, 1, false, false}; + case 23: return {CAST_TO_INT64, 1, false, false}; + case 24: return {CAST_TO_UINT8, 1, false, false}; + case 25: return {CAST_TO_UINT16, 1, false, false}; + case 26: return {CAST_TO_UINT32, 1, false, false}; + case 27: return {CAST_TO_UINT64, 1, false, false}; + case 28: return {CAST_TO_FLOAT32, 1, false, false}; + case 29: return {CAST_TO_FLOAT64, 1, false, false}; + case 30: return {CAST_TO_DECIMAL32, 1, false, false}; + case 31: return {CAST_TO_DECIMAL64, 1, false, false}; + case 32: return {CAST_TO_DECIMAL128, 1, false, false}; + case 33: return {RESCALE, 1, false, true}; + case 34: return {IF_ELSE, 3, false, false}; + default: throw std::invalid_argument("unexpected JNI AST JIT operator value"); } } -int32_t require_target_scale(std::optional const& target_scale) +/** + * Convert a serialized Java JIT error policy into its libcudf value. + * NOTE: This must be kept in sync with JitErrorPolicy.java! + */ +cudf::error_policy jni_to_jit_error_policy(jbyte jni_policy_value) { - if (!target_scale.has_value()) { - throw std::invalid_argument("expected target scale for JNI AST JIT operator"); + switch (jni_policy_value) { + case 0: return cudf::error_policy::PROPAGATE; + case 1: return cudf::error_policy::NULLIFY; + default: throw std::invalid_argument("unexpected JNI AST JIT error policy value"); } - return target_scale.value(); } /** @@ -449,13 +442,26 @@ cudf::ast::operation& compile_binary_expression(cudf::jni::ast::compiled_expr& c cudf::ast::expression const& compile_jit_expression(cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { - auto const op_byte = jni_ast.read_byte(); - auto const mode = jni_to_jit_compliance_mode(jni_ast.read_byte()); - auto const arity = static_cast(jni_ast.read_byte()); + auto const op_info = jni_to_jit_operator(jni_ast.read_byte()); + auto const error_policy = jni_to_jit_error_policy(jni_ast.read_byte()); + auto const arity = static_cast(jni_ast.read_byte()); if (arity < 0) { throw std::invalid_argument("unexpected JNI AST JIT operator arity"); } + if (static_cast(arity) != op_info.arity) { + throw std::invalid_argument("unexpected JNI AST JIT operator arity"); + } + if (error_policy == cudf::error_policy::NULLIFY && !op_info.is_fallible) { + throw std::invalid_argument("unexpected error policy for JNI AST JIT operator"); + } + auto const has_target_scale = jni_ast.read_byte(); + if (has_target_scale != 0 && has_target_scale != 1) { + throw std::invalid_argument("unexpected JNI AST JIT target scale flag"); + } std::optional target_scale; - if (has_target_scale != 0) { target_scale = jni_ast.read(); } + if (has_target_scale == 1) { target_scale = jni_ast.read(); } + if (target_scale.has_value() != op_info.requires_target_scale) { + throw std::invalid_argument("unexpected target scale for JNI AST JIT operator"); + } std::vector> args; args.reserve(arity); @@ -465,152 +471,7 @@ cudf::ast::expression const& compile_jit_expression(cudf::jni::ast::compiled_exp return compiled_expr.add_jit_expression( [&](cudf::ast::tree& tree) -> cudf::ast::expression const& { - switch (op_byte) { - case 0: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::ADD, cudf::ast::jit::op::ADD_OVERFLOW); - case 1: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::SUB, cudf::ast::jit::op::SUB_OVERFLOW); - case 2: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::MUL, cudf::ast::jit::op::MUL_OVERFLOW); - case 3: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::DIV, cudf::ast::jit::op::DIV_OVERFLOW); - case 4: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::MOD, cudf::ast::jit::op::MOD_OVERFLOW); - case 5: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::ABS, cudf::ast::jit::op::ABS_OVERFLOW); - case 6: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - return dispatch_jit_expression( - tree, args, mode, cudf::ast::jit::op::NEG, cudf::ast::jit::op::NEG_OVERFLOW); - case 7: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - return cudf::ast::jit::operation( - tree, cudf::ast::jit::op::CHECK_PRECISION, args, jni_to_jit_error_policy(mode)); - case 8: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::BITWISE_SHIFT_LEFT, args); - case 9: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::BITWISE_SHIFT_RIGHT, args); - case 10: - expect_jit_arity(args, 2); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::COALESCE, args); - case 11: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::PREDICATE, args); - case 12: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_BOOL8, args); - case 13: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT8, args); - case 14: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT16, args); - case 15: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT32, args); - case 16: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_INT64, args); - case 17: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT8, args); - case 18: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT16, args); - case 19: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT32, args); - case 20: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_UINT64, args); - case 21: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_FLOAT32, args); - case 22: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_FLOAT64, args); - case 23: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_DECIMAL32, args); - case 24: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_DECIMAL64, args); - case 25: - expect_jit_arity(args, 1); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::CAST_TO_DECIMAL128, args); - case 26: - expect_jit_arity(args, 1); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, - cudf::ast::jit::op::RESCALE, - args, - cudf::error_policy::PROPAGATE, - require_target_scale(target_scale)); - case 27: - expect_jit_arity(args, 3); - expect_no_target_scale(target_scale); - expect_default_jit_compliance_mode(mode); - return cudf::ast::jit::operation(tree, cudf::ast::jit::op::IF_ELSE, args); - default: throw std::invalid_argument("unexpected JNI AST JIT operator value"); - } + return cudf::ast::jit::operation(tree, op_info.op, args, error_policy, target_scale); }); } @@ -647,6 +508,19 @@ std::unique_ptr compile_serialized_ast(jni_serial return jni_expr_ptr; } +enum class execution_backend { DEFAULT, JIT }; + +jlong execute_compiled_expression(jlong j_ast, jlong j_table, execution_backend backend) +{ + auto compiled_expr_ptr = reinterpret_cast(j_ast); + auto tview_ptr = reinterpret_cast(j_table); + std::unique_ptr result = + backend == execution_backend::JIT + ? cudf::compute_column_jit(*tview_ptr, compiled_expr_ptr->get_top_expression()) + : cudf::compute_column(*tview_ptr, compiled_expr_ptr->get_top_expression()); + return reinterpret_cast(result.release()); +} + } // anonymous namespace extern "C" { @@ -678,11 +552,22 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ast_CompiledExpression_computeColumn JNI_TRY { cudf::jni::auto_set_device(env); - auto compiled_expr_ptr = reinterpret_cast(j_ast); - auto tview_ptr = reinterpret_cast(j_table); - std::unique_ptr result = - cudf::compute_column(*tview_ptr, compiled_expr_ptr->get_top_expression()); - return reinterpret_cast(result.release()); + 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); } JNI_CATCH(env, 0); } diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 38aecf64672e..3592a0394bd1 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #include "cudf_jni_apis.hpp" +#include #include #include #include @@ -191,6 +192,16 @@ JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_Cuda_isPtdsEnabled(JNIEnv* env, j return cudf::jni::is_ptds_enabled; } +JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cudf_initializeJitRuntime(JNIEnv* env, jclass) +{ + JNI_TRY + { + cudf::jni::auto_set_device(env); + cudf::initialize(cudf::init_flags::INIT_JIT_CACHE); + } + JNI_CATCH(env, ); +} + JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cudf_setKernelPinnedCopyThreshold(JNIEnv* env, jclass clazz, jlong jthreshold) diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index 7d0e7f6214f2..ed4f6990ebbf 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ 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 7a0b03fb8576..e8a3d1a1cff6 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -6,9 +6,11 @@ package ai.rapids.cudf.ast; import ai.rapids.cudf.ColumnVector; +import ai.rapids.cudf.Cudf; import ai.rapids.cudf.CudfException; import ai.rapids.cudf.CudfTestBase; import ai.rapids.cudf.DType; +import ai.rapids.cudf.Scalar; import ai.rapids.cudf.Table; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -27,9 +29,14 @@ import java.util.stream.Stream; import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual; -import static org.junit.jupiter.api.Assumptions.assumeTrue; public class CompiledExpressionTest extends CudfTestBase { + @Test + public void testInitializeJitRuntime() { + Assertions.assertDoesNotThrow(Cudf::initializeJitRuntime); + Assertions.assertDoesNotThrow(Cudf::initializeJitRuntime); + } + @Test public void testColumnReferenceTransform() { try (Table t = new Table.TestBuilder().column(5, 4, 3, 2, 1).column(6, 7, 8, null, 10).build()) { @@ -56,6 +63,7 @@ public void testInvalidColumnReferenceTransform() { try (Table t = new Table.TestBuilder().column(5, 4, 3, 2, 1).column(6, 7, 8, null, 10).build(); CompiledExpression compiledExpr = expr.compile()) { Assertions.assertThrows(CudfException.class, () -> compiledExpr.computeColumn(t).close()); + Assertions.assertThrows(CudfException.class, () -> compiledExpr.computeColumnJit(t).close()); } } @@ -297,19 +305,57 @@ public void testDurationNanoSecondsLiteralTransform(Long value) { } } - @Test - public void testDecimal128LiteralTransform() { - int scale = -4; - BigInteger value = new BigInteger("-123456789012345678901234567890"); - Literal expr = Literal.ofDecimal(DType.create(DType.DTypeEnum.DECIMAL128, scale), value); + private static Stream createDecimalLiteralParams() { + return Stream.of( + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, -2), + new BigInteger("1234567")), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, 0), (BigInteger) null), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL64, -4), + new BigInteger("-123456789012345678")), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL64, -18), (BigInteger) null), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, -4), + new BigInteger("-123456789012345678901234567890")), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, -38), BigInteger.ONE), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, 0), + BigInteger.ONE.shiftLeft(127).subtract(BigInteger.ONE)), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, 0), + BigInteger.ONE.shiftLeft(127).negate()), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, -10), (BigInteger) null)); + } + + @ParameterizedTest + @MethodSource("createDecimalLiteralParams") + 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.computeColumn(t); - ColumnVector expected = ColumnVector.decimalFromBigInt(scale, value, value, value)) { + ColumnVector actual = compiledExpr.computeColumnJit(t); + Scalar expectedScalar = value == null ? + Scalar.fromNull(type) : Scalar.fromDecimal(value, type); + ColumnVector expected = ColumnVector.fromScalar(expectedScalar, 3)) { assertColumnsAreEqual(expected, actual); } } + @Test + public void testDecimalLiteralValidation() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> Literal.ofDecimal(DType.INT32, BigInteger.ONE)); + Assertions.assertThrows(ArithmeticException.class, + () -> Literal.ofDecimal(DType.create(DType.DTypeEnum.DECIMAL32, 0), + BigInteger.ONE.shiftLeft(31))); + Assertions.assertThrows(ArithmeticException.class, + () -> Literal.ofDecimal(DType.create(DType.DTypeEnum.DECIMAL64, 0), + BigInteger.ONE.shiftLeft(63))); + + DType decimal128 = DType.create(DType.DTypeEnum.DECIMAL128, 0); + Assertions.assertThrows(ArithmeticException.class, + () -> Literal.ofDecimal(decimal128, BigInteger.ONE.shiftLeft(127))); + Assertions.assertThrows(ArithmeticException.class, + () -> Literal.ofDecimal(decimal128, + BigInteger.ONE.shiftLeft(127).negate().subtract(BigInteger.ONE))); + } + @Test void testJitOperationValidation() { AstExpression[] inputs = new AstExpression[] { @@ -318,14 +364,21 @@ void testJitOperationValidation() { new ColumnReference(2) }; for (JitOperator op : JitOperator.values()) { + AstExpression[] opInputs = Arrays.copyOf(inputs, op.getArity()); + JitOperation expr; if (op == JitOperator.RESCALE) { - Assertions.assertDoesNotThrow( - () -> new JitOperation(op, -2, Arrays.copyOf(inputs, op.getArity()))); + expr = new JitOperation(op, -2, opInputs); } else { - JitComplianceMode mode = op == JitOperator.PRECISION_CHECK - ? JitComplianceMode.ANSI : JitComplianceMode.DEFAULT; - Assertions.assertDoesNotThrow( - () -> new JitOperation(op, mode, Arrays.copyOf(inputs, op.getArity()))); + expr = new JitOperation(op, opInputs); + } + try (CompiledExpression ignored = expr.compile()) { + // The native decoder is part of the serialized operator contract. + } + if (op.isFallible()) { + JitOperation nullifyingExpr = new JitOperation(op, JitErrorPolicy.NULLIFY, opInputs); + try (CompiledExpression ignored = nullifyingExpr.compile()) { + // The native decoder is part of the serialized error-policy contract. + } } if (op.getArity() > 0) { Assertions.assertThrows( @@ -336,36 +389,160 @@ void testJitOperationValidation() { Assertions.assertThrows( IllegalArgumentException.class, - () -> new JitOperation(JitOperator.PRECISION_CHECK, + () -> new JitOperation(JitOperator.ADD, JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1))); Assertions.assertThrows( IllegalArgumentException.class, - () -> new JitOperation(JitOperator.IF_ELSE, JitComplianceMode.ANSI, - new ColumnReference(0), new ColumnReference(1), new ColumnReference(2))); + () -> new JitOperation(JitOperator.ADD, -2, + new ColumnReference(0), new ColumnReference(1))); Assertions.assertThrows( IllegalArgumentException.class, () -> new JitOperation(JitOperator.RESCALE, new ColumnReference(0))); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1), new ColumnReference(2))); + Assertions.assertThrows( + NullPointerException.class, + () -> new JitOperation(null, new ColumnReference(0))); + Assertions.assertThrows( + NullPointerException.class, + () -> new JitOperation(JitOperator.ADD, (JitErrorPolicy) null, + new ColumnReference(0), new ColumnReference(1))); + Assertions.assertThrows( + NullPointerException.class, + () -> new JitOperation(JitOperator.ADD, + new ColumnReference(0), (AstExpression) null)); + Assertions.assertThrows( + NullPointerException.class, + () -> new JitOperation(JitOperator.ADD, (AstExpression[]) null)); + } + + @Test + void testJitNestedArithmeticTransform() { + AstExpression expr = new JitOperation(JitOperator.ADD, new ColumnReference(0), Literal.ofInt(2)); + expr = new JitOperation(JitOperator.SUB, expr, Literal.ofInt(1)); + expr = new JitOperation(JitOperator.MUL, expr, Literal.ofInt(3)); + expr = new JitOperation(JitOperator.DIV, expr, Literal.ofInt(2)); + expr = new JitOperation(JitOperator.MOD, expr, Literal.ofInt(5)); + expr = new JitOperation(JitOperator.NEG, expr); + expr = new JitOperation(JitOperator.ABS, expr); + expr = new JitOperation(JitOperator.BITWISE_SHIFT_LEFT, expr, Literal.ofInt(2)); + 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); + ColumnVector expected = ColumnVector.fromInts(6, 8, 2, 4)) { + assertColumnsAreEqual(expected, actual); + } + } + + @Test + void testJitOverflowPolicies() { + try (Table t = new Table.TestBuilder() + .column(1, 3) + .column(10, 7) + .column(10, Integer.MAX_VALUE) + .build()) { + JitOperation successExpr = new JitOperation(JitOperator.ADD_OVERFLOW, + new ColumnReference(0), new ColumnReference(1)); + try (CompiledExpression compiledExpr = successExpr.compile(); + ColumnVector actual = compiledExpr.computeColumnJit(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()) { + Assertions.assertThrows(CudfException.class, + () -> compiledExpr.computeColumnJit(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); + ColumnVector expected = ColumnVector.fromBoxedInts(11, null)) { + assertColumnsAreEqual(expected, actual); + } + } + } + + @Test + void testJitFusedNullifyingOverflowTransform() { + try (Table t = new Table.TestBuilder() + .column(1, 3, 20, 1, 50, 10) + .column(1, 10, 7, 20, Integer.MAX_VALUE, 2) + .column(1, 5, 4, Integer.MAX_VALUE, 2, 5) + .column(0, 1, 0, 0, 1, 5) + .build()) { + AstExpression expr = new JitOperation(JitOperator.ADD_OVERFLOW, JitErrorPolicy.NULLIFY, + new ColumnReference(0), new ColumnReference(1)); + expr = new JitOperation(JitOperator.MUL_OVERFLOW, JitErrorPolicy.NULLIFY, + 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); + ColumnVector expected = ColumnVector.fromBoxedInts(null, 65, null, null, null, 12)) { + assertColumnsAreEqual(expected, actual); + } + } + } + + @Test + void testJitUnaryAndSubtractOverflowTransform() { + try (Table t = new Table.TestBuilder() + .column(10, Integer.MIN_VALUE, 1) + .column(3, 1, 0) + .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); + 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); + 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); + ColumnVector expected = ColumnVector.fromBoxedInts(10, null, 1)) { + assertColumnsAreEqual(expected, actual); + } + } } @Test void testJitTryDivModTransform() { - assumeTrue("1".equals(System.getenv("LIBCUDF_JIT_ENABLED"))); try (Table t = new Table.TestBuilder() .column(10, 7, null, 6) .column(2, 0, 3, null) .build()) { - JitOperation divExpr = new JitOperation(JitOperator.DIV, JitComplianceMode.ANSI_TRY, - new ColumnReference(0), new ColumnReference(1)); + JitOperation divExpr = new JitOperation(JitOperator.DIV_OVERFLOW, + JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); try (CompiledExpression compiledExpr = divExpr.compile(); - ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector actual = compiledExpr.computeColumnJit(t); ColumnVector expected = ColumnVector.fromBoxedInts(5, null, null, null)) { assertColumnsAreEqual(expected, actual); } - JitOperation modExpr = new JitOperation(JitOperator.MOD, JitComplianceMode.ANSI_TRY, - new ColumnReference(0), new ColumnReference(1)); + JitOperation modExpr = new JitOperation(JitOperator.MOD_OVERFLOW, + JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); try (CompiledExpression compiledExpr = modExpr.compile(); - ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector actual = compiledExpr.computeColumnJit(t); ColumnVector expected = ColumnVector.fromBoxedInts(0, null, null, null)) { assertColumnsAreEqual(expected, actual); } @@ -373,18 +550,144 @@ void testJitTryDivModTransform() { } @Test - void testJitIfElseTransform() { - assumeTrue("1".equals(System.getenv("LIBCUDF_JIT_ENABLED"))); + void testJitMixedConditionalTransform() { try (Table t = new Table.TestBuilder() - .column(1, 2, 3, 4) + .column(1, null, 3, null) .column(10, 20, 30, 40) - .column(true, false, true, false) .build()) { + AstExpression condition = new BinaryOperation(BinaryOperator.GREATER, + new ColumnReference(0), Literal.ofInt(2)); + AstExpression predicate = new JitOperation(JitOperator.PREDICATE, condition); + AstExpression coalesced = new JitOperation(JitOperator.COALESCE, + new ColumnReference(0), Literal.ofInt(99)); JitOperation expr = new JitOperation(JitOperator.IF_ELSE, - new ColumnReference(0), new ColumnReference(1), new ColumnReference(2)); + coalesced, new ColumnReference(1), predicate); try (CompiledExpression compiledExpr = expr.compile(); - ColumnVector actual = compiledExpr.computeColumn(t); - ColumnVector expected = ColumnVector.fromBoxedInts(1, 20, 3, 40)) { + ColumnVector actual = compiledExpr.computeColumnJit(t); + ColumnVector expected = ColumnVector.fromBoxedInts(10, 20, 3, 40)) { + assertColumnsAreEqual(expected, actual); + } + } + } + + private static Stream createJitNumericCastParams() { + return Stream.of( + JitOperator.CAST_TO_BOOL8, + JitOperator.CAST_TO_INT8, + JitOperator.CAST_TO_INT16, + JitOperator.CAST_TO_INT32, + JitOperator.CAST_TO_INT64, + JitOperator.CAST_TO_UINT8, + JitOperator.CAST_TO_UINT16, + JitOperator.CAST_TO_UINT32, + JitOperator.CAST_TO_UINT64, + JitOperator.CAST_TO_FLOAT32, + JitOperator.CAST_TO_FLOAT64); + } + + private static ColumnVector makeJitNumericCastExpected(JitOperator op) { + switch (op) { + case CAST_TO_BOOL8: + return ColumnVector.fromBooleans(false, true, true, true); + case CAST_TO_INT8: + return ColumnVector.fromBytes((byte) 0, (byte) 1, (byte) 2, (byte) 3); + case CAST_TO_INT16: + return ColumnVector.fromShorts((short) 0, (short) 1, (short) 2, (short) 3); + case CAST_TO_INT32: + return ColumnVector.fromInts(0, 1, 2, 3); + case CAST_TO_INT64: + return ColumnVector.fromLongs(0L, 1L, 2L, 3L); + case CAST_TO_UINT8: + return ColumnVector.fromUnsignedBytes((byte) 0, (byte) 1, (byte) 2, (byte) 3); + case CAST_TO_UINT16: + return ColumnVector.fromUnsignedShorts((short) 0, (short) 1, (short) 2, (short) 3); + case CAST_TO_UINT32: + return ColumnVector.fromUnsignedInts(0, 1, 2, 3); + case CAST_TO_UINT64: + return ColumnVector.fromUnsignedLongs(0L, 1L, 2L, 3L); + case CAST_TO_FLOAT32: + return ColumnVector.fromFloats(0.0f, 1.0f, 2.0f, 3.0f); + case CAST_TO_FLOAT64: + return ColumnVector.fromDoubles(0.0, 1.0, 2.0, 3.0); + default: + throw new IllegalArgumentException("Unexpected numeric cast operator " + op); + } + } + + @ParameterizedTest + @MethodSource("createJitNumericCastParams") + void testJitNumericCastTransform(JitOperator op) { + 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); + ColumnVector expected = makeJitNumericCastExpected(op)) { + assertColumnsAreEqual(expected, actual); + } + } + + private static Stream createJitDecimalCastParams() { + return Stream.of( + JitOperator.CAST_TO_DECIMAL32, + JitOperator.CAST_TO_DECIMAL64, + JitOperator.CAST_TO_DECIMAL128); + } + + private static ColumnVector makeJitDecimalCastExpected(JitOperator op) { + switch (op) { + case CAST_TO_DECIMAL32: + return ColumnVector.decimalFromInts(0, 0, 1, -2, 3); + case CAST_TO_DECIMAL64: + return ColumnVector.decimalFromLongs(0, 0L, 1L, -2L, 3L); + case CAST_TO_DECIMAL128: + return ColumnVector.decimalFromBigInt(0, + BigInteger.ZERO, BigInteger.ONE, BigInteger.valueOf(-2), BigInteger.valueOf(3)); + default: + throw new IllegalArgumentException("Unexpected decimal cast operator " + op); + } + } + + @ParameterizedTest + @MethodSource("createJitDecimalCastParams") + void testJitDecimalCastTransform(JitOperator op) { + 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); + ColumnVector expected = makeJitDecimalCastExpected(op)) { + assertColumnsAreEqual(expected, actual); + } + } + + @Test + void testJitDecimalRescaleTransform() { + JitOperation expr = new JitOperation(JitOperator.RESCALE, -2, new ColumnReference(0)); + try (Table t = new Table.TestBuilder() + .decimal32Column(0, 123, 1234, 12345, 123456, 1234567) + .build(); + CompiledExpression compiledExpr = expr.compile(); + ColumnVector actual = compiledExpr.computeColumnJit(t); + ColumnVector expected = ColumnVector.decimalFromInts( + -2, 12300, 123400, 1234500, 12345600, 123456700)) { + assertColumnsAreEqual(expected, actual); + } + } + + @Test + 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()) { + Assertions.assertThrows(CudfException.class, + () -> compiledExpr.computeColumnJit(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); + ColumnVector expected = ColumnVector.decimalFromBoxedInts(0, 3, 200, 250, null)) { assertColumnsAreEqual(expected, actual); } } From 4680b9aefd26d360eec9eb44a9be2289b03f7f59 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 7 Jul 2026 15:37:20 +0800 Subject: [PATCH 03/11] local review Signed-off-by: Haoyang Li --- java/src/main/java/ai/rapids/cudf/ast/Literal.java | 9 +++------ java/src/main/native/src/CompiledExpression.cpp | 2 +- .../java/ai/rapids/cudf/ast/CompiledExpressionTest.java | 1 + 3 files changed, 5 insertions(+), 7 deletions(-) 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 b8178fc4c306..f62e3f8f6487 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -278,10 +278,9 @@ private int getDataTypeSerializedSize() { int nativeTypeId = type.getTypeId().getNativeId(); assert nativeTypeId == (byte) nativeTypeId : "Type ID does not fit in a byte"; if (type.isDecimalType()) { - assert type.getScale() == (byte) type.getScale() : "Decimal scale does not fit in a byte"; - return 2; + return Byte.BYTES + Integer.BYTES; } - return 1; + return Byte.BYTES; } private void serializeDataType(ByteBuffer bb) { @@ -289,9 +288,7 @@ private void serializeDataType(ByteBuffer bb) { assert nativeTypeId == type.getTypeId().getNativeId() : "DType ID does not fit in a byte"; bb.put(nativeTypeId); if (type.isDecimalType()) { - byte scale = (byte) type.getScale(); - assert scale == (byte) type.getScale() : "Decimal scale does not fit in a byte"; - bb.put(scale); + bb.putInt(type.getScale()); } } diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index 3794f1567899..bc73c095bab5 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -103,7 +103,7 @@ class jni_serialized_ast { case cudf::type_id::DECIMAL32: case cudf::type_id::DECIMAL64: case cudf::type_id::DECIMAL128: { - int32_t const scale = read_byte(); + int32_t const scale = read(); return cudf::data_type(dtype_id, scale); } default: throw new std::invalid_argument("unrecognized cudf data type"); 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 e8a3d1a1cff6..df735b3eaf81 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -309,6 +309,7 @@ private static Stream createDecimalLiteralParams() { return Stream.of( Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, -2), new BigInteger("1234567")), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, 128), BigInteger.ONE), Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, 0), (BigInteger) null), Arguments.of(DType.create(DType.DTypeEnum.DECIMAL64, -4), new BigInteger("-123456789012345678")), From 6f79a23feb1aa18cd336087e399eb819cbdef9b1 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 7 Jul 2026 17:18:01 +0800 Subject: [PATCH 04/11] local review comments address Signed-off-by: Haoyang Li --- .../main/native/src/CompiledExpression.cpp | 98 +++++++++---------- .../src/main/native/src/jni_compiled_expr.hpp | 57 +++++------ .../cudf/ast/CompiledExpressionTest.java | 42 +++++++- 3 files changed, 109 insertions(+), 88 deletions(-) diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index bc73c095bab5..eebcb37b4afc 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -106,14 +107,14 @@ class jni_serialized_ast { int32_t const scale = read(); return cudf::data_type(dtype_id, scale); } - default: throw new std::invalid_argument("unrecognized cudf data type"); + default: throw std::invalid_argument("unrecognized cudf data type"); } } }; /** * Enumeration of the AST expression types that can appear in the serialized data. - * NOTE: This must be kept in sync with the NodeType enumeration in AstNode.java! + * NOTE: This must be kept in sync with the ExpressionType enumeration in AstExpression.java! */ enum class jni_serialized_expression_type : int8_t { VALID_LITERAL = 0, @@ -283,10 +284,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& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { std::unique_ptr scalar_ptr = cudf::make_numeric_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -297,16 +298,15 @@ struct make_literal { } auto& numeric_scalar = static_cast&>(*scalar_ptr); - return compiled_expr.add_literal(std::make_unique(numeric_scalar), - std::move(scalar_ptr)); + return compiled_expr.add_literal(numeric_scalar, std::move(scalar_ptr)); } /** Construct an AST literal from a timestamp value */ template ()>* = nullptr> - cudf::ast::literal& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { std::unique_ptr scalar_ptr = cudf::make_timestamp_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -317,16 +317,15 @@ struct make_literal { } auto& timestamp_scalar = static_cast&>(*scalar_ptr); - return compiled_expr.add_literal(std::make_unique(timestamp_scalar), - std::move(scalar_ptr)); + return compiled_expr.add_literal(timestamp_scalar, std::move(scalar_ptr)); } /** Construct an AST literal from a duration value */ template ()>* = nullptr> - cudf::ast::literal& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { std::unique_ptr scalar_ptr = cudf::make_duration_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -337,16 +336,15 @@ struct make_literal { } auto& duration_scalar = static_cast&>(*scalar_ptr); - return compiled_expr.add_literal(std::make_unique(duration_scalar), - std::move(scalar_ptr)); + return compiled_expr.add_literal(duration_scalar, std::move(scalar_ptr)); } /** Construct an AST literal from a string value */ template >* = nullptr> - cudf::ast::literal& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { std::unique_ptr scalar_ptr = [&]() { if (is_valid) { @@ -358,8 +356,7 @@ struct make_literal { }(); auto& str_scalar = static_cast(*scalar_ptr); - return compiled_expr.add_literal(std::make_unique(str_scalar), - std::move(scalar_ptr)); + return compiled_expr.add_literal(str_scalar, std::move(scalar_ptr)); } /** Default functor implementation to catch type dispatch errors */ @@ -367,20 +364,20 @@ 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& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { throw std::logic_error("Unsupported AST literal type"); } /** Construct an AST literal from a fixed-point value */ template ()>* = nullptr> - cudf::ast::literal& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) { using rep_type = typename T::rep; auto const val = is_valid ? jni_ast.read() : rep_type{}; @@ -389,28 +386,26 @@ struct make_literal { scalar_ptr->set_valid_async(is_valid); auto& fixed_point_scalar = static_cast&>(*scalar_ptr); - return compiled_expr.add_literal(std::make_unique(fixed_point_scalar), - std::move(scalar_ptr)); + return compiled_expr.add_literal(fixed_point_scalar, std::move(scalar_ptr)); } }; /** Decode a serialized AST literal */ -cudf::ast::literal& 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::ast::column_reference& compile_column_reference(cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) +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()); cudf::size_type const column_index = jni_ast.read(); - return compiled_expr.add_column_ref( - std::make_unique(column_index, table_ref)); + return compiled_expr.add_column_ref(column_index, table_ref); } // forward declaration @@ -418,24 +413,22 @@ cudf::ast::expression const& compile_expression(cudf::jni::ast::compiled_expr& c jni_serialized_ast& jni_ast); /** Decode a serialized AST unary expression */ -cudf::ast::operation& 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()); cudf::ast::expression const& child_expression = compile_expression(compiled_expr, jni_ast); - return compiled_expr.add_operation( - std::make_unique(ast_op, child_expression)); + return compiled_expr.add_operation(ast_op, child_expression); } /** Decode a serialized AST binary expression */ -cudf::ast::operation& 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()); 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( - std::make_unique(ast_op, left_child, right_child)); + return compiled_expr.add_operation(ast_op, left_child, right_child); } /** Decode a serialized JIT AST expression */ @@ -505,6 +498,9 @@ std::unique_ptr compile_serialized_ast(jni_serial 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(); } + return jni_expr_ptr; } diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index ed4f6990ebbf..4244cebef59c 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -16,61 +16,52 @@ namespace cudf { namespace jni { namespace ast { -/** - * A class to capture all of the resources associated with a compiled AST expression. - * AST nodes do not own their child nodes, so every node in the expression tree - * must be explicitly tracked in order to free the underlying resources for each node. - * - * This should be cleaned up a bit after the libcudf AST refactoring in - * https://github.com/rapidsai/cudf/pull/8815 when a virtual destructor is added to the - * base AST node type. Then we do not have to track every AST node type separately. - */ +/** A class to capture all resources associated with a compiled AST expression. */ class compiled_expr { - /** All expression nodes within the expression tree */ - std::vector> expressions; - - /** Expression tree for JIT helper-created nodes */ - cudf::ast::tree jit_expressions; - /** GPU scalar instances that correspond to literal nodes */ std::vector> scalars; - cudf::ast::expression const* top_expression = nullptr; + /** All expression nodes within the expression tree */ + cudf::ast::tree expressions; public: - cudf::ast::literal& add_literal(std::unique_ptr literal_ptr, - std::unique_ptr scalar_ptr) + template + cudf::ast::literal const& add_literal(ScalarType& scalar, + std::unique_ptr scalar_ptr) { - expressions.push_back(std::move(literal_ptr)); scalars.push_back(std::move(scalar_ptr)); - top_expression = expressions.back().get(); - return static_cast(*expressions.back()); + return expressions.emplace(scalar); } - cudf::ast::column_reference& add_column_ref(std::unique_ptr ref_ptr) + cudf::ast::column_reference const& add_column_ref(cudf::size_type column_index, + cudf::ast::table_reference table_ref) { - expressions.push_back(std::move(ref_ptr)); - top_expression = expressions.back().get(); - return static_cast(*expressions.back()); + return expressions.emplace(column_index, table_ref); } - cudf::ast::operation& add_operation(std::unique_ptr expr_ptr) + cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, + cudf::ast::expression const& child) { - expressions.push_back(std::move(expr_ptr)); - top_expression = expressions.back().get(); - return static_cast(*expressions.back()); + return expressions.emplace(op, child); + } + + 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, right); } template cudf::ast::expression const& add_jit_expression(F&& factory) { - auto& expr = factory(jit_expressions); - top_expression = &expr; - return expr; + 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 *top_expression; } + cudf::ast::expression const& get_top_expression() const { 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 df735b3eaf81..fbcb0fc95322 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -419,6 +419,29 @@ void testJitOperationValidation() { () -> new JitOperation(JitOperator.ADD, (AstExpression[]) null)); } + @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()) { + Assertions.assertThrows(CudfException.class, + () -> compiledExpr.computeColumnJit(t).close()); + } + } + + @Test + 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); + ColumnVector expected = ColumnVector.fromInts()) { + assertColumnsAreEqual(expected, actual); + } + } + @Test void testJitNestedArithmeticTransform() { AstExpression expr = new JitOperation(JitOperator.ADD, new ColumnReference(0), Literal.ofInt(2)); @@ -439,6 +462,17 @@ void testJitNestedArithmeticTransform() { } } + @Test + 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); + ColumnVector expected = ColumnVector.fromInts(5, 0, -7)) { + assertColumnsAreEqual(expected, actual); + } + } + @Test void testJitOverflowPolicies() { try (Table t = new Table.TestBuilder() @@ -529,14 +563,14 @@ void testJitUnaryAndSubtractOverflowTransform() { @Test void testJitTryDivModTransform() { try (Table t = new Table.TestBuilder() - .column(10, 7, null, 6) - .column(2, 0, 3, null) + .column(10, 7, null, 6, Integer.MIN_VALUE) + .column(2, 0, 3, null, -1) .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); - ColumnVector expected = ColumnVector.fromBoxedInts(5, null, null, null)) { + ColumnVector expected = ColumnVector.fromBoxedInts(5, null, null, null, null)) { assertColumnsAreEqual(expected, actual); } @@ -544,7 +578,7 @@ void testJitTryDivModTransform() { JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); try (CompiledExpression compiledExpr = modExpr.compile(); ColumnVector actual = compiledExpr.computeColumnJit(t); - ColumnVector expected = ColumnVector.fromBoxedInts(0, null, null, null)) { + ColumnVector expected = ColumnVector.fromBoxedInts(0, null, null, null, 0)) { assertColumnsAreEqual(expected, actual); } } From 19f974b998732971cd655908813b25bc8b1e789c Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Thu, 9 Jul 2026 13:19:25 +0800 Subject: [PATCH 05/11] address code rabbit comments Signed-off-by: Haoyang Li --- .../main/java/ai/rapids/cudf/ast/Literal.java | 7 ++--- .../cudf/ast/CompiledExpressionTest.java | 26 +++++++++++++++++-- 2 files changed, 28 insertions(+), 5 deletions(-) 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 f62e3f8f6487..54547fc16d3b 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -127,9 +127,10 @@ 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. - * A {@code DECIMAL128} literal used as the root expression must be evaluated with - * {@link CompiledExpression#computeColumnJit}; the legacy executor cannot materialize it - * directly. + * 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 directly. * * @param type decimal storage type and scale * @param unscaledValue unscaled decimal value, or null 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 a464dddb9e0c..f20ea6b37393 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -305,7 +305,7 @@ public void testDurationNanoSecondsLiteralTransform(Long value) { } } - private static Stream createDecimalLiteralParams() { + private static Stream createLegacyDecimalLiteralParams() { return Stream.of( Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, -2), new BigInteger("1234567")), @@ -313,7 +313,11 @@ private static Stream createDecimalLiteralParams() { Arguments.of(DType.create(DType.DTypeEnum.DECIMAL32, 0), (BigInteger) null), Arguments.of(DType.create(DType.DTypeEnum.DECIMAL64, -4), new BigInteger("-123456789012345678")), - Arguments.of(DType.create(DType.DTypeEnum.DECIMAL64, -18), (BigInteger) null), + Arguments.of(DType.create(DType.DTypeEnum.DECIMAL64, -18), (BigInteger) null)); + } + + private static Stream createDecimal128LiteralParams() { + return Stream.of( Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, -4), new BigInteger("-123456789012345678901234567890")), Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, -38), BigInteger.ONE), @@ -324,6 +328,24 @@ private static Stream createDecimalLiteralParams() { Arguments.of(DType.create(DType.DTypeEnum.DECIMAL128, -10), (BigInteger) null)); } + private static Stream createDecimalLiteralParams() { + return Stream.concat(createLegacyDecimalLiteralParams(), createDecimal128LiteralParams()); + } + + @ParameterizedTest + @MethodSource("createLegacyDecimalLiteralParams") + public void testDecimalLiteralTransform(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.computeColumn(t); + Scalar expectedScalar = value == null ? + Scalar.fromNull(type) : Scalar.fromDecimal(value, type); + ColumnVector expected = ColumnVector.fromScalar(expectedScalar, 3)) { + assertColumnsAreEqual(expected, actual); + } + } + @ParameterizedTest @MethodSource("createDecimalLiteralParams") public void testJitDecimalLiteralTransform(DType type, BigInteger value) { From 30cceaf4e20084376ab42af3a605f5192411d307 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 10 Jul 2026 11:18:47 +0800 Subject: [PATCH 06/11] address comments Signed-off-by: Haoyang Li --- .../java/ai/rapids/cudf/ast/JitOperation.java | 18 +--- .../java/ai/rapids/cudf/ast/JitOperator.java | 90 ++++++++----------- .../cudf/ast/CompiledExpressionTest.java | 58 ++++-------- 3 files changed, 54 insertions(+), 112 deletions(-) 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 2d2f99dbb0ce..867336efaeaa 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,10 @@ import java.util.Objects; /** - * A libcudf row-IR operation. Expressions containing a JIT operation must be evaluated with + * A libcudf JIT operation. Expressions containing a JIT operation must be evaluated with * {@link CompiledExpression#computeColumnJit}. + * Operator arity, error policy, and target-scale constraints are validated when the expression + * is compiled. */ public final class JitOperation extends AstExpression { private final JitOperator op; @@ -24,7 +26,6 @@ public final class JitOperation extends AstExpression { * @param op operator to apply * @param inputs operator inputs * @throws NullPointerException if {@code op}, {@code inputs}, or an input is null - * @throws IllegalArgumentException if the operator arity or target-scale usage is invalid */ public JitOperation(JitOperator op, AstExpression... inputs) { this(op, JitErrorPolicy.PROPAGATE, null, inputs); @@ -37,7 +38,6 @@ public JitOperation(JitOperator op, AstExpression... inputs) { * @param errorPolicy error handling policy * @param inputs operator inputs * @throws NullPointerException if any argument or input is null - * @throws IllegalArgumentException if the arity, policy, or target-scale usage is invalid */ public JitOperation(JitOperator op, JitErrorPolicy errorPolicy, AstExpression... inputs) { this(op, errorPolicy, null, inputs); @@ -51,8 +51,6 @@ public JitOperation(JitOperator op, JitErrorPolicy errorPolicy, AstExpression... * @param targetScale target fixed-point scale * @param inputs operator inputs * @throws NullPointerException if {@code op}, {@code inputs}, or an input is null - * @throws IllegalArgumentException if the operator is not {@link JitOperator#RESCALE} or its - * arity is invalid */ public JitOperation(JitOperator op, int targetScale, AstExpression... inputs) { this(op, JitErrorPolicy.PROPAGATE, Integer.valueOf(targetScale), inputs); @@ -67,16 +65,6 @@ private JitOperation( this.errorPolicy = Objects.requireNonNull(errorPolicy, "errorPolicy is null"); this.inputs = Objects.requireNonNull(inputs, "inputs is null").clone(); this.targetScale = targetScale; - if (this.inputs.length != op.getArity()) { - throw new IllegalArgumentException( - op + " requires " + op.getArity() + " inputs, found " + this.inputs.length); - } - if (!op.isFallible() && errorPolicy == JitErrorPolicy.NULLIFY) { - throw new IllegalArgumentException(op + " cannot nullify errors"); - } - if (op.requiresTargetScale() != (targetScale != null)) { - throw new IllegalArgumentException(op + " target scale usage is invalid"); - } for (AstExpression input : this.inputs) { Objects.requireNonNull(input, "input is null"); } diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java index 661a75689031..16419ea8956a 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java @@ -19,73 +19,55 @@ */ public enum JitOperator { /** Return the first non-null input. */ - COALESCE(0, 2, false, false), + COALESCE(0), /** Convert a nullable boolean input into an always-valid predicate. */ - PREDICATE(1, 1, false, false), - ADD(2, 2, false, false), - SUB(3, 2, false, false), - MUL(4, 2, false, false), + PREDICATE(1), + ADD(2), + SUB(3), + MUL(4), /** Divide without reporting arithmetic errors. */ - DIV(5, 2, false, false), - NEG(6, 1, false, false), - ABS(7, 1, false, false), - MOD(8, 2, false, false), - ADD_OVERFLOW(9, 2, true, false), - SUB_OVERFLOW(10, 2, true, false), - MUL_OVERFLOW(11, 2, true, false), + DIV(5), + NEG(6), + ABS(7), + MOD(8), + ADD_OVERFLOW(9), + SUB_OVERFLOW(10), + MUL_OVERFLOW(11), /** Divide and report division-by-zero and signed-overflow errors. */ - DIV_OVERFLOW(12, 2, true, false), - NEG_OVERFLOW(13, 1, true, false), - ABS_OVERFLOW(14, 1, true, false), - MOD_OVERFLOW(15, 2, true, false), + DIV_OVERFLOW(12), + NEG_OVERFLOW(13), + ABS_OVERFLOW(14), + MOD_OVERFLOW(15), /** Verify that decimal input 0 fits the INT32 precision supplied by input 1. */ - CHECK_PRECISION(16, 2, true, false), - BITWISE_SHIFT_LEFT(17, 2, false, false), - BITWISE_SHIFT_RIGHT(18, 2, false, false), - CAST_TO_BOOL8(19, 1, false, false), - CAST_TO_INT8(20, 1, false, false), - CAST_TO_INT16(21, 1, false, false), - CAST_TO_INT32(22, 1, false, false), - CAST_TO_INT64(23, 1, false, false), - CAST_TO_UINT8(24, 1, false, false), - CAST_TO_UINT16(25, 1, false, false), - CAST_TO_UINT32(26, 1, false, false), - CAST_TO_UINT64(27, 1, false, false), - CAST_TO_FLOAT32(28, 1, false, false), - CAST_TO_FLOAT64(29, 1, false, false), - CAST_TO_DECIMAL32(30, 1, false, false), - CAST_TO_DECIMAL64(31, 1, false, false), - CAST_TO_DECIMAL128(32, 1, false, false), + CHECK_PRECISION(16), + BITWISE_SHIFT_LEFT(17), + BITWISE_SHIFT_RIGHT(18), + CAST_TO_BOOL8(19), + CAST_TO_INT8(20), + CAST_TO_INT16(21), + CAST_TO_INT32(22), + CAST_TO_INT64(23), + CAST_TO_UINT8(24), + CAST_TO_UINT16(25), + CAST_TO_UINT32(26), + CAST_TO_UINT64(27), + CAST_TO_FLOAT32(28), + CAST_TO_FLOAT64(29), + CAST_TO_DECIMAL32(30), + CAST_TO_DECIMAL64(31), + CAST_TO_DECIMAL128(32), /** Rescale a decimal input to the target scale supplied to {@link JitOperation}. */ - RESCALE(33, 1, false, true), + RESCALE(33), /** Select input 0 when input 2 is true, otherwise input 1. */ - IF_ELSE(34, 3, false, false); + IF_ELSE(34); private final byte nativeId; - private final int arity; - private final boolean fallible; - private final boolean requiresTargetScale; - JitOperator(int nativeId, int arity, boolean fallible, boolean requiresTargetScale) { + JitOperator(int nativeId) { this.nativeId = (byte) nativeId; - this.arity = arity; - this.fallible = fallible; - this.requiresTargetScale = requiresTargetScale; assert this.nativeId == nativeId; } - int getArity() { - return arity; - } - - boolean isFallible() { - return fallible; - } - - boolean requiresTargetScale() { - return requiresTargetScale; - } - int getSerializedSize() { return Byte.BYTES; } 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 f20ea6b37393..e4fe81d084f0 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -381,50 +381,15 @@ public void testDecimalLiteralValidation() { @Test void testJitOperationValidation() { - AstExpression[] inputs = new AstExpression[] { - new ColumnReference(0), - new ColumnReference(1), - new ColumnReference(2) - }; - for (JitOperator op : JitOperator.values()) { - AstExpression[] opInputs = Arrays.copyOf(inputs, op.getArity()); - JitOperation expr; - if (op == JitOperator.RESCALE) { - expr = new JitOperation(op, -2, opInputs); - } else { - expr = new JitOperation(op, opInputs); - } - try (CompiledExpression ignored = expr.compile()) { - // The native decoder is part of the serialized operator contract. - } - if (op.isFallible()) { - JitOperation nullifyingExpr = new JitOperation(op, JitErrorPolicy.NULLIFY, opInputs); - try (CompiledExpression ignored = nullifyingExpr.compile()) { - // The native decoder is part of the serialized error-policy contract. - } - } - if (op.getArity() > 0) { - Assertions.assertThrows( - IllegalArgumentException.class, - () -> new JitOperation(op, Arrays.copyOf(inputs, op.getArity() - 1))); - } - } + assertJitCompileThrows(new JitOperation(JitOperator.ADD, new ColumnReference(0))); + assertJitCompileThrows(new JitOperation(JitOperator.ADD, + new ColumnReference(0), new ColumnReference(1), new ColumnReference(2))); + assertJitCompileThrows(new JitOperation(JitOperator.ADD, JitErrorPolicy.NULLIFY, + new ColumnReference(0), new ColumnReference(1))); + assertJitCompileThrows(new JitOperation(JitOperator.ADD, -2, + new ColumnReference(0), new ColumnReference(1))); + assertJitCompileThrows(new JitOperation(JitOperator.RESCALE, new ColumnReference(0))); - Assertions.assertThrows( - IllegalArgumentException.class, - () -> new JitOperation(JitOperator.ADD, JitErrorPolicy.NULLIFY, - new ColumnReference(0), new ColumnReference(1))); - Assertions.assertThrows( - IllegalArgumentException.class, - () -> new JitOperation(JitOperator.ADD, -2, - new ColumnReference(0), new ColumnReference(1))); - Assertions.assertThrows( - IllegalArgumentException.class, - () -> new JitOperation(JitOperator.RESCALE, new ColumnReference(0))); - Assertions.assertThrows( - IllegalArgumentException.class, - () -> new JitOperation(JitOperator.ADD, - new ColumnReference(0), new ColumnReference(1), new ColumnReference(2))); Assertions.assertThrows( NullPointerException.class, () -> new JitOperation(null, new ColumnReference(0))); @@ -441,6 +406,13 @@ void testJitOperationValidation() { () -> new JitOperation(JitOperator.ADD, (AstExpression[]) null)); } + private static void assertJitCompileThrows(JitOperation expr) { + Assertions.assertThrows(CudfException.class, () -> { + try (CompiledExpression ignored = expr.compile()) { + } + }); + } + @Test void testJitMismatchedOperandTypes() { JitOperation expr = new JitOperation(JitOperator.ADD, From 4f2dd1b2469820720d8da4887e6bd66317ba45de Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 15 Jul 2026 16:30:50 +0800 Subject: [PATCH 07/11] address comments Signed-off-by: Haoyang Li --- java/src/main/java/ai/rapids/cudf/Cudf.java | 10 -- .../ai/rapids/cudf/ast/AstExpression.java | 3 +- .../java/ai/rapids/cudf/ast/AstUtils.java | 19 +++ .../ai/rapids/cudf/ast/BinaryOperator.java | 3 +- .../ai/rapids/cudf/ast/JitErrorPolicy.java | 3 +- .../java/ai/rapids/cudf/ast/JitOperation.java | 10 +- .../java/ai/rapids/cudf/ast/JitOperator.java | 3 +- .../main/java/ai/rapids/cudf/ast/Literal.java | 34 ++-- .../ai/rapids/cudf/ast/TableReference.java | 3 +- .../ai/rapids/cudf/ast/UnaryOperator.java | 3 +- .../main/native/src/CompiledExpression.cpp | 107 ++++++++---- java/src/main/native/src/CudfJni.cpp | 11 -- .../cudf/ast/CompiledExpressionTest.java | 156 ++++++++++-------- 13 files changed, 209 insertions(+), 156 deletions(-) create mode 100644 java/src/main/java/ai/rapids/cudf/ast/AstUtils.java diff --git a/java/src/main/java/ai/rapids/cudf/Cudf.java b/java/src/main/java/ai/rapids/cudf/Cudf.java index 07cb67d1a519..9654b41741d5 100644 --- a/java/src/main/java/ai/rapids/cudf/Cudf.java +++ b/java/src/main/java/ai/rapids/cudf/Cudf.java @@ -11,16 +11,6 @@ public class Cudf { NativeDepsLoader.loadNativeDeps(); } - /** - * Initialize the libcudf JIT runtime and program cache. - * This method may be called repeatedly. It validates runtime dependencies but does not actively - * enable JIT evaluation; {@code CompiledExpression.computeColumn} continues to use the - * process-level libcudf configuration. - * - * @throws CudfException if the JIT runtime cannot be initialized - */ - public static native void initializeJitRuntime(); - /** * cuDF copies that are smaller than the threshold will use a kernel to copy, instead * of cudaMemcpyAsync. 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 143fcf4b6db9..503ba613d0ff 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java @@ -27,8 +27,7 @@ protected enum ExpressionType { private final byte nativeId; ExpressionType(int nativeId) { - this.nativeId = (byte) nativeId; - assert this.nativeId == nativeId; + this.nativeId = AstUtils.checkByte(nativeId); } /** Get the size in bytes to serialize this node type */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java b/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java new file mode 100644 index 000000000000..def8b1e4703e --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +final class AstUtils { + private AstUtils() { + } + + static byte checkByte(int value) { + byte result = (byte) value; + if (result != value) { + throw new IllegalArgumentException("value does not fit in a byte: " + value); + } + return result; + } +} diff --git a/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java b/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java index 91d208d3a3aa..85de7a165f25 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java @@ -40,8 +40,7 @@ public enum BinaryOperator { private final byte nativeId; BinaryOperator(int nativeId) { - this.nativeId = (byte) nativeId; - assert this.nativeId == nativeId; + this.nativeId = AstUtils.checkByte(nativeId); } /** Get the size in bytes to serialize this operator */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java index e5dfc0069980..99ef84c98097 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java @@ -21,8 +21,7 @@ public enum JitErrorPolicy { private final byte nativeId; JitErrorPolicy(int nativeId) { - this.nativeId = (byte) nativeId; - assert this.nativeId == nativeId; + this.nativeId = AstUtils.checkByte(nativeId); } int getSerializedSize() { 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 867336efaeaa..5194797eb866 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java @@ -65,8 +65,8 @@ private JitOperation( this.errorPolicy = Objects.requireNonNull(errorPolicy, "errorPolicy is null"); this.inputs = Objects.requireNonNull(inputs, "inputs is null").clone(); this.targetScale = targetScale; - for (AstExpression input : this.inputs) { - Objects.requireNonNull(input, "input is null"); + for (int i = 0; i < this.inputs.length; i++) { + Objects.requireNonNull(this.inputs[i], "input " + i + " is null"); } } @@ -75,8 +75,8 @@ int getSerializedSize() { int size = ExpressionType.JIT_EXPRESSION.getSerializedSize() + op.getSerializedSize() + errorPolicy.getSerializedSize() + - Byte.BYTES + - Byte.BYTES; + Byte.BYTES + // targetScale present + Byte.BYTES; // inputs.length if (targetScale != null) { size += Integer.BYTES; } @@ -91,11 +91,11 @@ void serialize(ByteBuffer bb) { ExpressionType.JIT_EXPRESSION.serialize(bb); op.serialize(bb); errorPolicy.serialize(bb); - bb.put((byte) inputs.length); bb.put((byte) (targetScale == null ? 0 : 1)); if (targetScale != null) { bb.putInt(targetScale); } + bb.put((byte) inputs.length); for (AstExpression input : inputs) { input.serialize(bb); } diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java index 16419ea8956a..b4d558a43dea 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java @@ -64,8 +64,7 @@ public enum JitOperator { private final byte nativeId; JitOperator(int nativeId) { - this.nativeId = (byte) nativeId; - assert this.nativeId == nativeId; + this.nativeId = AstUtils.checkByte(nativeId); } int getSerializedSize() { 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 54547fc16d3b..472d3d3e9973 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -130,7 +130,7 @@ public static Literal ofDouble(Double value) { * 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 directly. + * cannot materialize it correctly. * * @param type decimal storage type and scale * @param unscaledValue unscaled decimal value, or null @@ -151,7 +151,7 @@ public static Literal ofDecimal(DType type, BigInteger unscaledValue) { return ofLongBasedType(type, unscaledValue.longValueExact()); } else { if (unscaledValue.bitLength() > type.getSizeInBytes() * Byte.SIZE - 1) { - throw new ArithmeticException("BigInteger out of range for " + type); + throw new ArithmeticException("BigInteger out of DECIMAL128 range"); } return new Literal(type, convertDecimal128FromJavaToCudf(unscaledValue.toByteArray(), type)); } @@ -276,8 +276,7 @@ void serialize(ByteBuffer bb) { } private int getDataTypeSerializedSize() { - int nativeTypeId = type.getTypeId().getNativeId(); - assert nativeTypeId == (byte) nativeTypeId : "Type ID does not fit in a byte"; + AstUtils.checkByte(type.getTypeId().getNativeId()); if (type.isDecimalType()) { return Byte.BYTES + Integer.BYTES; } @@ -285,8 +284,7 @@ private int getDataTypeSerializedSize() { } private void serializeDataType(ByteBuffer bb) { - byte nativeTypeId = (byte) type.getTypeId().getNativeId(); - assert nativeTypeId == type.getTypeId().getNativeId() : "DType ID does not fit in a byte"; + byte nativeTypeId = AstUtils.checkByte(type.getTypeId().getNativeId()); bb.put(nativeTypeId); if (type.isDecimalType()) { bb.putInt(type.getScale()); @@ -308,13 +306,27 @@ private static Literal ofLongBasedType(DType type, long value) { } private static byte[] convertDecimal128FromJavaToCudf(byte[] bytes, DType type) { + return convertDecimal128FromJavaToCudf(bytes, type, ByteOrder.nativeOrder()); + } + + static byte[] convertDecimal128FromJavaToCudf( + byte[] bytes, DType type, ByteOrder byteOrder) { + // BigInteger uses big-endian bytes, while JNI reads the decimal128 value in native order. byte[] finalBytes = new byte[type.getSizeInBytes()]; byte signByte = (bytes[0] & 0x80) > 0 ? (byte) 0xff : (byte) 0x00; - for (int i = bytes.length; i < finalBytes.length; i++) { - finalBytes[i] = signByte; - } - for (int i = 0; i < bytes.length; i++) { - finalBytes[i] = bytes[bytes.length - i - 1]; + if (byteOrder == ByteOrder.BIG_ENDIAN) { + int offset = finalBytes.length - bytes.length; + for (int i = 0; i < offset; i++) { + finalBytes[i] = signByte; + } + System.arraycopy(bytes, 0, finalBytes, offset, bytes.length); + } else { + for (int i = bytes.length; i < finalBytes.length; i++) { + finalBytes[i] = signByte; + } + for (int i = 0; i < bytes.length; i++) { + finalBytes[i] = bytes[bytes.length - i - 1]; + } } return finalBytes; } diff --git a/java/src/main/java/ai/rapids/cudf/ast/TableReference.java b/java/src/main/java/ai/rapids/cudf/ast/TableReference.java index a7979e8c7bb9..8e61592a992c 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/TableReference.java +++ b/java/src/main/java/ai/rapids/cudf/ast/TableReference.java @@ -20,8 +20,7 @@ public enum TableReference { private final byte nativeId; TableReference(int nativeId) { - this.nativeId = (byte) nativeId; - assert this.nativeId == nativeId; + this.nativeId = AstUtils.checkByte(nativeId); } /** Get the size in bytes to serialize this table reference */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java b/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java index 9aee369fcaff..2ab63a8fb16d 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java @@ -43,8 +43,7 @@ public enum UnaryOperator { private final byte nativeId; UnaryOperator(int nativeId) { - this.nativeId = (byte) nativeId; - assert this.nativeId == nativeId; + this.nativeId = AstUtils.checkByte(nativeId); } /** Get the size in bytes to serialize this operator */ int getSerializedSize() { diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index ea65a6d0ac35..e35fbe2b40f6 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -15,10 +15,12 @@ #include #include +#include #include #include #include #include +#include #include namespace { @@ -123,7 +125,7 @@ enum class jni_serialized_expression_type : int8_t { UNARY_OPERATION = 3, BINARY_OPERATION = 4, COLUMN_NAME_REFERENCE = 5, - JIT_OPERATION = 6 + JIT_OPERATION = 6, }; /** @@ -250,7 +252,9 @@ jni_jit_operator_info jni_to_jit_operator(jbyte jni_op_value) case 32: return {CAST_TO_DECIMAL128, 1, false, false}; case 33: return {RESCALE, 1, false, true}; case 34: return {IF_ELSE, 3, false, false}; - default: throw std::invalid_argument("unexpected JNI AST JIT operator value"); + default: + throw std::invalid_argument(std::format("unexpected JNI AST JIT operator value {}", + static_cast(jni_op_value))); } } @@ -258,12 +262,16 @@ jni_jit_operator_info jni_to_jit_operator(jbyte jni_op_value) * Convert a serialized Java JIT error policy into its libcudf value. * NOTE: This must be kept in sync with JitErrorPolicy.java! */ -cudf::error_policy jni_to_jit_error_policy(jbyte jni_policy_value) +cudf::error_policy jni_to_jit_error_policy(jbyte jni_policy_value, jbyte jni_op_value) { switch (jni_policy_value) { case 0: return cudf::error_policy::PROPAGATE; case 1: return cudf::error_policy::NULLIFY; - default: throw std::invalid_argument("unexpected JNI AST JIT error policy value"); + default: + throw std::invalid_argument( + std::format("unexpected JNI AST JIT error policy {} for operator {}", + static_cast(jni_policy_value), + static_cast(jni_op_value))); } } @@ -288,7 +296,7 @@ struct make_literal { cudf::ast::literal const& operator()(cudf::data_type dtype, bool is_valid, cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_numeric_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -307,7 +315,7 @@ struct make_literal { cudf::ast::literal const& operator()(cudf::data_type dtype, bool is_valid, cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_timestamp_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -326,7 +334,7 @@ struct make_literal { cudf::ast::literal const& operator()(cudf::data_type dtype, bool is_valid, cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = cudf::make_duration_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -345,7 +353,7 @@ struct make_literal { cudf::ast::literal const& operator()(cudf::data_type dtype, bool is_valid, cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + jni_serialized_ast& jni_ast) const { std::unique_ptr scalar_ptr = [&]() { if (is_valid) { @@ -360,25 +368,12 @@ struct make_literal { return compiled_expr.add_literal(str_scalar, std::move(scalar_ptr)); } - /** Default functor implementation to catch type dispatch errors */ - template () && !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) - { - throw std::logic_error("Unsupported AST literal type"); - } - /** 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) + jni_serialized_ast& jni_ast) const { using rep_type = typename T::rep; auto const val = is_valid ? jni_ast.read() : rep_type{}; @@ -389,6 +384,19 @@ struct make_literal { auto& fixed_point_scalar = static_cast&>(*scalar_ptr); return compiled_expr.add_literal(fixed_point_scalar, std::move(scalar_ptr)); } + + /** Default functor implementation to catch type dispatch errors */ + template () && !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 + { + throw std::logic_error("Unsupported AST literal type"); + } }; /** Decode a serialized AST literal */ @@ -444,25 +452,45 @@ 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) { - auto const op_info = jni_to_jit_operator(jni_ast.read_byte()); - auto const error_policy = jni_to_jit_error_policy(jni_ast.read_byte()); - auto const arity = static_cast(jni_ast.read_byte()); - if (arity < 0) { throw std::invalid_argument("unexpected JNI AST JIT operator arity"); } - if (static_cast(arity) != op_info.arity) { - throw std::invalid_argument("unexpected JNI AST JIT operator arity"); - } + 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(); + auto const error_policy = jni_to_jit_error_policy(jni_policy_value, jni_op_value); if (error_policy == cudf::error_policy::NULLIFY && !op_info.is_fallible) { - throw std::invalid_argument("unexpected error policy for JNI AST JIT operator"); + throw std::invalid_argument( + std::format("unexpected error policy {} for non-fallible JNI AST JIT operator {}", + static_cast(jni_policy_value), + static_cast(jni_op_value))); } auto const has_target_scale = jni_ast.read_byte(); if (has_target_scale != 0 && has_target_scale != 1) { - throw std::invalid_argument("unexpected JNI AST JIT target scale flag"); + throw std::invalid_argument( + std::format("unexpected target scale flag {} for JNI AST JIT operator {}; expected 0 or 1", + static_cast(has_target_scale), + static_cast(jni_op_value))); } std::optional target_scale; if (has_target_scale == 1) { target_scale = jni_ast.read(); } if (target_scale.has_value() != op_info.requires_target_scale) { - throw std::invalid_argument("unexpected target scale for JNI AST JIT operator"); + auto const actual = target_scale.has_value() ? std::to_string(*target_scale) : "none"; + auto const expected = op_info.requires_target_scale ? "a value" : "none"; + throw std::invalid_argument( + std::format("unexpected target scale {} for JNI AST JIT operator " + "{}; expected {}", + actual, + static_cast(jni_op_value), + expected)); + } + + auto const arity = static_cast(jni_ast.read_byte()); + if (static_cast(arity) != op_info.arity) { + throw std::invalid_argument( + std::format("unexpected arity {} for JNI AST JIT operator {}; " + "expected {}", + arity, + static_cast(jni_op_value), + op_info.arity)); } std::vector> args; @@ -521,10 +549,17 @@ 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); - std::unique_ptr result = - backend == execution_backend::JIT - ? cudf::compute_column_jit(*tview_ptr, compiled_expr_ptr->get_top_expression()) - : cudf::compute_column(*tview_ptr, compiled_expr_ptr->get_top_expression()); + auto const& expression = compiled_expr_ptr->get_top_expression(); + if (backend == execution_backend::DEFAULT) { + auto const* literal = dynamic_cast(&expression); + // The legacy evaluator silently corrupts decimal128 literal outputs in release builds. + if (literal != nullptr && literal->get_data_type().id() == cudf::type_id::DECIMAL128) { + throw std::invalid_argument("DECIMAL128 root literals require JIT evaluation"); + } + } + std::unique_ptr result = backend == execution_backend::JIT + ? cudf::compute_column_jit(*tview_ptr, expression) + : cudf::compute_column(*tview_ptr, expression); return reinterpret_cast(result.release()); } diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 3592a0394bd1..7e648a9ee0e0 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -5,7 +5,6 @@ #include "cudf_jni_apis.hpp" -#include #include #include #include @@ -192,16 +191,6 @@ JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_Cuda_isPtdsEnabled(JNIEnv* env, j return cudf::jni::is_ptds_enabled; } -JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cudf_initializeJitRuntime(JNIEnv* env, jclass) -{ - JNI_TRY - { - cudf::jni::auto_set_device(env); - cudf::initialize(cudf::init_flags::INIT_JIT_CACHE); - } - JNI_CATCH(env, ); -} - JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cudf_setKernelPinnedCopyThreshold(JNIEnv* env, jclass clazz, jlong jthreshold) 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 e4fe81d084f0..1aaba18f3d29 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -6,7 +6,6 @@ package ai.rapids.cudf.ast; import ai.rapids.cudf.ColumnVector; -import ai.rapids.cudf.Cudf; import ai.rapids.cudf.CudfException; import ai.rapids.cudf.CudfTestBase; import ai.rapids.cudf.DType; @@ -21,22 +20,18 @@ import org.junit.jupiter.params.provider.NullSource; import java.math.BigInteger; +import java.nio.ByteOrder; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.function.BiFunction; import java.util.function.Function; +import java.util.function.Supplier; import java.util.stream.Stream; import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual; public class CompiledExpressionTest extends CudfTestBase { - @Test - public void testInitializeJitRuntime() { - Assertions.assertDoesNotThrow(Cudf::initializeJitRuntime); - Assertions.assertDoesNotThrow(Cudf::initializeJitRuntime); - } - @Test public void testColumnReferenceTransform() { try (Table t = new Table.TestBuilder().column(5, 4, 3, 2, 1).column(6, 7, 8, null, 10).build()) { @@ -346,6 +341,16 @@ public void testDecimalLiteralTransform(DType type, BigInteger value) { } } + @ParameterizedTest + @MethodSource("createDecimal128LiteralParams") + public void testDecimal128LiteralLegacyTransformFails(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()) { + Assertions.assertThrows(CudfException.class, () -> compiledExpr.computeColumn(t).close()); + } + } + @ParameterizedTest @MethodSource("createDecimalLiteralParams") public void testJitDecimalLiteralTransform(DType type, BigInteger value) { @@ -379,6 +384,35 @@ public void testDecimalLiteralValidation() { BigInteger.ONE.shiftLeft(127).negate().subtract(BigInteger.ONE))); } + @Test + public void testDecimal128LiteralByteOrderConversion() { + DType type = DType.create(DType.DTypeEnum.DECIMAL128, 0); + byte[] positiveBigEndian = new byte[type.getSizeInBytes()]; + positiveBigEndian[positiveBigEndian.length - 1] = 1; + byte[] positiveLittleEndian = new byte[type.getSizeInBytes()]; + positiveLittleEndian[0] = 1; + Assertions.assertArrayEquals(positiveBigEndian, + Literal.convertDecimal128FromJavaToCudf( + BigInteger.ONE.toByteArray(), type, ByteOrder.BIG_ENDIAN)); + Assertions.assertArrayEquals(positiveLittleEndian, + Literal.convertDecimal128FromJavaToCudf( + BigInteger.ONE.toByteArray(), type, ByteOrder.LITTLE_ENDIAN)); + + byte[] negativeBigEndian = new byte[type.getSizeInBytes()]; + Arrays.fill(negativeBigEndian, (byte) 0xff); + negativeBigEndian[negativeBigEndian.length - 1] = (byte) 0xfe; + byte[] negativeLittleEndian = new byte[type.getSizeInBytes()]; + Arrays.fill(negativeLittleEndian, (byte) 0xff); + negativeLittleEndian[0] = (byte) 0xfe; + byte[] negativeValue = BigInteger.valueOf(-2).toByteArray(); + Assertions.assertArrayEquals(negativeBigEndian, + Literal.convertDecimal128FromJavaToCudf( + negativeValue, type, ByteOrder.BIG_ENDIAN)); + Assertions.assertArrayEquals(negativeLittleEndian, + Literal.convertDecimal128FromJavaToCudf( + negativeValue, type, ByteOrder.LITTLE_ENDIAN)); + } + @Test void testJitOperationValidation() { assertJitCompileThrows(new JitOperation(JitOperator.ADD, new ColumnReference(0))); @@ -397,10 +431,11 @@ void testJitOperationValidation() { NullPointerException.class, () -> new JitOperation(JitOperator.ADD, (JitErrorPolicy) null, new ColumnReference(0), new ColumnReference(1))); - Assertions.assertThrows( + NullPointerException nullInputError = Assertions.assertThrows( NullPointerException.class, () -> new JitOperation(JitOperator.ADD, new ColumnReference(0), (AstExpression) null)); + Assertions.assertEquals("input 1 is null", nullInputError.getMessage()); Assertions.assertThrows( NullPointerException.class, () -> new JitOperation(JitOperator.ADD, (AstExpression[]) null)); @@ -599,91 +634,70 @@ void testJitMixedConditionalTransform() { } } - private static Stream createJitNumericCastParams() { + private static Arguments jitCastCase( + JitOperator op, Supplier expectedFactory) { + return Arguments.of(op, expectedFactory); + } + + private static Stream createJitNumericCastParams() { return Stream.of( - JitOperator.CAST_TO_BOOL8, - JitOperator.CAST_TO_INT8, - JitOperator.CAST_TO_INT16, - JitOperator.CAST_TO_INT32, - JitOperator.CAST_TO_INT64, - JitOperator.CAST_TO_UINT8, - JitOperator.CAST_TO_UINT16, - JitOperator.CAST_TO_UINT32, - JitOperator.CAST_TO_UINT64, - JitOperator.CAST_TO_FLOAT32, - JitOperator.CAST_TO_FLOAT64); - } - - private static ColumnVector makeJitNumericCastExpected(JitOperator op) { - switch (op) { - case CAST_TO_BOOL8: - return ColumnVector.fromBooleans(false, true, true, true); - case CAST_TO_INT8: - return ColumnVector.fromBytes((byte) 0, (byte) 1, (byte) 2, (byte) 3); - case CAST_TO_INT16: - return ColumnVector.fromShorts((short) 0, (short) 1, (short) 2, (short) 3); - case CAST_TO_INT32: - return ColumnVector.fromInts(0, 1, 2, 3); - case CAST_TO_INT64: - return ColumnVector.fromLongs(0L, 1L, 2L, 3L); - case CAST_TO_UINT8: - return ColumnVector.fromUnsignedBytes((byte) 0, (byte) 1, (byte) 2, (byte) 3); - case CAST_TO_UINT16: - return ColumnVector.fromUnsignedShorts((short) 0, (short) 1, (short) 2, (short) 3); - case CAST_TO_UINT32: - return ColumnVector.fromUnsignedInts(0, 1, 2, 3); - case CAST_TO_UINT64: - return ColumnVector.fromUnsignedLongs(0L, 1L, 2L, 3L); - case CAST_TO_FLOAT32: - return ColumnVector.fromFloats(0.0f, 1.0f, 2.0f, 3.0f); - case CAST_TO_FLOAT64: - return ColumnVector.fromDoubles(0.0, 1.0, 2.0, 3.0); - default: - throw new IllegalArgumentException("Unexpected numeric cast operator " + op); - } + jitCastCase(JitOperator.CAST_TO_BOOL8, + () -> ColumnVector.fromBooleans(false, true, true, true)), + jitCastCase(JitOperator.CAST_TO_INT8, + () -> ColumnVector.fromBytes((byte) 0, (byte) 1, (byte) 2, (byte) 3)), + jitCastCase(JitOperator.CAST_TO_INT16, + () -> ColumnVector.fromShorts((short) 0, (short) 1, (short) 2, (short) 3)), + jitCastCase(JitOperator.CAST_TO_INT32, + () -> ColumnVector.fromInts(0, 1, 2, 3)), + jitCastCase(JitOperator.CAST_TO_INT64, + () -> ColumnVector.fromLongs(0L, 1L, 2L, 3L)), + jitCastCase(JitOperator.CAST_TO_UINT8, + () -> ColumnVector.fromUnsignedBytes((byte) 0, (byte) 1, (byte) 2, (byte) 3)), + jitCastCase(JitOperator.CAST_TO_UINT16, + () -> ColumnVector.fromUnsignedShorts((short) 0, (short) 1, (short) 2, (short) 3)), + jitCastCase(JitOperator.CAST_TO_UINT32, + () -> ColumnVector.fromUnsignedInts(0, 1, 2, 3)), + jitCastCase(JitOperator.CAST_TO_UINT64, + () -> ColumnVector.fromUnsignedLongs(0L, 1L, 2L, 3L)), + jitCastCase(JitOperator.CAST_TO_FLOAT32, + () -> ColumnVector.fromFloats(0.0f, 1.0f, 2.0f, 3.0f)), + jitCastCase(JitOperator.CAST_TO_FLOAT64, + () -> ColumnVector.fromDoubles(0.0, 1.0, 2.0, 3.0))); } @ParameterizedTest @MethodSource("createJitNumericCastParams") - void testJitNumericCastTransform(JitOperator op) { + 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); - ColumnVector expected = makeJitNumericCastExpected(op)) { + ColumnVector expected = expectedFactory.get()) { assertColumnsAreEqual(expected, actual); } } - private static Stream createJitDecimalCastParams() { + private static Stream createJitDecimalCastParams() { return Stream.of( - JitOperator.CAST_TO_DECIMAL32, - JitOperator.CAST_TO_DECIMAL64, - JitOperator.CAST_TO_DECIMAL128); - } - - private static ColumnVector makeJitDecimalCastExpected(JitOperator op) { - switch (op) { - case CAST_TO_DECIMAL32: - return ColumnVector.decimalFromInts(0, 0, 1, -2, 3); - case CAST_TO_DECIMAL64: - return ColumnVector.decimalFromLongs(0, 0L, 1L, -2L, 3L); - case CAST_TO_DECIMAL128: - return ColumnVector.decimalFromBigInt(0, - BigInteger.ZERO, BigInteger.ONE, BigInteger.valueOf(-2), BigInteger.valueOf(3)); - default: - throw new IllegalArgumentException("Unexpected decimal cast operator " + op); - } + jitCastCase(JitOperator.CAST_TO_DECIMAL32, + () -> ColumnVector.decimalFromInts(0, 0, 1, -2, 3)), + jitCastCase(JitOperator.CAST_TO_DECIMAL64, + () -> ColumnVector.decimalFromLongs(0, 0L, 1L, -2L, 3L)), + jitCastCase(JitOperator.CAST_TO_DECIMAL128, + () -> ColumnVector.decimalFromBigInt(0, + BigInteger.ZERO, BigInteger.ONE, BigInteger.valueOf(-2), BigInteger.valueOf(3)))); } @ParameterizedTest @MethodSource("createJitDecimalCastParams") - void testJitDecimalCastTransform(JitOperator op) { + 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); - ColumnVector expected = makeJitDecimalCastExpected(op)) { + ColumnVector expected = expectedFactory.get()) { assertColumnsAreEqual(expected, actual); } } From 122192fa84ccd21a2059a2c9e63b5f99979e0fa4 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Thu, 16 Jul 2026 12:00:33 +0800 Subject: [PATCH 08/11] address comments Signed-off-by: Haoyang Li --- cpp/src/ast/expression_parser.cpp | 3 ++ cpp/tests/ast/transform_tests.cpp | 21 ++++++++++++++ .../main/java/ai/rapids/cudf/ast/Literal.java | 1 + .../main/native/src/CompiledExpression.cpp | 7 ----- .../cudf/ast/CompiledExpressionTest.java | 28 +++++++++++++++++++ 5 files changed, 53 insertions(+), 7 deletions(-) diff --git a/cpp/src/ast/expression_parser.cpp b/cpp/src/ast/expression_parser.cpp index 37b9da8f02ba..8392779e388d 100644 --- a/cpp/src/ast/expression_parser.cpp +++ b/cpp/src/ast/expression_parser.cpp @@ -265,6 +265,9 @@ cudf::size_type expression_parser::visit(operation const& expr) auto const output = [&]() { if (expression_index == 0) { // This expression is the root. Output should be directed to the output column. + CUDF_EXPECTS(data_type.id() != cudf::type_id::DECIMAL128, + "decimal128 is not supported as an AST expression output.", + cudf::data_type_error); return detail::device_data_reference( detail::device_data_reference_type::COLUMN, data_type, 0, table_reference::OUTPUT); } else { diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index 172c3cba0a82..a923c8e96c8c 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -1542,4 +1542,25 @@ TYPED_TEST(TransformTest, Decimal128Unsupported) } } +TYPED_TEST(TransformTest, Decimal128IdentityOutput) +{ + using Executor = TypeParam; + + auto const input = column_wrapper{0, 0}; + auto const table = cudf::table_view{{input}}; + auto const scale = numeric::scale_type{-2}; + auto literal_value = cudf::fixed_point_scalar(12345, scale, true); + auto literal = cudf::ast::literal(literal_value); + auto expression = cudf::ast::operation(cudf::ast::ast_operator::IDENTITY, literal); + + if constexpr (std::is_same_v) { + EXPECT_THROW(Executor::compute_column(table, expression), cudf::data_type_error); + } else { + auto result = Executor::compute_column(table, expression); + auto expected = cudf::test::fixed_point_column_wrapper<__int128_t>( + {__int128_t{12345}, __int128_t{12345}}, scale); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity); + } +} + CUDF_TEST_PROGRAM_MAIN() 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 472d3d3e9973..ead38ab0e560 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -309,6 +309,7 @@ private static byte[] convertDecimal128FromJavaToCudf(byte[] bytes, DType type) return convertDecimal128FromJavaToCudf(bytes, type, ByteOrder.nativeOrder()); } + // Visible for testing so both possible native byte orders can be exercised. static byte[] convertDecimal128FromJavaToCudf( byte[] bytes, DType type, ByteOrder byteOrder) { // BigInteger uses big-endian bytes, while JNI reads the decimal128 value in native order. diff --git a/java/src/main/native/src/CompiledExpression.cpp b/java/src/main/native/src/CompiledExpression.cpp index e35fbe2b40f6..44748ef6b9b9 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -550,13 +550,6 @@ 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(); - if (backend == execution_backend::DEFAULT) { - auto const* literal = dynamic_cast(&expression); - // The legacy evaluator silently corrupts decimal128 literal outputs in release builds. - if (literal != nullptr && literal->get_data_type().id() == cudf::type_id::DECIMAL128) { - throw std::invalid_argument("DECIMAL128 root literals require JIT evaluation"); - } - } std::unique_ptr result = backend == execution_backend::JIT ? cudf::compute_column_jit(*tview_ptr, expression) : cudf::compute_column(*tview_ptr, expression); 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 1aaba18f3d29..575e899bf32d 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -20,6 +20,7 @@ import org.junit.jupiter.params.provider.NullSource; import java.math.BigInteger; +import java.math.RoundingMode; import java.nio.ByteOrder; import java.util.ArrayList; import java.util.Arrays; @@ -351,6 +352,33 @@ public void testDecimal128LiteralLegacyTransformFails(DType type, BigInteger val } } + @ParameterizedTest + @MethodSource("createDecimal128LiteralParams") + public void testDecimal128IdentityLegacyTransformFails(DType type, BigInteger value) { + UnaryOperation expr = new UnaryOperation( + UnaryOperator.IDENTITY, Literal.ofDecimal(type, value)); + try (Table t = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiledExpr = expr.compile()) { + Assertions.assertThrows(CudfException.class, () -> compiledExpr.computeColumn(t).close()); + } + } + + @Test + public void testDecimal128LiteralComparisonLegacyTransform() { + DType type = DType.create(DType.DTypeEnum.DECIMAL128, 0); + BinaryOperation expr = new BinaryOperation(BinaryOperator.GREATER, + new ColumnReference(0), Literal.ofDecimal(type, BigInteger.ONE)); + try (Table t = new Table.TestBuilder() + .decimal128Column(0, RoundingMode.UNNECESSARY, + BigInteger.ZERO, BigInteger.ONE, BigInteger.valueOf(2), null) + .build(); + CompiledExpression compiledExpr = expr.compile(); + ColumnVector actual = compiledExpr.computeColumn(t); + ColumnVector expected = ColumnVector.fromBoxedBooleans(false, false, true, null)) { + assertColumnsAreEqual(expected, actual); + } + } + @ParameterizedTest @MethodSource("createDecimalLiteralParams") public void testJitDecimalLiteralTransform(DType type, BigInteger value) { From 27d63d860371cb5f832579960b305f2da4d39ad5 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 17 Jul 2026 09:39:08 +0800 Subject: [PATCH 09/11] address comments Signed-off-by: Haoyang Li --- cpp/tests/ast/transform_tests.cpp | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index a923c8e96c8c..890848987fc5 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -1514,9 +1514,8 @@ TYPED_TEST(TransformTest, Decimal128Unsupported) // column = {2000.00, 2000.00} (rep 200000 @ scale -2) auto const scale = numeric::scale_type{-2}; - auto const col = cudf::test::fixed_point_column_wrapper<__int128>{ - {__int128_t{200000}, __int128_t{200000}}, scale}; - auto table = cudf::table_view{{col}}; + auto const col = cudf::test::fixed_point_column_wrapper<__int128>{{200000, 200000}, scale}; + auto table = cudf::table_view{{col}}; // literal = 0.5 @ scale -2 (rep 50) auto half = numeric::decimal128{numeric::scaled_integer{50, scale}}; @@ -1526,18 +1525,15 @@ TYPED_TEST(TransformTest, Decimal128Unsupported) auto cr = cudf::ast::column_reference(0); auto ast = cudf::ast::operation(cudf::ast::ast_operator::MUL, lr, cr); - EXPECT_THROW(cudf::compute_column(table, ast), cudf::data_type_error); - if constexpr (std::is_same_v) { - auto ast2 = cudf::ast::operation(cudf::ast::ast_operator::MUL, lr, cr); - EXPECT_THROW(cudf::compute_column(table, ast), cudf::data_type_error); + EXPECT_THROW(Executor::compute_column(table, ast), cudf::data_type_error); } else { auto result = Executor::compute_column(table, ast); // Expected: 0.5 * 2000.00 = 1000.00 => rep 10000000 @ scale -4 EXPECT_EQ(result->type().id(), cudf::type_id::DECIMAL128); EXPECT_EQ(result->type().scale(), numeric::scale_type{-4}); - auto expected = cudf::test::fixed_point_column_wrapper<__int128>{ - {__int128_t{10000000}, __int128_t{10000000}}, numeric::scale_type{-4}}; + auto expected = cudf::test::fixed_point_column_wrapper<__int128>{{10000000, 10000000}, + numeric::scale_type{-4}}; CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, expected); } } @@ -1557,8 +1553,7 @@ TYPED_TEST(TransformTest, Decimal128IdentityOutput) EXPECT_THROW(Executor::compute_column(table, expression), cudf::data_type_error); } else { auto result = Executor::compute_column(table, expression); - auto expected = cudf::test::fixed_point_column_wrapper<__int128_t>( - {__int128_t{12345}, __int128_t{12345}}, scale); + auto expected = cudf::test::fixed_point_column_wrapper<__int128_t>({12345, 12345}, scale); CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity); } } From 842c4a0ee05a2484c5a0ad96204426f9875c6e7e Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 20 Jul 2026 10:10:48 +0800 Subject: [PATCH 10/11] address comments Signed-off-by: Haoyang Li --- cpp/tests/ast/transform_tests.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index 890848987fc5..fe4960b9dfec 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -1553,7 +1553,7 @@ TYPED_TEST(TransformTest, Decimal128IdentityOutput) EXPECT_THROW(Executor::compute_column(table, expression), cudf::data_type_error); } else { auto result = Executor::compute_column(table, expression); - auto expected = cudf::test::fixed_point_column_wrapper<__int128_t>({12345, 12345}, scale); + auto expected = cudf::test::fixed_point_column_wrapper<__int128>({12345, 12345}, scale); CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity); } } From 3aeba52574bf051eaf02ac2760237d757c405dbe Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 20 Jul 2026 10:29:41 +0800 Subject: [PATCH 11/11] style Signed-off-by: Haoyang Li --- java/src/main/java/ai/rapids/cudf/Cudf.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/AstUtils.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/JitOperation.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/JitOperator.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/Literal.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/TableReference.java | 2 +- java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java | 2 +- java/src/main/native/src/CudfJni.cpp | 2 +- 11 files changed, 11 insertions(+), 11 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/Cudf.java b/java/src/main/java/ai/rapids/cudf/Cudf.java index 9654b41741d5..bf26dda89eae 100644 --- a/java/src/main/java/ai/rapids/cudf/Cudf.java +++ b/java/src/main/java/ai/rapids/cudf/Cudf.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java b/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java index def8b1e4703e..332e874d82c7 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java +++ b/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java b/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java index 85de7a165f25..80276f2642c3 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/BinaryOperator.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ 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 80b7a1217392..52c3c3ac03d7 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java index 99ef84c98097..70b901999e37 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ 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 5194797eb866..7418413702af 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java index b4d558a43dea..845ac5d4050f 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ 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 ead38ab0e560..715fc1cce42e 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/Literal.java +++ b/java/src/main/java/ai/rapids/cudf/ast/Literal.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/TableReference.java b/java/src/main/java/ai/rapids/cudf/ast/TableReference.java index 8e61592a992c..69d25b7fcf9b 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/TableReference.java +++ b/java/src/main/java/ai/rapids/cudf/ast/TableReference.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java b/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java index 2ab63a8fb16d..2be73b961c75 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java +++ b/java/src/main/java/ai/rapids/cudf/ast/UnaryOperator.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 7e648a9ee0e0..65d01b41d4f7 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */