diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index 4ee7747dc055..868c166450b4 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -71,13 +71,15 @@ namespace io::parquet::experimental { * `desired_type`. * * @param values `list` 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` 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 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` diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index aabc55449405..c166ef873819 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -373,29 +373,54 @@ __device__ device_span locate_object_field(device_span constexpr bool is_variant_int = - cuda::std::is_same_v || cuda::std::is_same_v || - cuda::std::is_same_v || cuda::std::is_same_v; + cudf::is_integral_not_bool() && cudf::is_signed() && !cuda::std::is_same_v; -// 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 +constexpr bool is_variant_primitive = is_variant_int || cudf::is_floating_point(); + +// The output types a VARIANT value can be cast to: the fixed-width signed integers, floats, and +// strings. template constexpr bool is_variant_castable = - is_variant_int || cuda::std::is_same_v; + is_variant_primitive || cuda::std::is_same_v; -// 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 -__device__ inline cuda::std::optional decode_int(device_span enc) + requires(is_variant_primitive) +__device__ constexpr primitive_type primitive_type_for() { - static_assert(is_variant_int, "decode_int: T must be int8_t, int16_t, int32_t, or int64_t"); + if constexpr (cuda::std::is_same_v) { + return primitive_type::int8; + } else if constexpr (cuda::std::is_same_v) { + return primitive_type::int16; + } else if constexpr (cuda::std::is_same_v) { + return primitive_type::int32; + } else if constexpr (cuda::std::is_same_v) { + return primitive_type::int64; + } else if constexpr (cuda::std::is_same_v) { + return primitive_type::float32; + } else if constexpr (cuda::std::is_same_v) { + 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 +__device__ inline cuda::std::optional decode_primitive(device_span enc) +{ if (cuda::std::cmp_less(enc.size(), 1 + sizeof(T))) { return cuda::std::nullopt; } - constexpr primitive_type expected = cuda::std::is_same_v ? primitive_type::int8 - : cuda::std::is_same_v ? primitive_type::int16 - : cuda::std::is_same_v ? 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(expected)) { + variant_value_header(value_metadata) != static_cast(primitive_type_for())) { return cuda::std::nullopt; } return cudf::io::unaligned_load(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 -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 d_output, bitmask_type* d_null_mask) { auto const num_rows = static_cast(d_output.size()); @@ -533,7 +559,7 @@ CUDF_KERNEL __launch_bounds__(block_size) void cast_variant_int_kernel( device_span const val{val_child.data() + val_begin, static_cast(val_end - val_begin)}; - auto const decoded = decode_int(val); + auto const decoded = decode_primitive(val); if (decoded.has_value()) { d_output[row] = *decoded; } else { @@ -607,12 +633,12 @@ struct cast_variant_fn { template std::unique_ptr operator()() - requires(is_variant_int) + requires(is_variant_primitive) { rmm::device_buffer data{num_rows * sizeof(T), stream, mr}; auto grid = cudf::detail::grid_1d{num_rows, block_size}; - cast_variant_int_kernel<<>>( + cast_variant_primitive_kernel<<>>( values, {static_cast(data.data()), static_cast(num_rows)}, d_null_mask); CUDF_CUDA_TRY(cudaGetLastError()); diff --git a/cpp/src/io/utilities/block_utils.cuh b/cpp/src/io/utilities/block_utils.cuh index d00e91ef2bb0..d157378127d3 100644 --- a/cpp/src/io/utilities/block_utils.cuh +++ b/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 - requires(cuda::std::is_integral_v) + requires(cuda::std::is_trivially_copyable_v) inline __device__ T unaligned_load(uint8_t const* p) { T value; diff --git a/cpp/tests/io/experimental/variant_extract_test.cpp b/cpp/tests/io/experimental/variant_extract_test.cpp index 051c3d94e240..2e53fd92bceb 100644 --- a/cpp/tests/io/experimental/variant_extract_test.cpp +++ b/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 +std::unique_ptr cast_apache_primitive(avf::fixture 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()}, 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()}, stream); - cudf::test::fixed_width_column_wrapper expected{expected_val}; + { + auto got = cast_apache_primitive(avf::primitive_int8); + cudf::test::fixed_width_column_wrapper expected{int8_t{42}}; CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); - }; + } + { + auto got = cast_apache_primitive(avf::primitive_int16); + cudf::test::fixed_width_column_wrapper expected{int16_t{1234}}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); + } + { + auto got = cast_apache_primitive(avf::primitive_int32); + cudf::test::fixed_width_column_wrapper expected{int32_t{123456}}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); + } + { + auto got = cast_apache_primitive(avf::primitive_int64); + cudf::test::fixed_width_column_wrapper 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(avf::primitive_float); + cudf::test::fixed_width_column_wrapper expected{float{1234567936.0f}}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); + } + { + auto got = cast_apache_primitive(avf::primitive_double); + cudf::test::fixed_width_column_wrapper 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);