diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index 7c53a89e4e7..dca21e4c7da 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -109,6 +110,26 @@ namespace io::parquet::experimental { rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +/** + * @brief Return the logical type of each VARIANT value blob in a `list` column. + * + * Classifies only the value_metadata header byte; does not validate the remaining payload. + * A recognized header returns its logical type even when the payload is truncated. A null output + * row is produced when the input row is null, the blob is empty, or the header carries an + * unrecognized type. An encoded Variant null (NULLVAL) produces a valid `NULL_VALUE` row. + * + * @param values `list` column of VARIANT-encoded value bytes + * @param stream CUDA stream + * @param mr Device memory resource + * @return `UINT8` column of `variant_logical_type` values cast to `uint8_t` + * + * @throws std::invalid_argument if `values` is not a `list` column + */ +[[nodiscard]] std::unique_ptr get_variant_type_id( + column_view const& values, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + /** @} */ } // namespace io::parquet::experimental } // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/io/experimental/variant_spec.hpp b/cpp/include/cudf/io/experimental/variant_spec.hpp index 6b71dc57385..ed09a490fb4 100644 --- a/cpp/include/cudf/io/experimental/variant_spec.hpp +++ b/cpp/include/cudf/io/experimental/variant_spec.hpp @@ -46,4 +46,25 @@ enum class variant_primitive_type : uint8_t { UUID = 20, }; +/** + * @brief Logical type of a VARIANT value as returned by get_variant_type_id. + */ +enum class variant_logical_type : uint8_t { + OBJECT = 0, + ARRAY = 1, + NULL_VALUE = 2, + BOOLEAN = 3, + LONG_VALUE = 4, + STRING = 5, + DOUBLE_VALUE = 6, + DECIMAL = 7, + DATE = 8, + TIMESTAMP = 9, + TIMESTAMP_NTZ = 10, + FLOAT_VALUE = 11, + BINARY = 12, + UUID = 13, + TIME_NTZ = 14, +}; + } // namespace cudf::io::parquet::experimental diff --git a/cpp/src/io/parquet/experimental/variant_extract.cu b/cpp/src/io/parquet/experimental/variant_extract.cu index 8327e6c5cde..f380538696e 100644 --- a/cpp/src/io/parquet/experimental/variant_extract.cu +++ b/cpp/src/io/parquet/experimental/variant_extract.cu @@ -789,6 +789,47 @@ struct cast_variant_fn { } }; +/** + * @brief Classifies only the first (value_metadata) byte of enc; does not validate the remaining + * payload. A recognized header returns its logical type even when the payload is truncated. Returns + * nullopt for an empty blob or an unrecognized primitive type ID. + */ +__device__ cuda::std::optional logical_type_of(device_span enc) +{ + if (enc.empty()) { return cuda::std::nullopt; } + auto const value_metadata = enc[0]; + auto const btype = decode_basic_type(value_metadata); + + if (btype == basic_type::SHORT_STRING) { return variant_logical_type::STRING; } + if (btype == basic_type::OBJECT) { return variant_logical_type::OBJECT; } + if (btype == basic_type::ARRAY) { return variant_logical_type::ARRAY; } + + switch (static_cast(variant_value_header(value_metadata))) { + case primitive_type::NULLVAL: return variant_logical_type::NULL_VALUE; + case primitive_type::BOOLEAN_TRUE: + case primitive_type::BOOLEAN_FALSE: return variant_logical_type::BOOLEAN; + case primitive_type::INT8: + case primitive_type::INT16: + case primitive_type::INT32: + case primitive_type::INT64: return variant_logical_type::LONG_VALUE; + case primitive_type::FLOAT64: return variant_logical_type::DOUBLE_VALUE; + case primitive_type::DECIMAL4: + case primitive_type::DECIMAL8: + case primitive_type::DECIMAL16: return variant_logical_type::DECIMAL; + case primitive_type::DATE: return variant_logical_type::DATE; + case primitive_type::TIMESTAMP_MICROS: + case primitive_type::TIMESTAMP_NANOS: return variant_logical_type::TIMESTAMP; + case primitive_type::TIMESTAMP_NTZ_MICROS: + case primitive_type::TIMESTAMP_NTZ_NANOS: return variant_logical_type::TIMESTAMP_NTZ; + case primitive_type::FLOAT32: return variant_logical_type::FLOAT_VALUE; + case primitive_type::BINARY: return variant_logical_type::BINARY; + case primitive_type::LONG_STRING: return variant_logical_type::STRING; + case primitive_type::TIME_NTZ_MICROS: return variant_logical_type::TIME_NTZ; + case primitive_type::UUID: return variant_logical_type::UUID; + default: return cuda::std::nullopt; + } +} + std::unique_ptr build_path_column(cudf::host_span steps, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) @@ -958,6 +999,45 @@ std::unique_ptr cast_variant(column_view const& values, mr}); } +std::unique_ptr get_variant_type_id(column_view const& values, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + validate_variant_child(values); + size_type const num_rows = values.size(); + if (num_rows == 0) { return make_empty_column(data_type{type_id::UINT8}); } + + auto val_device_view = column_device_view::create(values, stream); + cudf::lists_column_device_view val_lists_device_view(*val_device_view); + + auto null_mask = values.nullable() + ? cudf::detail::copy_bitmask(values, stream, mr) + : cudf::create_null_mask(num_rows, mask_state::ALL_VALID, stream, mr); + auto* d_null_mask = static_cast(null_mask.data()); + + rmm::device_buffer data{static_cast(num_rows) * sizeof(uint8_t), stream, mr}; + + thrust::transform( + rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + cuda::counting_iterator(0), + cuda::counting_iterator(num_rows), + static_cast(data.data()), + [values = val_lists_device_view, d_null_mask] __device__(size_type row) -> uint8_t { + if (!cudf::bit_is_set(d_null_mask, row)) { return 0; } + auto const ltype = logical_type_of(list_row_span(values, row)); + if (ltype.has_value()) { return static_cast(ltype.value()); } + cudf::clear_bit(d_null_mask, row); + return 0; + }); + + auto const null_count = num_rows - cudf::detail::count_set_bits(d_null_mask, 0, num_rows, stream); + return std::make_unique(data_type{type_id::UINT8}, + num_rows, + std::move(data), + null_count > 0 ? std::move(null_mask) : rmm::device_buffer{}, + null_count); +} + } // namespace detail std::unique_ptr get_variant_field(column_view const& variant_column, @@ -978,6 +1058,14 @@ std::unique_ptr cast_variant(column_view const& values, return detail::cast_variant(values, desired_type, stream, mr); } +std::unique_ptr get_variant_type_id(column_view const& values, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + return detail::get_variant_type_id(values, stream, mr); +} + std::unique_ptr extract_variant_field(column_view const& variant_column, std::string_view path, data_type desired_type, diff --git a/cpp/tests/io/experimental/variant_extract_test.cpp b/cpp/tests/io/experimental/variant_extract_test.cpp index 03c334f250e..b6dd4d51554 100644 --- a/cpp/tests/io/experimental/variant_extract_test.cpp +++ b/cpp/tests/io/experimental/variant_extract_test.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -1392,3 +1393,363 @@ TEST_F(InvalidInputShapeTest, CastVariantRejectsMalformedInput) std::invalid_argument); } } + +TEST_F(InvalidInputShapeTest, GetVariantTypeIdRejectsMalformedInput) +{ + // get_variant_type_id requires a list input; every other shape must be rejected with + // std::invalid_argument. + auto stream = cudf::test::get_default_stream(); + + std::vector cases; + cases.push_back({"input is not a list", scalar_i32()}); + cases.push_back({"input list has wrong element type (not uint8)", list_i32({1, 2, 3})}); + + for (auto const& c : cases) { + SCOPED_TRACE(c.label); + EXPECT_THROW(static_cast( + cudf::io::parquet::experimental::get_variant_type_id(c.column->view(), stream)), + std::invalid_argument); + } +} + +namespace { + +/** + * @brief Helper: run get_variant_type_id on the value child of an apache fixture. + */ +template +[[nodiscard]] std::unique_ptr apache_type_id(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::get_variant_type_id(value, stream); +} + +/** + * @brief Build a list column from blobs with per-row validity. Rows where valid[i] is false + * are null at the list level (not an encoded Variant null — those are valid rows with a NULLVAL + * blob). + */ +[[nodiscard]] std::unique_ptr make_list_u8_nullable( + cudf::host_span const> blobs, std::vector const& valid) +{ + auto const num_rows = static_cast(blobs.size()); + std::vector offsets(num_rows + 1, 0); + std::vector flat; + for (cudf::size_type i = 0; i < num_rows; ++i) { + flat.insert(flat.end(), blobs[i].begin(), blobs[i].end()); + offsets[i + 1] = static_cast(flat.size()); + } + auto off_col = + cudf::test::fixed_width_column_wrapper(offsets.begin(), offsets.end()).release(); + auto dat_col = + cudf::test::fixed_width_column_wrapper(flat.begin(), flat.end()).release(); + auto [d_mask, null_count] = cudf::test::detail::make_null_mask(valid.begin(), valid.end()); + return cudf::make_lists_column( + num_rows, std::move(off_col), std::move(dat_col), null_count, std::move(d_mask)); +} + +} // namespace + +struct GetVariantTypeIdTest : public cudf::test::BaseFixture {}; + +using LT = cudf::io::parquet::experimental::variant_logical_type; + +// --------------------------------------------------------------------------- +// Apache fixtures: one test per logical-type category. +// --------------------------------------------------------------------------- + +TEST_F(GetVariantTypeIdTest, NullValue) +{ + auto got = apache_type_id(avf::primitive_null); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::NULL_VALUE)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, Boolean) +{ + cudf::test::fixed_width_column_wrapper const expected{static_cast(LT::BOOLEAN)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_boolean_true), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_boolean_false), expected); +} + +TEST_F(GetVariantTypeIdTest, LongValueAllIntWidths) +{ + // INT8, INT16, INT32, INT64 all map to long_value regardless of physical width. + cudf::test::fixed_width_column_wrapper const expected{ + static_cast(LT::LONG_VALUE)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_int8), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_int16), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_int32), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_int64), expected); +} + +TEST_F(GetVariantTypeIdTest, StringBothEncodings) +{ + // SHORT_STRING and primitive LONG_STRING both map to string. + cudf::test::fixed_width_column_wrapper const expected{static_cast(LT::STRING)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::short_string), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_string), expected); +} + +TEST_F(GetVariantTypeIdTest, FloatTypes) +{ + { + auto got = apache_type_id(avf::primitive_float); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::FLOAT_VALUE)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); + } + { + auto got = apache_type_id(avf::primitive_double); + cudf::test::fixed_width_column_wrapper expected{ + static_cast(LT::DOUBLE_VALUE)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); + } +} + +TEST_F(GetVariantTypeIdTest, Decimal) +{ + cudf::test::fixed_width_column_wrapper const expected{static_cast(LT::DECIMAL)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_decimal4), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_decimal8), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_decimal16), expected); +} + +TEST_F(GetVariantTypeIdTest, Date) +{ + auto got = apache_type_id(avf::primitive_date); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::DATE)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, TimestampBothNanos) +{ + // TIMESTAMP_MICROS and TIMESTAMP_NANOS both map to timestamp. + cudf::test::fixed_width_column_wrapper const expected{ + static_cast(LT::TIMESTAMP)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_timestamp), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_timestamp_nanos), expected); +} + +TEST_F(GetVariantTypeIdTest, TimestampNtzBothNanos) +{ + // TIMESTAMP_NTZ_MICROS and TIMESTAMP_NTZ_NANOS both map to timestamp_ntz. + cudf::test::fixed_width_column_wrapper const expected{ + static_cast(LT::TIMESTAMP_NTZ)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_timestampntz), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::primitive_timestampntz_nanos), expected); +} + +TEST_F(GetVariantTypeIdTest, Binary) +{ + auto got = apache_type_id(avf::primitive_binary); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::BINARY)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, Uuid) +{ + auto got = apache_type_id(avf::primitive_uuid); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::UUID)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, TimeNtz) +{ + auto got = apache_type_id(avf::primitive_time); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::TIME_NTZ)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, ObjectAndArray) +{ + { + cudf::test::fixed_width_column_wrapper const expected{ + static_cast(LT::OBJECT)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::object_primitive), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::object_nested), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::object_empty), expected); + } + { + cudf::test::fixed_width_column_wrapper const expected{static_cast(LT::ARRAY)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::array_primitive), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::array_nested), expected); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*apache_type_id(avf::array_empty), expected); + } +} + +// --------------------------------------------------------------------------- +// Null and unknown-type behavior. +// --------------------------------------------------------------------------- + +TEST_F(GetVariantTypeIdTest, UnknownPhysicalTypeProducesNull) +{ + // Primitive header byte 0xFC = (63 << 2) | 0: type_id 63 is not in the spec. + auto const stream = cudf::test::get_default_stream(); + auto values = make_list_u8_nullable(std::vector>{{0xFC}}, {true}); + auto got = cudf::io::parquet::experimental::get_variant_type_id(*values, stream); + ASSERT_EQ(got->size(), 1); + EXPECT_EQ(got->null_count(), 1); +} + +TEST_F(GetVariantTypeIdTest, InputNullRowPropagates) +{ + // A null row in the input list column propagates to the output. + auto const stream = cudf::test::get_default_stream(); + auto values = make_list_u8_nullable( + std::vector>{enc_int32(1), enc_int32(2), enc_int32(3)}, + {true, false, true}); + + auto got = cudf::io::parquet::experimental::get_variant_type_id(*values, stream); + + std::initializer_list expected_vals1{ + static_cast(LT::LONG_VALUE), 0, static_cast(LT::LONG_VALUE)}; + std::initializer_list validity1{true, false, true}; + cudf::test::fixed_width_column_wrapper expected(expected_vals1, validity1); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, EncodedNullIsNotInputNull) +{ + // An encoded Variant NULLVAL blob is a valid row whose type is null_value, not a null row. + auto const stream = cudf::test::get_default_stream(); + auto val = enc_null(); + cudf::test::lists_column_wrapper values(val.begin(), val.end()); + auto got = cudf::io::parquet::experimental::get_variant_type_id(values, stream); + + ASSERT_EQ(got->size(), 1); + EXPECT_EQ(got->null_count(), 0); + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::NULL_VALUE)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, EmptyValueBlobProducesNull) +{ + // An empty list row (zero bytes) has no header byte to decode → null. + auto const stream = cudf::test::get_default_stream(); + auto values = make_list_u8_nullable( + std::vector>{enc_int32(1), {}, enc_int32(3)}, {true, true, true}); + + auto got = cudf::io::parquet::experimental::get_variant_type_id(*values, stream); + + std::initializer_list expected_vals2{ + static_cast(LT::LONG_VALUE), 0, static_cast(LT::LONG_VALUE)}; + std::initializer_list validity2{true, false, true}; + cudf::test::fixed_width_column_wrapper expected(expected_vals2, validity2); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +// --------------------------------------------------------------------------- +// Multi-row and structural tests. +// --------------------------------------------------------------------------- + +TEST_F(GetVariantTypeIdTest, MixedTypesColumn) +{ + auto const stream = cudf::test::get_default_stream(); + + auto null_val = enc_null(); + auto bool_val = enc_bool(true); + auto int_val = enc_int64(999); + auto str_val = enc_short_string("hi"); + auto dbl_val = enc_float64(3.14); + + cudf::test::lists_column_wrapper values{ + {null_val.begin(), null_val.end()}, + {bool_val.begin(), bool_val.end()}, + {int_val.begin(), int_val.end()}, + {str_val.begin(), str_val.end()}, + {dbl_val.begin(), dbl_val.end()}, + }; + auto got = cudf::io::parquet::experimental::get_variant_type_id(values, stream); + + cudf::test::fixed_width_column_wrapper expected{ + static_cast(LT::NULL_VALUE), + static_cast(LT::BOOLEAN), + static_cast(LT::LONG_VALUE), + static_cast(LT::STRING), + static_cast(LT::DOUBLE_VALUE), + }; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, AllNullInputColumn) +{ + // All rows are null at the list level → all output rows are null. + auto const stream = cudf::test::get_default_stream(); + auto values = make_list_u8_nullable( + std::vector>{enc_int32(1), enc_int32(2), enc_int32(3)}, + {false, false, false}); + + auto got = cudf::io::parquet::experimental::get_variant_type_id(*values, stream); + + ASSERT_EQ(got->size(), 3); + EXPECT_EQ(got->null_count(), 3); +} + +TEST_F(GetVariantTypeIdTest, EmptyInput) +{ + auto const stream = cudf::test::get_default_stream(); + auto const values = + cudf::empty_like(cudf::structs_column_view{make_xyz_three_row_variant()}.child(1)); + + auto got = cudf::io::parquet::experimental::get_variant_type_id(*values, stream); + EXPECT_EQ(got->type().id(), cudf::type_id::UINT8); + EXPECT_EQ(got->size(), 0); + EXPECT_EQ(got->null_count(), 0); +} + +TEST_F(GetVariantTypeIdTest, SlicedValuesColumn) +{ + // Verify that a sliced input produces correct results for the slice only. + auto const stream = cudf::test::get_default_stream(); + auto col = make_xyz_three_row_variant(); + auto const value_child = cudf::structs_column_view{col}.get_sliced_child(1, stream); + + // The xyz variant has object rows; slicing [1,3) gives 2 object rows. + auto const sliced_values = cudf::slice(value_child, {1, 3}).front(); + auto got = cudf::io::parquet::experimental::get_variant_type_id(sliced_values, stream); + + cudf::test::fixed_width_column_wrapper expected{static_cast(LT::OBJECT), + static_cast(LT::OBJECT)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +} + +TEST_F(GetVariantTypeIdTest, LargeMultiRowColumn) +{ + // 600 rows cycling through all types that get_variant_type_id can classify. + // 600 > 512 (typical block size) so the kernel exercises the multi-block grid-stride path. + auto const stream = cudf::test::get_default_stream(); + + struct row_spec { + std::vector blob; + uint8_t expected_id; + }; + + std::vector const types{ + {enc_null(), static_cast(LT::NULL_VALUE)}, + {enc_bool(false), static_cast(LT::BOOLEAN)}, + {enc_int8(1), static_cast(LT::LONG_VALUE)}, + {enc_int16(2), static_cast(LT::LONG_VALUE)}, + {enc_int32(3), static_cast(LT::LONG_VALUE)}, + {enc_int64(4), static_cast(LT::LONG_VALUE)}, + {enc_float64(5.0), static_cast(LT::DOUBLE_VALUE)}, + {enc_short_string("x"), static_cast(LT::STRING)}, + {enc_long_string(std::string(70, 'z')), static_cast(LT::STRING)}, + }; + constexpr int num_rows = 600; + std::vector> blobs(num_rows); + std::vector expected_ids(num_rows); + for (int i = 0; i < num_rows; ++i) { + auto const& spec = types[i % types.size()]; + blobs[i] = spec.blob; + expected_ids[i] = spec.expected_id; + } + + auto values = make_list_u8_nullable(blobs, std::vector(num_rows, true)); + auto got = cudf::io::parquet::experimental::get_variant_type_id(*values, stream); + + cudf::test::fixed_width_column_wrapper expected(expected_ids.begin(), + expected_ids.end()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*got, expected); +}