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

@@ -438,11 +438,11 @@ constexpr bool is_variant_int =
 template <typename T>
 constexpr bool is_variant_primitive = 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_primitive<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>
@@ -485,6 +485,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.size() < 1) { return cuda::std::nullopt; }
+  uint8_t const value_metadata = enc[0];
+  if (variant_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).
@@ -666,6 +683,42 @@ CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_primitive_kernel(
   }
 }
 
+/**
+ * @brief Per-row kernel: decode each VARIANT value blob into a bool.
+ *
+ * Boolean values are encoded as a single-byte `primitive_type::boolean_true` or
+ * `primitive_type::boolean_false` header with no payload. Rows that are null, or whose value is
+ * not a boolean primitive, are marked null in `d_null_mask` with an output of false.
+ */
+CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_bool_kernel(
+  cudf::lists_column_device_view values, device_span<bool> d_output, bitmask_type* d_null_mask)
+{
+  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 (!cudf::bit_is_set(d_null_mask, row)) {
+      d_output[row] = false;
+      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 decoded = decode_bool(val);
+    if (decoded.has_value()) {
+      d_output[row] = *decoded;
+    } else {
+      d_output[row] = false;
+      cudf::clear_bit(d_null_mask, row);
+    }
+  }
+}
+
 /**
  * @brief Strings-children functor: decode each VARIANT value blob into a string.
  *
@@ -748,6 +801,26 @@ 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};
+
+    auto grid = cudf::detail::grid_1d{num_rows, block_size};
+    cast_variant_bool_kernel<<<grid.num_blocks, block_size, 0, stream.value()>>>(
+      values, {static_cast<bool*>(data.data()), static_cast<std::size_t>(num_rows)}, d_null_mask);
+    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, cudf::string_view>)

cpp/tests/io/experimental/variant_extract_test.cpp

@@ -810,14 +810,91 @@ 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.
   {
-    auto got = cast_apache_primitive<double>(avf::primitive_double);
-    cudf::test::fixed_width_column_wrapper<double> expected{double{1234567890.1234}};
+    constexpr int num_rows  = 130;
+    constexpr int slice_beg = 2;
+    constexpr int slice_end = 128;
+
+    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);
   }
 }
@@ -876,7 +953,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);