Add decimal support to VARIANT casting

cpp/benchmarks/io/parquet/experimental/variant/extract.cpp

@@ -35,7 +35,16 @@ using cudf::io::parquet::experimental::variant_basic_type;
 using cudf::io::parquet::experimental::variant_primitive_type;
 
 // The leaf value type exercised by the benchmark (nvbench "type" string axis).
-enum class bench_variant_type : uint8_t { INT32, FLOAT, BOOL, STRING, ARRAY };
+enum class bench_variant_type : uint8_t {
+  INT32,
+  FLOAT,
+  BOOL,
+  STRING,
+  ARRAY,
+  DECIMAL32,
+  DECIMAL64,
+  DECIMAL128
+};
 
 bench_variant_type parse_bench_variant_type(std::string const& type_str)
 {
@@ -44,9 +53,15 @@ bench_variant_type parse_bench_variant_type(std::string const& type_str)
   if (type_str == "bool") { return bench_variant_type::BOOL; }
   if (type_str == "string") { return bench_variant_type::STRING; }
   if (type_str == "array") { return bench_variant_type::ARRAY; }
+  if (type_str == "decimal32") { return bench_variant_type::DECIMAL32; }
+  if (type_str == "decimal64") { return bench_variant_type::DECIMAL64; }
+  if (type_str == "decimal128") { return bench_variant_type::DECIMAL128; }
   CUDF_FAIL("Unrecognized benchmark type: " + type_str);
 }
 
+// Shared by the encoded leaf values and the target column, so the cast measures decoding only.
+constexpr uint8_t bench_decimal_scale = 2;
+
 // Compose a value-metadata header byte from a basic type and its 6-bit value_header.
 // See cpp/tests/io/experimental/variant_extract_test.cpp for the header byte layout.
 constexpr uint8_t make_variant_header(variant_basic_type basic, uint8_t value_header)
@@ -168,6 +183,27 @@ std::vector<uint8_t> build_leaf_value(bench_variant_type type)
       out.insert(out.end(), s.begin(), s.end());
       return out;
     }
