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..fe4960b9dfec 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,20 +1525,37 @@ 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); } } +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>({12345, 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/Cudf.java b/java/src/main/java/ai/rapids/cudf/Cudf.java index 2d8cf6c88ff4..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, 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/AstExpression.java b/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java index 1eea40bca0f8..503ba613d0ff 100644 --- a/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java +++ b/java/src/main/java/ai/rapids/cudf/ast/AstExpression.java @@ -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_expression_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), @@ -20,13 +21,13 @@ protected enum ExpressionType { COLUMN_REFERENCE(2), UNARY_EXPRESSION(3), BINARY_EXPRESSION(4), - COLUMN_NAME_REFERENCE(5); + COLUMN_NAME_REFERENCE(5), + JIT_EXPRESSION(6); 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..332e874d82c7 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/AstUtils.java @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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..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 */ @@ -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/CompiledExpression.java b/java/src/main/java/ai/rapids/cudf/ast/CompiledExpression.java index 5d92d623b38b..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, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * 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/JitErrorPolicy.java b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java new file mode 100644 index 000000000000..70b901999e37 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/JitErrorPolicy.java @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +import java.nio.ByteBuffer; + +/** + * 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 JitErrorPolicy { + /** Propagate an evaluation error to the caller. */ + PROPAGATE(0), + /** Produce null for a row where evaluation fails. */ + NULLIFY(1); + + private final byte nativeId; + + JitErrorPolicy(int nativeId) { + this.nativeId = AstUtils.checkByte(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..7418413702af --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperation.java @@ -0,0 +1,125 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +import java.nio.ByteBuffer; +import java.util.Objects; + +/** + * 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; + 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 + */ + public JitOperation(JitOperator op, AstExpression... inputs) { + this(op, JitErrorPolicy.PROPAGATE, 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 + */ + 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 + */ + public JitOperation(JitOperator op, int targetScale, AstExpression... inputs) { + this(op, JitErrorPolicy.PROPAGATE, Integer.valueOf(targetScale), inputs); + } + + private JitOperation( + JitOperator op, + JitErrorPolicy errorPolicy, + Integer targetScale, + AstExpression... inputs) { + this.op = Objects.requireNonNull(op, "op is null"); + this.errorPolicy = Objects.requireNonNull(errorPolicy, "errorPolicy is null"); + this.inputs = Objects.requireNonNull(inputs, "inputs is null").clone(); + this.targetScale = targetScale; + for (int i = 0; i < this.inputs.length; i++) { + Objects.requireNonNull(this.inputs[i], "input " + i + " is null"); + } + } + + @Override + int getSerializedSize() { + int size = ExpressionType.JIT_EXPRESSION.getSerializedSize() + + op.getSerializedSize() + + errorPolicy.getSerializedSize() + + Byte.BYTES + // targetScale present + Byte.BYTES; // inputs.length + 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); + errorPolicy.serialize(bb); + 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); + } + } + + @Override + public String toString() { + StringBuilder ret = new StringBuilder(op.toString()); + if (errorPolicy != JitErrorPolicy.PROPAGATE) { + ret.append("[").append(errorPolicy).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..845ac5d4050f --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/ast/JitOperator.java @@ -0,0 +1,77 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf.ast; + +import java.nio.ByteBuffer; + +/** + * 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 { + /** Return the first non-null input. */ + COALESCE(0), + /** Convert a nullable boolean input into an always-valid predicate. */ + PREDICATE(1), + ADD(2), + SUB(3), + MUL(4), + /** Divide without reporting arithmetic errors. */ + 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), + 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), + 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), + /** Select input 0 when input 2 is true, otherwise input 1. */ + IF_ELSE(34); + + private final byte nativeId; + + JitOperator(int nativeId) { + this.nativeId = AstUtils.checkByte(nativeId); + } + + 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..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-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -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,39 @@ public static Literal ofDouble(Double value) { return ofDouble(value.doubleValue()); } + /** + * Construct a decimal literal with the specified type and unscaled value. + * A null {@code unscaledValue} produces a null literal of the requested type. + * Root literals of type {@code DECIMAL32} or {@code DECIMAL64} can be evaluated with either + * {@link CompiledExpression#computeColumn} or {@link CompiledExpression#computeColumnJit}. + * A {@code DECIMAL128} root literal must use {@code computeColumnJit}; the legacy executor + * cannot materialize it correctly. + * + * @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); + } + 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 { + if (unscaledValue.bitLength() > type.getSizeInBytes() * Byte.SIZE - 1) { + throw new ArithmeticException("BigInteger out of DECIMAL128 range"); + } + 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); @@ -242,23 +276,18 @@ 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()) { - 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) { - 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()) { - 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()); } } @@ -275,4 +304,31 @@ 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) { + 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. + byte[] finalBytes = new byte[type.getSizeInBytes()]; + byte signByte = (bytes[0] & 0x80) > 0 ? (byte) 0xff : (byte) 0x00; + 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..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 */ @@ -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..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 */ @@ -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 cca12b7b1f92..44748ef6b9b9 100644 --- a/java/src/main/native/src/CompiledExpression.cpp +++ b/java/src/main/native/src/CompiledExpression.cpp @@ -7,14 +7,20 @@ #include "jni_compiled_expr.hpp" #include +#include #include #include #include #include +#include #include +#include +#include #include +#include #include +#include #include namespace { @@ -98,11 +104,12 @@ class jni_serialized_ast { return cudf::data_type(dtype_id); } case cudf::type_id::DECIMAL32: - case cudf::type_id::DECIMAL64: { - int32_t const scale = read_byte(); + case cudf::type_id::DECIMAL64: + case cudf::type_id::DECIMAL128: { + 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"); } } }; @@ -117,7 +124,8 @@ enum class jni_serialized_expression_type : int8_t { COLUMN_REFERENCE = 2, UNARY_OPERATION = 3, BINARY_OPERATION = 4, - COLUMN_NAME_REFERENCE = 5 + COLUMN_NAME_REFERENCE = 5, + JIT_OPERATION = 6, }; /** @@ -194,6 +202,79 @@ cudf::ast::ast_operator jni_to_binary_operator(jbyte jni_op_value) } } +struct jni_jit_operator_info { + cudf::ast::jit::op op; + std::size_t arity; + bool is_fallible; + bool requires_target_scale; +}; + +/** + * Convert a serialized Java JIT operator into its libcudf definition. + * NOTE: This must be kept in sync with JitOperator.java! + */ +jni_jit_operator_info jni_to_jit_operator(jbyte jni_op_value) +{ + 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(std::format("unexpected JNI AST JIT operator value {}", + static_cast(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, 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( + std::format("unexpected JNI AST JIT error policy {} for operator {}", + static_cast(jni_policy_value), + static_cast(jni_op_value))); + } +} + /** * Convert a Java AST serialized byte representing an AST table reference into the * corresponding libcudf AST table reference. @@ -212,10 +293,10 @@ cudf::ast::table_reference jni_to_table_reference(jbyte jni_value) struct make_literal { /** Construct an AST literal from a numeric value */ template ()>* = nullptr> - cudf::ast::literal& 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) const { std::unique_ptr scalar_ptr = cudf::make_numeric_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -226,16 +307,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) const { std::unique_ptr scalar_ptr = cudf::make_timestamp_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -246,16 +326,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) const { std::unique_ptr scalar_ptr = cudf::make_duration_scalar(dtype); scalar_ptr->set_valid_async(is_valid); @@ -266,16 +345,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) const { std::unique_ptr scalar_ptr = [&]() { if (is_valid) { @@ -287,80 +365,149 @@ 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)); + } + + /** Construct an AST literal from a fixed-point value */ + template ()>* = nullptr> + cudf::ast::literal const& operator()(cudf::data_type dtype, + bool is_valid, + cudf::jni::ast::compiled_expr& compiled_expr, + jni_serialized_ast& jni_ast) const + { + 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(fixed_point_scalar, std::move(scalar_ptr)); } /** 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> - cudf::ast::literal& operator()(cudf::data_type dtype, - bool is_valid, - cudf::jni::ast::compiled_expr& compiled_expr, - jni_serialized_ast& jni_ast) + 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 */ -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); } /** Decode a serialized AST column name reference */ -cudf::ast::column_name_reference& compile_column_name_reference( +cudf::ast::column_name_reference const& compile_column_name_reference( cudf::jni::ast::compiled_expr& compiled_expr, jni_serialized_ast& jni_ast) { std::string column_name = jni_ast.read(); - return compiled_expr.add_column_name_ref( - std::make_unique(std::move(column_name))); + return compiled_expr.add_column_name_ref(std::move(column_name)); } // 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) +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& child_expression = compile_expression(compiled_expr, jni_ast); - return compiled_expr.add_operation( - std::make_unique(ast_op, child_expression)); + auto const ast_op = jni_to_unary_operator(jni_ast.read_byte()); + cudf::ast::expression const& child_expression = compile_expression(compiled_expr, jni_ast); + return compiled_expr.add_operation(ast_op, child_expression); } /** Decode a serialized AST binary expression */ -cudf::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(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 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); - return compiled_expr.add_operation( - std::make_unique(ast_op, left_child, right_child)); + 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( + 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( + 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) { + 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; + args.reserve(arity); + for (int32_t index = 0; index < arity; ++index) { + args.emplace_back(compile_expression(compiled_expr, jni_ast)); + } + + return compiled_expr.add_jit_expression( + [&](cudf::ast::tree& tree) -> cudf::ast::expression const& { + return cudf::ast::jit::operation(tree, op_info.op, args, error_policy, target_scale); + }); } /** 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) { @@ -376,6 +523,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"); } } @@ -388,9 +537,25 @@ 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; } +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); + auto const& expression = compiled_expr_ptr->get_top_expression(); + std::unique_ptr result = backend == execution_backend::JIT + ? cudf::compute_column_jit(*tview_ptr, expression) + : cudf::compute_column(*tview_ptr, expression); + return reinterpret_cast(result.release()); +} + } // anonymous namespace extern "C" { @@ -422,11 +587,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..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-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/java/src/main/native/src/jni_compiled_expr.hpp b/java/src/main/native/src/jni_compiled_expr.hpp index c0185565524b..f51457bd0606 100644 --- a/java/src/main/native/src/jni_compiled_expr.hpp +++ b/java/src/main/native/src/jni_compiled_expr.hpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -16,54 +17,57 @@ 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; - /** GPU scalar instances that correspond to literal nodes */ std::vector> scalars; + /** 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)); - return static_cast(*expressions.back()); + return expressions.emplace(scalar); + } + + cudf::ast::column_reference const& add_column_ref(cudf::size_type column_index, + cudf::ast::table_reference table_ref) + { + return expressions.emplace(column_index, table_ref); + } + + cudf::ast::column_name_reference const& add_column_name_ref(std::string column_name) + { + return expressions.emplace(std::move(column_name)); } - cudf::ast::column_reference& add_column_ref(std::unique_ptr ref_ptr) + cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, + cudf::ast::expression const& child) { - expressions.push_back(std::move(ref_ptr)); - return static_cast(*expressions.back()); + return expressions.emplace(op, child); } - /** @brief Take ownership of @p ref_ptr; returns a reference stable for the lifetime of this - * compiled_expr. */ - cudf::ast::column_name_reference& add_column_name_ref( - std::unique_ptr ref_ptr) + cudf::ast::operation const& add_operation(cudf::ast::ast_operator op, + cudf::ast::expression const& left, + cudf::ast::expression const& right) { - expressions.push_back(std::move(ref_ptr)); - return static_cast(*expressions.back()); + return expressions.emplace(op, left, right); } - cudf::ast::operation& add_operation(std::unique_ptr expr_ptr) + template + cudf::ast::expression const& add_jit_expression(F&& factory) { - expressions.push_back(std::move(expr_ptr)); - return static_cast(*expressions.back()); + return factory(expressions); } + [[nodiscard]] bool has_literals() const { return !scalars.empty(); } + /** 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 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 d32bd130bde8..575e899bf32d 100644 --- a/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java +++ b/java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java @@ -9,6 +9,7 @@ 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; @@ -18,11 +19,15 @@ import org.junit.jupiter.params.provider.ValueSource; 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; 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; @@ -54,6 +59,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()); } } @@ -295,6 +301,469 @@ public void testDurationNanoSecondsLiteralTransform(Long value) { } } + private static Stream createLegacyDecimalLiteralParams() { + 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")), + 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), + 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)); + } + + 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("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("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) { + Literal expr = Literal.ofDecimal(type, value); + try (Table t = new Table.TestBuilder().column(1, 2, 3).build(); + CompiledExpression compiledExpr = expr.compile(); + ColumnVector actual = compiledExpr.computeColumnJit(t); + 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 + 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))); + 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( + NullPointerException.class, + () -> new JitOperation(null, new ColumnReference(0))); + Assertions.assertThrows( + NullPointerException.class, + () -> new JitOperation(JitOperator.ADD, (JitErrorPolicy) null, + new ColumnReference(0), new ColumnReference(1))); + 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)); + } + + private static void assertJitCompileThrows(JitOperation expr) { + Assertions.assertThrows(CudfException.class, () -> { + try (CompiledExpression ignored = expr.compile()) { + } + }); + } + + @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)); + 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 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() + .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() { + try (Table t = new Table.TestBuilder() + .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, null)) { + assertColumnsAreEqual(expected, actual); + } + + JitOperation modExpr = new JitOperation(JitOperator.MOD_OVERFLOW, + JitErrorPolicy.NULLIFY, new ColumnReference(0), new ColumnReference(1)); + try (CompiledExpression compiledExpr = modExpr.compile(); + ColumnVector actual = compiledExpr.computeColumnJit(t); + ColumnVector expected = ColumnVector.fromBoxedInts(0, null, null, null, 0)) { + assertColumnsAreEqual(expected, actual); + } + } + } + + @Test + void testJitMixedConditionalTransform() { + try (Table t = new Table.TestBuilder() + .column(1, null, 3, null) + .column(10, 20, 30, 40) + .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, + coalesced, new ColumnReference(1), predicate); + try (CompiledExpression compiledExpr = expr.compile(); + ColumnVector actual = compiledExpr.computeColumnJit(t); + ColumnVector expected = ColumnVector.fromBoxedInts(10, 20, 3, 40)) { + assertColumnsAreEqual(expected, actual); + } + } + } + + private static Arguments jitCastCase( + JitOperator op, Supplier expectedFactory) { + return Arguments.of(op, expectedFactory); + } + + private static Stream createJitNumericCastParams() { + return Stream.of( + 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, 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 = expectedFactory.get()) { + assertColumnsAreEqual(expected, actual); + } + } + + private static Stream createJitDecimalCastParams() { + return Stream.of( + 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, 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 = expectedFactory.get()) { + 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); + } + } + } + private static ArrayList mapArray(T[] input, Function func) { ArrayList result = new ArrayList<>(input.length); for (T t : input) {