Add floating point type support to Parquet variant field extraction

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

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

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

@@ -373,29 +373,54 @@ __device__ device_span<uint8_t const> locate_object_field(device_span<uint8_t co
 // exact width types (not e.g. __int128) since those are the only variant primitive int headers.
 template <typename T>
 constexpr bool is_variant_int =
-  cuda::std::is_same_v<T, int8_t> || cuda::std::is_same_v<T, int16_t> ||
-  cuda::std::is_same_v<T, int32_t> || cuda::std::is_same_v<T, int64_t>;
+  cudf::is_integral_not_bool<T>() && cudf::is_signed<T>() && !cuda::std::is_same_v<T, __int128_t>;
 
-// The output types a VARIANT value can be cast to: the fixed-width signed integers plus strings.
+// 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>();
+
+// The output types a VARIANT value can be cast to: the fixed-width signed integers, floats, and
+// strings.
 template <typename T>
 constexpr bool is_variant_castable =
-  is_variant_int<T> || cuda::std::is_same_v<T, cudf::string_view>;
+  is_variant_primitive<T> || cuda::std::is_same_v<T, cudf::string_view>;
 
-// Variant primitive ints: basic_type == primitive, value_header maps INT{8,16,32,64}.
+// Maps a fixed-width output type to the VARIANT primitive type header id that encodes it.
 template <typename T>
-__device__ inline cuda::std::optional<T> decode_int(device_span<uint8_t const> enc)
+  requires(is_variant_primitive<T>)
+__device__ constexpr primitive_type primitive_type_for()
 {
-  static_assert(is_variant_int<T>, "decode_int: T must be int8_t, int16_t, int32_t, or int64_t");
+  if constexpr (cuda::std::is_same_v<T, int8_t>) {
+    return primitive_type::int8;
+  } else if constexpr (cuda::std::is_same_v<T, int16_t>) {
+    return primitive_type::int16;
+  } else if constexpr (cuda::std::is_same_v<T, int32_t>) {
+    return primitive_type::int32;
+  } else if constexpr (cuda::std::is_same_v<T, int64_t>) {
+    return primitive_type::int64;
+  } else if constexpr (cuda::std::is_same_v<T, float>) {
+    return primitive_type::float32;
+  } else if constexpr (cuda::std::is_same_v<T, double>) {
+    return primitive_type::float64;
+  } else {
+    CUDF_UNREACHABLE("primitive_type_for: T is not a supported variant primitive type");
+    return primitive_type::null;
+  }
+}
 
+/**
+ * @brief Decode a single VARIANT value blob into a fixed-width primitive of type `T`.
+ *
+ * Requires `basic_type == primitive` and a value header whose physical type id matches `T` exactly.
+ */
+template <typename T>
+__device__ inline cuda::std::optional<T> decode_primitive(device_span<uint8_t const> enc)
+{
   if (cuda::std::cmp_less(enc.size(), 1 + sizeof(T))) { return cuda::std::nullopt; }
 
-  constexpr primitive_type expected = cuda::std::is_same_v<T, int8_t>    ? primitive_type::int8
-                                      : cuda::std::is_same_v<T, int16_t> ? primitive_type::int16
-                                      : cuda::std::is_same_v<T, int32_t> ? primitive_type::int32
-                                                                         : primitive_type::int64;
-  uint8_t const value_metadata      = enc[0];
+  uint8_t const value_metadata = enc[0];
   if (variant_basic_type(value_metadata) != basic_type::primitive ||
-      variant_value_header(value_metadata) != static_cast<uint8_t>(expected)) {
+      variant_value_header(value_metadata) != static_cast<uint8_t>(primitive_type_for<T>())) {
     return cuda::std::nullopt;
   }
   return cudf::io::unaligned_load<T>(enc.data() + 1);
@@ -506,15 +531,16 @@ CUDF_KERNEL __launch_bounds__(block_size) void locate_variant_fields_kernel(
 }
 
 /**
- * @brief Per-row kernel: decode each VARIANT value blob into an integer of type `T`.
+ * @brief Per-row kernel: decode each VARIANT value blob into a fixed-width primitive of type `T`.
  *
  * Writes the decoded value to `d_output[row]` for non-null rows whose blob is a variant primitive
- * int whose physical type id matches `T` exactly (e.g. an int16 value does not decode into an
- * int32 output; there is no widening). Rows that are null, or whose value is not an exact-width
- * match for `T`, are marked null in `d_null_mask` with an output of 0.
+ * whose physical type id matches `T` exactly (e.g. an int16 value does not decode into an int32
+ * output, and a float32 value does not decode into a float64 output; there is no widening). Rows
+ * that are null, or whose value is not an exact-width match for `T`, are marked null in
+ * `d_null_mask` with an output of 0.
  */
 template <typename T>
-CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_int_kernel(
+CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_primitive_kernel(
   cudf::lists_column_device_view values, device_span<T> d_output, bitmask_type* d_null_mask)
 {
   auto const num_rows = static_cast<size_type>(d_output.size());
@@ -533,7 +559,7 @@ CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_int_kernel(
     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_int<T>(val);
+    auto const decoded = decode_primitive<T>(val);
     if (decoded.has_value()) {
       d_output[row] = *decoded;
     } else {
@@ -607,12 +633,12 @@ struct cast_variant_fn {
 
   template <typename T>
   std::unique_ptr<column> operator()()
-    requires(is_variant_int<T>)
+    requires(is_variant_primitive<T>)
   {
     rmm::device_buffer data{num_rows * sizeof(T), stream, mr};
 
     auto grid = cudf::detail::grid_1d{num_rows, block_size};
-    cast_variant_int_kernel<T><<<grid.num_blocks, block_size, 0, stream.value()>>>(
+    cast_variant_primitive_kernel<T><<<grid.num_blocks, block_size, 0, stream.value()>>>(
       values, {static_cast<T*>(data.data()), static_cast<std::size_t>(num_rows)}, d_null_mask);
     CUDF_CUDA_TRY(cudaGetLastError());
 

cpp/src/io/utilities/block_utils.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -58,7 +58,7 @@ inline __device__ T warp_reduce_pos(T pos, uint32_t t)
 }
 
 template <typename T>
-  requires(cuda::std::is_integral_v<T>)
+  requires(cuda::std::is_trivially_copyable_v<T>)
 inline __device__ T unaligned_load(uint8_t const* p)
 {
   T value;

cpp/tests/io/experimental/variant_extract_test.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -637,25 +637,54 @@ TEST_F(GetVariantFieldTest, EmptyInput)
   EXPECT_EQ(cudf::lists_column_view{got->view()}.child().type().id(), cudf::type_id::UINT8);
 }
 
+template <typename T, std::size_t M, std::size_t V>
+std::unique_ptr<cudf::column> cast_apache_primitive(avf::fixture<M, V> const& fixture)
+{
+  auto const stream = cudf::test::get_default_stream();
+  auto col          = make_apache_variant(fixture);
+  auto const value  = cudf::structs_column_view{col}.get_sliced_child(1, stream);
+  return cudf::io::parquet::experimental::cast_variant(
+    value, cudf::data_type{cudf::type_to_id<T>()}, stream);
+}
+
 struct CastVariantTest : public cudf::test::BaseFixture {};
 
 TEST_F(CastVariantTest, ApachePrimitiveInts)
 {
-  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};
+  {
+    auto got = cast_apache_primitive<int8_t>(avf::primitive_int8);
+    cudf::test::fixed_width_column_wrapper<int8_t> expected{int8_t{42}};
     CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
-  };
+  }
+  {
+    auto got = cast_apache_primitive<int16_t>(avf::primitive_int16);
+    cudf::test::fixed_width_column_wrapper<int16_t> expected{int16_t{1234}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  }
+  {
+    auto got = cast_apache_primitive<int32_t>(avf::primitive_int32);
+    cudf::test::fixed_width_column_wrapper<int32_t> expected{int32_t{123456}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  }
+  {
+    auto got = cast_apache_primitive<int64_t>(avf::primitive_int64);
+    cudf::test::fixed_width_column_wrapper<int64_t> expected{int64_t{1234567890123456789LL}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  }
+}
 
-  cast(avf::primitive_int8, int8_t{42});
-  cast(avf::primitive_int16, int16_t{1234});
-  cast(avf::primitive_int32, int32_t{123456});
-  cast(avf::primitive_int64, int64_t{1234567890123456789LL});
+TEST_F(CastVariantTest, ApachePrimitiveFloats)
+{
+  {
+    auto got = cast_apache_primitive<float>(avf::primitive_float);
+    cudf::test::fixed_width_column_wrapper<float> expected{float{1234567936.0f}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  }
+  {
+    auto got = cast_apache_primitive<double>(avf::primitive_double);
+    cudf::test::fixed_width_column_wrapper<double> expected{double{1234567890.1234}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected);
+  }
 }
 
 TEST_F(CastVariantTest, ApacheShortString)
@@ -709,7 +738,10 @@ TEST_F(CastVariantTest, EmptyInput)
   auto const values =
     cudf::empty_like(cudf::structs_column_view{make_xyz_three_row_variant()}.child(1));
 
-  for (auto const id : {cudf::type_id::INT32, cudf::type_id::STRING}) {
+  for (auto const id : {cudf::type_id::INT32,
+                        cudf::type_id::STRING,
+                        cudf::type_id::FLOAT32,
+                        cudf::type_id::FLOAT64}) {
     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);