Review ANSI row IR follow-up fixes

cpp/include/cudf/operators/ansi_arithmetic.cuh

@@ -311,9 +311,7 @@ template <typename T>
 __device__ inline errc ansi_mod(T* out, T const* a, T const* b)
 {
   if (*b == 0) { return errc::DIVISION_BY_ZERO; }
-  T r = *a % *b;
-  if (r != 0 && ((r > 0) != (*b > 0))) { r += *b; }
-  *out = r;
+  *out = *a % *b;
   return errc::OK;
 }
 
@@ -329,14 +327,14 @@ __device__ inline errc ansi_mod(T* out, T const* a, T const* b)
 __device__ inline errc ansi_mod(float* out, float const* a, float const* b)
 {
   if (*b == 0) { return errc::DIVISION_BY_ZERO; }
-  *out = (*a) - (*b) * ::floorf((*a) / (*b));
+  *out = ::fmodf(*a, *b);
   return errc::OK;
 }
 
 __device__ inline errc ansi_mod(double* out, double const* a, double const* b)
 {
   if (*b == 0) { return errc::DIVISION_BY_ZERO; }
-  *out = (*a) - (*b) * ::floor((*a) / (*b));
+  *out = ::fmod(*a, *b);
   return errc::OK;
 }
 
@@ -350,7 +348,11 @@ __device__ inline errc ansi_mod(decimal<R>* out, decimal<R> const* a, decimal<R>
   if (errc e = ansi_div(&div, a, b); e != errc::OK) { return e; }
 
   decimal<R> quotient;
-  floor(&quotient, &div);
+  if (div.value() < 0) {
+    ceil(&quotient, &div);
+  } else {
+    floor(&quotient, &div);
+  }
   *out = *a - *b * quotient;
   return errc::OK;
 }

cpp/include/cudf/operators/casts.cuh