+    case bench_variant_type::DECIMAL32: {
+      std::vector<uint8_t> out{make_variant_primitive_header(variant_primitive_type::DECIMAL4),
+                               bench_decimal_scale};
+      append_le(out, 1234u, 4);
+      return out;
+    }
+    case bench_variant_type::DECIMAL64: {
+      std::vector<uint8_t> out{make_variant_primitive_header(variant_primitive_type::DECIMAL8),
+                               bench_decimal_scale};
+      append_le(out, (static_cast<uint64_t>(1234u) << 32) | 5678u, 8);
+      return out;
+    }
+    case bench_variant_type::DECIMAL128: {
+      // The value needs more than 64 bits, so the decode is not measured on an all-zero high half.
+      std::vector<uint8_t> out{make_variant_primitive_header(variant_primitive_type::DECIMAL16),
+                               bench_decimal_scale};
+      auto const unscaled = (static_cast<__uint128_t>(1234u) << 64) | 5678u;
+      append_le(out, static_cast<uint64_t>(unscaled), 8);
+      append_le(out, static_cast<uint64_t>(unscaled >> 64), 8);
+      return out;
+    }
     case bench_variant_type::ARRAY: {
       // VARIANT array of two INT32 values [42, 99]; element [1] is accessed in the benchmark.
       // 2 elements, offsets [0, 5, 10], then INT32(42) and INT32(99) (5 bytes each).
@@ -368,6 +404,12 @@ cudf::data_type get_target_type(bench_variant_type type)
     case bench_variant_type::FLOAT: return cudf::data_type{cudf::type_id::FLOAT32};
     case bench_variant_type::BOOL: return cudf::data_type{cudf::type_id::BOOL8};
     case bench_variant_type::STRING: return cudf::data_type{cudf::type_id::STRING};
+    case bench_variant_type::DECIMAL32:
+      return cudf::data_type{cudf::type_id::DECIMAL32, -bench_decimal_scale};
+    case bench_variant_type::DECIMAL64:
+      return cudf::data_type{cudf::type_id::DECIMAL64, -bench_decimal_scale};
+    case bench_variant_type::DECIMAL128:
+      return cudf::data_type{cudf::type_id::DECIMAL128, -bench_decimal_scale};
     // "array": element access yields INT32.
     case bench_variant_type::INT32:
     case bench_variant_type::ARRAY: return cudf::data_type{cudf::type_id::INT32};
@@ -435,7 +477,8 @@ static void bench_variant_cast(nvbench::state& state)
 NVBENCH_BENCH(bench_variant_cast)
   .set_name("bench_variant_cast")
   .add_int64_axis("num_rows", {32768, 262144, 2097152})
-  .add_string_axis("type", {"string", "float", "bool", "int32_t"})
+  .add_string_axis("type",
+                   {"string", "float", "bool", "int32_t", "decimal32", "decimal64", "decimal128"})
   .add_int64_axis("hit_rate", {20, 80});
 
 // Benchmarks get_variant_field with varying path depth (nesting >= 1). Casting is exercised

cpp/include/cudf/io/experimental/variant.hpp

@@ -77,9 +77,13 @@ namespace io::parquet::experimental {
  * A null value is produced when the input row is null or the encoded type does not match
  * `desired_type`.
  *
+ * For a decimal `desired_type`, every encoded width is accepted and each value is rescaled from its
+ * own encoded scale to `desired_type.scale()`, truncating toward zero; a value that no longer fits
+ * produces a null row with `variant_operation_status::OVERFLOW`.
+ *
  * @param values `list<uint8>` column of VARIANT-encoded value bytes
  * @param desired_type Target cuDF type (`STRING`, `INT8`/`INT16`/`INT32`/`INT64`,
- *        `FLOAT32`/`FLOAT64`, or `BOOL8`)
+ *        `FLOAT32`/`FLOAT64`, `BOOL8`, or `DECIMAL32`/`DECIMAL64`/`DECIMAL128`)
  * @param status Optional in-out parameter, `variant_operation_status` values, one per row. Must be
  *        non-nullable, `UINT8`, and have the same row count as `values`. On input, its existing
  *        values are treated as status from a prior `get_variant_field` call: rows already marked
@@ -92,8 +96,8 @@ namespace io::parquet::experimental {
  *
  * @throws std::invalid_argument if `values` is not a `list<uint8>` column; if `desired_type`
  *         is not one of the supported types (`STRING`, `INT8`/`INT16`/`INT32`/`INT64`,
- *         `FLOAT32`/`FLOAT64`, or `BOOL8`); or if `status` is provided but is nullable, not
- *         `UINT8`, or has a different row count than `values`
+ *         `FLOAT32`/`FLOAT64`, `BOOL8`, or `DECIMAL32`/`DECIMAL64`/`DECIMAL128`); or if `status`
+ *         is provided but is nullable, not `UINT8`, or has a different row count than `values`
  */
 [[nodiscard]] std::unique_ptr<column> cast_variant(
   column_view const& values,
@@ -111,7 +115,8 @@ namespace io::parquet::experimental {
  * @param variant_column Struct column (VARIANT materialization)
  * @param path JSONPath-like path string (see `get_variant_field` for syntax)
  * @param desired_type Target type: `STRING`, `INT8`/`INT16`/`INT32`/`INT64`,
- *        `FLOAT32`/`FLOAT64`, or `BOOL8`
+ *        `FLOAT32`/`FLOAT64`, `BOOL8`, or `DECIMAL32`/`DECIMAL64`/`DECIMAL128` (see `cast_variant`
+ *        for decimal rescaling)
  * @param status Optional. When provided, filled with `variant_operation_status` values, one per
  *               row. Must be non-nullable, `UINT8`, and have the same row count as
  *               `variant_column`

cpp/src/io/parquet/experimental/variant_extract.cu

@@ -525,11 +525,12 @@ constexpr bool is_variant_int =
 template <typename T>
 constexpr bool is_variant_numerical = is_variant_int<T> || cudf::is_floating_point<T>();
 
-// The output types a VARIANT value can be cast to: the fixed-width signed integers, floats, bool,
-// and strings.
+// The output types a VARIANT value can be cast to: the fixed-width signed integers, floats,
+// decimals, bool, and strings.
 template <typename T>
-constexpr bool is_variant_castable = is_variant_numerical<T> || cuda::std::is_same_v<T, bool> ||
-                                     cuda::std::is_same_v<T, cudf::string_view>;
+constexpr bool is_variant_castable =
+  is_variant_numerical<T> || cudf::is_fixed_point<T>() || cuda::std::is_same_v<T, bool> ||
+  cuda::std::is_same_v<T, cudf::string_view>;
 
 // Maps a fixed-width output type to the VARIANT primitive type header id that encodes it.
 template <typename T>
@@ -798,6 +799,129 @@ __device__ op_status cast_status_for_primitive(device_span<uint8_t const> val)
                                              : op_status::MALFORMED_VARIANT;
 }
 
+// The spec allows a scale in [0, 38] for every decimal width.
+constexpr int variant_decimal_max_scale = 38;
+
+// Multiply `value` by 10^exp, or return nullopt if the result does not fit in `__int128_t`.
+__device__ cuda::std::optional<__int128_t> multiply_pow10(__int128_t value, int exp)
+{
+  constexpr __int128_t max_over_10 = cuda::std::numeric_limits<__int128_t>::max() / 10;
+  constexpr __int128_t min_over_10 = cuda::std::numeric_limits<__int128_t>::min() / 10;
+  for (int i = 0; i < exp && value != 0; ++i) {
+    if (value > max_over_10 || value < min_over_10) { return cuda::std::nullopt; }
+    value *= 10;
+  }
+  return value;
+}
+
+// Divide `value` by 10^exp, truncating toward zero. Iterating keeps an `exp` whose power of ten
+// would itself overflow exact, since every value truncates to zero there.
+__device__ __int128_t divide_pow10(__int128_t value, int exp)
+{
+  for (int i = 0; i < exp && value != 0; ++i) {
+    value /= 10;
+  }
+  return value;
+}
+
+// Returns 0 for a `ptype` that is not a decimal.
+__device__ int variant_decimal_unscaled_width(primitive_type ptype)
+{
+  switch (ptype) {
+    case primitive_type::DECIMAL4: return 4;
+    case primitive_type::DECIMAL8: return 8;
+    case primitive_type::DECIMAL16: return 16;
+    default: return 0;
+  }
+}
+
+/**
+ * @brief Decode a single VARIANT decimal value blob into the representation of a cuDF fixed-point
+ * type, rescaled to `desired_scale`.
+ *
+ * Encoded as the value metadata byte, a one-byte scale (fractional digit count), then the
+ * little-endian two's-complement unscaled integer. Any of the three widths decodes into any `Rep`
+ * whose range fits the rescaled value; digits below `desired_scale` truncate toward zero.
+ *
+ * @return The rescaled representation, valid only when the returned status is `SUCCESS`
+ */
+template <typename Rep>
+__device__ cuda::std::pair<Rep, op_status> decode_decimal(device_span<uint8_t const> enc,
+                                                          int desired_scale)
+{
+  auto const fail = [](op_status status) { return cuda::std::pair<Rep, op_status>{Rep{}, status}; };
+
+  if (enc.empty()) { return fail(op_status::MALFORMED_VARIANT); }
+  if (is_variant_null(enc)) { return fail(op_status::VARIANT_NULL); }
+  if (decode_basic_type(enc[0]) != basic_type::PRIMITIVE) { return fail(op_status::TYPE_MISMATCH); }
+
+  auto const ptype = static_cast<primitive_type>(variant_value_header(enc[0]));
+  auto const width = variant_decimal_unscaled_width(ptype);
+  if (width == 0) {
+    return fail(is_recognized_primitive_type(ptype) ? op_status::TYPE_MISMATCH
+                                                    : op_status::MALFORMED_VARIANT);
+  }
+
+  constexpr size_type scale_bytes = 1;
+  if (cuda::std::cmp_less(enc.size(), variant_header_bytes + scale_bytes + width)) {
+    return fail(op_status::MALFORMED_VARIANT);
+  }
+  int const encoded_scale = enc[variant_header_bytes];
+  if (encoded_scale > variant_decimal_max_scale) { return fail(op_status::MALFORMED_VARIANT); }
+
+  auto const* unscaled_data = enc.data() + variant_header_bytes + scale_bytes;
+  auto const unscaled       = [&]() -> __int128_t {
+    switch (width) {
+      case 4: return cudf::io::unaligned_load<int32_t>(unscaled_data);
+      case 8: return cudf::io::unaligned_load<int64_t>(unscaled_data);
+      default: return cudf::io::unaligned_load<__int128_t>(unscaled_data);
+    }
+  }();
+
+  // The encoded value is `unscaled * 10^-encoded_scale`; the output is `rep * 10^desired_scale`.
+  auto const shift = -encoded_scale - desired_scale;
+  __int128_t rescaled{};
+  if (shift >= 0) {
+    auto const scaled = multiply_pow10(unscaled, shift);
+    if (!scaled.has_value()) { return fail(op_status::OVERFLOW); }
+    rescaled = scaled.value();
+  } else {
+    rescaled = divide_pow10(unscaled, -shift);
+  }
+
+  if constexpr (!cuda::std::is_same_v<Rep, __int128_t>) {
+    if (rescaled < static_cast<__int128_t>(cuda::std::numeric_limits<Rep>::min()) ||
+        rescaled > static_cast<__int128_t>(cuda::std::numeric_limits<Rep>::max())) {
+      return fail(op_status::OVERFLOW);
+    }
+  }
+  return {static_cast<Rep>(rescaled), op_status::SUCCESS};
+}
+
+/**
+ * @brief Shared per-row preamble for the cast paths: decides whether row `row` should be decoded.
+ *
+ * A skipped row has its null bit cleared and its status recorded here, but its output element is
+ * left to the caller, which alone knows the output type.
+ *
+ * @return True when the row's value blob should be decoded
+ */
+__device__ bool should_decode_row(size_type row, bitmask_type* d_null_mask, op_status* d_status)
+{
+  if (d_status == nullptr) { return cudf::bit_is_set(d_null_mask, row); }
+
+  // Status column is always non-nullable; ROW_NULL replaces the null bit.
+  if (d_status[row] != op_status::SUCCESS) {
+    if (cudf::bit_is_set(d_null_mask, row)) { cudf::clear_bit(d_null_mask, row); }
+    return false;
+  }
+  if (!cudf::bit_is_set(d_null_mask, row)) {
+    d_status[row] = op_status::ROW_NULL;
+    return false;
+  }
+  return true;
+}
+
 /**
  * @brief Per-row kernel: decode each VARIANT value blob into a fixed-width primitive of type `T`.
  *
@@ -823,24 +947,9 @@ CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_primitive_kernel(
   auto const stride   = cudf::detail::grid_1d::grid_stride<block_size>();
 
   for (auto row = tid; row < num_rows; row += stride) {
-    if (d_status != nullptr) {
-      // Status column is always non-nullable; row_null replaces the null bit.
-      auto const s = d_status[row];
-      if (s != op_status::SUCCESS) {
-        d_output[row] = T{};
-        if (cudf::bit_is_set(d_null_mask, row)) { cudf::clear_bit(d_null_mask, row); }
-        continue;
-      }
-      if (!cudf::bit_is_set(d_null_mask, row)) {
-        d_output[row] = T{};
-        d_status[row] = op_status::ROW_NULL;
-        continue;
-      }
-    } else {
-      if (!cudf::bit_is_set(d_null_mask, row)) {
-        d_output[row] = T{};
-        continue;
-      }
+    if (!should_decode_row(row, d_null_mask, d_status)) {
+      d_output[row] = T{};
+      continue;
     }
 
     auto const val     = list_row_span(values, row);
@@ -856,6 +965,40 @@ CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_primitive_kernel(
   }
 }
 
+/**
+ * @brief Per-row kernel: decode each VARIANT decimal value blob into a fixed-point representation
+ * of type `Rep`, rescaled to `desired_scale`. Same null and status protocol as
+ * `cast_variant_primitive_kernel`.
+ */
+template <typename Rep>
+CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_decimal_kernel(
+  cudf::lists_column_device_view values,
+  device_span<Rep> d_output,
+  int desired_scale,
+  bitmask_type* d_null_mask,
+  op_status* d_status)  // nullptr when no status was requested
+{
+  auto const num_rows = static_cast<size_type>(d_output.size());
+  auto const tid      = cudf::detail::grid_1d::global_thread_id<block_size>();
+  auto const stride   = cudf::detail::grid_1d::grid_stride<block_size>();
+
+  for (auto row = tid; row < num_rows; row += stride) {
+    if (!should_decode_row(row, d_null_mask, d_status)) {
+      d_output[row] = Rep{};
+      continue;
+    }
+
+    auto const [value, status] = decode_decimal<Rep>(list_row_span(values, row), desired_scale);
+    if (status == op_status::SUCCESS) {
+      d_output[row] = value;
+    } else {
+      d_output[row] = Rep{};
+      cudf::clear_bit(d_null_mask, row);
+    }
+    if (d_status != nullptr) { d_status[row] = status; }
+  }
+}
+
 __device__ op_status cast_status_for_bool(device_span<uint8_t const> val)
 {
   if (val.empty()) { return op_status::MALFORMED_VARIANT; }
@@ -997,6 +1140,28 @@ struct cast_variant_fn {
                                     null_count);
   }
 
+  template <typename T>
+  std::unique_ptr<column> operator()()
+    requires(cudf::is_fixed_point<T>())
+  {
+    using Rep = typename T::rep;
+    rmm::device_buffer data{num_rows * sizeof(Rep), stream, mr};
+    auto const grid = cudf::detail::grid_1d{num_rows, block_size};
+    auto const d_out =
+      device_span<Rep>{static_cast<Rep*>(data.data()), static_cast<std::size_t>(num_rows)};
+    cast_variant_decimal_kernel<Rep><<<grid.num_blocks, block_size, 0, stream.get()>>>(
+      values, d_out, desired_type.scale(), d_null_mask, d_status);
+    CUDF_CUDA_TRY(cudaGetLastError());
+
+    auto const null_count =
+      num_rows - cudf::detail::count_set_bits(d_null_mask, 0, num_rows, stream);
+    return std::make_unique<column>(desired_type,
+                                    num_rows,
+                                    std::move(data),
+                                    null_count > 0 ? std::move(null_mask) : rmm::device_buffer{},
+                                    null_count);
+  }
+
   template <typename T>
   std::unique_ptr<column> operator()()
     requires(cuda::std::is_same_v<T, bool>)
@@ -1012,14 +1177,16 @@ struct cast_variant_fn {
                       d_out = static_cast<bool*>(data.data()),
                       dnm   = this->d_null_mask,
                       dp_s] __device__(size_type row) {
+                       if (!should_decode_row(row, dnm, dp_s)) {
+                         d_out[row] = false;
+                         return;
+                       }
+                       // The row is known live here, so a decode failure always clears its bit.
                        auto const fail = [&](op_status s) {
                          d_out[row] = false;
-                         if (cudf::bit_is_set(dnm, row)) { cudf::clear_bit(dnm, row); }
+                         cudf::clear_bit(dnm, row);
                          if (dp_s) { dp_s[row] = s; }
                        };
-                       if (dp_s and dp_s[row] != op_status::SUCCESS) { return fail(dp_s[row]); }
-                       // Status column is always non-nullable; ROW_NULL replaces the null bit.
-                       if (!cudf::bit_is_set(dnm, row)) { return fail(op_status::ROW_NULL); }
                        auto const val     = list_row_span(vals, row);
                        auto const decoded = decode_bool(val);
                        if (!decoded) { return fail(cast_status_for_bool(val)); }
@@ -1262,7 +1429,10 @@ std::unique_ptr<column> cast_variant(column_view const& values,
     case type_id::FLOAT32:
     case type_id::FLOAT64:
     case type_id::BOOL8:
-    case type_id::STRING: break;
+    case type_id::STRING:
+    case type_id::DECIMAL32:
+    case type_id::DECIMAL64:
+    case type_id::DECIMAL128: break;
     default: CUDF_FAIL("unsupported type for variant cast", std::invalid_argument);
   }