Add boolean type support to Parquet variant field extraction

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

@@ -70,15 +70,15 @@ namespace io::parquet::experimental {
  * `desired_type`.
  *
  * @param values `list<uint8>` column of VARIANT-encoded value bytes
- * @param desired_type Target cuDF type (`STRING`, `INT8`/`INT16`/`INT32`/`INT64`, or
- *        `FLOAT32`/`FLOAT64`)
+ * @param desired_type Target cuDF type (`STRING`, `INT8`/`INT16`/`INT32`/`INT64`,
+ *        `FLOAT32`/`FLOAT64`, or `BOOL8`)
  * @param stream CUDA stream
  * @param mr Device memory resource
  * @return Typed column decoded from the VARIANT value blobs
  *
  * @throws std::invalid_argument if `values` is not a `list<uint8>` column, or if `desired_type`
- *         is not one of the supported types (`STRING`, `INT8`/`INT16`/`INT32`/`INT64`, or
- *         `FLOAT32`/`FLOAT64`)
+ *         is not one of the supported types (`STRING`, `INT8`/`INT16`/`INT32`/`INT64`,
+ *         `FLOAT32`/`FLOAT64`, or `BOOL8`)
  */
 [[nodiscard]] std::unique_ptr<column> cast_variant(
   column_view const& values,
@@ -94,8 +94,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`, or
- *        `FLOAT32`/`FLOAT64`
+ * @param desired_type Target type: `STRING`, `INT8`/`INT16`/`INT32`/`INT64`,
+ *        `FLOAT32`/`FLOAT64`, or `BOOL8`
  * @param stream CUDA stream
  * @param mr Device memory resource
  * @return Column of `desired_type`

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

@@ -34,6 +34,7 @@
 #include <rmm/exec_policy.hpp>
 
 #include <cuda/functional>
+#include <cuda/iterator>
 #include <cuda/numeric>
 #include <cuda/std/cstring>
 #include <cuda/std/limits>
@@ -415,17 +416,17 @@ constexpr bool is_variant_int =
 
 // The fixed-width primitive types (signed integers and floats) a VARIANT value can be decoded into.
 template <typename T>
-constexpr bool is_variant_primitive = is_variant_int<T> || cudf::is_floating_point<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, and
-// strings.
+// The output types a VARIANT value can be cast to: the fixed-width signed integers, floats, bool,
+// and strings.
 template <typename T>
-constexpr bool is_variant_castable =
-  is_variant_primitive<T> || cuda::std::is_same_v<T, cudf::string_view>;
+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>;
 
 // Maps a fixed-width output type to the VARIANT primitive type header id that encodes it.
 template <typename T>
-  requires(is_variant_primitive<T>)
+  requires(is_variant_numerical<T>)
 __device__ constexpr primitive_type primitive_type_for()
 {
   if constexpr (cuda::std::is_same_v<T, int8_t>) {
@@ -464,6 +465,23 @@ __device__ inline cuda::std::optional<T> decode_primitive(device_span<uint8_t co
   return cudf::io::unaligned_load<T>(enc.data() + 1);
 }
 
+/**
+ * @brief Decode a single VARIANT value blob into a bool.
+ *
+ * Boolean values carry no payload: the distinction between true and false is encoded entirely in
+ * the primitive type header (`boolean_true` vs `boolean_false`).
+ */
+__device__ inline cuda::std::optional<bool> decode_bool(device_span<uint8_t const> enc)
+{
+  if (enc.empty()) { return cuda::std::nullopt; }
+  uint8_t const value_metadata = enc[0];
+  if (decode_basic_type(value_metadata) != basic_type::PRIMITIVE) { return cuda::std::nullopt; }
+  auto const value_header = variant_value_header(value_metadata);
+  if (value_header == static_cast<uint8_t>(primitive_type::BOOLEAN_TRUE)) { return true; }
+  if (value_header == static_cast<uint8_t>(primitive_type::BOOLEAN_FALSE)) { return false; }
+  return cuda::std::nullopt;
+}
+
 // Parse an array-index step token of the form "[<N>]" into its zero-based index. Returns nullopt
 // for any malformed token or an index that does not fit in `size_type` (such an index is out of
 // range for any array, so the caller treats it as a missing element).
@@ -547,20 +565,21 @@ __device__ cuda::std::optional<device_span<uint8_t const>> decode_string(
   return cuda::std::nullopt;
 }
 
+__device__ device_span<uint8_t const> list_row_span(cudf::lists_column_device_view const& col,
+                                                    size_type row)
+{
+  auto const begin = col.offset_at(row);
+  auto const end   = col.offset_at(row + 1);
+  return {col.child().data<uint8_t>() + begin, static_cast<std::size_t>(end - begin)};
+}
+
 // Returns the metadata and value list bytes for a given row from device views
 __device__ cuda::std::pair<device_span<uint8_t const>, device_span<uint8_t const>>
 metadata_and_value_at(cudf::lists_column_device_view const& metadata,
                       cudf::lists_column_device_view const& values,
                       size_type row)
 {
-  auto const meta_begin = metadata.offset_at(row);
-  auto const meta_end   = metadata.offset_at(row + 1);
-  auto const val_begin  = values.offset_at(row);
-  auto const val_end    = values.offset_at(row + 1);
-  return {
-    {metadata.child().data<uint8_t>() + meta_begin,
-     static_cast<std::size_t>(meta_end - meta_begin)},
-    {values.child().data<uint8_t>() + val_begin, static_cast<std::size_t>(val_end - val_begin)}};
+  return {list_row_span(metadata, row), list_row_span(values, row)};
 }
 
 constexpr int block_size = 256;
@@ -629,11 +648,7 @@ CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_primitive_kernel(
       continue;
     }
 
-    auto const val_begin = values.offset_at(row);
-    auto const val_end   = values.offset_at(row + 1);
-    auto const val_child = values.child();
-    device_span<uint8_t const> const val{val_child.data<uint8_t>() + val_begin,
-                                         static_cast<std::size_t>(val_end - val_begin)};
+    auto const val = list_row_span(values, row);
 
     auto const decoded = decode_primitive<T>(val);
     if (decoded.has_value()) {
@@ -667,11 +682,7 @@ struct cast_variant_string_fn {
       return;
     }
 
-    auto const val_begin = d_values.offset_at(row);
-    auto const val_end   = d_values.offset_at(row + 1);
-    auto const val_child = d_values.child();
-    device_span<uint8_t const> const val{val_child.data<uint8_t>() + val_begin,
-                                         static_cast<std::size_t>(val_end - val_begin)};
+    auto const val = list_row_span(d_values, row);
 
     auto const str = decode_string(val);
     if (!str) {
@@ -709,7 +720,7 @@ struct cast_variant_fn {
 
   template <typename T>
   std::unique_ptr<column> operator()()
-    requires(is_variant_primitive<T>)
+    requires(is_variant_numerical<T>)
   {
     rmm::device_buffer data{num_rows * sizeof(T), stream, mr};
 
@@ -727,6 +738,35 @@ struct cast_variant_fn {
                                     null_count);
   }
 
+  template <typename T>
+  std::unique_ptr<column> operator()()
+    requires(cuda::std::is_same_v<T, bool>)
+  {
+    rmm::device_buffer data{num_rows * sizeof(bool), stream, mr};
+
+    thrust::transform(
+      rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+      cuda::counting_iterator<size_type>(0),
+      cuda::counting_iterator<size_type>(num_rows),
+      static_cast<bool*>(data.data()),
+      [values = this->values, d_null_mask = this->d_null_mask] __device__(size_type row) -> bool {
+        if (!cudf::bit_is_set(d_null_mask, row)) { return false; }
+        auto const val     = list_row_span(values, row);
+        auto const decoded = decode_bool(val);
+        if (decoded.has_value()) { return *decoded; }
+        cudf::clear_bit(d_null_mask, row);
+        return false;
+      });
+
+    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, cudf::string_view>)

cpp/tests/io/experimental/variant_extract_test.cpp

@@ -958,14 +958,93 @@ TEST_F(CastVariantTest, ApachePrimitiveInts)
 
 TEST_F(CastVariantTest, ApachePrimitiveFloats)
 {
+  auto stream     = cudf::test::get_default_stream();
+  auto const cast = [&](auto const& fixture, auto expected_val) {
+    using T          = decltype(expected_val);
+    auto col         = make_apache_variant(fixture);
+    auto const value = cudf::structs_column_view{col}.get_sliced_child(1, stream);
+    auto got         = cudf::io::parquet::experimental::cast_variant(
+      value, cudf::data_type{cudf::type_to_id<T>()}, stream);
+    cudf::test::fixed_width_column_wrapper<T> expected{expected_val};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  };
+
+  cast(avf::primitive_float, float{1234567936.0f});
+  cast(avf::primitive_double, double{1234567890.1234});
+}
+
+TEST_F(CastVariantTest, ApachePrimitiveBooleans)
+{
+  auto stream     = cudf::test::get_default_stream();
+  auto const cast = [&](auto const& fixture, bool expected_val) {
+    auto col         = make_apache_variant(fixture);
+    auto const value = cudf::structs_column_view{col}.get_sliced_child(1, stream);
+    auto got         = cudf::io::parquet::experimental::cast_variant(
+      value, cudf::data_type{cudf::type_id::BOOL8}, stream);
+    cudf::test::fixed_width_column_wrapper<bool> expected{expected_val};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  };
+
+  cast(avf::primitive_boolean_true, true);
+  cast(avf::primitive_boolean_false, false);
+
+  // Null variant value must cast to a null BOOL8, not false.
   {
-    auto got = cast_apache_primitive<float>(avf::primitive_float);
-    cudf::test::fixed_width_column_wrapper<float> expected{float{1234567936.0f}};
+    auto col         = make_apache_variant(avf::primitive_null);
+    auto const value = cudf::structs_column_view{col}.get_sliced_child(1, stream);
+    auto got         = cudf::io::parquet::experimental::cast_variant(
+      value, cudf::data_type{cudf::type_id::BOOL8}, stream);
+    cudf::test::fixed_width_column_wrapper<bool> expected({false}, {false});
     CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
   }
+
+  // Sliced multi-row column exercises grid-stride paths with a non-zero slice offset.
+  // num_rows / slice range is chosen so the sliced window (slice_end - slice_beg = 512)
+  // spans more than one cast_variant_bool_kernel block (block_size = 256).
   {
-    auto got = cast_apache_primitive<double>(avf::primitive_double);
-    cudf::test::fixed_width_column_wrapper<double> expected{double{1234567890.1234}};
+    constexpr int num_rows  = 516;
+    constexpr int slice_beg = 3;
+    constexpr int slice_end = 515;
+
+    std::vector<uint8_t> const true_bytes{avf::primitive_boolean_true.value.begin(),
+                                          avf::primitive_boolean_true.value.end()};
+    std::vector<uint8_t> const false_bytes{avf::primitive_boolean_false.value.begin(),
+                                           avf::primitive_boolean_false.value.end()};
+    std::vector<uint8_t> const null_bytes{avf::primitive_null.value.begin(),
+                                          avf::primitive_null.value.end()};
+    std::vector<uint8_t> const meta_bytes{avf::primitive_boolean_true.metadata.begin(),
+                                          avf::primitive_boolean_true.metadata.end()};
+
+    std::vector<std::vector<uint8_t>> metas(num_rows, meta_bytes);
+    std::vector<std::vector<uint8_t>> vals(num_rows);
+    std::vector<bool> exp_vals(num_rows);
+    std::vector<bool> exp_valid(num_rows);
+
+    for (int i = 0; i < num_rows; ++i) {
+      int const pat = i % 3;
+      if (pat == 0) {
+        vals[i]      = true_bytes;
+        exp_vals[i]  = true;
+        exp_valid[i] = true;
+      } else if (pat == 1) {
+        vals[i]      = false_bytes;
+        exp_vals[i]  = false;
+        exp_valid[i] = true;
+      } else {
+        vals[i]      = null_bytes;
+        exp_vals[i]  = false;
+        exp_valid[i] = false;
+      }
+    }
+
+    auto col          = wrap_multi_row_variant(metas, vals);
+    auto const sliced = cudf::slice(col, {slice_beg, slice_end}).front();
+    auto const value  = cudf::structs_column_view{sliced}.get_sliced_child(1, stream);
+    auto got          = cudf::io::parquet::experimental::cast_variant(
+      value, cudf::data_type{cudf::type_id::BOOL8}, stream);
+
+    cudf::test::fixed_width_column_wrapper<bool> expected(
+      exp_vals.begin() + slice_beg, exp_vals.begin() + slice_end, exp_valid.begin() + slice_beg);
     CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
   }
 }
@@ -1024,7 +1103,8 @@ TEST_F(CastVariantTest, EmptyInput)
   for (auto const id : {cudf::type_id::INT32,
                         cudf::type_id::STRING,
                         cudf::type_id::FLOAT32,
-                        cudf::type_id::FLOAT64}) {
+                        cudf::type_id::FLOAT64,
+                        cudf::type_id::BOOL8}) {
     auto got = cudf::io::parquet::experimental::cast_variant(*values, cudf::data_type{id}, stream);
     EXPECT_EQ(got->type().id(), id);
     EXPECT_EQ(got->size(), 0);