@@ -333,7 +333,8 @@ __device__ inline errc rescale(optional<decimal<R>>* out,
 {
   if (a->has_value() && new_scale->has_value()) {
     decimal<R> r;
-    rescale(&r, &a->value(), new_scale->value());
+    auto const scale = new_scale->value();
+    rescale(&r, &a->value(), &scale);
     *out = r;
   } else {
     *out = nullopt;

cpp/include/cudf/operators/logic.cuh

@@ -138,8 +138,13 @@ __device__ inline errc if_else(optional<T>* out,
                                optional<T> const* false_value,
                                optional<bool> const* pred)
 {
-  if (pred->has_value() && true_value->has_value() && false_value->has_value()) {
-    if_else<T>(&out->value(), &pred->value(), &true_value->value(), &false_value->value());
+  if (pred->has_value()) {
+    auto const* value = pred->value() ? true_value : false_value;
+    if (value->has_value()) {
+      *out = value->value();
+    } else {
+      *out = nullopt;
+    }
   } else {
     *out = nullopt;
   }

cpp/src/jit/helpers.cpp

@@ -9,8 +9,18 @@
 
 #include <jit/cache.hpp>
 
+#include <string>
+
 namespace cudf {
 namespace jit {
+namespace {
+
+bool is_jitify_deserialization_failure(jitify2::Kernel const& kernel)
+{
+  return !kernel && kernel.error().find("Deserialization failed") != std::string::npos;
+}
+
+}  // namespace
 
 bool is_scalar(cudf::size_type base_column_size, cudf::size_type column_size)
 {
@@ -111,8 +121,18 @@ jitify2::Kernel get_udf_kernel(jitify2::PreprocessedProgramData const& preproces
     options.push_back(opt);
   }
 
-  return cudf::jit::get_program_cache(preprocessed_program_data)
-    .get_kernel(kernel_name, {}, {{"cudf/detail/operation-udf.hpp", cuda_source}}, options);
+  auto& cache = cudf::jit::get_program_cache(preprocessed_program_data);
+  auto const get_kernel = [&] {
+    return cache.get_kernel(
+      kernel_name, {}, {{"cudf/detail/operation-udf.hpp", cuda_source}}, options);
+  };
+
+  auto kernel = get_kernel();
+  if (is_jitify_deserialization_failure(kernel)) {
+    // Corrupt file-cache entries otherwise poison all later runs until manual cleanup.
+    if (cache.clear()) { kernel = get_kernel(); }
+  }
+  return kernel;
 }
 
 }  // namespace jit

cpp/src/transform/jit/kernel.cu

@@ -46,29 +46,35 @@ namespace cudf {
 namespace jit {
 
 template <ops::error_mode mode, bool has_user_data, typename Args>
-__device__ void execute_transform_op(error_sink* __restrict__ error_sink,
+__device__ bool execute_transform_op(error_sink* __restrict__ error_sink,
                                      void* user_data,
                                      size_type element_idx,
                                      Args args)
 {
   // TODO: static assert invocable
   if constexpr (has_user_data) {
-    cuda::std::apply(
+    return cuda::std::apply(
       [&](auto... a) {
         if constexpr (mode == ops::error_mode::IGNORE) {
           GENERIC_TRANSFORM_OP(a...);
+          return true;
         } else {
-          error_sink->report<mode>(GENERIC_TRANSFORM_OP(a...));
+          auto const error = GENERIC_TRANSFORM_OP(a...);
+          error_sink->report<mode>(error);
+          return error == ops::errc::OK;
         }
       },
       cuda::std::tuple_cat(cuda::std::tuple{user_data, element_idx}, args));
   } else {
-    cuda::std::apply(
+    return cuda::std::apply(
       [&](auto... a) {
         if constexpr (mode == ops::error_mode::IGNORE) {
           GENERIC_TRANSFORM_OP(a...);
+          return true;
         } else {
-          error_sink->report<mode>(GENERIC_TRANSFORM_OP(a...));
+          auto const error = GENERIC_TRANSFORM_OP(a...);
+          error_sink->report<mode>(error);
+          return error == ops::errc::OK;
         }
       },
       args);
@@ -106,8 +112,9 @@ CUDF_KERNEL void transform_kernel(size_type row_size,
       auto out_ptrs =
         cuda::std::apply([&](auto&... args) { return cuda::std::tuple{&args...}; }, outs);
 
-      execute_transform_op<error_mode, has_user_data>(
+      auto const success = execute_transform_op<error_mode, has_user_data>(
         error_sink, user_data, element_idx, cuda::std::tuple_cat(out_ptrs, ins));
+      if (!success) { continue; }
 
       OutputAccessors::map([&]<typename... A>() {
         (A::assign(output_cols, element_idx, cuda::std::get<A::index>(outs)), ...);
@@ -127,13 +134,18 @@ CUDF_KERNEL void transform_kernel(size_type row_size,
       auto out_ptrs =
         cuda::std::apply([&](auto&... args) { return cuda::std::tuple{&args...}; }, outs);
 
-      execute_transform_op<error_mode, has_user_data>(
+      auto const success = execute_transform_op<error_mode, has_user_data>(
         error_sink, user_data, element_idx, cuda::std::tuple_cat(out_ptrs, ins));
 
       OutputAccessors::map([&]<typename... A>() {
-        (A::assign(output_cols, element_idx, *cuda::std::get<A::index>(outs)), ...);
+        if (success) {
+          (A::assign(output_cols, element_idx, *cuda::std::get<A::index>(outs)), ...);
+        }
         (warp_compact_validity<A>(
-           active_mask, output_cols, element_idx, cuda::std::get<A::index>(outs).has_value()),
+           active_mask,
+           output_cols,
+           element_idx,
+           success && cuda::std::get<A::index>(outs).has_value()),
          ...);
       });
     }

cpp/src/transform/transform.cu

@@ -802,6 +802,25 @@ auto finalize_outputs(null_aware is_null_aware,
   return results;
 }
 
+void check_transform_error(ops::error_mode error_handling_mode,
+                           std::optional<rmm::device_scalar<jit::error_sink>> const& d_error_sink,
+                           rmm::cuda_stream_view stream)
+{
+  switch (error_handling_mode) {
+    case ops::error_mode::IGNORE: break;
+    case ops::error_mode::ANY_ROW: {
+      auto error = d_error_sink->value(stream).any_error();
+      switch (error) {
+        case ops::errc::OK: break;
+        case ops::errc::OVERFLOW: CUDF_FAIL("Overflow error in transform UDF", std::overflow_error);
+        case ops::errc::DIVISION_BY_ZERO:
+          CUDF_FAIL("Division by zero error in transform UDF", std::overflow_error);
+        default: CUDF_FAIL("Unknown error in transform UDF", std::runtime_error);
+      }
+    } break;
+  }
+}
+
 std::unique_ptr<table> execute_transform(std::string const& udf,
                                          udf_source_type source_type,
                                          ops::error_mode error_handling_mode,
@@ -851,22 +870,9 @@ std::unique_ptr<table> execute_transform(std::string const& udf,
                      stream,
                      mr);
 
-  auto finalized = finalize_outputs(is_null_aware, row_size, std::move(output_columns), stream, mr);
+  check_transform_error(error_handling_mode, d_error_sink, stream);
 
-  switch (error_handling_mode) {
-    case ops::error_mode::IGNORE: {
-    } break;
-    case ops::error_mode::ANY_ROW: {
-      auto error = d_error_sink->value(stream).any_error();
-      switch (error) {
-        case ops::errc::OK: break;
-        case ops::errc::OVERFLOW: CUDF_FAIL("Overflow error in transform UDF", std::overflow_error);
-        case ops::errc::DIVISION_BY_ZERO:
-          CUDF_FAIL("Division by zero error in transform UDF", std::overflow_error);
-        default: CUDF_FAIL("Unknown error in transform UDF", std::runtime_error);
-      }
-    } break;
-  }
+  auto finalized = finalize_outputs(is_null_aware, row_size, std::move(output_columns), stream, mr);
 
   return std::make_unique<table>(std::move(finalized));
 }

cpp/tests/ast/jit_ast_tests.cpp

@@ -47,13 +47,17 @@ struct JITIntegerArithmeticTest : public cudf::test::BaseFixture {
 template <typename T>
 struct JITSignedIntegerArithmeticTest : public JITIntegerArithmeticTest<T> {};
 
+template <typename T>
+struct JITFloatingPointArithmeticTest : public cudf::test::BaseFixture {};
+
 template <typename T>
 struct JITDecimalArithmeticTest : public JITIntegerArithmeticTest<typename T::rep> {};
 
 using SignedIntegralTypesNotBool = cudf::test::Types<int8_t, int16_t, int32_t, int64_t>;
 
 TYPED_TEST_SUITE(JITIntegerArithmeticTest, cudf::test::IntegralTypesNotBool);
 TYPED_TEST_SUITE(JITSignedIntegerArithmeticTest, SignedIntegralTypesNotBool);
+TYPED_TEST_SUITE(JITFloatingPointArithmeticTest, cudf::test::FloatingPointTypes);
 TYPED_TEST_SUITE(JITDecimalArithmeticTest, cudf::test::FixedPointTypes);
 
 TEST_F(JITExpressionTest, NullifyIf)
@@ -330,15 +334,58 @@ TYPED_TEST(JITIntegerArithmeticTest, AnsiMod)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY);
 }
 
+TYPED_TEST(JITSignedIntegerArithmeticTest, AnsiModSignedRemainder)
+{
+  using T       = TypeParam;
+  auto a        = column_wrapper<T>{{T{-5}, T{5}, T{-5}, T{5}}};
+  auto b        = column_wrapper<T>{{T{3}, T{-3}, T{-3}, T{3}}};
+  auto expected = column_wrapper<T>{{T{-2}, T{2}, T{-2}, T{2}}};
+  auto table    = cudf::table_view{{a, b}};
+  auto a_ref    = cudf::ast::column_reference(0);
+  auto b_ref    = cudf::ast::column_reference(1);
+  auto tree     = cudf::ast::tree{};
+  auto& mod     = cudf::ast::jit::ansi_mod(tree, a_ref, b_ref);
+  auto result   = cudf::compute_column_jit(table, mod);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY);
+}
+
+TYPED_TEST(JITFloatingPointArithmeticTest, AnsiMod)
+{
+  using T            = TypeParam;
+  auto a             = column_wrapper<T>{{T{3.0}, T{20.0}, T{-5.5}, T{5.5}, T{-5.5}}};
+  auto b             = column_wrapper<T>{{T{10.0}, T{7.0}, T{2.0}, T{-2.0}, T{-2.0}}};
+  auto b_fail        = column_wrapper<T>{{T{10.0}, T{0.0}, T{2.0}, T{0.0}, T{-2.0}}};
+  auto expected      = column_wrapper<T>{{T{3.0}, T{6.0}, T{-1.5}, T{1.5}, T{-1.5}}};
+  auto expected_fail = column_wrapper<T>{
+    {T{3.0}, T{0.0}, T{-1.5}, T{0.0}, T{-1.5}}, {1, 0, 1, 0, 1}};
+  auto table         = cudf::table_view{{a, b, b_fail}};
+  auto a_ref         = cudf::ast::column_reference(0);
+  auto b_ref         = cudf::ast::column_reference(1);
+  auto b_fail_ref    = cudf::ast::column_reference(2);
+  auto tree          = cudf::ast::tree{};
+  auto& mod          = cudf::ast::jit::ansi_mod(tree, a_ref, b_ref);
+  auto& mod_fail     = cudf::ast::jit::ansi_mod(tree, a_ref, b_fail_ref);
+  auto& try_mod_fail = cudf::ast::jit::ansi_try_mod(tree, a_ref, b_fail_ref);
+  auto result        = cudf::compute_column_jit(table, mod);
+  auto result_fail   = cudf::compute_column_jit(table, try_mod_fail);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY);
+
+  EXPECT_THROW(result = cudf::compute_column_jit(table, mod_fail), std::overflow_error);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY);
+}
+
 TYPED_TEST(JITDecimalArithmeticTest, AnsiMod)
 {
   using T       = TypeParam;
-  auto a        = decimal_column_wrapper<T>{{3, 20, 1, 50}, numeric::scale_type{0}};
-  auto b        = decimal_column_wrapper<T>{{10, 7, 2, 1}, numeric::scale_type{0}};
-  auto b_fail   = decimal_column_wrapper<T>{{10, 1, 20, 0}, numeric::scale_type{0}};
-  auto expected = decimal_column_wrapper<T>{{3, 6, 1, 0}, numeric::scale_type{0}};
+  auto a        = decimal_column_wrapper<T>{{3, 20, 1, 50, -5, 5, -5}, numeric::scale_type{0}};
+  auto b        = decimal_column_wrapper<T>{{10, 7, 2, 1, 3, -3, -3}, numeric::scale_type{0}};
+  auto b_fail   = decimal_column_wrapper<T>{{10, 1, 20, 0, 1, 1, 1}, numeric::scale_type{0}};
+  auto expected = decimal_column_wrapper<T>{{3, 6, 1, 0, -2, 2, -2}, numeric::scale_type{0}};
   auto expected_fail =
-    decimal_column_wrapper<T>{{3, 0, 1, 0}, {1, 1, 1, 0}, numeric::scale_type{0}};
+    decimal_column_wrapper<T>{{3, 0, 1, 0, 0, 0, 0}, {1, 1, 1, 0, 1, 1, 1}, numeric::scale_type{0}};
   auto table         = cudf::table_view{{a, b, b_fail}};
   auto a_ref         = cudf::ast::column_reference(0);
   auto b_ref         = cudf::ast::column_reference(1);

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

@@ -12,10 +12,20 @@
 public final class JitOperation extends AstExpression {
   private final JitOperator op;
   private final AstExpression[] inputs;
+  private final Integer targetScale;
 
   public JitOperation(JitOperator op, AstExpression... inputs) {
+    this(op, null, inputs);
+  }
+
+  public JitOperation(JitOperator op, int targetScale, AstExpression... inputs) {
+    this(op, Integer.valueOf(targetScale), inputs);
+  }
+
+  private JitOperation(JitOperator op, Integer targetScale, AstExpression... inputs) {
     this.op = Objects.requireNonNull(op, "op 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);
@@ -30,6 +40,10 @@ int getSerializedSize() {
     int size = ExpressionType.JIT_EXPRESSION.getSerializedSize() +
         op.getSerializedSize() +
         Byte.BYTES;
+    size += Byte.BYTES;
+    if (targetScale != null) {
+      size += Integer.BYTES;
+    }
     for (AstExpression input : inputs) {
       size += input.getSerializedSize();
     }
@@ -41,6 +55,10 @@ void serialize(ByteBuffer bb) {
     ExpressionType.JIT_EXPRESSION.serialize(bb);
     op.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);
     }

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

@@ -21,7 +21,34 @@ public enum JitOperator {
   BIT_SHIFT_RIGHT(6, 2),
   COALESCE(7, 2),
   NULLIFY_IF(8, 2),
-  PREDICATE(9, 1);
+  PREDICATE(9, 1),
+  ANSI_PRECISION_CHECK(10, 2),
+  ANSI_TRY_PRECISION_CHECK(11, 2),
+  CAST_TO_DEC32(12, 1),
+  CAST_TO_DEC64(13, 1),
+  CAST_TO_DEC128(14, 1),
+  RESCALE(15, 1),
+  ANSI_DIV(16, 2),
+  ANSI_MOD(17, 2),
+  CAST_TO_I64(18, 1),
+  ANSI_TRY_ADD(19, 2),
+  ANSI_TRY_SUB(20, 2),
+  ANSI_TRY_MUL(21, 2),
+  ANSI_TRY_DIV(22, 2),
+  ANSI_TRY_MOD(23, 2),
+  ANSI_TRY_ABS(24, 1),
+  ANSI_TRY_NEG(25, 1),
+  CAST_TO_B8(26, 1),
+  CAST_TO_I8(27, 1),
+  CAST_TO_I16(28, 1),
+  CAST_TO_I32(29, 1),
+  CAST_TO_U8(30, 1),
+  CAST_TO_U16(31, 1),
+  CAST_TO_U32(32, 1),
+  CAST_TO_U64(33, 1),
+  CAST_TO_F32(34, 1),
+  CAST_TO_F64(35, 1),
+  IF_ELSE(36, 3);
 
   private final byte nativeId;
   private final int arity;