From aa844fbf8ec1abbbda93ad45631c3621f544fce3 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Thu, 25 Jun 2026 17:46:45 +0800 Subject: [PATCH 01/10] support to encode, write and validate variant --- cpp/src/arrow/extension/parquet_variant.cc | 188 ++- cpp/src/arrow/extension/parquet_variant.h | 22 +- cpp/src/parquet/CMakeLists.txt | 11 +- .../parquet/arrow/arrow_reader_writer_test.cc | 7 +- cpp/src/parquet/arrow/arrow_schema_test.cc | 356 ++++- cpp/src/parquet/arrow/schema.cc | 336 ++++- cpp/src/parquet/arrow/variant_test.cc | 374 ++++- cpp/src/parquet/arrow/writer.cc | 18 +- cpp/src/parquet/meson.build | 5 + cpp/src/parquet/properties.h | 45 +- cpp/src/parquet/variant/CMakeLists.txt | 26 + cpp/src/parquet/variant/builder.cc | 1273 +++++++++++++++++ cpp/src/parquet/variant/builder.h | 264 ++++ cpp/src/parquet/variant/builder_test.cc | 310 ++++ cpp/src/parquet/variant/encoding.cc | 470 ++++++ cpp/src/parquet/variant/encoding.h | 170 +++ cpp/src/parquet/variant/encoding_internal.h | 147 ++ cpp/src/parquet/variant/encoding_test.cc | 190 +++ cpp/src/parquet/variant/meson.build | 31 + cpp/src/parquet/variant/test_util_internal.cc | 197 +++ cpp/src/parquet/variant/test_util_internal.h | 83 ++ cpp/src/parquet/variant/type_test.cc | 167 +++ cpp/src/parquet/variant/validate.cc | 473 ++++++ cpp/src/parquet/variant/validate.h | 36 + cpp/src/parquet/variant/validate_test.cc | 190 +++ 25 files changed, 5166 insertions(+), 223 deletions(-) create mode 100644 cpp/src/parquet/variant/CMakeLists.txt create mode 100644 cpp/src/parquet/variant/builder.cc create mode 100644 cpp/src/parquet/variant/builder.h create mode 100644 cpp/src/parquet/variant/builder_test.cc create mode 100644 cpp/src/parquet/variant/encoding.cc create mode 100644 cpp/src/parquet/variant/encoding.h create mode 100644 cpp/src/parquet/variant/encoding_internal.h create mode 100644 cpp/src/parquet/variant/encoding_test.cc create mode 100644 cpp/src/parquet/variant/meson.build create mode 100644 cpp/src/parquet/variant/test_util_internal.cc create mode 100644 cpp/src/parquet/variant/test_util_internal.h create mode 100644 cpp/src/parquet/variant/type_test.cc create mode 100644 cpp/src/parquet/variant/validate.cc create mode 100644 cpp/src/parquet/variant/validate.h create mode 100644 cpp/src/parquet/variant/validate_test.cc diff --git a/cpp/src/arrow/extension/parquet_variant.cc b/cpp/src/arrow/extension/parquet_variant.cc index 95aa5a0eb68e..4630c1d39321 100644 --- a/cpp/src/arrow/extension/parquet_variant.cc +++ b/cpp/src/arrow/extension/parquet_variant.cc @@ -17,28 +17,30 @@ #include "arrow/extension/parquet_variant.h" +#include #include #include "arrow/extension_type.h" #include "arrow/result.h" #include "arrow/status.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" #include "arrow/util/logging_internal.h" namespace arrow::extension { VariantExtensionType::VariantExtensionType(const std::shared_ptr& storage_type) : ExtensionType(storage_type) { - // GH-45948: Shredded variants will need to handle an optional shredded_value as - // well as value_ becoming optional. - - // IsSupportedStorageType should have been called already, asserting that both - // metadata and value are present. - if (storage_type->field(0)->name() == "metadata") { - metadata_ = storage_type->field(0); - value_ = storage_type->field(1); - } else { - value_ = storage_type->field(0); - metadata_ = storage_type->field(1); + // IsSupportedStorageType should have been called already, asserting that + // metadata is present and at least one of value / typed_value is present. + for (const auto& field : storage_type->fields()) { + if (field->name() == "metadata") { + metadata_ = field; + } else if (field->name() == "value") { + value_ = field; + } else if (field->name() == "typed_value") { + typed_value_ = field; + } } } @@ -49,6 +51,9 @@ bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const { Result> VariantExtensionType::Deserialize( std::shared_ptr storage_type, const std::string& serialized) const { + if (!serialized.empty()) { + return Status::Invalid("Unexpected serialized metadata: '", serialized, "'"); + } return VariantExtensionType::Make(std::move(storage_type)); } @@ -57,50 +62,148 @@ std::string VariantExtensionType::Serialize() const { return ""; } std::shared_ptr VariantExtensionType::MakeArray( std::shared_ptr data) const { DCHECK_EQ(data->type->id(), Type::EXTENSION); - DCHECK_EQ("arrow.parquet.variant", + DCHECK_EQ(kVariantExtensionName, internal::checked_cast(*data->type).extension_name()); return std::make_shared(data); } namespace { -bool IsBinaryField(const std::shared_ptr field) { - return field->type()->storage_id() == Type::BINARY || - field->type()->storage_id() == Type::LARGE_BINARY; + +bool IsSupportedPrimitiveTypedValue(const std::shared_ptr& type) { + switch (type->id()) { + case Type::BOOL: + case Type::INT8: + case Type::INT16: + case Type::INT32: + case Type::INT64: + case Type::FLOAT: + case Type::DOUBLE: + case Type::DATE32: + case Type::BINARY: + case Type::LARGE_BINARY: + case Type::BINARY_VIEW: + case Type::STRING: + case Type::LARGE_STRING: + case Type::STRING_VIEW: + return true; + case Type::DECIMAL32: + case Type::DECIMAL64: + case Type::DECIMAL128: { + const auto& decimal = internal::checked_cast(*type); + return decimal.scale() >= 0 && decimal.scale() <= decimal.precision(); + } + case Type::TIME64: + return internal::checked_cast(*type).unit() == TimeUnit::MICRO; + case Type::TIMESTAMP: { + const auto unit = internal::checked_cast(*type).unit(); + return unit == TimeUnit::MICRO || unit == TimeUnit::NANO; + } + case Type::FIXED_SIZE_BINARY: + return internal::checked_cast(*type).byte_width() == 16; + case Type::EXTENSION: { + const auto& ext_type = internal::checked_cast(*type); + return ext_type.extension_name() == "arrow.uuid"; + } + default: + return false; + } } + +bool IsSupportedTypedValue(const std::shared_ptr& field); + +bool IsVariantFieldGroup(const std::shared_ptr& type) { + if (type->id() != Type::STRUCT) { + return false; + } + + std::shared_ptr value; + std::shared_ptr typed_value; + for (const auto& field : type->fields()) { + if (field->name() == "value") { + if (value != nullptr || !field->nullable() || + !is_binary_or_binary_view(field->type()->storage_id())) { + return false; + } + value = field; + } else if (field->name() == "typed_value") { + if (typed_value != nullptr || !IsSupportedTypedValue(field)) { + return false; + } + typed_value = field; + } else { + return false; + } + } + return value != nullptr || typed_value != nullptr; +} + +bool IsSupportedTypedValue(const std::shared_ptr& field) { + if (!field->nullable()) { + return false; + } + auto is_variant_field_group = [](const auto& field) { + return !field->nullable() && IsVariantFieldGroup(field->type()); + }; + + switch (field->type()->id()) { + case Type::STRUCT: + return field->type()->num_fields() > 0 && + std::ranges::all_of(field->type()->fields(), is_variant_field_group); + case Type::LIST: + case Type::LARGE_LIST: + case Type::LIST_VIEW: + case Type::LARGE_LIST_VIEW: + case Type::FIXED_SIZE_LIST: + return is_variant_field_group(field->type()->field(0)); + default: + return IsSupportedPrimitiveTypedValue(field->type()); + } +} + } // namespace bool VariantExtensionType::IsSupportedStorageType( const std::shared_ptr& storage_type) { - // For now we only supported unshredded variants. Unshredded variant storage - // type should be a struct with a binary metadata and binary value. - // - // GH-45948: In shredded variants, the binary value field can be replaced - // with one or more of the following: object, array, typed_value, and - // variant_value. - if (storage_type->id() == Type::STRUCT) { - if (storage_type->num_fields() == 2) { - // Ordering of metadata and value fields does not matter, as we will assign - // these to the VariantExtensionType's member shared_ptrs in the constructor. - // Here we just need to check that they are both present. - - const auto& field0 = storage_type->field(0); - const auto& field1 = storage_type->field(1); - - bool metadata_and_value_present = - (field0->name() == "metadata" && field1->name() == "value") || - (field1->name() == "metadata" && field0->name() == "value"); - - if (metadata_and_value_present) { - // Both metadata and value must be non-nullable binary types for unshredded - // variants. This will change in GH-46948, when we will require a Visitor - // to traverse the structure of the variant. - return IsBinaryField(field0) && IsBinaryField(field1) && !field0->nullable() && - !field1->nullable(); + if (storage_type->id() != Type::STRUCT) { + return false; + } + + std::shared_ptr metadata; + std::shared_ptr value; + std::shared_ptr typed_value; + + for (const auto& field : storage_type->fields()) { + if (field->name() == "metadata") { + if (metadata != nullptr || !is_binary_or_binary_view(field->type()->storage_id()) || + field->nullable()) { + return false; + } + metadata = field; + } else if (field->name() == "value") { + if (value != nullptr || !is_binary_or_binary_view(field->type()->storage_id())) { + return false; + } + value = field; + } else if (field->name() == "typed_value") { + if (typed_value != nullptr || !IsSupportedTypedValue(field)) { + return false; } + typed_value = field; + } else { + return false; } } - return false; + if (metadata == nullptr || (value == nullptr && typed_value == nullptr)) { + return false; + } + if (value == nullptr) { + return true; + } + if (typed_value == nullptr) { + return !value->nullable(); + } + return value->nullable(); } Result> VariantExtensionType::Make( @@ -113,9 +216,6 @@ Result> VariantExtensionType::Make( return std::make_shared(std::move(storage_type)); } -/// NOTE: this is still experimental. GH-45948 will add shredding support, at which point -/// we need to separate this into unshredded_variant and shredded_variant helper -/// functions. std::shared_ptr variant(std::shared_ptr storage_type) { return VariantExtensionType::Make(std::move(storage_type)).ValueOrDie(); } diff --git a/cpp/src/arrow/extension/parquet_variant.h b/cpp/src/arrow/extension/parquet_variant.h index be90923f14e6..c74a794f6a01 100644 --- a/cpp/src/arrow/extension/parquet_variant.h +++ b/cpp/src/arrow/extension/parquet_variant.h @@ -18,12 +18,16 @@ #pragma once #include +#include #include "arrow/extension_type.h" #include "arrow/util/visibility.h" namespace arrow::extension { +/// \brief The extension name for the Variant extension type. +inline constexpr std::string_view kVariantExtensionName = "arrow.parquet.variant"; + class ARROW_EXPORT VariantArray : public ExtensionArray { public: using ExtensionArray::ExtensionArray; @@ -43,13 +47,25 @@ class ARROW_EXPORT VariantArray : public ExtensionArray { /// To read more about variant encoding, see the variant encoding spec at /// https://github.com/apache/parquet-format/blob/master/VariantEncoding.md /// +/// Shredded variant representation: +/// optional group shredded_variant_name (VARIANT) { +/// required binary metadata; +/// optional binary value; +/// optional typed_value; +/// } +/// +/// The value and typed_value fields are optional in the schema, but at least one +/// must be present. +/// /// To read more about variant shredding, see the variant shredding spec at /// https://github.com/apache/parquet-format/blob/master/VariantShredding.md class ARROW_EXPORT VariantExtensionType : public ExtensionType { public: explicit VariantExtensionType(const std::shared_ptr& storage_type); - std::string extension_name() const override { return "arrow.parquet.variant"; } + std::string extension_name() const override { + return std::string(kVariantExtensionName); + } bool ExtensionEquals(const ExtensionType& other) const override; @@ -69,10 +85,12 @@ class ARROW_EXPORT VariantExtensionType : public ExtensionType { std::shared_ptr value() const { return value_; } + std::shared_ptr typed_value() const { return typed_value_; } + private: - // TODO GH-45948 added shredded_value std::shared_ptr metadata_; std::shared_ptr value_; + std::shared_ptr typed_value_; }; /// \brief Return a VariantExtensionType instance. diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index f8a42b5b96bf..11da3eccecab 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -180,7 +180,6 @@ set(PARQUET_SRCS level_comparison.cc level_conversion.cc metadata.cc - xxhasher.cc page_index.cc "${PARQUET_THRIFT_SOURCE_DIR}/parquet_types.cpp" platform.cc @@ -191,7 +190,11 @@ set(PARQUET_SRCS statistics.cc stream_reader.cc stream_writer.cc - types.cc) + types.cc + variant/builder.cc + variant/encoding.cc + variant/validate.cc + xxhasher.cc) if(ARROW_HAVE_RUNTIME_AVX2) # AVX2 is used as a proxy for BMI2. @@ -363,6 +366,7 @@ add_subdirectory(api) add_subdirectory(arrow) add_subdirectory(encryption) add_subdirectory(geospatial) +add_subdirectory(variant) arrow_install_all_headers("parquet") @@ -409,7 +413,8 @@ add_parquet_test(arrow-reader-writer-test SOURCES arrow/arrow_reader_writer_test.cc arrow/arrow_statistics_test.cc - arrow/variant_test.cc) + arrow/variant_test.cc + variant/test_util_internal.cc) add_parquet_test(arrow-index-test SOURCES arrow/index_test.cc) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 8735aea731ce..917f2fe95a64 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -5999,11 +5999,8 @@ TEST(TestArrowReadWrite, AllNulls) { auto schema = ::arrow::schema({::arrow::field("all_nulls", ::arrow::int8())}); constexpr int64_t length = 3; - ASSERT_OK_AND_ASSIGN(auto null_bitmap, ::arrow::AllocateEmptyBitmap(length)); - auto array_data = ::arrow::ArrayData::Make( - ::arrow::int8(), length, {null_bitmap, /*values=*/nullptr}, /*null_count=*/length); - auto array = ::arrow::MakeArray(array_data); - auto record_batch = ::arrow::RecordBatch::Make(schema, length, {array}); + ASSERT_OK_AND_ASSIGN(auto array, MakeArrayOfNull(::arrow::int8(), length)); + auto record_batch = ::arrow::RecordBatch::Make(schema, length, {std::move(array)}); auto sink = CreateOutputStream(); ASSERT_OK_AND_ASSIGN(auto writer, parquet::arrow::FileWriter::Open( diff --git a/cpp/src/parquet/arrow/arrow_schema_test.cc b/cpp/src/parquet/arrow/arrow_schema_test.cc index 7a7b5a336939..f8196240648f 100644 --- a/cpp/src/parquet/arrow/arrow_schema_test.cc +++ b/cpp/src/parquet/arrow/arrow_schema_test.cc @@ -31,7 +31,7 @@ #include "parquet/test_util.h" #include "parquet/thrift_internal.h" -#include "arrow/array.h" +#include "arrow/array.h" // IWYU pragma: keep #include "arrow/extension/json.h" #include "arrow/extension/parquet_variant.h" #include "arrow/extension/uuid.h" @@ -86,6 +86,9 @@ static const std::vector kListCases = { [](std::shared_ptr<::arrow::Field> field) { return ::arrow::large_list(field); }}, }; +Status ArrowSchemaToParquetMetadata(std::shared_ptr<::arrow::Schema>& arrow_schema, + std::shared_ptr& metadata); + class TestConvertParquetSchema : public ::testing::Test { public: virtual void SetUp() {} @@ -111,6 +114,55 @@ class TestConvertParquetSchema : public ::testing::Test { return FromParquetSchema(&descr_, props, key_value_metadata, &result_schema_); } + void CheckParquetVariantSchema(const std::string& name, + std::vector parquet_children, + const std::shared_ptr<::arrow::DataType>& storage_type, + bool check_serialized_arrow_schema = false) { + SCOPED_TRACE(name); + auto variant = GroupNode::Make(name, Repetition::OPTIONAL, + std::move(parquet_children), LogicalType::Variant()); + std::vector parquet_fields = {variant}; + auto variant_extension = ::arrow::extension::variant(storage_type); + + { + auto arrow_schema = ::arrow::schema({::arrow::field(name, storage_type)}); + ASSERT_OK(ConvertSchema(parquet_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); + } + + for (bool register_extension : {true, false}) { + ::arrow::ExtensionTypeGuard guard(register_extension + ? ::arrow::DataTypeVector{variant_extension} + : ::arrow::DataTypeVector{}); + + ArrowReaderProperties props; + props.set_arrow_extensions_enabled(true); + auto arrow_schema = ::arrow::schema( + {::arrow::field(name, register_extension ? variant_extension : storage_type)}); + + ASSERT_OK(ConvertSchema(parquet_fields, /*metadata=*/nullptr, props)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); + } + + if (check_serialized_arrow_schema) { + for (bool register_extension : {true, false}) { + ::arrow::ExtensionTypeGuard guard(register_extension + ? ::arrow::DataTypeVector{variant_extension} + : ::arrow::DataTypeVector{}); + + ArrowReaderProperties props; + props.set_arrow_extensions_enabled(false); + auto arrow_schema = ::arrow::schema({::arrow::field( + name, register_extension ? variant_extension : storage_type)}); + + std::shared_ptr metadata; + ASSERT_OK(ArrowSchemaToParquetMetadata(arrow_schema, metadata)); + ASSERT_OK(ConvertSchema(parquet_fields, metadata, props)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); + } + } + } + protected: SchemaDescriptor descr_; std::shared_ptr<::arrow::Schema> result_schema_; @@ -992,73 +1044,43 @@ Status ArrowSchemaToParquetMetadata(std::shared_ptr<::arrow::Schema>& arrow_sche TEST_F(TestConvertParquetSchema, ParquetVariant) { // Unshredded variant - // optional group variant_col { + // optional group variant_unshredded { // required binary metadata; // required binary value; // } - // - // GH-45948: add shredded variants - std::vector parquet_fields; auto metadata = PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY); auto value = PrimitiveNode::Make("value", Repetition::REQUIRED, ParquetType::BYTE_ARRAY); - auto variant = GroupNode::Make("variant_unshredded", Repetition::OPTIONAL, - {metadata, value}, LogicalType::Variant()); - parquet_fields.push_back(variant); - - // Arrow schema for unshredded variant struct. - auto arrow_metadata = ::arrow::field("metadata", ::arrow::binary(), /*nullable=*/false); - auto arrow_value = ::arrow::field("value", ::arrow::binary(), /*nullable=*/false); - auto arrow_variant = ::arrow::struct_({arrow_metadata, arrow_value}); - auto variant_extension = ::arrow::extension::variant(arrow_variant); - - { - // Parquet file does not contain Arrow schema. - // By default, field should be treated as a normal struct in Arrow. - auto arrow_schema = - ::arrow::schema({::arrow::field("variant_unshredded", arrow_variant)}); - ASSERT_OK(ConvertSchema(parquet_fields)); - ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); - } - - for (bool register_extension : {true, false}) { - ::arrow::ExtensionTypeGuard guard(register_extension - ? ::arrow::DataTypeVector{variant_extension} - : ::arrow::DataTypeVector{}); - - // Parquet file does not contain Arrow schema. - // If Arrow extensions are enabled, field should be interpreted as Parquet Variant - // extension type if registered. - ArrowReaderProperties props; - props.set_arrow_extensions_enabled(true); - - auto arrow_schema = ::arrow::schema({::arrow::field( - "variant_unshredded", register_extension ? variant_extension : arrow_variant)}); - - ASSERT_OK(ConvertSchema(parquet_fields, /*metadata=*/nullptr, props)); - ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); - } - - for (bool register_extension : {true, false}) { - ::arrow::ExtensionTypeGuard guard(register_extension - ? ::arrow::DataTypeVector{variant_extension} - : ::arrow::DataTypeVector{}); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), /*nullable=*/false), + ::arrow::field("value", ::arrow::binary(), /*nullable=*/false)}); - // Parquet file does contain Arrow schema. - // Field should be interpreted as Parquet Variant extension, if registered, - // even though extensions are not enabled. - ArrowReaderProperties props; - props.set_arrow_extensions_enabled(false); - - auto arrow_schema = ::arrow::schema({::arrow::field( - "variant_unshredded", register_extension ? variant_extension : arrow_variant)}); + ASSERT_NO_FATAL_FAILURE( + CheckParquetVariantSchema("variant_unshredded", {metadata, value}, storage_type, + /*check_serialized_arrow_schema=*/true)); +} - std::shared_ptr metadata; - ASSERT_OK(ArrowSchemaToParquetMetadata(arrow_schema, metadata)); - ASSERT_OK(ConvertSchema(parquet_fields, metadata, props)); - ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); - } +TEST_F(TestConvertParquetSchema, ParquetVariantShredded) { + // Shredded variant + // optional group variant_shredded { + // required binary metadata; + // optional binary value; + // optional int64 typed_value; + // } + auto metadata = + PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY); + auto value = + PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY); + auto typed_value = + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, ParquetType::INT64); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", ::arrow::int64())}); + + ASSERT_NO_FATAL_FAILURE(CheckParquetVariantSchema( + "variant_shredded", {metadata, value, typed_value}, storage_type)); } TEST_F(TestConvertParquetSchema, ParquetSchemaArrowJsonExtension) { @@ -1534,6 +1556,222 @@ TEST_F(TestConvertArrowSchema, ParquetFlatPrimitivesAsDictionaries) { ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); } +TEST_F(TestConvertArrowSchema, ParquetVariantShredded) { + auto storage_type = + ::arrow::struct_({::arrow::field("typed_value", ::arrow::int64()), + ::arrow::field("value", ::arrow::binary()), + ::arrow::field("metadata", ::arrow::binary(), false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + + std::vector> arrow_fields = { + ::arrow::field("variant", variant_type)}; + + std::vector parquet_fields = {GroupNode::Make( + "variant", Repetition::OPTIONAL, + {PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, ParquetType::INT64)}, + LogicalType::Variant())}; + + ASSERT_OK(ConvertSchema(arrow_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetVariantDictionary) { + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary(), false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto dictionary_variant = ::arrow::dictionary(::arrow::int32(), variant_type); + auto no_validation = + ArrowWriterProperties::Builder().set_variant_validation_enabled(false)->build(); + + ASSERT_RAISES( + NotImplemented, + ConvertSchema({::arrow::field("variant", dictionary_variant)}, no_validation)); +} + +TEST_F(TestConvertArrowSchema, ParquetVariantTypedSchema) { + auto ShreddedType = [](std::shared_ptr<::arrow::DataType> type) { + return ::arrow::struct_({::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", std::move(type))}); + }; + auto typed_value = ::arrow::struct_( + {::arrow::field("d4", ShreddedType(::arrow::decimal32(8, 2)), + /*nullable=*/false), + ::arrow::field("d8", ShreddedType(::arrow::decimal64(16, 4)), + /*nullable=*/false), + ::arrow::field("d64p8", ShreddedType(::arrow::decimal64(8, 2)), + /*nullable=*/false), + ::arrow::field("d128p8", ShreddedType(::arrow::decimal128(8, 2)), + /*nullable=*/false), + ::arrow::field("d128p16", ShreddedType(::arrow::decimal128(16, 4)), + /*nullable=*/false), + ::arrow::field("d128p32", ShreddedType(::arrow::decimal128(32, 8)), + /*nullable=*/false), + ::arrow::field("ts", ShreddedType(::arrow::timestamp(TimeUnit::NANO, "UTC")), + /*nullable=*/false), + ::arrow::field("time", ShreddedType(::arrow::time64(TimeUnit::MICRO)), + /*nullable=*/false), + ::arrow::field("id", ShreddedType(::arrow::fixed_size_binary(16)), + /*nullable=*/false), + ::arrow::field("uuid", ShreddedType(::arrow::extension::uuid()), + /*nullable=*/false)}); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", typed_value)}); + auto variant_type = ::arrow::extension::variant(storage_type); + + std::vector> arrow_fields = { + ::arrow::field("variant", variant_type)}; + + auto ShreddedField = [](const std::string& name, NodePtr typed) { + return GroupNode::Make( + name, Repetition::REQUIRED, + {PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + std::move(typed)}); + }; + std::vector parquet_fields = {GroupNode::Make( + "variant", Repetition::OPTIONAL, + {PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + GroupNode::Make( + "typed_value", Repetition::OPTIONAL, + {ShreddedField("d4", PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, + LogicalType::Decimal(8, 2), + ParquetType::INT32)), + ShreddedField("d8", PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, + LogicalType::Decimal(16, 4), + ParquetType::INT64)), + ShreddedField("d64p8", PrimitiveNode::Make( + "typed_value", Repetition::OPTIONAL, + LogicalType::Decimal(8, 2), ParquetType::INT32)), + ShreddedField("d128p8", PrimitiveNode::Make( + "typed_value", Repetition::OPTIONAL, + LogicalType::Decimal(8, 2), ParquetType::INT32)), + ShreddedField( + "d128p16", + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, + LogicalType::Decimal(16, 4), ParquetType::INT64)), + ShreddedField("d128p32", + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, + LogicalType::Decimal(32, 8), + ParquetType::FIXED_LEN_BYTE_ARRAY, 14)), + ShreddedField("ts", + PrimitiveNode::Make( + "typed_value", Repetition::OPTIONAL, + LogicalType::Timestamp(true, LogicalType::TimeUnit::NANOS), + ParquetType::INT64)), + ShreddedField("time", + PrimitiveNode::Make( + "typed_value", Repetition::OPTIONAL, + LogicalType::Time(false, LogicalType::TimeUnit::MICROS), + ParquetType::INT64)), + ShreddedField("id", + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, + LogicalType::UUID(), + ParquetType::FIXED_LEN_BYTE_ARRAY, 16)), + ShreddedField("uuid", + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, + LogicalType::UUID(), + ParquetType::FIXED_LEN_BYTE_ARRAY, 16))})}, + LogicalType::Variant())}; + + ASSERT_OK(ConvertSchema(arrow_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetVariantShreddedObject) { + auto field_group = ::arrow::struct_({::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", ::arrow::int64())}); + auto typed_value = + ::arrow::struct_({::arrow::field("a", field_group, /*nullable=*/false), + ::arrow::field("b", field_group, /*nullable=*/false)}); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", typed_value)}); + auto variant_type = ::arrow::extension::variant(storage_type); + + std::vector> arrow_fields = { + ::arrow::field("variant", variant_type)}; + + auto ShreddedField = [](const std::string& name) { + return GroupNode::Make( + name, Repetition::REQUIRED, + {PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, ParquetType::INT64)}); + }; + std::vector parquet_fields = {GroupNode::Make( + "variant", Repetition::OPTIONAL, + {PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + GroupNode::Make("typed_value", Repetition::OPTIONAL, + {ShreddedField("a"), ShreddedField("b")})}, + LogicalType::Variant())}; + + ASSERT_OK(ConvertSchema(arrow_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetVariantEmptyObject) { + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", ::arrow::struct_({}))}); + auto variant_type = + std::make_shared<::arrow::extension::VariantExtensionType>(storage_type); + + ASSERT_RAISES(Invalid, + ConvertSchema({::arrow::field("variant", std::move(variant_type))})); +} + +TEST_F(TestConvertArrowSchema, ParquetVariantShreddedList) { + auto field_group = ::arrow::struct_({::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", ::arrow::int64())}); + + auto element = GroupNode::Make( + "element", Repetition::REQUIRED, + {PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, ParquetType::INT64)}); + auto list = GroupNode::Make("list", Repetition::REPEATED, {element}); + const auto expected = GroupNode::Make( + "variant", Repetition::OPTIONAL, + {PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + GroupNode::Make("typed_value", Repetition::OPTIONAL, {list}, ConvertedType::LIST)}, + LogicalType::Variant()); + + const std::vector< + std::function(std::shared_ptr<::arrow::Field>)>> + make_lists = { + [](std::shared_ptr<::arrow::Field> field) { return ::arrow::list(field); }, + [](std::shared_ptr<::arrow::Field> field) { + return ::arrow::large_list(field); + }, + [](std::shared_ptr<::arrow::Field> field) { return ::arrow::list_view(field); }, + [](std::shared_ptr<::arrow::Field> field) { + return ::arrow::large_list_view(field); + }}; + for (const auto& make_list : make_lists) { + auto typed_value = + make_list(::arrow::field("element", field_group, /*nullable=*/false)); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", typed_value)}); + auto variant_type = ::arrow::extension::variant(storage_type); + + std::vector> arrow_fields = { + ::arrow::field("variant", variant_type)}; + std::vector parquet_fields = {expected}; + + ASSERT_OK(ConvertSchema(arrow_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); + } +} + TEST_F(TestConvertArrowSchema, ParquetGeoArrowCrsLonLat) { // All the Arrow Schemas below should convert to the type defaults for GEOMETRY // and GEOGRAPHY when GeoArrow extension types are registered and the appropriate diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index bc4de6c39b5e..9d7d71d64aa6 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -50,6 +50,8 @@ using arrow::Field; using arrow::FieldVector; using arrow::KeyValueMetadata; using arrow::Status; +using arrow::extension::kVariantExtensionName; +using arrow::extension::VariantExtensionType; using arrow::internal::checked_cast; using arrow::internal::ToChars; @@ -112,6 +114,239 @@ Status ListToNode(const std::shared_ptr<::arrow::BaseListType>& type, return Status::OK(); } +static constexpr char FIELD_ID_KEY[] = "PARQUET:field_id"; + +int FieldIdFromMetadata( + const std::shared_ptr& metadata) { + if (!metadata) { + return -1; + } + int key = metadata->FindKey(FIELD_ID_KEY); + if (key < 0) { + return -1; + } + const std::string& field_id_str = metadata->value(key); + int field_id; + if (::arrow::internal::ParseValue<::arrow::Int32Type>( + field_id_str.c_str(), field_id_str.length(), &field_id)) { + if (field_id < 0) { + // Thrift should convert any negative value to null but normalize to -1 here in + // case we later check this in logic. + return -1; + } + return field_id; + } else { + return -1; + } +} + +Status VariantTypedValueToNode(const std::shared_ptr& field, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + NodePtr* out); + +Status VariantFieldGroupToNode(const std::shared_ptr& field, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + NodePtr* out) { + const int field_id = FieldIdFromMetadata(field->metadata()); + const auto& type = field->type(); + if (type->id() != ArrowTypeId::STRUCT) { + return Status::Invalid("Invalid Variant shredded field group: ", type->ToString()); + } + + std::vector children; + children.reserve(type->num_fields()); + for (const auto& child : type->fields()) { + NodePtr node; + if (child->name() == "value") { + RETURN_NOT_OK( + FieldToNode(child->name(), child, properties, arrow_properties, &node)); + } else if (child->name() == "typed_value") { + RETURN_NOT_OK(VariantTypedValueToNode(child, properties, arrow_properties, &node)); + } else { + return Status::Invalid("Invalid Variant shredded field group child: ", + child->name()); + } + children.emplace_back(std::move(node)); + } + + *out = GroupNode::Make(field->name(), RepetitionFromNullable(field->nullable()), + std::move(children), nullptr, field_id); + return Status::OK(); +} + +Status VariantObjectToNode(const std::shared_ptr& field, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, NodePtr* out) { + const int field_id = FieldIdFromMetadata(field->metadata()); + const auto& type = field->type(); + if (type->num_fields() == 0) { + return Status::Invalid("Invalid Variant object typed_value: expected at least one ", + "shredded field"); + } + + std::vector children(type->num_fields()); + for (int i = 0; i < type->num_fields(); ++i) { + RETURN_NOT_OK(VariantFieldGroupToNode(type->field(i), properties, arrow_properties, + &children[i])); + } + + *out = GroupNode::Make(field->name(), RepetitionFromNullable(field->nullable()), + std::move(children), nullptr, field_id); + return Status::OK(); +} + +Status VariantListToNode(const std::shared_ptr& field, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, NodePtr* out) { + const int field_id = FieldIdFromMetadata(field->metadata()); + const auto list_type = std::static_pointer_cast<::arrow::BaseListType>(field->type()); + + NodePtr element; + RETURN_NOT_OK(VariantFieldGroupToNode(list_type->value_field()->WithName("element"), + properties, arrow_properties, &element)); + + NodePtr list = GroupNode::Make("list", Repetition::REPEATED, {element}); + *out = GroupNode::Make(field->name(), RepetitionFromNullable(field->nullable()), {list}, + LogicalType::List(), field_id); + return Status::OK(); +} + +Status VariantPrimitiveToNode(const std::shared_ptr& field, NodePtr* out) { + std::shared_ptr logical_type = LogicalType::None(); + ParquetType::type type = ParquetType::UNDEFINED; + const int field_id = FieldIdFromMetadata(field->metadata()); + + int length = -1; + int precision = -1; + int scale = -1; + + switch (field->type()->id()) { + case ArrowTypeId::BOOL: + type = ParquetType::BOOLEAN; + break; + case ArrowTypeId::INT8: + type = ParquetType::INT32; + logical_type = LogicalType::Int(8, true); + break; + case ArrowTypeId::INT16: + type = ParquetType::INT32; + logical_type = LogicalType::Int(16, true); + break; + case ArrowTypeId::INT32: + type = ParquetType::INT32; + break; + case ArrowTypeId::INT64: + type = ParquetType::INT64; + break; + case ArrowTypeId::FLOAT: + type = ParquetType::FLOAT; + break; + case ArrowTypeId::DOUBLE: + type = ParquetType::DOUBLE; + break; + case ArrowTypeId::DECIMAL32: + case ArrowTypeId::DECIMAL64: + case ArrowTypeId::DECIMAL128: { + const auto& decimal = checked_cast(*field->type()); + precision = decimal.precision(); + scale = decimal.scale(); + if (precision <= 9) { + type = ParquetType::INT32; + } else if (precision <= 18) { + type = ParquetType::INT64; + } else { + type = ParquetType::FIXED_LEN_BYTE_ARRAY; + length = ::arrow::DecimalType::DecimalSize(precision); + } + PARQUET_CATCH_NOT_OK(logical_type = LogicalType::Decimal(precision, scale)); + } break; + case ArrowTypeId::DATE32: + type = ParquetType::INT32; + logical_type = LogicalType::Date(); + break; + case ArrowTypeId::TIME64: + type = ParquetType::INT64; + logical_type = + LogicalType::Time(/*is_adjusted_to_utc=*/false, LogicalType::TimeUnit::MICROS); + break; + case ArrowTypeId::TIMESTAMP: { + type = ParquetType::INT64; + const auto& timestamp = checked_cast(*field->type()); + const bool utc = !timestamp.timezone().empty(); + switch (timestamp.unit()) { + case ::arrow::TimeUnit::MICRO: + logical_type = LogicalType::Timestamp(utc, LogicalType::TimeUnit::MICROS, + /*is_from_converted_type=*/false, + /*force_set_converted_type=*/true); + break; + case ::arrow::TimeUnit::NANO: + logical_type = LogicalType::Timestamp(utc, LogicalType::TimeUnit::NANOS, + /*is_from_converted_type=*/false, + /*force_set_converted_type=*/false); + break; + default: + return Status::Invalid("Invalid Variant typed_value timestamp unit: ", + field->type()->ToString()); + } + } break; + case ArrowTypeId::LARGE_BINARY: + case ArrowTypeId::BINARY: + case ArrowTypeId::BINARY_VIEW: + type = ParquetType::BYTE_ARRAY; + break; + case ArrowTypeId::LARGE_STRING: + case ArrowTypeId::STRING: + case ArrowTypeId::STRING_VIEW: + type = ParquetType::BYTE_ARRAY; + logical_type = LogicalType::String(); + break; + case ArrowTypeId::FIXED_SIZE_BINARY: + type = ParquetType::FIXED_LEN_BYTE_ARRAY; + logical_type = LogicalType::UUID(); + length = 16; + break; + case ArrowTypeId::EXTENSION: { + const auto& ext_type = checked_cast(*field->type()); + if (ext_type.extension_name() == "arrow.uuid") { + type = ParquetType::FIXED_LEN_BYTE_ARRAY; + logical_type = LogicalType::UUID(); + length = 16; + break; + } + return Status::Invalid("Invalid Variant typed_value type: ", + field->type()->ToString()); + } + default: + return Status::Invalid("Invalid Variant typed_value type: ", + field->type()->ToString()); + } + + PARQUET_CATCH_NOT_OK(*out = PrimitiveNode::Make( + field->name(), RepetitionFromNullable(field->nullable()), + std::move(logical_type), type, length, field_id)); + return Status::OK(); +} + +Status VariantTypedValueToNode(const std::shared_ptr& field, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + NodePtr* out) { + switch (field->type()->id()) { + case ArrowTypeId::STRUCT: + return VariantObjectToNode(field, properties, arrow_properties, out); + case ArrowTypeId::FIXED_SIZE_LIST: + case ArrowTypeId::LARGE_LIST: + case ArrowTypeId::LIST: + case ArrowTypeId::LIST_VIEW: + case ArrowTypeId::LARGE_LIST_VIEW: + return VariantListToNode(field, properties, arrow_properties, out); + default: + return VariantPrimitiveToNode(field, out); + } +} + Status MapToNode(const std::shared_ptr<::arrow::MapType>& type, const std::string& name, bool nullable, int field_id, const WriterProperties& properties, const ArrowWriterProperties& arrow_properties, NodePtr* out) { @@ -131,21 +366,33 @@ Status MapToNode(const std::shared_ptr<::arrow::MapType>& type, const std::strin return Status::OK(); } -Status VariantToNode( - const std::shared_ptr<::arrow::extension::VariantExtensionType>& type, - const std::string& name, bool nullable, int field_id, - const WriterProperties& properties, const ArrowWriterProperties& arrow_properties, - NodePtr* out) { - NodePtr metadata_node; - RETURN_NOT_OK(FieldToNode("metadata", type->metadata(), properties, arrow_properties, - &metadata_node)); +Status VariantToNode(const std::shared_ptr& type, + const std::string& name, bool nullable, int field_id, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, NodePtr* out) { + std::vector children; + children.reserve(type->storage_type()->num_fields()); - NodePtr value_node; - RETURN_NOT_OK( - FieldToNode("value", type->value(), properties, arrow_properties, &value_node)); + auto AddChild = [&](const std::shared_ptr<::arrow::Field>& field) { + if (field == nullptr) { + return Status::OK(); + } + NodePtr child; + if (field->name() == "typed_value") { + RETURN_NOT_OK(VariantTypedValueToNode(field, properties, arrow_properties, &child)); + } else { + RETURN_NOT_OK( + FieldToNode(field->name(), field, properties, arrow_properties, &child)); + } + children.emplace_back(std::move(child)); + return Status::OK(); + }; + + RETURN_NOT_OK(AddChild(type->metadata())); + RETURN_NOT_OK(AddChild(type->value())); + RETURN_NOT_OK(AddChild(type->typed_value())); - *out = GroupNode::Make(name, RepetitionFromNullable(nullable), - {std::move(metadata_node), std::move(value_node)}, + *out = GroupNode::Make(name, RepetitionFromNullable(nullable), std::move(children), LogicalType::Variant(), field_id); return Status::OK(); @@ -280,8 +527,6 @@ static Status GetTimestampMetadata(const ::arrow::TimestampType& type, return Status::OK(); } -static constexpr char FIELD_ID_KEY[] = "PARQUET:field_id"; - std::shared_ptr<::arrow::KeyValueMetadata> FieldIdMetadata(int field_id) { if (field_id >= 0) { return ::arrow::key_value_metadata({FIELD_ID_KEY}, {ToChars(field_id)}); @@ -290,30 +535,6 @@ std::shared_ptr<::arrow::KeyValueMetadata> FieldIdMetadata(int field_id) { } } -int FieldIdFromMetadata( - const std::shared_ptr& metadata) { - if (!metadata) { - return -1; - } - int key = metadata->FindKey(FIELD_ID_KEY); - if (key < 0) { - return -1; - } - std::string field_id_str = metadata->value(key); - int field_id; - if (::arrow::internal::ParseValue<::arrow::Int32Type>( - field_id_str.c_str(), field_id_str.length(), &field_id)) { - if (field_id < 0) { - // Thrift should convert any negative value to null but normalize to -1 here in - // case we later check this in logic. - return -1; - } - return field_id; - } else { - return -1; - } -} - Status FieldToNode(const std::string& name, const std::shared_ptr& field, const WriterProperties& properties, const ArrowWriterProperties& arrow_properties, NodePtr* out) { @@ -466,8 +687,15 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, case ArrowTypeId::DICTIONARY: { // Parquet has no Dictionary type, dictionary-encoded is handled on // the encoding, not the schema level. - const ::arrow::DictionaryType& dict_type = - static_cast(*field->type()); + const auto& dict_type = static_cast(*field->type()); + if (dict_type.value_type()->id() == ArrowTypeId::EXTENSION) { + const auto& ext_type = + checked_cast(*dict_type.value_type()); + if (ext_type.extension_name() == kVariantExtensionName) { + return Status::NotImplemented( + "Dictionary-encoded Variant arrays are not supported"); + } + } std::shared_ptr<::arrow::Field> unpacked_field = ::arrow::field( name, dict_type.value_type(), field->nullable(), field->metadata()); return FieldToNode(name, unpacked_field, properties, arrow_properties, out); @@ -490,10 +718,8 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, ARROW_ASSIGN_OR_RAISE(logical_type, LogicalTypeFromGeoArrowMetadata(ext_type->Serialize())); break; - } else if (ext_type->extension_name() == std::string("arrow.parquet.variant")) { - auto variant_type = - std::static_pointer_cast<::arrow::extension::VariantExtensionType>( - field->type()); + } else if (ext_type->extension_name() == kVariantExtensionName) { + auto variant_type = std::static_pointer_cast(field->type()); return VariantToNode(variant_type, name, field->nullable(), field_id, properties, arrow_properties, out); @@ -543,6 +769,17 @@ bool IsDictionaryReadSupported(const ArrowType& type) { return type.id() == ::arrow::Type::BINARY || type.id() == ::arrow::Type::STRING; } +bool IsInsideVariantLogicalType(const Node& node) { + const Node* current = node.parent(); + while (current != nullptr) { + if (current->logical_type()->is_variant()) { + return true; + } + current = current->parent(); + } + return false; +} + // ---------------------------------------------------------------------- // Schema logic @@ -552,7 +789,8 @@ ::arrow::Result> GetTypeForNode( ARROW_ASSIGN_OR_RAISE(std::shared_ptr storage_type, GetArrowType(primitive_node, ctx->properties, ctx->metadata)); if (ctx->properties.read_dictionary(column_index) && - IsDictionaryReadSupported(*storage_type)) { + IsDictionaryReadSupported(*storage_type) && + !IsInsideVariantLogicalType(primitive_node)) { return ::arrow::dictionary(::arrow::int32(), storage_type); } return storage_type; @@ -604,7 +842,7 @@ Status GroupToStruct(const GroupNode& node, LevelInfo current_levels, auto struct_type = ::arrow::struct_(arrow_fields); if (ctx->properties.get_arrow_extensions_enabled() && node.logical_type()->is_variant()) { - auto extension_type = ::arrow::GetExtensionType("arrow.parquet.variant"); + auto extension_type = ::arrow::GetExtensionType(std::string(kVariantExtensionName)); if (extension_type) { ARROW_ASSIGN_OR_RAISE( struct_type, @@ -1194,10 +1432,10 @@ Result ApplyOriginalMetadata(const Field& origin_field, SchemaField* infer extension_supports_inferred_storage = arrow_extension_inferred || ::arrow::extension::UuidType::IsSupportedStorageType(inferred_type); - } else if (origin_extension_name == "arrow.parquet.variant") { + } else if (origin_extension_name == kVariantExtensionName) { extension_supports_inferred_storage = arrow_extension_inferred || - ::arrow::extension::VariantExtensionType::IsSupportedStorageType(inferred_type); + VariantExtensionType::IsSupportedStorageType(inferred_type); } else { extension_supports_inferred_storage = origin_extension_type.storage_type()->Equals(*inferred_type); diff --git a/cpp/src/parquet/arrow/variant_test.cc b/cpp/src/parquet/arrow/variant_test.cc index 04f46d2e444d..62e68f3e3ad5 100644 --- a/cpp/src/parquet/arrow/variant_test.cc +++ b/cpp/src/parquet/arrow/variant_test.cc @@ -15,58 +15,348 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/array/validate.h" +#include "arrow/array.h" // IWYU pragma: keep #include "arrow/extension/parquet_variant.h" -#include "arrow/ipc/test_common.h" +#include "arrow/extension/uuid.h" +#include "arrow/io/memory.h" #include "arrow/record_batch.h" +#include "arrow/table.h" +#include "arrow/testing/extension_type.h" #include "arrow/testing/gtest_util.h" +#include "parquet/arrow/reader.h" +#include "parquet/arrow/writer.h" #include "parquet/exception.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/test_util_internal.h" + +#include +#include +#include namespace parquet::arrow { using ::arrow::binary; +using ::arrow::field; using ::arrow::struct_; +using variant::internal::BinaryArrayFromValues; +using variant::internal::BinaryViewArrayFromValues; +using variant::internal::EmptyVariantMetadata; +using variant::internal::Int32ArrayFromValues; +using variant::internal::Int64ArrayFromValues; +using variant::internal::Int8Variant; +using variant::internal::StringArrayFromValues; +using variant::internal::UuidArrayFromValues; +using variant::internal::VariantTable; +using variant::internal::WriteVariantRecordBatch; +using variant::internal::WriteVariantTable; + +TEST(TestVariantExtensionType, WriterValidatesUnshreddedVariantBytes) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues({std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues({std::string_view{*encoded.value}}); + auto table = + VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); + ASSERT_OK(WriteVariantTable(table)); + + auto invalid_value = BinaryArrayFromValues({std::string_view("\xff", 1)}); + auto invalid_table = + VariantTable(variant_type, {metadata_array, invalid_value}, storage_type->fields()); + ASSERT_RAISES(Invalid, WriteVariantTable(invalid_table)); + + auto no_validation = + ArrowWriterProperties::Builder().set_variant_validation_enabled(false)->build(); + ASSERT_OK(WriteVariantTable(invalid_table, default_writer_properties(), no_validation)); +} + +TEST(TestVariantExtensionType, WriteRecordBatchValidatesVariantBytes) { + ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + + auto metadata_array = BinaryArrayFromValues({std::string_view{*metadata}}); + auto invalid_value = BinaryArrayFromValues({std::string_view("\xff", 1)}); + auto table = + VariantTable(variant_type, {metadata_array, invalid_value}, storage_type->fields()); + + ASSERT_RAISES(Invalid, WriteVariantRecordBatch(table)); + + auto no_validation = + ArrowWriterProperties::Builder().set_variant_validation_enabled(false)->build(); + ASSERT_OK(WriteVariantRecordBatch(table, no_validation)); +} + +TEST(TestVariantExtensionType, WriteRecordBatchValidatesBatch) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = StringArrayFromValues({"not binary"}); + auto value_array = BinaryArrayFromValues({std::string_view{*encoded.value}}); + ASSERT_OK_AND_ASSIGN( + auto storage, + ::arrow::StructArray::Make({metadata_array, value_array}, storage_type->fields())); + auto variant_array = ::arrow::ExtensionType::WrapArray(variant_type, storage); + auto schema = ::arrow::schema({field("variant", variant_type)}); + auto batch = ::arrow::RecordBatch::Make(schema, 1, {variant_array}); + + ASSERT_OK_AND_ASSIGN(auto sink, ::arrow::io::BufferOutputStream::Create( + 1024, ::arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN( + auto writer, + FileWriter::Open(*schema, ::arrow::default_memory_pool(), sink, + default_writer_properties(), default_arrow_writer_properties())); + + const auto status = writer->WriteRecordBatch(*batch); + ASSERT_TRUE(status.IsInvalid()) << status; + ASSERT_NE(std::string::npos, status.ToStringWithoutContextLines().find( + "Struct child array #0 does not match type field")) + << status; +} + +TEST(TestVariantExtensionType, WriterValidatesBinaryViewVariantBytes) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = + struct_({field("metadata", ::arrow::binary_view(), /*nullable=*/false), + field("value", ::arrow::binary_view(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryViewArrayFromValues({std::string_view{*encoded.metadata}}); + auto value_array = BinaryViewArrayFromValues({std::string_view{*encoded.value}}); + auto table = + VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); + ASSERT_OK(WriteVariantTable(table)); +} + +TEST(TestVariantExtensionType, WriterSkipsNullParents) { + ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues({std::string_view{*metadata}}); + auto invalid_value = BinaryArrayFromValues({std::string_view("\xff", 1)}); + PARQUET_ASSIGN_OR_THROW(auto storage, + ::arrow::StructArray::Make({metadata_array, invalid_value}, + storage_type->fields())); + auto variant_array = ::arrow::ExtensionType::WrapArray(variant_type, storage); + + auto parent_type = struct_({field("child", variant_type)}); + PARQUET_ASSIGN_OR_THROW( + auto parent_array, + ::arrow::StructArray::Make({variant_array}, parent_type->fields(), + ::arrow::Buffer::FromString(std::string("\0", 1)))); + auto parent_table = ::arrow::Table::Make( + ::arrow::schema({field("parent", parent_type)}), {parent_array}); + ASSERT_OK(WriteVariantTable(parent_table)); + + auto map_type = ::arrow::map(::arrow::utf8(), field("item", variant_type)); + ASSERT_OK_AND_ASSIGN( + auto map_array, + ::arrow::MapArray::FromArrays(map_type, Int32ArrayFromValues({0, 1}), + StringArrayFromValues({"hidden"}), variant_array, + ::arrow::default_memory_pool(), + ::arrow::Buffer::FromString(std::string("\0", 1)))); + auto map_table = ::arrow::Table::Make(::arrow::schema({field("variant_map", map_type)}), + {map_array}); + ASSERT_OK(WriteVariantTable(map_table)); +} + +TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", ::arrow::int64())}); + auto variant_type = ::arrow::extension::variant(storage_type); + + auto metadata_array = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto value_array = + BinaryArrayFromValues({std::nullopt, std::string_view{*encoded.value}}); + auto typed_array = Int64ArrayFromValues({34, 100}); + auto table = VariantTable(variant_type, {metadata_array, value_array, typed_array}, + storage_type->fields()); + ASSERT_RAISES(Invalid, WriteVariantTable(table)); + + auto empty_value = BinaryArrayFromValues({std::string_view{}, std::nullopt}); + auto empty_table = VariantTable( + variant_type, {metadata_array, empty_value, typed_array}, storage_type->fields()); + ASSERT_RAISES(Invalid, WriteVariantTable(empty_table)); + + auto valid_values = + BinaryArrayFromValues({std::nullopt, std::string_view{*encoded.value}}); + auto valid_typed = Int64ArrayFromValues({34, std::nullopt}); + auto valid_table = VariantTable( + variant_type, {metadata_array, valid_values, valid_typed}, storage_type->fields()); + ASSERT_OK(WriteVariantTable(valid_table)); +} + +TEST(TestVariantExtensionType, WriterValidatesShreddedWithoutValue) { + ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("typed_value", ::arrow::int64())}); + ASSERT_OK_AND_ASSIGN(auto variant_type, + ::arrow::extension::VariantExtensionType::Make(storage_type)); + + auto metadata_array = BinaryArrayFromValues({std::string_view{*metadata}}); + auto typed_array = Int64ArrayFromValues({34}); + auto table = + VariantTable(variant_type, {metadata_array, typed_array}, storage_type->fields()); + ASSERT_OK(WriteVariantTable(table)); +} + +TEST(TestVariantExtensionType, ReadsDictionaryEncodedMetadata) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues( + {std::string_view{*encoded.value}, std::string_view{*encoded.value}}); + auto table = + VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); + + ASSERT_OK_AND_ASSIGN( + auto buffer, + WriteVariantTable(table, WriterProperties::Builder().enable_dictionary()->build())); + + auto buffer_reader = std::make_shared<::arrow::io::BufferReader>(buffer); + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + ::arrow::ExtensionTypeGuard guard(::arrow::extension::variant(storage_type)); + FileReaderBuilder builder; + ASSERT_OK(builder.Open(buffer_reader)); + builder.properties(reader_properties); + ASSERT_OK_AND_ASSIGN(auto reader, builder.Build()); + + ASSERT_TRUE(reader->parquet_reader() + ->metadata() + ->RowGroup(0) + ->ColumnChunk(0) + ->has_dictionary_page()); + + ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); + ASSERT_OK(read_table->ValidateFull()); + + auto field = read_table->schema()->GetFieldByName("variant"); + ASSERT_NE(nullptr, field); + auto read_variant_type = + std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>(field->type()); + ASSERT_NE(nullptr, read_variant_type); + ASSERT_EQ(::arrow::Type::BINARY, read_variant_type->metadata()->type()->id()); + + ASSERT_NE(nullptr, read_table->GetColumnByName(field->name())); +} + +TEST(TestVariantExtensionType, ReadsWithDictionaryOption) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues( + {std::string_view{*encoded.value}, std::string_view{*encoded.value}}); + auto table = + VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); + + ASSERT_OK_AND_ASSIGN(auto buffer, WriteVariantTable(table)); + + auto buffer_reader = std::make_shared<::arrow::io::BufferReader>(buffer); + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + reader_properties.set_read_dictionary(0, true); + reader_properties.set_read_dictionary(1, true); + ::arrow::ExtensionTypeGuard guard(::arrow::extension::variant(storage_type)); + FileReaderBuilder builder; + ASSERT_OK(builder.Open(buffer_reader)); + builder.properties(reader_properties); + ASSERT_OK_AND_ASSIGN(auto reader, builder.Build()); + + ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); + ASSERT_OK(read_table->ValidateFull()); + + auto field = read_table->schema()->GetFieldByName("variant"); + ASSERT_NE(nullptr, field); + auto read_variant_type = + std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>(field->type()); + ASSERT_NE(nullptr, read_variant_type); + ASSERT_EQ(::arrow::Type::BINARY, read_variant_type->metadata()->type()->id()); + ASSERT_EQ(::arrow::Type::BINARY, read_variant_type->value()->type()->id()); + + ASSERT_NE(nullptr, read_table->GetColumnByName(field->name())); +} + +TEST(TestVariantExtensionType, WriterWritesUuid) { + ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto storage_type = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", ::arrow::extension::uuid())}); + auto variant_type = ::arrow::extension::variant(storage_type); + + auto metadata_array = BinaryArrayFromValues({std::string_view{*metadata}}); + auto value_array = BinaryArrayFromValues({std::nullopt}); + auto typed_array = UuidArrayFromValues({std::string_view("0123456789abcdef", 16)}); + auto table = VariantTable(variant_type, {metadata_array, value_array, typed_array}, + storage_type->fields()); + ASSERT_OK(WriteVariantTable(table)); +} + +TEST(TestVariantExtensionType, WriterValidatesShreddedObjectConflicts) { + variant::VariantBuilder object_builder; + ASSERT_OK_AND_ASSIGN(auto object, object_builder.StartObject()); + ASSERT_OK(object.AppendShortString("event_type", "login")); + ASSERT_OK(object.Finish()); + ASSERT_OK_AND_ASSIGN(auto encoded, object_builder.Finish()); + + auto field_group_type = + struct_({field("value", binary()), field("typed_value", ::arrow::utf8())}); + auto typed_value_type = + struct_({field("event_type", field_group_type, /*nullable=*/false)}); + auto storage_type = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", typed_value_type)}); + auto variant_type = ::arrow::extension::variant(storage_type); + + auto metadata_array = BinaryArrayFromValues({std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues({std::string_view{*encoded.value}}); + ASSERT_OK_AND_ASSIGN(auto event_type_group, + ::arrow::StructArray::Make({BinaryArrayFromValues({std::nullopt}), + StringArrayFromValues({"login"})}, + field_group_type->fields())); + ASSERT_OK_AND_ASSIGN( + auto typed_array, + ::arrow::StructArray::Make({event_type_group}, typed_value_type->fields())); + + auto table = VariantTable(variant_type, {metadata_array, value_array, typed_array}, + storage_type->fields()); + ASSERT_RAISES(Invalid, WriteVariantTable(table)); + + auto valid_value_array = BinaryArrayFromValues({std::nullopt}); + auto valid_table = + VariantTable(variant_type, {metadata_array, valid_value_array, typed_array}, + storage_type->fields()); + ASSERT_OK(WriteVariantTable(valid_table)); -TEST(TestVariantExtensionType, StorageTypeValidation) { - auto variant1 = ::arrow::extension::variant( - struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary(), /*nullable=*/false)})); - auto variant2 = ::arrow::extension::variant( - struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary(), /*nullable=*/false)})); - - ASSERT_TRUE(variant1->Equals(variant2)); - - // Metadata and value fields can be provided in either order - auto variantFieldsFlipped = - std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>( - ::arrow::extension::variant( - struct_({field("value", binary(), /*nullable=*/false), - field("metadata", binary(), /*nullable=*/false)}))); - - ASSERT_EQ("metadata", variantFieldsFlipped->metadata()->name()); - ASSERT_EQ("value", variantFieldsFlipped->value()->name()); - - auto missing_value = struct_({field("metadata", binary(), /*nullable=*/false)}); - auto missing_metadata = struct_({field("value", binary(), /*nullable=*/false)}); - auto bad_value_type = struct_({field("metadata", binary(), /*nullable=*/false), - field("value", ::arrow::int32(), /*nullable=*/false)}); - auto extra_field = struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary(), /*nullable=*/false), - field("extra", binary(), /*nullable=*/false)}); - auto nullable_metadata = struct_( - {field("metadata", binary()), field("value", binary(), /*nullable=*/false)}); - auto nullable_value = struct_( - {field("metadata", binary(), /*nullable=*/false), field("value", binary())}); - - for (const auto& storage_type : {missing_value, missing_metadata, bad_value_type, - extra_field, nullable_metadata, nullable_value}) { - ASSERT_RAISES_WITH_MESSAGE( - Invalid, - "Invalid: Invalid storage type for VariantExtensionType: " + - storage_type->ToString(), - ::arrow::extension::VariantExtensionType::Make(storage_type)); - } + ASSERT_OK_AND_ASSIGN(auto missing_event_type_group, + ::arrow::StructArray::Make({BinaryArrayFromValues({std::nullopt}), + StringArrayFromValues({std::nullopt})}, + field_group_type->fields())); + ASSERT_OK_AND_ASSIGN( + auto missing_typed_array, + ::arrow::StructArray::Make({missing_event_type_group}, typed_value_type->fields())); + auto missing_table = + VariantTable(variant_type, {metadata_array, valid_value_array, missing_typed_array}, + storage_type->fields()); + ASSERT_OK(WriteVariantTable(missing_table)); } } // namespace parquet::arrow diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index e0fbe308219c..024f5f6e7699 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -18,14 +18,12 @@ #include "parquet/arrow/writer.h" #include -#include #include #include -#include #include #include -#include "arrow/array.h" +#include "arrow/array.h" // IWYU pragma: keep #include "arrow/array/concatenate.h" #include "arrow/extension_type.h" #include "arrow/ipc/writer.h" @@ -46,6 +44,7 @@ #include "parquet/file_writer.h" #include "parquet/platform.h" #include "parquet/schema.h" +#include "parquet/variant/validate.h" using arrow::Array; using arrow::BinaryArray; @@ -62,7 +61,6 @@ using arrow::MemoryPool; using arrow::NumericArray; using arrow::PrimitiveArray; using arrow::RecordBatch; -using arrow::ResizableBuffer; using arrow::Result; using arrow::Status; using arrow::Table; @@ -323,6 +321,7 @@ class FileWriterImpl : public FileWriter { std::unique_ptr writer, std::shared_ptr arrow_properties) : schema_(std::move(schema)), + pool_(pool), writer_(std::move(writer)), row_group_writer_(nullptr), column_write_context_(pool, arrow_properties.get()), @@ -382,6 +381,9 @@ class FileWriterImpl : public FileWriter { Status WriteColumnChunk(const std::shared_ptr& data, int64_t offset, int64_t size) override { RETURN_NOT_OK(CheckClosed()); + if (arrow_properties_->variant_validation_enabled()) { + RETURN_NOT_OK(variant::ValidateVariants(*data->Slice(offset, size), pool_)); + } if (arrow_properties_->engine_version() == ArrowWriterProperties::V2 || arrow_properties_->engine_version() == ArrowWriterProperties::V1) { if (row_group_writer_->buffered()) { @@ -450,6 +452,9 @@ class FileWriterImpl : public FileWriter { Status WriteRecordBatch(const RecordBatch& batch) override { RETURN_NOT_OK(CheckClosed()); + if (arrow_properties_->variant_validation_enabled()) { + RETURN_NOT_OK(batch.Validate()); + } if (batch.num_rows() == 0) { return Status::OK(); } @@ -469,6 +474,10 @@ class FileWriterImpl : public FileWriter { for (int i = 0; i < batch.num_columns(); i++) { ChunkedArray chunked_array{batch.column(i)}; + if (arrow_properties_->variant_validation_enabled()) { + RETURN_NOT_OK( + variant::ValidateVariants(*chunked_array.Slice(offset, size), pool_)); + } ARROW_ASSIGN_OR_RAISE( std::unique_ptr writer, ArrowColumnWriterV2::Make(chunked_array, offset, size, schema_manifest_, @@ -532,6 +541,7 @@ class FileWriterImpl : public FileWriter { friend class FileWriter; std::shared_ptr<::arrow::Schema> schema_; + MemoryPool* pool_; SchemaManifest schema_manifest_; diff --git a/cpp/src/parquet/meson.build b/cpp/src/parquet/meson.build index 9069ccb5fd1a..d7737f7db0f9 100644 --- a/cpp/src/parquet/meson.build +++ b/cpp/src/parquet/meson.build @@ -55,6 +55,9 @@ parquet_srcs = files( 'stream_reader.cc', 'stream_writer.cc', 'types.cc', + 'variant/builder.cc', + 'variant/encoding.cc', + 'variant/validate.cc', 'xxhasher.cc', ) @@ -130,6 +133,7 @@ subdir('api') subdir('arrow') subdir('encryption') subdir('geospatial') +subdir('variant') install_headers( [ @@ -229,6 +233,7 @@ parquet_tests = { 'arrow/arrow_reader_writer_test.cc', 'arrow/arrow_statistics_test.cc', 'arrow/variant_test.cc', + 'variant/test_util_internal.cc', ), }, 'arrow-index-test': {'sources': files('arrow/index_test.cc')}, diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index e2244a1176e3..e00ec2fe7dd6 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1259,10 +1259,9 @@ class PARQUET_EXPORT ArrowReaderProperties { /// Enable Parquet-supported Arrow extension types. /// /// When enabled, Parquet logical types will be mapped to their corresponding Arrow - /// extension types at read time, if such exist. Currently only arrow::extension::json() - /// extension type is supported. Columns whose LogicalType is JSON will be interpreted - /// as arrow::extension::json(), with storage type inferred from the serialized Arrow - /// schema if present, or `utf8` by default. + /// extension types at read time, if such exist. For example, columns whose LogicalType + /// is JSON will be interpreted as arrow::extension::json(), with storage type inferred + /// from the serialized Arrow schema if present, or `utf8` by default. void set_arrow_extensions_enabled(bool extensions_enabled) { arrow_extensions_enabled_ = extensions_enabled; } @@ -1332,7 +1331,8 @@ class PARQUET_EXPORT ArrowWriterProperties { engine_version_(V2), use_threads_(kArrowDefaultUseThreads), executor_(NULLPTR), - write_time_adjusted_to_utc_(false) {} + write_time_adjusted_to_utc_(false), + variant_validation_enabled_(true) {} /// \brief Disable writing legacy int96 timestamps (default disabled). Builder* disable_deprecated_int96_timestamps() { @@ -1436,12 +1436,23 @@ class PARQUET_EXPORT ArrowWriterProperties { return this; } + /// \brief Set whether to validate Parquet Variant binary values before writing. + /// + /// This is enabled by default. When enabled, Variant metadata/value bytes + /// are checked against the Parquet Variant encoding, and shredded value / + /// typed_value combinations are checked for conflicts. + Builder* set_variant_validation_enabled(bool enabled) { + variant_validation_enabled_ = enabled; + return this; + } + /// Create the final properties. std::shared_ptr build() { return std::shared_ptr(new ArrowWriterProperties( write_timestamps_as_int96_, coerce_timestamps_enabled_, coerce_timestamps_unit_, truncated_timestamps_allowed_, store_schema_, compliant_nested_types_, - engine_version_, use_threads_, executor_, write_time_adjusted_to_utc_)); + engine_version_, use_threads_, executor_, write_time_adjusted_to_utc_, + variant_validation_enabled_)); } private: @@ -1459,6 +1470,7 @@ class PARQUET_EXPORT ArrowWriterProperties { ::arrow::internal::Executor* executor_; bool write_time_adjusted_to_utc_; + bool variant_validation_enabled_; }; bool support_deprecated_int96_timestamps() const { return write_timestamps_as_int96_; } @@ -1497,15 +1509,16 @@ class PARQUET_EXPORT ArrowWriterProperties { /// Note this setting doesn't affect TIMESTAMP data. bool write_time_adjusted_to_utc() const { return write_time_adjusted_to_utc_; } + /// \brief Returns whether Parquet Variant binary values are validated before writing. + bool variant_validation_enabled() const { return variant_validation_enabled_; } + private: - explicit ArrowWriterProperties(bool write_nanos_as_int96, - bool coerce_timestamps_enabled, - ::arrow::TimeUnit::type coerce_timestamps_unit, - bool truncated_timestamps_allowed, bool store_schema, - bool compliant_nested_types, - EngineVersion engine_version, bool use_threads, - ::arrow::internal::Executor* executor, - bool write_time_adjusted_to_utc) + explicit ArrowWriterProperties( + bool write_nanos_as_int96, bool coerce_timestamps_enabled, + ::arrow::TimeUnit::type coerce_timestamps_unit, bool truncated_timestamps_allowed, + bool store_schema, bool compliant_nested_types, EngineVersion engine_version, + bool use_threads, ::arrow::internal::Executor* executor, + bool write_time_adjusted_to_utc, bool variant_validation_enabled) : write_timestamps_as_int96_(write_nanos_as_int96), coerce_timestamps_enabled_(coerce_timestamps_enabled), coerce_timestamps_unit_(coerce_timestamps_unit), @@ -1515,7 +1528,8 @@ class PARQUET_EXPORT ArrowWriterProperties { engine_version_(engine_version), use_threads_(use_threads), executor_(executor), - write_time_adjusted_to_utc_(write_time_adjusted_to_utc) {} + write_time_adjusted_to_utc_(write_time_adjusted_to_utc), + variant_validation_enabled_(variant_validation_enabled) {} const bool write_timestamps_as_int96_; const bool coerce_timestamps_enabled_; @@ -1527,6 +1541,7 @@ class PARQUET_EXPORT ArrowWriterProperties { const bool use_threads_; ::arrow::internal::Executor* executor_; const bool write_time_adjusted_to_utc_; + const bool variant_validation_enabled_; }; /// \brief State object used for writing Arrow data directly to a Parquet diff --git a/cpp/src/parquet/variant/CMakeLists.txt b/cpp/src/parquet/variant/CMakeLists.txt new file mode 100644 index 000000000000..8c9b99266331 --- /dev/null +++ b/cpp/src/parquet/variant/CMakeLists.txt @@ -0,0 +1,26 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +add_parquet_test(variant-test + SOURCES + builder_test.cc + encoding_test.cc + test_util_internal.cc + type_test.cc + validate_test.cc) + +arrow_install_all_headers("parquet/variant") diff --git a/cpp/src/parquet/variant/builder.cc b/cpp/src/parquet/variant/builder.cc new file mode 100644 index 000000000000..bd8540b368ee --- /dev/null +++ b/cpp/src/parquet/variant/builder.cc @@ -0,0 +1,1273 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/builder.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/buffer_builder.h" +#include "arrow/builder.h" // IWYU pragma: keep +#include "arrow/util/checked_cast.h" +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" +#include "parquet/variant/encoding.h" +#include "parquet/variant/encoding_internal.h" + +namespace parquet::variant { + +using ::arrow::binary; +using ::arrow::BinaryBuilder; +using ::arrow::BooleanArray; +using ::arrow::BooleanBuilder; +using ::arrow::BufferBuilder; +using ::arrow::ExtensionType; +using ::arrow::field; +using ::arrow::struct_; +using ::arrow::StructType; +using ::arrow::Type; +using ::arrow::extension::VariantExtensionType; + +namespace bit_util = ::arrow::bit_util; + +namespace { + +Status AppendLittleEndian(BufferBuilder& out, uint32_t value, uint8_t width) { + DCHECK_LE(width, sizeof(uint32_t)); + const auto little_endian = bit_util::ToLittleEndian(value); + return out.Append(&little_endian, width); +} + +template + requires(std::is_arithmetic_v) +Status AppendFixedLittleEndian(BufferBuilder& out, T value) { + const auto little_endian = bit_util::ToLittleEndian(value); + return out.Append(&little_endian, sizeof(T)); +} + +void AppendLittleEndianToString(std::string& out, uint32_t value, uint8_t width) { + DCHECK_LE(width, sizeof(uint32_t)); + const auto little_endian = bit_util::ToLittleEndian(value); + out.append(reinterpret_cast(&little_endian), width); +} + +Status InsertBytes(BufferBuilder& out, int64_t offset, std::string_view bytes) { + const int64_t old_size = out.length(); + const auto insert_size = bytes.size(); + ARROW_RETURN_NOT_OK(out.Reserve(insert_size)); + uint8_t* data = out.mutable_data(); + std::memmove(data + offset + insert_size, data + offset, old_size - offset); + std::memcpy(data + offset, bytes.data(), insert_size); + out.UnsafeAdvance(insert_size); + return Status::OK(); +} + +uint8_t WidthForValue(uint64_t value) { + if (value <= std::numeric_limits::max()) { + return 1; + } + if (value <= std::numeric_limits::max()) { + return 2; + } + if (value <= 0xFFFFFFU) { + return 3; + } + return 4; +} + +class VariantMetadataBuilder { + public: + explicit VariantMetadataBuilder(MemoryPool* pool) : pool_(pool) {} + + Status Reserve(int64_t capacity) { + if (capacity < 0) { + return Status::Invalid("Variant metadata capacity must be non-negative"); + } + field_ids_.reserve(capacity); + return Status::OK(); + } + + Result Upsert(std::string_view name) { + auto it = field_ids_.find(name); + if (it != field_ids_.end()) { + return it->second; + } + if (field_names_.size() >= std::numeric_limits::max()) { + return Status::Invalid("Variant metadata dictionary is too large"); + } + ARROW_RETURN_NOT_OK(internal::ValidateUtf8(name, "metadata dictionary string")); + + const auto id = field_names_.size(); + if (field_names_.empty()) { + is_sorted_ = true; + } else if (is_sorted_ && !(field_names_.back() < name)) { + is_sorted_ = false; + } + field_names_.emplace_back(name); + field_ids_.emplace(field_names_.back(), id); + return static_cast(id); + } + + std::string_view FieldName(uint32_t id) const { + DCHECK_LT(id, field_names_.size()); + return field_names_[id]; + } + + size_t size() const { return field_names_.size(); } + + void Truncate(size_t size) { + DCHECK_LE(size, field_names_.size()); + field_names_.resize(size); + RebuildIndex(); + } + + Result> Finish() const { + uint64_t bytes_size = 0; + for (const auto& string : field_names_) { + bytes_size += string.size(); + } + if (field_names_.size() > std::numeric_limits::max() || + bytes_size > std::numeric_limits::max()) { + return Status::Invalid("Variant metadata dictionary is too large"); + } + + const uint8_t offset_size = + WidthForValue(std::max(field_names_.size(), bytes_size)); + BufferBuilder out(pool_); + ARROW_RETURN_NOT_OK(out.Reserve( + 1 + offset_size + (field_names_.size() + 1) * offset_size + bytes_size)); + + const bool sorted_strings = !field_names_.empty() && is_sorted_; + const auto header = + static_cast(internal::kVariantVersion | + (sorted_strings ? internal::kMetadataSortedStringsMask : 0) | + ((offset_size - 1) << 6)); + ARROW_RETURN_NOT_OK(out.Append(&header, sizeof(header))); + ARROW_RETURN_NOT_OK( + AppendLittleEndian(out, static_cast(field_names_.size()), offset_size)); + + uint32_t offset = 0; + ARROW_RETURN_NOT_OK(AppendLittleEndian(out, offset, offset_size)); + for (const auto& string : field_names_) { + offset += static_cast(string.size()); + ARROW_RETURN_NOT_OK(AppendLittleEndian(out, offset, offset_size)); + } + for (const auto& string : field_names_) { + ARROW_RETURN_NOT_OK(out.Append(string)); + } + return out.Finish(); + } + + private: + void RebuildIndex() { + field_ids_.clear(); + field_ids_.reserve(field_names_.size()); + is_sorted_ = false; + for (uint32_t i = 0; i < field_names_.size(); ++i) { + field_ids_.emplace(field_names_[i], i); + if (i == 0) { + is_sorted_ = true; + } else if (is_sorted_ && !(field_names_[i - 1] < field_names_[i])) { + is_sorted_ = false; + } + } + } + + MemoryPool* pool_; + std::deque field_names_; + std::unordered_map field_ids_; + bool is_sorted_ = false; +}; + +class VariantValueWriter { + public: + explicit VariantValueWriter(BufferBuilder& out) : out_(out) {} + + template + requires internal::HeaderOnlyVariantPrimitive + Status Append() { + return AppendPrimitiveHeader(); + } + + template + requires internal::FixedVariantPrimitive + Status Append(typename internal::VariantFixedPrimitiveTraits::CType value) { + using CType = typename internal::VariantFixedPrimitiveTraits::CType; + ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); + if constexpr (sizeof(CType) == 1) { + const auto byte = static_cast(value); + return out_.Append(&byte, sizeof(byte)); + } else { + return AppendFixedLittleEndian(out_, value); + } + } + + template + requires internal::DecimalVariantPrimitive + Status Append( + typename internal::VariantDecimalPrimitiveTraits::CType unscaled_value, + uint8_t scale) { + ARROW_RETURN_NOT_OK(internal::ValidateDecimalScale(scale)); + ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); + ARROW_RETURN_NOT_OK(out_.Append(&scale, sizeof(scale))); + return AppendFixedLittleEndian(out_, unscaled_value); + } + + template + requires internal::Decimal16VariantPrimitive + Status Append(std::string_view little_endian_unscaled_value, uint8_t scale) { + if (little_endian_unscaled_value.size() != 16) { + return Status::Invalid("Variant Decimal16 values must be 16 bytes"); + } + ARROW_RETURN_NOT_OK(internal::ValidateDecimalScale(scale)); + ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); + ARROW_RETURN_NOT_OK(out_.Append(&scale, sizeof(scale))); + return out_.Append(little_endian_unscaled_value); + } + + template + requires internal::LengthPrefixedVariantPrimitive + Status Append(std::string_view value) { + if (value.size() > std::numeric_limits::max()) { + return Status::Invalid("Variant ", + type == VariantPrimitiveType::kBinary ? "binary" : "string", + " value is too large"); + } + if constexpr (type == VariantPrimitiveType::kString) { + ARROW_RETURN_NOT_OK(internal::ValidateUtf8(value, "primitive string value")); + } + ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); + ARROW_RETURN_NOT_OK(AppendLittleEndian(out_, static_cast(value.size()), 4)); + return out_.Append(value); + } + + template + requires internal::UuidVariantPrimitive + Status Append(std::string_view big_endian_bytes) { + if (big_endian_bytes.size() != 16) { + return Status::Invalid("Variant UUID values must be 16 bytes"); + } + ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); + return out_.Append(big_endian_bytes); + } + + Status AppendShortString(std::string_view value) { + if (value.size() >= 64) { + return Status::Invalid("Variant short string value must be shorter than 64 bytes"); + } + ARROW_RETURN_NOT_OK(internal::ValidateUtf8(value, "short string value")); + const auto header = static_cast( + (value.size() << 2) | static_cast(VariantBasicType::kShortString)); + ARROW_RETURN_NOT_OK(out_.Append(&header, sizeof(header))); + return out_.Append(value); + } + + private: + template + Status AppendPrimitiveHeader() { + const auto header = + static_cast((static_cast(type) << 2) | + static_cast(VariantBasicType::kPrimitive)); + return out_.Append(&header, sizeof(header)); + } + + BufferBuilder& out_; +}; + +enum class VariantContainerKind { Object, List }; + +struct VariantFieldDescriptor { + uint32_t field_id = 0; + uint32_t offset = 0; +}; + +struct VariantBuildFrame { + VariantContainerKind kind; + int64_t value_start = 0; + size_t metadata_size = 0; + size_t parent_frame = 0; + size_t parent_entry_count = 0; + bool has_parent = false; + std::vector fields{}; + std::vector offsets{}; + std::unordered_set object_field_ids{}; + bool finished = false; +}; + +struct VariantBuildState { + explicit VariantBuildState(MemoryPool* pool) + : pool(pool), value(pool), metadata(pool) {} + + MemoryPool* pool; + BufferBuilder value; + VariantMetadataBuilder metadata; + std::vector frames; + bool root_has_value = false; +}; + +Status CheckRootWritable(const VariantBuildState& state) { + if (!state.frames.empty()) { + return Status::Invalid("VariantBuilder has an active container"); + } + if (state.root_has_value) { + return Status::Invalid("VariantBuilder already has a root value"); + } + return Status::OK(); +} + +Status CheckTopFrame(const VariantBuildState& state, size_t frame_index, + VariantContainerKind kind) { + if (state.frames.empty() || frame_index + 1 != state.frames.size()) { + return Status::Invalid("Variant nested builder is not the active container"); + } + const auto& frame = state.frames.back(); + if (frame.finished || frame.kind != kind) { + return Status::Invalid("Variant nested builder has invalid state"); + } + return Status::OK(); +} + +void TruncateFrameEntries(VariantBuildFrame& frame, size_t entry_count) { + if (frame.kind == VariantContainerKind::Object) { + frame.fields.resize(entry_count); + frame.object_field_ids.clear(); + for (const auto& field : frame.fields) { + frame.object_field_ids.insert(field.field_id); + } + } else { + frame.offsets.resize(entry_count); + } +} + +template +void RollbackIfActive(Impl* impl) { + if (impl != nullptr && impl->active) { + if (impl->state == nullptr || impl->frame_index >= impl->state->frames.size()) { + return; + } + const auto frame = impl->state->frames[impl->frame_index]; + impl->state->value.Rewind(frame.value_start); + impl->state->metadata.Truncate(frame.metadata_size); + if (frame.has_parent && frame.parent_frame < impl->state->frames.size()) { + TruncateFrameEntries(impl->state->frames[frame.parent_frame], + frame.parent_entry_count); + } + impl->state->frames.resize(impl->frame_index); + impl->active = false; + } +} + +Result BuildObjectHeader(const VariantBuildState& state, + const VariantBuildFrame& frame, + uint32_t values_size) { + if (frame.fields.size() > std::numeric_limits::max()) { + return Status::Invalid("Variant object has too many fields"); + } + + std::vector fields = frame.fields; + std::ranges::sort(fields, [&](const VariantFieldDescriptor& left, + const VariantFieldDescriptor& right) { + return state.metadata.FieldName(left.field_id) < + state.metadata.FieldName(right.field_id); + }); + + uint32_t max_field_id = 0; + for (const auto& field : fields) { + max_field_id = std::max(max_field_id, field.field_id); + } + const uint8_t id_size = WidthForValue(max_field_id); + const uint8_t offset_size = WidthForValue(values_size); + const bool is_large = fields.size() > std::numeric_limits::max(); + const auto header = static_cast(((is_large ? 1 : 0) << 4) | + ((id_size - 1) << 2) | (offset_size - 1)); + + std::string out; + out.push_back( + static_cast((header << 2) | static_cast(VariantBasicType::kObject))); + AppendLittleEndianToString(out, static_cast(fields.size()), is_large ? 4 : 1); + for (const auto& field : fields) { + AppendLittleEndianToString(out, field.field_id, id_size); + } + for (const auto& field : fields) { + AppendLittleEndianToString(out, field.offset, offset_size); + } + AppendLittleEndianToString(out, values_size, offset_size); + return out; +} + +Result BuildListHeader(const VariantBuildFrame& frame, + uint32_t values_size) { + if (frame.offsets.size() > std::numeric_limits::max()) { + return Status::Invalid("Variant array has too many elements"); + } + + const uint8_t offset_size = WidthForValue(values_size); + const bool is_large = frame.offsets.size() > std::numeric_limits::max(); + const auto header = static_cast(((is_large ? 1 : 0) << 2) | (offset_size - 1)); + + std::string out; + out.push_back( + static_cast((header << 2) | static_cast(VariantBasicType::kArray))); + AppendLittleEndianToString(out, static_cast(frame.offsets.size()), + is_large ? 4 : 1); + for (const auto offset : frame.offsets) { + AppendLittleEndianToString(out, offset, offset_size); + } + AppendLittleEndianToString(out, values_size, offset_size); + return out; +} + +using RootFinishCallback = std::function; + +Status FinishFrame(const std::shared_ptr& state, size_t frame_index, + VariantContainerKind kind, const RootFinishCallback& callback) { + ARROW_RETURN_NOT_OK(CheckTopFrame(*state, frame_index, kind)); + + auto& frame = state->frames.back(); + const auto values_size = state->value.length() - frame.value_start; + DCHECK_GE(values_size, 0); + if (values_size > std::numeric_limits::max()) { + return Status::Invalid("Variant container values are too large"); + } + + ARROW_ASSIGN_OR_RAISE( + auto header, + kind == VariantContainerKind::Object + ? BuildObjectHeader(*state, frame, static_cast(values_size)) + : BuildListHeader(frame, static_cast(values_size))); + ARROW_RETURN_NOT_OK(InsertBytes(state->value, frame.value_start, header)); + + const bool is_root = !frame.has_parent; + frame.finished = true; + state->frames.pop_back(); + if (!is_root) { + return Status::OK(); + } + + state->root_has_value = true; + if (!callback) { + return Status::OK(); + } + + ARROW_ASSIGN_OR_RAISE(auto metadata, state->metadata.Finish()); + ARROW_ASSIGN_OR_RAISE(auto value, state->value.Finish()); + return callback({.metadata = std::move(metadata), .value = std::move(value)}); +} + +template +Status AppendRootPrimitiveWith(const std::shared_ptr& state, + Write&& write) { + ARROW_RETURN_NOT_OK(CheckRootWritable(*state)); + const auto value_size = state->value.length(); + VariantValueWriter writer(state->value); + const Status status = std::invoke(std::forward(write), writer); + if (!status.ok()) { + state->value.Rewind(value_size); + return status; + } + state->root_has_value = true; + return Status::OK(); +} + +template +Status AppendRootPrimitive(const std::shared_ptr& state, + Args&&... args) { + return AppendRootPrimitiveWith(state, [&](VariantValueWriter& writer) { + return writer.template Append(std::forward(args)...); + }); +} + +template +Status AppendObjectPrimitiveWith(const std::shared_ptr& state, + size_t frame_index, std::string_view field_name, + Write&& write) { + ARROW_RETURN_NOT_OK(CheckTopFrame(*state, frame_index, VariantContainerKind::Object)); + auto& frame = state->frames.back(); + const auto metadata_size = state->metadata.size(); + const auto value_size = state->value.length(); + const auto field_count = frame.fields.size(); + + ARROW_ASSIGN_OR_RAISE(auto field_id, state->metadata.Upsert(field_name)); + if (!frame.object_field_ids.insert(field_id).second) { + state->metadata.Truncate(metadata_size); + return Status::Invalid("Duplicate Variant object field: ", field_name); + } + const auto offset = state->value.length() - frame.value_start; + DCHECK_GE(offset, 0); + if (offset > std::numeric_limits::max()) { + state->metadata.Truncate(metadata_size); + TruncateFrameEntries(frame, field_count); + return Status::Invalid("Variant object values are too large"); + } + frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, + .offset = static_cast(offset)}); + + VariantValueWriter writer(state->value); + const Status status = std::invoke(std::forward(write), writer); + if (!status.ok()) { + state->value.Rewind(value_size); + state->metadata.Truncate(metadata_size); + TruncateFrameEntries(frame, field_count); + return status; + } + return Status::OK(); +} + +template +Status AppendObjectPrimitive(const std::shared_ptr& state, + size_t frame_index, std::string_view field_name, + Args&&... args) { + return AppendObjectPrimitiveWith( + state, frame_index, field_name, [&](VariantValueWriter& writer) { + return writer.template Append(std::forward(args)...); + }); +} + +template +Status AppendListPrimitiveWith(const std::shared_ptr& state, + size_t frame_index, Write&& write) { + ARROW_RETURN_NOT_OK(CheckTopFrame(*state, frame_index, VariantContainerKind::List)); + auto& frame = state->frames.back(); + const auto value_size = state->value.length(); + const auto element_count = frame.offsets.size(); + const auto offset = state->value.length() - frame.value_start; + DCHECK_GE(offset, 0); + if (offset > std::numeric_limits::max()) { + return Status::Invalid("Variant array values are too large"); + } + frame.offsets.push_back(static_cast(offset)); + + VariantValueWriter writer(state->value); + const Status status = std::invoke(std::forward(write), writer); + if (!status.ok()) { + state->value.Rewind(value_size); + TruncateFrameEntries(frame, element_count); + return status; + } + return Status::OK(); +} + +template +Status AppendListPrimitive(const std::shared_ptr& state, + size_t frame_index, Args&&... args) { + return AppendListPrimitiveWith(state, frame_index, [&](VariantValueWriter& writer) { + return writer.template Append(std::forward(args)...); + }); +} + +} // namespace + +namespace internal { + +struct NestedVariantBuilderImpl { + NestedVariantBuilderImpl(std::shared_ptr state, size_t frame_index, + RootFinishCallback callback) + : state(std::move(state)), + frame_index(frame_index), + callback(std::move(callback)) {} + + std::shared_ptr state; + size_t frame_index = 0; + RootFinishCallback callback; + bool active = true; +}; + +} // namespace internal + +struct VariantBuilder::Impl { + explicit Impl(MemoryPool* pool) + : pool(pool), state(std::make_shared(pool)) {} + + MemoryPool* pool; + std::shared_ptr state; +}; + +VariantBuilder::VariantBuilder(MemoryPool* pool) : impl_(std::make_unique(pool)) {} +VariantBuilder::~VariantBuilder() = default; +VariantBuilder::VariantBuilder(VariantBuilder&&) noexcept = default; +VariantBuilder& VariantBuilder::operator=(VariantBuilder&&) noexcept = default; + +Status VariantBuilder::ReserveFieldNames(int64_t capacity) { + return impl_->state->metadata.Reserve(capacity); +} + +Result VariantBuilder::AddFieldName(std::string_view name) { + return impl_->state->metadata.Upsert(name); +} + +Status VariantBuilder::AppendVariantNull() { + return AppendRootPrimitive(impl_->state); +} + +Status VariantBuilder::AppendBoolean(bool value) { + if (value) { + return AppendRootPrimitive(impl_->state); + } + return AppendRootPrimitive(impl_->state); +} + +#define VARIANT_ROOT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + Status VariantBuilder::Append##NAME(C_TYPE value) { \ + return AppendRootPrimitive(impl_->state, value); \ + } + +VARIANT_ROOT_APPEND_ONE_ARG(Int8, Int8, int8_t) +VARIANT_ROOT_APPEND_ONE_ARG(Int16, Int16, int16_t) +VARIANT_ROOT_APPEND_ONE_ARG(Int32, Int32, int32_t) +VARIANT_ROOT_APPEND_ONE_ARG(Int64, Int64, int64_t) +VARIANT_ROOT_APPEND_ONE_ARG(Float, Float, float) +VARIANT_ROOT_APPEND_ONE_ARG(Double, Double, double) +VARIANT_ROOT_APPEND_ONE_ARG(Binary, Binary, std::string_view) +VARIANT_ROOT_APPEND_ONE_ARG(String, String, std::string_view) +VARIANT_ROOT_APPEND_ONE_ARG(Date, Date, int32_t) +VARIANT_ROOT_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) +VARIANT_ROOT_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) + +#undef VARIANT_ROOT_APPEND_ONE_ARG + +Status VariantBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { + return AppendRootPrimitive(impl_->state, + unscaled_value, scale); +} + +Status VariantBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { + return AppendRootPrimitive(impl_->state, + unscaled_value, scale); +} + +Status VariantBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, + uint8_t scale) { + return AppendRootPrimitive( + impl_->state, little_endian_unscaled_value, scale); +} + +Status VariantBuilder::AppendShortString(std::string_view value) { + return AppendRootPrimitiveWith(impl_->state, [&](VariantValueWriter& writer) { + return writer.AppendShortString(value); + }); +} + +Status VariantBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return AppendRootPrimitive(impl_->state, + micros); + } + return AppendRootPrimitive(impl_->state, + micros); +} + +Status VariantBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return AppendRootPrimitive(impl_->state, + nanos); + } + return AppendRootPrimitive(impl_->state, + nanos); +} + +Result VariantBuilder::StartObject() { + ARROW_RETURN_NOT_OK(CheckRootWritable(*impl_->state)); + impl_->state->frames.push_back( + VariantBuildFrame{.kind = VariantContainerKind::Object, + .value_start = impl_->state->value.length(), + .metadata_size = impl_->state->metadata.size()}); + return VariantObjectBuilder(std::make_unique( + impl_->state, impl_->state->frames.size() - 1, RootFinishCallback{})); +} + +Result VariantBuilder::StartList() { + ARROW_RETURN_NOT_OK(CheckRootWritable(*impl_->state)); + impl_->state->frames.push_back( + VariantBuildFrame{.kind = VariantContainerKind::List, + .value_start = impl_->state->value.length(), + .metadata_size = impl_->state->metadata.size()}); + return VariantListBuilder(std::make_unique( + impl_->state, impl_->state->frames.size() - 1, RootFinishCallback{})); +} + +Result VariantBuilder::Finish() { + if (!impl_->state->frames.empty()) { + return Status::Invalid("Cannot finish VariantBuilder with active containers"); + } + if (!impl_->state->root_has_value) { + return Status::Invalid("Cannot finish empty VariantBuilder"); + } + + ARROW_ASSIGN_OR_RAISE(auto metadata, impl_->state->metadata.Finish()); + ARROW_ASSIGN_OR_RAISE(auto value, impl_->state->value.Finish()); + EncodedVariantValue out{.metadata = std::move(metadata), .value = std::move(value)}; + Reset(); + return out; +} + +void VariantBuilder::Reset() { + impl_->state = std::make_shared(impl_->pool); +} + +VariantObjectBuilder::VariantObjectBuilder( + std::unique_ptr impl) + : impl_(std::move(impl)) {} +VariantObjectBuilder::~VariantObjectBuilder() { RollbackIfActive(impl_.get()); } +VariantObjectBuilder::VariantObjectBuilder(VariantObjectBuilder&&) noexcept = default; +VariantObjectBuilder& VariantObjectBuilder::operator=(VariantObjectBuilder&&) noexcept = + default; + +Status VariantObjectBuilder::AppendVariantNull(std::string_view field_name) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name); +} + +Status VariantObjectBuilder::AppendBoolean(std::string_view field_name, bool value) { + if (value) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name); + } + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name); +} + +#define VARIANT_OBJECT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + Status VariantObjectBuilder::Append##NAME(std::string_view field_name, C_TYPE value) { \ + return AppendObjectPrimitive( \ + impl_->state, impl_->frame_index, field_name, value); \ + } + +VARIANT_OBJECT_APPEND_ONE_ARG(Int8, Int8, int8_t) +VARIANT_OBJECT_APPEND_ONE_ARG(Int16, Int16, int16_t) +VARIANT_OBJECT_APPEND_ONE_ARG(Int32, Int32, int32_t) +VARIANT_OBJECT_APPEND_ONE_ARG(Int64, Int64, int64_t) +VARIANT_OBJECT_APPEND_ONE_ARG(Float, Float, float) +VARIANT_OBJECT_APPEND_ONE_ARG(Double, Double, double) +VARIANT_OBJECT_APPEND_ONE_ARG(Binary, Binary, std::string_view) +VARIANT_OBJECT_APPEND_ONE_ARG(String, String, std::string_view) +VARIANT_OBJECT_APPEND_ONE_ARG(Date, Date, int32_t) +VARIANT_OBJECT_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) +VARIANT_OBJECT_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) + +#undef VARIANT_OBJECT_APPEND_ONE_ARG + +Status VariantObjectBuilder::AppendDecimal4(std::string_view field_name, + int32_t unscaled_value, uint8_t scale) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, unscaled_value, scale); +} + +Status VariantObjectBuilder::AppendDecimal8(std::string_view field_name, + int64_t unscaled_value, uint8_t scale) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, unscaled_value, scale); +} + +Status VariantObjectBuilder::AppendDecimal16( + std::string_view field_name, std::string_view little_endian_unscaled_value, + uint8_t scale) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, little_endian_unscaled_value, scale); +} + +Status VariantObjectBuilder::AppendShortString(std::string_view field_name, + std::string_view value) { + return AppendObjectPrimitiveWith( + impl_->state, impl_->frame_index, field_name, + [&](VariantValueWriter& writer) { return writer.AppendShortString(value); }); +} + +Status VariantObjectBuilder::AppendTimestampMicros(std::string_view field_name, + int64_t micros, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, micros); + } + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, micros); +} + +Status VariantObjectBuilder::AppendTimestampNanos(std::string_view field_name, + int64_t nanos, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, nanos); + } + return AppendObjectPrimitive( + impl_->state, impl_->frame_index, field_name, nanos); +} + +Result VariantObjectBuilder::StartObject( + std::string_view field_name) { + ARROW_RETURN_NOT_OK( + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object)); + auto& frame = impl_->state->frames.back(); + const auto metadata_size = impl_->state->metadata.size(); + const auto field_count = frame.fields.size(); + + ARROW_ASSIGN_OR_RAISE(auto field_id, impl_->state->metadata.Upsert(field_name)); + if (!frame.object_field_ids.insert(field_id).second) { + impl_->state->metadata.Truncate(metadata_size); + return Status::Invalid("Duplicate Variant object field: ", field_name); + } + const auto offset = impl_->state->value.length() - frame.value_start; + DCHECK_GE(offset, 0); + if (offset > std::numeric_limits::max()) { + impl_->state->metadata.Truncate(metadata_size); + TruncateFrameEntries(frame, field_count); + return Status::Invalid("Variant object values are too large"); + } + + frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, + .offset = static_cast(offset)}); + const auto value_start = impl_->state->value.length(); + impl_->state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::Object, + .value_start = value_start, + .metadata_size = metadata_size, + .parent_frame = impl_->frame_index, + .parent_entry_count = field_count, + .has_parent = true}); + return VariantObjectBuilder(std::make_unique( + impl_->state, impl_->state->frames.size() - 1, impl_->callback)); +} + +Result VariantObjectBuilder::StartList(std::string_view field_name) { + ARROW_RETURN_NOT_OK( + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object)); + auto& frame = impl_->state->frames.back(); + const auto metadata_size = impl_->state->metadata.size(); + const auto field_count = frame.fields.size(); + + ARROW_ASSIGN_OR_RAISE(auto field_id, impl_->state->metadata.Upsert(field_name)); + if (!frame.object_field_ids.insert(field_id).second) { + impl_->state->metadata.Truncate(metadata_size); + return Status::Invalid("Duplicate Variant object field: ", field_name); + } + const auto offset = impl_->state->value.length() - frame.value_start; + DCHECK_GE(offset, 0); + if (offset > std::numeric_limits::max()) { + impl_->state->metadata.Truncate(metadata_size); + TruncateFrameEntries(frame, field_count); + return Status::Invalid("Variant object values are too large"); + } + + frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, + .offset = static_cast(offset)}); + const auto value_start = impl_->state->value.length(); + impl_->state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::List, + .value_start = value_start, + .metadata_size = metadata_size, + .parent_frame = impl_->frame_index, + .parent_entry_count = field_count, + .has_parent = true}); + return VariantListBuilder(std::make_unique( + impl_->state, impl_->state->frames.size() - 1, impl_->callback)); +} + +Status VariantObjectBuilder::Finish() { + ARROW_RETURN_NOT_OK(FinishFrame(impl_->state, impl_->frame_index, + VariantContainerKind::Object, impl_->callback)); + impl_->active = false; + return Status::OK(); +} + +VariantListBuilder::VariantListBuilder( + std::unique_ptr impl) + : impl_(std::move(impl)) {} +VariantListBuilder::~VariantListBuilder() { RollbackIfActive(impl_.get()); } +VariantListBuilder::VariantListBuilder(VariantListBuilder&&) noexcept = default; +VariantListBuilder& VariantListBuilder::operator=(VariantListBuilder&&) noexcept = + default; + +Status VariantListBuilder::AppendVariantNull() { + return AppendListPrimitive(impl_->state, + impl_->frame_index); +} + +Status VariantListBuilder::AppendBoolean(bool value) { + if (value) { + return AppendListPrimitive(impl_->state, + impl_->frame_index); + } + return AppendListPrimitive(impl_->state, + impl_->frame_index); +} + +#define VARIANT_LIST_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + Status VariantListBuilder::Append##NAME(C_TYPE value) { \ + return AppendListPrimitive( \ + impl_->state, impl_->frame_index, value); \ + } + +VARIANT_LIST_APPEND_ONE_ARG(Int8, Int8, int8_t) +VARIANT_LIST_APPEND_ONE_ARG(Int16, Int16, int16_t) +VARIANT_LIST_APPEND_ONE_ARG(Int32, Int32, int32_t) +VARIANT_LIST_APPEND_ONE_ARG(Int64, Int64, int64_t) +VARIANT_LIST_APPEND_ONE_ARG(Float, Float, float) +VARIANT_LIST_APPEND_ONE_ARG(Double, Double, double) +VARIANT_LIST_APPEND_ONE_ARG(Binary, Binary, std::string_view) +VARIANT_LIST_APPEND_ONE_ARG(String, String, std::string_view) +VARIANT_LIST_APPEND_ONE_ARG(Date, Date, int32_t) +VARIANT_LIST_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) +VARIANT_LIST_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) + +#undef VARIANT_LIST_APPEND_ONE_ARG + +Status VariantListBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { + return AppendListPrimitive( + impl_->state, impl_->frame_index, unscaled_value, scale); +} + +Status VariantListBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { + return AppendListPrimitive( + impl_->state, impl_->frame_index, unscaled_value, scale); +} + +Status VariantListBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, + uint8_t scale) { + return AppendListPrimitive( + impl_->state, impl_->frame_index, little_endian_unscaled_value, scale); +} + +Status VariantListBuilder::AppendShortString(std::string_view value) { + return AppendListPrimitiveWith( + impl_->state, impl_->frame_index, + [&](VariantValueWriter& writer) { return writer.AppendShortString(value); }); +} + +Status VariantListBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return AppendListPrimitive( + impl_->state, impl_->frame_index, micros); + } + return AppendListPrimitive( + impl_->state, impl_->frame_index, micros); +} + +Status VariantListBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return AppendListPrimitive( + impl_->state, impl_->frame_index, nanos); + } + return AppendListPrimitive( + impl_->state, impl_->frame_index, nanos); +} + +Result VariantListBuilder::StartObject() { + ARROW_RETURN_NOT_OK( + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List)); + auto& frame = impl_->state->frames.back(); + const auto element_count = frame.offsets.size(); + const auto offset = impl_->state->value.length() - frame.value_start; + DCHECK_GE(offset, 0); + if (offset > std::numeric_limits::max()) { + return Status::Invalid("Variant array values are too large"); + } + + frame.offsets.push_back(static_cast(offset)); + const auto value_start = impl_->state->value.length(); + impl_->state->frames.push_back( + VariantBuildFrame{.kind = VariantContainerKind::Object, + .value_start = value_start, + .metadata_size = impl_->state->metadata.size(), + .parent_frame = impl_->frame_index, + .parent_entry_count = element_count, + .has_parent = true}); + return VariantObjectBuilder(std::make_unique( + impl_->state, impl_->state->frames.size() - 1, impl_->callback)); +} + +Result VariantListBuilder::StartList() { + ARROW_RETURN_NOT_OK( + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List)); + auto& frame = impl_->state->frames.back(); + const auto element_count = frame.offsets.size(); + const auto offset = impl_->state->value.length() - frame.value_start; + DCHECK_GE(offset, 0); + if (offset > std::numeric_limits::max()) { + return Status::Invalid("Variant array values are too large"); + } + + frame.offsets.push_back(static_cast(offset)); + const auto value_start = impl_->state->value.length(); + impl_->state->frames.push_back( + VariantBuildFrame{.kind = VariantContainerKind::List, + .value_start = value_start, + .metadata_size = impl_->state->metadata.size(), + .parent_frame = impl_->frame_index, + .parent_entry_count = element_count, + .has_parent = true}); + return VariantListBuilder(std::make_unique( + impl_->state, impl_->state->frames.size() - 1, impl_->callback)); +} + +Status VariantListBuilder::Finish() { + ARROW_RETURN_NOT_OK(FinishFrame(impl_->state, impl_->frame_index, + VariantContainerKind::List, impl_->callback)); + impl_->active = false; + return Status::OK(); +} + +struct VariantArrayBuilder::Impl { + explicit Impl(MemoryPool* pool) + : pool(pool), metadata_builder(pool), value_builder(pool), validity_builder(pool) {} + + template + Status AppendValue(Write&& write) { + auto state = std::make_shared(pool); + ARROW_RETURN_NOT_OK(AppendRootPrimitiveWith(state, std::forward(write))); + ARROW_ASSIGN_OR_RAISE(auto metadata, state->metadata.Finish()); + ARROW_ASSIGN_OR_RAISE(auto value, state->value.Finish()); + return AppendEncoded({.metadata = std::move(metadata), .value = std::move(value)}); + } + + template + Status AppendPrimitive(Args&&... args) { + return AppendValue([&](VariantValueWriter& writer) { + return writer.template Append(std::forward(args)...); + }); + } + + Status AppendShortString(std::string_view value) { + return AppendValue( + [&](VariantValueWriter& writer) { return writer.AppendShortString(value); }); + } + + Status AppendEncoded(const EncodedVariantValue& value) { + if (value.metadata == nullptr || value.value == nullptr) { + return Status::Invalid( + "Encoded Variant metadata and value buffers must be non-null"); + } + ARROW_ASSIGN_OR_RAISE(auto metadata, + VariantMetadataView::Make(std::string_view{*value.metadata})); + ARROW_RETURN_NOT_OK( + VariantValueView::Validate(std::string_view{*value.value}, metadata)); + ARROW_RETURN_NOT_OK(metadata_builder.Append(std::string_view{*value.metadata})); + ARROW_RETURN_NOT_OK(value_builder.Append(std::string_view{*value.value})); + return validity_builder.Append(true); + } + + MemoryPool* pool; + BinaryBuilder metadata_builder; + BinaryBuilder value_builder; + BooleanBuilder validity_builder; +}; + +VariantArrayBuilder::VariantArrayBuilder(MemoryPool* pool) + : impl_(std::make_unique(pool)) {} +VariantArrayBuilder::~VariantArrayBuilder() = default; +VariantArrayBuilder::VariantArrayBuilder(VariantArrayBuilder&&) noexcept = default; +VariantArrayBuilder& VariantArrayBuilder::operator=(VariantArrayBuilder&&) noexcept = + default; + +Status VariantArrayBuilder::AppendNull() { + ARROW_RETURN_NOT_OK(impl_->metadata_builder.Append("")); + ARROW_RETURN_NOT_OK(impl_->value_builder.Append("")); + return impl_->validity_builder.Append(false); +} + +Status VariantArrayBuilder::AppendVariantNull() { + return impl_->AppendPrimitive(); +} + +Status VariantArrayBuilder::AppendBoolean(bool value) { + if (value) { + return impl_->AppendPrimitive(); + } + return impl_->AppendPrimitive(); +} + +#define VARIANT_ARRAY_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + Status VariantArrayBuilder::Append##NAME(C_TYPE value) { \ + return impl_->AppendPrimitive(value); \ + } + +VARIANT_ARRAY_APPEND_ONE_ARG(Int8, Int8, int8_t) +VARIANT_ARRAY_APPEND_ONE_ARG(Int16, Int16, int16_t) +VARIANT_ARRAY_APPEND_ONE_ARG(Int32, Int32, int32_t) +VARIANT_ARRAY_APPEND_ONE_ARG(Int64, Int64, int64_t) +VARIANT_ARRAY_APPEND_ONE_ARG(Float, Float, float) +VARIANT_ARRAY_APPEND_ONE_ARG(Double, Double, double) +VARIANT_ARRAY_APPEND_ONE_ARG(Binary, Binary, std::string_view) +VARIANT_ARRAY_APPEND_ONE_ARG(String, String, std::string_view) +VARIANT_ARRAY_APPEND_ONE_ARG(Date, Date, int32_t) +VARIANT_ARRAY_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) +VARIANT_ARRAY_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) + +#undef VARIANT_ARRAY_APPEND_ONE_ARG + +Status VariantArrayBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { + return impl_->AppendPrimitive(unscaled_value, scale); +} + +Status VariantArrayBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { + return impl_->AppendPrimitive(unscaled_value, scale); +} + +Status VariantArrayBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, + uint8_t scale) { + return impl_->AppendPrimitive( + little_endian_unscaled_value, scale); +} + +Status VariantArrayBuilder::AppendShortString(std::string_view value) { + return impl_->AppendShortString(value); +} + +Status VariantArrayBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return impl_->AppendPrimitive(micros); + } + return impl_->AppendPrimitive(micros); +} + +Status VariantArrayBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { + if (adjusted_to_utc) { + return impl_->AppendPrimitive(nanos); + } + return impl_->AppendPrimitive(nanos); +} + +Status VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { + return impl_->AppendEncoded(value); +} + +Result VariantArrayBuilder::StartObject() { + auto state = std::make_shared(impl_->pool); + state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::Object, + .value_start = state->value.length(), + .metadata_size = state->metadata.size()}); + auto callback = [this](EncodedVariantValue encoded) { + return impl_->AppendEncoded(encoded); + }; + return VariantObjectBuilder(std::make_unique( + state, state->frames.size() - 1, std::move(callback))); +} + +Result VariantArrayBuilder::StartList() { + auto state = std::make_shared(impl_->pool); + state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::List, + .value_start = state->value.length(), + .metadata_size = state->metadata.size()}); + auto callback = [this](EncodedVariantValue encoded) { + return impl_->AppendEncoded(encoded); + }; + return VariantListBuilder(std::make_unique( + state, state->frames.size() - 1, std::move(callback))); +} + +Result> VariantArrayBuilder::Finish() { + std::shared_ptr metadata; + std::shared_ptr value; + std::shared_ptr validity; + ARROW_RETURN_NOT_OK(impl_->metadata_builder.Finish(&metadata)); + ARROW_RETURN_NOT_OK(impl_->value_builder.Finish(&value)); + ARROW_RETURN_NOT_OK(impl_->validity_builder.Finish(&validity)); + + auto null_bitmap = validity->data()->buffers[1]; + const int64_t null_count = validity->false_count(); + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + ARROW_ASSIGN_OR_RAISE(auto storage, + StructArray::Make({metadata, value}, storage_type->fields(), + null_bitmap, null_count)); + return MakeVariantArrayFromStorage(storage); +} + +void VariantArrayBuilder::Reset() { impl_ = std::make_unique(impl_->pool); } + +struct VariantValueArrayBuilder::Impl { + explicit Impl(MemoryPool* pool) : value_builder(pool) {} + + BinaryBuilder value_builder; +}; + +VariantValueArrayBuilder::VariantValueArrayBuilder(MemoryPool* pool) + : impl_(std::make_unique(pool)) {} +VariantValueArrayBuilder::VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept = + default; +VariantValueArrayBuilder& VariantValueArrayBuilder::operator=( + VariantValueArrayBuilder&&) noexcept = default; +VariantValueArrayBuilder::~VariantValueArrayBuilder() = default; + +Status VariantValueArrayBuilder::AppendNull() { + return impl_->value_builder.AppendNull(); +} + +Status VariantValueArrayBuilder::AppendEncodedValue(std::string_view metadata, + std::string_view value) { + ARROW_ASSIGN_OR_RAISE(auto metadata_view, VariantMetadataView::Make(metadata)); + ARROW_RETURN_NOT_OK(VariantValueView::Validate(value, metadata_view)); + return impl_->value_builder.Append(value); +} + +Result> VariantValueArrayBuilder::Finish() { + std::shared_ptr out; + ARROW_RETURN_NOT_OK(impl_->value_builder.Finish(&out)); + return out; +} + +Result> MakeVariantArrayFromStorage( + std::shared_ptr storage) { + if (storage == nullptr) { + return Status::Invalid("Variant storage array must be non-null"); + } + ARROW_ASSIGN_OR_RAISE(auto type, VariantExtensionType::Make(storage->type())); + auto array = ExtensionType::WrapArray(type, std::move(storage)); + return std::static_pointer_cast(array); +} + +Result> MakeVariantArrayFromChildren( + std::shared_ptr storage_type, std::vector> children, + std::shared_ptr null_bitmap) { + if (storage_type->id() != Type::STRUCT) { + return Status::Invalid("Variant storage type must be struct, got ", + storage_type->ToString()); + } + + const auto& struct_type = + ::arrow::internal::checked_cast(*storage_type); + if (children.size() != static_cast(struct_type.num_fields())) { + return Status::Invalid("Variant storage expected ", struct_type.num_fields(), + " children, got ", children.size()); + } + + const int64_t length = children.empty() ? 0 : children[0]->length(); + for (int i = 0; i < struct_type.num_fields(); ++i) { + if (children[i] == nullptr) { + return Status::Invalid("Variant storage child ", i, " is null"); + } + if (!children[i]->type()->Equals(struct_type.field(i)->type())) { + return Status::Invalid("Variant storage child ", i, " has type ", + children[i]->type()->ToString(), ", expected ", + struct_type.field(i)->type()->ToString()); + } + if (children[i]->length() != length) { + return Status::Invalid("Variant storage child lengths must match"); + } + } + + ARROW_ASSIGN_OR_RAISE(auto storage, + StructArray::Make(std::move(children), struct_type.fields(), + std::move(null_bitmap))); + return MakeVariantArrayFromStorage(storage); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/builder.h b/cpp/src/parquet/variant/builder.h new file mode 100644 index 000000000000..cc0507bfa9f8 --- /dev/null +++ b/cpp/src/parquet/variant/builder.h @@ -0,0 +1,264 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include + +#include "arrow/buffer.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/memory_pool.h" +#include "arrow/result.h" +#include "arrow/status.h" +#include "arrow/type.h" +#include "parquet/platform.h" + +namespace parquet::variant { + +using ::arrow::Array; +using ::arrow::BinaryArray; +using ::arrow::Buffer; +using ::arrow::DataType; +using ::arrow::MemoryPool; +using ::arrow::Result; +using ::arrow::Status; +using ::arrow::StructArray; +using ::arrow::extension::VariantArray; + +struct PARQUET_EXPORT EncodedVariantValue { + std::shared_ptr metadata; + std::shared_ptr value; +}; + +class VariantObjectBuilder; +class VariantListBuilder; + +class PARQUET_EXPORT VariantBuilder { + public: + explicit VariantBuilder(MemoryPool* pool = ::arrow::default_memory_pool()); + ~VariantBuilder(); + VariantBuilder(const VariantBuilder&) = delete; + VariantBuilder& operator=(const VariantBuilder&) = delete; + VariantBuilder(VariantBuilder&&) noexcept; + VariantBuilder& operator=(VariantBuilder&&) noexcept; + + Status ReserveFieldNames(int64_t capacity); + Result AddFieldName(std::string_view name); + + Status AppendVariantNull(); + Status AppendBoolean(bool value); + Status AppendInt8(int8_t value); + Status AppendInt16(int16_t value); + Status AppendInt32(int32_t value); + Status AppendInt64(int64_t value); + Status AppendFloat(float value); + Status AppendDouble(double value); + Status AppendBinary(std::string_view value); + Status AppendString(std::string_view value); + Status AppendDate(int32_t days); + Status AppendTimeNTZMicros(int64_t micros); + Status AppendUuid(std::string_view big_endian_bytes); + Status AppendDecimal4(int32_t unscaled_value, uint8_t scale); + Status AppendDecimal8(int64_t unscaled_value, uint8_t scale); + Status AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + Status AppendShortString(std::string_view value); + Status AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + Status AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + + Result StartObject(); + Result StartList(); + + Result Finish(); + void Reset(); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +namespace internal { +struct NestedVariantBuilderImpl; +} + +class PARQUET_EXPORT VariantObjectBuilder { + public: + ~VariantObjectBuilder(); + VariantObjectBuilder(const VariantObjectBuilder&) = delete; + VariantObjectBuilder& operator=(const VariantObjectBuilder&) = delete; + VariantObjectBuilder(VariantObjectBuilder&&) noexcept; + VariantObjectBuilder& operator=(VariantObjectBuilder&&) noexcept; + + Status AppendVariantNull(std::string_view field_name); + Status AppendBoolean(std::string_view field_name, bool value); + Status AppendInt8(std::string_view field_name, int8_t value); + Status AppendInt16(std::string_view field_name, int16_t value); + Status AppendInt32(std::string_view field_name, int32_t value); + Status AppendInt64(std::string_view field_name, int64_t value); + Status AppendFloat(std::string_view field_name, float value); + Status AppendDouble(std::string_view field_name, double value); + Status AppendBinary(std::string_view field_name, std::string_view value); + Status AppendString(std::string_view field_name, std::string_view value); + Status AppendDate(std::string_view field_name, int32_t days); + Status AppendTimeNTZMicros(std::string_view field_name, int64_t micros); + Status AppendUuid(std::string_view field_name, std::string_view big_endian_bytes); + Status AppendDecimal4(std::string_view field_name, int32_t unscaled_value, + uint8_t scale); + Status AppendDecimal8(std::string_view field_name, int64_t unscaled_value, + uint8_t scale); + Status AppendDecimal16(std::string_view field_name, + std::string_view little_endian_unscaled_value, uint8_t scale); + Status AppendShortString(std::string_view field_name, std::string_view value); + Status AppendTimestampMicros(std::string_view field_name, int64_t micros, + bool adjusted_to_utc); + Status AppendTimestampNanos(std::string_view field_name, int64_t nanos, + bool adjusted_to_utc); + + Result StartObject(std::string_view field_name); + Result StartList(std::string_view field_name); + + /// Commit this nested object into its parent builder. Destroying an unfinished + /// nested builder rolls back the object contents written through this builder. + Status Finish(); + + private: + friend class VariantBuilder; + friend class VariantListBuilder; + friend class VariantArrayBuilder; + + explicit VariantObjectBuilder(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +class PARQUET_EXPORT VariantListBuilder { + public: + ~VariantListBuilder(); + VariantListBuilder(const VariantListBuilder&) = delete; + VariantListBuilder& operator=(const VariantListBuilder&) = delete; + VariantListBuilder(VariantListBuilder&&) noexcept; + VariantListBuilder& operator=(VariantListBuilder&&) noexcept; + + Status AppendVariantNull(); + Status AppendBoolean(bool value); + Status AppendInt8(int8_t value); + Status AppendInt16(int16_t value); + Status AppendInt32(int32_t value); + Status AppendInt64(int64_t value); + Status AppendFloat(float value); + Status AppendDouble(double value); + Status AppendBinary(std::string_view value); + Status AppendString(std::string_view value); + Status AppendDate(int32_t days); + Status AppendTimeNTZMicros(int64_t micros); + Status AppendUuid(std::string_view big_endian_bytes); + Status AppendDecimal4(int32_t unscaled_value, uint8_t scale); + Status AppendDecimal8(int64_t unscaled_value, uint8_t scale); + Status AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + Status AppendShortString(std::string_view value); + Status AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + Status AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + + Result StartObject(); + Result StartList(); + + /// Commit this nested list into its parent builder. Destroying an unfinished nested + /// builder rolls back the list contents written through this builder. + Status Finish(); + + private: + friend class VariantBuilder; + friend class VariantObjectBuilder; + friend class VariantArrayBuilder; + + explicit VariantListBuilder(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +class PARQUET_EXPORT VariantArrayBuilder { + public: + explicit VariantArrayBuilder(MemoryPool* pool = ::arrow::default_memory_pool()); + ~VariantArrayBuilder(); + VariantArrayBuilder(const VariantArrayBuilder&) = delete; + VariantArrayBuilder& operator=(const VariantArrayBuilder&) = delete; + VariantArrayBuilder(VariantArrayBuilder&&) noexcept; + VariantArrayBuilder& operator=(VariantArrayBuilder&&) noexcept; + + Status AppendNull(); + Status AppendVariantNull(); + Status AppendBoolean(bool value); + Status AppendInt8(int8_t value); + Status AppendInt16(int16_t value); + Status AppendInt32(int32_t value); + Status AppendInt64(int64_t value); + Status AppendFloat(float value); + Status AppendDouble(double value); + Status AppendBinary(std::string_view value); + Status AppendString(std::string_view value); + Status AppendDate(int32_t days); + Status AppendTimeNTZMicros(int64_t micros); + Status AppendUuid(std::string_view big_endian_bytes); + Status AppendDecimal4(int32_t unscaled_value, uint8_t scale); + Status AppendDecimal8(int64_t unscaled_value, uint8_t scale); + Status AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + Status AppendShortString(std::string_view value); + Status AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + Status AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + Status AppendEncoded(const EncodedVariantValue& value); + + Result StartObject(); + Result StartList(); + + Result> Finish(); + void Reset(); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +class PARQUET_EXPORT VariantValueArrayBuilder { + public: + explicit VariantValueArrayBuilder(MemoryPool* pool = ::arrow::default_memory_pool()); + ~VariantValueArrayBuilder(); + VariantValueArrayBuilder(const VariantValueArrayBuilder&) = delete; + VariantValueArrayBuilder& operator=(const VariantValueArrayBuilder&) = delete; + VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept; + VariantValueArrayBuilder& operator=(VariantValueArrayBuilder&&) noexcept; + + Status AppendNull(); + Status AppendEncodedValue(std::string_view metadata, std::string_view value); + Result> Finish(); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +PARQUET_EXPORT +Result> MakeVariantArrayFromStorage( + std::shared_ptr storage); + +PARQUET_EXPORT +Result> MakeVariantArrayFromChildren( + std::shared_ptr storage_type, std::vector> children, + std::shared_ptr null_bitmap = nullptr); + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/builder_test.cc b/cpp/src/parquet/variant/builder_test.cc new file mode 100644 index 000000000000..e82addb27a81 --- /dev/null +++ b/cpp/src/parquet/variant/builder_test.cc @@ -0,0 +1,310 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/builder.h" +#include "parquet/variant/encoding.h" +#include "parquet/variant/test_util_internal.h" + +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/testing/gtest_util.h" +#include "arrow/util/checked_cast.h" + +namespace parquet::variant { + +namespace { + +using ::arrow::ArrayFromJSON; +using ::arrow::binary; +using ::arrow::default_memory_pool; +using ::arrow::field; +using ::arrow::int64; +using ::arrow::ProxyMemoryPool; +using ::arrow::struct_; +using ::arrow::StructArray; +using ::arrow::Type; +using ::arrow::extension::variant; +using internal::BinaryArrayFromValues; +using internal::MakeVariantValueView; + +void AssertPrimitiveType(std::string_view value, const VariantMetadataView& metadata, + VariantPrimitiveType expected) { + ASSERT_OK_AND_ASSIGN(auto view, VariantValueView::Make(value, metadata)); + ASSERT_EQ(VariantBasicType::kPrimitive, view.basic_type()); + ASSERT_EQ(expected, std::get(view.data()).type()); +} + +void AssertPrimitiveFieldType(const VariantObjectView& object, std::string_view name, + const VariantMetadataView& metadata, + VariantPrimitiveType expected) { + const auto* field = object.FindField(name); + ASSERT_NE(nullptr, field) << "Missing Variant object field: " << name; + AssertPrimitiveType(field->value, metadata, expected); +} + +} // namespace + +TEST(TestVariantBuilder, Object) { + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + ASSERT_OK(object.AppendVariantNull("c")); + ASSERT_OK(object.AppendVariantNull("b")); + ASSERT_OK(object.AppendVariantNull("a")); + ASSERT_OK(object.Finish()); + + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto view, MakeVariantValueView(encoded)); + ASSERT_EQ(VariantBasicType::kObject, view.basic_type()); + const auto& fields = std::get(view.data()).fields(); + ASSERT_EQ(3, fields.size()); + ASSERT_EQ("a", fields[0].name); + ASSERT_EQ("b", fields[1].name); + ASSERT_EQ("c", fields[2].name); +} + +TEST(TestVariantBuilder, List) { + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto list, builder.StartList()); + ASSERT_OK(list.AppendVariantNull()); + ASSERT_OK(list.AppendInt32(42)); + ASSERT_OK(list.AppendString("x")); + ASSERT_OK(list.Finish()); + + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + ASSERT_OK_AND_ASSIGN( + auto view, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + ASSERT_EQ(VariantBasicType::kArray, view.basic_type()); + const auto& elements = std::get(view.data()).elements(); + ASSERT_EQ(3, elements.size()); + + ASSERT_OK_AND_ASSIGN(auto element, VariantValueView::Make(elements[1], metadata)); + ASSERT_EQ(VariantPrimitiveType::kInt32, + std::get(element.data()).type()); +} + +TEST(TestVariantBuilder, Nested) { + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + ASSERT_OK_AND_ASSIGN(auto list, object.StartList("items")); + ASSERT_OK(list.AppendInt32(1)); + ASSERT_OK_AND_ASSIGN(auto child, list.StartObject()); + ASSERT_OK(child.AppendString("name", "x")); + ASSERT_OK(child.Finish()); + ASSERT_OK(list.Finish()); + ASSERT_OK(object.Finish()); + + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + ASSERT_OK_AND_ASSIGN( + auto root, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + const auto& root_fields = std::get(root.data()).fields(); + ASSERT_EQ(1, root_fields.size()); + ASSERT_EQ("items", root_fields[0].name); + + ASSERT_OK_AND_ASSIGN(auto items, + VariantValueView::Make(root_fields[0].value, metadata)); + const auto& item_values = std::get(items.data()).elements(); + ASSERT_EQ(2, item_values.size()); + ASSERT_OK_AND_ASSIGN(auto item, VariantValueView::Make(item_values[1], metadata)); + ASSERT_EQ("name", std::get(item.data()).fields()[0].name); +} + +TEST(TestVariantBuilder, ObjectAppends) { + const std::string uuid(16, '\1'); + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + ASSERT_OK(object.AppendInt8("int8", 1)); + ASSERT_OK(object.AppendInt64("int64", 2)); + ASSERT_OK(object.AppendDouble("double", 3)); + ASSERT_OK(object.AppendDecimal4("decimal", 1234, 2)); + ASSERT_OK(object.AppendBinary("binary", "abc")); + ASSERT_OK(object.AppendDate("date", 1)); + ASSERT_OK(object.AppendTimestampNanos("ts", 2, true)); + ASSERT_OK(object.AppendUuid("uuid", uuid)); + ASSERT_OK(object.Finish()); + + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + ASSERT_OK_AND_ASSIGN( + auto root, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + const auto& object_view = std::get(root.data()); + ASSERT_EQ(8, object_view.fields().size()); + AssertPrimitiveFieldType(object_view, "int8", metadata, VariantPrimitiveType::kInt8); + AssertPrimitiveFieldType(object_view, "int64", metadata, VariantPrimitiveType::kInt64); + AssertPrimitiveFieldType(object_view, "double", metadata, + VariantPrimitiveType::kDouble); + AssertPrimitiveFieldType(object_view, "decimal", metadata, + VariantPrimitiveType::kDecimal4); + AssertPrimitiveFieldType(object_view, "binary", metadata, + VariantPrimitiveType::kBinary); + AssertPrimitiveFieldType(object_view, "date", metadata, VariantPrimitiveType::kDate); + AssertPrimitiveFieldType(object_view, "ts", metadata, + VariantPrimitiveType::kTimestampNanos); + AssertPrimitiveFieldType(object_view, "uuid", metadata, VariantPrimitiveType::kUuid); +} + +TEST(TestVariantBuilder, ListAppends) { + const std::string uuid(16, '\2'); + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto list, builder.StartList()); + ASSERT_OK(list.AppendInt8(1)); + ASSERT_OK(list.AppendInt64(2)); + ASSERT_OK(list.AppendDouble(3)); + ASSERT_OK(list.AppendDecimal4(1234, 2)); + ASSERT_OK(list.AppendBinary("abc")); + ASSERT_OK(list.AppendDate(1)); + ASSERT_OK(list.AppendTimestampNanos(2, true)); + ASSERT_OK(list.AppendUuid(uuid)); + ASSERT_OK(list.Finish()); + + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + ASSERT_OK_AND_ASSIGN( + auto root, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + const auto& elements = std::get(root.data()).elements(); + ASSERT_EQ(8, elements.size()); + AssertPrimitiveType(elements[0], metadata, VariantPrimitiveType::kInt8); + AssertPrimitiveType(elements[1], metadata, VariantPrimitiveType::kInt64); + AssertPrimitiveType(elements[2], metadata, VariantPrimitiveType::kDouble); + AssertPrimitiveType(elements[3], metadata, VariantPrimitiveType::kDecimal4); + AssertPrimitiveType(elements[4], metadata, VariantPrimitiveType::kBinary); + AssertPrimitiveType(elements[5], metadata, VariantPrimitiveType::kDate); + AssertPrimitiveType(elements[6], metadata, VariantPrimitiveType::kTimestampNanos); + AssertPrimitiveType(elements[7], metadata, VariantPrimitiveType::kUuid); +} + +TEST(TestVariantBuilder, Rollback) { + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + { + ASSERT_OK_AND_ASSIGN(auto child, object.StartObject("drop")); + ASSERT_OK(child.AppendString("nested", "x")); + } + ASSERT_OK(object.AppendInt32("ok", 1)); + ASSERT_OK(object.Finish()); + + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto view, MakeVariantValueView(encoded)); + const auto& fields = std::get(view.data()).fields(); + ASSERT_EQ(1, fields.size()); + ASSERT_EQ("ok", fields[0].name); +} + +TEST(TestVariantBuilder, Duplicate) { + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + ASSERT_OK(object.AppendInt32("a", 1)); + ASSERT_RAISES(Invalid, object.AppendString("a", "x")); +} + +TEST(TestVariantBuilder, UsesMemoryPool) { + ProxyMemoryPool pool(default_memory_pool()); + VariantBuilder builder(&pool); + ASSERT_OK(builder.AppendString(std::string(128, 'x'))); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_GT(pool.total_bytes_allocated(), 0); + + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + ASSERT_OK(VariantValueView::Validate(std::string_view{*encoded.value}, metadata)); +} + +TEST(TestVariantBuilder, ArrayBuilder) { + VariantArrayBuilder builder; + ASSERT_OK(builder.AppendNull()); + ASSERT_OK(builder.AppendVariantNull()); + ASSERT_OK(builder.AppendInt32(42)); + ASSERT_OK(builder.AppendInt8(1)); + ASSERT_OK(builder.AppendDouble(2)); + ASSERT_OK(builder.AppendBinary("abc")); + ASSERT_OK(builder.AppendTimestampMicros(3, true)); + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + ASSERT_OK(object.AppendString("a", "x")); + ASSERT_OK(object.Finish()); + + ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + ASSERT_EQ(8, array->length()); + ASSERT_TRUE(array->IsNull(0)); + ASSERT_FALSE(array->IsNull(1)); + + const auto& storage = + ::arrow::internal::checked_cast(*array->storage()); + ASSERT_FALSE(storage.type()->field(0)->nullable()); + ASSERT_FALSE(storage.type()->field(1)->nullable()); + ASSERT_EQ(Type::BINARY, storage.field(0)->type_id()); + ASSERT_EQ(Type::BINARY, storage.field(1)->type_id()); +} + +TEST(TestVariantBuilder, FromStorage) { + VariantBuilder value; + ASSERT_OK(value.AppendInt32(1)); + ASSERT_OK_AND_ASSIGN(auto encoded, value.Finish()); + + auto metadata = BinaryArrayFromValues({std::string_view{*encoded.metadata}}); + auto values = BinaryArrayFromValues({std::string_view{*encoded.value}}); + ASSERT_OK_AND_ASSIGN( + auto storage, + StructArray::Make({metadata, values}, {field("metadata", binary(), false), + field("value", binary(), false)})); + + ASSERT_OK_AND_ASSIGN(auto array, MakeVariantArrayFromStorage(storage)); + ASSERT_EQ(1, array->length()); + ASSERT_TRUE(array->type()->Equals(variant(storage->type()))); +} + +TEST(TestVariantBuilder, FromShredded) { + VariantBuilder value; + ASSERT_OK(value.AppendVariantNull()); + ASSERT_OK_AND_ASSIGN(auto encoded, value.Finish()); + + auto metadata = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto values = BinaryArrayFromValues({std::nullopt, std::string_view{*encoded.value}}); + auto typed = ArrayFromJSON(int64(), "[1, null]"); + auto storage_type = struct_({field("metadata", binary(), false), + field("value", binary()), field("typed_value", int64())}); + + ASSERT_OK_AND_ASSIGN( + auto array, MakeVariantArrayFromChildren(storage_type, {metadata, values, typed})); + ASSERT_EQ(2, array->length()); + ASSERT_TRUE(array->type()->Equals(variant(storage_type))); +} + +TEST(TestVariantBuilder, FallbackValue) { + VariantBuilder value; + ASSERT_OK(value.AppendString("x")); + ASSERT_OK_AND_ASSIGN(auto encoded, value.Finish()); + + VariantValueArrayBuilder builder; + ASSERT_OK(builder.AppendNull()); + ASSERT_OK(builder.AppendEncodedValue(std::string_view{*encoded.metadata}, + std::string_view{*encoded.value})); + ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + ASSERT_EQ(2, array->length()); + ASSERT_TRUE(array->IsNull(0)); + ASSERT_EQ(std::string_view{*encoded.value}, array->GetView(1)); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/encoding.cc b/cpp/src/parquet/variant/encoding.cc new file mode 100644 index 000000000000..95c61a235589 --- /dev/null +++ b/cpp/src/parquet/variant/encoding.cc @@ -0,0 +1,470 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/encoding.h" + +#include +#include +#include + +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" +#include "parquet/variant/encoding_internal.h" + +namespace parquet::variant { + +namespace bit_util = ::arrow::bit_util; + +namespace { + +uint32_t ReadLittleEndian(std::string_view data, size_t offset, size_t width) { + DCHECK_LE(width, sizeof(uint32_t)); + uint32_t value = 0; + std::memcpy(&value, data.data() + offset, width); + return bit_util::FromLittleEndian(value); +} + +Status CheckAvailable(std::string_view data, size_t offset, size_t size, + std::string_view context) { + if (offset > data.size() || data.size() - offset < size) { + return Status::Invalid("Invalid Variant encoding: truncated ", context); + } + return Status::OK(); +} + +Status PrimitivePayloadSize(std::string_view value, size_t offset, + VariantPrimitiveType primitive, size_t* size) { + switch (primitive) { + case VariantPrimitiveType::kNull: + case VariantPrimitiveType::kBooleanTrue: + case VariantPrimitiveType::kBooleanFalse: + *size = 0; + return Status::OK(); + case VariantPrimitiveType::kInt8: + *size = 1; + return Status::OK(); + case VariantPrimitiveType::kInt16: + *size = 2; + return Status::OK(); + case VariantPrimitiveType::kInt32: + case VariantPrimitiveType::kDate: + case VariantPrimitiveType::kFloat: + *size = 4; + return Status::OK(); + case VariantPrimitiveType::kInt64: + case VariantPrimitiveType::kDouble: + case VariantPrimitiveType::kTimestampMicros: + case VariantPrimitiveType::kTimestampNTZMicros: + case VariantPrimitiveType::kTimeNTZMicros: + case VariantPrimitiveType::kTimestampNanos: + case VariantPrimitiveType::kTimestampNTZNanos: + *size = 8; + return Status::OK(); + case VariantPrimitiveType::kDecimal4: + *size = 5; + return Status::OK(); + case VariantPrimitiveType::kDecimal8: + *size = 9; + return Status::OK(); + case VariantPrimitiveType::kDecimal16: + *size = 17; + return Status::OK(); + case VariantPrimitiveType::kUuid: + *size = 16; + return Status::OK(); + case VariantPrimitiveType::kBinary: + case VariantPrimitiveType::kString: { + ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, 4, "variable-length size")); + const uint32_t length = ReadLittleEndian(value, offset, 4); + *size = 4 + static_cast(length); + return Status::OK(); + } + } + return Status::Invalid("Invalid Variant encoding: unknown primitive type"); +} + +Status ParsePrimitive(std::string_view value, size_t offset, + VariantPrimitiveType primitive, size_t& consumed) { + if (!internal::IsKnownVariantPrimitive(primitive)) { + return Status::Invalid("Invalid Variant encoding: unknown primitive type ", + static_cast(primitive)); + } + + size_t payload_size = 0; + ARROW_RETURN_NOT_OK(PrimitivePayloadSize(value, offset, primitive, &payload_size)); + ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, payload_size, "primitive value")); + + if (internal::IsDecimalVariantPrimitive(primitive)) { + const auto scale = static_cast(value[offset]); + ARROW_RETURN_NOT_OK(internal::ValidateDecimalScale(scale)); + } + + if (primitive == VariantPrimitiveType::kString) { + const uint32_t length = ReadLittleEndian(value, offset, 4); + ARROW_RETURN_NOT_OK(internal::ValidateUtf8(value.substr(offset + 4, length), + "primitive string value")); + } + + consumed = payload_size; + return Status::OK(); +} + +Status ParseValue(std::string_view value, const VariantMetadataView& metadata, + size_t& consumed, VariantValueView* out); + +Status ParseArray(std::string_view value, const VariantMetadataView& metadata, + uint8_t header, size_t& consumed, VariantValueView* out) { + const auto offset_size = static_cast((header & 0x03) + 1); + const bool is_large = (header & 0x04) != 0; + const size_t count_size = is_large ? 4 : 1; + + size_t offset = 1; + ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, count_size, "array size")); + const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); + offset += count_size; + + ARROW_RETURN_NOT_OK( + CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, + "array offsets")); + + std::vector offsets(num_elements + 1); + for (uint32_t i = 0; i <= num_elements; ++i) { + offsets[i] = ReadLittleEndian(value, offset, offset_size); + offset += offset_size; + } + + if (offsets[0] != 0) { + return Status::Invalid("Invalid Variant encoding: first array offset must be 0"); + } + for (uint32_t i = 0; i < num_elements; ++i) { + if (offsets[i] > offsets[i + 1]) { + return Status::Invalid("Invalid Variant encoding: array offsets must be monotonic"); + } + } + + const size_t values_start = offset; + const size_t total_value_size = offsets[num_elements]; + ARROW_RETURN_NOT_OK( + CheckAvailable(value, values_start, total_value_size, "array values")); + + size_t current = 0; + std::vector array_elements; + if (out != nullptr) { + array_elements.reserve(num_elements); + } + for (uint32_t i = 0; i < num_elements; ++i) { + if (offsets[i] != current) { + return Status::Invalid( + "Invalid Variant encoding: array offset does not match value boundary"); + } + size_t child_consumed = 0; + ARROW_RETURN_NOT_OK(ParseValue(value.substr(values_start + current), metadata, + child_consumed, + /*out=*/nullptr)); + current += child_consumed; + if (current != offsets[i + 1]) { + return Status::Invalid( + "Invalid Variant encoding: array value does not end at next offset"); + } + if (out != nullptr) { + array_elements.push_back(value.substr(values_start + offsets[i], child_consumed)); + } + } + + if (current != total_value_size) { + return Status::Invalid("Invalid Variant encoding: array values have trailing data"); + } + + consumed = values_start + total_value_size; + if (out != nullptr) { + *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kArray, + VariantArrayView(std::move(array_elements))); + } + return Status::OK(); +} + +Status ParseObject(std::string_view value, const VariantMetadataView& metadata, + uint8_t header, size_t& consumed, VariantValueView* out) { + const auto offset_size = static_cast((header & 0x03) + 1); + const auto id_size = static_cast(((header >> 2) & 0x03) + 1); + const bool is_large = (header & 0x10) != 0; + const size_t count_size = is_large ? 4 : 1; + + size_t offset = 1; + ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, count_size, "object size")); + const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); + offset += count_size; + + ARROW_RETURN_NOT_OK(CheckAvailable( + value, offset, static_cast(num_elements) * id_size, "object field ids")); + std::vector field_ids(num_elements); + for (uint32_t i = 0; i < num_elements; ++i) { + field_ids[i] = ReadLittleEndian(value, offset, id_size); + offset += id_size; + } + + ARROW_RETURN_NOT_OK( + CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, + "object field offsets")); + std::vector field_offsets(num_elements + 1); + for (uint32_t i = 0; i <= num_elements; ++i) { + field_offsets[i] = ReadLittleEndian(value, offset, offset_size); + offset += offset_size; + } + + const size_t values_start = offset; + const size_t total_value_size = field_offsets[num_elements]; + if (num_elements == 0 && total_value_size != 0) { + return Status::Invalid( + "Invalid Variant encoding: empty object must have zero value size"); + } + ARROW_RETURN_NOT_OK( + CheckAvailable(value, values_start, total_value_size, "object values")); + + std::vector object_fields; + if (out != nullptr) { + object_fields.reserve(num_elements); + } + + std::string_view previous_name; + for (uint32_t i = 0; i < num_elements; ++i) { + if (field_ids[i] >= metadata.dictionary_size()) { + return Status::Invalid("Invalid Variant encoding: object field id ", field_ids[i], + " is outside metadata dictionary of size ", + metadata.dictionary_size()); + } + const auto name = metadata.string(field_ids[i]); + if (i > 0 && !(previous_name < name)) { + return Status::Invalid( + "Invalid Variant encoding: object field names must be sorted and unique"); + } + previous_name = name; + + const auto field_offset = field_offsets[i]; + if (field_offset >= total_value_size) { + return Status::Invalid( + "Invalid Variant encoding: object field offset is outside values"); + } + } + + std::vector value_offsets = field_offsets; + std::ranges::sort(value_offsets); + if (std::ranges::adjacent_find(value_offsets) != value_offsets.end()) { + return Status::Invalid( + "Invalid Variant encoding: object field offsets must be unique"); + } + if (value_offsets.front() != 0) { + return Status::Invalid("Invalid Variant encoding: object values have leading data"); + } + + for (uint32_t i = 0; i < num_elements; ++i) { + const uint32_t start = value_offsets[i]; + const uint32_t end = value_offsets[i + 1]; + size_t child_consumed = 0; + ARROW_RETURN_NOT_OK(ParseValue(value.substr(values_start + start), metadata, + child_consumed, + /*out=*/nullptr)); + if (child_consumed != end - start) { + return Status::Invalid( + "Invalid Variant encoding: object value does not end at next value boundary"); + } + } + + if (out != nullptr) { + for (uint32_t i = 0; i < num_elements; ++i) { + const auto field_offset = field_offsets[i]; + auto offset_it = std::ranges::lower_bound(value_offsets, field_offset); + DCHECK(offset_it != value_offsets.end()); + DCHECK(offset_it + 1 != value_offsets.end()); + DCHECK(*offset_it == field_offset); + const auto end = *(offset_it + 1); + object_fields.push_back(VariantObjectField{ + .name = metadata.string(field_ids[i]), + .field_id = field_ids[i], + .value = value.substr(values_start + field_offset, end - field_offset)}); + } + } + + consumed = values_start + total_value_size; + if (out != nullptr) { + *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kObject, + VariantObjectView(std::move(object_fields))); + } + return Status::OK(); +} + +Status ParseValue(std::string_view value, const VariantMetadataView& metadata, + size_t& consumed, VariantValueView* out) { + ARROW_RETURN_NOT_OK(CheckAvailable(value, 0, 1, "value header")); + + const auto metadata_byte = static_cast(value[0]); + const auto basic_type = static_cast(metadata_byte & 0x03); + const auto header = static_cast(metadata_byte >> 2); + + switch (basic_type) { + case VariantBasicType::kPrimitive: { + const auto primitive = static_cast(header); + size_t payload_size = 0; + ARROW_RETURN_NOT_OK(ParsePrimitive(value, 1, primitive, payload_size)); + consumed = 1 + payload_size; + if (out != nullptr) { + *out = VariantValueView( + value.substr(0, consumed), VariantBasicType::kPrimitive, + VariantPrimitiveView(primitive, value.substr(1, payload_size))); + } + return Status::OK(); + } + case VariantBasicType::kShortString: { + ARROW_RETURN_NOT_OK(CheckAvailable(value, 1, header, "short string value")); + ARROW_RETURN_NOT_OK( + internal::ValidateUtf8(value.substr(1, header), "short string value")); + consumed = 1 + header; + if (out != nullptr) { + *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kShortString, + VariantShortStringView(value.substr(1, header))); + } + return Status::OK(); + } + case VariantBasicType::kObject: + return ParseObject(value, metadata, header, consumed, out); + case VariantBasicType::kArray: + return ParseArray(value, metadata, header, consumed, out); + } + return Status::Invalid("Invalid Variant encoding: unknown basic type"); +} + +} // namespace + +Result VariantMetadataView::Make(std::string_view metadata) { + ARROW_RETURN_NOT_OK(CheckAvailable(metadata, 0, 1, "metadata header")); + const auto header = static_cast(metadata[0]); + const auto version = static_cast(header & internal::kMetadataVersionMask); + if (version != internal::kVariantVersion) { + return Status::Invalid("Invalid Variant metadata: expected version 1, got ", + static_cast(version)); + } + + VariantMetadataView view; + view.metadata_ = metadata; + view.sorted_strings_ = (header & internal::kMetadataSortedStringsMask) != 0; + view.offset_size_ = static_cast(((header >> 6) & 0x03) + 1); + + ARROW_RETURN_NOT_OK( + CheckAvailable(metadata, 1, view.offset_size_, "metadata dictionary size")); + const uint32_t dictionary_size = ReadLittleEndian(metadata, 1, view.offset_size_); + const size_t offsets_offset = 1 + view.offset_size_; + ARROW_RETURN_NOT_OK( + CheckAvailable(metadata, offsets_offset, + (static_cast(dictionary_size) + 1) * view.offset_size_, + "metadata dictionary offsets")); + + std::vector offsets(dictionary_size + 1); + for (uint32_t i = 0; i <= dictionary_size; ++i) { + offsets[i] = ReadLittleEndian(metadata, offsets_offset + i * view.offset_size_, + view.offset_size_); + } + + if (offsets[0] != 0) { + return Status::Invalid("Invalid Variant metadata: first dictionary offset must be 0"); + } + for (uint32_t i = 0; i < dictionary_size; ++i) { + if (offsets[i] > offsets[i + 1]) { + return Status::Invalid( + "Invalid Variant metadata: dictionary offsets must be monotonic"); + } + } + + const size_t bytes_offset = + offsets_offset + (static_cast(dictionary_size) + 1) * view.offset_size_; + const size_t bytes_size = offsets[dictionary_size]; + ARROW_RETURN_NOT_OK( + CheckAvailable(metadata, bytes_offset, bytes_size, "metadata dictionary bytes")); + if (metadata.size() != bytes_offset + bytes_size) { + return Status::Invalid("Invalid Variant metadata: trailing bytes after dictionary"); + } + + view.strings_.reserve(dictionary_size); + for (uint32_t i = 0; i < dictionary_size; ++i) { + auto string = metadata.substr(bytes_offset + offsets[i], offsets[i + 1] - offsets[i]); + ARROW_RETURN_NOT_OK(internal::ValidateUtf8(string, "metadata dictionary string")); + if (view.sorted_strings_ && i > 0 && !(view.strings_.back() < string)) { + return Status::Invalid( + "Invalid Variant metadata: sorted dictionary strings must be unique and " + "lexicographically sorted"); + } + view.strings_.push_back(string); + } + + return view; +} + +std::string_view VariantMetadataView::string(uint32_t id) const { + DCHECK_LT(id, strings_.size()); + return strings_[id]; +} + +std::optional VariantMetadataView::FindString(std::string_view value) const { + if (sorted_strings_) { + const auto it = std::ranges::lower_bound(strings_, value); + if (it != strings_.end() && *it == value) { + return static_cast(it - strings_.begin()); + } + return std::nullopt; + } + + for (uint32_t i = 0; i < strings_.size(); ++i) { + if (strings_[i] == value) { + return i; + } + } + return std::nullopt; +} + +const VariantObjectField* VariantObjectView::FindField(std::string_view name) const { + const auto it = std::ranges::lower_bound(fields_, name, {}, &VariantObjectField::name); + return (it == fields_.end() || it->name != name) ? nullptr : &*it; +} + +bool VariantObjectView::ContainsField(std::string_view name) const { + // The Parquet Variant encoding requires object fields to be sorted by name, so field + // lookup can use binary search. + return std::ranges::binary_search(fields_, name, {}, &VariantObjectField::name); +} + +Result VariantValueView::Make(std::string_view value, + const VariantMetadataView& metadata) { + size_t consumed = 0; + VariantValueView view({}, VariantBasicType::kPrimitive, + VariantPrimitiveView(VariantPrimitiveType::kNull, {})); + ARROW_RETURN_NOT_OK(ParseValue(value, metadata, consumed, &view)); + if (consumed != value.size()) { + return Status::Invalid("Invalid Variant encoding: trailing bytes after value"); + } + return view; +} + +Status VariantValueView::Validate(std::string_view value, + const VariantMetadataView& metadata) { + size_t consumed = 0; + ARROW_RETURN_NOT_OK(ParseValue(value, metadata, consumed, /*out=*/nullptr)); + if (consumed != value.size()) { + return Status::Invalid("Invalid Variant encoding: trailing bytes after value"); + } + return Status::OK(); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/encoding.h b/cpp/src/parquet/variant/encoding.h new file mode 100644 index 000000000000..1dbf3d9a104d --- /dev/null +++ b/cpp/src/parquet/variant/encoding.h @@ -0,0 +1,170 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/result.h" +#include "arrow/status.h" +#include "parquet/platform.h" + +namespace parquet::variant { + +using ::arrow::Result; +using ::arrow::Status; + +enum class VariantBasicType : uint8_t { + kPrimitive = 0, + kShortString = 1, + kObject = 2, + kArray = 3, +}; + +enum class VariantPrimitiveType : uint8_t { + kNull = 0, + kBooleanTrue = 1, + kBooleanFalse = 2, + kInt8 = 3, + kInt16 = 4, + kInt32 = 5, + kInt64 = 6, + kDouble = 7, + kDecimal4 = 8, + kDecimal8 = 9, + kDecimal16 = 10, + kDate = 11, + kTimestampMicros = 12, + kTimestampNTZMicros = 13, + kFloat = 14, + kBinary = 15, + kString = 16, + kTimeNTZMicros = 17, + kTimestampNanos = 18, + kTimestampNTZNanos = 19, + kUuid = 20, +}; + +struct PARQUET_EXPORT VariantObjectField { + std::string_view name; + uint32_t field_id = 0; + std::string_view value; +}; + +class PARQUET_EXPORT VariantMetadataView { + public: + /// Parse Variant metadata bytes without copying them. + /// + /// The returned view and all string views borrowed from it are valid only while the + /// input metadata bytes remain alive and unchanged. + static Result Make(std::string_view metadata); + + std::string_view metadata() const { return metadata_; } + bool sorted_strings() const { return sorted_strings_; } + uint8_t offset_size() const { return offset_size_; } + uint32_t dictionary_size() const { return static_cast(strings_.size()); } + + std::string_view string(uint32_t id) const; + std::optional FindString(std::string_view value) const; + + private: + std::string_view metadata_; + bool sorted_strings_ = false; + uint8_t offset_size_ = 0; + std::vector strings_; +}; + +class PARQUET_EXPORT VariantPrimitiveView { + public: + VariantPrimitiveView(VariantPrimitiveType type, std::string_view payload) + : type_(type), payload_(payload) {} + + VariantPrimitiveType type() const { return type_; } + std::string_view payload() const { return payload_; } + + private: + VariantPrimitiveType type_ = VariantPrimitiveType::kNull; + std::string_view payload_; +}; + +class PARQUET_EXPORT VariantShortStringView { + public: + explicit VariantShortStringView(std::string_view string) : string_(string) {} + + std::string_view string() const { return string_; } + + private: + std::string_view string_; +}; + +class PARQUET_EXPORT VariantObjectView { + public: + explicit VariantObjectView(std::vector fields) + : fields_(std::move(fields)) {} + + const std::vector& fields() const { return fields_; } + const VariantObjectField* FindField(std::string_view name) const; + bool ContainsField(std::string_view name) const; + + private: + std::vector fields_; +}; + +class PARQUET_EXPORT VariantArrayView { + public: + explicit VariantArrayView(std::vector elements) + : elements_(std::move(elements)) {} + + const std::vector& elements() const { return elements_; } + + private: + std::vector elements_; +}; + +class PARQUET_EXPORT VariantValueView { + public: + using Data = std::variant; + + VariantValueView(std::string_view value, VariantBasicType basic_type, Data data) + : value_(value), basic_type_(basic_type), data_(std::move(data)) {} + + /// Parse Variant value bytes without copying them. + /// + /// The returned view and any nested object/list/string views are valid only while the + /// input value bytes remain alive and unchanged. The metadata view must also remain + /// valid while object field names are accessed. + static Result Make(std::string_view value, + const VariantMetadataView& metadata); + static Status Validate(std::string_view value, const VariantMetadataView& metadata); + + std::string_view value() const { return value_; } + VariantBasicType basic_type() const { return basic_type_; } + const Data& data() const { return data_; } + + private: + std::string_view value_; + VariantBasicType basic_type_; + Data data_; +}; + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/encoding_internal.h b/cpp/src/parquet/variant/encoding_internal.h new file mode 100644 index 000000000000..f34e088a2ef3 --- /dev/null +++ b/cpp/src/parquet/variant/encoding_internal.h @@ -0,0 +1,147 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include + +#include "arrow/status.h" +#include "arrow/util/utf8.h" +#include "parquet/variant/encoding.h" + +namespace parquet::variant::internal { + +using ::arrow::Status; + +namespace util = ::arrow::util; + +inline constexpr uint8_t kVariantVersion = 1; +inline constexpr uint8_t kMetadataVersionMask = 0x0F; +inline constexpr uint8_t kMetadataSortedStringsMask = 0x10; +inline constexpr uint8_t kMetadataReservedMask = 0x20; + +template +struct VariantBasicTypeTraits {}; + +template <> +struct VariantBasicTypeTraits { + using ViewType = VariantPrimitiveView; +}; + +template <> +struct VariantBasicTypeTraits { + using ViewType = VariantShortStringView; +}; + +template <> +struct VariantBasicTypeTraits { + using ViewType = VariantObjectView; +}; + +template <> +struct VariantBasicTypeTraits { + using ViewType = VariantArrayView; +}; + +inline bool IsKnownVariantPrimitive(VariantPrimitiveType type) { + return type <= VariantPrimitiveType::kUuid; +} + +inline bool IsDecimalVariantPrimitive(VariantPrimitiveType type) { + return type == VariantPrimitiveType::kDecimal4 || + type == VariantPrimitiveType::kDecimal8 || + type == VariantPrimitiveType::kDecimal16; +} + +inline Status ValidateDecimalScale(uint8_t scale) { + if (scale > 38) { + return Status::Invalid("Invalid Variant decimal scale ", scale, " exceeds 38"); + } + return Status::OK(); +} + +inline Status ValidateUtf8(std::string_view value, std::string_view context) { + util::InitializeUTF8(); + if (!util::ValidateUTF8(reinterpret_cast(value.data()), value.size())) { + return Status::Invalid("Invalid Variant encoding: ", context, " is not valid UTF-8"); + } + return Status::OK(); +} + +template +concept HeaderOnlyVariantPrimitive = + type == VariantPrimitiveType::kNull || type == VariantPrimitiveType::kBooleanTrue || + type == VariantPrimitiveType::kBooleanFalse; + +template +struct VariantFixedPrimitiveTraits {}; + +#define VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(TYPE, C_TYPE) \ + template <> \ + struct VariantFixedPrimitiveTraits { \ + using CType = C_TYPE; \ + }; + +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Int8, int8_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Int16, int16_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Int32, int32_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Int64, int64_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Float, float) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Double, double) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(Date, int32_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(TimeNTZMicros, int64_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(TimestampMicros, int64_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(TimestampNTZMicros, int64_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(TimestampNanos, int64_t) +VARIANT_FIXED_PRIMITIVE_TRAITS_DEF(TimestampNTZNanos, int64_t) + +#undef VARIANT_FIXED_PRIMITIVE_TRAITS_DEF + +template +concept FixedVariantPrimitive = + requires { typename VariantFixedPrimitiveTraits::CType; }; + +template +struct VariantDecimalPrimitiveTraits {}; + +#define VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(TYPE, C_TYPE) \ + template <> \ + struct VariantDecimalPrimitiveTraits { \ + using CType = C_TYPE; \ + }; + +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal4, int32_t) +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal8, int64_t) + +#undef VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF + +template +concept DecimalVariantPrimitive = + requires { typename VariantDecimalPrimitiveTraits::CType; }; + +template +concept LengthPrefixedVariantPrimitive = + type == VariantPrimitiveType::kBinary || type == VariantPrimitiveType::kString; + +template +concept Decimal16VariantPrimitive = type == VariantPrimitiveType::kDecimal16; + +template +concept UuidVariantPrimitive = type == VariantPrimitiveType::kUuid; + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/encoding_test.cc b/cpp/src/parquet/variant/encoding_test.cc new file mode 100644 index 000000000000..e68b15692394 --- /dev/null +++ b/cpp/src/parquet/variant/encoding_test.cc @@ -0,0 +1,190 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/encoding.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/test_util_internal.h" + +#include +#include +#include + +#include "arrow/testing/gtest_util.h" + +namespace parquet::variant { + +TEST(TestVariantEncoding, EmptyMetadata) { + VariantBuilder builder; + ASSERT_OK(builder.AppendVariantNull()); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_EQ(std::string_view("\x01\x00\x00", 3), std::string_view{*encoded.metadata}); + + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + ASSERT_FALSE(metadata.sorted_strings()); + ASSERT_EQ(0, metadata.dictionary_size()); +} + +TEST(TestVariantEncoding, Primitive) { + std::array decimal16 = {0}; + std::array uuid = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; + + auto check = [](auto append, VariantPrimitiveType expected) { + VariantBuilder builder; + ASSERT_OK(append(builder)); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto view, internal::MakeVariantValueView(encoded)); + ASSERT_EQ(VariantBasicType::kPrimitive, view.basic_type()); + ASSERT_EQ(expected, std::get(view.data()).type()); + }; + + check([](VariantBuilder& b) { return b.AppendVariantNull(); }, + VariantPrimitiveType::kNull); + check([](VariantBuilder& b) { return b.AppendBoolean(true); }, + VariantPrimitiveType::kBooleanTrue); + check([](VariantBuilder& b) { return b.AppendInt8(1); }, VariantPrimitiveType::kInt8); + check([](VariantBuilder& b) { return b.AppendInt16(2); }, VariantPrimitiveType::kInt16); + check([](VariantBuilder& b) { return b.AppendInt32(3); }, VariantPrimitiveType::kInt32); + check([](VariantBuilder& b) { return b.AppendInt64(4); }, VariantPrimitiveType::kInt64); + check([](VariantBuilder& b) { return b.AppendFloat(1.5F); }, + VariantPrimitiveType::kFloat); + check([](VariantBuilder& b) { return b.AppendDouble(2.5); }, + VariantPrimitiveType::kDouble); + check([](VariantBuilder& b) { return b.AppendDecimal4(123, 2); }, + VariantPrimitiveType::kDecimal4); + check([](VariantBuilder& b) { return b.AppendDecimal8(123, 2); }, + VariantPrimitiveType::kDecimal8); + check( + [&](VariantBuilder& b) { + return b.AppendDecimal16(std::string_view(decimal16.data(), decimal16.size()), 2); + }, + VariantPrimitiveType::kDecimal16); + check([](VariantBuilder& b) { return b.AppendBinary("abc"); }, + VariantPrimitiveType::kBinary); + check([](VariantBuilder& b) { return b.AppendString("abc"); }, + VariantPrimitiveType::kString); + check([](VariantBuilder& b) { return b.AppendDate(1); }, VariantPrimitiveType::kDate); + check([](VariantBuilder& b) { return b.AppendTimeNTZMicros(1); }, + VariantPrimitiveType::kTimeNTZMicros); + check([](VariantBuilder& b) { return b.AppendTimestampMicros(1, true); }, + VariantPrimitiveType::kTimestampMicros); + check([](VariantBuilder& b) { return b.AppendTimestampMicros(1, false); }, + VariantPrimitiveType::kTimestampNTZMicros); + check([](VariantBuilder& b) { return b.AppendTimestampNanos(1, true); }, + VariantPrimitiveType::kTimestampNanos); + check([](VariantBuilder& b) { return b.AppendTimestampNanos(1, false); }, + VariantPrimitiveType::kTimestampNTZNanos); + check( + [&](VariantBuilder& b) { + return b.AppendUuid(std::string_view(uuid.data(), uuid.size())); + }, + VariantPrimitiveType::kUuid); +} + +TEST(TestVariantEncoding, ShortString) { + VariantBuilder builder; + ASSERT_OK(builder.AppendShortString("abc")); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto view, internal::MakeVariantValueView(encoded)); + ASSERT_EQ(VariantBasicType::kShortString, view.basic_type()); + ASSERT_EQ("abc", std::get(view.data()).string()); +} + +TEST(TestVariantEncoding, InvalidMetadata) { + // Metadata version is 2 instead of 1. + ASSERT_RAISES(Invalid, VariantMetadataView::Make(std::string("\x02\x00\x00", 3))); + // Dictionary offsets are truncated. + ASSERT_RAISES(Invalid, VariantMetadataView::Make(std::string("\x01\x01\x00", 3))); + // Dictionary string payload is not valid UTF-8. + ASSERT_RAISES(Invalid, + VariantMetadataView::Make(std::string("\x01\x01\x00\x01\xff", 5))); +} + +TEST(TestVariantEncoding, ReservedBits) { + // Metadata bytes: + // - 0x21: version 1 with reserved bit 0x20 set + // - 0x00: zero dictionary entries, encoded with one-byte width + // - 0x00: first and final dictionary byte offset + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string("\x21\x00\x00", 3))); + ASSERT_EQ(0, metadata.dictionary_size()); + + // Object value bytes: + // - 0x82: basic type Object with reserved header bit 0x20 set + // - 0x00: zero fields + // - 0x00: first and final value offset + ASSERT_OK(VariantValueView::Validate(std::string("\x82\x00\x00", 3), metadata)); + // Array value bytes: + // - 0xe3: basic type Array with reserved header bits 0x38 set + // - 0x00: zero elements + // - 0x00: first and final value offset + ASSERT_OK(VariantValueView::Validate(std::string("\xe3\x00\x00", 3), metadata)); +} + +TEST(TestVariantEncoding, ValidateValue) { + VariantBuilder builder; + ASSERT_OK(builder.AppendVariantNull()); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + + ASSERT_OK(VariantValueView::Validate(std::string_view{*encoded.value}, metadata)); + std::string invalid_value(std::string_view{*encoded.value}); + invalid_value.push_back('\0'); + ASSERT_RAISES(Invalid, VariantValueView::Validate(invalid_value, metadata)); +} + +TEST(TestVariantEncoding, InvalidValue) { + VariantBuilder builder; + ASSERT_OK(builder.AppendVariantNull()); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + + // Header byte 0x54 encodes Primitive with primitive tag 21, which is unknown. + ASSERT_RAISES( + Invalid, VariantValueView::Make(std::string(1, static_cast(0x54)), metadata)); + // Short string payload is not valid UTF-8. + ASSERT_RAISES(Invalid, VariantValueView::Make(std::string("\x05\xff", 2), metadata)); + // Null primitive has trailing bytes after the encoded value. + ASSERT_RAISES(Invalid, VariantValueView::Make(std::string("\x00\x00", 2), metadata)); + // Object references field id 0, but the metadata dictionary is empty. + ASSERT_RAISES(Invalid, VariantValueView::Make( + std::string("\x02\x01\x00\x00\x01\x00", 6), metadata)); +} + +TEST(TestVariantEncoding, InvalidObject) { + VariantBuilder builder; + ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + ASSERT_OK(object.AppendVariantNull("a")); + ASSERT_OK(object.AppendVariantNull("b")); + ASSERT_OK(object.Finish()); + ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + + // Object value bytes: + // - 0x02: basic type Object, one-byte field ids, one-byte offsets, one-byte count + // - 0x02: two fields + // - 0x00 0x01: field ids for "a" and "b" + // - 0x00 0x00 0x01: duplicate field start offsets 0/0, final value size 1 + // - 0x00: encoded Variant null primitive value + const std::string duplicate_offsets("\x02\x02\x00\x01\x00\x00\x01\x00", 8); + ASSERT_RAISES(Invalid, VariantValueView::Validate(duplicate_offsets, metadata)); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/meson.build b/cpp/src/parquet/variant/meson.build new file mode 100644 index 000000000000..218ac39aec79 --- /dev/null +++ b/cpp/src/parquet/variant/meson.build @@ -0,0 +1,31 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +exc = executable( + 'parquet-variant-test', + sources: [ + 'builder_test.cc', + 'encoding_test.cc', + 'test_util_internal.cc', + 'type_test.cc', + 'validate_test.cc', + ], + dependencies: parquet_test_dep, +) +test('parquet-variant-test', exc) + +install_headers(['builder.h', 'encoding.h', 'validate.h'], subdir: 'parquet/variant') diff --git a/cpp/src/parquet/variant/test_util_internal.cc b/cpp/src/parquet/variant/test_util_internal.cc new file mode 100644 index 000000000000..048b20271202 --- /dev/null +++ b/cpp/src/parquet/variant/test_util_internal.cc @@ -0,0 +1,197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/test_util_internal.h" + +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/array/builder_binary.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/extension/uuid.h" +#include "arrow/extension_type.h" +#include "arrow/io/memory.h" +#include "arrow/table.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/type.h" +#include "parquet/arrow/reader.h" +#include "parquet/arrow/writer.h" +#include "parquet/exception.h" + +namespace parquet::variant::internal { + +Result MakeVariantValueView(const EncodedVariantValue& encoded) { + ARROW_ASSIGN_OR_RAISE(auto metadata, + VariantMetadataView::Make(std::string_view{*encoded.metadata})); + return VariantValueView::Make(std::string_view{*encoded.value}, metadata); +} + +std::shared_ptr<::arrow::Table> VariantTable( + const std::shared_ptr& variant_type, + const std::vector>& storage_children, + const FieldVector& storage_fields) { + PARQUET_ASSIGN_OR_THROW(auto storage, + ::arrow::StructArray::Make(storage_children, storage_fields)); + auto array = ::arrow::ExtensionType::WrapArray(variant_type, storage); + return ::arrow::Table::Make(::arrow::schema({::arrow::field("variant", variant_type)}), + {array}); +} + +Result> WriteVariantTable( + const std::shared_ptr<::arrow::Table>& table, + std::shared_ptr writer_properties, + std::shared_ptr arrow_properties) { + ARROW_ASSIGN_OR_RAISE(auto sink, ::arrow::io::BufferOutputStream::Create( + 1024, ::arrow::default_memory_pool())); + RETURN_NOT_OK(parquet::arrow::WriteTable(*table, ::arrow::default_memory_pool(), sink, + /*chunk_size=*/table->num_rows(), + std::move(writer_properties), + std::move(arrow_properties))); + return sink->Finish(); +} + +::arrow::Status WriteVariantRecordBatch( + const std::shared_ptr<::arrow::Table>& table, + std::shared_ptr arrow_properties) { + ARROW_ASSIGN_OR_RAISE(auto sink, ::arrow::io::BufferOutputStream::Create( + 1024, ::arrow::default_memory_pool())); + ARROW_ASSIGN_OR_RAISE(auto writer, + parquet::arrow::FileWriter::Open( + *table->schema(), ::arrow::default_memory_pool(), sink, + default_writer_properties(), std::move(arrow_properties))); + ARROW_ASSIGN_OR_RAISE(auto batch, table->CombineChunksToBatch()); + RETURN_NOT_OK(writer->WriteRecordBatch(*batch)); + return writer->Close(); +} + +std::optional ShreddedVariantTestingDir() { + const char* data = std::getenv("PARQUET_TEST_DATA"); + if (data == nullptr) { + return std::nullopt; + } + std::string path(data); + const auto pos = path.find_last_of("/\\"); + if (pos == std::string::npos) { + return std::nullopt; + } + return path.substr(0, pos) + "/shredded_variant"; +} + +Result> ReadVariantTestingTable(const std::string& path) { + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + + parquet::arrow::FileReaderBuilder builder; + RETURN_NOT_OK(builder.OpenFile(path, /*memory_map=*/false)); + builder.properties(reader_properties); + ARROW_ASSIGN_OR_RAISE(auto reader, builder.Build()); + return reader->ReadTable(); +} + +Result> EmptyVariantMetadata() { + VariantBuilder builder; + ARROW_RETURN_NOT_OK(builder.AppendVariantNull()); + ARROW_ASSIGN_OR_RAISE(auto encoded, builder.Finish()); + return encoded.metadata; +} + +Result Int8Variant(int8_t value) { + VariantBuilder builder; + ARROW_RETURN_NOT_OK(builder.AppendInt8(value)); + return builder.Finish(); +} + +std::shared_ptr BinaryArrayFromValues( + const std::vector>& values) { + ::arrow::BinaryBuilder builder; + for (const auto& value : values) { + if (value.has_value()) { + ARROW_EXPECT_OK(builder.Append(*value)); + } else { + ARROW_EXPECT_OK(builder.AppendNull()); + } + } + std::shared_ptr out; + ARROW_EXPECT_OK(builder.Finish(&out)); + return out; +} + +std::shared_ptr BinaryViewArrayFromValues( + const std::vector>& values) { + ::arrow::BinaryViewBuilder builder; + for (const auto& value : values) { + if (value.has_value()) { + ARROW_EXPECT_OK(builder.Append(*value)); + } else { + ARROW_EXPECT_OK(builder.AppendNull()); + } + } + std::shared_ptr out; + ARROW_EXPECT_OK(builder.Finish(&out)); + return out; +} + +std::shared_ptr Int64ArrayFromValues( + const std::vector>& values) { + ::arrow::Int64Builder builder; + for (const auto& value : values) { + if (value.has_value()) { + ARROW_EXPECT_OK(builder.Append(*value)); + } else { + ARROW_EXPECT_OK(builder.AppendNull()); + } + } + std::shared_ptr out; + ARROW_EXPECT_OK(builder.Finish(&out)); + return out; +} + +std::shared_ptr Int32ArrayFromValues(const std::vector& values) { + ::arrow::Int32Builder builder; + ARROW_EXPECT_OK(builder.AppendValues(values)); + std::shared_ptr out; + ARROW_EXPECT_OK(builder.Finish(&out)); + return out; +} + +std::shared_ptr StringArrayFromValues( + const std::vector>& values) { + ::arrow::StringBuilder builder; + for (const auto& value : values) { + if (value.has_value()) { + ARROW_EXPECT_OK(builder.Append(*value)); + } else { + ARROW_EXPECT_OK(builder.AppendNull()); + } + } + std::shared_ptr out; + ARROW_EXPECT_OK(builder.Finish(&out)); + return out; +} + +std::shared_ptr UuidArrayFromValues(const std::vector& values) { + ::arrow::FixedSizeBinaryBuilder builder(::arrow::fixed_size_binary(16)); + for (const auto& value : values) { + ARROW_EXPECT_OK(builder.Append(value)); + } + std::shared_ptr storage; + ARROW_EXPECT_OK(builder.Finish(&storage)); + return ::arrow::ExtensionType::WrapArray(::arrow::extension::uuid(), storage); +} + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/test_util_internal.h b/cpp/src/parquet/variant/test_util_internal.h new file mode 100644 index 000000000000..f36e92d01a4b --- /dev/null +++ b/cpp/src/parquet/variant/test_util_internal.h @@ -0,0 +1,83 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/result.h" +#include "arrow/type_fwd.h" +#include "parquet/properties.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/encoding.h" + +namespace parquet::variant::internal { + +using ::arrow::Array; +using ::arrow::Buffer; +using ::arrow::DataType; +using ::arrow::FieldVector; +using ::arrow::Result; + +Result MakeVariantValueView(const EncodedVariantValue& encoded); + +std::shared_ptr<::arrow::Table> VariantTable( + const std::shared_ptr& variant_type, + const std::vector>& storage_children, + const FieldVector& storage_fields); + +Result> WriteVariantTable( + const std::shared_ptr<::arrow::Table>& table, + std::shared_ptr writer_properties = default_writer_properties(), + std::shared_ptr arrow_properties = + default_arrow_writer_properties()); + +::arrow::Status WriteVariantRecordBatch( + const std::shared_ptr<::arrow::Table>& table, + std::shared_ptr arrow_properties = + default_arrow_writer_properties()); + +std::optional ShreddedVariantTestingDir(); + +Result> ReadVariantTestingTable(const std::string& path); + +Result> EmptyVariantMetadata(); + +Result Int8Variant(int8_t value); + +std::shared_ptr BinaryArrayFromValues( + const std::vector>& values); + +std::shared_ptr BinaryViewArrayFromValues( + const std::vector>& values); + +std::shared_ptr Int64ArrayFromValues( + const std::vector>& values); + +std::shared_ptr Int32ArrayFromValues(const std::vector& values); + +std::shared_ptr StringArrayFromValues( + const std::vector>& values); + +std::shared_ptr UuidArrayFromValues(const std::vector& values); + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/type_test.cc b/cpp/src/parquet/variant/type_test.cc new file mode 100644 index 000000000000..c6e93263ab83 --- /dev/null +++ b/cpp/src/parquet/variant/type_test.cc @@ -0,0 +1,167 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension/uuid.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/util/checked_cast.h" + +#include +#include + +namespace parquet::variant { + +using ::arrow::binary; +using ::arrow::binary_view; +using ::arrow::decimal128; +using ::arrow::decimal256; +using ::arrow::decimal32; +using ::arrow::decimal64; +using ::arrow::dictionary; +using ::arrow::duration; +using ::arrow::ExtensionType; +using ::arrow::field; +using ::arrow::fixed_size_binary; +using ::arrow::fixed_size_list; +using ::arrow::int32; +using ::arrow::int64; +using ::arrow::large_list_view; +using ::arrow::list; +using ::arrow::list_view; +using ::arrow::struct_; +using ::arrow::timestamp; +using ::arrow::TimeUnit; +using ::arrow::uint64; +using ::arrow::utf8; +using ::arrow::utf8_view; +using ::arrow::extension::kVariantExtensionName; +using ::arrow::extension::VariantExtensionType; + +TEST(TestVariantType, Storage) { + auto unshredded = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN(auto type, VariantExtensionType::Make(unshredded)); + ASSERT_EQ( + kVariantExtensionName, + ::arrow::internal::checked_cast(*type).extension_name()); + + auto shredded = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary()), field("typed_value", int64())}); + ASSERT_OK_AND_ASSIGN(auto shredded_type, VariantExtensionType::Make(shredded)); + auto variant_type = + ::arrow::internal::checked_pointer_cast(shredded_type); + ASSERT_EQ("metadata", variant_type->metadata()->name()); + ASSERT_EQ("value", variant_type->value()->name()); + ASSERT_EQ("typed_value", variant_type->typed_value()->name()); + + auto typed_only = struct_( + {field("metadata", binary(), /*nullable=*/false), field("typed_value", int64())}); + ASSERT_OK(VariantExtensionType::Make(typed_only)); + + auto flipped = + std::dynamic_pointer_cast(::arrow::extension::variant( + struct_({field("value", binary(), /*nullable=*/false), + field("metadata", binary(), /*nullable=*/false)}))); + ASSERT_EQ("metadata", flipped->metadata()->name()); + ASSERT_EQ("value", flipped->value()->name()); + + ASSERT_OK(VariantExtensionType::Make( + struct_({field("metadata", binary_view(), /*nullable=*/false), + field("value", binary_view(), /*nullable=*/false)}))); + + auto shredded_field_group = + struct_({field("value", binary()), field("typed_value", int64())}); + auto shredded_object = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", + struct_({field("a", shredded_field_group, /*nullable=*/false)}))}); + auto shredded_list = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", + list(field("element", shredded_field_group, /*nullable=*/false)))}); + ASSERT_OK(VariantExtensionType::Make(shredded_object)); + ASSERT_OK(VariantExtensionType::Make(shredded_list)); + + for (const auto& typed_value_type : + {binary_view(), utf8_view(), ::arrow::extension::uuid(), + list_view(field("element", shredded_field_group, /*nullable=*/false)), + large_list_view(field("element", shredded_field_group, /*nullable=*/false)), + fixed_size_list(field("element", shredded_field_group, /*nullable=*/false), + /*list_size=*/2)}) { + auto valid_shredded_type = + struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary()), field("typed_value", typed_value_type)}); + ASSERT_OK(VariantExtensionType::Make(valid_shredded_type)); + } +} + +TEST(TestVariantType, InvalidStorage) { + auto missing_value = struct_({field("metadata", binary(), /*nullable=*/false)}); + auto missing_metadata = struct_({field("value", binary(), /*nullable=*/false)}); + auto nullable_metadata = + struct_({field("metadata", binary()), field("value", binary(), false)}); + auto nullable_unshredded = + struct_({field("metadata", binary(), false), field("value", binary())}); + auto bad_value_type = + struct_({field("metadata", binary(), false), field("value", int32(), false)}); + auto extra = struct_({field("metadata", binary(), false), + field("value", binary(), false), field("extra", binary())}); + auto dictionary_typed_value = + struct_({field("metadata", binary(), false), field("value", binary()), + field("typed_value", dictionary(int32(), utf8()))}); + auto dictionary_value = struct_({field("metadata", binary(), false), + field("value", dictionary(int32(), binary()), false)}); + auto dictionary_metadata = + struct_({field("metadata", dictionary(int32(), binary()), false), + field("value", binary(), false)}); + auto required_nested_value = + struct_({field("metadata", binary(), false), field("value", binary()), + field("typed_value", + struct_({field("a", struct_({field("value", binary(), false)}), + /*nullable=*/false)}))}); + + ASSERT_RAISES(Invalid, VariantExtensionType::Make(missing_value)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(missing_metadata)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(nullable_metadata)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(nullable_unshredded)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(bad_value_type)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(extra)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(dictionary_typed_value)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(dictionary_value)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(dictionary_metadata)); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(required_nested_value)); + + std::array invalid_typed_value_types{ + uint64(), + duration(TimeUnit::MICRO), + timestamp(TimeUnit::MILLI), + struct_({}), + fixed_size_binary(/*byte_width=*/8), + decimal32(/*precision=*/8, /*scale=*/9), + decimal64(/*precision=*/16, /*scale=*/17), + decimal128(/*precision=*/32, /*scale=*/-1), + decimal256(/*precision=*/39, /*scale=*/0), + }; + for (const auto& typed_value_type : invalid_typed_value_types) { + auto invalid_shredded_type = + struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary()), field("typed_value", typed_value_type)}); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(std::move(invalid_shredded_type))); + } +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate.cc b/cpp/src/parquet/variant/validate.cc new file mode 100644 index 000000000000..8ff53d4e58be --- /dev/null +++ b/cpp/src/parquet/variant/validate.cc @@ -0,0 +1,473 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/validate.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/buffer.h" +#include "arrow/chunked_array.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension_type.h" +#include "arrow/result.h" +#include "arrow/type.h" +#include "arrow/util/bit_block_counter.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/bitmap_ops.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/logging_internal.h" +#include "parquet/variant/encoding.h" + +namespace parquet::variant { + +namespace { + +using ::arrow::Array; +using ::arrow::Buffer; +using ::arrow::DataType; +using ::arrow::ExtensionArray; +using ::arrow::ExtensionType; +using ::arrow::MemoryPool; +using ::arrow::Result; +using ::arrow::Status; +using ::arrow::extension::kVariantExtensionName; +using ::arrow::internal::checked_cast; + +enum class VariantShreddedTypeKind { + None, + Primitive, + Object, + Array, +}; + +struct VariantValidationPlan { + std::shared_ptr array; + std::vector children; +}; + +Result> ValuesArray(const Array& array) { + switch (array.type_id()) { + case ::arrow::Type::LIST_VIEW: + return checked_cast(array).values(); + case ::arrow::Type::LARGE_LIST_VIEW: + return checked_cast(array).values(); + case ::arrow::Type::LIST: + return checked_cast(array).values(); + case ::arrow::Type::LARGE_LIST: + return checked_cast(array).values(); + case ::arrow::Type::FIXED_SIZE_LIST: + return checked_cast(array).values(); + case ::arrow::Type::MAP: + return checked_cast(array).values(); + default: + return Status::Invalid("Expected list or map storage, got ", + array.type()->ToString()); + } +} + +Result> BuildVariantValidationPlan( + std::shared_ptr array) { + switch (array->type_id()) { + case ::arrow::Type::EXTENSION: { + const auto& ext_array = checked_cast(*array); + const auto& ext_type = checked_cast(*array->type()); + if (ext_type.extension_name() == kVariantExtensionName) { + return VariantValidationPlan{.array = std::move(array), .children = {}}; + } + return BuildVariantValidationPlan(ext_array.storage()); + } + case ::arrow::Type::STRUCT: { + const auto& struct_array = checked_cast(*array); + std::vector children; + for (auto field : struct_array.fields()) { + ARROW_ASSIGN_OR_RAISE(auto child_plan, + BuildVariantValidationPlan(std::move(field))); + if (child_plan.has_value()) { + children.push_back(std::move(*child_plan)); + } + } + if (children.empty()) { + return std::nullopt; + } + return VariantValidationPlan{.array = std::move(array), + .children = std::move(children)}; + } + case ::arrow::Type::LIST_VIEW: + case ::arrow::Type::LARGE_LIST_VIEW: + case ::arrow::Type::LIST: + case ::arrow::Type::LARGE_LIST: + case ::arrow::Type::FIXED_SIZE_LIST: + case ::arrow::Type::MAP: { + ARROW_ASSIGN_OR_RAISE(auto values, ValuesArray(*array)); + ARROW_ASSIGN_OR_RAISE(auto child_plan, + BuildVariantValidationPlan(std::move(values))); + if (!child_plan.has_value()) { + return std::nullopt; + } + std::vector children; + children.push_back(std::move(*child_plan)); + return VariantValidationPlan{.array = std::move(array), + .children = std::move(children)}; + } + default: { + std::vector children; + for (const auto& child_data : array->data()->child_data) { + if (child_data != nullptr) { + ARROW_ASSIGN_OR_RAISE(auto child_plan, BuildVariantValidationPlan( + ::arrow::MakeArray(child_data))); + if (child_plan.has_value()) { + children.push_back(std::move(*child_plan)); + } + } + } + if (children.empty()) { + return std::nullopt; + } + return VariantValidationPlan{.array = std::move(array), + .children = std::move(children)}; + } + } +} + +Result BinaryFieldView(const Array& array, int64_t row) { + switch (array.type_id()) { + case ::arrow::Type::BINARY: + return checked_cast(array).GetView(row); + case ::arrow::Type::LARGE_BINARY: + return checked_cast(array).GetView(row); + case ::arrow::Type::BINARY_VIEW: + return checked_cast(array).GetView(row); + default: + return Status::Invalid("Expected binary Variant field, got ", + array.type()->ToString()); + } +} + +VariantShreddedTypeKind ShreddedTypeKind(const DataType& type) { + switch (type.id()) { + case ::arrow::Type::STRUCT: + return VariantShreddedTypeKind::Object; + case ::arrow::Type::LIST: + case ::arrow::Type::LARGE_LIST: + case ::arrow::Type::LIST_VIEW: + case ::arrow::Type::LARGE_LIST_VIEW: + case ::arrow::Type::FIXED_SIZE_LIST: + return VariantShreddedTypeKind::Array; + default: + return VariantShreddedTypeKind::Primitive; + } +} + +Status ValidateVariantShredding( + const VariantMetadataView& metadata, std::string_view value, bool typed_value_present, + VariantShreddedTypeKind typed_value_kind, + std::span shredded_field_names = {}) { + if (typed_value_kind == VariantShreddedTypeKind::None) { + if (value.empty()) { + return Status::OK(); + } + return VariantValueView::Validate(value, metadata); + } + + if (typed_value_kind == VariantShreddedTypeKind::Object) { + for (const auto& name : shredded_field_names) { + if (!metadata.FindString(name).has_value()) { + return Status::Invalid("Invalid shredded Variant: shredded field '", name, + "' is not in metadata dictionary"); + } + } + } + + if (value.empty()) { + return Status::OK(); + } + ARROW_ASSIGN_OR_RAISE(auto value_view, VariantValueView::Make(value, metadata)); + + switch (typed_value_kind) { + case VariantShreddedTypeKind::Primitive: + if (typed_value_present) { + return Status::Invalid( + "Invalid shredded Variant: value and primitive typed_value are both " + "non-null"); + } + break; + case VariantShreddedTypeKind::Array: + if (typed_value_present) { + return Status::Invalid( + "Invalid shredded Variant: value and array typed_value are both " + "non-null"); + } + if (value_view.basic_type() == VariantBasicType::kArray) { + return Status::Invalid( + "Invalid shredded Variant: array value must be stored in typed_value"); + } + break; + case VariantShreddedTypeKind::Object: + default: + if (value_view.basic_type() == VariantBasicType::kObject) { + if (!typed_value_present) { + return Status::Invalid( + "Invalid shredded Variant: object value requires object typed_value"); + } + for (const auto& name : shredded_field_names) { + if (std::get(value_view.data()).ContainsField(name)) { + return Status::Invalid( + "Invalid shredded Variant: value object contains shredded field: ", name); + } + } + } else if (typed_value_present) { + return Status::Invalid( + "Invalid shredded Variant: partially shredded value must be an object"); + } + break; + } + return Status::OK(); +} + +template +std::pair ValueOffsetAndLength(const Array& array, int64_t row) { + const auto& typed_array = checked_cast(array); + return {typed_array.value_offset(row), typed_array.value_length(row)}; +} + +Result> ValuesRangeAt(const Array& array, int64_t row) { + switch (array.type_id()) { + case ::arrow::Type::LIST: + return ValueOffsetAndLength<::arrow::ListArray>(array, row); + case ::arrow::Type::LARGE_LIST: + return ValueOffsetAndLength<::arrow::LargeListArray>(array, row); + case ::arrow::Type::LIST_VIEW: + return ValueOffsetAndLength<::arrow::ListViewArray>(array, row); + case ::arrow::Type::LARGE_LIST_VIEW: + return ValueOffsetAndLength<::arrow::LargeListViewArray>(array, row); + case ::arrow::Type::MAP: + return ValueOffsetAndLength<::arrow::MapArray>(array, row); + case ::arrow::Type::FIXED_SIZE_LIST: + return ValueOffsetAndLength<::arrow::FixedSizeListArray>(array, row); + default: + return Status::Invalid("Expected list or map storage, got ", + array.type()->ToString()); + } +} + +Status ValidateShreddedVariantSlot(const VariantMetadataView& metadata, + const std::shared_ptr& value_array, + const std::shared_ptr& typed_array, int64_t row, + bool allow_missing, std::string_view path); + +Status ValidateShreddedVariantField(const VariantMetadataView& metadata, + const Array& field_array, int64_t row, + bool allow_missing, std::string_view path) { + if (field_array.type_id() != ::arrow::Type::STRUCT) { + return Status::Invalid("Invalid shredded Variant field at ", path, + ": expected struct storage, got ", + field_array.type()->ToString()); + } + if (field_array.IsNull(row)) { + return Status::Invalid("Invalid shredded Variant field at ", path, + ": field group must be required"); + } + const auto& field_struct = checked_cast(field_array); + auto value_array = field_struct.GetFieldByName("value"); + auto typed_array = field_struct.GetFieldByName("typed_value"); + return ValidateShreddedVariantSlot(metadata, value_array, typed_array, row, + allow_missing, path); +} + +Status ValidateShreddedVariantSlot(const VariantMetadataView& metadata, + const std::shared_ptr& value_array, + const std::shared_ptr& typed_array, int64_t row, + bool allow_missing, std::string_view path) { + std::string_view value; + if (value_array != nullptr && !value_array->IsNull(row)) { + ARROW_ASSIGN_OR_RAISE(value, BinaryFieldView(*value_array, row)); + if (value.empty()) { + ARROW_RETURN_NOT_OK(VariantValueView::Validate(value, metadata)); + } + } + + const bool typed_value_present = typed_array != nullptr && !typed_array->IsNull(row); + if (value.empty() && !typed_value_present) { + if (allow_missing) { + return Status::OK(); + } + return Status::Invalid("Invalid shredded Variant at ", path, + ": value and typed_value are both null"); + } + + const auto typed_kind = typed_array == nullptr ? VariantShreddedTypeKind::None + : ShreddedTypeKind(*typed_array->type()); + std::vector field_names; + if (typed_kind == VariantShreddedTypeKind::Object) { + field_names.reserve(typed_array->type()->num_fields()); + for (const auto& field : typed_array->type()->fields()) { + field_names.emplace_back(field->name()); + } + } + ARROW_RETURN_NOT_OK(ValidateVariantShredding(metadata, value, typed_value_present, + typed_kind, field_names)); + + if (!typed_value_present) { + return Status::OK(); + } + + if (typed_kind == VariantShreddedTypeKind::Object) { + const auto& typed_struct = checked_cast(*typed_array); + for (int i = 0; i < typed_struct.struct_type()->num_fields(); ++i) { + ARROW_RETURN_NOT_OK(ValidateShreddedVariantField( + metadata, *typed_struct.field(i), row, /*allow_missing=*/true, + typed_struct.struct_type()->field(i)->name())); + } + } else if (typed_kind == VariantShreddedTypeKind::Array) { + ARROW_ASSIGN_OR_RAISE(auto values, ValuesArray(*typed_array)); + ARROW_ASSIGN_OR_RAISE(auto range, ValuesRangeAt(*typed_array, row)); + const auto [offset, length] = range; + auto elements = values->Slice(offset, length); + for (int64_t i = 0; i < elements->length(); ++i) { + ARROW_RETURN_NOT_OK(ValidateShreddedVariantField(metadata, *elements, i, + /*allow_missing=*/false, path)); + } + } + return Status::OK(); +} + +template + requires requires(VisitVisible& visit_visible, int64_t row) { + ::arrow::ToStatus(visit_visible(row)); + } +Status VisitVisibleRows(const std::shared_ptr& valid_rows, const Array& array, + VisitVisible&& visit_visible) { + if (valid_rows == nullptr && !array.data()->MayHaveNulls()) { + for (int64_t row = 0; row < array.length(); ++row) { + ARROW_RETURN_NOT_OK(visit_visible(row)); + } + return Status::OK(); + } + return ::arrow::internal::VisitTwoBitBlocks( + valid_rows != nullptr ? valid_rows->data() : nullptr, /*left_offset=*/0, + array.data()->MayHaveNulls() ? array.null_bitmap_data() : nullptr, array.offset(), + array.length(), std::forward(visit_visible), + [] { return Status::OK(); }); +} + +Status ValidateVariantExtensionArray(const ExtensionArray& array, + const std::shared_ptr& valid_rows) { + const auto& storage = checked_cast(*array.storage()); + auto metadata_array = storage.GetFieldByName("metadata"); + auto value_array = storage.GetFieldByName("value"); + auto typed_array = storage.GetFieldByName("typed_value"); + + std::string_view last_metadata_bytes; + std::optional last_metadata; + return VisitVisibleRows(valid_rows, array, [&](int64_t row) { + if (metadata_array->IsNull(row)) { + return Status::Invalid("Invalid Variant extension storage: metadata is null"); + } + + ARROW_ASSIGN_OR_RAISE(auto metadata_value, BinaryFieldView(*metadata_array, row)); + if (!last_metadata.has_value() || last_metadata_bytes != metadata_value) { + ARROW_ASSIGN_OR_RAISE(last_metadata, VariantMetadataView::Make(metadata_value)); + last_metadata_bytes = metadata_value; + } + + if (typed_array == nullptr) { + if (value_array->IsNull(row)) { + return Status::Invalid( + "Invalid Variant extension storage: unshredded value is null"); + } + ARROW_ASSIGN_OR_RAISE(auto value, BinaryFieldView(*value_array, row)); + ARROW_RETURN_NOT_OK(VariantValueView::Validate(value, *last_metadata)); + } else { + ARROW_RETURN_NOT_OK( + ValidateShreddedVariantSlot(*last_metadata, value_array, typed_array, row, + /*allow_missing=*/false, "variant")); + } + return Status::OK(); + }); +} + +Status ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, + const std::shared_ptr& valid_rows) { + switch (plan.array->type_id()) { + case ::arrow::Type::EXTENSION: { + return ValidateVariantExtensionArray( + checked_cast(*plan.array), valid_rows); + } + case ::arrow::Type::STRUCT: { + const auto& struct_array = checked_cast(*plan.array); + std::shared_ptr child_valid_rows = valid_rows; + if (struct_array.data()->MayHaveNulls()) { + ARROW_ASSIGN_OR_RAISE( + child_valid_rows, + ::arrow::internal::OptionalBitmapAnd( + pool, valid_rows, /*left_offset=*/0, struct_array.null_bitmap(), + struct_array.offset(), struct_array.length(), /*out_offset=*/0)); + } + for (const auto& child : plan.children) { + ARROW_RETURN_NOT_OK(ValidateVariantPlan(child, pool, child_valid_rows)); + } + return Status::OK(); + } + case ::arrow::Type::LIST_VIEW: + case ::arrow::Type::LARGE_LIST_VIEW: + case ::arrow::Type::LIST: + case ::arrow::Type::LARGE_LIST: + case ::arrow::Type::MAP: + case ::arrow::Type::FIXED_SIZE_LIST: { + DCHECK_EQ(plan.children.size(), 1); + ARROW_ASSIGN_OR_RAISE(auto values, ValuesArray(*plan.array)); + ARROW_ASSIGN_OR_RAISE(auto values_valid_rows, + ::arrow::AllocateEmptyBitmap(values->length(), pool)); + ARROW_RETURN_NOT_OK( + VisitVisibleRows(valid_rows, *plan.array, [&](int64_t row) -> Status { + ARROW_ASSIGN_OR_RAISE(auto range, ValuesRangeAt(*plan.array, row)); + const auto [offset, length] = range; + ::arrow::bit_util::SetBitsTo(values_valid_rows->mutable_data(), offset, + length, true); + return Status::OK(); + })); + return ValidateVariantPlan(plan.children[0], pool, values_valid_rows); + } + default: { + for (const auto& child : plan.children) { + ARROW_RETURN_NOT_OK(ValidateVariantPlan(child, pool, valid_rows)); + } + return Status::OK(); + } + } +} + +} // namespace + +Status ValidateVariants(const ::arrow::ChunkedArray& data, MemoryPool* pool) { + for (const auto& chunk : data.chunks()) { + ARROW_ASSIGN_OR_RAISE(auto plan, BuildVariantValidationPlan(chunk)); + if (plan.has_value()) { + ARROW_RETURN_NOT_OK(ValidateVariantPlan(*plan, pool, nullptr)); + } + } + return Status::OK(); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate.h b/cpp/src/parquet/variant/validate.h new file mode 100644 index 000000000000..2afc808b9182 --- /dev/null +++ b/cpp/src/parquet/variant/validate.h @@ -0,0 +1,36 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include "arrow/status.h" +#include "parquet/platform.h" + +namespace arrow { + +class ChunkedArray; +class MemoryPool; + +} // namespace arrow + +namespace parquet::variant { + +PARQUET_EXPORT +::arrow::Status ValidateVariants(const ::arrow::ChunkedArray& data, + ::arrow::MemoryPool* pool); + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate_test.cc b/cpp/src/parquet/variant/validate_test.cc new file mode 100644 index 000000000000..425cf79df137 --- /dev/null +++ b/cpp/src/parquet/variant/validate_test.cc @@ -0,0 +1,190 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/validate.h" + +#include +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/chunked_array.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension_type.h" +#include "arrow/io/memory.h" +#include "arrow/table.h" +#include "arrow/testing/extension_type.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/type.h" +#include "parquet/arrow/reader.h" +#include "parquet/variant/test_util_internal.h" + +namespace parquet::variant { + +using ::arrow::binary; +using ::arrow::field; +using ::arrow::struct_; +using internal::BinaryArrayFromValues; +using internal::Int32ArrayFromValues; +using internal::Int8Variant; +using internal::ReadVariantTestingTable; +using internal::ShreddedVariantTestingDir; +using internal::VariantTable; +using internal::WriteVariantTable; + +TEST(TestVariantValidate, ListView) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues( + {std::string_view{*encoded.value}, std::string_view("\xff", 1)}); + ASSERT_OK_AND_ASSIGN( + auto storage, + ::arrow::StructArray::Make({metadata_array, value_array}, storage_type->fields())); + auto variant_array = ::arrow::ExtensionType::WrapArray(variant_type, storage); + + ASSERT_OK_AND_ASSIGN(auto valid_list, + ::arrow::ListViewArray::FromArrays( + *Int32ArrayFromValues({0}), *Int32ArrayFromValues({1}), + *variant_array, ::arrow::default_memory_pool())); + ::arrow::ChunkedArray valid_data{valid_list}; + ASSERT_OK(ValidateVariants(valid_data, ::arrow::default_memory_pool())); + + ASSERT_OK_AND_ASSIGN(auto invalid_list, + ::arrow::ListViewArray::FromArrays( + *Int32ArrayFromValues({1}), *Int32ArrayFromValues({1}), + *variant_array, ::arrow::default_memory_pool())); + ::arrow::ChunkedArray invalid_data{invalid_list}; + ASSERT_RAISES(Invalid, ValidateVariants(invalid_data, ::arrow::default_memory_pool())); +} + +TEST(TestVariantValidate, ParquetTestingShredded) { + auto maybe_dir = ShreddedVariantTestingDir(); + if (!maybe_dir.has_value()) { + GTEST_SKIP() << "PARQUET_TEST_DATA not set"; + } + if (!std::filesystem::exists(*maybe_dir)) { + GTEST_SKIP() << *maybe_dir << " does not exist"; + } + + auto registered_storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + ::arrow::ExtensionTypeGuard guard(::arrow::extension::variant(registered_storage_type)); + + struct ShreddedCase { + std::string file_name; + bool has_top_level_value; + }; + const std::vector cases = { + {.file_name = "case-041.parquet", .has_top_level_value = false}, + {.file_name = "case-088.parquet", .has_top_level_value = true}, + {.file_name = "case-131.parquet", .has_top_level_value = false}, + {.file_name = "case-132.parquet", .has_top_level_value = true}, + {.file_name = "case-138.parquet", .has_top_level_value = false}, + }; + + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.file_name); + ASSERT_OK_AND_ASSIGN(auto table, + ReadVariantTestingTable(*maybe_dir + "/" + test_case.file_name)); + ASSERT_OK(table->ValidateFull()); + + auto field = table->schema()->GetFieldByName("var"); + ASSERT_NE(nullptr, field); + auto variant_type = + std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>( + field->type()); + ASSERT_NE(nullptr, variant_type); + ASSERT_EQ(test_case.has_top_level_value, variant_type->value() != nullptr); + + auto column = table->GetColumnByName("var"); + ASSERT_NE(nullptr, column); + ASSERT_OK(ValidateVariants(*column, ::arrow::default_memory_pool())); + } +} + +TEST(TestVariantValidate, DictionaryMetadata) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues( + {std::string_view{*encoded.value}, std::string_view{*encoded.value}}); + auto table = + VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); + + ASSERT_OK_AND_ASSIGN( + auto buffer, + WriteVariantTable(table, WriterProperties::Builder().enable_dictionary()->build())); + + auto buffer_reader = std::make_shared<::arrow::io::BufferReader>(buffer); + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + ::arrow::ExtensionTypeGuard guard(::arrow::extension::variant(storage_type)); + parquet::arrow::FileReaderBuilder builder; + ASSERT_OK(builder.Open(buffer_reader)); + builder.properties(reader_properties); + ASSERT_OK_AND_ASSIGN(auto reader, builder.Build()); + + ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); + auto column = read_table->GetColumnByName("variant"); + ASSERT_NE(nullptr, column); + ASSERT_OK(ValidateVariants(*column, ::arrow::default_memory_pool())); +} + +TEST(TestVariantValidate, ReadDictionaryOption) { + ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + + auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + auto variant_type = ::arrow::extension::variant(storage_type); + auto metadata_array = BinaryArrayFromValues( + {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); + auto value_array = BinaryArrayFromValues( + {std::string_view{*encoded.value}, std::string_view{*encoded.value}}); + auto table = + VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); + + ASSERT_OK_AND_ASSIGN(auto buffer, WriteVariantTable(table)); + + auto buffer_reader = std::make_shared<::arrow::io::BufferReader>(buffer); + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + reader_properties.set_read_dictionary(0, true); + reader_properties.set_read_dictionary(1, true); + ::arrow::ExtensionTypeGuard guard(::arrow::extension::variant(storage_type)); + parquet::arrow::FileReaderBuilder builder; + ASSERT_OK(builder.Open(buffer_reader)); + builder.properties(reader_properties); + ASSERT_OK_AND_ASSIGN(auto reader, builder.Build()); + + ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); + auto column = read_table->GetColumnByName("variant"); + ASSERT_NE(nullptr, column); + ASSERT_OK(ValidateVariants(*column, ::arrow::default_memory_pool())); +} + +} // namespace parquet::variant From bec504b842da91091db2bf7ec8b4cd3af5b62e39 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Fri, 26 Jun 2026 20:07:26 +0800 Subject: [PATCH 02/10] fix cpplint ci --- cpp/src/parquet/variant/meson.build | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/src/parquet/variant/meson.build b/cpp/src/parquet/variant/meson.build index 218ac39aec79..7a3f74014ad5 100644 --- a/cpp/src/parquet/variant/meson.build +++ b/cpp/src/parquet/variant/meson.build @@ -28,4 +28,7 @@ exc = executable( ) test('parquet-variant-test', exc) -install_headers(['builder.h', 'encoding.h', 'validate.h'], subdir: 'parquet/variant') +install_headers( + ['builder.h', 'encoding.h', 'validate.h'], + subdir: 'parquet/variant', +) From 945403ca9052086ee78d58c1f9e97575fdf863f7 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Thu, 2 Jul 2026 20:19:15 +0800 Subject: [PATCH 03/10] change status/result to exception --- cpp/src/parquet/arrow/variant_test.cc | 28 +- cpp/src/parquet/arrow/writer.cc | 4 +- cpp/src/parquet/variant/builder.cc | 694 +++++++++--------- cpp/src/parquet/variant/builder.h | 220 +++--- cpp/src/parquet/variant/builder_test.cc | 209 +++--- cpp/src/parquet/variant/encoding.cc | 240 +++--- cpp/src/parquet/variant/encoding.h | 13 +- cpp/src/parquet/variant/encoding_internal.h | 16 +- cpp/src/parquet/variant/encoding_test.cc | 136 ++-- cpp/src/parquet/variant/test_util_internal.cc | 17 +- cpp/src/parquet/variant/test_util_internal.h | 6 +- cpp/src/parquet/variant/validate.cc | 221 +++--- cpp/src/parquet/variant/validate.h | 4 +- cpp/src/parquet/variant/validate_test.cc | 18 +- 14 files changed, 884 insertions(+), 942 deletions(-) diff --git a/cpp/src/parquet/arrow/variant_test.cc b/cpp/src/parquet/arrow/variant_test.cc index 62e68f3e3ad5..95b1ebad5c26 100644 --- a/cpp/src/parquet/arrow/variant_test.cc +++ b/cpp/src/parquet/arrow/variant_test.cc @@ -51,7 +51,7 @@ using variant::internal::WriteVariantRecordBatch; using variant::internal::WriteVariantTable; TEST(TestVariantExtensionType, WriterValidatesUnshreddedVariantBytes) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -73,7 +73,7 @@ TEST(TestVariantExtensionType, WriterValidatesUnshreddedVariantBytes) { } TEST(TestVariantExtensionType, WriteRecordBatchValidatesVariantBytes) { - ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto metadata = EmptyVariantMetadata(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); auto variant_type = ::arrow::extension::variant(storage_type); @@ -91,7 +91,7 @@ TEST(TestVariantExtensionType, WriteRecordBatchValidatesVariantBytes) { } TEST(TestVariantExtensionType, WriteRecordBatchValidatesBatch) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -120,7 +120,7 @@ TEST(TestVariantExtensionType, WriteRecordBatchValidatesBatch) { } TEST(TestVariantExtensionType, WriterValidatesBinaryViewVariantBytes) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", ::arrow::binary_view(), /*nullable=*/false), @@ -134,7 +134,7 @@ TEST(TestVariantExtensionType, WriterValidatesBinaryViewVariantBytes) { } TEST(TestVariantExtensionType, WriterSkipsNullParents) { - ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto metadata = EmptyVariantMetadata(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); auto variant_type = ::arrow::extension::variant(storage_type); @@ -167,7 +167,7 @@ TEST(TestVariantExtensionType, WriterSkipsNullParents) { } TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), @@ -197,7 +197,7 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { } TEST(TestVariantExtensionType, WriterValidatesShreddedWithoutValue) { - ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto metadata = EmptyVariantMetadata(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("typed_value", ::arrow::int64())}); ASSERT_OK_AND_ASSIGN(auto variant_type, @@ -211,7 +211,7 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedWithoutValue) { } TEST(TestVariantExtensionType, ReadsDictionaryEncodedMetadata) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -256,7 +256,7 @@ TEST(TestVariantExtensionType, ReadsDictionaryEncodedMetadata) { } TEST(TestVariantExtensionType, ReadsWithDictionaryOption) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -296,7 +296,7 @@ TEST(TestVariantExtensionType, ReadsWithDictionaryOption) { } TEST(TestVariantExtensionType, WriterWritesUuid) { - ASSERT_OK_AND_ASSIGN(auto metadata, EmptyVariantMetadata()); + auto metadata = EmptyVariantMetadata(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), field("typed_value", ::arrow::extension::uuid())}); @@ -312,10 +312,10 @@ TEST(TestVariantExtensionType, WriterWritesUuid) { TEST(TestVariantExtensionType, WriterValidatesShreddedObjectConflicts) { variant::VariantBuilder object_builder; - ASSERT_OK_AND_ASSIGN(auto object, object_builder.StartObject()); - ASSERT_OK(object.AppendShortString("event_type", "login")); - ASSERT_OK(object.Finish()); - ASSERT_OK_AND_ASSIGN(auto encoded, object_builder.Finish()); + auto object = object_builder.StartObject(); + object.AppendShortString("event_type", "login"); + object.Finish(); + auto encoded = object_builder.Finish(); auto field_group_type = struct_({field("value", binary()), field("typed_value", ::arrow::utf8())}); diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index 024f5f6e7699..47f50062f109 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -382,7 +382,7 @@ class FileWriterImpl : public FileWriter { int64_t size) override { RETURN_NOT_OK(CheckClosed()); if (arrow_properties_->variant_validation_enabled()) { - RETURN_NOT_OK(variant::ValidateVariants(*data->Slice(offset, size), pool_)); + PARQUET_CATCH_NOT_OK(variant::ValidateVariants(*data->Slice(offset, size), pool_)); } if (arrow_properties_->engine_version() == ArrowWriterProperties::V2 || arrow_properties_->engine_version() == ArrowWriterProperties::V1) { @@ -475,7 +475,7 @@ class FileWriterImpl : public FileWriter { for (int i = 0; i < batch.num_columns(); i++) { ChunkedArray chunked_array{batch.column(i)}; if (arrow_properties_->variant_validation_enabled()) { - RETURN_NOT_OK( + PARQUET_CATCH_NOT_OK( variant::ValidateVariants(*chunked_array.Slice(offset, size), pool_)); } ARROW_ASSIGN_OR_RAISE( diff --git a/cpp/src/parquet/variant/builder.cc b/cpp/src/parquet/variant/builder.cc index bd8540b368ee..e167c4f7cbc0 100644 --- a/cpp/src/parquet/variant/builder.cc +++ b/cpp/src/parquet/variant/builder.cc @@ -33,6 +33,7 @@ #include "arrow/util/checked_cast.h" #include "arrow/util/endian.h" #include "arrow/util/logging_internal.h" +#include "parquet/exception.h" #include "parquet/variant/encoding.h" #include "parquet/variant/encoding_internal.h" @@ -54,17 +55,17 @@ namespace bit_util = ::arrow::bit_util; namespace { -Status AppendLittleEndian(BufferBuilder& out, uint32_t value, uint8_t width) { +void AppendLittleEndian(BufferBuilder& out, uint32_t value, uint8_t width) { DCHECK_LE(width, sizeof(uint32_t)); const auto little_endian = bit_util::ToLittleEndian(value); - return out.Append(&little_endian, width); + PARQUET_THROW_NOT_OK(out.Append(&little_endian, width)); } template requires(std::is_arithmetic_v) -Status AppendFixedLittleEndian(BufferBuilder& out, T value) { +void AppendFixedLittleEndian(BufferBuilder& out, T value) { const auto little_endian = bit_util::ToLittleEndian(value); - return out.Append(&little_endian, sizeof(T)); + PARQUET_THROW_NOT_OK(out.Append(&little_endian, sizeof(T))); } void AppendLittleEndianToString(std::string& out, uint32_t value, uint8_t width) { @@ -73,15 +74,14 @@ void AppendLittleEndianToString(std::string& out, uint32_t value, uint8_t width) out.append(reinterpret_cast(&little_endian), width); } -Status InsertBytes(BufferBuilder& out, int64_t offset, std::string_view bytes) { +void InsertBytes(BufferBuilder& out, int64_t offset, std::string_view bytes) { const int64_t old_size = out.length(); const auto insert_size = bytes.size(); - ARROW_RETURN_NOT_OK(out.Reserve(insert_size)); + PARQUET_THROW_NOT_OK(out.Reserve(insert_size)); uint8_t* data = out.mutable_data(); std::memmove(data + offset + insert_size, data + offset, old_size - offset); std::memcpy(data + offset, bytes.data(), insert_size); out.UnsafeAdvance(insert_size); - return Status::OK(); } uint8_t WidthForValue(uint64_t value) { @@ -101,23 +101,22 @@ class VariantMetadataBuilder { public: explicit VariantMetadataBuilder(MemoryPool* pool) : pool_(pool) {} - Status Reserve(int64_t capacity) { + void Reserve(int64_t capacity) { if (capacity < 0) { - return Status::Invalid("Variant metadata capacity must be non-negative"); + throw ParquetException("Variant metadata capacity must be non-negative"); } field_ids_.reserve(capacity); - return Status::OK(); } - Result Upsert(std::string_view name) { + uint32_t Upsert(std::string_view name) { auto it = field_ids_.find(name); if (it != field_ids_.end()) { return it->second; } if (field_names_.size() >= std::numeric_limits::max()) { - return Status::Invalid("Variant metadata dictionary is too large"); + throw ParquetException("Variant metadata dictionary is too large"); } - ARROW_RETURN_NOT_OK(internal::ValidateUtf8(name, "metadata dictionary string")); + internal::ValidateUtf8(name, "metadata dictionary string"); const auto id = field_names_.size(); if (field_names_.empty()) { @@ -143,20 +142,20 @@ class VariantMetadataBuilder { RebuildIndex(); } - Result> Finish() const { + std::shared_ptr Finish() const { uint64_t bytes_size = 0; for (const auto& string : field_names_) { bytes_size += string.size(); } if (field_names_.size() > std::numeric_limits::max() || bytes_size > std::numeric_limits::max()) { - return Status::Invalid("Variant metadata dictionary is too large"); + throw ParquetException("Variant metadata dictionary is too large"); } const uint8_t offset_size = WidthForValue(std::max(field_names_.size(), bytes_size)); BufferBuilder out(pool_); - ARROW_RETURN_NOT_OK(out.Reserve( + PARQUET_THROW_NOT_OK(out.Reserve( 1 + offset_size + (field_names_.size() + 1) * offset_size + bytes_size)); const bool sorted_strings = !field_names_.empty() && is_sorted_; @@ -164,20 +163,21 @@ class VariantMetadataBuilder { static_cast(internal::kVariantVersion | (sorted_strings ? internal::kMetadataSortedStringsMask : 0) | ((offset_size - 1) << 6)); - ARROW_RETURN_NOT_OK(out.Append(&header, sizeof(header))); - ARROW_RETURN_NOT_OK( - AppendLittleEndian(out, static_cast(field_names_.size()), offset_size)); + PARQUET_THROW_NOT_OK(out.Append(&header, sizeof(header))); + AppendLittleEndian(out, static_cast(field_names_.size()), offset_size); uint32_t offset = 0; - ARROW_RETURN_NOT_OK(AppendLittleEndian(out, offset, offset_size)); + AppendLittleEndian(out, offset, offset_size); for (const auto& string : field_names_) { offset += static_cast(string.size()); - ARROW_RETURN_NOT_OK(AppendLittleEndian(out, offset, offset_size)); + AppendLittleEndian(out, offset, offset_size); } for (const auto& string : field_names_) { - ARROW_RETURN_NOT_OK(out.Append(string)); + PARQUET_THROW_NOT_OK(out.Append(string)); } - return out.Finish(); + std::shared_ptr buffer; + PARQUET_ASSIGN_OR_THROW(buffer, out.Finish()); + return buffer; } private: @@ -207,90 +207,90 @@ class VariantValueWriter { template requires internal::HeaderOnlyVariantPrimitive - Status Append() { - return AppendPrimitiveHeader(); + void Append() { + AppendPrimitiveHeader(); } template requires internal::FixedVariantPrimitive - Status Append(typename internal::VariantFixedPrimitiveTraits::CType value) { + void Append(typename internal::VariantFixedPrimitiveTraits::CType value) { using CType = typename internal::VariantFixedPrimitiveTraits::CType; - ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); + AppendPrimitiveHeader(); if constexpr (sizeof(CType) == 1) { const auto byte = static_cast(value); - return out_.Append(&byte, sizeof(byte)); + PARQUET_THROW_NOT_OK(out_.Append(&byte, sizeof(byte))); } else { - return AppendFixedLittleEndian(out_, value); + AppendFixedLittleEndian(out_, value); } } template requires internal::DecimalVariantPrimitive - Status Append( + void Append( typename internal::VariantDecimalPrimitiveTraits::CType unscaled_value, uint8_t scale) { - ARROW_RETURN_NOT_OK(internal::ValidateDecimalScale(scale)); - ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); - ARROW_RETURN_NOT_OK(out_.Append(&scale, sizeof(scale))); - return AppendFixedLittleEndian(out_, unscaled_value); + internal::ValidateDecimalScale(scale); + AppendPrimitiveHeader(); + PARQUET_THROW_NOT_OK(out_.Append(&scale, sizeof(scale))); + AppendFixedLittleEndian(out_, unscaled_value); } template requires internal::Decimal16VariantPrimitive - Status Append(std::string_view little_endian_unscaled_value, uint8_t scale) { + void Append(std::string_view little_endian_unscaled_value, uint8_t scale) { if (little_endian_unscaled_value.size() != 16) { - return Status::Invalid("Variant Decimal16 values must be 16 bytes"); + throw ParquetException("Variant Decimal16 values must be 16 bytes"); } - ARROW_RETURN_NOT_OK(internal::ValidateDecimalScale(scale)); - ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); - ARROW_RETURN_NOT_OK(out_.Append(&scale, sizeof(scale))); - return out_.Append(little_endian_unscaled_value); + internal::ValidateDecimalScale(scale); + AppendPrimitiveHeader(); + PARQUET_THROW_NOT_OK(out_.Append(&scale, sizeof(scale))); + PARQUET_THROW_NOT_OK(out_.Append(little_endian_unscaled_value)); } template requires internal::LengthPrefixedVariantPrimitive - Status Append(std::string_view value) { + void Append(std::string_view value) { if (value.size() > std::numeric_limits::max()) { - return Status::Invalid("Variant ", + throw ParquetException("Variant ", type == VariantPrimitiveType::kBinary ? "binary" : "string", " value is too large"); } if constexpr (type == VariantPrimitiveType::kString) { - ARROW_RETURN_NOT_OK(internal::ValidateUtf8(value, "primitive string value")); + internal::ValidateUtf8(value, "primitive string value"); } - ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); - ARROW_RETURN_NOT_OK(AppendLittleEndian(out_, static_cast(value.size()), 4)); - return out_.Append(value); + AppendPrimitiveHeader(); + AppendLittleEndian(out_, static_cast(value.size()), 4); + PARQUET_THROW_NOT_OK(out_.Append(value)); } template requires internal::UuidVariantPrimitive - Status Append(std::string_view big_endian_bytes) { + void Append(std::string_view big_endian_bytes) { if (big_endian_bytes.size() != 16) { - return Status::Invalid("Variant UUID values must be 16 bytes"); + throw ParquetException("Variant UUID values must be 16 bytes"); } - ARROW_RETURN_NOT_OK(AppendPrimitiveHeader()); - return out_.Append(big_endian_bytes); + AppendPrimitiveHeader(); + PARQUET_THROW_NOT_OK(out_.Append(big_endian_bytes)); } - Status AppendShortString(std::string_view value) { + void AppendShortString(std::string_view value) { if (value.size() >= 64) { - return Status::Invalid("Variant short string value must be shorter than 64 bytes"); + throw ParquetException("Variant short string value must be shorter than 64 bytes"); } - ARROW_RETURN_NOT_OK(internal::ValidateUtf8(value, "short string value")); + internal::ValidateUtf8(value, "short string value"); const auto header = static_cast( (value.size() << 2) | static_cast(VariantBasicType::kShortString)); - ARROW_RETURN_NOT_OK(out_.Append(&header, sizeof(header))); - return out_.Append(value); + PARQUET_THROW_NOT_OK(out_.Append(&header, sizeof(header))); + PARQUET_THROW_NOT_OK(out_.Append(value)); } private: template - Status AppendPrimitiveHeader() { + void AppendPrimitiveHeader() { const auto header = static_cast((static_cast(type) << 2) | static_cast(VariantBasicType::kPrimitive)); - return out_.Append(&header, sizeof(header)); + PARQUET_THROW_NOT_OK(out_.Append(&header, sizeof(header))); } BufferBuilder& out_; @@ -327,26 +327,24 @@ struct VariantBuildState { bool root_has_value = false; }; -Status CheckRootWritable(const VariantBuildState& state) { +void CheckRootWritable(const VariantBuildState& state) { if (!state.frames.empty()) { - return Status::Invalid("VariantBuilder has an active container"); + throw ParquetException("VariantBuilder has an active container"); } if (state.root_has_value) { - return Status::Invalid("VariantBuilder already has a root value"); + throw ParquetException("VariantBuilder already has a root value"); } - return Status::OK(); } -Status CheckTopFrame(const VariantBuildState& state, size_t frame_index, - VariantContainerKind kind) { +void CheckTopFrame(const VariantBuildState& state, size_t frame_index, + VariantContainerKind kind) { if (state.frames.empty() || frame_index + 1 != state.frames.size()) { - return Status::Invalid("Variant nested builder is not the active container"); + throw ParquetException("Variant nested builder is not the active container"); } const auto& frame = state.frames.back(); if (frame.finished || frame.kind != kind) { - return Status::Invalid("Variant nested builder has invalid state"); + throw ParquetException("Variant nested builder has invalid state"); } - return Status::OK(); } void TruncateFrameEntries(VariantBuildFrame& frame, size_t entry_count) { @@ -379,11 +377,10 @@ void RollbackIfActive(Impl* impl) { } } -Result BuildObjectHeader(const VariantBuildState& state, - const VariantBuildFrame& frame, - uint32_t values_size) { +std::string BuildObjectHeader(const VariantBuildState& state, + const VariantBuildFrame& frame, uint32_t values_size) { if (frame.fields.size() > std::numeric_limits::max()) { - return Status::Invalid("Variant object has too many fields"); + throw ParquetException("Variant object has too many fields"); } std::vector fields = frame.fields; @@ -417,10 +414,9 @@ Result BuildObjectHeader(const VariantBuildState& state, return out; } -Result BuildListHeader(const VariantBuildFrame& frame, - uint32_t values_size) { +std::string BuildListHeader(const VariantBuildFrame& frame, uint32_t values_size) { if (frame.offsets.size() > std::numeric_limits::max()) { - return Status::Invalid("Variant array has too many elements"); + throw ParquetException("Variant array has too many elements"); } const uint8_t offset_size = WidthForValue(values_size); @@ -439,141 +435,140 @@ Result BuildListHeader(const VariantBuildFrame& frame, return out; } -using RootFinishCallback = std::function; +using RootFinishCallback = std::function; -Status FinishFrame(const std::shared_ptr& state, size_t frame_index, - VariantContainerKind kind, const RootFinishCallback& callback) { - ARROW_RETURN_NOT_OK(CheckTopFrame(*state, frame_index, kind)); +void FinishFrame(const std::shared_ptr& state, size_t frame_index, + VariantContainerKind kind, const RootFinishCallback& callback) { + CheckTopFrame(*state, frame_index, kind); auto& frame = state->frames.back(); const auto values_size = state->value.length() - frame.value_start; DCHECK_GE(values_size, 0); if (values_size > std::numeric_limits::max()) { - return Status::Invalid("Variant container values are too large"); + throw ParquetException("Variant container values are too large"); } - ARROW_ASSIGN_OR_RAISE( - auto header, - kind == VariantContainerKind::Object - ? BuildObjectHeader(*state, frame, static_cast(values_size)) - : BuildListHeader(frame, static_cast(values_size))); - ARROW_RETURN_NOT_OK(InsertBytes(state->value, frame.value_start, header)); + auto header = kind == VariantContainerKind::Object + ? BuildObjectHeader(*state, frame, static_cast(values_size)) + : BuildListHeader(frame, static_cast(values_size)); + InsertBytes(state->value, frame.value_start, header); const bool is_root = !frame.has_parent; frame.finished = true; state->frames.pop_back(); if (!is_root) { - return Status::OK(); + return; } state->root_has_value = true; if (!callback) { - return Status::OK(); + return; } - ARROW_ASSIGN_OR_RAISE(auto metadata, state->metadata.Finish()); - ARROW_ASSIGN_OR_RAISE(auto value, state->value.Finish()); - return callback({.metadata = std::move(metadata), .value = std::move(value)}); + auto metadata = state->metadata.Finish(); + std::shared_ptr value; + PARQUET_ASSIGN_OR_THROW(value, state->value.Finish()); + callback({.metadata = std::move(metadata), .value = std::move(value)}); } template -Status AppendRootPrimitiveWith(const std::shared_ptr& state, - Write&& write) { - ARROW_RETURN_NOT_OK(CheckRootWritable(*state)); +void AppendRootPrimitiveWith(const std::shared_ptr& state, + Write&& write) { + CheckRootWritable(*state); const auto value_size = state->value.length(); VariantValueWriter writer(state->value); - const Status status = std::invoke(std::forward(write), writer); - if (!status.ok()) { + try { + std::invoke(std::forward(write), writer); + } catch (...) { state->value.Rewind(value_size); - return status; + throw; } state->root_has_value = true; - return Status::OK(); } template -Status AppendRootPrimitive(const std::shared_ptr& state, - Args&&... args) { - return AppendRootPrimitiveWith(state, [&](VariantValueWriter& writer) { - return writer.template Append(std::forward(args)...); +void AppendRootPrimitive(const std::shared_ptr& state, + Args&&... args) { + AppendRootPrimitiveWith(state, [&](VariantValueWriter& writer) { + writer.template Append(std::forward(args)...); }); } template -Status AppendObjectPrimitiveWith(const std::shared_ptr& state, - size_t frame_index, std::string_view field_name, - Write&& write) { - ARROW_RETURN_NOT_OK(CheckTopFrame(*state, frame_index, VariantContainerKind::Object)); +void AppendObjectPrimitiveWith(const std::shared_ptr& state, + size_t frame_index, std::string_view field_name, + Write&& write) { + CheckTopFrame(*state, frame_index, VariantContainerKind::Object); auto& frame = state->frames.back(); const auto metadata_size = state->metadata.size(); const auto value_size = state->value.length(); const auto field_count = frame.fields.size(); - ARROW_ASSIGN_OR_RAISE(auto field_id, state->metadata.Upsert(field_name)); + auto field_id = state->metadata.Upsert(field_name); if (!frame.object_field_ids.insert(field_id).second) { state->metadata.Truncate(metadata_size); - return Status::Invalid("Duplicate Variant object field: ", field_name); + throw ParquetException("Duplicate Variant object field: ", field_name); } const auto offset = state->value.length() - frame.value_start; DCHECK_GE(offset, 0); if (offset > std::numeric_limits::max()) { state->metadata.Truncate(metadata_size); TruncateFrameEntries(frame, field_count); - return Status::Invalid("Variant object values are too large"); + throw ParquetException("Variant object values are too large"); } frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, .offset = static_cast(offset)}); VariantValueWriter writer(state->value); - const Status status = std::invoke(std::forward(write), writer); - if (!status.ok()) { + try { + std::invoke(std::forward(write), writer); + } catch (...) { state->value.Rewind(value_size); state->metadata.Truncate(metadata_size); TruncateFrameEntries(frame, field_count); - return status; + throw; } - return Status::OK(); } template -Status AppendObjectPrimitive(const std::shared_ptr& state, - size_t frame_index, std::string_view field_name, - Args&&... args) { - return AppendObjectPrimitiveWith( - state, frame_index, field_name, [&](VariantValueWriter& writer) { - return writer.template Append(std::forward(args)...); - }); +void AppendObjectPrimitive(const std::shared_ptr& state, + size_t frame_index, std::string_view field_name, + Args&&... args) { + AppendObjectPrimitiveWith(state, frame_index, field_name, + [&](VariantValueWriter& writer) { + writer.template Append(std::forward(args)...); + }); } template -Status AppendListPrimitiveWith(const std::shared_ptr& state, - size_t frame_index, Write&& write) { - ARROW_RETURN_NOT_OK(CheckTopFrame(*state, frame_index, VariantContainerKind::List)); +void AppendListPrimitiveWith(const std::shared_ptr& state, + size_t frame_index, Write&& write) { + CheckTopFrame(*state, frame_index, VariantContainerKind::List); auto& frame = state->frames.back(); const auto value_size = state->value.length(); const auto element_count = frame.offsets.size(); const auto offset = state->value.length() - frame.value_start; DCHECK_GE(offset, 0); if (offset > std::numeric_limits::max()) { - return Status::Invalid("Variant array values are too large"); + throw ParquetException("Variant array values are too large"); } frame.offsets.push_back(static_cast(offset)); VariantValueWriter writer(state->value); - const Status status = std::invoke(std::forward(write), writer); - if (!status.ok()) { + try { + std::invoke(std::forward(write), writer); + } catch (...) { state->value.Rewind(value_size); TruncateFrameEntries(frame, element_count); - return status; + throw; } - return Status::OK(); } template -Status AppendListPrimitive(const std::shared_ptr& state, - size_t frame_index, Args&&... args) { - return AppendListPrimitiveWith(state, frame_index, [&](VariantValueWriter& writer) { - return writer.template Append(std::forward(args)...); +void AppendListPrimitive(const std::shared_ptr& state, + size_t frame_index, Args&&... args) { + AppendListPrimitiveWith(state, frame_index, [&](VariantValueWriter& writer) { + writer.template Append(std::forward(args)...); }); } @@ -609,28 +604,29 @@ VariantBuilder::~VariantBuilder() = default; VariantBuilder::VariantBuilder(VariantBuilder&&) noexcept = default; VariantBuilder& VariantBuilder::operator=(VariantBuilder&&) noexcept = default; -Status VariantBuilder::ReserveFieldNames(int64_t capacity) { - return impl_->state->metadata.Reserve(capacity); +void VariantBuilder::ReserveFieldNames(int64_t capacity) { + impl_->state->metadata.Reserve(capacity); } -Result VariantBuilder::AddFieldName(std::string_view name) { +uint32_t VariantBuilder::AddFieldName(std::string_view name) { return impl_->state->metadata.Upsert(name); } -Status VariantBuilder::AppendVariantNull() { - return AppendRootPrimitive(impl_->state); +void VariantBuilder::AppendVariantNull() { + AppendRootPrimitive(impl_->state); } -Status VariantBuilder::AppendBoolean(bool value) { +void VariantBuilder::AppendBoolean(bool value) { if (value) { - return AppendRootPrimitive(impl_->state); + AppendRootPrimitive(impl_->state); + return; } - return AppendRootPrimitive(impl_->state); + AppendRootPrimitive(impl_->state); } -#define VARIANT_ROOT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - Status VariantBuilder::Append##NAME(C_TYPE value) { \ - return AppendRootPrimitive(impl_->state, value); \ +#define VARIANT_ROOT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + void VariantBuilder::Append##NAME(C_TYPE value) { \ + AppendRootPrimitive(impl_->state, value); \ } VARIANT_ROOT_APPEND_ONE_ARG(Int8, Int8, int8_t) @@ -647,48 +643,45 @@ VARIANT_ROOT_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) #undef VARIANT_ROOT_APPEND_ONE_ARG -Status VariantBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { - return AppendRootPrimitive(impl_->state, - unscaled_value, scale); +void VariantBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { + AppendRootPrimitive(impl_->state, unscaled_value, + scale); } -Status VariantBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { - return AppendRootPrimitive(impl_->state, - unscaled_value, scale); +void VariantBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { + AppendRootPrimitive(impl_->state, unscaled_value, + scale); } -Status VariantBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, - uint8_t scale) { - return AppendRootPrimitive( +void VariantBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, + uint8_t scale) { + AppendRootPrimitive( impl_->state, little_endian_unscaled_value, scale); } -Status VariantBuilder::AppendShortString(std::string_view value) { - return AppendRootPrimitiveWith(impl_->state, [&](VariantValueWriter& writer) { - return writer.AppendShortString(value); - }); +void VariantBuilder::AppendShortString(std::string_view value) { + AppendRootPrimitiveWith( + impl_->state, [&](VariantValueWriter& writer) { writer.AppendShortString(value); }); } -Status VariantBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { +void VariantBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { if (adjusted_to_utc) { - return AppendRootPrimitive(impl_->state, - micros); + AppendRootPrimitive(impl_->state, micros); + return; } - return AppendRootPrimitive(impl_->state, - micros); + AppendRootPrimitive(impl_->state, micros); } -Status VariantBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { +void VariantBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { if (adjusted_to_utc) { - return AppendRootPrimitive(impl_->state, - nanos); + AppendRootPrimitive(impl_->state, nanos); + return; } - return AppendRootPrimitive(impl_->state, - nanos); + AppendRootPrimitive(impl_->state, nanos); } -Result VariantBuilder::StartObject() { - ARROW_RETURN_NOT_OK(CheckRootWritable(*impl_->state)); +VariantObjectBuilder VariantBuilder::StartObject() { + CheckRootWritable(*impl_->state); impl_->state->frames.push_back( VariantBuildFrame{.kind = VariantContainerKind::Object, .value_start = impl_->state->value.length(), @@ -697,8 +690,8 @@ Result VariantBuilder::StartObject() { impl_->state, impl_->state->frames.size() - 1, RootFinishCallback{})); } -Result VariantBuilder::StartList() { - ARROW_RETURN_NOT_OK(CheckRootWritable(*impl_->state)); +VariantListBuilder VariantBuilder::StartList() { + CheckRootWritable(*impl_->state); impl_->state->frames.push_back( VariantBuildFrame{.kind = VariantContainerKind::List, .value_start = impl_->state->value.length(), @@ -707,16 +700,17 @@ Result VariantBuilder::StartList() { impl_->state, impl_->state->frames.size() - 1, RootFinishCallback{})); } -Result VariantBuilder::Finish() { +EncodedVariantValue VariantBuilder::Finish() { if (!impl_->state->frames.empty()) { - return Status::Invalid("Cannot finish VariantBuilder with active containers"); + throw ParquetException("Cannot finish VariantBuilder with active containers"); } if (!impl_->state->root_has_value) { - return Status::Invalid("Cannot finish empty VariantBuilder"); + throw ParquetException("Cannot finish empty VariantBuilder"); } - ARROW_ASSIGN_OR_RAISE(auto metadata, impl_->state->metadata.Finish()); - ARROW_ASSIGN_OR_RAISE(auto value, impl_->state->value.Finish()); + auto metadata = impl_->state->metadata.Finish(); + std::shared_ptr value; + PARQUET_ASSIGN_OR_THROW(value, impl_->state->value.Finish()); EncodedVariantValue out{.metadata = std::move(metadata), .value = std::move(value)}; Reset(); return out; @@ -734,24 +728,25 @@ VariantObjectBuilder::VariantObjectBuilder(VariantObjectBuilder&&) noexcept = de VariantObjectBuilder& VariantObjectBuilder::operator=(VariantObjectBuilder&&) noexcept = default; -Status VariantObjectBuilder::AppendVariantNull(std::string_view field_name) { - return AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name); +void VariantObjectBuilder::AppendVariantNull(std::string_view field_name) { + AppendObjectPrimitive(impl_->state, impl_->frame_index, + field_name); } -Status VariantObjectBuilder::AppendBoolean(std::string_view field_name, bool value) { +void VariantObjectBuilder::AppendBoolean(std::string_view field_name, bool value) { if (value) { - return AppendObjectPrimitive( + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name); + return; } - return AppendObjectPrimitive( + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name); } -#define VARIANT_OBJECT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - Status VariantObjectBuilder::Append##NAME(std::string_view field_name, C_TYPE value) { \ - return AppendObjectPrimitive( \ - impl_->state, impl_->frame_index, field_name, value); \ +#define VARIANT_OBJECT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + void VariantObjectBuilder::Append##NAME(std::string_view field_name, C_TYPE value) { \ + AppendObjectPrimitive( \ + impl_->state, impl_->frame_index, field_name, value); \ } VARIANT_OBJECT_APPEND_ONE_ARG(Int8, Int8, int8_t) @@ -768,71 +763,71 @@ VARIANT_OBJECT_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) #undef VARIANT_OBJECT_APPEND_ONE_ARG -Status VariantObjectBuilder::AppendDecimal4(std::string_view field_name, - int32_t unscaled_value, uint8_t scale) { - return AppendObjectPrimitive( +void VariantObjectBuilder::AppendDecimal4(std::string_view field_name, + int32_t unscaled_value, uint8_t scale) { + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, unscaled_value, scale); } -Status VariantObjectBuilder::AppendDecimal8(std::string_view field_name, - int64_t unscaled_value, uint8_t scale) { - return AppendObjectPrimitive( +void VariantObjectBuilder::AppendDecimal8(std::string_view field_name, + int64_t unscaled_value, uint8_t scale) { + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, unscaled_value, scale); } -Status VariantObjectBuilder::AppendDecimal16( - std::string_view field_name, std::string_view little_endian_unscaled_value, - uint8_t scale) { - return AppendObjectPrimitive( +void VariantObjectBuilder::AppendDecimal16(std::string_view field_name, + std::string_view little_endian_unscaled_value, + uint8_t scale) { + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, little_endian_unscaled_value, scale); } -Status VariantObjectBuilder::AppendShortString(std::string_view field_name, - std::string_view value) { - return AppendObjectPrimitiveWith( +void VariantObjectBuilder::AppendShortString(std::string_view field_name, + std::string_view value) { + AppendObjectPrimitiveWith( impl_->state, impl_->frame_index, field_name, - [&](VariantValueWriter& writer) { return writer.AppendShortString(value); }); + [&](VariantValueWriter& writer) { writer.AppendShortString(value); }); } -Status VariantObjectBuilder::AppendTimestampMicros(std::string_view field_name, - int64_t micros, bool adjusted_to_utc) { +void VariantObjectBuilder::AppendTimestampMicros(std::string_view field_name, + int64_t micros, bool adjusted_to_utc) { if (adjusted_to_utc) { - return AppendObjectPrimitive( + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, micros); + return; } - return AppendObjectPrimitive( + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, micros); } -Status VariantObjectBuilder::AppendTimestampNanos(std::string_view field_name, - int64_t nanos, bool adjusted_to_utc) { +void VariantObjectBuilder::AppendTimestampNanos(std::string_view field_name, + int64_t nanos, bool adjusted_to_utc) { if (adjusted_to_utc) { - return AppendObjectPrimitive( + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, nanos); + return; } - return AppendObjectPrimitive( + AppendObjectPrimitive( impl_->state, impl_->frame_index, field_name, nanos); } -Result VariantObjectBuilder::StartObject( - std::string_view field_name) { - ARROW_RETURN_NOT_OK( - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object)); +VariantObjectBuilder VariantObjectBuilder::StartObject(std::string_view field_name) { + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object); auto& frame = impl_->state->frames.back(); const auto metadata_size = impl_->state->metadata.size(); const auto field_count = frame.fields.size(); - ARROW_ASSIGN_OR_RAISE(auto field_id, impl_->state->metadata.Upsert(field_name)); + auto field_id = impl_->state->metadata.Upsert(field_name); if (!frame.object_field_ids.insert(field_id).second) { impl_->state->metadata.Truncate(metadata_size); - return Status::Invalid("Duplicate Variant object field: ", field_name); + throw ParquetException("Duplicate Variant object field: ", field_name); } const auto offset = impl_->state->value.length() - frame.value_start; DCHECK_GE(offset, 0); if (offset > std::numeric_limits::max()) { impl_->state->metadata.Truncate(metadata_size); TruncateFrameEntries(frame, field_count); - return Status::Invalid("Variant object values are too large"); + throw ParquetException("Variant object values are too large"); } frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, @@ -848,24 +843,23 @@ Result VariantObjectBuilder::StartObject( impl_->state, impl_->state->frames.size() - 1, impl_->callback)); } -Result VariantObjectBuilder::StartList(std::string_view field_name) { - ARROW_RETURN_NOT_OK( - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object)); +VariantListBuilder VariantObjectBuilder::StartList(std::string_view field_name) { + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object); auto& frame = impl_->state->frames.back(); const auto metadata_size = impl_->state->metadata.size(); const auto field_count = frame.fields.size(); - ARROW_ASSIGN_OR_RAISE(auto field_id, impl_->state->metadata.Upsert(field_name)); + auto field_id = impl_->state->metadata.Upsert(field_name); if (!frame.object_field_ids.insert(field_id).second) { impl_->state->metadata.Truncate(metadata_size); - return Status::Invalid("Duplicate Variant object field: ", field_name); + throw ParquetException("Duplicate Variant object field: ", field_name); } const auto offset = impl_->state->value.length() - frame.value_start; DCHECK_GE(offset, 0); if (offset > std::numeric_limits::max()) { impl_->state->metadata.Truncate(metadata_size); TruncateFrameEntries(frame, field_count); - return Status::Invalid("Variant object values are too large"); + throw ParquetException("Variant object values are too large"); } frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, @@ -881,11 +875,10 @@ Result VariantObjectBuilder::StartList(std::string_view fiel impl_->state, impl_->state->frames.size() - 1, impl_->callback)); } -Status VariantObjectBuilder::Finish() { - ARROW_RETURN_NOT_OK(FinishFrame(impl_->state, impl_->frame_index, - VariantContainerKind::Object, impl_->callback)); +void VariantObjectBuilder::Finish() { + FinishFrame(impl_->state, impl_->frame_index, VariantContainerKind::Object, + impl_->callback); impl_->active = false; - return Status::OK(); } VariantListBuilder::VariantListBuilder( @@ -896,24 +889,24 @@ VariantListBuilder::VariantListBuilder(VariantListBuilder&&) noexcept = default; VariantListBuilder& VariantListBuilder::operator=(VariantListBuilder&&) noexcept = default; -Status VariantListBuilder::AppendVariantNull() { - return AppendListPrimitive(impl_->state, - impl_->frame_index); +void VariantListBuilder::AppendVariantNull() { + AppendListPrimitive(impl_->state, impl_->frame_index); } -Status VariantListBuilder::AppendBoolean(bool value) { +void VariantListBuilder::AppendBoolean(bool value) { if (value) { - return AppendListPrimitive(impl_->state, - impl_->frame_index); + AppendListPrimitive(impl_->state, + impl_->frame_index); + return; } - return AppendListPrimitive(impl_->state, - impl_->frame_index); + AppendListPrimitive(impl_->state, + impl_->frame_index); } -#define VARIANT_LIST_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - Status VariantListBuilder::Append##NAME(C_TYPE value) { \ - return AppendListPrimitive( \ - impl_->state, impl_->frame_index, value); \ +#define VARIANT_LIST_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + void VariantListBuilder::Append##NAME(C_TYPE value) { \ + AppendListPrimitive(impl_->state, impl_->frame_index, \ + value); \ } VARIANT_LIST_APPEND_ONE_ARG(Int8, Int8, int8_t) @@ -930,55 +923,56 @@ VARIANT_LIST_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) #undef VARIANT_LIST_APPEND_ONE_ARG -Status VariantListBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { - return AppendListPrimitive( - impl_->state, impl_->frame_index, unscaled_value, scale); +void VariantListBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { + AppendListPrimitive(impl_->state, impl_->frame_index, + unscaled_value, scale); } -Status VariantListBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { - return AppendListPrimitive( - impl_->state, impl_->frame_index, unscaled_value, scale); +void VariantListBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { + AppendListPrimitive(impl_->state, impl_->frame_index, + unscaled_value, scale); } -Status VariantListBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, - uint8_t scale) { - return AppendListPrimitive( +void VariantListBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, + uint8_t scale) { + AppendListPrimitive( impl_->state, impl_->frame_index, little_endian_unscaled_value, scale); } -Status VariantListBuilder::AppendShortString(std::string_view value) { - return AppendListPrimitiveWith( +void VariantListBuilder::AppendShortString(std::string_view value) { + AppendListPrimitiveWith( impl_->state, impl_->frame_index, - [&](VariantValueWriter& writer) { return writer.AppendShortString(value); }); + [&](VariantValueWriter& writer) { writer.AppendShortString(value); }); } -Status VariantListBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { +void VariantListBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { if (adjusted_to_utc) { - return AppendListPrimitive( + AppendListPrimitive( impl_->state, impl_->frame_index, micros); + return; } - return AppendListPrimitive( + AppendListPrimitive( impl_->state, impl_->frame_index, micros); } -Status VariantListBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { +void VariantListBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { if (adjusted_to_utc) { - return AppendListPrimitive( - impl_->state, impl_->frame_index, nanos); + AppendListPrimitive(impl_->state, + impl_->frame_index, nanos); + return; } - return AppendListPrimitive( + AppendListPrimitive( impl_->state, impl_->frame_index, nanos); } -Result VariantListBuilder::StartObject() { - ARROW_RETURN_NOT_OK( - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List)); +VariantObjectBuilder VariantListBuilder::StartObject() { + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List); auto& frame = impl_->state->frames.back(); const auto element_count = frame.offsets.size(); const auto offset = impl_->state->value.length() - frame.value_start; DCHECK_GE(offset, 0); if (offset > std::numeric_limits::max()) { - return Status::Invalid("Variant array values are too large"); + throw ParquetException("Variant array values are too large"); } frame.offsets.push_back(static_cast(offset)); @@ -994,15 +988,14 @@ Result VariantListBuilder::StartObject() { impl_->state, impl_->state->frames.size() - 1, impl_->callback)); } -Result VariantListBuilder::StartList() { - ARROW_RETURN_NOT_OK( - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List)); +VariantListBuilder VariantListBuilder::StartList() { + CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List); auto& frame = impl_->state->frames.back(); const auto element_count = frame.offsets.size(); const auto offset = impl_->state->value.length() - frame.value_start; DCHECK_GE(offset, 0); if (offset > std::numeric_limits::max()) { - return Status::Invalid("Variant array values are too large"); + throw ParquetException("Variant array values are too large"); } frame.offsets.push_back(static_cast(offset)); @@ -1018,11 +1011,10 @@ Result VariantListBuilder::StartList() { impl_->state, impl_->state->frames.size() - 1, impl_->callback)); } -Status VariantListBuilder::Finish() { - ARROW_RETURN_NOT_OK(FinishFrame(impl_->state, impl_->frame_index, - VariantContainerKind::List, impl_->callback)); +void VariantListBuilder::Finish() { + FinishFrame(impl_->state, impl_->frame_index, VariantContainerKind::List, + impl_->callback); impl_->active = false; - return Status::OK(); } struct VariantArrayBuilder::Impl { @@ -1030,38 +1022,36 @@ struct VariantArrayBuilder::Impl { : pool(pool), metadata_builder(pool), value_builder(pool), validity_builder(pool) {} template - Status AppendValue(Write&& write) { + void AppendValue(Write&& write) { auto state = std::make_shared(pool); - ARROW_RETURN_NOT_OK(AppendRootPrimitiveWith(state, std::forward(write))); - ARROW_ASSIGN_OR_RAISE(auto metadata, state->metadata.Finish()); - ARROW_ASSIGN_OR_RAISE(auto value, state->value.Finish()); - return AppendEncoded({.metadata = std::move(metadata), .value = std::move(value)}); + AppendRootPrimitiveWith(state, std::forward(write)); + auto metadata = state->metadata.Finish(); + std::shared_ptr value; + PARQUET_ASSIGN_OR_THROW(value, state->value.Finish()); + AppendEncoded({.metadata = std::move(metadata), .value = std::move(value)}); } template - Status AppendPrimitive(Args&&... args) { - return AppendValue([&](VariantValueWriter& writer) { - return writer.template Append(std::forward(args)...); + void AppendPrimitive(Args&&... args) { + AppendValue([&](VariantValueWriter& writer) { + writer.template Append(std::forward(args)...); }); } - Status AppendShortString(std::string_view value) { - return AppendValue( - [&](VariantValueWriter& writer) { return writer.AppendShortString(value); }); + void AppendShortString(std::string_view value) { + AppendValue([&](VariantValueWriter& writer) { writer.AppendShortString(value); }); } - Status AppendEncoded(const EncodedVariantValue& value) { + void AppendEncoded(const EncodedVariantValue& value) { if (value.metadata == nullptr || value.value == nullptr) { - return Status::Invalid( + throw ParquetException( "Encoded Variant metadata and value buffers must be non-null"); } - ARROW_ASSIGN_OR_RAISE(auto metadata, - VariantMetadataView::Make(std::string_view{*value.metadata})); - ARROW_RETURN_NOT_OK( - VariantValueView::Validate(std::string_view{*value.value}, metadata)); - ARROW_RETURN_NOT_OK(metadata_builder.Append(std::string_view{*value.metadata})); - ARROW_RETURN_NOT_OK(value_builder.Append(std::string_view{*value.value})); - return validity_builder.Append(true); + auto metadata = VariantMetadataView::Make(std::string_view{*value.metadata}); + VariantValueView::Validate(std::string_view{*value.value}, metadata); + PARQUET_THROW_NOT_OK(metadata_builder.Append(std::string_view{*value.metadata})); + PARQUET_THROW_NOT_OK(value_builder.Append(std::string_view{*value.value})); + PARQUET_THROW_NOT_OK(validity_builder.Append(true)); } MemoryPool* pool; @@ -1077,26 +1067,27 @@ VariantArrayBuilder::VariantArrayBuilder(VariantArrayBuilder&&) noexcept = defau VariantArrayBuilder& VariantArrayBuilder::operator=(VariantArrayBuilder&&) noexcept = default; -Status VariantArrayBuilder::AppendNull() { - ARROW_RETURN_NOT_OK(impl_->metadata_builder.Append("")); - ARROW_RETURN_NOT_OK(impl_->value_builder.Append("")); - return impl_->validity_builder.Append(false); +void VariantArrayBuilder::AppendNull() { + PARQUET_THROW_NOT_OK(impl_->metadata_builder.Append("")); + PARQUET_THROW_NOT_OK(impl_->value_builder.Append("")); + PARQUET_THROW_NOT_OK(impl_->validity_builder.Append(false)); } -Status VariantArrayBuilder::AppendVariantNull() { - return impl_->AppendPrimitive(); +void VariantArrayBuilder::AppendVariantNull() { + impl_->AppendPrimitive(); } -Status VariantArrayBuilder::AppendBoolean(bool value) { +void VariantArrayBuilder::AppendBoolean(bool value) { if (value) { - return impl_->AppendPrimitive(); + impl_->AppendPrimitive(); + return; } - return impl_->AppendPrimitive(); + impl_->AppendPrimitive(); } -#define VARIANT_ARRAY_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - Status VariantArrayBuilder::Append##NAME(C_TYPE value) { \ - return impl_->AppendPrimitive(value); \ +#define VARIANT_ARRAY_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ + void VariantArrayBuilder::Append##NAME(C_TYPE value) { \ + impl_->AppendPrimitive(value); \ } VARIANT_ARRAY_APPEND_ONE_ARG(Int8, Int8, int8_t) @@ -1113,81 +1104,80 @@ VARIANT_ARRAY_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) #undef VARIANT_ARRAY_APPEND_ONE_ARG -Status VariantArrayBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { - return impl_->AppendPrimitive(unscaled_value, scale); +void VariantArrayBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { + impl_->AppendPrimitive(unscaled_value, scale); } -Status VariantArrayBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { - return impl_->AppendPrimitive(unscaled_value, scale); +void VariantArrayBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { + impl_->AppendPrimitive(unscaled_value, scale); } -Status VariantArrayBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, - uint8_t scale) { - return impl_->AppendPrimitive( - little_endian_unscaled_value, scale); +void VariantArrayBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, + uint8_t scale) { + impl_->AppendPrimitive(little_endian_unscaled_value, + scale); } -Status VariantArrayBuilder::AppendShortString(std::string_view value) { - return impl_->AppendShortString(value); +void VariantArrayBuilder::AppendShortString(std::string_view value) { + impl_->AppendShortString(value); } -Status VariantArrayBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { +void VariantArrayBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { if (adjusted_to_utc) { - return impl_->AppendPrimitive(micros); + impl_->AppendPrimitive(micros); + return; } - return impl_->AppendPrimitive(micros); + impl_->AppendPrimitive(micros); } -Status VariantArrayBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { +void VariantArrayBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { if (adjusted_to_utc) { - return impl_->AppendPrimitive(nanos); + impl_->AppendPrimitive(nanos); + return; } - return impl_->AppendPrimitive(nanos); + impl_->AppendPrimitive(nanos); } -Status VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { - return impl_->AppendEncoded(value); +void VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { + impl_->AppendEncoded(value); } -Result VariantArrayBuilder::StartObject() { +VariantObjectBuilder VariantArrayBuilder::StartObject() { auto state = std::make_shared(impl_->pool); state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::Object, .value_start = state->value.length(), .metadata_size = state->metadata.size()}); - auto callback = [this](EncodedVariantValue encoded) { - return impl_->AppendEncoded(encoded); - }; + auto callback = [this](EncodedVariantValue encoded) { impl_->AppendEncoded(encoded); }; return VariantObjectBuilder(std::make_unique( state, state->frames.size() - 1, std::move(callback))); } -Result VariantArrayBuilder::StartList() { +VariantListBuilder VariantArrayBuilder::StartList() { auto state = std::make_shared(impl_->pool); state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::List, .value_start = state->value.length(), .metadata_size = state->metadata.size()}); - auto callback = [this](EncodedVariantValue encoded) { - return impl_->AppendEncoded(encoded); - }; + auto callback = [this](EncodedVariantValue encoded) { impl_->AppendEncoded(encoded); }; return VariantListBuilder(std::make_unique( state, state->frames.size() - 1, std::move(callback))); } -Result> VariantArrayBuilder::Finish() { +std::shared_ptr VariantArrayBuilder::Finish() { std::shared_ptr metadata; std::shared_ptr value; std::shared_ptr validity; - ARROW_RETURN_NOT_OK(impl_->metadata_builder.Finish(&metadata)); - ARROW_RETURN_NOT_OK(impl_->value_builder.Finish(&value)); - ARROW_RETURN_NOT_OK(impl_->validity_builder.Finish(&validity)); + PARQUET_THROW_NOT_OK(impl_->metadata_builder.Finish(&metadata)); + PARQUET_THROW_NOT_OK(impl_->value_builder.Finish(&value)); + PARQUET_THROW_NOT_OK(impl_->validity_builder.Finish(&validity)); auto null_bitmap = validity->data()->buffers[1]; const int64_t null_count = validity->false_count(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); - ARROW_ASSIGN_OR_RAISE(auto storage, - StructArray::Make({metadata, value}, storage_type->fields(), - null_bitmap, null_count)); + std::shared_ptr storage; + PARQUET_ASSIGN_OR_THROW( + storage, StructArray::Make({metadata, value}, storage_type->fields(), null_bitmap, + null_count)); return MakeVariantArrayFromStorage(storage); } @@ -1207,66 +1197,68 @@ VariantValueArrayBuilder& VariantValueArrayBuilder::operator=( VariantValueArrayBuilder&&) noexcept = default; VariantValueArrayBuilder::~VariantValueArrayBuilder() = default; -Status VariantValueArrayBuilder::AppendNull() { - return impl_->value_builder.AppendNull(); +void VariantValueArrayBuilder::AppendNull() { + PARQUET_THROW_NOT_OK(impl_->value_builder.AppendNull()); } -Status VariantValueArrayBuilder::AppendEncodedValue(std::string_view metadata, - std::string_view value) { - ARROW_ASSIGN_OR_RAISE(auto metadata_view, VariantMetadataView::Make(metadata)); - ARROW_RETURN_NOT_OK(VariantValueView::Validate(value, metadata_view)); - return impl_->value_builder.Append(value); +void VariantValueArrayBuilder::AppendEncodedValue(std::string_view metadata, + std::string_view value) { + auto metadata_view = VariantMetadataView::Make(metadata); + VariantValueView::Validate(value, metadata_view); + PARQUET_THROW_NOT_OK(impl_->value_builder.Append(value)); } -Result> VariantValueArrayBuilder::Finish() { +std::shared_ptr VariantValueArrayBuilder::Finish() { std::shared_ptr out; - ARROW_RETURN_NOT_OK(impl_->value_builder.Finish(&out)); + PARQUET_THROW_NOT_OK(impl_->value_builder.Finish(&out)); return out; } -Result> MakeVariantArrayFromStorage( +std::shared_ptr MakeVariantArrayFromStorage( std::shared_ptr storage) { if (storage == nullptr) { - return Status::Invalid("Variant storage array must be non-null"); + throw ParquetException("Variant storage array must be non-null"); } - ARROW_ASSIGN_OR_RAISE(auto type, VariantExtensionType::Make(storage->type())); + std::shared_ptr type; + PARQUET_ASSIGN_OR_THROW(type, VariantExtensionType::Make(storage->type())); auto array = ExtensionType::WrapArray(type, std::move(storage)); return std::static_pointer_cast(array); } -Result> MakeVariantArrayFromChildren( +std::shared_ptr MakeVariantArrayFromChildren( std::shared_ptr storage_type, std::vector> children, std::shared_ptr null_bitmap) { if (storage_type->id() != Type::STRUCT) { - return Status::Invalid("Variant storage type must be struct, got ", + throw ParquetException("Variant storage type must be struct, got ", storage_type->ToString()); } const auto& struct_type = ::arrow::internal::checked_cast(*storage_type); if (children.size() != static_cast(struct_type.num_fields())) { - return Status::Invalid("Variant storage expected ", struct_type.num_fields(), + throw ParquetException("Variant storage expected ", struct_type.num_fields(), " children, got ", children.size()); } const int64_t length = children.empty() ? 0 : children[0]->length(); for (int i = 0; i < struct_type.num_fields(); ++i) { if (children[i] == nullptr) { - return Status::Invalid("Variant storage child ", i, " is null"); + throw ParquetException("Variant storage child ", i, " is null"); } if (!children[i]->type()->Equals(struct_type.field(i)->type())) { - return Status::Invalid("Variant storage child ", i, " has type ", + throw ParquetException("Variant storage child ", i, " has type ", children[i]->type()->ToString(), ", expected ", struct_type.field(i)->type()->ToString()); } if (children[i]->length() != length) { - return Status::Invalid("Variant storage child lengths must match"); + throw ParquetException("Variant storage child lengths must match"); } } - ARROW_ASSIGN_OR_RAISE(auto storage, - StructArray::Make(std::move(children), struct_type.fields(), - std::move(null_bitmap))); + std::shared_ptr storage; + PARQUET_ASSIGN_OR_THROW(storage, + StructArray::Make(std::move(children), struct_type.fields(), + std::move(null_bitmap))); return MakeVariantArrayFromStorage(storage); } diff --git a/cpp/src/parquet/variant/builder.h b/cpp/src/parquet/variant/builder.h index cc0507bfa9f8..1fd0a03843e0 100644 --- a/cpp/src/parquet/variant/builder.h +++ b/cpp/src/parquet/variant/builder.h @@ -25,8 +25,6 @@ #include "arrow/buffer.h" #include "arrow/extension/parquet_variant.h" #include "arrow/memory_pool.h" -#include "arrow/result.h" -#include "arrow/status.h" #include "arrow/type.h" #include "parquet/platform.h" @@ -37,8 +35,6 @@ using ::arrow::BinaryArray; using ::arrow::Buffer; using ::arrow::DataType; using ::arrow::MemoryPool; -using ::arrow::Result; -using ::arrow::Status; using ::arrow::StructArray; using ::arrow::extension::VariantArray; @@ -59,33 +55,33 @@ class PARQUET_EXPORT VariantBuilder { VariantBuilder(VariantBuilder&&) noexcept; VariantBuilder& operator=(VariantBuilder&&) noexcept; - Status ReserveFieldNames(int64_t capacity); - Result AddFieldName(std::string_view name); - - Status AppendVariantNull(); - Status AppendBoolean(bool value); - Status AppendInt8(int8_t value); - Status AppendInt16(int16_t value); - Status AppendInt32(int32_t value); - Status AppendInt64(int64_t value); - Status AppendFloat(float value); - Status AppendDouble(double value); - Status AppendBinary(std::string_view value); - Status AppendString(std::string_view value); - Status AppendDate(int32_t days); - Status AppendTimeNTZMicros(int64_t micros); - Status AppendUuid(std::string_view big_endian_bytes); - Status AppendDecimal4(int32_t unscaled_value, uint8_t scale); - Status AppendDecimal8(int64_t unscaled_value, uint8_t scale); - Status AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); - Status AppendShortString(std::string_view value); - Status AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); - Status AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); - - Result StartObject(); - Result StartList(); - - Result Finish(); + void ReserveFieldNames(int64_t capacity); + uint32_t AddFieldName(std::string_view name); + + void AppendVariantNull(); + void AppendBoolean(bool value); + void AppendInt8(int8_t value); + void AppendInt16(int16_t value); + void AppendInt32(int32_t value); + void AppendInt64(int64_t value); + void AppendFloat(float value); + void AppendDouble(double value); + void AppendBinary(std::string_view value); + void AppendString(std::string_view value); + void AppendDate(int32_t days); + void AppendTimeNTZMicros(int64_t micros); + void AppendUuid(std::string_view big_endian_bytes); + void AppendDecimal4(int32_t unscaled_value, uint8_t scale); + void AppendDecimal8(int64_t unscaled_value, uint8_t scale); + void AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendShortString(std::string_view value); + void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + + VariantObjectBuilder StartObject(); + VariantListBuilder StartList(); + + EncodedVariantValue Finish(); void Reset(); private: @@ -105,37 +101,35 @@ class PARQUET_EXPORT VariantObjectBuilder { VariantObjectBuilder(VariantObjectBuilder&&) noexcept; VariantObjectBuilder& operator=(VariantObjectBuilder&&) noexcept; - Status AppendVariantNull(std::string_view field_name); - Status AppendBoolean(std::string_view field_name, bool value); - Status AppendInt8(std::string_view field_name, int8_t value); - Status AppendInt16(std::string_view field_name, int16_t value); - Status AppendInt32(std::string_view field_name, int32_t value); - Status AppendInt64(std::string_view field_name, int64_t value); - Status AppendFloat(std::string_view field_name, float value); - Status AppendDouble(std::string_view field_name, double value); - Status AppendBinary(std::string_view field_name, std::string_view value); - Status AppendString(std::string_view field_name, std::string_view value); - Status AppendDate(std::string_view field_name, int32_t days); - Status AppendTimeNTZMicros(std::string_view field_name, int64_t micros); - Status AppendUuid(std::string_view field_name, std::string_view big_endian_bytes); - Status AppendDecimal4(std::string_view field_name, int32_t unscaled_value, - uint8_t scale); - Status AppendDecimal8(std::string_view field_name, int64_t unscaled_value, - uint8_t scale); - Status AppendDecimal16(std::string_view field_name, - std::string_view little_endian_unscaled_value, uint8_t scale); - Status AppendShortString(std::string_view field_name, std::string_view value); - Status AppendTimestampMicros(std::string_view field_name, int64_t micros, - bool adjusted_to_utc); - Status AppendTimestampNanos(std::string_view field_name, int64_t nanos, - bool adjusted_to_utc); - - Result StartObject(std::string_view field_name); - Result StartList(std::string_view field_name); + void AppendVariantNull(std::string_view field_name); + void AppendBoolean(std::string_view field_name, bool value); + void AppendInt8(std::string_view field_name, int8_t value); + void AppendInt16(std::string_view field_name, int16_t value); + void AppendInt32(std::string_view field_name, int32_t value); + void AppendInt64(std::string_view field_name, int64_t value); + void AppendFloat(std::string_view field_name, float value); + void AppendDouble(std::string_view field_name, double value); + void AppendBinary(std::string_view field_name, std::string_view value); + void AppendString(std::string_view field_name, std::string_view value); + void AppendDate(std::string_view field_name, int32_t days); + void AppendTimeNTZMicros(std::string_view field_name, int64_t micros); + void AppendUuid(std::string_view field_name, std::string_view big_endian_bytes); + void AppendDecimal4(std::string_view field_name, int32_t unscaled_value, uint8_t scale); + void AppendDecimal8(std::string_view field_name, int64_t unscaled_value, uint8_t scale); + void AppendDecimal16(std::string_view field_name, + std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendShortString(std::string_view field_name, std::string_view value); + void AppendTimestampMicros(std::string_view field_name, int64_t micros, + bool adjusted_to_utc); + void AppendTimestampNanos(std::string_view field_name, int64_t nanos, + bool adjusted_to_utc); + + VariantObjectBuilder StartObject(std::string_view field_name); + VariantListBuilder StartList(std::string_view field_name); /// Commit this nested object into its parent builder. Destroying an unfinished /// nested builder rolls back the object contents written through this builder. - Status Finish(); + void Finish(); private: friend class VariantBuilder; @@ -155,32 +149,32 @@ class PARQUET_EXPORT VariantListBuilder { VariantListBuilder(VariantListBuilder&&) noexcept; VariantListBuilder& operator=(VariantListBuilder&&) noexcept; - Status AppendVariantNull(); - Status AppendBoolean(bool value); - Status AppendInt8(int8_t value); - Status AppendInt16(int16_t value); - Status AppendInt32(int32_t value); - Status AppendInt64(int64_t value); - Status AppendFloat(float value); - Status AppendDouble(double value); - Status AppendBinary(std::string_view value); - Status AppendString(std::string_view value); - Status AppendDate(int32_t days); - Status AppendTimeNTZMicros(int64_t micros); - Status AppendUuid(std::string_view big_endian_bytes); - Status AppendDecimal4(int32_t unscaled_value, uint8_t scale); - Status AppendDecimal8(int64_t unscaled_value, uint8_t scale); - Status AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); - Status AppendShortString(std::string_view value); - Status AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); - Status AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); - - Result StartObject(); - Result StartList(); + void AppendVariantNull(); + void AppendBoolean(bool value); + void AppendInt8(int8_t value); + void AppendInt16(int16_t value); + void AppendInt32(int32_t value); + void AppendInt64(int64_t value); + void AppendFloat(float value); + void AppendDouble(double value); + void AppendBinary(std::string_view value); + void AppendString(std::string_view value); + void AppendDate(int32_t days); + void AppendTimeNTZMicros(int64_t micros); + void AppendUuid(std::string_view big_endian_bytes); + void AppendDecimal4(int32_t unscaled_value, uint8_t scale); + void AppendDecimal8(int64_t unscaled_value, uint8_t scale); + void AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendShortString(std::string_view value); + void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + + VariantObjectBuilder StartObject(); + VariantListBuilder StartList(); /// Commit this nested list into its parent builder. Destroying an unfinished nested /// builder rolls back the list contents written through this builder. - Status Finish(); + void Finish(); private: friend class VariantBuilder; @@ -201,32 +195,32 @@ class PARQUET_EXPORT VariantArrayBuilder { VariantArrayBuilder(VariantArrayBuilder&&) noexcept; VariantArrayBuilder& operator=(VariantArrayBuilder&&) noexcept; - Status AppendNull(); - Status AppendVariantNull(); - Status AppendBoolean(bool value); - Status AppendInt8(int8_t value); - Status AppendInt16(int16_t value); - Status AppendInt32(int32_t value); - Status AppendInt64(int64_t value); - Status AppendFloat(float value); - Status AppendDouble(double value); - Status AppendBinary(std::string_view value); - Status AppendString(std::string_view value); - Status AppendDate(int32_t days); - Status AppendTimeNTZMicros(int64_t micros); - Status AppendUuid(std::string_view big_endian_bytes); - Status AppendDecimal4(int32_t unscaled_value, uint8_t scale); - Status AppendDecimal8(int64_t unscaled_value, uint8_t scale); - Status AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); - Status AppendShortString(std::string_view value); - Status AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); - Status AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); - Status AppendEncoded(const EncodedVariantValue& value); - - Result StartObject(); - Result StartList(); - - Result> Finish(); + void AppendNull(); + void AppendVariantNull(); + void AppendBoolean(bool value); + void AppendInt8(int8_t value); + void AppendInt16(int16_t value); + void AppendInt32(int32_t value); + void AppendInt64(int64_t value); + void AppendFloat(float value); + void AppendDouble(double value); + void AppendBinary(std::string_view value); + void AppendString(std::string_view value); + void AppendDate(int32_t days); + void AppendTimeNTZMicros(int64_t micros); + void AppendUuid(std::string_view big_endian_bytes); + void AppendDecimal4(int32_t unscaled_value, uint8_t scale); + void AppendDecimal8(int64_t unscaled_value, uint8_t scale); + void AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendShortString(std::string_view value); + void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + void AppendEncoded(const EncodedVariantValue& value); + + VariantObjectBuilder StartObject(); + VariantListBuilder StartList(); + + std::shared_ptr Finish(); void Reset(); private: @@ -243,9 +237,9 @@ class PARQUET_EXPORT VariantValueArrayBuilder { VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept; VariantValueArrayBuilder& operator=(VariantValueArrayBuilder&&) noexcept; - Status AppendNull(); - Status AppendEncodedValue(std::string_view metadata, std::string_view value); - Result> Finish(); + void AppendNull(); + void AppendEncodedValue(std::string_view metadata, std::string_view value); + std::shared_ptr Finish(); private: struct Impl; @@ -253,11 +247,11 @@ class PARQUET_EXPORT VariantValueArrayBuilder { }; PARQUET_EXPORT -Result> MakeVariantArrayFromStorage( +std::shared_ptr MakeVariantArrayFromStorage( std::shared_ptr storage); PARQUET_EXPORT -Result> MakeVariantArrayFromChildren( +std::shared_ptr MakeVariantArrayFromChildren( std::shared_ptr storage_type, std::vector> children, std::shared_ptr null_bitmap = nullptr); diff --git a/cpp/src/parquet/variant/builder_test.cc b/cpp/src/parquet/variant/builder_test.cc index e82addb27a81..b3f069dde7ce 100644 --- a/cpp/src/parquet/variant/builder_test.cc +++ b/cpp/src/parquet/variant/builder_test.cc @@ -46,7 +46,7 @@ using internal::MakeVariantValueView; void AssertPrimitiveType(std::string_view value, const VariantMetadataView& metadata, VariantPrimitiveType expected) { - ASSERT_OK_AND_ASSIGN(auto view, VariantValueView::Make(value, metadata)); + auto view = VariantValueView::Make(value, metadata); ASSERT_EQ(VariantBasicType::kPrimitive, view.basic_type()); ASSERT_EQ(expected, std::get(view.data()).type()); } @@ -63,14 +63,14 @@ void AssertPrimitiveFieldType(const VariantObjectView& object, std::string_view TEST(TestVariantBuilder, Object) { VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); - ASSERT_OK(object.AppendVariantNull("c")); - ASSERT_OK(object.AppendVariantNull("b")); - ASSERT_OK(object.AppendVariantNull("a")); - ASSERT_OK(object.Finish()); - - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto view, MakeVariantValueView(encoded)); + auto object = builder.StartObject(); + object.AppendVariantNull("c"); + object.AppendVariantNull("b"); + object.AppendVariantNull("a"); + object.Finish(); + + auto encoded = builder.Finish(); + auto view = MakeVariantValueView(encoded); ASSERT_EQ(VariantBasicType::kObject, view.basic_type()); const auto& fields = std::get(view.data()).fields(); ASSERT_EQ(3, fields.size()); @@ -81,73 +81,66 @@ TEST(TestVariantBuilder, Object) { TEST(TestVariantBuilder, List) { VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto list, builder.StartList()); - ASSERT_OK(list.AppendVariantNull()); - ASSERT_OK(list.AppendInt32(42)); - ASSERT_OK(list.AppendString("x")); - ASSERT_OK(list.Finish()); - - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); - ASSERT_OK_AND_ASSIGN( - auto view, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + auto list = builder.StartList(); + list.AppendVariantNull(); + list.AppendInt32(42); + list.AppendString("x"); + list.Finish(); + + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); ASSERT_EQ(VariantBasicType::kArray, view.basic_type()); const auto& elements = std::get(view.data()).elements(); ASSERT_EQ(3, elements.size()); - ASSERT_OK_AND_ASSIGN(auto element, VariantValueView::Make(elements[1], metadata)); + auto element = VariantValueView::Make(elements[1], metadata); ASSERT_EQ(VariantPrimitiveType::kInt32, std::get(element.data()).type()); } TEST(TestVariantBuilder, Nested) { VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); - ASSERT_OK_AND_ASSIGN(auto list, object.StartList("items")); - ASSERT_OK(list.AppendInt32(1)); - ASSERT_OK_AND_ASSIGN(auto child, list.StartObject()); - ASSERT_OK(child.AppendString("name", "x")); - ASSERT_OK(child.Finish()); - ASSERT_OK(list.Finish()); - ASSERT_OK(object.Finish()); - - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); - ASSERT_OK_AND_ASSIGN( - auto root, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + auto object = builder.StartObject(); + auto list = object.StartList("items"); + list.AppendInt32(1); + auto child = list.StartObject(); + child.AppendString("name", "x"); + child.Finish(); + list.Finish(); + object.Finish(); + + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto root = VariantValueView::Make(std::string_view{*encoded.value}, metadata); const auto& root_fields = std::get(root.data()).fields(); ASSERT_EQ(1, root_fields.size()); ASSERT_EQ("items", root_fields[0].name); - ASSERT_OK_AND_ASSIGN(auto items, - VariantValueView::Make(root_fields[0].value, metadata)); + auto items = VariantValueView::Make(root_fields[0].value, metadata); const auto& item_values = std::get(items.data()).elements(); ASSERT_EQ(2, item_values.size()); - ASSERT_OK_AND_ASSIGN(auto item, VariantValueView::Make(item_values[1], metadata)); + auto item = VariantValueView::Make(item_values[1], metadata); ASSERT_EQ("name", std::get(item.data()).fields()[0].name); } TEST(TestVariantBuilder, ObjectAppends) { const std::string uuid(16, '\1'); VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); - ASSERT_OK(object.AppendInt8("int8", 1)); - ASSERT_OK(object.AppendInt64("int64", 2)); - ASSERT_OK(object.AppendDouble("double", 3)); - ASSERT_OK(object.AppendDecimal4("decimal", 1234, 2)); - ASSERT_OK(object.AppendBinary("binary", "abc")); - ASSERT_OK(object.AppendDate("date", 1)); - ASSERT_OK(object.AppendTimestampNanos("ts", 2, true)); - ASSERT_OK(object.AppendUuid("uuid", uuid)); - ASSERT_OK(object.Finish()); - - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); - ASSERT_OK_AND_ASSIGN( - auto root, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + auto object = builder.StartObject(); + object.AppendInt8("int8", 1); + object.AppendInt64("int64", 2); + object.AppendDouble("double", 3); + object.AppendDecimal4("decimal", 1234, 2); + object.AppendBinary("binary", "abc"); + object.AppendDate("date", 1); + object.AppendTimestampNanos("ts", 2, true); + object.AppendUuid("uuid", uuid); + object.Finish(); + + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto root = VariantValueView::Make(std::string_view{*encoded.value}, metadata); const auto& object_view = std::get(root.data()); ASSERT_EQ(8, object_view.fields().size()); AssertPrimitiveFieldType(object_view, "int8", metadata, VariantPrimitiveType::kInt8); @@ -167,22 +160,20 @@ TEST(TestVariantBuilder, ObjectAppends) { TEST(TestVariantBuilder, ListAppends) { const std::string uuid(16, '\2'); VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto list, builder.StartList()); - ASSERT_OK(list.AppendInt8(1)); - ASSERT_OK(list.AppendInt64(2)); - ASSERT_OK(list.AppendDouble(3)); - ASSERT_OK(list.AppendDecimal4(1234, 2)); - ASSERT_OK(list.AppendBinary("abc")); - ASSERT_OK(list.AppendDate(1)); - ASSERT_OK(list.AppendTimestampNanos(2, true)); - ASSERT_OK(list.AppendUuid(uuid)); - ASSERT_OK(list.Finish()); - - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); - ASSERT_OK_AND_ASSIGN( - auto root, VariantValueView::Make(std::string_view{*encoded.value}, metadata)); + auto list = builder.StartList(); + list.AppendInt8(1); + list.AppendInt64(2); + list.AppendDouble(3); + list.AppendDecimal4(1234, 2); + list.AppendBinary("abc"); + list.AppendDate(1); + list.AppendTimestampNanos(2, true); + list.AppendUuid(uuid); + list.Finish(); + + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto root = VariantValueView::Make(std::string_view{*encoded.value}, metadata); const auto& elements = std::get(root.data()).elements(); ASSERT_EQ(8, elements.size()); AssertPrimitiveType(elements[0], metadata, VariantPrimitiveType::kInt8); @@ -197,16 +188,16 @@ TEST(TestVariantBuilder, ListAppends) { TEST(TestVariantBuilder, Rollback) { VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); + auto object = builder.StartObject(); { - ASSERT_OK_AND_ASSIGN(auto child, object.StartObject("drop")); - ASSERT_OK(child.AppendString("nested", "x")); + auto child = object.StartObject("drop"); + child.AppendString("nested", "x"); } - ASSERT_OK(object.AppendInt32("ok", 1)); - ASSERT_OK(object.Finish()); + object.AppendInt32("ok", 1); + object.Finish(); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto view, MakeVariantValueView(encoded)); + auto encoded = builder.Finish(); + auto view = MakeVariantValueView(encoded); const auto& fields = std::get(view.data()).fields(); ASSERT_EQ(1, fields.size()); ASSERT_EQ("ok", fields[0].name); @@ -214,37 +205,36 @@ TEST(TestVariantBuilder, Rollback) { TEST(TestVariantBuilder, Duplicate) { VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); - ASSERT_OK(object.AppendInt32("a", 1)); - ASSERT_RAISES(Invalid, object.AppendString("a", "x")); + auto object = builder.StartObject(); + object.AppendInt32("a", 1); + ASSERT_THROW(object.AppendString("a", "x"), ParquetException); } TEST(TestVariantBuilder, UsesMemoryPool) { ProxyMemoryPool pool(default_memory_pool()); VariantBuilder builder(&pool); - ASSERT_OK(builder.AppendString(std::string(128, 'x'))); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + builder.AppendString(std::string(128, 'x')); + auto encoded = builder.Finish(); ASSERT_GT(pool.total_bytes_allocated(), 0); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); - ASSERT_OK(VariantValueView::Validate(std::string_view{*encoded.value}, metadata)); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + VariantValueView::Validate(std::string_view{*encoded.value}, metadata); } TEST(TestVariantBuilder, ArrayBuilder) { VariantArrayBuilder builder; - ASSERT_OK(builder.AppendNull()); - ASSERT_OK(builder.AppendVariantNull()); - ASSERT_OK(builder.AppendInt32(42)); - ASSERT_OK(builder.AppendInt8(1)); - ASSERT_OK(builder.AppendDouble(2)); - ASSERT_OK(builder.AppendBinary("abc")); - ASSERT_OK(builder.AppendTimestampMicros(3, true)); - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); - ASSERT_OK(object.AppendString("a", "x")); - ASSERT_OK(object.Finish()); - - ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + builder.AppendNull(); + builder.AppendVariantNull(); + builder.AppendInt32(42); + builder.AppendInt8(1); + builder.AppendDouble(2); + builder.AppendBinary("abc"); + builder.AppendTimestampMicros(3, true); + auto object = builder.StartObject(); + object.AppendString("a", "x"); + object.Finish(); + + auto array = builder.Finish(); ASSERT_EQ(8, array->length()); ASSERT_TRUE(array->IsNull(0)); ASSERT_FALSE(array->IsNull(1)); @@ -259,8 +249,8 @@ TEST(TestVariantBuilder, ArrayBuilder) { TEST(TestVariantBuilder, FromStorage) { VariantBuilder value; - ASSERT_OK(value.AppendInt32(1)); - ASSERT_OK_AND_ASSIGN(auto encoded, value.Finish()); + value.AppendInt32(1); + auto encoded = value.Finish(); auto metadata = BinaryArrayFromValues({std::string_view{*encoded.metadata}}); auto values = BinaryArrayFromValues({std::string_view{*encoded.value}}); @@ -269,15 +259,15 @@ TEST(TestVariantBuilder, FromStorage) { StructArray::Make({metadata, values}, {field("metadata", binary(), false), field("value", binary(), false)})); - ASSERT_OK_AND_ASSIGN(auto array, MakeVariantArrayFromStorage(storage)); + auto array = MakeVariantArrayFromStorage(storage); ASSERT_EQ(1, array->length()); ASSERT_TRUE(array->type()->Equals(variant(storage->type()))); } TEST(TestVariantBuilder, FromShredded) { VariantBuilder value; - ASSERT_OK(value.AppendVariantNull()); - ASSERT_OK_AND_ASSIGN(auto encoded, value.Finish()); + value.AppendVariantNull(); + auto encoded = value.Finish(); auto metadata = BinaryArrayFromValues( {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); @@ -286,22 +276,21 @@ TEST(TestVariantBuilder, FromShredded) { auto storage_type = struct_({field("metadata", binary(), false), field("value", binary()), field("typed_value", int64())}); - ASSERT_OK_AND_ASSIGN( - auto array, MakeVariantArrayFromChildren(storage_type, {metadata, values, typed})); + auto array = MakeVariantArrayFromChildren(storage_type, {metadata, values, typed}); ASSERT_EQ(2, array->length()); ASSERT_TRUE(array->type()->Equals(variant(storage_type))); } TEST(TestVariantBuilder, FallbackValue) { VariantBuilder value; - ASSERT_OK(value.AppendString("x")); - ASSERT_OK_AND_ASSIGN(auto encoded, value.Finish()); + value.AppendString("x"); + auto encoded = value.Finish(); VariantValueArrayBuilder builder; - ASSERT_OK(builder.AppendNull()); - ASSERT_OK(builder.AppendEncodedValue(std::string_view{*encoded.metadata}, - std::string_view{*encoded.value})); - ASSERT_OK_AND_ASSIGN(auto array, builder.Finish()); + builder.AppendNull(); + builder.AppendEncodedValue(std::string_view{*encoded.metadata}, + std::string_view{*encoded.value}); + auto array = builder.Finish(); ASSERT_EQ(2, array->length()); ASSERT_TRUE(array->IsNull(0)); ASSERT_EQ(std::string_view{*encoded.value}, array->GetView(1)); diff --git a/cpp/src/parquet/variant/encoding.cc b/cpp/src/parquet/variant/encoding.cc index 95c61a235589..7b2c98441066 100644 --- a/cpp/src/parquet/variant/encoding.cc +++ b/cpp/src/parquet/variant/encoding.cc @@ -23,6 +23,7 @@ #include "arrow/util/endian.h" #include "arrow/util/logging_internal.h" +#include "parquet/exception.h" #include "parquet/variant/encoding_internal.h" namespace parquet::variant { @@ -38,33 +39,29 @@ uint32_t ReadLittleEndian(std::string_view data, size_t offset, size_t width) { return bit_util::FromLittleEndian(value); } -Status CheckAvailable(std::string_view data, size_t offset, size_t size, - std::string_view context) { +void CheckAvailable(std::string_view data, size_t offset, size_t size, + std::string_view context) { if (offset > data.size() || data.size() - offset < size) { - return Status::Invalid("Invalid Variant encoding: truncated ", context); + throw ParquetInvalidOrCorruptedFileException("Invalid Variant encoding: truncated ", + context); } - return Status::OK(); } -Status PrimitivePayloadSize(std::string_view value, size_t offset, - VariantPrimitiveType primitive, size_t* size) { +size_t PrimitivePayloadSize(std::string_view value, size_t offset, + VariantPrimitiveType primitive) { switch (primitive) { case VariantPrimitiveType::kNull: case VariantPrimitiveType::kBooleanTrue: case VariantPrimitiveType::kBooleanFalse: - *size = 0; - return Status::OK(); + return 0; case VariantPrimitiveType::kInt8: - *size = 1; - return Status::OK(); + return 1; case VariantPrimitiveType::kInt16: - *size = 2; - return Status::OK(); + return 2; case VariantPrimitiveType::kInt32: case VariantPrimitiveType::kDate: case VariantPrimitiveType::kFloat: - *size = 4; - return Status::OK(); + return 4; case VariantPrimitiveType::kInt64: case VariantPrimitiveType::kDouble: case VariantPrimitiveType::kTimestampMicros: @@ -72,74 +69,65 @@ Status PrimitivePayloadSize(std::string_view value, size_t offset, case VariantPrimitiveType::kTimeNTZMicros: case VariantPrimitiveType::kTimestampNanos: case VariantPrimitiveType::kTimestampNTZNanos: - *size = 8; - return Status::OK(); + return 8; case VariantPrimitiveType::kDecimal4: - *size = 5; - return Status::OK(); + return 5; case VariantPrimitiveType::kDecimal8: - *size = 9; - return Status::OK(); + return 9; case VariantPrimitiveType::kDecimal16: - *size = 17; - return Status::OK(); + return 17; case VariantPrimitiveType::kUuid: - *size = 16; - return Status::OK(); + return 16; case VariantPrimitiveType::kBinary: case VariantPrimitiveType::kString: { - ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, 4, "variable-length size")); + CheckAvailable(value, offset, 4, "variable-length size"); const uint32_t length = ReadLittleEndian(value, offset, 4); - *size = 4 + static_cast(length); - return Status::OK(); + return 4 + static_cast(length); } } - return Status::Invalid("Invalid Variant encoding: unknown primitive type"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown primitive type"); } -Status ParsePrimitive(std::string_view value, size_t offset, - VariantPrimitiveType primitive, size_t& consumed) { +size_t ParsePrimitive(std::string_view value, size_t offset, + VariantPrimitiveType primitive) { if (!internal::IsKnownVariantPrimitive(primitive)) { - return Status::Invalid("Invalid Variant encoding: unknown primitive type ", - static_cast(primitive)); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown primitive type ", static_cast(primitive)); } - size_t payload_size = 0; - ARROW_RETURN_NOT_OK(PrimitivePayloadSize(value, offset, primitive, &payload_size)); - ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, payload_size, "primitive value")); + const size_t payload_size = PrimitivePayloadSize(value, offset, primitive); + CheckAvailable(value, offset, payload_size, "primitive value"); if (internal::IsDecimalVariantPrimitive(primitive)) { const auto scale = static_cast(value[offset]); - ARROW_RETURN_NOT_OK(internal::ValidateDecimalScale(scale)); + internal::ValidateDecimalScale(scale); } if (primitive == VariantPrimitiveType::kString) { const uint32_t length = ReadLittleEndian(value, offset, 4); - ARROW_RETURN_NOT_OK(internal::ValidateUtf8(value.substr(offset + 4, length), - "primitive string value")); + internal::ValidateUtf8(value.substr(offset + 4, length), "primitive string value"); } - consumed = payload_size; - return Status::OK(); + return payload_size; } -Status ParseValue(std::string_view value, const VariantMetadataView& metadata, - size_t& consumed, VariantValueView* out); +size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, + VariantValueView* out); -Status ParseArray(std::string_view value, const VariantMetadataView& metadata, - uint8_t header, size_t& consumed, VariantValueView* out) { +size_t ParseArray(std::string_view value, const VariantMetadataView& metadata, + uint8_t header, VariantValueView* out) { const auto offset_size = static_cast((header & 0x03) + 1); const bool is_large = (header & 0x04) != 0; const size_t count_size = is_large ? 4 : 1; size_t offset = 1; - ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, count_size, "array size")); + CheckAvailable(value, offset, count_size, "array size"); const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); offset += count_size; - ARROW_RETURN_NOT_OK( - CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, - "array offsets")); + CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, + "array offsets"); std::vector offsets(num_elements + 1); for (uint32_t i = 0; i <= num_elements; ++i) { @@ -148,18 +136,19 @@ Status ParseArray(std::string_view value, const VariantMetadataView& metadata, } if (offsets[0] != 0) { - return Status::Invalid("Invalid Variant encoding: first array offset must be 0"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: first array offset must be 0"); } for (uint32_t i = 0; i < num_elements; ++i) { if (offsets[i] > offsets[i + 1]) { - return Status::Invalid("Invalid Variant encoding: array offsets must be monotonic"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: array offsets must be monotonic"); } } const size_t values_start = offset; const size_t total_value_size = offsets[num_elements]; - ARROW_RETURN_NOT_OK( - CheckAvailable(value, values_start, total_value_size, "array values")); + CheckAvailable(value, values_start, total_value_size, "array values"); size_t current = 0; std::vector array_elements; @@ -168,16 +157,14 @@ Status ParseArray(std::string_view value, const VariantMetadataView& metadata, } for (uint32_t i = 0; i < num_elements; ++i) { if (offsets[i] != current) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: array offset does not match value boundary"); } - size_t child_consumed = 0; - ARROW_RETURN_NOT_OK(ParseValue(value.substr(values_start + current), metadata, - child_consumed, - /*out=*/nullptr)); + const size_t child_consumed = + ParseValue(value.substr(values_start + current), metadata, /*out=*/nullptr); current += child_consumed; if (current != offsets[i + 1]) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: array value does not end at next offset"); } if (out != nullptr) { @@ -186,40 +173,40 @@ Status ParseArray(std::string_view value, const VariantMetadataView& metadata, } if (current != total_value_size) { - return Status::Invalid("Invalid Variant encoding: array values have trailing data"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: array values have trailing data"); } - consumed = values_start + total_value_size; + const size_t consumed = values_start + total_value_size; if (out != nullptr) { *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kArray, VariantArrayView(std::move(array_elements))); } - return Status::OK(); + return consumed; } -Status ParseObject(std::string_view value, const VariantMetadataView& metadata, - uint8_t header, size_t& consumed, VariantValueView* out) { +size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, + uint8_t header, VariantValueView* out) { const auto offset_size = static_cast((header & 0x03) + 1); const auto id_size = static_cast(((header >> 2) & 0x03) + 1); const bool is_large = (header & 0x10) != 0; const size_t count_size = is_large ? 4 : 1; size_t offset = 1; - ARROW_RETURN_NOT_OK(CheckAvailable(value, offset, count_size, "object size")); + CheckAvailable(value, offset, count_size, "object size"); const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); offset += count_size; - ARROW_RETURN_NOT_OK(CheckAvailable( - value, offset, static_cast(num_elements) * id_size, "object field ids")); + CheckAvailable(value, offset, static_cast(num_elements) * id_size, + "object field ids"); std::vector field_ids(num_elements); for (uint32_t i = 0; i < num_elements; ++i) { field_ids[i] = ReadLittleEndian(value, offset, id_size); offset += id_size; } - ARROW_RETURN_NOT_OK( - CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, - "object field offsets")); + CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, + "object field offsets"); std::vector field_offsets(num_elements + 1); for (uint32_t i = 0; i <= num_elements; ++i) { field_offsets[i] = ReadLittleEndian(value, offset, offset_size); @@ -229,11 +216,10 @@ Status ParseObject(std::string_view value, const VariantMetadataView& metadata, const size_t values_start = offset; const size_t total_value_size = field_offsets[num_elements]; if (num_elements == 0 && total_value_size != 0) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: empty object must have zero value size"); } - ARROW_RETURN_NOT_OK( - CheckAvailable(value, values_start, total_value_size, "object values")); + CheckAvailable(value, values_start, total_value_size, "object values"); std::vector object_fields; if (out != nullptr) { @@ -243,20 +229,20 @@ Status ParseObject(std::string_view value, const VariantMetadataView& metadata, std::string_view previous_name; for (uint32_t i = 0; i < num_elements; ++i) { if (field_ids[i] >= metadata.dictionary_size()) { - return Status::Invalid("Invalid Variant encoding: object field id ", field_ids[i], - " is outside metadata dictionary of size ", - metadata.dictionary_size()); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object field id ", field_ids[i], + " is outside metadata dictionary of size ", metadata.dictionary_size()); } const auto name = metadata.string(field_ids[i]); if (i > 0 && !(previous_name < name)) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: object field names must be sorted and unique"); } previous_name = name; const auto field_offset = field_offsets[i]; if (field_offset >= total_value_size) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: object field offset is outside values"); } } @@ -264,22 +250,21 @@ Status ParseObject(std::string_view value, const VariantMetadataView& metadata, std::vector value_offsets = field_offsets; std::ranges::sort(value_offsets); if (std::ranges::adjacent_find(value_offsets) != value_offsets.end()) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: object field offsets must be unique"); } if (value_offsets.front() != 0) { - return Status::Invalid("Invalid Variant encoding: object values have leading data"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object values have leading data"); } for (uint32_t i = 0; i < num_elements; ++i) { const uint32_t start = value_offsets[i]; const uint32_t end = value_offsets[i + 1]; - size_t child_consumed = 0; - ARROW_RETURN_NOT_OK(ParseValue(value.substr(values_start + start), metadata, - child_consumed, - /*out=*/nullptr)); + const size_t child_consumed = + ParseValue(value.substr(values_start + start), metadata, /*out=*/nullptr); if (child_consumed != end - start) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: object value does not end at next value boundary"); } } @@ -299,17 +284,17 @@ Status ParseObject(std::string_view value, const VariantMetadataView& metadata, } } - consumed = values_start + total_value_size; + const size_t consumed = values_start + total_value_size; if (out != nullptr) { *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kObject, VariantObjectView(std::move(object_fields))); } - return Status::OK(); + return consumed; } -Status ParseValue(std::string_view value, const VariantMetadataView& metadata, - size_t& consumed, VariantValueView* out) { - ARROW_RETURN_NOT_OK(CheckAvailable(value, 0, 1, "value header")); +size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, + VariantValueView* out) { + CheckAvailable(value, 0, 1, "value header"); const auto metadata_byte = static_cast(value[0]); const auto basic_type = static_cast(metadata_byte & 0x03); @@ -318,44 +303,43 @@ Status ParseValue(std::string_view value, const VariantMetadataView& metadata, switch (basic_type) { case VariantBasicType::kPrimitive: { const auto primitive = static_cast(header); - size_t payload_size = 0; - ARROW_RETURN_NOT_OK(ParsePrimitive(value, 1, primitive, payload_size)); - consumed = 1 + payload_size; + const size_t payload_size = ParsePrimitive(value, 1, primitive); + const size_t consumed = 1 + payload_size; if (out != nullptr) { *out = VariantValueView( value.substr(0, consumed), VariantBasicType::kPrimitive, VariantPrimitiveView(primitive, value.substr(1, payload_size))); } - return Status::OK(); + return consumed; } case VariantBasicType::kShortString: { - ARROW_RETURN_NOT_OK(CheckAvailable(value, 1, header, "short string value")); - ARROW_RETURN_NOT_OK( - internal::ValidateUtf8(value.substr(1, header), "short string value")); - consumed = 1 + header; + CheckAvailable(value, 1, header, "short string value"); + internal::ValidateUtf8(value.substr(1, header), "short string value"); + const size_t consumed = 1 + header; if (out != nullptr) { *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kShortString, VariantShortStringView(value.substr(1, header))); } - return Status::OK(); + return consumed; } case VariantBasicType::kObject: - return ParseObject(value, metadata, header, consumed, out); + return ParseObject(value, metadata, header, out); case VariantBasicType::kArray: - return ParseArray(value, metadata, header, consumed, out); + return ParseArray(value, metadata, header, out); } - return Status::Invalid("Invalid Variant encoding: unknown basic type"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown basic type"); } } // namespace -Result VariantMetadataView::Make(std::string_view metadata) { - ARROW_RETURN_NOT_OK(CheckAvailable(metadata, 0, 1, "metadata header")); +VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { + CheckAvailable(metadata, 0, 1, "metadata header"); const auto header = static_cast(metadata[0]); const auto version = static_cast(header & internal::kMetadataVersionMask); if (version != internal::kVariantVersion) { - return Status::Invalid("Invalid Variant metadata: expected version 1, got ", - static_cast(version)); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant metadata: expected version 1, got ", static_cast(version)); } VariantMetadataView view; @@ -363,14 +347,12 @@ Result VariantMetadataView::Make(std::string_view metadata) view.sorted_strings_ = (header & internal::kMetadataSortedStringsMask) != 0; view.offset_size_ = static_cast(((header >> 6) & 0x03) + 1); - ARROW_RETURN_NOT_OK( - CheckAvailable(metadata, 1, view.offset_size_, "metadata dictionary size")); + CheckAvailable(metadata, 1, view.offset_size_, "metadata dictionary size"); const uint32_t dictionary_size = ReadLittleEndian(metadata, 1, view.offset_size_); const size_t offsets_offset = 1 + view.offset_size_; - ARROW_RETURN_NOT_OK( - CheckAvailable(metadata, offsets_offset, - (static_cast(dictionary_size) + 1) * view.offset_size_, - "metadata dictionary offsets")); + CheckAvailable(metadata, offsets_offset, + (static_cast(dictionary_size) + 1) * view.offset_size_, + "metadata dictionary offsets"); std::vector offsets(dictionary_size + 1); for (uint32_t i = 0; i <= dictionary_size; ++i) { @@ -379,11 +361,12 @@ Result VariantMetadataView::Make(std::string_view metadata) } if (offsets[0] != 0) { - return Status::Invalid("Invalid Variant metadata: first dictionary offset must be 0"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant metadata: first dictionary offset must be 0"); } for (uint32_t i = 0; i < dictionary_size; ++i) { if (offsets[i] > offsets[i + 1]) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant metadata: dictionary offsets must be monotonic"); } } @@ -391,18 +374,18 @@ Result VariantMetadataView::Make(std::string_view metadata) const size_t bytes_offset = offsets_offset + (static_cast(dictionary_size) + 1) * view.offset_size_; const size_t bytes_size = offsets[dictionary_size]; - ARROW_RETURN_NOT_OK( - CheckAvailable(metadata, bytes_offset, bytes_size, "metadata dictionary bytes")); + CheckAvailable(metadata, bytes_offset, bytes_size, "metadata dictionary bytes"); if (metadata.size() != bytes_offset + bytes_size) { - return Status::Invalid("Invalid Variant metadata: trailing bytes after dictionary"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant metadata: trailing bytes after dictionary"); } view.strings_.reserve(dictionary_size); for (uint32_t i = 0; i < dictionary_size; ++i) { auto string = metadata.substr(bytes_offset + offsets[i], offsets[i + 1] - offsets[i]); - ARROW_RETURN_NOT_OK(internal::ValidateUtf8(string, "metadata dictionary string")); + internal::ValidateUtf8(string, "metadata dictionary string"); if (view.sorted_strings_ && i > 0 && !(view.strings_.back() < string)) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant metadata: sorted dictionary strings must be unique and " "lexicographically sorted"); } @@ -445,26 +428,25 @@ bool VariantObjectView::ContainsField(std::string_view name) const { return std::ranges::binary_search(fields_, name, {}, &VariantObjectField::name); } -Result VariantValueView::Make(std::string_view value, - const VariantMetadataView& metadata) { - size_t consumed = 0; +VariantValueView VariantValueView::Make(std::string_view value, + const VariantMetadataView& metadata) { VariantValueView view({}, VariantBasicType::kPrimitive, VariantPrimitiveView(VariantPrimitiveType::kNull, {})); - ARROW_RETURN_NOT_OK(ParseValue(value, metadata, consumed, &view)); + const size_t consumed = ParseValue(value, metadata, &view); if (consumed != value.size()) { - return Status::Invalid("Invalid Variant encoding: trailing bytes after value"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: trailing bytes after value"); } return view; } -Status VariantValueView::Validate(std::string_view value, - const VariantMetadataView& metadata) { - size_t consumed = 0; - ARROW_RETURN_NOT_OK(ParseValue(value, metadata, consumed, /*out=*/nullptr)); +void VariantValueView::Validate(std::string_view value, + const VariantMetadataView& metadata) { + const size_t consumed = ParseValue(value, metadata, /*out=*/nullptr); if (consumed != value.size()) { - return Status::Invalid("Invalid Variant encoding: trailing bytes after value"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: trailing bytes after value"); } - return Status::OK(); } } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/encoding.h b/cpp/src/parquet/variant/encoding.h index 1dbf3d9a104d..b24dd86a2e92 100644 --- a/cpp/src/parquet/variant/encoding.h +++ b/cpp/src/parquet/variant/encoding.h @@ -24,15 +24,10 @@ #include #include -#include "arrow/result.h" -#include "arrow/status.h" #include "parquet/platform.h" namespace parquet::variant { -using ::arrow::Result; -using ::arrow::Status; - enum class VariantBasicType : uint8_t { kPrimitive = 0, kShortString = 1, @@ -76,7 +71,7 @@ class PARQUET_EXPORT VariantMetadataView { /// /// The returned view and all string views borrowed from it are valid only while the /// input metadata bytes remain alive and unchanged. - static Result Make(std::string_view metadata); + static VariantMetadataView Make(std::string_view metadata); std::string_view metadata() const { return metadata_; } bool sorted_strings() const { return sorted_strings_; } @@ -153,9 +148,9 @@ class PARQUET_EXPORT VariantValueView { /// The returned view and any nested object/list/string views are valid only while the /// input value bytes remain alive and unchanged. The metadata view must also remain /// valid while object field names are accessed. - static Result Make(std::string_view value, - const VariantMetadataView& metadata); - static Status Validate(std::string_view value, const VariantMetadataView& metadata); + static VariantValueView Make(std::string_view value, + const VariantMetadataView& metadata); + static void Validate(std::string_view value, const VariantMetadataView& metadata); std::string_view value() const { return value_; } VariantBasicType basic_type() const { return basic_type_; } diff --git a/cpp/src/parquet/variant/encoding_internal.h b/cpp/src/parquet/variant/encoding_internal.h index f34e088a2ef3..3c647e21d5e5 100644 --- a/cpp/src/parquet/variant/encoding_internal.h +++ b/cpp/src/parquet/variant/encoding_internal.h @@ -20,14 +20,12 @@ #include #include -#include "arrow/status.h" #include "arrow/util/utf8.h" +#include "parquet/exception.h" #include "parquet/variant/encoding.h" namespace parquet::variant::internal { -using ::arrow::Status; - namespace util = ::arrow::util; inline constexpr uint8_t kVariantVersion = 1; @@ -68,19 +66,19 @@ inline bool IsDecimalVariantPrimitive(VariantPrimitiveType type) { type == VariantPrimitiveType::kDecimal16; } -inline Status ValidateDecimalScale(uint8_t scale) { +inline void ValidateDecimalScale(uint8_t scale) { if (scale > 38) { - return Status::Invalid("Invalid Variant decimal scale ", scale, " exceeds 38"); + throw ParquetInvalidOrCorruptedFileException("Invalid Variant decimal scale ", scale, + " exceeds 38"); } - return Status::OK(); } -inline Status ValidateUtf8(std::string_view value, std::string_view context) { +inline void ValidateUtf8(std::string_view value, std::string_view context) { util::InitializeUTF8(); if (!util::ValidateUTF8(reinterpret_cast(value.data()), value.size())) { - return Status::Invalid("Invalid Variant encoding: ", context, " is not valid UTF-8"); + throw ParquetInvalidOrCorruptedFileException("Invalid Variant encoding: ", context, + " is not valid UTF-8"); } - return Status::OK(); } template diff --git a/cpp/src/parquet/variant/encoding_test.cc b/cpp/src/parquet/variant/encoding_test.cc index e68b15692394..e9b1bdd8a95e 100644 --- a/cpp/src/parquet/variant/encoding_test.cc +++ b/cpp/src/parquet/variant/encoding_test.cc @@ -16,25 +16,25 @@ // under the License. #include "parquet/variant/encoding.h" +#include "parquet/exception.h" #include "parquet/variant/builder.h" #include "parquet/variant/test_util_internal.h" +#include + #include #include #include -#include "arrow/testing/gtest_util.h" - namespace parquet::variant { TEST(TestVariantEncoding, EmptyMetadata) { VariantBuilder builder; - ASSERT_OK(builder.AppendVariantNull()); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); + builder.AppendVariantNull(); + auto encoded = builder.Finish(); ASSERT_EQ(std::string_view("\x01\x00\x00", 3), std::string_view{*encoded.metadata}); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); ASSERT_FALSE(metadata.sorted_strings()); ASSERT_EQ(0, metadata.dictionary_size()); } @@ -45,73 +45,70 @@ TEST(TestVariantEncoding, Primitive) { auto check = [](auto append, VariantPrimitiveType expected) { VariantBuilder builder; - ASSERT_OK(append(builder)); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto view, internal::MakeVariantValueView(encoded)); + append(builder); + auto encoded = builder.Finish(); + auto view = internal::MakeVariantValueView(encoded); ASSERT_EQ(VariantBasicType::kPrimitive, view.basic_type()); ASSERT_EQ(expected, std::get(view.data()).type()); }; - check([](VariantBuilder& b) { return b.AppendVariantNull(); }, - VariantPrimitiveType::kNull); - check([](VariantBuilder& b) { return b.AppendBoolean(true); }, + check([](VariantBuilder& b) { b.AppendVariantNull(); }, VariantPrimitiveType::kNull); + check([](VariantBuilder& b) { b.AppendBoolean(true); }, VariantPrimitiveType::kBooleanTrue); - check([](VariantBuilder& b) { return b.AppendInt8(1); }, VariantPrimitiveType::kInt8); - check([](VariantBuilder& b) { return b.AppendInt16(2); }, VariantPrimitiveType::kInt16); - check([](VariantBuilder& b) { return b.AppendInt32(3); }, VariantPrimitiveType::kInt32); - check([](VariantBuilder& b) { return b.AppendInt64(4); }, VariantPrimitiveType::kInt64); - check([](VariantBuilder& b) { return b.AppendFloat(1.5F); }, - VariantPrimitiveType::kFloat); - check([](VariantBuilder& b) { return b.AppendDouble(2.5); }, - VariantPrimitiveType::kDouble); - check([](VariantBuilder& b) { return b.AppendDecimal4(123, 2); }, + check([](VariantBuilder& b) { b.AppendInt8(1); }, VariantPrimitiveType::kInt8); + check([](VariantBuilder& b) { b.AppendInt16(2); }, VariantPrimitiveType::kInt16); + check([](VariantBuilder& b) { b.AppendInt32(3); }, VariantPrimitiveType::kInt32); + check([](VariantBuilder& b) { b.AppendInt64(4); }, VariantPrimitiveType::kInt64); + check([](VariantBuilder& b) { b.AppendFloat(1.5F); }, VariantPrimitiveType::kFloat); + check([](VariantBuilder& b) { b.AppendDouble(2.5); }, VariantPrimitiveType::kDouble); + check([](VariantBuilder& b) { b.AppendDecimal4(123, 2); }, VariantPrimitiveType::kDecimal4); - check([](VariantBuilder& b) { return b.AppendDecimal8(123, 2); }, + check([](VariantBuilder& b) { b.AppendDecimal8(123, 2); }, VariantPrimitiveType::kDecimal8); check( [&](VariantBuilder& b) { - return b.AppendDecimal16(std::string_view(decimal16.data(), decimal16.size()), 2); + b.AppendDecimal16(std::string_view(decimal16.data(), decimal16.size()), 2); }, VariantPrimitiveType::kDecimal16); - check([](VariantBuilder& b) { return b.AppendBinary("abc"); }, - VariantPrimitiveType::kBinary); - check([](VariantBuilder& b) { return b.AppendString("abc"); }, - VariantPrimitiveType::kString); - check([](VariantBuilder& b) { return b.AppendDate(1); }, VariantPrimitiveType::kDate); - check([](VariantBuilder& b) { return b.AppendTimeNTZMicros(1); }, + check([](VariantBuilder& b) { b.AppendBinary("abc"); }, VariantPrimitiveType::kBinary); + check([](VariantBuilder& b) { b.AppendString("abc"); }, VariantPrimitiveType::kString); + check([](VariantBuilder& b) { b.AppendDate(1); }, VariantPrimitiveType::kDate); + check([](VariantBuilder& b) { b.AppendTimeNTZMicros(1); }, VariantPrimitiveType::kTimeNTZMicros); - check([](VariantBuilder& b) { return b.AppendTimestampMicros(1, true); }, + check([](VariantBuilder& b) { b.AppendTimestampMicros(1, true); }, VariantPrimitiveType::kTimestampMicros); - check([](VariantBuilder& b) { return b.AppendTimestampMicros(1, false); }, + check([](VariantBuilder& b) { b.AppendTimestampMicros(1, false); }, VariantPrimitiveType::kTimestampNTZMicros); - check([](VariantBuilder& b) { return b.AppendTimestampNanos(1, true); }, + check([](VariantBuilder& b) { b.AppendTimestampNanos(1, true); }, VariantPrimitiveType::kTimestampNanos); - check([](VariantBuilder& b) { return b.AppendTimestampNanos(1, false); }, + check([](VariantBuilder& b) { b.AppendTimestampNanos(1, false); }, VariantPrimitiveType::kTimestampNTZNanos); check( [&](VariantBuilder& b) { - return b.AppendUuid(std::string_view(uuid.data(), uuid.size())); + b.AppendUuid(std::string_view(uuid.data(), uuid.size())); }, VariantPrimitiveType::kUuid); } TEST(TestVariantEncoding, ShortString) { VariantBuilder builder; - ASSERT_OK(builder.AppendShortString("abc")); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto view, internal::MakeVariantValueView(encoded)); + builder.AppendShortString("abc"); + auto encoded = builder.Finish(); + auto view = internal::MakeVariantValueView(encoded); ASSERT_EQ(VariantBasicType::kShortString, view.basic_type()); ASSERT_EQ("abc", std::get(view.data()).string()); } TEST(TestVariantEncoding, InvalidMetadata) { // Metadata version is 2 instead of 1. - ASSERT_RAISES(Invalid, VariantMetadataView::Make(std::string("\x02\x00\x00", 3))); + ASSERT_THROW(VariantMetadataView::Make(std::string("\x02\x00\x00", 3)), + ParquetInvalidOrCorruptedFileException); // Dictionary offsets are truncated. - ASSERT_RAISES(Invalid, VariantMetadataView::Make(std::string("\x01\x01\x00", 3))); + ASSERT_THROW(VariantMetadataView::Make(std::string("\x01\x01\x00", 3)), + ParquetInvalidOrCorruptedFileException); // Dictionary string payload is not valid UTF-8. - ASSERT_RAISES(Invalid, - VariantMetadataView::Make(std::string("\x01\x01\x00\x01\xff", 5))); + ASSERT_THROW(VariantMetadataView::Make(std::string("\x01\x01\x00\x01\xff", 5)), + ParquetInvalidOrCorruptedFileException); } TEST(TestVariantEncoding, ReservedBits) { @@ -119,63 +116,63 @@ TEST(TestVariantEncoding, ReservedBits) { // - 0x21: version 1 with reserved bit 0x20 set // - 0x00: zero dictionary entries, encoded with one-byte width // - 0x00: first and final dictionary byte offset - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string("\x21\x00\x00", 3))); + auto metadata = VariantMetadataView::Make(std::string("\x21\x00\x00", 3)); ASSERT_EQ(0, metadata.dictionary_size()); // Object value bytes: // - 0x82: basic type Object with reserved header bit 0x20 set // - 0x00: zero fields // - 0x00: first and final value offset - ASSERT_OK(VariantValueView::Validate(std::string("\x82\x00\x00", 3), metadata)); + VariantValueView::Validate(std::string("\x82\x00\x00", 3), metadata); // Array value bytes: // - 0xe3: basic type Array with reserved header bits 0x38 set // - 0x00: zero elements // - 0x00: first and final value offset - ASSERT_OK(VariantValueView::Validate(std::string("\xe3\x00\x00", 3), metadata)); + VariantValueView::Validate(std::string("\xe3\x00\x00", 3), metadata); } TEST(TestVariantEncoding, ValidateValue) { VariantBuilder builder; - ASSERT_OK(builder.AppendVariantNull()); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); + builder.AppendVariantNull(); + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); - ASSERT_OK(VariantValueView::Validate(std::string_view{*encoded.value}, metadata)); + VariantValueView::Validate(std::string_view{*encoded.value}, metadata); std::string invalid_value(std::string_view{*encoded.value}); invalid_value.push_back('\0'); - ASSERT_RAISES(Invalid, VariantValueView::Validate(invalid_value, metadata)); + ASSERT_THROW(VariantValueView::Validate(invalid_value, metadata), + ParquetInvalidOrCorruptedFileException); } TEST(TestVariantEncoding, InvalidValue) { VariantBuilder builder; - ASSERT_OK(builder.AppendVariantNull()); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); + builder.AppendVariantNull(); + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); // Header byte 0x54 encodes Primitive with primitive tag 21, which is unknown. - ASSERT_RAISES( - Invalid, VariantValueView::Make(std::string(1, static_cast(0x54)), metadata)); + ASSERT_THROW(VariantValueView::Make(std::string(1, static_cast(0x54)), metadata), + ParquetInvalidOrCorruptedFileException); // Short string payload is not valid UTF-8. - ASSERT_RAISES(Invalid, VariantValueView::Make(std::string("\x05\xff", 2), metadata)); + ASSERT_THROW(VariantValueView::Make(std::string("\x05\xff", 2), metadata), + ParquetInvalidOrCorruptedFileException); // Null primitive has trailing bytes after the encoded value. - ASSERT_RAISES(Invalid, VariantValueView::Make(std::string("\x00\x00", 2), metadata)); + ASSERT_THROW(VariantValueView::Make(std::string("\x00\x00", 2), metadata), + ParquetInvalidOrCorruptedFileException); // Object references field id 0, but the metadata dictionary is empty. - ASSERT_RAISES(Invalid, VariantValueView::Make( - std::string("\x02\x01\x00\x00\x01\x00", 6), metadata)); + ASSERT_THROW( + VariantValueView::Make(std::string("\x02\x01\x00\x00\x01\x00", 6), metadata), + ParquetInvalidOrCorruptedFileException); } TEST(TestVariantEncoding, InvalidObject) { VariantBuilder builder; - ASSERT_OK_AND_ASSIGN(auto object, builder.StartObject()); - ASSERT_OK(object.AppendVariantNull("a")); - ASSERT_OK(object.AppendVariantNull("b")); - ASSERT_OK(object.Finish()); - ASSERT_OK_AND_ASSIGN(auto encoded, builder.Finish()); - ASSERT_OK_AND_ASSIGN(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); + auto object = builder.StartObject(); + object.AppendVariantNull("a"); + object.AppendVariantNull("b"); + object.Finish(); + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); // Object value bytes: // - 0x02: basic type Object, one-byte field ids, one-byte offsets, one-byte count @@ -184,7 +181,8 @@ TEST(TestVariantEncoding, InvalidObject) { // - 0x00 0x00 0x01: duplicate field start offsets 0/0, final value size 1 // - 0x00: encoded Variant null primitive value const std::string duplicate_offsets("\x02\x02\x00\x01\x00\x00\x01\x00", 8); - ASSERT_RAISES(Invalid, VariantValueView::Validate(duplicate_offsets, metadata)); + ASSERT_THROW(VariantValueView::Validate(duplicate_offsets, metadata), + ParquetInvalidOrCorruptedFileException); } } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/test_util_internal.cc b/cpp/src/parquet/variant/test_util_internal.cc index 048b20271202..064954f39af4 100644 --- a/cpp/src/parquet/variant/test_util_internal.cc +++ b/cpp/src/parquet/variant/test_util_internal.cc @@ -35,9 +35,8 @@ namespace parquet::variant::internal { -Result MakeVariantValueView(const EncodedVariantValue& encoded) { - ARROW_ASSIGN_OR_RAISE(auto metadata, - VariantMetadataView::Make(std::string_view{*encoded.metadata})); +VariantValueView MakeVariantValueView(const EncodedVariantValue& encoded) { + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); return VariantValueView::Make(std::string_view{*encoded.value}, metadata); } @@ -103,16 +102,16 @@ Result> ReadVariantTestingTable(const std::strin return reader->ReadTable(); } -Result> EmptyVariantMetadata() { +std::shared_ptr EmptyVariantMetadata() { VariantBuilder builder; - ARROW_RETURN_NOT_OK(builder.AppendVariantNull()); - ARROW_ASSIGN_OR_RAISE(auto encoded, builder.Finish()); - return encoded.metadata; + builder.AppendVariantNull(); + auto encoded = builder.Finish(); + return std::move(encoded.metadata); } -Result Int8Variant(int8_t value) { +EncodedVariantValue Int8Variant(int8_t value) { VariantBuilder builder; - ARROW_RETURN_NOT_OK(builder.AppendInt8(value)); + builder.AppendInt8(value); return builder.Finish(); } diff --git a/cpp/src/parquet/variant/test_util_internal.h b/cpp/src/parquet/variant/test_util_internal.h index f36e92d01a4b..30186c2f533c 100644 --- a/cpp/src/parquet/variant/test_util_internal.h +++ b/cpp/src/parquet/variant/test_util_internal.h @@ -38,7 +38,7 @@ using ::arrow::DataType; using ::arrow::FieldVector; using ::arrow::Result; -Result MakeVariantValueView(const EncodedVariantValue& encoded); +VariantValueView MakeVariantValueView(const EncodedVariantValue& encoded); std::shared_ptr<::arrow::Table> VariantTable( const std::shared_ptr& variant_type, @@ -60,9 +60,9 @@ std::optional ShreddedVariantTestingDir(); Result> ReadVariantTestingTable(const std::string& path); -Result> EmptyVariantMetadata(); +std::shared_ptr EmptyVariantMetadata(); -Result Int8Variant(int8_t value); +EncodedVariantValue Int8Variant(int8_t value); std::shared_ptr BinaryArrayFromValues( const std::vector>& values); diff --git a/cpp/src/parquet/variant/validate.cc b/cpp/src/parquet/variant/validate.cc index 8ff53d4e58be..e9bb5f5f47e1 100644 --- a/cpp/src/parquet/variant/validate.cc +++ b/cpp/src/parquet/variant/validate.cc @@ -30,13 +30,14 @@ #include "arrow/chunked_array.h" #include "arrow/extension/parquet_variant.h" #include "arrow/extension_type.h" -#include "arrow/result.h" +#include "arrow/status.h" #include "arrow/type.h" #include "arrow/util/bit_block_counter.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" #include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" +#include "parquet/exception.h" #include "parquet/variant/encoding.h" namespace parquet::variant { @@ -49,7 +50,6 @@ using ::arrow::DataType; using ::arrow::ExtensionArray; using ::arrow::ExtensionType; using ::arrow::MemoryPool; -using ::arrow::Result; using ::arrow::Status; using ::arrow::extension::kVariantExtensionName; using ::arrow::internal::checked_cast; @@ -66,7 +66,7 @@ struct VariantValidationPlan { std::vector children; }; -Result> ValuesArray(const Array& array) { +std::shared_ptr ValuesArray(const Array& array) { switch (array.type_id()) { case ::arrow::Type::LIST_VIEW: return checked_cast(array).values(); @@ -81,12 +81,12 @@ Result> ValuesArray(const Array& array) { case ::arrow::Type::MAP: return checked_cast(array).values(); default: - return Status::Invalid("Expected list or map storage, got ", + throw ParquetException("Expected list or map storage, got ", array.type()->ToString()); } } -Result> BuildVariantValidationPlan( +std::optional BuildVariantValidationPlan( std::shared_ptr array) { switch (array->type_id()) { case ::arrow::Type::EXTENSION: { @@ -101,8 +101,7 @@ Result> BuildVariantValidationPlan( const auto& struct_array = checked_cast(*array); std::vector children; for (auto field : struct_array.fields()) { - ARROW_ASSIGN_OR_RAISE(auto child_plan, - BuildVariantValidationPlan(std::move(field))); + auto child_plan = BuildVariantValidationPlan(std::move(field)); if (child_plan.has_value()) { children.push_back(std::move(*child_plan)); } @@ -119,9 +118,8 @@ Result> BuildVariantValidationPlan( case ::arrow::Type::LARGE_LIST: case ::arrow::Type::FIXED_SIZE_LIST: case ::arrow::Type::MAP: { - ARROW_ASSIGN_OR_RAISE(auto values, ValuesArray(*array)); - ARROW_ASSIGN_OR_RAISE(auto child_plan, - BuildVariantValidationPlan(std::move(values))); + auto values = ValuesArray(*array); + auto child_plan = BuildVariantValidationPlan(std::move(values)); if (!child_plan.has_value()) { return std::nullopt; } @@ -134,8 +132,7 @@ Result> BuildVariantValidationPlan( std::vector children; for (const auto& child_data : array->data()->child_data) { if (child_data != nullptr) { - ARROW_ASSIGN_OR_RAISE(auto child_plan, BuildVariantValidationPlan( - ::arrow::MakeArray(child_data))); + auto child_plan = BuildVariantValidationPlan(::arrow::MakeArray(child_data)); if (child_plan.has_value()) { children.push_back(std::move(*child_plan)); } @@ -150,7 +147,7 @@ Result> BuildVariantValidationPlan( } } -Result BinaryFieldView(const Array& array, int64_t row) { +std::string_view BinaryFieldView(const Array& array, int64_t row) { switch (array.type_id()) { case ::arrow::Type::BINARY: return checked_cast(array).GetView(row); @@ -159,8 +156,8 @@ Result BinaryFieldView(const Array& array, int64_t row) { case ::arrow::Type::BINARY_VIEW: return checked_cast(array).GetView(row); default: - return Status::Invalid("Expected binary Variant field, got ", - array.type()->ToString()); + throw ParquetInvalidOrCorruptedFileException("Expected binary Variant field, got ", + array.type()->ToString()); } } @@ -179,47 +176,49 @@ VariantShreddedTypeKind ShreddedTypeKind(const DataType& type) { } } -Status ValidateVariantShredding( +void ValidateVariantShredding( const VariantMetadataView& metadata, std::string_view value, bool typed_value_present, VariantShreddedTypeKind typed_value_kind, std::span shredded_field_names = {}) { if (typed_value_kind == VariantShreddedTypeKind::None) { if (value.empty()) { - return Status::OK(); + return; } - return VariantValueView::Validate(value, metadata); + VariantValueView::Validate(value, metadata); + return; } if (typed_value_kind == VariantShreddedTypeKind::Object) { for (const auto& name : shredded_field_names) { if (!metadata.FindString(name).has_value()) { - return Status::Invalid("Invalid shredded Variant: shredded field '", name, - "' is not in metadata dictionary"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: shredded field '", name, + "' is not in metadata dictionary"); } } } if (value.empty()) { - return Status::OK(); + return; } - ARROW_ASSIGN_OR_RAISE(auto value_view, VariantValueView::Make(value, metadata)); + auto value_view = VariantValueView::Make(value, metadata); switch (typed_value_kind) { case VariantShreddedTypeKind::Primitive: if (typed_value_present) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid shredded Variant: value and primitive typed_value are both " "non-null"); } break; case VariantShreddedTypeKind::Array: if (typed_value_present) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid shredded Variant: value and array typed_value are both " "non-null"); } if (value_view.basic_type() == VariantBasicType::kArray) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid shredded Variant: array value must be stored in typed_value"); } break; @@ -227,22 +226,21 @@ Status ValidateVariantShredding( default: if (value_view.basic_type() == VariantBasicType::kObject) { if (!typed_value_present) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid shredded Variant: object value requires object typed_value"); } for (const auto& name : shredded_field_names) { if (std::get(value_view.data()).ContainsField(name)) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid shredded Variant: value object contains shredded field: ", name); } } } else if (typed_value_present) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid shredded Variant: partially shredded value must be an object"); } break; } - return Status::OK(); } template @@ -251,7 +249,7 @@ std::pair ValueOffsetAndLength(const Array& array, int64_t row return {typed_array.value_offset(row), typed_array.value_length(row)}; } -Result> ValuesRangeAt(const Array& array, int64_t row) { +std::pair ValuesRangeAt(const Array& array, int64_t row) { switch (array.type_id()) { case ::arrow::Type::LIST: return ValueOffsetAndLength<::arrow::ListArray>(array, row); @@ -266,54 +264,54 @@ Result> ValuesRangeAt(const Array& array, int64_t ro case ::arrow::Type::FIXED_SIZE_LIST: return ValueOffsetAndLength<::arrow::FixedSizeListArray>(array, row); default: - return Status::Invalid("Expected list or map storage, got ", + throw ParquetException("Expected list or map storage, got ", array.type()->ToString()); } } -Status ValidateShreddedVariantSlot(const VariantMetadataView& metadata, - const std::shared_ptr& value_array, - const std::shared_ptr& typed_array, int64_t row, - bool allow_missing, std::string_view path); +void ValidateShreddedVariantSlot(const VariantMetadataView& metadata, + const std::shared_ptr& value_array, + const std::shared_ptr& typed_array, int64_t row, + bool allow_missing, std::string_view path); -Status ValidateShreddedVariantField(const VariantMetadataView& metadata, - const Array& field_array, int64_t row, - bool allow_missing, std::string_view path) { +void ValidateShreddedVariantField(const VariantMetadataView& metadata, + const Array& field_array, int64_t row, + bool allow_missing, std::string_view path) { if (field_array.type_id() != ::arrow::Type::STRUCT) { - return Status::Invalid("Invalid shredded Variant field at ", path, - ": expected struct storage, got ", - field_array.type()->ToString()); + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + path, ": expected struct storage, got ", + field_array.type()->ToString()); } if (field_array.IsNull(row)) { - return Status::Invalid("Invalid shredded Variant field at ", path, - ": field group must be required"); + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + path, ": field group must be required"); } const auto& field_struct = checked_cast(field_array); auto value_array = field_struct.GetFieldByName("value"); auto typed_array = field_struct.GetFieldByName("typed_value"); - return ValidateShreddedVariantSlot(metadata, value_array, typed_array, row, - allow_missing, path); + ValidateShreddedVariantSlot(metadata, value_array, typed_array, row, allow_missing, + path); } -Status ValidateShreddedVariantSlot(const VariantMetadataView& metadata, - const std::shared_ptr& value_array, - const std::shared_ptr& typed_array, int64_t row, - bool allow_missing, std::string_view path) { +void ValidateShreddedVariantSlot(const VariantMetadataView& metadata, + const std::shared_ptr& value_array, + const std::shared_ptr& typed_array, int64_t row, + bool allow_missing, std::string_view path) { std::string_view value; if (value_array != nullptr && !value_array->IsNull(row)) { - ARROW_ASSIGN_OR_RAISE(value, BinaryFieldView(*value_array, row)); + value = BinaryFieldView(*value_array, row); if (value.empty()) { - ARROW_RETURN_NOT_OK(VariantValueView::Validate(value, metadata)); + VariantValueView::Validate(value, metadata); } } const bool typed_value_present = typed_array != nullptr && !typed_array->IsNull(row); if (value.empty() && !typed_value_present) { if (allow_missing) { - return Status::OK(); + return; } - return Status::Invalid("Invalid shredded Variant at ", path, - ": value and typed_value are both null"); + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant at ", path, + ": value and typed_value are both null"); } const auto typed_kind = typed_array == nullptr ? VariantShreddedTypeKind::None @@ -325,54 +323,53 @@ Status ValidateShreddedVariantSlot(const VariantMetadataView& metadata, field_names.emplace_back(field->name()); } } - ARROW_RETURN_NOT_OK(ValidateVariantShredding(metadata, value, typed_value_present, - typed_kind, field_names)); + ValidateVariantShredding(metadata, value, typed_value_present, typed_kind, field_names); if (!typed_value_present) { - return Status::OK(); + return; } if (typed_kind == VariantShreddedTypeKind::Object) { const auto& typed_struct = checked_cast(*typed_array); for (int i = 0; i < typed_struct.struct_type()->num_fields(); ++i) { - ARROW_RETURN_NOT_OK(ValidateShreddedVariantField( - metadata, *typed_struct.field(i), row, /*allow_missing=*/true, - typed_struct.struct_type()->field(i)->name())); + ValidateShreddedVariantField(metadata, *typed_struct.field(i), row, + /*allow_missing=*/true, + typed_struct.struct_type()->field(i)->name()); } } else if (typed_kind == VariantShreddedTypeKind::Array) { - ARROW_ASSIGN_OR_RAISE(auto values, ValuesArray(*typed_array)); - ARROW_ASSIGN_OR_RAISE(auto range, ValuesRangeAt(*typed_array, row)); + auto values = ValuesArray(*typed_array); + auto range = ValuesRangeAt(*typed_array, row); const auto [offset, length] = range; auto elements = values->Slice(offset, length); for (int64_t i = 0; i < elements->length(); ++i) { - ARROW_RETURN_NOT_OK(ValidateShreddedVariantField(metadata, *elements, i, - /*allow_missing=*/false, path)); + ValidateShreddedVariantField(metadata, *elements, i, + /*allow_missing=*/false, path); } } - return Status::OK(); } template - requires requires(VisitVisible& visit_visible, int64_t row) { - ::arrow::ToStatus(visit_visible(row)); - } -Status VisitVisibleRows(const std::shared_ptr& valid_rows, const Array& array, - VisitVisible&& visit_visible) { +void VisitVisibleRows(const std::shared_ptr& valid_rows, const Array& array, + VisitVisible&& visit_visible) { if (valid_rows == nullptr && !array.data()->MayHaveNulls()) { for (int64_t row = 0; row < array.length(); ++row) { - ARROW_RETURN_NOT_OK(visit_visible(row)); + visit_visible(row); } - return Status::OK(); + return; } - return ::arrow::internal::VisitTwoBitBlocks( + PARQUET_THROW_NOT_OK(::arrow::internal::VisitTwoBitBlocks( valid_rows != nullptr ? valid_rows->data() : nullptr, /*left_offset=*/0, array.data()->MayHaveNulls() ? array.null_bitmap_data() : nullptr, array.offset(), - array.length(), std::forward(visit_visible), - [] { return Status::OK(); }); + array.length(), + [&](int64_t row) { + visit_visible(row); + return Status::OK(); + }, + [] { return Status::OK(); })); } -Status ValidateVariantExtensionArray(const ExtensionArray& array, - const std::shared_ptr& valid_rows) { +void ValidateVariantExtensionArray(const ExtensionArray& array, + const std::shared_ptr& valid_rows) { const auto& storage = checked_cast(*array.storage()); auto metadata_array = storage.GetFieldByName("metadata"); auto value_array = storage.GetFieldByName("value"); @@ -380,54 +377,54 @@ Status ValidateVariantExtensionArray(const ExtensionArray& array, std::string_view last_metadata_bytes; std::optional last_metadata; - return VisitVisibleRows(valid_rows, array, [&](int64_t row) { + VisitVisibleRows(valid_rows, array, [&](int64_t row) { if (metadata_array->IsNull(row)) { - return Status::Invalid("Invalid Variant extension storage: metadata is null"); + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant extension storage: metadata is null"); } - ARROW_ASSIGN_OR_RAISE(auto metadata_value, BinaryFieldView(*metadata_array, row)); + auto metadata_value = BinaryFieldView(*metadata_array, row); if (!last_metadata.has_value() || last_metadata_bytes != metadata_value) { - ARROW_ASSIGN_OR_RAISE(last_metadata, VariantMetadataView::Make(metadata_value)); + last_metadata = VariantMetadataView::Make(metadata_value); last_metadata_bytes = metadata_value; } if (typed_array == nullptr) { if (value_array->IsNull(row)) { - return Status::Invalid( + throw ParquetInvalidOrCorruptedFileException( "Invalid Variant extension storage: unshredded value is null"); } - ARROW_ASSIGN_OR_RAISE(auto value, BinaryFieldView(*value_array, row)); - ARROW_RETURN_NOT_OK(VariantValueView::Validate(value, *last_metadata)); + auto value = BinaryFieldView(*value_array, row); + VariantValueView::Validate(value, *last_metadata); } else { - ARROW_RETURN_NOT_OK( - ValidateShreddedVariantSlot(*last_metadata, value_array, typed_array, row, - /*allow_missing=*/false, "variant")); + ValidateShreddedVariantSlot(*last_metadata, value_array, typed_array, row, + /*allow_missing=*/false, "variant"); } - return Status::OK(); }); } -Status ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, - const std::shared_ptr& valid_rows) { +void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, + const std::shared_ptr& valid_rows) { switch (plan.array->type_id()) { case ::arrow::Type::EXTENSION: { - return ValidateVariantExtensionArray( - checked_cast(*plan.array), valid_rows); + ValidateVariantExtensionArray(checked_cast(*plan.array), + valid_rows); + return; } case ::arrow::Type::STRUCT: { const auto& struct_array = checked_cast(*plan.array); std::shared_ptr child_valid_rows = valid_rows; if (struct_array.data()->MayHaveNulls()) { - ARROW_ASSIGN_OR_RAISE( + PARQUET_ASSIGN_OR_THROW( child_valid_rows, ::arrow::internal::OptionalBitmapAnd( pool, valid_rows, /*left_offset=*/0, struct_array.null_bitmap(), struct_array.offset(), struct_array.length(), /*out_offset=*/0)); } for (const auto& child : plan.children) { - ARROW_RETURN_NOT_OK(ValidateVariantPlan(child, pool, child_valid_rows)); + ValidateVariantPlan(child, pool, child_valid_rows); } - return Status::OK(); + return; } case ::arrow::Type::LIST_VIEW: case ::arrow::Type::LARGE_LIST_VIEW: @@ -436,38 +433,36 @@ Status ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, case ::arrow::Type::MAP: case ::arrow::Type::FIXED_SIZE_LIST: { DCHECK_EQ(plan.children.size(), 1); - ARROW_ASSIGN_OR_RAISE(auto values, ValuesArray(*plan.array)); - ARROW_ASSIGN_OR_RAISE(auto values_valid_rows, - ::arrow::AllocateEmptyBitmap(values->length(), pool)); - ARROW_RETURN_NOT_OK( - VisitVisibleRows(valid_rows, *plan.array, [&](int64_t row) -> Status { - ARROW_ASSIGN_OR_RAISE(auto range, ValuesRangeAt(*plan.array, row)); - const auto [offset, length] = range; - ::arrow::bit_util::SetBitsTo(values_valid_rows->mutable_data(), offset, - length, true); - return Status::OK(); - })); - return ValidateVariantPlan(plan.children[0], pool, values_valid_rows); + auto values = ValuesArray(*plan.array); + std::shared_ptr values_valid_rows; + PARQUET_ASSIGN_OR_THROW(values_valid_rows, + ::arrow::AllocateEmptyBitmap(values->length(), pool)); + VisitVisibleRows(valid_rows, *plan.array, [&](int64_t row) { + const auto [offset, length] = ValuesRangeAt(*plan.array, row); + ::arrow::bit_util::SetBitsTo(values_valid_rows->mutable_data(), offset, length, + true); + }); + ValidateVariantPlan(plan.children[0], pool, values_valid_rows); + return; } default: { for (const auto& child : plan.children) { - ARROW_RETURN_NOT_OK(ValidateVariantPlan(child, pool, valid_rows)); + ValidateVariantPlan(child, pool, valid_rows); } - return Status::OK(); + return; } } } } // namespace -Status ValidateVariants(const ::arrow::ChunkedArray& data, MemoryPool* pool) { +void ValidateVariants(const ::arrow::ChunkedArray& data, MemoryPool* pool) { for (const auto& chunk : data.chunks()) { - ARROW_ASSIGN_OR_RAISE(auto plan, BuildVariantValidationPlan(chunk)); + auto plan = BuildVariantValidationPlan(chunk); if (plan.has_value()) { - ARROW_RETURN_NOT_OK(ValidateVariantPlan(*plan, pool, nullptr)); + ValidateVariantPlan(*plan, pool, nullptr); } } - return Status::OK(); } } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate.h b/cpp/src/parquet/variant/validate.h index 2afc808b9182..9abbc7bfd3b6 100644 --- a/cpp/src/parquet/variant/validate.h +++ b/cpp/src/parquet/variant/validate.h @@ -17,7 +17,6 @@ #pragma once -#include "arrow/status.h" #include "parquet/platform.h" namespace arrow { @@ -30,7 +29,6 @@ class MemoryPool; namespace parquet::variant { PARQUET_EXPORT -::arrow::Status ValidateVariants(const ::arrow::ChunkedArray& data, - ::arrow::MemoryPool* pool); +void ValidateVariants(const ::arrow::ChunkedArray& data, ::arrow::MemoryPool* pool); } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate_test.cc b/cpp/src/parquet/variant/validate_test.cc index 425cf79df137..5aa52f537392 100644 --- a/cpp/src/parquet/variant/validate_test.cc +++ b/cpp/src/parquet/variant/validate_test.cc @@ -33,6 +33,7 @@ #include "arrow/testing/gtest_util.h" #include "arrow/type.h" #include "parquet/arrow/reader.h" +#include "parquet/exception.h" #include "parquet/variant/test_util_internal.h" namespace parquet::variant { @@ -49,7 +50,7 @@ using internal::VariantTable; using internal::WriteVariantTable; TEST(TestVariantValidate, ListView) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -68,14 +69,15 @@ TEST(TestVariantValidate, ListView) { *Int32ArrayFromValues({0}), *Int32ArrayFromValues({1}), *variant_array, ::arrow::default_memory_pool())); ::arrow::ChunkedArray valid_data{valid_list}; - ASSERT_OK(ValidateVariants(valid_data, ::arrow::default_memory_pool())); + ValidateVariants(valid_data, ::arrow::default_memory_pool()); ASSERT_OK_AND_ASSIGN(auto invalid_list, ::arrow::ListViewArray::FromArrays( *Int32ArrayFromValues({1}), *Int32ArrayFromValues({1}), *variant_array, ::arrow::default_memory_pool())); ::arrow::ChunkedArray invalid_data{invalid_list}; - ASSERT_RAISES(Invalid, ValidateVariants(invalid_data, ::arrow::default_memory_pool())); + ASSERT_THROW(ValidateVariants(invalid_data, ::arrow::default_memory_pool()), + ParquetInvalidOrCorruptedFileException); } TEST(TestVariantValidate, ParquetTestingShredded) { @@ -119,12 +121,12 @@ TEST(TestVariantValidate, ParquetTestingShredded) { auto column = table->GetColumnByName("var"); ASSERT_NE(nullptr, column); - ASSERT_OK(ValidateVariants(*column, ::arrow::default_memory_pool())); + ValidateVariants(*column, ::arrow::default_memory_pool()); } } TEST(TestVariantValidate, DictionaryMetadata) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -152,11 +154,11 @@ TEST(TestVariantValidate, DictionaryMetadata) { ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); auto column = read_table->GetColumnByName("variant"); ASSERT_NE(nullptr, column); - ASSERT_OK(ValidateVariants(*column, ::arrow::default_memory_pool())); + ValidateVariants(*column, ::arrow::default_memory_pool()); } TEST(TestVariantValidate, ReadDictionaryOption) { - ASSERT_OK_AND_ASSIGN(auto encoded, Int8Variant(42)); + auto encoded = Int8Variant(42); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); @@ -184,7 +186,7 @@ TEST(TestVariantValidate, ReadDictionaryOption) { ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); auto column = read_table->GetColumnByName("variant"); ASSERT_NE(nullptr, column); - ASSERT_OK(ValidateVariants(*column, ::arrow::default_memory_pool())); + ValidateVariants(*column, ::arrow::default_memory_pool()); } } // namespace parquet::variant From 1112c59a220f0ee3137d972c2640dc751c5461cc Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Thu, 2 Jul 2026 20:45:48 +0800 Subject: [PATCH 04/10] address review --- cpp/src/parquet/variant/encoding.cc | 6 +++--- cpp/src/parquet/variant/encoding.h | 5 ++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/src/parquet/variant/encoding.cc b/cpp/src/parquet/variant/encoding.cc index 7b2c98441066..1c38951cd97a 100644 --- a/cpp/src/parquet/variant/encoding.cc +++ b/cpp/src/parquet/variant/encoding.cc @@ -395,9 +395,9 @@ VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { return view; } -std::string_view VariantMetadataView::string(uint32_t id) const { - DCHECK_LT(id, strings_.size()); - return strings_[id]; +std::string_view VariantMetadataView::string(uint32_t field_id) const { + DCHECK_LT(field_id, strings_.size()); + return strings_[field_id]; } std::optional VariantMetadataView::FindString(std::string_view value) const { diff --git a/cpp/src/parquet/variant/encoding.h b/cpp/src/parquet/variant/encoding.h index b24dd86a2e92..ca7e93b034c8 100644 --- a/cpp/src/parquet/variant/encoding.h +++ b/cpp/src/parquet/variant/encoding.h @@ -76,9 +76,12 @@ class PARQUET_EXPORT VariantMetadataView { std::string_view metadata() const { return metadata_; } bool sorted_strings() const { return sorted_strings_; } uint8_t offset_size() const { return offset_size_; } + /// Number of strings in the metadata dictionary used by object field ids. uint32_t dictionary_size() const { return static_cast(strings_.size()); } - std::string_view string(uint32_t id) const; + /// Return the metadata dictionary string for an object field id. + std::string_view string(uint32_t field_id) const; + /// Return the metadata dictionary id for a string, if present. std::optional FindString(std::string_view value) const; private: From 6c2006b3a03c360e94e40fb41d8a19d52797341f Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Mon, 6 Jul 2026 11:49:08 +0800 Subject: [PATCH 05/10] reject to write typed-only variant --- cpp/src/parquet/arrow/arrow_schema_test.cc | 22 ++++++++++++++++++++++ cpp/src/parquet/arrow/schema.cc | 7 +++++++ cpp/src/parquet/arrow/variant_test.cc | 4 ++-- 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_schema_test.cc b/cpp/src/parquet/arrow/arrow_schema_test.cc index f8196240648f..603ab3333601 100644 --- a/cpp/src/parquet/arrow/arrow_schema_test.cc +++ b/cpp/src/parquet/arrow/arrow_schema_test.cc @@ -1083,6 +1083,19 @@ TEST_F(TestConvertParquetSchema, ParquetVariantShredded) { "variant_shredded", {metadata, value, typed_value}, storage_type)); } +TEST_F(TestConvertParquetSchema, ParquetVariantShreddedWithoutValue) { + auto metadata = + PrimitiveNode::Make("metadata", Repetition::REQUIRED, ParquetType::BYTE_ARRAY); + auto typed_value = + PrimitiveNode::Make("typed_value", Repetition::OPTIONAL, ParquetType::INT64); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("typed_value", ::arrow::int64())}); + + ASSERT_NO_FATAL_FAILURE(CheckParquetVariantSchema( + "variant_shredded", {metadata, typed_value}, storage_type)); +} + TEST_F(TestConvertParquetSchema, ParquetSchemaArrowJsonExtension) { std::vector parquet_fields; parquet_fields.push_back(PrimitiveNode::Make( @@ -1577,6 +1590,15 @@ TEST_F(TestConvertArrowSchema, ParquetVariantShredded) { ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); } +TEST_F(TestConvertArrowSchema, ParquetVariantShreddedWithoutValue) { + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("typed_value", ::arrow::int64())}); + auto variant_type = ::arrow::extension::variant(storage_type); + + ASSERT_RAISES(Invalid, ConvertSchema({::arrow::field("variant", variant_type)})); +} + TEST_F(TestConvertArrowSchema, ParquetVariantDictionary) { auto storage_type = ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index 9d7d71d64aa6..1c5296206c25 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -370,6 +370,13 @@ Status VariantToNode(const std::shared_ptr& type, const std::string& name, bool nullable, int field_id, const WriterProperties& properties, const ArrowWriterProperties& arrow_properties, NodePtr* out) { + // `VariantExtensionType` already validates the general storage shape. Writing is + // stricter than reading: we keep accepting typed-only Variant schemas on read, + // but must not write them. + if (type->value() == nullptr) { + return Status::Invalid("Invalid Variant write schema: missing top-level value field"); + } + std::vector children; children.reserve(type->storage_type()->num_fields()); diff --git a/cpp/src/parquet/arrow/variant_test.cc b/cpp/src/parquet/arrow/variant_test.cc index 95b1ebad5c26..981997e14152 100644 --- a/cpp/src/parquet/arrow/variant_test.cc +++ b/cpp/src/parquet/arrow/variant_test.cc @@ -196,7 +196,7 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { ASSERT_OK(WriteVariantTable(valid_table)); } -TEST(TestVariantExtensionType, WriterValidatesShreddedWithoutValue) { +TEST(TestVariantExtensionType, WriterRejectsShreddedWithoutValue) { auto metadata = EmptyVariantMetadata(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("typed_value", ::arrow::int64())}); @@ -207,7 +207,7 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedWithoutValue) { auto typed_array = Int64ArrayFromValues({34}); auto table = VariantTable(variant_type, {metadata_array, typed_array}, storage_type->fields()); - ASSERT_OK(WriteVariantTable(table)); + ASSERT_RAISES(Invalid, WriteVariantTable(table)); } TEST(TestVariantExtensionType, ReadsDictionaryEncodedMetadata) { From 75a30531f1f4fc801bbf7eaeff5c6a6778fd1a09 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Wed, 15 Jul 2026 00:06:05 +0800 Subject: [PATCH 06/10] avoid buffer copy --- cpp/src/parquet/variant/builder.cc | 320 +++++++++++++++++------- cpp/src/parquet/variant/builder.h | 23 +- cpp/src/parquet/variant/builder_test.cc | 47 ++-- 3 files changed, 263 insertions(+), 127 deletions(-) diff --git a/cpp/src/parquet/variant/builder.cc b/cpp/src/parquet/variant/builder.cc index e167c4f7cbc0..8fdeffd6855c 100644 --- a/cpp/src/parquet/variant/builder.cc +++ b/cpp/src/parquet/variant/builder.cc @@ -28,8 +28,11 @@ #include #include +#include "arrow/array/array_binary.h" #include "arrow/buffer_builder.h" #include "arrow/builder.h" // IWYU pragma: keep +#include "arrow/type.h" +#include "arrow/util/binary_view_util.h" #include "arrow/util/checked_cast.h" #include "arrow/util/endian.h" #include "arrow/util/logging_internal.h" @@ -39,16 +42,17 @@ namespace parquet::variant { -using ::arrow::binary; -using ::arrow::BinaryBuilder; -using ::arrow::BooleanArray; -using ::arrow::BooleanBuilder; +using ::arrow::binary_view; +using ::arrow::BinaryViewArray; +using ::arrow::BinaryViewType; using ::arrow::BufferBuilder; +using ::arrow::BufferVector; using ::arrow::ExtensionType; using ::arrow::field; using ::arrow::struct_; using ::arrow::StructType; using ::arrow::Type; +using ::arrow::TypedBufferBuilder; using ::arrow::extension::VariantExtensionType; namespace bit_util = ::arrow::bit_util; @@ -142,7 +146,7 @@ class VariantMetadataBuilder { RebuildIndex(); } - std::shared_ptr Finish() const { + void AppendTo(BufferBuilder& out) const { uint64_t bytes_size = 0; for (const auto& string : field_names_) { bytes_size += string.size(); @@ -154,7 +158,6 @@ class VariantMetadataBuilder { const uint8_t offset_size = WidthForValue(std::max(field_names_.size(), bytes_size)); - BufferBuilder out(pool_); PARQUET_THROW_NOT_OK(out.Reserve( 1 + offset_size + (field_names_.size() + 1) * offset_size + bytes_size)); @@ -175,6 +178,11 @@ class VariantMetadataBuilder { for (const auto& string : field_names_) { PARQUET_THROW_NOT_OK(out.Append(string)); } + } + + std::shared_ptr Finish() const { + BufferBuilder out(pool_); + AppendTo(out); std::shared_ptr buffer; PARQUET_ASSIGN_OR_THROW(buffer, out.Finish()); return buffer; @@ -317,11 +325,11 @@ struct VariantBuildFrame { }; struct VariantBuildState { - explicit VariantBuildState(MemoryPool* pool) - : pool(pool), value(pool), metadata(pool) {} + VariantBuildState(MemoryPool* pool, BufferBuilder& value) + : value(value), root_value_start(value.length()), metadata(pool) {} - MemoryPool* pool; - BufferBuilder value; + BufferBuilder& value; + int64_t root_value_start = 0; VariantMetadataBuilder metadata; std::vector frames; bool root_has_value = false; @@ -435,7 +443,7 @@ std::string BuildListHeader(const VariantBuildFrame& frame, uint32_t values_size return out; } -using RootFinishCallback = std::function; +using RootFinishCallback = std::function&)>; void FinishFrame(const std::shared_ptr& state, size_t frame_index, VariantContainerKind kind, const RootFinishCallback& callback) { @@ -465,10 +473,7 @@ void FinishFrame(const std::shared_ptr& state, size_t frame_i return; } - auto metadata = state->metadata.Finish(); - std::shared_ptr value; - PARQUET_ASSIGN_OR_THROW(value, state->value.Finish()); - callback({.metadata = std::move(metadata), .value = std::move(value)}); + callback(state); } template @@ -572,6 +577,95 @@ void AppendListPrimitive(const std::shared_ptr& state, }); } +// Local helper for array builders that owns one shared byte arena plus the +// corresponding `binary_view` slots. This lets row commit write bytes once, +// create inline or out-of-line views in place, and roll back both buffers +// together on failure without routing hot-path appends through `BinaryViewBuilder`. +class BinaryViewColumnBuilder { + public: + struct Mark { + int64_t bytes_length = 0; + int64_t views_length = 0; + bool has_out_of_line_data = false; + }; + + explicit BinaryViewColumnBuilder(MemoryPool* pool) : bytes_(pool), views_(pool) {} + + [[nodiscard]] Mark mark() const { + return Mark{.bytes_length = bytes_.length(), + .views_length = views_.length(), + .has_out_of_line_data = has_out_of_line_data_}; + } + + void Rollback(const Mark& mark) { + bytes_.Rewind(mark.bytes_length); + views_.bytes_builder()->Rewind(mark.views_length * sizeof(BinaryViewType::c_type)); + has_out_of_line_data_ = mark.has_out_of_line_data; + } + + int64_t length() const { return views_.length(); } + + std::string_view SliceFrom(int64_t start) const { + DCHECK_LE(start, bytes_.length()); + const auto size = bytes_.length() - start; + DCHECK_GE(size, 0); + if (size == 0) { + return std::string_view{}; + } + return std::string_view{reinterpret_cast(bytes_.data() + start), + static_cast(size)}; + } + + void AppendView(int64_t start) { + const auto slice = SliceFrom(start); + if (slice.size() > static_cast(std::numeric_limits::max())) { + throw ParquetException("Binary view value is too large"); + } + PARQUET_THROW_NOT_OK(views_.Reserve(1)); + if (slice.size() <= BinaryViewType::kInlineSize) { + views_.UnsafeAppend(::arrow::util::ToInlineBinaryView(slice)); + bytes_.Rewind(start); + return; + } + if (start > std::numeric_limits::max()) { + throw ParquetException("Binary view offset is too large"); + } + has_out_of_line_data_ = true; + views_.UnsafeAppend(::arrow::util::ToNonInlineBinaryView( + slice.data(), static_cast(slice.size()), + /*buffer_index=*/0, static_cast(start))); + } + + void AppendEmptyValue() { + PARQUET_THROW_NOT_OK(views_.Reserve(1)); + views_.UnsafeAppend(BinaryViewType::c_type{}); + } + + BufferBuilder& bytes_builder() { return bytes_; } + + std::shared_ptr Finish(std::shared_ptr null_bitmap = nullptr, + int64_t null_count = 0) { + const auto array_length = views_.length(); + std::shared_ptr views_buffer; + PARQUET_THROW_NOT_OK(views_.Finish(&views_buffer)); + + std::shared_ptr data_buffer; + PARQUET_ASSIGN_OR_THROW(data_buffer, bytes_.Finish()); + BufferVector data_buffers; + if (has_out_of_line_data_) { + data_buffers.push_back(std::move(data_buffer)); + } + return std::make_shared( + binary_view(), array_length, std::move(views_buffer), std::move(data_buffers), + std::move(null_bitmap), null_count); + } + + private: + BufferBuilder bytes_; + TypedBufferBuilder views_; + bool has_out_of_line_data_ = false; +}; + } // namespace namespace internal { @@ -593,9 +687,17 @@ struct NestedVariantBuilderImpl { struct VariantBuilder::Impl { explicit Impl(MemoryPool* pool) - : pool(pool), state(std::make_shared(pool)) {} + : pool(pool), + value_builder(pool), + state(std::make_shared(pool, value_builder)) {} + + void Reset() { + value_builder.Reset(); + state = std::make_shared(pool, value_builder); + } MemoryPool* pool; + BufferBuilder value_builder; std::shared_ptr state; }; @@ -716,9 +818,7 @@ EncodedVariantValue VariantBuilder::Finish() { return out; } -void VariantBuilder::Reset() { - impl_->state = std::make_shared(impl_->pool); -} +void VariantBuilder::Reset() { impl_->Reset(); } VariantObjectBuilder::VariantObjectBuilder( std::unique_ptr impl) @@ -1019,16 +1119,13 @@ void VariantListBuilder::Finish() { struct VariantArrayBuilder::Impl { explicit Impl(MemoryPool* pool) - : pool(pool), metadata_builder(pool), value_builder(pool), validity_builder(pool) {} + : pool(pool), metadata_column(pool), value_column(pool), validity_builder(pool) {} template void AppendValue(Write&& write) { - auto state = std::make_shared(pool); + auto state = std::make_shared(pool, value_column.bytes_builder()); AppendRootPrimitiveWith(state, std::forward(write)); - auto metadata = state->metadata.Finish(); - std::shared_ptr value; - PARQUET_ASSIGN_OR_THROW(value, state->value.Finish()); - AppendEncoded({.metadata = std::move(metadata), .value = std::move(value)}); + CommitBuiltRow(state); } template @@ -1047,17 +1144,104 @@ struct VariantArrayBuilder::Impl { throw ParquetException( "Encoded Variant metadata and value buffers must be non-null"); } - auto metadata = VariantMetadataView::Make(std::string_view{*value.metadata}); - VariantValueView::Validate(std::string_view{*value.value}, metadata); - PARQUET_THROW_NOT_OK(metadata_builder.Append(std::string_view{*value.metadata})); - PARQUET_THROW_NOT_OK(value_builder.Append(std::string_view{*value.value})); - PARQUET_THROW_NOT_OK(validity_builder.Append(true)); + const auto metadata_bytes = std::string_view{*value.metadata}; + const auto value_bytes = std::string_view{*value.value}; + auto metadata = VariantMetadataView::Make(metadata_bytes); + VariantValueView::Validate(value_bytes, metadata); + CommitEncodedRow(metadata_bytes, value_bytes); + } + + void AppendNull() { + const auto metadata_mark = metadata_column.mark(); + const auto value_mark = value_column.mark(); + try { + metadata_column.AppendEmptyValue(); + value_column.AppendEmptyValue(); + PARQUET_THROW_NOT_OK(validity_builder.Append(false)); + } catch (...) { + metadata_column.Rollback(metadata_mark); + value_column.Rollback(value_mark); + throw; + } + } + + void CommitBuiltRow(const std::shared_ptr& state) { + if (!state->frames.empty()) { + throw ParquetException("Cannot commit Variant row with active containers"); + } + if (!state->root_has_value) { + throw ParquetException("Cannot commit empty Variant row"); + } + + const auto value_start = state->root_value_start; + const auto metadata_mark = metadata_column.mark(); + auto value_mark = value_column.mark(); + DCHECK_LE(value_start, value_mark.bytes_length); + value_mark.bytes_length = value_start; + try { + const auto metadata_start = metadata_column.bytes_builder().length(); + state->metadata.AppendTo(metadata_column.bytes_builder()); + const auto metadata_bytes = metadata_column.SliceFrom(metadata_start); + auto metadata = VariantMetadataView::Make(metadata_bytes); + + const auto value_bytes = value_column.SliceFrom(value_start); + VariantValueView::Validate(value_bytes, metadata); + + metadata_column.AppendView(metadata_start); + value_column.AppendView(value_start); + PARQUET_THROW_NOT_OK(validity_builder.Append(true)); + } catch (...) { + metadata_column.Rollback(metadata_mark); + value_column.Rollback(value_mark); + throw; + } + } + + std::shared_ptr Finish() { + DCHECK_EQ(metadata_column.length(), value_column.length()); + DCHECK_EQ(metadata_column.length(), validity_builder.length()); + + const int64_t null_count = validity_builder.false_count(); + std::shared_ptr null_bitmap; + if (null_count > 0) { + PARQUET_THROW_NOT_OK(validity_builder.Finish(&null_bitmap)); + } + + auto metadata = metadata_column.Finish(); + auto value = value_column.Finish(); + auto storage_type = struct_({field("metadata", binary_view(), /*nullable=*/false), + field("value", binary_view(), /*nullable=*/false)}); + std::shared_ptr storage; + PARQUET_ASSIGN_OR_THROW(storage, + StructArray::Make({metadata, value}, storage_type->fields(), + std::move(null_bitmap), null_count)); + return MakeVariantArrayFromStorage(storage); } MemoryPool* pool; - BinaryBuilder metadata_builder; - BinaryBuilder value_builder; - BooleanBuilder validity_builder; + BinaryViewColumnBuilder metadata_column; + BinaryViewColumnBuilder value_column; + TypedBufferBuilder validity_builder; + + private: + void CommitEncodedRow(std::string_view metadata_bytes, std::string_view value_bytes) { + const auto metadata_mark = metadata_column.mark(); + const auto value_mark = value_column.mark(); + try { + const auto metadata_start = metadata_column.bytes_builder().length(); + PARQUET_THROW_NOT_OK(metadata_column.bytes_builder().Append(metadata_bytes)); + const auto value_start = value_column.bytes_builder().length(); + PARQUET_THROW_NOT_OK(value_column.bytes_builder().Append(value_bytes)); + + metadata_column.AppendView(metadata_start); + value_column.AppendView(value_start); + PARQUET_THROW_NOT_OK(validity_builder.Append(true)); + } catch (...) { + metadata_column.Rollback(metadata_mark); + value_column.Rollback(value_mark); + throw; + } + } }; VariantArrayBuilder::VariantArrayBuilder(MemoryPool* pool) @@ -1067,11 +1251,7 @@ VariantArrayBuilder::VariantArrayBuilder(VariantArrayBuilder&&) noexcept = defau VariantArrayBuilder& VariantArrayBuilder::operator=(VariantArrayBuilder&&) noexcept = default; -void VariantArrayBuilder::AppendNull() { - PARQUET_THROW_NOT_OK(impl_->metadata_builder.Append("")); - PARQUET_THROW_NOT_OK(impl_->value_builder.Append("")); - PARQUET_THROW_NOT_OK(impl_->validity_builder.Append(false)); -} +void VariantArrayBuilder::AppendNull() { impl_->AppendNull(); } void VariantArrayBuilder::AppendVariantNull() { impl_->AppendPrimitive(); @@ -1143,77 +1323,35 @@ void VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { } VariantObjectBuilder VariantArrayBuilder::StartObject() { - auto state = std::make_shared(impl_->pool); + auto state = std::make_shared(impl_->pool, + impl_->value_column.bytes_builder()); state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::Object, .value_start = state->value.length(), .metadata_size = state->metadata.size()}); - auto callback = [this](EncodedVariantValue encoded) { impl_->AppendEncoded(encoded); }; return VariantObjectBuilder(std::make_unique( - state, state->frames.size() - 1, std::move(callback))); + state, state->frames.size() - 1, + std::bind_front(&Impl::CommitBuiltRow, impl_.get()))); } VariantListBuilder VariantArrayBuilder::StartList() { - auto state = std::make_shared(impl_->pool); + auto state = std::make_shared(impl_->pool, + impl_->value_column.bytes_builder()); state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::List, .value_start = state->value.length(), .metadata_size = state->metadata.size()}); - auto callback = [this](EncodedVariantValue encoded) { impl_->AppendEncoded(encoded); }; return VariantListBuilder(std::make_unique( - state, state->frames.size() - 1, std::move(callback))); + state, state->frames.size() - 1, + std::bind_front(&Impl::CommitBuiltRow, impl_.get()))); } std::shared_ptr VariantArrayBuilder::Finish() { - std::shared_ptr metadata; - std::shared_ptr value; - std::shared_ptr validity; - PARQUET_THROW_NOT_OK(impl_->metadata_builder.Finish(&metadata)); - PARQUET_THROW_NOT_OK(impl_->value_builder.Finish(&value)); - PARQUET_THROW_NOT_OK(impl_->validity_builder.Finish(&validity)); - - auto null_bitmap = validity->data()->buffers[1]; - const int64_t null_count = validity->false_count(); - auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary(), /*nullable=*/false)}); - std::shared_ptr storage; - PARQUET_ASSIGN_OR_THROW( - storage, StructArray::Make({metadata, value}, storage_type->fields(), null_bitmap, - null_count)); - return MakeVariantArrayFromStorage(storage); + auto out = impl_->Finish(); + Reset(); + return out; } void VariantArrayBuilder::Reset() { impl_ = std::make_unique(impl_->pool); } -struct VariantValueArrayBuilder::Impl { - explicit Impl(MemoryPool* pool) : value_builder(pool) {} - - BinaryBuilder value_builder; -}; - -VariantValueArrayBuilder::VariantValueArrayBuilder(MemoryPool* pool) - : impl_(std::make_unique(pool)) {} -VariantValueArrayBuilder::VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept = - default; -VariantValueArrayBuilder& VariantValueArrayBuilder::operator=( - VariantValueArrayBuilder&&) noexcept = default; -VariantValueArrayBuilder::~VariantValueArrayBuilder() = default; - -void VariantValueArrayBuilder::AppendNull() { - PARQUET_THROW_NOT_OK(impl_->value_builder.AppendNull()); -} - -void VariantValueArrayBuilder::AppendEncodedValue(std::string_view metadata, - std::string_view value) { - auto metadata_view = VariantMetadataView::Make(metadata); - VariantValueView::Validate(value, metadata_view); - PARQUET_THROW_NOT_OK(impl_->value_builder.Append(value)); -} - -std::shared_ptr VariantValueArrayBuilder::Finish() { - std::shared_ptr out; - PARQUET_THROW_NOT_OK(impl_->value_builder.Finish(&out)); - return out; -} - std::shared_ptr MakeVariantArrayFromStorage( std::shared_ptr storage) { if (storage == nullptr) { diff --git a/cpp/src/parquet/variant/builder.h b/cpp/src/parquet/variant/builder.h index 1fd0a03843e0..d45a4a77bcfd 100644 --- a/cpp/src/parquet/variant/builder.h +++ b/cpp/src/parquet/variant/builder.h @@ -31,7 +31,7 @@ namespace parquet::variant { using ::arrow::Array; -using ::arrow::BinaryArray; +using ::arrow::BinaryViewArray; using ::arrow::Buffer; using ::arrow::DataType; using ::arrow::MemoryPool; @@ -217,6 +217,9 @@ class PARQUET_EXPORT VariantArrayBuilder { void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); void AppendEncoded(const EncodedVariantValue& value); + /// The returned nested builder borrows this builder's internal buffers. Finish or + /// destroy it before calling any other VariantArrayBuilder method, including Finish() + /// or Reset(). VariantObjectBuilder StartObject(); VariantListBuilder StartList(); @@ -228,24 +231,6 @@ class PARQUET_EXPORT VariantArrayBuilder { std::unique_ptr impl_; }; -class PARQUET_EXPORT VariantValueArrayBuilder { - public: - explicit VariantValueArrayBuilder(MemoryPool* pool = ::arrow::default_memory_pool()); - ~VariantValueArrayBuilder(); - VariantValueArrayBuilder(const VariantValueArrayBuilder&) = delete; - VariantValueArrayBuilder& operator=(const VariantValueArrayBuilder&) = delete; - VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept; - VariantValueArrayBuilder& operator=(VariantValueArrayBuilder&&) noexcept; - - void AppendNull(); - void AppendEncodedValue(std::string_view metadata, std::string_view value); - std::shared_ptr Finish(); - - private: - struct Impl; - std::unique_ptr impl_; -}; - PARQUET_EXPORT std::shared_ptr MakeVariantArrayFromStorage( std::shared_ptr storage); diff --git a/cpp/src/parquet/variant/builder_test.cc b/cpp/src/parquet/variant/builder_test.cc index b3f069dde7ce..54471d9a3b6f 100644 --- a/cpp/src/parquet/variant/builder_test.cc +++ b/cpp/src/parquet/variant/builder_test.cc @@ -33,6 +33,7 @@ namespace { using ::arrow::ArrayFromJSON; using ::arrow::binary; +using ::arrow::BinaryViewArray; using ::arrow::default_memory_pool; using ::arrow::field; using ::arrow::int64; @@ -243,8 +244,35 @@ TEST(TestVariantBuilder, ArrayBuilder) { ::arrow::internal::checked_cast(*array->storage()); ASSERT_FALSE(storage.type()->field(0)->nullable()); ASSERT_FALSE(storage.type()->field(1)->nullable()); - ASSERT_EQ(Type::BINARY, storage.field(0)->type_id()); - ASSERT_EQ(Type::BINARY, storage.field(1)->type_id()); + ASSERT_EQ(Type::BINARY_VIEW, storage.field(0)->type_id()); + ASSERT_EQ(Type::BINARY_VIEW, storage.field(1)->type_id()); +} + +TEST(TestVariantBuilder, ArrayBuilderViews) { + VariantBuilder value; + value.AppendString("x"); + auto short_encoded = value.Finish(); + + VariantArrayBuilder builder; + builder.AppendEncoded(short_encoded); + auto object = builder.StartObject(); + object.AppendString("abcdefghijklm", std::string(32, 'y')); + object.Finish(); + + auto array = builder.Finish(); + const auto& storage = + ::arrow::internal::checked_cast(*array->storage()); + const auto& metadata = + ::arrow::internal::checked_cast(*storage.field(0)); + const auto& values = + ::arrow::internal::checked_cast(*storage.field(1)); + + auto* metadata_views = metadata.data()->GetValues<::arrow::BinaryViewType::c_type>(1); + auto* value_views = values.data()->GetValues<::arrow::BinaryViewType::c_type>(1); + ASSERT_TRUE(metadata_views[0].is_inline()); + ASSERT_TRUE(value_views[0].is_inline()); + ASSERT_FALSE(metadata_views[1].is_inline()); + ASSERT_FALSE(value_views[1].is_inline()); } TEST(TestVariantBuilder, FromStorage) { @@ -281,19 +309,4 @@ TEST(TestVariantBuilder, FromShredded) { ASSERT_TRUE(array->type()->Equals(variant(storage_type))); } -TEST(TestVariantBuilder, FallbackValue) { - VariantBuilder value; - value.AppendString("x"); - auto encoded = value.Finish(); - - VariantValueArrayBuilder builder; - builder.AppendNull(); - builder.AppendEncodedValue(std::string_view{*encoded.metadata}, - std::string_view{*encoded.value}); - auto array = builder.Finish(); - ASSERT_EQ(2, array->length()); - ASSERT_TRUE(array->IsNull(0)); - ASSERT_EQ(std::string_view{*encoded.value}, array->GetView(1)); -} - } // namespace parquet::variant From 663e0f72b3e1a06dff3749cdebac0cf6c76795e5 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Tue, 21 Jul 2026 20:17:12 +0800 Subject: [PATCH 07/10] add unshred and shredded reader and roundtrip --- cpp/src/arrow/extension/parquet_variant.cc | 137 +- cpp/src/parquet/CMakeLists.txt | 5 +- cpp/src/parquet/arrow/arrow_schema_test.cc | 44 +- cpp/src/parquet/arrow/reader.cc | 32 +- cpp/src/parquet/arrow/variant_test.cc | 66 +- cpp/src/parquet/arrow/writer.cc | 5 +- cpp/src/parquet/meson.build | 5 +- cpp/src/parquet/properties.h | 12 + cpp/src/parquet/variant/CMakeLists.txt | 4 +- .../parquet/variant/append_table_internal.h | 41 + cpp/src/parquet/variant/array_internal.cc | 90 ++ cpp/src/parquet/variant/array_internal.h | 53 + .../binary_view_column_builder_internal.h | 129 ++ cpp/src/parquet/variant/builder.cc | 1356 ++++++++--------- cpp/src/parquet/variant/builder.h | 124 +- cpp/src/parquet/variant/builder_test.cc | 381 ++++- .../variant/{encoding.cc => decoding.cc} | 223 ++- .../variant/{encoding.h => decoding.h} | 114 +- .../{encoding_test.cc => decoding_test.cc} | 116 +- cpp/src/parquet/variant/format.h | 55 + ...{encoding_internal.h => format_internal.h} | 29 +- cpp/src/parquet/variant/meson.build | 6 +- .../parquet/variant/parquet_testing_test.cc | 474 ++++++ cpp/src/parquet/variant/slot_internal.cc | 703 +++++++++ cpp/src/parquet/variant/slot_internal.h | 130 ++ cpp/src/parquet/variant/test_util_internal.cc | 102 +- cpp/src/parquet/variant/test_util_internal.h | 21 - cpp/src/parquet/variant/type_test.cc | 75 +- cpp/src/parquet/variant/unshred.cc | 141 ++ cpp/src/parquet/variant/unshred.h | 51 + cpp/src/parquet/variant/unshred_test.cc | 355 +++++ cpp/src/parquet/variant/validate.cc | 378 +---- cpp/src/parquet/variant/validate.h | 17 +- cpp/src/parquet/variant/validate_test.cc | 77 +- 34 files changed, 3934 insertions(+), 1617 deletions(-) create mode 100644 cpp/src/parquet/variant/append_table_internal.h create mode 100644 cpp/src/parquet/variant/array_internal.cc create mode 100644 cpp/src/parquet/variant/array_internal.h create mode 100644 cpp/src/parquet/variant/binary_view_column_builder_internal.h rename cpp/src/parquet/variant/{encoding.cc => decoding.cc} (70%) rename cpp/src/parquet/variant/{encoding.h => decoding.h} (59%) rename cpp/src/parquet/variant/{encoding_test.cc => decoding_test.cc} (60%) create mode 100644 cpp/src/parquet/variant/format.h rename cpp/src/parquet/variant/{encoding_internal.h => format_internal.h} (87%) create mode 100644 cpp/src/parquet/variant/parquet_testing_test.cc create mode 100644 cpp/src/parquet/variant/slot_internal.cc create mode 100644 cpp/src/parquet/variant/slot_internal.h create mode 100644 cpp/src/parquet/variant/unshred.cc create mode 100644 cpp/src/parquet/variant/unshred.h create mode 100644 cpp/src/parquet/variant/unshred_test.cc diff --git a/cpp/src/arrow/extension/parquet_variant.cc b/cpp/src/arrow/extension/parquet_variant.cc index 4630c1d39321..3d865fbe2ccc 100644 --- a/cpp/src/arrow/extension/parquet_variant.cc +++ b/cpp/src/arrow/extension/parquet_variant.cc @@ -29,44 +29,6 @@ namespace arrow::extension { -VariantExtensionType::VariantExtensionType(const std::shared_ptr& storage_type) - : ExtensionType(storage_type) { - // IsSupportedStorageType should have been called already, asserting that - // metadata is present and at least one of value / typed_value is present. - for (const auto& field : storage_type->fields()) { - if (field->name() == "metadata") { - metadata_ = field; - } else if (field->name() == "value") { - value_ = field; - } else if (field->name() == "typed_value") { - typed_value_ = field; - } - } -} - -bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const { - return other.extension_name() == this->extension_name() && - other.storage_type()->Equals(this->storage_type()); -} - -Result> VariantExtensionType::Deserialize( - std::shared_ptr storage_type, const std::string& serialized) const { - if (!serialized.empty()) { - return Status::Invalid("Unexpected serialized metadata: '", serialized, "'"); - } - return VariantExtensionType::Make(std::move(storage_type)); -} - -std::string VariantExtensionType::Serialize() const { return ""; } - -std::shared_ptr VariantExtensionType::MakeArray( - std::shared_ptr data) const { - DCHECK_EQ(data->type->id(), Type::EXTENSION); - DCHECK_EQ(kVariantExtensionName, - internal::checked_cast(*data->type).extension_name()); - return std::make_shared(data); -} - namespace { bool IsSupportedPrimitiveTypedValue(const std::shared_ptr& type) { @@ -109,8 +71,10 @@ bool IsSupportedPrimitiveTypedValue(const std::shared_ptr& type) { } } +template bool IsSupportedTypedValue(const std::shared_ptr& field); +template bool IsVariantFieldGroup(const std::shared_ptr& type) { if (type->id() != Type::STRUCT) { return false; @@ -126,7 +90,7 @@ bool IsVariantFieldGroup(const std::shared_ptr& type) { } value = field; } else if (field->name() == "typed_value") { - if (typed_value != nullptr || !IsSupportedTypedValue(field)) { + if (typed_value != nullptr || !IsSupportedTypedValue(field)) { return false; } typed_value = field; @@ -134,36 +98,43 @@ bool IsVariantFieldGroup(const std::shared_ptr& type) { return false; } } - return value != nullptr || typed_value != nullptr; + + if constexpr (strict && in_object) { + return value != nullptr; + } else { + return value != nullptr || typed_value != nullptr; + } } +template bool IsSupportedTypedValue(const std::shared_ptr& field) { if (!field->nullable()) { return false; } - auto is_variant_field_group = [](const auto& field) { - return !field->nullable() && IsVariantFieldGroup(field->type()); - }; switch (field->type()->id()) { case Type::STRUCT: return field->type()->num_fields() > 0 && - std::ranges::all_of(field->type()->fields(), is_variant_field_group); + std::ranges::all_of(field->type()->fields(), [](const auto& field) { + return (!strict || !field->nullable()) && + IsVariantFieldGroup(field->type()); + }); case Type::LIST: case Type::LARGE_LIST: case Type::LIST_VIEW: case Type::LARGE_LIST_VIEW: - case Type::FIXED_SIZE_LIST: - return is_variant_field_group(field->type()->field(0)); + case Type::FIXED_SIZE_LIST: { + auto& inner_field = field->type()->field(0); + return !inner_field->nullable() && + IsVariantFieldGroup(inner_field->type()); + } default: return IsSupportedPrimitiveTypedValue(field->type()); } } -} // namespace - -bool VariantExtensionType::IsSupportedStorageType( - const std::shared_ptr& storage_type) { +template +bool IsSupportedStorageTypeImpl(const std::shared_ptr& storage_type) { if (storage_type->id() != Type::STRUCT) { return false; } @@ -185,7 +156,7 @@ bool VariantExtensionType::IsSupportedStorageType( } value = field; } else if (field->name() == "typed_value") { - if (typed_value != nullptr || !IsSupportedTypedValue(field)) { + if (typed_value != nullptr || !IsSupportedTypedValue(field)) { return false; } typed_value = field; @@ -194,16 +165,68 @@ bool VariantExtensionType::IsSupportedStorageType( } } - if (metadata == nullptr || (value == nullptr && typed_value == nullptr)) { + if (metadata == nullptr) { return false; } - if (value == nullptr) { - return true; + + if constexpr (strict) { + if (value == nullptr) { + return false; + } + return (typed_value == nullptr) != value->nullable(); + } else { + bool value_nullable = (value == nullptr) || value->nullable(); + return (typed_value == nullptr) != value_nullable; } - if (typed_value == nullptr) { - return !value->nullable(); +} + +} // namespace + +VariantExtensionType::VariantExtensionType(const std::shared_ptr& storage_type) + : ExtensionType(storage_type) { + // IsSupportedStorageType should have been called already, asserting that + // metadata is present and at least one of value / typed_value is present. + for (const auto& field : storage_type->fields()) { + if (field->name() == "metadata") { + metadata_ = field; + } else if (field->name() == "value") { + value_ = field; + } else if (field->name() == "typed_value") { + typed_value_ = field; + } + } +} + +bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const { + return other.extension_name() == this->extension_name() && + other.storage_type()->Equals(this->storage_type()); +} + +Result> VariantExtensionType::Deserialize( + std::shared_ptr storage_type, const std::string& serialized) const { + if (!serialized.empty()) { + return Status::Invalid("Unexpected serialized metadata: '", serialized, "'"); + } + if (!IsSupportedStorageTypeImpl(storage_type)) { + return Status::Invalid("Invalid storage type for VariantExtensionType: ", + storage_type->ToString()); } - return value->nullable(); + return std::make_shared(std::move(storage_type)); +} + +std::string VariantExtensionType::Serialize() const { return ""; } + +std::shared_ptr VariantExtensionType::MakeArray( + std::shared_ptr data) const { + DCHECK_EQ(data->type->id(), Type::EXTENSION); + DCHECK_EQ(kVariantExtensionName, + internal::checked_cast(*data->type).extension_name()); + return std::make_shared(data); +} + +bool VariantExtensionType::IsSupportedStorageType( + const std::shared_ptr& storage_type) { + return IsSupportedStorageTypeImpl(storage_type); } Result> VariantExtensionType::Make( diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index 11da3eccecab..a2d6cfdd8f74 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -191,8 +191,11 @@ set(PARQUET_SRCS stream_reader.cc stream_writer.cc types.cc + variant/array_internal.cc variant/builder.cc - variant/encoding.cc + variant/decoding.cc + variant/slot_internal.cc + variant/unshred.cc variant/validate.cc xxhasher.cc) diff --git a/cpp/src/parquet/arrow/arrow_schema_test.cc b/cpp/src/parquet/arrow/arrow_schema_test.cc index 603ab3333601..b96790ecd059 100644 --- a/cpp/src/parquet/arrow/arrow_schema_test.cc +++ b/cpp/src/parquet/arrow/arrow_schema_test.cc @@ -1092,8 +1092,36 @@ TEST_F(TestConvertParquetSchema, ParquetVariantShreddedWithoutValue) { ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), ::arrow::field("typed_value", ::arrow::int64())}); - ASSERT_NO_FATAL_FAILURE(CheckParquetVariantSchema( - "variant_shredded", {metadata, typed_value}, storage_type)); + { + auto arrow_schema = + ::arrow::schema({::arrow::field("variant_shredded", storage_type)}); + ASSERT_OK(ConvertSchema( + {GroupNode::Make("variant_shredded", Repetition::OPTIONAL, + {metadata, typed_value}, LogicalType::Variant())})); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true)); + } + + auto registered_storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary(), false)}); + auto registered_variant = ::arrow::extension::variant(registered_storage_type); + ::arrow::ExtensionTypeGuard guard({registered_variant}); + + ArrowReaderProperties props; + props.set_arrow_extensions_enabled(true); + ASSERT_OK( + ConvertSchema({GroupNode::Make("variant_shredded", Repetition::OPTIONAL, + {metadata, typed_value}, LogicalType::Variant())}, + /*metadata=*/nullptr, props)); + + auto field = result_schema_->field(0); + ASSERT_EQ(::arrow::Type::EXTENSION, field->type()->id()); + auto variant_type = + std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>(field->type()); + ASSERT_NE(nullptr, variant_type); + ASSERT_EQ(nullptr, variant_type->value()); + ASSERT_NE(nullptr, variant_type->typed_value()); + ASSERT_TRUE(variant_type->storage_type()->Equals(storage_type)); } TEST_F(TestConvertParquetSchema, ParquetSchemaArrowJsonExtension) { @@ -1594,7 +1622,17 @@ TEST_F(TestConvertArrowSchema, ParquetVariantShreddedWithoutValue) { auto storage_type = ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), ::arrow::field("typed_value", ::arrow::int64())}); - auto variant_type = ::arrow::extension::variant(storage_type); + auto registered_storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary(), false)}); + ASSERT_OK_AND_ASSIGN( + auto registered_variant_type, + ::arrow::extension::VariantExtensionType::Make(registered_storage_type)); + auto registered_variant = + ::arrow::internal::checked_pointer_cast<::arrow::extension::VariantExtensionType>( + registered_variant_type); + ASSERT_OK_AND_ASSIGN(auto variant_type, + registered_variant->Deserialize(storage_type, "")); ASSERT_RAISES(Invalid, ConvertSchema({::arrow::field("variant", variant_type)})); } diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index cc107c1802e3..ab62ad3692fb 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -27,6 +27,7 @@ #include "arrow/array.h" #include "arrow/buffer.h" +#include "arrow/extension/parquet_variant.h" #include "arrow/extension_type.h" #include "arrow/io/memory.h" #include "arrow/memory_pool.h" @@ -54,6 +55,7 @@ #include "parquet/page_index.h" #include "parquet/properties.h" #include "parquet/schema.h" +#include "parquet/variant/validate.h" using arrow::Array; using arrow::ArrayData; @@ -583,6 +585,26 @@ class ExtensionReader : public ColumnReaderImpl { std::unique_ptr storage_reader_; }; +class VariantReader : public ExtensionReader { + public: + VariantReader(std::shared_ptr ctx, std::shared_ptr field, + std::unique_ptr storage_reader) + : ExtensionReader(std::move(field), std::move(storage_reader)), + ctx_(std::move(ctx)) {} + + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override { + RETURN_NOT_OK(ExtensionReader::BuildArray(length_upper_bound, out)); + if (ctx_->reader_properties->variant_validation_enabled()) { + PARQUET_CATCH_NOT_OK(variant::ValidateVariants(**out, ctx_->pool)); + } + return Status::OK(); + } + + private: + std::shared_ptr ctx_; +}; + template class ListReader : public ColumnReaderImpl { public: @@ -892,8 +914,8 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& arrow_f auto type_id = arrow_field->type()->id(); if (type_id == ::arrow::Type::EXTENSION) { - auto storage_field = arrow_field->WithType( - checked_cast(*arrow_field->type()).storage_type()); + const auto& ext_type = checked_cast(*arrow_field->type()); + auto storage_field = arrow_field->WithType(ext_type.storage_type()); RETURN_NOT_OK(GetReader(field, storage_field, ctx, out)); if (*out) { auto storage_type = (*out)->field()->type(); @@ -902,7 +924,11 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& arrow_f "Due to column pruning only part of an extension's storage type was loaded. " "An extension type cannot be created without all of its fields"); } - *out = std::make_unique(arrow_field, std::move(*out)); + if (ext_type.extension_name() == ::arrow::extension::kVariantExtensionName) { + *out = std::make_unique(ctx, arrow_field, std::move(*out)); + } else { + *out = std::make_unique(arrow_field, std::move(*out)); + } } return Status::OK(); } diff --git a/cpp/src/parquet/arrow/variant_test.cc b/cpp/src/parquet/arrow/variant_test.cc index 981997e14152..03dda1d12720 100644 --- a/cpp/src/parquet/arrow/variant_test.cc +++ b/cpp/src/parquet/arrow/variant_test.cc @@ -21,6 +21,7 @@ #include "arrow/io/memory.h" #include "arrow/record_batch.h" #include "arrow/table.h" +#include "arrow/testing/builder.h" #include "arrow/testing/extension_type.h" #include "arrow/testing/gtest_util.h" #include "parquet/arrow/reader.h" @@ -29,6 +30,7 @@ #include "parquet/variant/builder.h" #include "parquet/variant/test_util_internal.h" +#include #include #include #include @@ -39,13 +41,8 @@ using ::arrow::binary; using ::arrow::field; using ::arrow::struct_; using variant::internal::BinaryArrayFromValues; -using variant::internal::BinaryViewArrayFromValues; using variant::internal::EmptyVariantMetadata; -using variant::internal::Int32ArrayFromValues; -using variant::internal::Int64ArrayFromValues; using variant::internal::Int8Variant; -using variant::internal::StringArrayFromValues; -using variant::internal::UuidArrayFromValues; using variant::internal::VariantTable; using variant::internal::WriteVariantRecordBatch; using variant::internal::WriteVariantTable; @@ -96,7 +93,7 @@ TEST(TestVariantExtensionType, WriteRecordBatchValidatesBatch) { auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary(), /*nullable=*/false)}); auto variant_type = ::arrow::extension::variant(storage_type); - auto metadata_array = StringArrayFromValues({"not binary"}); + auto metadata_array = ::arrow::ArrayFromJSON(::arrow::utf8(), R"(["not binary"])"); auto value_array = BinaryArrayFromValues({std::string_view{*encoded.value}}); ASSERT_OK_AND_ASSIGN( auto storage, @@ -105,8 +102,7 @@ TEST(TestVariantExtensionType, WriteRecordBatchValidatesBatch) { auto schema = ::arrow::schema({field("variant", variant_type)}); auto batch = ::arrow::RecordBatch::Make(schema, 1, {variant_array}); - ASSERT_OK_AND_ASSIGN(auto sink, ::arrow::io::BufferOutputStream::Create( - 1024, ::arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN(auto sink, ::arrow::io::BufferOutputStream::Create(1024)); ASSERT_OK_AND_ASSIGN( auto writer, FileWriter::Open(*schema, ::arrow::default_memory_pool(), sink, @@ -126,8 +122,12 @@ TEST(TestVariantExtensionType, WriterValidatesBinaryViewVariantBytes) { struct_({field("metadata", ::arrow::binary_view(), /*nullable=*/false), field("value", ::arrow::binary_view(), /*nullable=*/false)}); auto variant_type = ::arrow::extension::variant(storage_type); - auto metadata_array = BinaryViewArrayFromValues({std::string_view{*encoded.metadata}}); - auto value_array = BinaryViewArrayFromValues({std::string_view{*encoded.value}}); + std::shared_ptr<::arrow::Array> metadata_array; + ::arrow::ArrayFromVector<::arrow::BinaryViewType, std::string_view>( + {std::string_view{*encoded.metadata}}, &metadata_array); + std::shared_ptr<::arrow::Array> value_array; + ::arrow::ArrayFromVector<::arrow::BinaryViewType, std::string_view>( + {std::string_view{*encoded.value}}, &value_array); auto table = VariantTable(variant_type, {metadata_array, value_array}, storage_type->fields()); ASSERT_OK(WriteVariantTable(table)); @@ -155,12 +155,12 @@ TEST(TestVariantExtensionType, WriterSkipsNullParents) { ASSERT_OK(WriteVariantTable(parent_table)); auto map_type = ::arrow::map(::arrow::utf8(), field("item", variant_type)); - ASSERT_OK_AND_ASSIGN( - auto map_array, - ::arrow::MapArray::FromArrays(map_type, Int32ArrayFromValues({0, 1}), - StringArrayFromValues({"hidden"}), variant_array, - ::arrow::default_memory_pool(), - ::arrow::Buffer::FromString(std::string("\0", 1)))); + ASSERT_OK_AND_ASSIGN(auto map_array, + ::arrow::MapArray::FromArrays( + map_type, ::arrow::ArrayFromJSON(::arrow::int32(), "[0, 1]"), + ::arrow::ArrayFromJSON(::arrow::utf8(), R"(["hidden"])"), + variant_array, ::arrow::default_memory_pool(), + ::arrow::Buffer::FromString(std::string("\0", 1)))); auto map_table = ::arrow::Table::Make(::arrow::schema({field("variant_map", map_type)}), {map_array}); ASSERT_OK(WriteVariantTable(map_table)); @@ -178,7 +178,7 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { {std::string_view{*encoded.metadata}, std::string_view{*encoded.metadata}}); auto value_array = BinaryArrayFromValues({std::nullopt, std::string_view{*encoded.value}}); - auto typed_array = Int64ArrayFromValues({34, 100}); + auto typed_array = ::arrow::ArrayFromJSON(::arrow::int64(), "[34, 100]"); auto table = VariantTable(variant_type, {metadata_array, value_array, typed_array}, storage_type->fields()); ASSERT_RAISES(Invalid, WriteVariantTable(table)); @@ -190,24 +190,16 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { auto valid_values = BinaryArrayFromValues({std::nullopt, std::string_view{*encoded.value}}); - auto valid_typed = Int64ArrayFromValues({34, std::nullopt}); + auto valid_typed = ::arrow::ArrayFromJSON(::arrow::int64(), "[34, null]"); auto valid_table = VariantTable( variant_type, {metadata_array, valid_values, valid_typed}, storage_type->fields()); ASSERT_OK(WriteVariantTable(valid_table)); } TEST(TestVariantExtensionType, WriterRejectsShreddedWithoutValue) { - auto metadata = EmptyVariantMetadata(); auto storage_type = struct_({field("metadata", binary(), /*nullable=*/false), field("typed_value", ::arrow::int64())}); - ASSERT_OK_AND_ASSIGN(auto variant_type, - ::arrow::extension::VariantExtensionType::Make(storage_type)); - - auto metadata_array = BinaryArrayFromValues({std::string_view{*metadata}}); - auto typed_array = Int64ArrayFromValues({34}); - auto table = - VariantTable(variant_type, {metadata_array, typed_array}, storage_type->fields()); - ASSERT_RAISES(Invalid, WriteVariantTable(table)); + ASSERT_RAISES(Invalid, ::arrow::extension::VariantExtensionType::Make(storage_type)); } TEST(TestVariantExtensionType, ReadsDictionaryEncodedMetadata) { @@ -304,7 +296,9 @@ TEST(TestVariantExtensionType, WriterWritesUuid) { auto metadata_array = BinaryArrayFromValues({std::string_view{*metadata}}); auto value_array = BinaryArrayFromValues({std::nullopt}); - auto typed_array = UuidArrayFromValues({std::string_view("0123456789abcdef", 16)}); + auto typed_array = ::arrow::ExtensionType::WrapArray( + ::arrow::extension::uuid(), + ::arrow::ArrayFromJSON(::arrow::fixed_size_binary(16), R"(["0123456789abcdef"])")); auto table = VariantTable(variant_type, {metadata_array, value_array, typed_array}, storage_type->fields()); ASSERT_OK(WriteVariantTable(table)); @@ -329,9 +323,10 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedObjectConflicts) { auto metadata_array = BinaryArrayFromValues({std::string_view{*encoded.metadata}}); auto value_array = BinaryArrayFromValues({std::string_view{*encoded.value}}); ASSERT_OK_AND_ASSIGN(auto event_type_group, - ::arrow::StructArray::Make({BinaryArrayFromValues({std::nullopt}), - StringArrayFromValues({"login"})}, - field_group_type->fields())); + ::arrow::StructArray::Make( + {BinaryArrayFromValues({std::nullopt}), + ::arrow::ArrayFromJSON(::arrow::utf8(), R"(["login"])")}, + field_group_type->fields())); ASSERT_OK_AND_ASSIGN( auto typed_array, ::arrow::StructArray::Make({event_type_group}, typed_value_type->fields())); @@ -346,10 +341,11 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedObjectConflicts) { storage_type->fields()); ASSERT_OK(WriteVariantTable(valid_table)); - ASSERT_OK_AND_ASSIGN(auto missing_event_type_group, - ::arrow::StructArray::Make({BinaryArrayFromValues({std::nullopt}), - StringArrayFromValues({std::nullopt})}, - field_group_type->fields())); + ASSERT_OK_AND_ASSIGN( + auto missing_event_type_group, + ::arrow::StructArray::Make({BinaryArrayFromValues({std::nullopt}), + ::arrow::ArrayFromJSON(::arrow::utf8(), "[null]")}, + field_group_type->fields())); ASSERT_OK_AND_ASSIGN( auto missing_typed_array, ::arrow::StructArray::Make({missing_event_type_group}, typed_value_type->fields())); diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index 47f50062f109..ae04ceea6702 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -382,7 +382,8 @@ class FileWriterImpl : public FileWriter { int64_t size) override { RETURN_NOT_OK(CheckClosed()); if (arrow_properties_->variant_validation_enabled()) { - PARQUET_CATCH_NOT_OK(variant::ValidateVariants(*data->Slice(offset, size), pool_)); + PARQUET_CATCH_NOT_OK( + variant::ValidateVariants(*data->Slice(offset, size), pool_)); } if (arrow_properties_->engine_version() == ArrowWriterProperties::V2 || arrow_properties_->engine_version() == ArrowWriterProperties::V1) { @@ -476,7 +477,7 @@ class FileWriterImpl : public FileWriter { ChunkedArray chunked_array{batch.column(i)}; if (arrow_properties_->variant_validation_enabled()) { PARQUET_CATCH_NOT_OK( - variant::ValidateVariants(*chunked_array.Slice(offset, size), pool_)); + variant::ValidateVariants(*chunked_array.Slice(offset, size), pool_)); } ARROW_ASSIGN_OR_RAISE( std::unique_ptr writer, diff --git a/cpp/src/parquet/meson.build b/cpp/src/parquet/meson.build index d7737f7db0f9..6cbe589c5c67 100644 --- a/cpp/src/parquet/meson.build +++ b/cpp/src/parquet/meson.build @@ -55,8 +55,11 @@ parquet_srcs = files( 'stream_reader.cc', 'stream_writer.cc', 'types.cc', + 'variant/array_internal.cc', 'variant/builder.cc', - 'variant/encoding.cc', + 'variant/decoding.cc', + 'variant/slot_internal.cc', + 'variant/unshred.cc', 'variant/validate.cc', 'xxhasher.cc', ) diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index e00ec2fe7dd6..9ff4edea27d2 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1156,6 +1156,7 @@ class PARQUET_EXPORT ArrowReaderProperties { binary_type_(kArrowDefaultBinaryType), list_type_(kArrowDefaultListType), arrow_extensions_enabled_(false), + variant_validation_enabled_(true), should_load_statistics_(false), smallest_decimal_enabled_(false) {} @@ -1267,6 +1268,16 @@ class PARQUET_EXPORT ArrowReaderProperties { } bool get_arrow_extensions_enabled() const { return arrow_extensions_enabled_; } + /// Enable read-side row-level validation for Parquet Variant extension arrays. + /// + /// When enabled, Variant extension arrays read from Parquet are checked for + /// shredded value / typed_value semantic conflicts after the physical Arrow + /// storage has been materialized. + void set_variant_validation_enabled(bool enabled) { + variant_validation_enabled_ = enabled; + } + bool variant_validation_enabled() const { return variant_validation_enabled_; } + /// \brief Set whether to load statistics as much as possible. /// /// Default is false. @@ -1305,6 +1316,7 @@ class PARQUET_EXPORT ArrowReaderProperties { ::arrow::Type::type binary_type_; ::arrow::Type::type list_type_; bool arrow_extensions_enabled_; + bool variant_validation_enabled_; bool should_load_statistics_; bool smallest_decimal_enabled_; }; diff --git a/cpp/src/parquet/variant/CMakeLists.txt b/cpp/src/parquet/variant/CMakeLists.txt index 8c9b99266331..036b5b744b27 100644 --- a/cpp/src/parquet/variant/CMakeLists.txt +++ b/cpp/src/parquet/variant/CMakeLists.txt @@ -18,9 +18,11 @@ add_parquet_test(variant-test SOURCES builder_test.cc - encoding_test.cc + decoding_test.cc + parquet_testing_test.cc test_util_internal.cc type_test.cc + unshred_test.cc validate_test.cc) arrow_install_all_headers("parquet/variant") diff --git a/cpp/src/parquet/variant/append_table_internal.h b/cpp/src/parquet/variant/append_table_internal.h new file mode 100644 index 000000000000..862a77c6421f --- /dev/null +++ b/cpp/src/parquet/variant/append_table_internal.h @@ -0,0 +1,41 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#define PARQUET_VARIANT_DIRECT_APPEND_LIST(V) \ + V(Int8, (int8_t value), value) \ + V(Int16, (int16_t value), value) \ + V(Int32, (int32_t value), value) \ + V(Int64, (int64_t value), value) \ + V(Float, (float value), value) \ + V(Double, (double value), value) \ + V(Binary, (std::string_view value), value) \ + V(String, (std::string_view value), value) \ + V(Date, (int32_t value), value) \ + V(TimeNTZMicros, (int64_t value), value) \ + V(Uuid, (std::string_view value), value) \ + V(Decimal4, (const ::arrow::Decimal32& value, uint8_t scale), value, scale) \ + V(Decimal8, (const ::arrow::Decimal64& value, uint8_t scale), value, scale) \ + V(Decimal16, (const ::arrow::Decimal128& value, uint8_t scale), value, scale) + +#define PARQUET_VARIANT_SPECIAL_APPEND_LIST(V) \ + V(VariantNull, ()) \ + V(Boolean, (bool value), value) \ + V(ShortString, (std::string_view value), value) \ + V(TimestampMicros, (int64_t value, bool adjusted_to_utc), value, adjusted_to_utc) \ + V(TimestampNanos, (int64_t value, bool adjusted_to_utc), value, adjusted_to_utc) diff --git a/cpp/src/parquet/variant/array_internal.cc b/cpp/src/parquet/variant/array_internal.cc new file mode 100644 index 000000000000..c71fc31f4024 --- /dev/null +++ b/cpp/src/parquet/variant/array_internal.cc @@ -0,0 +1,90 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/array_internal.h" + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "parquet/exception.h" + +namespace parquet::variant::internal { + +using ::arrow::Array; +using ::arrow::internal::checked_cast; + +std::shared_ptr ValuesArray(const Array& array) { + switch (array.type_id()) { + case ::arrow::Type::LIST_VIEW: + return checked_cast(array).values(); + case ::arrow::Type::LARGE_LIST_VIEW: + return checked_cast(array).values(); + case ::arrow::Type::LIST: + return checked_cast(array).values(); + case ::arrow::Type::LARGE_LIST: + return checked_cast(array).values(); + case ::arrow::Type::FIXED_SIZE_LIST: + return checked_cast(array).values(); + case ::arrow::Type::MAP: + return checked_cast(array).values(); + default: + throw ParquetException("Expected list or map storage, got ", + array.type()->ToString()); + } +} + +template +std::pair ValueOffsetAndLength(const Array& array, int64_t row) { + const auto& typed_array = checked_cast(array); + return {typed_array.value_offset(row), typed_array.value_length(row)}; +} + +std::pair ValuesRangeAt(const Array& array, int64_t row) { + switch (array.type_id()) { + case ::arrow::Type::LIST: + return ValueOffsetAndLength<::arrow::ListArray>(array, row); + case ::arrow::Type::LARGE_LIST: + return ValueOffsetAndLength<::arrow::LargeListArray>(array, row); + case ::arrow::Type::LIST_VIEW: + return ValueOffsetAndLength<::arrow::ListViewArray>(array, row); + case ::arrow::Type::LARGE_LIST_VIEW: + return ValueOffsetAndLength<::arrow::LargeListViewArray>(array, row); + case ::arrow::Type::MAP: + return ValueOffsetAndLength<::arrow::MapArray>(array, row); + case ::arrow::Type::FIXED_SIZE_LIST: + return ValueOffsetAndLength<::arrow::FixedSizeListArray>(array, row); + default: + throw ParquetException("Expected list or map storage, got ", + array.type()->ToString()); + } +} + +std::string_view BinaryFieldView(const Array& array, int64_t row) { + switch (array.type_id()) { + case ::arrow::Type::BINARY: + return checked_cast(array).GetView(row); + case ::arrow::Type::LARGE_BINARY: + return checked_cast(array).GetView(row); + case ::arrow::Type::BINARY_VIEW: + return checked_cast(array).GetView(row); + default: + throw ParquetInvalidOrCorruptedFileException("Expected binary Variant field, got ", + array.type()->ToString()); + } +} + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/array_internal.h b/cpp/src/parquet/variant/array_internal.h new file mode 100644 index 000000000000..26f80471f35a --- /dev/null +++ b/cpp/src/parquet/variant/array_internal.h @@ -0,0 +1,53 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: export +#include "arrow/buffer.h" +#include "arrow/util/bit_block_counter.h" + +namespace parquet::variant::internal { + +std::shared_ptr<::arrow::Array> ValuesArray(const ::arrow::Array& array); + +std::pair ValuesRangeAt(const ::arrow::Array& array, int64_t row); + +std::string_view BinaryFieldView(const ::arrow::Array& array, int64_t row); + +template +void VisitVisibleRows(const std::shared_ptr<::arrow::Buffer>& valid_rows, + const ::arrow::Array& array, VisitVisible&& visit_visible) { + if (valid_rows == nullptr && !array.data()->MayHaveNulls()) { + for (int64_t row = 0; row < array.length(); ++row) { + visit_visible(row); + } + return; + } + + ::arrow::internal::VisitTwoBitBlocksVoid( + valid_rows == nullptr ? nullptr : valid_rows->data(), /*left_offset=*/0, + array.null_bitmap_data(), array.offset(), array.length(), + std::forward(visit_visible), [] {}); +} + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/binary_view_column_builder_internal.h b/cpp/src/parquet/variant/binary_view_column_builder_internal.h new file mode 100644 index 000000000000..deb5ea9f914e --- /dev/null +++ b/cpp/src/parquet/variant/binary_view_column_builder_internal.h @@ -0,0 +1,129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_binary.h" +#include "arrow/buffer.h" +#include "arrow/buffer_builder.h" +#include "arrow/type.h" +#include "arrow/util/binary_view_util.h" +#include "arrow/util/logging_internal.h" +#include "parquet/exception.h" + +namespace parquet::variant::internal { + +// Builds binary-view columns from Variant bytes encoded directly into a shared arena. +// Arrow's BinaryViewBuilder accepts complete values and copies them into its private +// StringHeapBuilder, so it cannot expose the writable arena used by nested Variant +// encoders or roll back a partially committed row. Keeping bytes and views separate lets +// callers write each value once, create its view afterward, discard inlined bytes, and +// roll back both buffers together. +class BinaryViewColumnBuilder { + public: + struct Mark { + int64_t bytes_length = 0; + int64_t views_length = 0; + bool has_out_of_line_data = false; + }; + + explicit BinaryViewColumnBuilder(::arrow::MemoryPool* pool) + : bytes_(pool), views_(pool) {} + + [[nodiscard]] Mark mark() const { + return Mark{.bytes_length = bytes_.length(), + .views_length = views_.length(), + .has_out_of_line_data = has_out_of_line_data_}; + } + + void Rollback(const Mark& mark) { + bytes_.Rewind(mark.bytes_length); + views_.bytes_builder()->Rewind(mark.views_length * + sizeof(::arrow::BinaryViewType::c_type)); + has_out_of_line_data_ = mark.has_out_of_line_data; + } + + int64_t length() const { return views_.length(); } + + std::string_view SliceFrom(int64_t start) const { + DCHECK_LE(start, bytes_.length()); + const auto size = bytes_.length() - start; + DCHECK_GE(size, 0); + if (size == 0) { + return std::string_view{}; + } + return std::string_view{reinterpret_cast(bytes_.data() + start), + static_cast(size)}; + } + + void AppendView(int64_t start) { + const auto slice = SliceFrom(start); + if (slice.size() > static_cast(std::numeric_limits::max())) { + throw ParquetException("Binary view value is too large"); + } + PARQUET_THROW_NOT_OK(views_.Reserve(1)); + if (slice.size() <= ::arrow::BinaryViewType::kInlineSize) { + views_.UnsafeAppend(::arrow::util::ToInlineBinaryView(slice)); + bytes_.Rewind(start); + return; + } + if (start > std::numeric_limits::max()) { + throw ParquetException("Binary view offset is too large"); + } + has_out_of_line_data_ = true; + views_.UnsafeAppend(::arrow::util::ToNonInlineBinaryView( + slice.data(), static_cast(slice.size()), + /*buffer_index=*/0, static_cast(start))); + } + + void AppendEmptyValue() { + PARQUET_THROW_NOT_OK(views_.Reserve(1)); + views_.UnsafeAppend(::arrow::BinaryViewType::c_type{}); + } + + ::arrow::BufferBuilder& bytes_builder() { return bytes_; } + + std::shared_ptr<::arrow::BinaryViewArray> Finish( + std::shared_ptr<::arrow::Buffer> null_bitmap = nullptr, int64_t null_count = 0) { + const auto array_length = views_.length(); + std::shared_ptr<::arrow::Buffer> views_buffer; + PARQUET_THROW_NOT_OK(views_.Finish(&views_buffer)); + + PARQUET_ASSIGN_OR_THROW(auto data_buffer, bytes_.Finish()); + ::arrow::BufferVector data_buffers; + if (has_out_of_line_data_) { + data_buffers.push_back(std::move(data_buffer)); + } + return std::make_shared<::arrow::BinaryViewArray>( + ::arrow::binary_view(), array_length, std::move(views_buffer), + std::move(data_buffers), std::move(null_bitmap), null_count); + } + + private: + ::arrow::BufferBuilder bytes_; + ::arrow::TypedBufferBuilder<::arrow::BinaryViewType::c_type> views_; + bool has_out_of_line_data_ = false; +}; + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/builder.cc b/cpp/src/parquet/variant/builder.cc index 8fdeffd6855c..29c516e59f11 100644 --- a/cpp/src/parquet/variant/builder.cc +++ b/cpp/src/parquet/variant/builder.cc @@ -32,21 +32,20 @@ #include "arrow/buffer_builder.h" #include "arrow/builder.h" // IWYU pragma: keep #include "arrow/type.h" -#include "arrow/util/binary_view_util.h" #include "arrow/util/checked_cast.h" #include "arrow/util/endian.h" #include "arrow/util/logging_internal.h" #include "parquet/exception.h" -#include "parquet/variant/encoding.h" -#include "parquet/variant/encoding_internal.h" +#include "parquet/variant/append_table_internal.h" +#include "parquet/variant/binary_view_column_builder_internal.h" +#include "parquet/variant/decoding.h" +#include "parquet/variant/format_internal.h" namespace parquet::variant { using ::arrow::binary_view; using ::arrow::BinaryViewArray; -using ::arrow::BinaryViewType; using ::arrow::BufferBuilder; -using ::arrow::BufferVector; using ::arrow::ExtensionType; using ::arrow::field; using ::arrow::struct_; @@ -54,6 +53,7 @@ using ::arrow::StructType; using ::arrow::Type; using ::arrow::TypedBufferBuilder; using ::arrow::extension::VariantExtensionType; +using internal::BinaryViewColumnBuilder; namespace bit_util = ::arrow::bit_util; @@ -72,20 +72,18 @@ void AppendFixedLittleEndian(BufferBuilder& out, T value) { PARQUET_THROW_NOT_OK(out.Append(&little_endian, sizeof(T))); } -void AppendLittleEndianToString(std::string& out, uint32_t value, uint8_t width) { - DCHECK_LE(width, sizeof(uint32_t)); - const auto little_endian = bit_util::ToLittleEndian(value); - out.append(reinterpret_cast(&little_endian), width); -} - -void InsertBytes(BufferBuilder& out, int64_t offset, std::string_view bytes) { - const int64_t old_size = out.length(); - const auto insert_size = bytes.size(); - PARQUET_THROW_NOT_OK(out.Reserve(insert_size)); - uint8_t* data = out.mutable_data(); - std::memmove(data + offset + insert_size, data + offset, old_size - offset); - std::memcpy(data + offset, bytes.data(), insert_size); - out.UnsafeAdvance(insert_size); +template +std::array ToLittleEndianBytes( + const DecimalValue& value) { + auto words = value.little_endian_array(); + std::array out{}; + size_t offset = 0; + for (const auto& word : words) { + const auto little_endian_word = bit_util::ToLittleEndian(word); + std::memcpy(out.data() + offset, &little_endian_word, sizeof(little_endian_word)); + offset += sizeof(little_endian_word); + } + return out; } uint8_t WidthForValue(uint64_t value) { @@ -101,18 +99,30 @@ uint8_t WidthForValue(uint64_t value) { return 4; } -class VariantMetadataBuilder { +class MetadataBuilder { public: - explicit VariantMetadataBuilder(MemoryPool* pool) : pool_(pool) {} + virtual ~MetadataBuilder() = default; + + virtual void Reserve(int64_t capacity) = 0; + // Returns an existing or newly assigned field ID, or throws if the field + // cannot be added or resolved. + virtual uint32_t TryUpsert(std::string_view name) = 0; + virtual std::string_view FieldName(uint32_t id) const = 0; + virtual size_t size() const = 0; + virtual void Truncate(size_t size) = 0; + virtual void AppendTo(BufferBuilder& out) const = 0; +}; - void Reserve(int64_t capacity) { +class WritableMetadataBuilder : public MetadataBuilder { + public: + void Reserve(int64_t capacity) override { if (capacity < 0) { throw ParquetException("Variant metadata capacity must be non-negative"); } field_ids_.reserve(capacity); } - uint32_t Upsert(std::string_view name) { + uint32_t TryUpsert(std::string_view name) override { auto it = field_ids_.find(name); if (it != field_ids_.end()) { return it->second; @@ -133,20 +143,28 @@ class VariantMetadataBuilder { return static_cast(id); } - std::string_view FieldName(uint32_t id) const { + std::string_view FieldName(uint32_t id) const override { DCHECK_LT(id, field_names_.size()); return field_names_[id]; } - size_t size() const { return field_names_.size(); } + size_t size() const override { return field_names_.size(); } - void Truncate(size_t size) { - DCHECK_LE(size, field_names_.size()); - field_names_.resize(size); - RebuildIndex(); + void Truncate(size_t size) override { + if (size == field_names_.size()) { + return; + } + DCHECK_LT(size, field_names_.size()); + while (field_names_.size() > size) { + const auto field_name = std::string_view{field_names_.back()}; + const auto erased = field_ids_.erase(field_name); + DCHECK_EQ(erased, 1); + field_names_.pop_back(); + } + is_sorted_ = std::ranges::is_sorted(field_names_); } - void AppendTo(BufferBuilder& out) const { + void AppendTo(BufferBuilder& out) const override { uint64_t bytes_size = 0; for (const auto& string : field_names_) { bytes_size += string.size(); @@ -180,33 +198,47 @@ class VariantMetadataBuilder { } } - std::shared_ptr Finish() const { - BufferBuilder out(pool_); - AppendTo(out); - std::shared_ptr buffer; - PARQUET_ASSIGN_OR_THROW(buffer, out.Finish()); - return buffer; - } - private: - void RebuildIndex() { - field_ids_.clear(); - field_ids_.reserve(field_names_.size()); - is_sorted_ = false; - for (uint32_t i = 0; i < field_names_.size(); ++i) { - field_ids_.emplace(field_names_[i], i); - if (i == 0) { - is_sorted_ = true; - } else if (is_sorted_ && !(field_names_[i - 1] < field_names_[i])) { - is_sorted_ = false; - } + std::deque field_names_; + std::unordered_map field_ids_; + bool is_sorted_ = false; +}; + +class ReadOnlyMetadataBuilder : public MetadataBuilder { + public: + explicit ReadOnlyMetadataBuilder(const VariantMetadataView& metadata) + : metadata_(metadata) {} + + void Reserve(int64_t) override {} + + uint32_t TryUpsert(std::string_view name) override { + auto it = field_ids_.find(name); + if (it != field_ids_.end()) { + return it->second; + } + auto field_id = metadata_.FindString(name); + if (!field_id.has_value()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: field name '", name, + "' is not in metadata dictionary"); } + field_ids_.emplace(metadata_.string(*field_id), *field_id); + return *field_id; } - MemoryPool* pool_; - std::deque field_names_; + std::string_view FieldName(uint32_t id) const override { return metadata_.string(id); } + + size_t size() const override { return metadata_.dictionary_size(); } + + void Truncate(size_t size) override { DCHECK_EQ(size, metadata_.dictionary_size()); } + + void AppendTo(BufferBuilder& out) const override { + PARQUET_THROW_NOT_OK(out.Append(metadata_.metadata())); + } + + private: + const VariantMetadataView& metadata_; std::unordered_map field_ids_; - bool is_sorted_ = false; }; class VariantValueWriter { @@ -221,10 +253,9 @@ class VariantValueWriter { template requires internal::FixedVariantPrimitive - void Append(typename internal::VariantFixedPrimitiveTraits::CType value) { - using CType = typename internal::VariantFixedPrimitiveTraits::CType; + void Append(internal::VariantFixedPrimitiveTraits::CType value) { AppendPrimitiveHeader(); - if constexpr (sizeof(CType) == 1) { + if constexpr (sizeof(value) == 1) { const auto byte = static_cast(value); PARQUET_THROW_NOT_OK(out_.Append(&byte, sizeof(byte))); } else { @@ -234,25 +265,12 @@ class VariantValueWriter { template requires internal::DecimalVariantPrimitive - void Append( - typename internal::VariantDecimalPrimitiveTraits::CType unscaled_value, - uint8_t scale) { - internal::ValidateDecimalScale(scale); - AppendPrimitiveHeader(); - PARQUET_THROW_NOT_OK(out_.Append(&scale, sizeof(scale))); - AppendFixedLittleEndian(out_, unscaled_value); - } - - template - requires internal::Decimal16VariantPrimitive - void Append(std::string_view little_endian_unscaled_value, uint8_t scale) { - if (little_endian_unscaled_value.size() != 16) { - throw ParquetException("Variant Decimal16 values must be 16 bytes"); - } + void Append(internal::VariantDecimalPrimitiveTraits::CType value, uint8_t scale) { internal::ValidateDecimalScale(scale); AppendPrimitiveHeader(); PARQUET_THROW_NOT_OK(out_.Append(&scale, sizeof(scale))); - PARQUET_THROW_NOT_OK(out_.Append(little_endian_unscaled_value)); + const auto bytes = ToLittleEndianBytes(value); + PARQUET_THROW_NOT_OK(out_.Append(bytes.data(), static_cast(bytes.size()))); } template @@ -281,6 +299,16 @@ class VariantValueWriter { PARQUET_THROW_NOT_OK(out_.Append(big_endian_bytes)); } + void AppendVariantNull() { Append(); } + + void AppendBoolean(bool value) { + if (value) { + Append(); + } else { + Append(); + } + } + void AppendShortString(std::string_view value) { if (value.size() >= 64) { throw ParquetException("Variant short string value must be shorter than 64 bytes"); @@ -292,6 +320,29 @@ class VariantValueWriter { PARQUET_THROW_NOT_OK(out_.Append(value)); } + void AppendTimestampMicros(int64_t value, bool adjusted_to_utc) { + if (adjusted_to_utc) { + Append(value); + } else { + Append(value); + } + } + + void AppendTimestampNanos(int64_t value, bool adjusted_to_utc) { + if (adjusted_to_utc) { + Append(value); + } else { + Append(value); + } + } + + void AppendEncodedValue(std::string_view value) { + if (value.empty()) { + throw ParquetException("Encoded Variant value length cannot be 0"); + } + PARQUET_THROW_NOT_OK(out_.Append(value)); + } + private: template void AppendPrimitiveHeader() { @@ -325,12 +376,19 @@ struct VariantBuildFrame { }; struct VariantBuildState { - VariantBuildState(MemoryPool* pool, BufferBuilder& value) - : value(value), root_value_start(value.length()), metadata(pool) {} + VariantBuildState(BufferBuilder& value, std::unique_ptr metadata) + : value(value), root_value_start(value.length()), metadata(std::move(metadata)) {} + + void Reset(std::unique_ptr new_metadata) { + root_value_start = value.length(); + metadata = std::move(new_metadata); + frames.clear(); + root_has_value = false; + } BufferBuilder& value; - int64_t root_value_start = 0; - VariantMetadataBuilder metadata; + int64_t root_value_start; + std::unique_ptr metadata; std::vector frames; bool root_has_value = false; }; @@ -344,8 +402,8 @@ void CheckRootWritable(const VariantBuildState& state) { } } -void CheckTopFrame(const VariantBuildState& state, size_t frame_index, - VariantContainerKind kind) { +template +void CheckTopFrame(const VariantBuildState& state, size_t frame_index) { if (state.frames.empty() || frame_index + 1 != state.frames.size()) { throw ParquetException("Variant nested builder is not the active container"); } @@ -367,26 +425,17 @@ void TruncateFrameEntries(VariantBuildFrame& frame, size_t entry_count) { } } -template -void RollbackIfActive(Impl* impl) { - if (impl != nullptr && impl->active) { - if (impl->state == nullptr || impl->frame_index >= impl->state->frames.size()) { - return; - } - const auto frame = impl->state->frames[impl->frame_index]; - impl->state->value.Rewind(frame.value_start); - impl->state->metadata.Truncate(frame.metadata_size); - if (frame.has_parent && frame.parent_frame < impl->state->frames.size()) { - TruncateFrameEntries(impl->state->frames[frame.parent_frame], - frame.parent_entry_count); - } - impl->state->frames.resize(impl->frame_index); - impl->active = false; - } +void MakeRoomForHeader(BufferBuilder& out, int64_t offset, int64_t header_size) { + const int64_t old_size = out.length(); + DCHECK_LE(offset, old_size); + PARQUET_THROW_NOT_OK(out.Reserve(header_size)); + uint8_t* data = out.mutable_data(); + std::memmove(data + offset + header_size, data + offset, old_size - offset); + out.Rewind(offset); } -std::string BuildObjectHeader(const VariantBuildState& state, - const VariantBuildFrame& frame, uint32_t values_size) { +void BuildObjectHeader(VariantBuildState& state, const VariantBuildFrame& frame, + uint32_t values_size) { if (frame.fields.size() > std::numeric_limits::max()) { throw ParquetException("Variant object has too many fields"); } @@ -394,8 +443,8 @@ std::string BuildObjectHeader(const VariantBuildState& state, std::vector fields = frame.fields; std::ranges::sort(fields, [&](const VariantFieldDescriptor& left, const VariantFieldDescriptor& right) { - return state.metadata.FieldName(left.field_id) < - state.metadata.FieldName(right.field_id); + return state.metadata->FieldName(left.field_id) < + state.metadata->FieldName(right.field_id); }); uint32_t max_field_id = 0; @@ -405,279 +454,196 @@ std::string BuildObjectHeader(const VariantBuildState& state, const uint8_t id_size = WidthForValue(max_field_id); const uint8_t offset_size = WidthForValue(values_size); const bool is_large = fields.size() > std::numeric_limits::max(); - const auto header = static_cast(((is_large ? 1 : 0) << 4) | - ((id_size - 1) << 2) | (offset_size - 1)); - - std::string out; - out.push_back( - static_cast((header << 2) | static_cast(VariantBasicType::kObject))); - AppendLittleEndianToString(out, static_cast(fields.size()), is_large ? 4 : 1); + const auto field_count = static_cast(fields.size()); + const uint8_t count_size = is_large ? 4 : 1; + const int64_t header_size = + 1 + count_size + field_count * id_size + (field_count + 1) * offset_size; + + MakeRoomForHeader(state.value, frame.value_start, header_size); + const auto header = static_cast( + ((((is_large ? 1 : 0) << 4) | ((id_size - 1) << 2) | (offset_size - 1)) << 2) | + static_cast(VariantBasicType::kObject)); + PARQUET_THROW_NOT_OK(state.value.Append(&header, sizeof(header))); + AppendLittleEndian(state.value, static_cast(fields.size()), count_size); for (const auto& field : fields) { - AppendLittleEndianToString(out, field.field_id, id_size); + AppendLittleEndian(state.value, field.field_id, id_size); } for (const auto& field : fields) { - AppendLittleEndianToString(out, field.offset, offset_size); + AppendLittleEndian(state.value, field.offset, offset_size); } - AppendLittleEndianToString(out, values_size, offset_size); - return out; + AppendLittleEndian(state.value, values_size, offset_size); + DCHECK_EQ(state.value.length(), frame.value_start + header_size); + state.value.UnsafeAdvance(values_size); } -std::string BuildListHeader(const VariantBuildFrame& frame, uint32_t values_size) { +void BuildListHeader(VariantBuildState& state, const VariantBuildFrame& frame, + uint32_t values_size) { if (frame.offsets.size() > std::numeric_limits::max()) { throw ParquetException("Variant array has too many elements"); } const uint8_t offset_size = WidthForValue(values_size); const bool is_large = frame.offsets.size() > std::numeric_limits::max(); - const auto header = static_cast(((is_large ? 1 : 0) << 2) | (offset_size - 1)); - - std::string out; - out.push_back( - static_cast((header << 2) | static_cast(VariantBasicType::kArray))); - AppendLittleEndianToString(out, static_cast(frame.offsets.size()), - is_large ? 4 : 1); + const auto element_count = static_cast(frame.offsets.size()); + const uint8_t count_size = is_large ? 4 : 1; + const int64_t header_size = 1 + count_size + (element_count + 1) * offset_size; + + MakeRoomForHeader(state.value, frame.value_start, header_size); + const auto header = + static_cast(((((is_large ? 1 : 0) << 2) | (offset_size - 1)) << 2) | + static_cast(VariantBasicType::kArray)); + PARQUET_THROW_NOT_OK(state.value.Append(&header, sizeof(header))); + AppendLittleEndian(state.value, static_cast(frame.offsets.size()), + count_size); for (const auto offset : frame.offsets) { - AppendLittleEndianToString(out, offset, offset_size); + AppendLittleEndian(state.value, offset, offset_size); } - AppendLittleEndianToString(out, values_size, offset_size); - return out; + AppendLittleEndian(state.value, values_size, offset_size); + DCHECK_EQ(state.value.length(), frame.value_start + header_size); + state.value.UnsafeAdvance(values_size); } -using RootFinishCallback = std::function&)>; +using RootFinishCallback = std::function; -void FinishFrame(const std::shared_ptr& state, size_t frame_index, - VariantContainerKind kind, const RootFinishCallback& callback) { - CheckTopFrame(*state, frame_index, kind); +template +void FinishFrame(VariantBuildState& state, size_t frame_index, + const RootFinishCallback& callback) { + CheckTopFrame(state, frame_index); - auto& frame = state->frames.back(); - const auto values_size = state->value.length() - frame.value_start; + auto& frame = state.frames.back(); + const auto values_size = state.value.length() - frame.value_start; DCHECK_GE(values_size, 0); - if (values_size > std::numeric_limits::max()) { + if (std::cmp_greater(values_size, std::numeric_limits::max())) { throw ParquetException("Variant container values are too large"); } - auto header = kind == VariantContainerKind::Object - ? BuildObjectHeader(*state, frame, static_cast(values_size)) - : BuildListHeader(frame, static_cast(values_size)); - InsertBytes(state->value, frame.value_start, header); + if constexpr (kind == VariantContainerKind::Object) { + BuildObjectHeader(state, frame, static_cast(values_size)); + } else { + BuildListHeader(state, frame, static_cast(values_size)); + } const bool is_root = !frame.has_parent; frame.finished = true; - state->frames.pop_back(); + state.frames.pop_back(); if (!is_root) { return; } - state->root_has_value = true; - if (!callback) { - return; + state.root_has_value = true; + if (callback) { + callback(state); } - - callback(state); } template -void AppendRootPrimitiveWith(const std::shared_ptr& state, - Write&& write) { - CheckRootWritable(*state); - const auto value_size = state->value.length(); - VariantValueWriter writer(state->value); +void AppendValueWith(VariantBuildState& state, Write&& write) { + CheckRootWritable(state); + const auto value_size = state.value.length(); + VariantValueWriter writer(state.value); try { std::invoke(std::forward(write), writer); } catch (...) { - state->value.Rewind(value_size); + state.value.Rewind(value_size); throw; } - state->root_has_value = true; + state.root_has_value = true; } -template -void AppendRootPrimitive(const std::shared_ptr& state, - Args&&... args) { - AppendRootPrimitiveWith(state, [&](VariantValueWriter& writer) { - writer.template Append(std::forward(args)...); - }); -} +std::pair PrepareObjectField(VariantBuildState& state, size_t frame_index, + std::string_view field_name) { + CheckTopFrame(state, frame_index); + auto& frame = state.frames.back(); + const auto metadata_size = state.metadata->size(); + const auto entry_count = frame.fields.size(); -template -void AppendObjectPrimitiveWith(const std::shared_ptr& state, - size_t frame_index, std::string_view field_name, - Write&& write) { - CheckTopFrame(*state, frame_index, VariantContainerKind::Object); - auto& frame = state->frames.back(); - const auto metadata_size = state->metadata.size(); - const auto value_size = state->value.length(); - const auto field_count = frame.fields.size(); - - auto field_id = state->metadata.Upsert(field_name); + const auto field_id = state.metadata->TryUpsert(field_name); if (!frame.object_field_ids.insert(field_id).second) { - state->metadata.Truncate(metadata_size); + state.metadata->Truncate(metadata_size); throw ParquetException("Duplicate Variant object field: ", field_name); } - const auto offset = state->value.length() - frame.value_start; + const auto offset = state.value.length() - frame.value_start; DCHECK_GE(offset, 0); - if (offset > std::numeric_limits::max()) { - state->metadata.Truncate(metadata_size); - TruncateFrameEntries(frame, field_count); + if (std::cmp_greater(offset, std::numeric_limits::max())) { + state.metadata->Truncate(metadata_size); + TruncateFrameEntries(frame, entry_count); throw ParquetException("Variant object values are too large"); } frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, .offset = static_cast(offset)}); + return {metadata_size, entry_count}; +} + +template +void AppendObjectValueWith(VariantBuildState& state, size_t frame_index, + std::string_view field_name, Write&& write) { + const auto value_size = state.value.length(); + const auto [metadata_size, entry_count] = + PrepareObjectField(state, frame_index, field_name); - VariantValueWriter writer(state->value); + VariantValueWriter writer(state.value); try { std::invoke(std::forward(write), writer); } catch (...) { - state->value.Rewind(value_size); - state->metadata.Truncate(metadata_size); - TruncateFrameEntries(frame, field_count); + state.value.Rewind(value_size); + state.metadata->Truncate(metadata_size); + TruncateFrameEntries(state.frames[frame_index], entry_count); throw; } } -template -void AppendObjectPrimitive(const std::shared_ptr& state, - size_t frame_index, std::string_view field_name, - Args&&... args) { - AppendObjectPrimitiveWith(state, frame_index, field_name, - [&](VariantValueWriter& writer) { - writer.template Append(std::forward(args)...); - }); -} +size_t PrepareListElement(VariantBuildState& state, size_t frame_index) { + CheckTopFrame(state, frame_index); + auto& frame = state.frames.back(); + const auto entry_count = frame.offsets.size(); -template -void AppendListPrimitiveWith(const std::shared_ptr& state, - size_t frame_index, Write&& write) { - CheckTopFrame(*state, frame_index, VariantContainerKind::List); - auto& frame = state->frames.back(); - const auto value_size = state->value.length(); - const auto element_count = frame.offsets.size(); - const auto offset = state->value.length() - frame.value_start; + const auto offset = state.value.length() - frame.value_start; DCHECK_GE(offset, 0); - if (offset > std::numeric_limits::max()) { + if (std::cmp_greater(offset, std::numeric_limits::max())) { throw ParquetException("Variant array values are too large"); } frame.offsets.push_back(static_cast(offset)); + return entry_count; +} + +template +void AppendListValueWith(VariantBuildState& state, size_t frame_index, Write&& write) { + const auto value_size = state.value.length(); + const auto entry_count = PrepareListElement(state, frame_index); - VariantValueWriter writer(state->value); + VariantValueWriter writer(state.value); try { std::invoke(std::forward(write), writer); } catch (...) { - state->value.Rewind(value_size); - TruncateFrameEntries(frame, element_count); + state.value.Rewind(value_size); + TruncateFrameEntries(state.frames[frame_index], entry_count); throw; } } -template -void AppendListPrimitive(const std::shared_ptr& state, - size_t frame_index, Args&&... args) { - AppendListPrimitiveWith(state, frame_index, [&](VariantValueWriter& writer) { - writer.template Append(std::forward(args)...); - }); -} - -// Local helper for array builders that owns one shared byte arena plus the -// corresponding `binary_view` slots. This lets row commit write bytes once, -// create inline or out-of-line views in place, and roll back both buffers -// together on failure without routing hot-path appends through `BinaryViewBuilder`. -class BinaryViewColumnBuilder { - public: - struct Mark { - int64_t bytes_length = 0; - int64_t views_length = 0; - bool has_out_of_line_data = false; - }; - - explicit BinaryViewColumnBuilder(MemoryPool* pool) : bytes_(pool), views_(pool) {} - - [[nodiscard]] Mark mark() const { - return Mark{.bytes_length = bytes_.length(), - .views_length = views_.length(), - .has_out_of_line_data = has_out_of_line_data_}; - } - - void Rollback(const Mark& mark) { - bytes_.Rewind(mark.bytes_length); - views_.bytes_builder()->Rewind(mark.views_length * sizeof(BinaryViewType::c_type)); - has_out_of_line_data_ = mark.has_out_of_line_data; - } +} // namespace - int64_t length() const { return views_.length(); } +namespace internal { - std::string_view SliceFrom(int64_t start) const { - DCHECK_LE(start, bytes_.length()); - const auto size = bytes_.length() - start; - DCHECK_GE(size, 0); - if (size == 0) { - return std::string_view{}; - } - return std::string_view{reinterpret_cast(bytes_.data() + start), - static_cast(size)}; - } +struct NestedVariantBuilderImpl { + NestedVariantBuilderImpl(VariantBuildState& state, size_t frame_index, + RootFinishCallback callback) + : state(state), frame_index(frame_index), callback(std::move(callback)) {} - void AppendView(int64_t start) { - const auto slice = SliceFrom(start); - if (slice.size() > static_cast(std::numeric_limits::max())) { - throw ParquetException("Binary view value is too large"); - } - PARQUET_THROW_NOT_OK(views_.Reserve(1)); - if (slice.size() <= BinaryViewType::kInlineSize) { - views_.UnsafeAppend(::arrow::util::ToInlineBinaryView(slice)); - bytes_.Rewind(start); + ~NestedVariantBuilderImpl() { + if (!active || frame_index >= state.frames.size()) { return; } - if (start > std::numeric_limits::max()) { - throw ParquetException("Binary view offset is too large"); + const auto frame = state.frames[frame_index]; + state.value.Rewind(frame.value_start); + state.metadata->Truncate(frame.metadata_size); + if (frame.has_parent && frame.parent_frame < state.frames.size()) { + TruncateFrameEntries(state.frames[frame.parent_frame], frame.parent_entry_count); } - has_out_of_line_data_ = true; - views_.UnsafeAppend(::arrow::util::ToNonInlineBinaryView( - slice.data(), static_cast(slice.size()), - /*buffer_index=*/0, static_cast(start))); + state.frames.resize(frame_index); } - void AppendEmptyValue() { - PARQUET_THROW_NOT_OK(views_.Reserve(1)); - views_.UnsafeAppend(BinaryViewType::c_type{}); - } - - BufferBuilder& bytes_builder() { return bytes_; } - - std::shared_ptr Finish(std::shared_ptr null_bitmap = nullptr, - int64_t null_count = 0) { - const auto array_length = views_.length(); - std::shared_ptr views_buffer; - PARQUET_THROW_NOT_OK(views_.Finish(&views_buffer)); - - std::shared_ptr data_buffer; - PARQUET_ASSIGN_OR_THROW(data_buffer, bytes_.Finish()); - BufferVector data_buffers; - if (has_out_of_line_data_) { - data_buffers.push_back(std::move(data_buffer)); - } - return std::make_shared( - binary_view(), array_length, std::move(views_buffer), std::move(data_buffers), - std::move(null_bitmap), null_count); - } - - private: - BufferBuilder bytes_; - TypedBufferBuilder views_; - bool has_out_of_line_data_ = false; -}; - -} // namespace - -namespace internal { - -struct NestedVariantBuilderImpl { - NestedVariantBuilderImpl(std::shared_ptr state, size_t frame_index, - RootFinishCallback callback) - : state(std::move(state)), - frame_index(frame_index), - callback(std::move(callback)) {} - - std::shared_ptr state; + VariantBuildState& state; size_t frame_index = 0; RootFinishCallback callback; bool active = true; @@ -686,135 +652,93 @@ struct NestedVariantBuilderImpl { } // namespace internal struct VariantBuilder::Impl { - explicit Impl(MemoryPool* pool) + explicit Impl(MemoryPool* pool, bool validate) : pool(pool), value_builder(pool), - state(std::make_shared(pool, value_builder)) {} + validate(validate), + state(value_builder, std::make_unique()) {} void Reset() { value_builder.Reset(); - state = std::make_shared(pool, value_builder); + state.Reset(std::make_unique()); } MemoryPool* pool; BufferBuilder value_builder; - std::shared_ptr state; + bool validate; + VariantBuildState state; }; -VariantBuilder::VariantBuilder(MemoryPool* pool) : impl_(std::make_unique(pool)) {} +VariantBuilder::VariantBuilder(MemoryPool* pool, bool validate) + : impl_(std::make_unique(pool, validate)) {} VariantBuilder::~VariantBuilder() = default; VariantBuilder::VariantBuilder(VariantBuilder&&) noexcept = default; VariantBuilder& VariantBuilder::operator=(VariantBuilder&&) noexcept = default; void VariantBuilder::ReserveFieldNames(int64_t capacity) { - impl_->state->metadata.Reserve(capacity); + impl_->state.metadata->Reserve(capacity); } uint32_t VariantBuilder::AddFieldName(std::string_view name) { - return impl_->state->metadata.Upsert(name); -} - -void VariantBuilder::AppendVariantNull() { - AppendRootPrimitive(impl_->state); -} - -void VariantBuilder::AppendBoolean(bool value) { - if (value) { - AppendRootPrimitive(impl_->state); - return; - } - AppendRootPrimitive(impl_->state); + return impl_->state.metadata->TryUpsert(name); } -#define VARIANT_ROOT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - void VariantBuilder::Append##NAME(C_TYPE value) { \ - AppendRootPrimitive(impl_->state, value); \ +#define DEFINE_DIRECT_APPEND(NAME, decl_args, ...) \ + void VariantBuilder::Append##NAME decl_args { \ + AppendValueWith(impl_->state, [&](VariantValueWriter& writer) { \ + writer.Append(__VA_ARGS__); \ + }); \ } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_DIRECT_APPEND) +#undef DEFINE_DIRECT_APPEND -VARIANT_ROOT_APPEND_ONE_ARG(Int8, Int8, int8_t) -VARIANT_ROOT_APPEND_ONE_ARG(Int16, Int16, int16_t) -VARIANT_ROOT_APPEND_ONE_ARG(Int32, Int32, int32_t) -VARIANT_ROOT_APPEND_ONE_ARG(Int64, Int64, int64_t) -VARIANT_ROOT_APPEND_ONE_ARG(Float, Float, float) -VARIANT_ROOT_APPEND_ONE_ARG(Double, Double, double) -VARIANT_ROOT_APPEND_ONE_ARG(Binary, Binary, std::string_view) -VARIANT_ROOT_APPEND_ONE_ARG(String, String, std::string_view) -VARIANT_ROOT_APPEND_ONE_ARG(Date, Date, int32_t) -VARIANT_ROOT_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) -VARIANT_ROOT_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) - -#undef VARIANT_ROOT_APPEND_ONE_ARG - -void VariantBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { - AppendRootPrimitive(impl_->state, unscaled_value, - scale); -} - -void VariantBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { - AppendRootPrimitive(impl_->state, unscaled_value, - scale); -} - -void VariantBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, - uint8_t scale) { - AppendRootPrimitive( - impl_->state, little_endian_unscaled_value, scale); -} - -void VariantBuilder::AppendShortString(std::string_view value) { - AppendRootPrimitiveWith( - impl_->state, [&](VariantValueWriter& writer) { writer.AppendShortString(value); }); -} - -void VariantBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { - if (adjusted_to_utc) { - AppendRootPrimitive(impl_->state, micros); - return; +#define DEFINE_SPECIAL_APPEND(NAME, decl_args, ...) \ + void VariantBuilder::Append##NAME decl_args { \ + AppendValueWith(impl_->state, [&](VariantValueWriter& writer) { \ + writer.Append##NAME(__VA_ARGS__); \ + }); \ } - AppendRootPrimitive(impl_->state, micros); -} +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_SPECIAL_APPEND) +DEFINE_SPECIAL_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_SPECIAL_APPEND -void VariantBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { - if (adjusted_to_utc) { - AppendRootPrimitive(impl_->state, nanos); - return; - } - AppendRootPrimitive(impl_->state, nanos); +template +std::unique_ptr StartContainer( + VariantBuildState& state) { + CheckRootWritable(state); + state.frames.push_back(VariantBuildFrame{.kind = kind, + .value_start = state.value.length(), + .metadata_size = state.metadata->size()}); + return std::make_unique( + state, state.frames.size() - 1, RootFinishCallback{}); } VariantObjectBuilder VariantBuilder::StartObject() { - CheckRootWritable(*impl_->state); - impl_->state->frames.push_back( - VariantBuildFrame{.kind = VariantContainerKind::Object, - .value_start = impl_->state->value.length(), - .metadata_size = impl_->state->metadata.size()}); - return VariantObjectBuilder(std::make_unique( - impl_->state, impl_->state->frames.size() - 1, RootFinishCallback{})); + return VariantObjectBuilder(StartContainer(impl_->state)); } VariantListBuilder VariantBuilder::StartList() { - CheckRootWritable(*impl_->state); - impl_->state->frames.push_back( - VariantBuildFrame{.kind = VariantContainerKind::List, - .value_start = impl_->state->value.length(), - .metadata_size = impl_->state->metadata.size()}); - return VariantListBuilder(std::make_unique( - impl_->state, impl_->state->frames.size() - 1, RootFinishCallback{})); + return VariantListBuilder(StartContainer(impl_->state)); } EncodedVariantValue VariantBuilder::Finish() { - if (!impl_->state->frames.empty()) { + if (!impl_->state.frames.empty()) { throw ParquetException("Cannot finish VariantBuilder with active containers"); } - if (!impl_->state->root_has_value) { + if (!impl_->state.root_has_value) { throw ParquetException("Cannot finish empty VariantBuilder"); } - auto metadata = impl_->state->metadata.Finish(); - std::shared_ptr value; - PARQUET_ASSIGN_OR_THROW(value, impl_->state->value.Finish()); + BufferBuilder metadata_builder(impl_->pool); + impl_->state.metadata->AppendTo(metadata_builder); + PARQUET_ASSIGN_OR_THROW(auto metadata, metadata_builder.Finish()); + PARQUET_ASSIGN_OR_THROW(auto value, impl_->state.value.Finish()); EncodedVariantValue out{.metadata = std::move(metadata), .value = std::move(value)}; Reset(); + if (impl_->validate) { + auto metadata_view = VariantMetadataView::Make(std::string_view{*out.metadata}); + VariantValueView::Validate(std::string_view{*out.value}, metadata_view); + } return out; } @@ -823,332 +747,133 @@ void VariantBuilder::Reset() { impl_->Reset(); } VariantObjectBuilder::VariantObjectBuilder( std::unique_ptr impl) : impl_(std::move(impl)) {} -VariantObjectBuilder::~VariantObjectBuilder() { RollbackIfActive(impl_.get()); } +VariantObjectBuilder::~VariantObjectBuilder() = default; VariantObjectBuilder::VariantObjectBuilder(VariantObjectBuilder&&) noexcept = default; VariantObjectBuilder& VariantObjectBuilder::operator=(VariantObjectBuilder&&) noexcept = default; -void VariantObjectBuilder::AppendVariantNull(std::string_view field_name) { - AppendObjectPrimitive(impl_->state, impl_->frame_index, - field_name); -} - -void VariantObjectBuilder::AppendBoolean(std::string_view field_name, bool value) { - if (value) { - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name); - return; - } - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name); -} - -#define VARIANT_OBJECT_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - void VariantObjectBuilder::Append##NAME(std::string_view field_name, C_TYPE value) { \ - AppendObjectPrimitive( \ - impl_->state, impl_->frame_index, field_name, value); \ - } - -VARIANT_OBJECT_APPEND_ONE_ARG(Int8, Int8, int8_t) -VARIANT_OBJECT_APPEND_ONE_ARG(Int16, Int16, int16_t) -VARIANT_OBJECT_APPEND_ONE_ARG(Int32, Int32, int32_t) -VARIANT_OBJECT_APPEND_ONE_ARG(Int64, Int64, int64_t) -VARIANT_OBJECT_APPEND_ONE_ARG(Float, Float, float) -VARIANT_OBJECT_APPEND_ONE_ARG(Double, Double, double) -VARIANT_OBJECT_APPEND_ONE_ARG(Binary, Binary, std::string_view) -VARIANT_OBJECT_APPEND_ONE_ARG(String, String, std::string_view) -VARIANT_OBJECT_APPEND_ONE_ARG(Date, Date, int32_t) -VARIANT_OBJECT_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) -VARIANT_OBJECT_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) - -#undef VARIANT_OBJECT_APPEND_ONE_ARG - -void VariantObjectBuilder::AppendDecimal4(std::string_view field_name, - int32_t unscaled_value, uint8_t scale) { - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, unscaled_value, scale); -} - -void VariantObjectBuilder::AppendDecimal8(std::string_view field_name, - int64_t unscaled_value, uint8_t scale) { - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, unscaled_value, scale); -} - -void VariantObjectBuilder::AppendDecimal16(std::string_view field_name, - std::string_view little_endian_unscaled_value, - uint8_t scale) { - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, little_endian_unscaled_value, scale); -} - -void VariantObjectBuilder::AppendShortString(std::string_view field_name, - std::string_view value) { - AppendObjectPrimitiveWith( - impl_->state, impl_->frame_index, field_name, - [&](VariantValueWriter& writer) { writer.AppendShortString(value); }); -} - -void VariantObjectBuilder::AppendTimestampMicros(std::string_view field_name, - int64_t micros, bool adjusted_to_utc) { - if (adjusted_to_utc) { - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, micros); - return; - } - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, micros); -} - -void VariantObjectBuilder::AppendTimestampNanos(std::string_view field_name, - int64_t nanos, bool adjusted_to_utc) { - if (adjusted_to_utc) { - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, nanos); - return; - } - AppendObjectPrimitive( - impl_->state, impl_->frame_index, field_name, nanos); +#ifdef _MSC_VER +# define VARIANT_OBJECT_DECL(...) (std::string_view field_name, ##__VA_ARGS__) +#else +# define VARIANT_OBJECT_DECL(...) \ + (std::string_view field_name __VA_OPT__(, ) __VA_ARGS__) // NOLINT +#endif + +#define DEFINE_OBJECT_DIRECT_APPEND(NAME, decl_args, ...) \ + void VariantObjectBuilder::Append##NAME VARIANT_OBJECT_DECL decl_args { \ + AppendObjectValueWith(impl_->state, impl_->frame_index, field_name, \ + [&](VariantValueWriter& writer) { \ + writer.Append(__VA_ARGS__); \ + }); \ + } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_OBJECT_DIRECT_APPEND) +#undef DEFINE_OBJECT_DIRECT_APPEND + +#define DEFINE_OBJECT_SPECIAL_APPEND(NAME, decl_args, ...) \ + void VariantObjectBuilder::Append##NAME VARIANT_OBJECT_DECL decl_args { \ + AppendObjectValueWith( \ + impl_->state, impl_->frame_index, field_name, \ + [&](VariantValueWriter& writer) { writer.Append##NAME(__VA_ARGS__); }); \ + } +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_OBJECT_SPECIAL_APPEND) +DEFINE_OBJECT_SPECIAL_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_OBJECT_SPECIAL_APPEND + +#undef VARIANT_OBJECT_DECL + +template +std::unique_ptr StartChildContainer( + internal::NestedVariantBuilderImpl& impl, size_t metadata_size, size_t entry_count) { + impl.state.frames.push_back(VariantBuildFrame{.kind = kind, + .value_start = impl.state.value.length(), + .metadata_size = metadata_size, + .parent_frame = impl.frame_index, + .parent_entry_count = entry_count, + .has_parent = true}); + return std::make_unique( + impl.state, impl.state.frames.size() - 1, impl.callback); +} + +template +std::unique_ptr ObjectStartContainer( + internal::NestedVariantBuilderImpl& impl, std::string_view field_name) { + const auto [metadata_size, entry_count] = + PrepareObjectField(impl.state, impl.frame_index, field_name); + return StartChildContainer(impl, metadata_size, entry_count); } VariantObjectBuilder VariantObjectBuilder::StartObject(std::string_view field_name) { - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object); - auto& frame = impl_->state->frames.back(); - const auto metadata_size = impl_->state->metadata.size(); - const auto field_count = frame.fields.size(); - - auto field_id = impl_->state->metadata.Upsert(field_name); - if (!frame.object_field_ids.insert(field_id).second) { - impl_->state->metadata.Truncate(metadata_size); - throw ParquetException("Duplicate Variant object field: ", field_name); - } - const auto offset = impl_->state->value.length() - frame.value_start; - DCHECK_GE(offset, 0); - if (offset > std::numeric_limits::max()) { - impl_->state->metadata.Truncate(metadata_size); - TruncateFrameEntries(frame, field_count); - throw ParquetException("Variant object values are too large"); - } - - frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, - .offset = static_cast(offset)}); - const auto value_start = impl_->state->value.length(); - impl_->state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::Object, - .value_start = value_start, - .metadata_size = metadata_size, - .parent_frame = impl_->frame_index, - .parent_entry_count = field_count, - .has_parent = true}); - return VariantObjectBuilder(std::make_unique( - impl_->state, impl_->state->frames.size() - 1, impl_->callback)); + return VariantObjectBuilder( + ObjectStartContainer(*impl_, field_name)); } VariantListBuilder VariantObjectBuilder::StartList(std::string_view field_name) { - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::Object); - auto& frame = impl_->state->frames.back(); - const auto metadata_size = impl_->state->metadata.size(); - const auto field_count = frame.fields.size(); - - auto field_id = impl_->state->metadata.Upsert(field_name); - if (!frame.object_field_ids.insert(field_id).second) { - impl_->state->metadata.Truncate(metadata_size); - throw ParquetException("Duplicate Variant object field: ", field_name); - } - const auto offset = impl_->state->value.length() - frame.value_start; - DCHECK_GE(offset, 0); - if (offset > std::numeric_limits::max()) { - impl_->state->metadata.Truncate(metadata_size); - TruncateFrameEntries(frame, field_count); - throw ParquetException("Variant object values are too large"); - } - - frame.fields.push_back(VariantFieldDescriptor{.field_id = field_id, - .offset = static_cast(offset)}); - const auto value_start = impl_->state->value.length(); - impl_->state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::List, - .value_start = value_start, - .metadata_size = metadata_size, - .parent_frame = impl_->frame_index, - .parent_entry_count = field_count, - .has_parent = true}); - return VariantListBuilder(std::make_unique( - impl_->state, impl_->state->frames.size() - 1, impl_->callback)); + return VariantListBuilder( + ObjectStartContainer(*impl_, field_name)); } void VariantObjectBuilder::Finish() { - FinishFrame(impl_->state, impl_->frame_index, VariantContainerKind::Object, - impl_->callback); + FinishFrame(impl_->state, impl_->frame_index, + impl_->callback); impl_->active = false; } VariantListBuilder::VariantListBuilder( std::unique_ptr impl) : impl_(std::move(impl)) {} -VariantListBuilder::~VariantListBuilder() { RollbackIfActive(impl_.get()); } +VariantListBuilder::~VariantListBuilder() = default; VariantListBuilder::VariantListBuilder(VariantListBuilder&&) noexcept = default; VariantListBuilder& VariantListBuilder::operator=(VariantListBuilder&&) noexcept = default; -void VariantListBuilder::AppendVariantNull() { - AppendListPrimitive(impl_->state, impl_->frame_index); -} - -void VariantListBuilder::AppendBoolean(bool value) { - if (value) { - AppendListPrimitive(impl_->state, - impl_->frame_index); - return; +#define DEFINE_LIST_DIRECT_APPEND(NAME, decl_args, ...) \ + void VariantListBuilder::Append##NAME decl_args { \ + AppendListValueWith(impl_->state, impl_->frame_index, \ + [&](VariantValueWriter& writer) { \ + writer.Append(__VA_ARGS__); \ + }); \ } - AppendListPrimitive(impl_->state, - impl_->frame_index); -} - -#define VARIANT_LIST_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - void VariantListBuilder::Append##NAME(C_TYPE value) { \ - AppendListPrimitive(impl_->state, impl_->frame_index, \ - value); \ - } - -VARIANT_LIST_APPEND_ONE_ARG(Int8, Int8, int8_t) -VARIANT_LIST_APPEND_ONE_ARG(Int16, Int16, int16_t) -VARIANT_LIST_APPEND_ONE_ARG(Int32, Int32, int32_t) -VARIANT_LIST_APPEND_ONE_ARG(Int64, Int64, int64_t) -VARIANT_LIST_APPEND_ONE_ARG(Float, Float, float) -VARIANT_LIST_APPEND_ONE_ARG(Double, Double, double) -VARIANT_LIST_APPEND_ONE_ARG(Binary, Binary, std::string_view) -VARIANT_LIST_APPEND_ONE_ARG(String, String, std::string_view) -VARIANT_LIST_APPEND_ONE_ARG(Date, Date, int32_t) -VARIANT_LIST_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) -VARIANT_LIST_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) - -#undef VARIANT_LIST_APPEND_ONE_ARG - -void VariantListBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { - AppendListPrimitive(impl_->state, impl_->frame_index, - unscaled_value, scale); -} - -void VariantListBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { - AppendListPrimitive(impl_->state, impl_->frame_index, - unscaled_value, scale); -} +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_LIST_DIRECT_APPEND) +#undef DEFINE_LIST_DIRECT_APPEND -void VariantListBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, - uint8_t scale) { - AppendListPrimitive( - impl_->state, impl_->frame_index, little_endian_unscaled_value, scale); -} - -void VariantListBuilder::AppendShortString(std::string_view value) { - AppendListPrimitiveWith( - impl_->state, impl_->frame_index, - [&](VariantValueWriter& writer) { writer.AppendShortString(value); }); -} - -void VariantListBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { - if (adjusted_to_utc) { - AppendListPrimitive( - impl_->state, impl_->frame_index, micros); - return; +#define DEFINE_LIST_SPECIAL_APPEND(NAME, decl_args, ...) \ + void VariantListBuilder::Append##NAME decl_args { \ + AppendListValueWith( \ + impl_->state, impl_->frame_index, \ + [&](VariantValueWriter& writer) { writer.Append##NAME(__VA_ARGS__); }); \ } - AppendListPrimitive( - impl_->state, impl_->frame_index, micros); -} +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_LIST_SPECIAL_APPEND) +DEFINE_LIST_SPECIAL_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_LIST_SPECIAL_APPEND -void VariantListBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { - if (adjusted_to_utc) { - AppendListPrimitive(impl_->state, - impl_->frame_index, nanos); - return; - } - AppendListPrimitive( - impl_->state, impl_->frame_index, nanos); +template +std::unique_ptr ListStartContainer( + internal::NestedVariantBuilderImpl& impl) { + const auto entry_count = PrepareListElement(impl.state, impl.frame_index); + return StartChildContainer(impl, impl.state.metadata->size(), entry_count); } VariantObjectBuilder VariantListBuilder::StartObject() { - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List); - auto& frame = impl_->state->frames.back(); - const auto element_count = frame.offsets.size(); - const auto offset = impl_->state->value.length() - frame.value_start; - DCHECK_GE(offset, 0); - if (offset > std::numeric_limits::max()) { - throw ParquetException("Variant array values are too large"); - } - - frame.offsets.push_back(static_cast(offset)); - const auto value_start = impl_->state->value.length(); - impl_->state->frames.push_back( - VariantBuildFrame{.kind = VariantContainerKind::Object, - .value_start = value_start, - .metadata_size = impl_->state->metadata.size(), - .parent_frame = impl_->frame_index, - .parent_entry_count = element_count, - .has_parent = true}); - return VariantObjectBuilder(std::make_unique( - impl_->state, impl_->state->frames.size() - 1, impl_->callback)); + return VariantObjectBuilder(ListStartContainer(*impl_)); } VariantListBuilder VariantListBuilder::StartList() { - CheckTopFrame(*impl_->state, impl_->frame_index, VariantContainerKind::List); - auto& frame = impl_->state->frames.back(); - const auto element_count = frame.offsets.size(); - const auto offset = impl_->state->value.length() - frame.value_start; - DCHECK_GE(offset, 0); - if (offset > std::numeric_limits::max()) { - throw ParquetException("Variant array values are too large"); - } - - frame.offsets.push_back(static_cast(offset)); - const auto value_start = impl_->state->value.length(); - impl_->state->frames.push_back( - VariantBuildFrame{.kind = VariantContainerKind::List, - .value_start = value_start, - .metadata_size = impl_->state->metadata.size(), - .parent_frame = impl_->frame_index, - .parent_entry_count = element_count, - .has_parent = true}); - return VariantListBuilder(std::make_unique( - impl_->state, impl_->state->frames.size() - 1, impl_->callback)); + return VariantListBuilder(ListStartContainer(*impl_)); } void VariantListBuilder::Finish() { - FinishFrame(impl_->state, impl_->frame_index, VariantContainerKind::List, - impl_->callback); + FinishFrame(impl_->state, impl_->frame_index, + impl_->callback); impl_->active = false; } struct VariantArrayBuilder::Impl { explicit Impl(MemoryPool* pool) - : pool(pool), metadata_column(pool), value_column(pool), validity_builder(pool) {} - - template - void AppendValue(Write&& write) { - auto state = std::make_shared(pool, value_column.bytes_builder()); - AppendRootPrimitiveWith(state, std::forward(write)); - CommitBuiltRow(state); - } - - template - void AppendPrimitive(Args&&... args) { - AppendValue([&](VariantValueWriter& writer) { - writer.template Append(std::forward(args)...); - }); - } - - void AppendShortString(std::string_view value) { - AppendValue([&](VariantValueWriter& writer) { writer.AppendShortString(value); }); - } - - void AppendEncoded(const EncodedVariantValue& value) { - if (value.metadata == nullptr || value.value == nullptr) { - throw ParquetException( - "Encoded Variant metadata and value buffers must be non-null"); - } - const auto metadata_bytes = std::string_view{*value.metadata}; - const auto value_bytes = std::string_view{*value.value}; - auto metadata = VariantMetadataView::Make(metadata_bytes); - VariantValueView::Validate(value_bytes, metadata); - CommitEncodedRow(metadata_bytes, value_bytes); + : pool(pool), + metadata_column(pool), + value_column(pool), + validity_builder(pool), + state(value_column.bytes_builder(), std::make_unique()) { } void AppendNull() { @@ -1165,28 +890,39 @@ struct VariantArrayBuilder::Impl { } } - void CommitBuiltRow(const std::shared_ptr& state) { - if (!state->frames.empty()) { + template + void AppendValue(Write&& write) { + state.Reset(std::make_unique()); + AppendValueWith(state, std::forward(write)); + CommitBuiltRow(state); + } + + void AppendEncoded(const EncodedVariantValue& value) { + if (value.metadata == nullptr || value.value == nullptr) { + throw ParquetException( + "Encoded Variant metadata and value buffers must be non-null"); + } + const auto metadata_bytes = std::string_view{*value.metadata}; + const auto value_bytes = std::string_view{*value.value}; + CommitEncodedRow(metadata_bytes, value_bytes); + } + + void CommitBuiltRow(VariantBuildState& state) { + if (!state.frames.empty()) { throw ParquetException("Cannot commit Variant row with active containers"); } - if (!state->root_has_value) { + if (!state.root_has_value) { throw ParquetException("Cannot commit empty Variant row"); } - const auto value_start = state->root_value_start; + const auto value_start = state.root_value_start; const auto metadata_mark = metadata_column.mark(); auto value_mark = value_column.mark(); - DCHECK_LE(value_start, value_mark.bytes_length); + DCHECK_LT(value_start, value_mark.bytes_length); value_mark.bytes_length = value_start; try { const auto metadata_start = metadata_column.bytes_builder().length(); - state->metadata.AppendTo(metadata_column.bytes_builder()); - const auto metadata_bytes = metadata_column.SliceFrom(metadata_start); - auto metadata = VariantMetadataView::Make(metadata_bytes); - - const auto value_bytes = value_column.SliceFrom(value_start); - VariantValueView::Validate(value_bytes, metadata); - + state.metadata->AppendTo(metadata_column.bytes_builder()); metadata_column.AppendView(metadata_start); value_column.AppendView(value_start); PARQUET_THROW_NOT_OK(validity_builder.Append(true)); @@ -1211,17 +947,17 @@ struct VariantArrayBuilder::Impl { auto value = value_column.Finish(); auto storage_type = struct_({field("metadata", binary_view(), /*nullable=*/false), field("value", binary_view(), /*nullable=*/false)}); - std::shared_ptr storage; - PARQUET_ASSIGN_OR_THROW(storage, + PARQUET_ASSIGN_OR_THROW(auto storage, StructArray::Make({metadata, value}, storage_type->fields(), std::move(null_bitmap), null_count)); - return MakeVariantArrayFromStorage(storage); + return MakeVariantArrayFromStorage(std::move(storage)); } MemoryPool* pool; BinaryViewColumnBuilder metadata_column; BinaryViewColumnBuilder value_column; TypedBufferBuilder validity_builder; + VariantBuildState state; private: void CommitEncodedRow(std::string_view metadata_bytes, std::string_view value_bytes) { @@ -1253,112 +989,221 @@ VariantArrayBuilder& VariantArrayBuilder::operator=(VariantArrayBuilder&&) noexc void VariantArrayBuilder::AppendNull() { impl_->AppendNull(); } -void VariantArrayBuilder::AppendVariantNull() { - impl_->AppendPrimitive(); -} - -void VariantArrayBuilder::AppendBoolean(bool value) { - if (value) { - impl_->AppendPrimitive(); - return; +#define DEFINE_ARRAY_DIRECT_APPEND(NAME, decl_args, ...) \ + void VariantArrayBuilder::Append##NAME decl_args { \ + impl_->AppendValue([&](VariantValueWriter& writer) { \ + writer.Append(__VA_ARGS__); \ + }); \ } - impl_->AppendPrimitive(); -} +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_ARRAY_DIRECT_APPEND) +#undef DEFINE_ARRAY_DIRECT_APPEND -#define VARIANT_ARRAY_APPEND_ONE_ARG(NAME, TYPE, C_TYPE) \ - void VariantArrayBuilder::Append##NAME(C_TYPE value) { \ - impl_->AppendPrimitive(value); \ +#define DEFINE_ARRAY_SPECIAL_APPEND(NAME, decl_args, ...) \ + void VariantArrayBuilder::Append##NAME decl_args { \ + impl_->AppendValue( \ + [&](VariantValueWriter& writer) { writer.Append##NAME(__VA_ARGS__); }); \ } +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_ARRAY_SPECIAL_APPEND) +#undef DEFINE_ARRAY_SPECIAL_APPEND -VARIANT_ARRAY_APPEND_ONE_ARG(Int8, Int8, int8_t) -VARIANT_ARRAY_APPEND_ONE_ARG(Int16, Int16, int16_t) -VARIANT_ARRAY_APPEND_ONE_ARG(Int32, Int32, int32_t) -VARIANT_ARRAY_APPEND_ONE_ARG(Int64, Int64, int64_t) -VARIANT_ARRAY_APPEND_ONE_ARG(Float, Float, float) -VARIANT_ARRAY_APPEND_ONE_ARG(Double, Double, double) -VARIANT_ARRAY_APPEND_ONE_ARG(Binary, Binary, std::string_view) -VARIANT_ARRAY_APPEND_ONE_ARG(String, String, std::string_view) -VARIANT_ARRAY_APPEND_ONE_ARG(Date, Date, int32_t) -VARIANT_ARRAY_APPEND_ONE_ARG(TimeNTZMicros, TimeNTZMicros, int64_t) -VARIANT_ARRAY_APPEND_ONE_ARG(Uuid, Uuid, std::string_view) - -#undef VARIANT_ARRAY_APPEND_ONE_ARG +void VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { + impl_->AppendEncoded(value); +} -void VariantArrayBuilder::AppendDecimal4(int32_t unscaled_value, uint8_t scale) { - impl_->AppendPrimitive(unscaled_value, scale); +template +std::unique_ptr ArrayStartContainer(ImplType& impl) { + impl.state.Reset(std::make_unique()); + impl.state.frames.push_back( + VariantBuildFrame{.kind = kind, + .value_start = impl.state.value.length(), + .metadata_size = impl.state.metadata->size()}); + return std::make_unique( + impl.state, impl.state.frames.size() - 1, + std::bind_front(&ImplType::CommitBuiltRow, &impl)); } -void VariantArrayBuilder::AppendDecimal8(int64_t unscaled_value, uint8_t scale) { - impl_->AppendPrimitive(unscaled_value, scale); +VariantObjectBuilder VariantArrayBuilder::StartObject() { + return VariantObjectBuilder(ArrayStartContainer(*impl_)); } -void VariantArrayBuilder::AppendDecimal16(std::string_view little_endian_unscaled_value, - uint8_t scale) { - impl_->AppendPrimitive(little_endian_unscaled_value, - scale); +VariantListBuilder VariantArrayBuilder::StartList() { + return VariantListBuilder(ArrayStartContainer(*impl_)); } -void VariantArrayBuilder::AppendShortString(std::string_view value) { - impl_->AppendShortString(value); +std::shared_ptr VariantArrayBuilder::Finish() { + auto out = impl_->Finish(); + Reset(); + return out; } -void VariantArrayBuilder::AppendTimestampMicros(int64_t micros, bool adjusted_to_utc) { - if (adjusted_to_utc) { - impl_->AppendPrimitive(micros); - return; +void VariantArrayBuilder::Reset() { impl_ = std::make_unique(impl_->pool); } + +struct VariantValueArrayBuilder::Impl { + explicit Impl(MemoryPool* pool) + : pool(pool), value_column(pool), validity_builder(pool) {} + + void AppendNull() { + const auto value_mark = value_column.mark(); + try { + value_column.AppendEmptyValue(); + PARQUET_THROW_NOT_OK(validity_builder.Append(false)); + } catch (...) { + value_column.Rollback(value_mark); + throw; + } } - impl_->AppendPrimitive(micros); -} -void VariantArrayBuilder::AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc) { - if (adjusted_to_utc) { - impl_->AppendPrimitive(nanos); - return; + void AppendEncodedValue(std::string_view value) { + if (value.empty()) { + throw ParquetException("Encoded Variant value length cannot be 0"); + } + const auto value_mark = value_column.mark(); + try { + const auto value_start = value_column.bytes_builder().length(); + PARQUET_THROW_NOT_OK(value_column.bytes_builder().Append(value)); + value_column.AppendView(value_start); + PARQUET_THROW_NOT_OK(validity_builder.Append(true)); + } catch (...) { + value_column.Rollback(value_mark); + throw; + } } - impl_->AppendPrimitive(nanos); -} -void VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { - impl_->AppendEncoded(value); -} + void CommitBuiltRow(VariantBuildState& state) { + if (!state.frames.empty()) { + throw ParquetException("Cannot commit Variant value row with active containers"); + } + if (!state.root_has_value) { + throw ParquetException("Cannot commit empty Variant value row"); + } -VariantObjectBuilder VariantArrayBuilder::StartObject() { - auto state = std::make_shared(impl_->pool, - impl_->value_column.bytes_builder()); - state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::Object, - .value_start = state->value.length(), - .metadata_size = state->metadata.size()}); - return VariantObjectBuilder(std::make_unique( - state, state->frames.size() - 1, - std::bind_front(&Impl::CommitBuiltRow, impl_.get()))); + const auto value_start = state.root_value_start; + auto value_mark = value_column.mark(); + DCHECK_LT(value_start, value_mark.bytes_length); + value_mark.bytes_length = value_start; + try { + value_column.AppendView(value_start); + PARQUET_THROW_NOT_OK(validity_builder.Append(true)); + } catch (...) { + value_column.Rollback(value_mark); + throw; + } + } + + void DiscardBuiltRow(VariantBuildState& state) { + value_column.bytes_builder().Rewind(state.root_value_start); + state.frames.clear(); + state.root_has_value = false; + } + + std::shared_ptr Finish() { + const int64_t null_count = validity_builder.false_count(); + std::shared_ptr null_bitmap; + if (null_count > 0) { + PARQUET_THROW_NOT_OK(validity_builder.Finish(&null_bitmap)); + } + return value_column.Finish(std::move(null_bitmap), null_count); + } + + MemoryPool* pool; + BinaryViewColumnBuilder value_column; + TypedBufferBuilder validity_builder; +}; + +VariantValueArrayBuilder::VariantValueArrayBuilder(MemoryPool* pool) + : impl_(std::make_unique(pool)) {} +VariantValueArrayBuilder::~VariantValueArrayBuilder() = default; +VariantValueArrayBuilder::VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept = + default; +VariantValueArrayBuilder& VariantValueArrayBuilder::operator=( + VariantValueArrayBuilder&&) noexcept = default; + +void VariantValueArrayBuilder::AppendNull() { impl_->AppendNull(); } + +void VariantValueArrayBuilder::AppendEncodedValue(std::string_view value) { + impl_->AppendEncodedValue(value); } -VariantListBuilder VariantArrayBuilder::StartList() { - auto state = std::make_shared(impl_->pool, - impl_->value_column.bytes_builder()); - state->frames.push_back(VariantBuildFrame{.kind = VariantContainerKind::List, - .value_start = state->value.length(), - .metadata_size = state->metadata.size()}); - return VariantListBuilder(std::make_unique( - state, state->frames.size() - 1, - std::bind_front(&Impl::CommitBuiltRow, impl_.get()))); +VariantValueRowBuilder VariantValueArrayBuilder::BindMetadata( + const VariantMetadataView& metadata) { + return {*this, metadata}; } -std::shared_ptr VariantArrayBuilder::Finish() { +std::shared_ptr VariantValueArrayBuilder::Finish() { auto out = impl_->Finish(); Reset(); return out; } -void VariantArrayBuilder::Reset() { impl_ = std::make_unique(impl_->pool); } +void VariantValueArrayBuilder::Reset() { impl_ = std::make_unique(impl_->pool); } + +struct VariantValueRowBuilder::Impl { + Impl(VariantValueArrayBuilder::Impl& parent, const VariantMetadataView& metadata) + : parent(parent), + state(parent.value_column.bytes_builder(), + std::make_unique(metadata)) {} + + ~Impl() { + if (!finished) { + parent.DiscardBuiltRow(state); + } + } + + VariantValueArrayBuilder::Impl& parent; + VariantBuildState state; + bool finished = false; +}; + +VariantValueRowBuilder::VariantValueRowBuilder(VariantValueArrayBuilder& parent, + const VariantMetadataView& metadata) + : impl_(std::make_unique(*parent.impl_, metadata)) {} +VariantValueRowBuilder::~VariantValueRowBuilder() = default; +VariantValueRowBuilder::VariantValueRowBuilder(VariantValueRowBuilder&&) noexcept = + default; +VariantValueRowBuilder& VariantValueRowBuilder::operator=( + VariantValueRowBuilder&&) noexcept = default; + +#define DEFINE_VALUE_ROW_DIRECT_APPEND(NAME, decl_args, ...) \ + void VariantValueRowBuilder::Append##NAME decl_args { \ + AppendValueWith(impl_->state, [&](VariantValueWriter& writer) { \ + writer.Append(__VA_ARGS__); \ + }); \ + } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_VALUE_ROW_DIRECT_APPEND) +#undef DEFINE_VALUE_ROW_DIRECT_APPEND + +#define DEFINE_VALUE_ROW_SPECIAL_APPEND(NAME, decl_args, ...) \ + void VariantValueRowBuilder::Append##NAME decl_args { \ + AppendValueWith(impl_->state, [&](VariantValueWriter& writer) { \ + writer.Append##NAME(__VA_ARGS__); \ + }); \ + } +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_VALUE_ROW_SPECIAL_APPEND) +DEFINE_VALUE_ROW_SPECIAL_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_VALUE_ROW_SPECIAL_APPEND + +VariantObjectBuilder VariantValueRowBuilder::StartObject() { + return VariantObjectBuilder(StartContainer(impl_->state)); +} + +VariantListBuilder VariantValueRowBuilder::StartList() { + return VariantListBuilder(StartContainer(impl_->state)); +} + +void VariantValueRowBuilder::Finish() { + if (impl_->finished) { + throw ParquetException("Variant value row is already finished"); + } + impl_->parent.CommitBuiltRow(impl_->state); + impl_->finished = true; +} std::shared_ptr MakeVariantArrayFromStorage( std::shared_ptr storage) { if (storage == nullptr) { throw ParquetException("Variant storage array must be non-null"); } - std::shared_ptr type; - PARQUET_ASSIGN_OR_THROW(type, VariantExtensionType::Make(storage->type())); + PARQUET_ASSIGN_OR_THROW(auto type, VariantExtensionType::Make(storage->type())); auto array = ExtensionType::WrapArray(type, std::move(storage)); return std::static_pointer_cast(array); } @@ -1393,11 +1238,10 @@ std::shared_ptr MakeVariantArrayFromChildren( } } - std::shared_ptr storage; - PARQUET_ASSIGN_OR_THROW(storage, + PARQUET_ASSIGN_OR_THROW(auto storage, StructArray::Make(std::move(children), struct_type.fields(), std::move(null_bitmap))); - return MakeVariantArrayFromStorage(storage); + return MakeVariantArrayFromStorage(std::move(storage)); } } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/builder.h b/cpp/src/parquet/variant/builder.h index d45a4a77bcfd..380a5eafa2c6 100644 --- a/cpp/src/parquet/variant/builder.h +++ b/cpp/src/parquet/variant/builder.h @@ -25,7 +25,7 @@ #include "arrow/buffer.h" #include "arrow/extension/parquet_variant.h" #include "arrow/memory_pool.h" -#include "arrow/type.h" +#include "arrow/type_fwd.h" #include "parquet/platform.h" namespace parquet::variant { @@ -38,17 +38,20 @@ using ::arrow::MemoryPool; using ::arrow::StructArray; using ::arrow::extension::VariantArray; +class VariantMetadataView; +class VariantObjectBuilder; +class VariantListBuilder; +class VariantValueRowBuilder; + struct PARQUET_EXPORT EncodedVariantValue { std::shared_ptr metadata; std::shared_ptr value; }; -class VariantObjectBuilder; -class VariantListBuilder; - class PARQUET_EXPORT VariantBuilder { public: - explicit VariantBuilder(MemoryPool* pool = ::arrow::default_memory_pool()); + explicit VariantBuilder(MemoryPool* pool = ::arrow::default_memory_pool(), + bool validate = true); ~VariantBuilder(); VariantBuilder(const VariantBuilder&) = delete; VariantBuilder& operator=(const VariantBuilder&) = delete; @@ -71,13 +74,17 @@ class PARQUET_EXPORT VariantBuilder { void AppendDate(int32_t days); void AppendTimeNTZMicros(int64_t micros); void AppendUuid(std::string_view big_endian_bytes); - void AppendDecimal4(int32_t unscaled_value, uint8_t scale); - void AppendDecimal8(int64_t unscaled_value, uint8_t scale); - void AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendDecimal4(const ::arrow::Decimal32& value, uint8_t scale); + void AppendDecimal8(const ::arrow::Decimal64& value, uint8_t scale); + void AppendDecimal16(const ::arrow::Decimal128& value, uint8_t scale); void AppendShortString(std::string_view value); void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + void AppendEncodedValue(std::string_view value); + /// The returned nested builder borrows this builder's state. This builder must + /// outlive it. Finish or destroy the nested builder before calling any other method + /// on this builder, including Finish() or Reset(). VariantObjectBuilder StartObject(); VariantListBuilder StartList(); @@ -114,15 +121,18 @@ class PARQUET_EXPORT VariantObjectBuilder { void AppendDate(std::string_view field_name, int32_t days); void AppendTimeNTZMicros(std::string_view field_name, int64_t micros); void AppendUuid(std::string_view field_name, std::string_view big_endian_bytes); - void AppendDecimal4(std::string_view field_name, int32_t unscaled_value, uint8_t scale); - void AppendDecimal8(std::string_view field_name, int64_t unscaled_value, uint8_t scale); - void AppendDecimal16(std::string_view field_name, - std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendDecimal4(std::string_view field_name, const ::arrow::Decimal32& value, + uint8_t scale); + void AppendDecimal8(std::string_view field_name, const ::arrow::Decimal64& value, + uint8_t scale); + void AppendDecimal16(std::string_view field_name, const ::arrow::Decimal128& value, + uint8_t scale); void AppendShortString(std::string_view field_name, std::string_view value); void AppendTimestampMicros(std::string_view field_name, int64_t micros, bool adjusted_to_utc); void AppendTimestampNanos(std::string_view field_name, int64_t nanos, bool adjusted_to_utc); + void AppendEncodedValue(std::string_view field_name, std::string_view value); VariantObjectBuilder StartObject(std::string_view field_name); VariantListBuilder StartList(std::string_view field_name); @@ -135,6 +145,7 @@ class PARQUET_EXPORT VariantObjectBuilder { friend class VariantBuilder; friend class VariantListBuilder; friend class VariantArrayBuilder; + friend class VariantValueRowBuilder; explicit VariantObjectBuilder(std::unique_ptr impl); @@ -162,12 +173,13 @@ class PARQUET_EXPORT VariantListBuilder { void AppendDate(int32_t days); void AppendTimeNTZMicros(int64_t micros); void AppendUuid(std::string_view big_endian_bytes); - void AppendDecimal4(int32_t unscaled_value, uint8_t scale); - void AppendDecimal8(int64_t unscaled_value, uint8_t scale); - void AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendDecimal4(const ::arrow::Decimal32& value, uint8_t scale); + void AppendDecimal8(const ::arrow::Decimal64& value, uint8_t scale); + void AppendDecimal16(const ::arrow::Decimal128& value, uint8_t scale); void AppendShortString(std::string_view value); void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + void AppendEncodedValue(std::string_view value); VariantObjectBuilder StartObject(); VariantListBuilder StartList(); @@ -180,6 +192,7 @@ class PARQUET_EXPORT VariantListBuilder { friend class VariantBuilder; friend class VariantObjectBuilder; friend class VariantArrayBuilder; + friend class VariantValueRowBuilder; explicit VariantListBuilder(std::unique_ptr impl); @@ -209,9 +222,9 @@ class PARQUET_EXPORT VariantArrayBuilder { void AppendDate(int32_t days); void AppendTimeNTZMicros(int64_t micros); void AppendUuid(std::string_view big_endian_bytes); - void AppendDecimal4(int32_t unscaled_value, uint8_t scale); - void AppendDecimal8(int64_t unscaled_value, uint8_t scale); - void AppendDecimal16(std::string_view little_endian_unscaled_value, uint8_t scale); + void AppendDecimal4(const ::arrow::Decimal32& value, uint8_t scale); + void AppendDecimal8(const ::arrow::Decimal64& value, uint8_t scale); + void AppendDecimal16(const ::arrow::Decimal128& value, uint8_t scale); void AppendShortString(std::string_view value); void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); @@ -231,6 +244,81 @@ class PARQUET_EXPORT VariantArrayBuilder { std::unique_ptr impl_; }; +class PARQUET_EXPORT VariantValueArrayBuilder { + public: + explicit VariantValueArrayBuilder(MemoryPool* pool = ::arrow::default_memory_pool()); + ~VariantValueArrayBuilder(); + VariantValueArrayBuilder(const VariantValueArrayBuilder&) = delete; + VariantValueArrayBuilder& operator=(const VariantValueArrayBuilder&) = delete; + VariantValueArrayBuilder(VariantValueArrayBuilder&&) noexcept; + VariantValueArrayBuilder& operator=(VariantValueArrayBuilder&&) noexcept; + + void AppendNull(); + void AppendEncodedValue(std::string_view value); + + /// The returned row builder borrows this builder's internal value buffer. This + /// builder, the metadata view, and its underlying bytes must remain valid until the + /// row builder is finished or destroyed. Do not call any other method on this builder, + /// including Finish() or Reset(), while the row builder is unfinished. + VariantValueRowBuilder BindMetadata(const VariantMetadataView& metadata); + VariantValueRowBuilder BindMetadata(VariantMetadataView&& metadata) = delete; + + std::shared_ptr Finish(); + void Reset(); + + private: + friend class VariantValueRowBuilder; + + struct Impl; + std::unique_ptr impl_; +}; + +class PARQUET_EXPORT VariantValueRowBuilder { + public: + ~VariantValueRowBuilder(); + VariantValueRowBuilder(const VariantValueRowBuilder&) = delete; + VariantValueRowBuilder& operator=(const VariantValueRowBuilder&) = delete; + VariantValueRowBuilder(VariantValueRowBuilder&&) noexcept; + VariantValueRowBuilder& operator=(VariantValueRowBuilder&&) noexcept; + + void AppendVariantNull(); + void AppendBoolean(bool value); + void AppendInt8(int8_t value); + void AppendInt16(int16_t value); + void AppendInt32(int32_t value); + void AppendInt64(int64_t value); + void AppendFloat(float value); + void AppendDouble(double value); + void AppendBinary(std::string_view value); + void AppendString(std::string_view value); + void AppendDate(int32_t days); + void AppendTimeNTZMicros(int64_t micros); + void AppendUuid(std::string_view big_endian_bytes); + void AppendDecimal4(const ::arrow::Decimal32& value, uint8_t scale); + void AppendDecimal8(const ::arrow::Decimal64& value, uint8_t scale); + void AppendDecimal16(const ::arrow::Decimal128& value, uint8_t scale); + void AppendShortString(std::string_view value); + void AppendTimestampMicros(int64_t micros, bool adjusted_to_utc); + void AppendTimestampNanos(int64_t nanos, bool adjusted_to_utc); + void AppendEncodedValue(std::string_view value); + + /// The returned nested builder borrows this row builder's state. This row builder + /// must outlive it. Finish or destroy the nested builder before calling any other + /// method on this row builder, including Finish(). + VariantObjectBuilder StartObject(); + VariantListBuilder StartList(); + void Finish(); + + private: + friend class VariantValueArrayBuilder; + + VariantValueRowBuilder(VariantValueArrayBuilder& parent, + const VariantMetadataView& metadata); + + struct Impl; + std::unique_ptr impl_; +}; + PARQUET_EXPORT std::shared_ptr MakeVariantArrayFromStorage( std::shared_ptr storage); diff --git a/cpp/src/parquet/variant/builder_test.cc b/cpp/src/parquet/variant/builder_test.cc index 54471d9a3b6f..ec0312e9fb41 100644 --- a/cpp/src/parquet/variant/builder_test.cc +++ b/cpp/src/parquet/variant/builder_test.cc @@ -16,16 +16,23 @@ // under the License. #include "parquet/variant/builder.h" -#include "parquet/variant/encoding.h" +#include "parquet/variant/decoding.h" #include "parquet/variant/test_util_internal.h" +#include #include #include #include +#include #include "arrow/array.h" // IWYU pragma: keep +#include "arrow/array/util.h" +#include "arrow/scalar.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/checked_cast.h" +#include "arrow/util/decimal.h" + +using namespace std::string_view_literals; // NOLINT namespace parquet::variant { @@ -33,17 +40,28 @@ namespace { using ::arrow::ArrayFromJSON; using ::arrow::binary; +using ::arrow::binary_view; using ::arrow::BinaryViewArray; +using ::arrow::BinaryViewScalar; using ::arrow::default_memory_pool; using ::arrow::field; using ::arrow::int64; +using ::arrow::MakeArrayFromScalar; using ::arrow::ProxyMemoryPool; using ::arrow::struct_; using ::arrow::StructArray; using ::arrow::Type; using ::arrow::extension::variant; using internal::BinaryArrayFromValues; -using internal::MakeVariantValueView; + +template +concept CanBindMetadata = + (requires(VariantValueArrayBuilder& builder, Metadata&& metadata) { + builder.BindMetadata(std::forward(metadata)); + }); + +static_assert(CanBindMetadata); +static_assert(!CanBindMetadata); void AssertPrimitiveType(std::string_view value, const VariantMetadataView& metadata, VariantPrimitiveType expected) { @@ -53,11 +71,19 @@ void AssertPrimitiveType(std::string_view value, const VariantMetadataView& meta } void AssertPrimitiveFieldType(const VariantObjectView& object, std::string_view name, - const VariantMetadataView& metadata, VariantPrimitiveType expected) { - const auto* field = object.FindField(name); - ASSERT_NE(nullptr, field) << "Missing Variant object field: " << name; - AssertPrimitiveType(field->value, metadata, expected); + auto field = object.GetField(name); + ASSERT_TRUE(field.has_value()) << "Missing Variant object field: " << name; + ASSERT_EQ(VariantBasicType::kPrimitive, field->basic_type()); + ASSERT_EQ(expected, std::get(field->data()).type()); +} + +std::pair MakeEmptyMetadata() { + VariantBuilder builder; + builder.AppendVariantNull(); + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + return {std::move(encoded), std::move(metadata)}; } } // namespace @@ -71,7 +97,8 @@ TEST(TestVariantBuilder, Object) { object.Finish(); auto encoded = builder.Finish(); - auto view = MakeVariantValueView(encoded); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); ASSERT_EQ(VariantBasicType::kObject, view.basic_type()); const auto& fields = std::get(view.data()).fields(); ASSERT_EQ(3, fields.size()); @@ -92,12 +119,38 @@ TEST(TestVariantBuilder, List) { auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); ASSERT_EQ(VariantBasicType::kArray, view.basic_type()); - const auto& elements = std::get(view.data()).elements(); - ASSERT_EQ(3, elements.size()); + const auto& array = std::get(view.data()); + ASSERT_EQ(3, array.elements().size()); - auto element = VariantValueView::Make(elements[1], metadata); + auto element = array.GetElement(1); + ASSERT_TRUE(element.has_value()); ASSERT_EQ(VariantPrimitiveType::kInt32, - std::get(element.data()).type()); + std::get(element->data()).type()); +} + +TEST(TestVariantBuilder, Reset) { + VariantBuilder builder; + auto object = builder.StartObject(); + object.AppendInt32("old", 1); + object.Finish(); + auto first = builder.Finish(); + + auto list = builder.StartList(); + list.AppendInt32(2); + list.Finish(); + auto second = builder.Finish(); + + auto first_metadata = VariantMetadataView::Make(std::string_view{*first.metadata}); + ASSERT_TRUE(first_metadata.FindString("old").has_value()); + + auto second_metadata = VariantMetadataView::Make(std::string_view{*second.metadata}); + ASSERT_FALSE(second_metadata.FindString("old").has_value()); + + auto second_view = + VariantValueView::Make(std::string_view{*second.value}, second_metadata); + const auto& elements = std::get(second_view.data()).elements(); + ASSERT_EQ(1, elements.size()); + AssertPrimitiveType(elements[0], second_metadata, VariantPrimitiveType::kInt32); } TEST(TestVariantBuilder, Nested) { @@ -114,15 +167,17 @@ TEST(TestVariantBuilder, Nested) { auto encoded = builder.Finish(); auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); auto root = VariantValueView::Make(std::string_view{*encoded.value}, metadata); - const auto& root_fields = std::get(root.data()).fields(); - ASSERT_EQ(1, root_fields.size()); - ASSERT_EQ("items", root_fields[0].name); - - auto items = VariantValueView::Make(root_fields[0].value, metadata); - const auto& item_values = std::get(items.data()).elements(); - ASSERT_EQ(2, item_values.size()); - auto item = VariantValueView::Make(item_values[1], metadata); - ASSERT_EQ("name", std::get(item.data()).fields()[0].name); + const auto& root_object = std::get(root.data()); + ASSERT_EQ(1, root_object.fields().size()); + ASSERT_EQ("items", root_object.fields()[0].name); + + auto items = root_object.GetField("items"); + ASSERT_TRUE(items.has_value()); + const auto& items_array = std::get(items->data()); + ASSERT_EQ(2, items_array.elements().size()); + auto item = items_array.GetElement(1); + ASSERT_TRUE(item.has_value()); + ASSERT_EQ("name", std::get(item->data()).fields()[0].name); } TEST(TestVariantBuilder, ObjectAppends) { @@ -132,7 +187,7 @@ TEST(TestVariantBuilder, ObjectAppends) { object.AppendInt8("int8", 1); object.AppendInt64("int64", 2); object.AppendDouble("double", 3); - object.AppendDecimal4("decimal", 1234, 2); + object.AppendDecimal4("decimal", ::arrow::Decimal32(1234), 2); object.AppendBinary("binary", "abc"); object.AppendDate("date", 1); object.AppendTimestampNanos("ts", 2, true); @@ -144,18 +199,14 @@ TEST(TestVariantBuilder, ObjectAppends) { auto root = VariantValueView::Make(std::string_view{*encoded.value}, metadata); const auto& object_view = std::get(root.data()); ASSERT_EQ(8, object_view.fields().size()); - AssertPrimitiveFieldType(object_view, "int8", metadata, VariantPrimitiveType::kInt8); - AssertPrimitiveFieldType(object_view, "int64", metadata, VariantPrimitiveType::kInt64); - AssertPrimitiveFieldType(object_view, "double", metadata, - VariantPrimitiveType::kDouble); - AssertPrimitiveFieldType(object_view, "decimal", metadata, - VariantPrimitiveType::kDecimal4); - AssertPrimitiveFieldType(object_view, "binary", metadata, - VariantPrimitiveType::kBinary); - AssertPrimitiveFieldType(object_view, "date", metadata, VariantPrimitiveType::kDate); - AssertPrimitiveFieldType(object_view, "ts", metadata, - VariantPrimitiveType::kTimestampNanos); - AssertPrimitiveFieldType(object_view, "uuid", metadata, VariantPrimitiveType::kUuid); + AssertPrimitiveFieldType(object_view, "int8", VariantPrimitiveType::kInt8); + AssertPrimitiveFieldType(object_view, "int64", VariantPrimitiveType::kInt64); + AssertPrimitiveFieldType(object_view, "double", VariantPrimitiveType::kDouble); + AssertPrimitiveFieldType(object_view, "decimal", VariantPrimitiveType::kDecimal4); + AssertPrimitiveFieldType(object_view, "binary", VariantPrimitiveType::kBinary); + AssertPrimitiveFieldType(object_view, "date", VariantPrimitiveType::kDate); + AssertPrimitiveFieldType(object_view, "ts", VariantPrimitiveType::kTimestampNanos); + AssertPrimitiveFieldType(object_view, "uuid", VariantPrimitiveType::kUuid); } TEST(TestVariantBuilder, ListAppends) { @@ -165,7 +216,7 @@ TEST(TestVariantBuilder, ListAppends) { list.AppendInt8(1); list.AppendInt64(2); list.AppendDouble(3); - list.AppendDecimal4(1234, 2); + list.AppendDecimal4(::arrow::Decimal32(1234), 2); list.AppendBinary("abc"); list.AppendDate(1); list.AppendTimestampNanos(2, true); @@ -187,21 +238,63 @@ TEST(TestVariantBuilder, ListAppends) { AssertPrimitiveType(elements[7], metadata, VariantPrimitiveType::kUuid); } -TEST(TestVariantBuilder, Rollback) { +TEST(TestVariantBuilder, ObjectRollback) { VariantBuilder builder; auto object = builder.StartObject(); { - auto child = object.StartObject("drop"); - child.AppendString("nested", "x"); + auto child = object.StartObject("value"); + child.AppendString("drop", "x"); } - object.AppendInt32("ok", 1); + object.AppendInt32("value", 1); object.Finish(); auto encoded = builder.Finish(); - auto view = MakeVariantValueView(encoded); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + ASSERT_FALSE(metadata.FindString("drop").has_value()); + ASSERT_TRUE(metadata.sorted_strings()); + auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); const auto& fields = std::get(view.data()).fields(); ASSERT_EQ(1, fields.size()); - ASSERT_EQ("ok", fields[0].name); + ASSERT_EQ("value", fields[0].name); +} + +TEST(TestVariantBuilder, ListRollback) { + VariantBuilder builder; + auto list = builder.StartList(); + { + auto child = list.StartObject(); + child.AppendString("drop", "x"); + } + list.AppendInt32(1); + list.Finish(); + + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + ASSERT_FALSE(metadata.FindString("drop").has_value()); + auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); + const auto& elements = std::get(view.data()).elements(); + ASSERT_EQ(1, elements.size()); + AssertPrimitiveType(elements[0], metadata, VariantPrimitiveType::kInt32); +} + +TEST(TestVariantBuilder, MoveRollback) { + VariantBuilder replacement_builder; + VariantBuilder builder; + auto object = builder.StartObject(); + object.AppendString("drop", "x"); + + auto replacement = replacement_builder.StartObject(); + object = std::move(replacement); + + auto restored = builder.StartObject(); + restored.AppendInt32("value", 1); + restored.Finish(); + auto encoded = builder.Finish(); + + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + ASSERT_FALSE(metadata.FindString("drop").has_value()); + object.Finish(); + replacement_builder.Finish(); } TEST(TestVariantBuilder, Duplicate) { @@ -222,6 +315,31 @@ TEST(TestVariantBuilder, UsesMemoryPool) { VariantValueView::Validate(std::string_view{*encoded.value}, metadata); } +TEST(TestVariantBuilder, FinishValidates) { + const std::string invalid_value(1, '\x54'); + + VariantBuilder builder1; + builder1.AppendEncodedValue(invalid_value); + ASSERT_THROW(builder1.Finish(), ParquetInvalidOrCorruptedFileException); + + // Invalid case, skip validation + VariantBuilder builder2(default_memory_pool(), /*validate=*/false); + builder2.AppendEncodedValue(invalid_value); + auto encoded = builder2.Finish(); + ASSERT_EQ("\x01\x00\x00"sv, std::string_view{*encoded.metadata}); + ASSERT_EQ(invalid_value, std::string_view{*encoded.value}); +} + +TEST(TestVariantBuilder, EmptyEncodedValue) { + VariantBuilder builder; + ASSERT_THROW(builder.AppendEncodedValue(""), ParquetException); + auto [metadata_owner, metadata] = MakeEmptyMetadata(); + VariantValueArrayBuilder value_builder; + ASSERT_THROW(value_builder.AppendEncodedValue(""), ParquetException); + auto row = value_builder.BindMetadata(metadata); + ASSERT_THROW(row.AppendEncodedValue(""), ParquetException); +} + TEST(TestVariantBuilder, ArrayBuilder) { VariantArrayBuilder builder; builder.AppendNull(); @@ -248,6 +366,34 @@ TEST(TestVariantBuilder, ArrayBuilder) { ASSERT_EQ(Type::BINARY_VIEW, storage.field(1)->type_id()); } +TEST(TestVariantBuilder, Tags) { + VariantArrayBuilder builder; + builder.AppendBoolean(true); + builder.AppendBoolean(false); + builder.AppendTimestampMicros(1, true); + builder.AppendTimestampMicros(2, false); + builder.AppendTimestampNanos(3, true); + builder.AppendTimestampNanos(4, false); + + auto array = builder.Finish(); + const auto& storage = + ::arrow::internal::checked_cast(*array->storage()); + const auto& metadata_values = + ::arrow::internal::checked_cast(*storage.field(0)); + const auto& values = + ::arrow::internal::checked_cast(*storage.field(1)); + const std::array expected = { + VariantPrimitiveType::kBooleanTrue, VariantPrimitiveType::kBooleanFalse, + VariantPrimitiveType::kTimestampMicros, VariantPrimitiveType::kTimestampNTZMicros, + VariantPrimitiveType::kTimestampNanos, VariantPrimitiveType::kTimestampNTZNanos, + }; + ASSERT_EQ(static_cast(expected.size()), array->length()); + for (int64_t row = 0; row < array->length(); ++row) { + auto metadata = VariantMetadataView::Make(metadata_values.GetView(row)); + AssertPrimitiveType(values.GetView(row), metadata, expected[row]); + } +} + TEST(TestVariantBuilder, ArrayBuilderViews) { VariantBuilder value; value.AppendString("x"); @@ -275,6 +421,161 @@ TEST(TestVariantBuilder, ArrayBuilderViews) { ASSERT_FALSE(value_views[1].is_inline()); } +TEST(TestVariantBuilder, ValueArrayBuilder) { + auto [metadata_owner, metadata] = MakeEmptyMetadata(); + + VariantBuilder encoded_builder; + encoded_builder.AppendString("hello"); + auto encoded = encoded_builder.Finish(); + + VariantValueArrayBuilder builder; + builder.AppendNull(); + builder.AppendEncodedValue(std::string_view{*encoded.value}); + auto row = builder.BindMetadata(metadata); + row.AppendInt32(42); + row.Finish(); + + auto array = builder.Finish(); + ASSERT_EQ(3, array->length()); + ASSERT_TRUE(array->IsNull(0)); + AssertPrimitiveType(array->GetView(1), metadata, VariantPrimitiveType::kString); + AssertPrimitiveType(array->GetView(2), metadata, VariantPrimitiveType::kInt32); +} + +TEST(TestVariantBuilder, SharedMetadata) { + constexpr std::string_view kId = "customer_identifier"; + constexpr std::string_view kName = "customer_name"; + + VariantBuilder metadata_builder; + metadata_builder.AddFieldName(kId); + metadata_builder.AddFieldName(kName); + metadata_builder.AppendVariantNull(); + auto metadata_encoded = metadata_builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*metadata_encoded.metadata}); + + VariantValueArrayBuilder value_builder; + { + auto row = value_builder.BindMetadata(metadata); + auto object = row.StartObject(); + object.AppendInt32(kId, 1); + object.Finish(); + row.Finish(); + } + { + auto row = value_builder.BindMetadata(metadata); + auto object = row.StartObject(); + object.AppendString(kName, "Alice"); + object.Finish(); + row.Finish(); + } + auto values = value_builder.Finish(); + + ASSERT_OK_AND_ASSIGN( + auto metadata_values, + MakeArrayFromScalar(BinaryViewScalar(metadata_encoded.metadata), values->length())); + auto storage_type = struct_({field("metadata", binary_view(), /*nullable=*/false), + field("value", binary_view(), /*nullable=*/false)}); + auto array = MakeVariantArrayFromChildren( + std::move(storage_type), {std::move(metadata_values), std::move(values)}); + + ASSERT_EQ(2, array->length()); + const auto& storage = + ::arrow::internal::checked_cast(*array->storage()); + const auto& metadata_array = + ::arrow::internal::checked_cast(*storage.field(0)); + const auto& value_array = + ::arrow::internal::checked_cast(*storage.field(1)); + ASSERT_EQ(std::string_view{*metadata_encoded.metadata}, metadata_array.GetView(0)); + ASSERT_EQ(metadata_array.GetView(0), metadata_array.GetView(1)); + ASSERT_EQ(metadata_encoded.metadata, metadata_array.data()->buffers[2]); + + auto first = VariantValueView::Make(value_array.GetView(0), metadata); + AssertPrimitiveFieldType(std::get(first.data()), kId, + VariantPrimitiveType::kInt32); + + auto second = VariantValueView::Make(value_array.GetView(1), metadata); + AssertPrimitiveFieldType(std::get(second.data()), kName, + VariantPrimitiveType::kString); +} + +TEST(TestVariantBuilder, ValueRowFinish) { + auto [metadata_owner, metadata] = MakeEmptyMetadata(); + + VariantValueArrayBuilder builder; + auto row = builder.BindMetadata(metadata); + row.AppendInt32(42); + row.Finish(); + ASSERT_THROW(row.Finish(), ParquetException); + + auto array = builder.Finish(); + ASSERT_EQ(1, array->length()); + AssertPrimitiveType(array->GetView(0), metadata, VariantPrimitiveType::kInt32); +} + +TEST(TestVariantBuilder, ValueRowMove) { + auto [metadata_owner, metadata] = MakeEmptyMetadata(); + + VariantValueArrayBuilder replacement_builder; + VariantValueArrayBuilder builder; + auto row = builder.BindMetadata(metadata); + row.AppendInt32(1); + auto replacement = replacement_builder.BindMetadata(metadata); + row = std::move(replacement); + + auto restored = builder.BindMetadata(metadata); + restored.AppendString(std::string(32, 'x')); + restored.Finish(); + auto values = builder.Finish(); + ASSERT_EQ(1, values->length()); + ASSERT_EQ(values->GetView(0).size(), values->data()->buffers[2]->size()); + + row.AppendVariantNull(); + row.Finish(); + replacement_builder.Finish(); +} + +TEST(TestVariantBuilder, ValueArrayBuilderObject) { + VariantBuilder metadata_builder; + auto metadata_object = metadata_builder.StartObject(); + metadata_object.AppendInt32("a", 1); + auto metadata_list = metadata_object.StartList("b"); + metadata_list.AppendVariantNull(); + metadata_list.Finish(); + metadata_object.Finish(); + auto metadata_encoded = metadata_builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*metadata_encoded.metadata}); + + VariantValueArrayBuilder builder; + auto row = builder.BindMetadata(metadata); + auto object = row.StartObject(); + object.AppendInt32("a", 42); + auto list = object.StartList("b"); + list.AppendString("x"); + list.Finish(); + object.Finish(); + row.Finish(); + + auto array = builder.Finish(); + ASSERT_EQ(1, array->length()); + auto view = VariantValueView::Make(array->GetView(0), metadata); + ASSERT_EQ(VariantBasicType::kObject, view.basic_type()); + const auto& object_view = std::get(view.data()); + ASSERT_EQ(2, object_view.fields().size()); + AssertPrimitiveFieldType(object_view, "a", VariantPrimitiveType::kInt32); + auto list_view = object_view.GetField("b"); + ASSERT_TRUE(list_view.has_value()); + ASSERT_EQ(VariantBasicType::kArray, list_view->basic_type()); +} + +TEST(TestVariantBuilder, ValueArrayBuilderMissingField) { + auto [metadata_owner, metadata] = MakeEmptyMetadata(); + + VariantValueArrayBuilder builder; + auto row = builder.BindMetadata(metadata); + auto object = row.StartObject(); + ASSERT_THROW(object.AppendInt32("missing", 1), ParquetInvalidOrCorruptedFileException); +} + TEST(TestVariantBuilder, FromStorage) { VariantBuilder value; value.AppendInt32(1); diff --git a/cpp/src/parquet/variant/encoding.cc b/cpp/src/parquet/variant/decoding.cc similarity index 70% rename from cpp/src/parquet/variant/encoding.cc rename to cpp/src/parquet/variant/decoding.cc index 1c38951cd97a..e6e00eb69c88 100644 --- a/cpp/src/parquet/variant/encoding.cc +++ b/cpp/src/parquet/variant/decoding.cc @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -#include "parquet/variant/encoding.h" +#include "parquet/variant/decoding.h" #include #include @@ -24,14 +24,28 @@ #include "arrow/util/endian.h" #include "arrow/util/logging_internal.h" #include "parquet/exception.h" -#include "parquet/variant/encoding_internal.h" +#include "parquet/variant/format_internal.h" namespace parquet::variant { namespace bit_util = ::arrow::bit_util; +namespace internal { + +class ViewAccess { + public: + template + static View Make(Args&&... args) { + return View(std::forward(args)...); + } +}; + +} // namespace internal + namespace { +using internal::ViewAccess; + uint32_t ReadLittleEndian(std::string_view data, size_t offset, size_t width) { DCHECK_LE(width, sizeof(uint32_t)); uint32_t value = 0; @@ -112,9 +126,11 @@ size_t ParsePrimitive(std::string_view value, size_t offset, return payload_size; } +template size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, VariantValueView* out); +template size_t ParseArray(std::string_view value, const VariantMetadataView& metadata, uint8_t header, VariantValueView* out) { const auto offset_size = static_cast((header & 0x03) + 1); @@ -150,41 +166,33 @@ size_t ParseArray(std::string_view value, const VariantMetadataView& metadata, const size_t total_value_size = offsets[num_elements]; CheckAvailable(value, values_start, total_value_size, "array values"); - size_t current = 0; - std::vector array_elements; - if (out != nullptr) { - array_elements.reserve(num_elements); - } - for (uint32_t i = 0; i < num_elements; ++i) { - if (offsets[i] != current) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant encoding: array offset does not match value boundary"); - } - const size_t child_consumed = - ParseValue(value.substr(values_start + current), metadata, /*out=*/nullptr); - current += child_consumed; - if (current != offsets[i + 1]) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant encoding: array value does not end at next offset"); - } - if (out != nullptr) { - array_elements.push_back(value.substr(values_start + offsets[i], child_consumed)); + if constexpr (validate_children) { + for (uint32_t i = 0; i < num_elements; ++i) { + const size_t child_consumed = ParseValue( + value.substr(values_start + offsets[i]), metadata, /*out=*/nullptr); + if (child_consumed != offsets[i + 1] - offsets[i]) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: array value does not end at next offset"); + } } } - if (current != total_value_size) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant encoding: array values have trailing data"); - } - const size_t consumed = values_start + total_value_size; if (out != nullptr) { - *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kArray, - VariantArrayView(std::move(array_elements))); + std::vector array_elements; + array_elements.reserve(num_elements); + for (uint32_t i = 0; i < num_elements; ++i) { + array_elements.push_back( + value.substr(values_start + offsets[i], offsets[i + 1] - offsets[i])); + } + *out = ViewAccess::Make( + value.substr(0, consumed), + ViewAccess::Make(std::move(array_elements), metadata)); } return consumed; } +template size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, uint8_t header, VariantValueView* out) { const auto offset_size = static_cast((header & 0x03) + 1); @@ -226,19 +234,18 @@ size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, object_fields.reserve(num_elements); } - std::string_view previous_name; for (uint32_t i = 0; i < num_elements; ++i) { if (field_ids[i] >= metadata.dictionary_size()) { throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: object field id ", field_ids[i], " is outside metadata dictionary of size ", metadata.dictionary_size()); } - const auto name = metadata.string(field_ids[i]); - if (i > 0 && !(previous_name < name)) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant encoding: object field names must be sorted and unique"); + if constexpr (validate_children) { + if (i > 0 && !(metadata.string(field_ids[i - 1]) < metadata.string(field_ids[i]))) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object field names must be sorted and unique"); + } } - previous_name = name; const auto field_offset = field_offsets[i]; if (field_offset >= total_value_size) { @@ -258,14 +265,16 @@ size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, "Invalid Variant encoding: object values have leading data"); } - for (uint32_t i = 0; i < num_elements; ++i) { - const uint32_t start = value_offsets[i]; - const uint32_t end = value_offsets[i + 1]; - const size_t child_consumed = - ParseValue(value.substr(values_start + start), metadata, /*out=*/nullptr); - if (child_consumed != end - start) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant encoding: object value does not end at next value boundary"); + if constexpr (validate_children) { + for (uint32_t i = 0; i < num_elements; ++i) { + const uint32_t start = value_offsets[i]; + const uint32_t end = value_offsets[i + 1]; + const size_t child_consumed = + ParseValue(value.substr(values_start + start), metadata, /*out=*/nullptr); + if (child_consumed != end - start) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object value does not end at next value boundary"); + } } } @@ -286,19 +295,18 @@ size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, const size_t consumed = values_start + total_value_size; if (out != nullptr) { - *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kObject, - VariantObjectView(std::move(object_fields))); + *out = ViewAccess::Make( + value.substr(0, consumed), + ViewAccess::Make(std::move(object_fields), metadata)); } return consumed; } +template size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, VariantValueView* out) { - CheckAvailable(value, 0, 1, "value header"); - - const auto metadata_byte = static_cast(value[0]); - const auto basic_type = static_cast(metadata_byte & 0x03); - const auto header = static_cast(metadata_byte >> 2); + const auto basic_type = VariantValueView::PeekBasicType(value); + const auto header = static_cast(static_cast(value[0]) >> 2); switch (basic_type) { case VariantBasicType::kPrimitive: { @@ -306,9 +314,9 @@ size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, const size_t payload_size = ParsePrimitive(value, 1, primitive); const size_t consumed = 1 + payload_size; if (out != nullptr) { - *out = VariantValueView( - value.substr(0, consumed), VariantBasicType::kPrimitive, - VariantPrimitiveView(primitive, value.substr(1, payload_size))); + *out = ViewAccess::Make( + value.substr(0, consumed), ViewAccess::Make( + primitive, value.substr(1, payload_size))); } return consumed; } @@ -317,25 +325,41 @@ size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, internal::ValidateUtf8(value.substr(1, header), "short string value"); const size_t consumed = 1 + header; if (out != nullptr) { - *out = VariantValueView(value.substr(0, consumed), VariantBasicType::kShortString, - VariantShortStringView(value.substr(1, header))); + *out = ViewAccess::Make( + value.substr(0, consumed), + ViewAccess::Make(value.substr(1, header))); } return consumed; } case VariantBasicType::kObject: - return ParseObject(value, metadata, header, out); + return ParseObject(value, metadata, header, out); case VariantBasicType::kArray: - return ParseArray(value, metadata, header, out); + return ParseArray(value, metadata, header, out); } throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: unknown basic type"); } +template +VariantValueView ParseCompleteValue(std::string_view value, + const VariantMetadataView& metadata) { + auto view = ViewAccess::Make( + std::string_view{}, ViewAccess::Make( + VariantPrimitiveType::kNull, std::string_view{})); + const size_t consumed = ParseValue(value, metadata, &view); + if (consumed != value.size()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: trailing bytes after value"); + } + return view; +} + } // namespace -VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { - CheckAvailable(metadata, 0, 1, "metadata header"); - const auto header = static_cast(metadata[0]); +VariantMetadataView VariantMetadataView::ParsePrefix(std::string_view data, + size_t* consumed) { + CheckAvailable(data, 0, 1, "metadata header"); + const auto header = static_cast(data[0]); const auto version = static_cast(header & internal::kMetadataVersionMask); if (version != internal::kVariantVersion) { throw ParquetInvalidOrCorruptedFileException( @@ -343,21 +367,20 @@ VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { } VariantMetadataView view; - view.metadata_ = metadata; view.sorted_strings_ = (header & internal::kMetadataSortedStringsMask) != 0; view.offset_size_ = static_cast(((header >> 6) & 0x03) + 1); - CheckAvailable(metadata, 1, view.offset_size_, "metadata dictionary size"); - const uint32_t dictionary_size = ReadLittleEndian(metadata, 1, view.offset_size_); + CheckAvailable(data, 1, view.offset_size_, "metadata dictionary size"); + const uint32_t dictionary_size = ReadLittleEndian(data, 1, view.offset_size_); const size_t offsets_offset = 1 + view.offset_size_; - CheckAvailable(metadata, offsets_offset, + CheckAvailable(data, offsets_offset, (static_cast(dictionary_size) + 1) * view.offset_size_, "metadata dictionary offsets"); std::vector offsets(dictionary_size + 1); for (uint32_t i = 0; i <= dictionary_size; ++i) { - offsets[i] = ReadLittleEndian(metadata, offsets_offset + i * view.offset_size_, - view.offset_size_); + offsets[i] = + ReadLittleEndian(data, offsets_offset + i * view.offset_size_, view.offset_size_); } if (offsets[0] != 0) { @@ -374,15 +397,17 @@ VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { const size_t bytes_offset = offsets_offset + (static_cast(dictionary_size) + 1) * view.offset_size_; const size_t bytes_size = offsets[dictionary_size]; - CheckAvailable(metadata, bytes_offset, bytes_size, "metadata dictionary bytes"); - if (metadata.size() != bytes_offset + bytes_size) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant metadata: trailing bytes after dictionary"); + CheckAvailable(data, bytes_offset, bytes_size, "metadata dictionary bytes"); + const size_t metadata_size = bytes_offset + bytes_size; + if (consumed != nullptr) { + *consumed = metadata_size; } + view.metadata_ = data.substr(0, metadata_size); view.strings_.reserve(dictionary_size); for (uint32_t i = 0; i < dictionary_size; ++i) { - auto string = metadata.substr(bytes_offset + offsets[i], offsets[i + 1] - offsets[i]); + auto string = + view.metadata_.substr(bytes_offset + offsets[i], offsets[i + 1] - offsets[i]); internal::ValidateUtf8(string, "metadata dictionary string"); if (view.sorted_strings_ && i > 0 && !(view.strings_.back() < string)) { throw ParquetInvalidOrCorruptedFileException( @@ -395,6 +420,16 @@ VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { return view; } +VariantMetadataView VariantMetadataView::Make(std::string_view metadata) { + size_t consumed; + auto view = ParsePrefix(metadata, &consumed); + if (consumed != metadata.size()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant metadata: trailing bytes after dictionary"); + } + return view; +} + std::string_view VariantMetadataView::string(uint32_t field_id) const { DCHECK_LT(field_id, strings_.size()); return strings_[field_id]; @@ -417,32 +452,52 @@ std::optional VariantMetadataView::FindString(std::string_view value) return std::nullopt; } -const VariantObjectField* VariantObjectView::FindField(std::string_view name) const { - const auto it = std::ranges::lower_bound(fields_, name, {}, &VariantObjectField::name); - return (it == fields_.end() || it->name != name) ? nullptr : &*it; -} - bool VariantObjectView::ContainsField(std::string_view name) const { // The Parquet Variant encoding requires object fields to be sorted by name, so field // lookup can use binary search. return std::ranges::binary_search(fields_, name, {}, &VariantObjectField::name); } +std::optional VariantObjectView::GetField(std::string_view name) const { + const auto it = std::ranges::lower_bound(fields_, name, {}, &VariantObjectField::name); + if (it == fields_.end() || it->name != name) { + return std::nullopt; + } + return VariantValueView::MakeShallow(it->value, metadata_.get()); +} + +std::optional VariantObjectView::GetField(size_t index) const { + if (index >= fields_.size()) { + return std::nullopt; + } + return VariantValueView::MakeShallow(fields_[index].value, metadata_.get()); +} + +std::optional VariantArrayView::GetElement(size_t index) const { + if (index >= elements_.size()) { + return std::nullopt; + } + return VariantValueView::MakeShallow(elements_[index], metadata_.get()); +} + VariantValueView VariantValueView::Make(std::string_view value, const VariantMetadataView& metadata) { - VariantValueView view({}, VariantBasicType::kPrimitive, - VariantPrimitiveView(VariantPrimitiveType::kNull, {})); - const size_t consumed = ParseValue(value, metadata, &view); - if (consumed != value.size()) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant encoding: trailing bytes after value"); - } - return view; + return ParseCompleteValue(value, metadata); +} + +VariantValueView VariantValueView::MakeShallow(std::string_view value, + const VariantMetadataView& metadata) { + return ParseCompleteValue(value, metadata); +} + +VariantBasicType VariantValueView::PeekBasicType(std::string_view value) { + CheckAvailable(value, 0, 1, "value header"); + return static_cast(static_cast(value[0]) & 0x03); } void VariantValueView::Validate(std::string_view value, const VariantMetadataView& metadata) { - const size_t consumed = ParseValue(value, metadata, /*out=*/nullptr); + const size_t consumed = ParseValue(value, metadata, /*out=*/nullptr); if (consumed != value.size()) { throw ParquetInvalidOrCorruptedFileException( "Invalid Variant encoding: trailing bytes after value"); diff --git a/cpp/src/parquet/variant/encoding.h b/cpp/src/parquet/variant/decoding.h similarity index 59% rename from cpp/src/parquet/variant/encoding.h rename to cpp/src/parquet/variant/decoding.h index ca7e93b034c8..6ca3f5d21bfc 100644 --- a/cpp/src/parquet/variant/encoding.h +++ b/cpp/src/parquet/variant/decoding.h @@ -17,7 +17,7 @@ #pragma once -#include +#include #include #include #include @@ -25,39 +25,15 @@ #include #include "parquet/platform.h" +#include "parquet/variant/format.h" namespace parquet::variant { -enum class VariantBasicType : uint8_t { - kPrimitive = 0, - kShortString = 1, - kObject = 2, - kArray = 3, -}; +namespace internal { -enum class VariantPrimitiveType : uint8_t { - kNull = 0, - kBooleanTrue = 1, - kBooleanFalse = 2, - kInt8 = 3, - kInt16 = 4, - kInt32 = 5, - kInt64 = 6, - kDouble = 7, - kDecimal4 = 8, - kDecimal8 = 9, - kDecimal16 = 10, - kDate = 11, - kTimestampMicros = 12, - kTimestampNTZMicros = 13, - kFloat = 14, - kBinary = 15, - kString = 16, - kTimeNTZMicros = 17, - kTimestampNanos = 18, - kTimestampNTZNanos = 19, - kUuid = 20, -}; +class ViewAccess; + +} struct PARQUET_EXPORT VariantObjectField { std::string_view name; @@ -73,6 +49,13 @@ class PARQUET_EXPORT VariantMetadataView { /// input metadata bytes remain alive and unchanged. static VariantMetadataView Make(std::string_view metadata); + /// Parse a Variant metadata prefix without copying it and return bytes consumed. + /// The returned view contains exactly the consumed metadata bytes. + /// + /// \param[out] consumed receives the metadata size when non-null + static VariantMetadataView ParsePrefix(std::string_view data, + size_t* consumed = nullptr); + std::string_view metadata() const { return metadata_; } bool sorted_strings() const { return sorted_strings_; } uint8_t offset_size() const { return offset_size_; } @@ -93,49 +76,76 @@ class PARQUET_EXPORT VariantMetadataView { class PARQUET_EXPORT VariantPrimitiveView { public: - VariantPrimitiveView(VariantPrimitiveType type, std::string_view payload) - : type_(type), payload_(payload) {} + static constexpr VariantBasicType kBasicType = VariantBasicType::kPrimitive; VariantPrimitiveType type() const { return type_; } std::string_view payload() const { return payload_; } private: + friend class internal::ViewAccess; + + VariantPrimitiveView(VariantPrimitiveType type, std::string_view payload) + : type_(type), payload_(payload) {} + VariantPrimitiveType type_ = VariantPrimitiveType::kNull; std::string_view payload_; }; class PARQUET_EXPORT VariantShortStringView { public: - explicit VariantShortStringView(std::string_view string) : string_(string) {} + static constexpr VariantBasicType kBasicType = VariantBasicType::kShortString; std::string_view string() const { return string_; } private: + friend class internal::ViewAccess; + + explicit VariantShortStringView(std::string_view string) : string_(string) {} + std::string_view string_; }; +class VariantValueView; + class PARQUET_EXPORT VariantObjectView { public: - explicit VariantObjectView(std::vector fields) - : fields_(std::move(fields)) {} + static constexpr VariantBasicType kBasicType = VariantBasicType::kObject; const std::vector& fields() const { return fields_; } - const VariantObjectField* FindField(std::string_view name) const; bool ContainsField(std::string_view name) const; + /// Return a field by name, or nullopt if it does not exist. + std::optional GetField(std::string_view name) const; + /// Return a field by index, or nullopt if the index is out of bounds. + std::optional GetField(size_t index) const; private: + friend class internal::ViewAccess; + + VariantObjectView(std::vector fields, + const VariantMetadataView& metadata) + : fields_(std::move(fields)), metadata_(metadata) {} + std::vector fields_; + std::reference_wrapper metadata_; }; class PARQUET_EXPORT VariantArrayView { public: - explicit VariantArrayView(std::vector elements) - : elements_(std::move(elements)) {} + static constexpr VariantBasicType kBasicType = VariantBasicType::kArray; const std::vector& elements() const { return elements_; } + /// Return an element by index, or nullopt if the index is out of bounds. + std::optional GetElement(size_t index) const; private: + friend class internal::ViewAccess; + + VariantArrayView(std::vector elements, + const VariantMetadataView& metadata) + : elements_(std::move(elements)), metadata_(metadata) {} + std::vector elements_; + std::reference_wrapper metadata_; }; class PARQUET_EXPORT VariantValueView { @@ -143,25 +153,39 @@ class PARQUET_EXPORT VariantValueView { using Data = std::variant; - VariantValueView(std::string_view value, VariantBasicType basic_type, Data data) - : value_(value), basic_type_(basic_type), data_(std::move(data)) {} - /// Parse Variant value bytes without copying them. /// /// The returned view and any nested object/list/string views are valid only while the - /// input value bytes remain alive and unchanged. The metadata view must also remain - /// valid while object field names are accessed. + /// input value bytes remain alive and unchanged. The metadata view object and its + /// underlying bytes must also remain alive and unchanged. static VariantValueView Make(std::string_view value, const VariantMetadataView& metadata); + static VariantValueView Make(std::string_view value, + VariantMetadataView&& metadata) = delete; + + /// Read the basic type from the value header without validating the remaining bytes. + static VariantBasicType PeekBasicType(std::string_view value); + static void Validate(std::string_view value, const VariantMetadataView& metadata); std::string_view value() const { return value_; } - VariantBasicType basic_type() const { return basic_type_; } + VariantBasicType basic_type() const { + return std::visit([](const auto& view) { return view.kBasicType; }, data_); + } const Data& data() const { return data_; } private: + friend class internal::ViewAccess; + friend class VariantArrayView; + friend class VariantObjectView; + + VariantValueView(std::string_view value, Data data) + : value_(value), data_(std::move(data)) {} + + static VariantValueView MakeShallow(std::string_view value, + const VariantMetadataView& metadata); + std::string_view value_; - VariantBasicType basic_type_; Data data_; }; diff --git a/cpp/src/parquet/variant/encoding_test.cc b/cpp/src/parquet/variant/decoding_test.cc similarity index 60% rename from cpp/src/parquet/variant/encoding_test.cc rename to cpp/src/parquet/variant/decoding_test.cc index e9b1bdd8a95e..25bb6e71bd4c 100644 --- a/cpp/src/parquet/variant/encoding_test.cc +++ b/cpp/src/parquet/variant/decoding_test.cc @@ -15,39 +15,62 @@ // specific language governing permissions and limitations // under the License. -#include "parquet/variant/encoding.h" -#include "parquet/exception.h" +#include "parquet/variant/decoding.h" #include "parquet/variant/builder.h" -#include "parquet/variant/test_util_internal.h" #include #include #include #include +#include + +#include "arrow/util/decimal.h" +#include "parquet/exception.h" + +using namespace std::string_view_literals; // NOLINT namespace parquet::variant { +template +concept CanMakeVariantValueView = (requires(std::string_view value, Metadata&& metadata) { + VariantValueView::Make(value, std::forward(metadata)); +}); + +static_assert(CanMakeVariantValueView); +static_assert(!CanMakeVariantValueView); + TEST(TestVariantEncoding, EmptyMetadata) { VariantBuilder builder; builder.AppendVariantNull(); auto encoded = builder.Finish(); - ASSERT_EQ(std::string_view("\x01\x00\x00", 3), std::string_view{*encoded.metadata}); + ASSERT_EQ("\x01\x00\x00"sv, std::string_view{*encoded.metadata}); auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); ASSERT_FALSE(metadata.sorted_strings()); ASSERT_EQ(0, metadata.dictionary_size()); } +TEST(TestVariantEncoding, MetadataPrefix) { + constexpr auto encoded = "\x01\x00\x00\x00"sv; + size_t metadata_size; + auto metadata = VariantMetadataView::ParsePrefix(encoded, &metadata_size); + ASSERT_EQ(3, metadata_size); + ASSERT_EQ("\x01\x00\x00"sv, metadata.metadata()); + ASSERT_THROW(VariantMetadataView::Make(encoded), + ParquetInvalidOrCorruptedFileException); + ASSERT_NO_THROW(VariantValueView::Validate(encoded.substr(metadata_size), metadata)); +} + TEST(TestVariantEncoding, Primitive) { - std::array decimal16 = {0}; std::array uuid = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; auto check = [](auto append, VariantPrimitiveType expected) { VariantBuilder builder; append(builder); auto encoded = builder.Finish(); - auto view = internal::MakeVariantValueView(encoded); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); ASSERT_EQ(VariantBasicType::kPrimitive, view.basic_type()); ASSERT_EQ(expected, std::get(view.data()).type()); }; @@ -61,15 +84,12 @@ TEST(TestVariantEncoding, Primitive) { check([](VariantBuilder& b) { b.AppendInt64(4); }, VariantPrimitiveType::kInt64); check([](VariantBuilder& b) { b.AppendFloat(1.5F); }, VariantPrimitiveType::kFloat); check([](VariantBuilder& b) { b.AppendDouble(2.5); }, VariantPrimitiveType::kDouble); - check([](VariantBuilder& b) { b.AppendDecimal4(123, 2); }, + check([](VariantBuilder& b) { b.AppendDecimal4(::arrow::Decimal32(123), 2); }, VariantPrimitiveType::kDecimal4); - check([](VariantBuilder& b) { b.AppendDecimal8(123, 2); }, + check([](VariantBuilder& b) { b.AppendDecimal8(::arrow::Decimal64(123), 2); }, VariantPrimitiveType::kDecimal8); - check( - [&](VariantBuilder& b) { - b.AppendDecimal16(std::string_view(decimal16.data(), decimal16.size()), 2); - }, - VariantPrimitiveType::kDecimal16); + check([&](VariantBuilder& b) { b.AppendDecimal16(::arrow::Decimal128(1, 0), 2); }, + VariantPrimitiveType::kDecimal16); check([](VariantBuilder& b) { b.AppendBinary("abc"); }, VariantPrimitiveType::kBinary); check([](VariantBuilder& b) { b.AppendString("abc"); }, VariantPrimitiveType::kString); check([](VariantBuilder& b) { b.AppendDate(1); }, VariantPrimitiveType::kDate); @@ -94,20 +114,21 @@ TEST(TestVariantEncoding, ShortString) { VariantBuilder builder; builder.AppendShortString("abc"); auto encoded = builder.Finish(); - auto view = internal::MakeVariantValueView(encoded); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto view = VariantValueView::Make(std::string_view{*encoded.value}, metadata); ASSERT_EQ(VariantBasicType::kShortString, view.basic_type()); ASSERT_EQ("abc", std::get(view.data()).string()); } TEST(TestVariantEncoding, InvalidMetadata) { // Metadata version is 2 instead of 1. - ASSERT_THROW(VariantMetadataView::Make(std::string("\x02\x00\x00", 3)), + ASSERT_THROW(VariantMetadataView::Make("\x02\x00\x00"sv), ParquetInvalidOrCorruptedFileException); // Dictionary offsets are truncated. - ASSERT_THROW(VariantMetadataView::Make(std::string("\x01\x01\x00", 3)), + ASSERT_THROW(VariantMetadataView::Make("\x01\x01\x00"sv), ParquetInvalidOrCorruptedFileException); // Dictionary string payload is not valid UTF-8. - ASSERT_THROW(VariantMetadataView::Make(std::string("\x01\x01\x00\x01\xff", 5)), + ASSERT_THROW(VariantMetadataView::Make("\x01\x01\x00\x01\xff"sv), ParquetInvalidOrCorruptedFileException); } @@ -116,19 +137,28 @@ TEST(TestVariantEncoding, ReservedBits) { // - 0x21: version 1 with reserved bit 0x20 set // - 0x00: zero dictionary entries, encoded with one-byte width // - 0x00: first and final dictionary byte offset - auto metadata = VariantMetadataView::Make(std::string("\x21\x00\x00", 3)); + auto metadata = VariantMetadataView::Make("\x21\x00\x00"sv); ASSERT_EQ(0, metadata.dictionary_size()); // Object value bytes: // - 0x82: basic type Object with reserved header bit 0x20 set // - 0x00: zero fields // - 0x00: first and final value offset - VariantValueView::Validate(std::string("\x82\x00\x00", 3), metadata); + VariantValueView::Validate("\x82\x00\x00"sv, metadata); // Array value bytes: // - 0xe3: basic type Array with reserved header bits 0x38 set // - 0x00: zero elements // - 0x00: first and final value offset - VariantValueView::Validate(std::string("\xe3\x00\x00", 3), metadata); + VariantValueView::Validate("\xe3\x00\x00"sv, metadata); +} + +TEST(TestVariantEncoding, PeekBasicType) { + ASSERT_EQ(VariantBasicType::kPrimitive, VariantValueView::PeekBasicType("\x00"sv)); + ASSERT_EQ(VariantBasicType::kShortString, VariantValueView::PeekBasicType("\x01"sv)); + ASSERT_EQ(VariantBasicType::kObject, VariantValueView::PeekBasicType("\x02"sv)); + ASSERT_EQ(VariantBasicType::kArray, VariantValueView::PeekBasicType("\x03"sv)); + ASSERT_THROW(VariantValueView::PeekBasicType({}), + ParquetInvalidOrCorruptedFileException); } TEST(TestVariantEncoding, ValidateValue) { @@ -144,6 +174,35 @@ TEST(TestVariantEncoding, ValidateValue) { ParquetInvalidOrCorruptedFileException); } +TEST(TestVariantEncoding, ChildAccess) { + VariantBuilder builder; + auto object = builder.StartObject(); + auto list = object.StartList("items"); + list.AppendInt32(42); + list.Finish(); + object.Finish(); + + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + auto root = VariantValueView::Make(std::string_view{*encoded.value}, metadata); + + const auto& object_view = std::get(root.data()); + auto items_by_name = object_view.GetField("items"); + auto items_by_index = object_view.GetField(0); + ASSERT_TRUE(items_by_name.has_value()); + ASSERT_TRUE(items_by_index.has_value()); + ASSERT_EQ(items_by_name->value(), items_by_index->value()); + ASSERT_FALSE(object_view.GetField("missing").has_value()); + ASSERT_FALSE(object_view.GetField(1).has_value()); + + const auto& array = std::get(items_by_name->data()); + auto element = array.GetElement(0); + ASSERT_TRUE(element.has_value()); + ASSERT_EQ(VariantPrimitiveType::kInt32, + std::get(element->data()).type()); + ASSERT_FALSE(array.GetElement(1).has_value()); +} + TEST(TestVariantEncoding, InvalidValue) { VariantBuilder builder; builder.AppendVariantNull(); @@ -151,18 +210,21 @@ TEST(TestVariantEncoding, InvalidValue) { auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); // Header byte 0x54 encodes Primitive with primitive tag 21, which is unknown. - ASSERT_THROW(VariantValueView::Make(std::string(1, static_cast(0x54)), metadata), + ASSERT_THROW(VariantValueView::Make("\x54"sv, metadata), ParquetInvalidOrCorruptedFileException); // Short string payload is not valid UTF-8. - ASSERT_THROW(VariantValueView::Make(std::string("\x05\xff", 2), metadata), + ASSERT_THROW(VariantValueView::Make("\x05\xff"sv, metadata), ParquetInvalidOrCorruptedFileException); // Null primitive has trailing bytes after the encoded value. - ASSERT_THROW(VariantValueView::Make(std::string("\x00\x00", 2), metadata), + ASSERT_THROW(VariantValueView::Make("\x00\x00"sv, metadata), ParquetInvalidOrCorruptedFileException); // Object references field id 0, but the metadata dictionary is empty. - ASSERT_THROW( - VariantValueView::Make(std::string("\x02\x01\x00\x00\x01\x00", 6), metadata), - ParquetInvalidOrCorruptedFileException); + ASSERT_THROW(VariantValueView::Make("\x02\x01\x00\x00\x01\x00"sv, metadata), + ParquetInvalidOrCorruptedFileException); + // The outer and inner array offsets are valid, but the grandchild primitive tag is + // unknown. Make must recursively validate the entire value before returning. + ASSERT_THROW(VariantValueView::Make("\x03\x01\x00\x05\x03\x01\x00\x01\x54"sv, metadata), + ParquetInvalidOrCorruptedFileException); } TEST(TestVariantEncoding, InvalidObject) { @@ -180,7 +242,7 @@ TEST(TestVariantEncoding, InvalidObject) { // - 0x00 0x01: field ids for "a" and "b" // - 0x00 0x00 0x01: duplicate field start offsets 0/0, final value size 1 // - 0x00: encoded Variant null primitive value - const std::string duplicate_offsets("\x02\x02\x00\x01\x00\x00\x01\x00", 8); + const auto duplicate_offsets = "\x02\x02\x00\x01\x00\x00\x01\x00"sv; ASSERT_THROW(VariantValueView::Validate(duplicate_offsets, metadata), ParquetInvalidOrCorruptedFileException); } diff --git a/cpp/src/parquet/variant/format.h b/cpp/src/parquet/variant/format.h new file mode 100644 index 000000000000..67a910a21838 --- /dev/null +++ b/cpp/src/parquet/variant/format.h @@ -0,0 +1,55 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include + +namespace parquet::variant { + +enum class VariantBasicType : uint8_t { + kPrimitive = 0, + kShortString = 1, + kObject = 2, + kArray = 3, +}; + +enum class VariantPrimitiveType : uint8_t { + kNull = 0, + kBooleanTrue = 1, + kBooleanFalse = 2, + kInt8 = 3, + kInt16 = 4, + kInt32 = 5, + kInt64 = 6, + kDouble = 7, + kDecimal4 = 8, + kDecimal8 = 9, + kDecimal16 = 10, + kDate = 11, + kTimestampMicros = 12, + kTimestampNTZMicros = 13, + kFloat = 14, + kBinary = 15, + kString = 16, + kTimeNTZMicros = 17, + kTimestampNanos = 18, + kTimestampNTZNanos = 19, + kUuid = 20, +}; + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/encoding_internal.h b/cpp/src/parquet/variant/format_internal.h similarity index 87% rename from cpp/src/parquet/variant/encoding_internal.h rename to cpp/src/parquet/variant/format_internal.h index 3c647e21d5e5..ea74fda27117 100644 --- a/cpp/src/parquet/variant/encoding_internal.h +++ b/cpp/src/parquet/variant/format_internal.h @@ -20,13 +20,19 @@ #include #include +#include "arrow/type_fwd.h" #include "arrow/util/utf8.h" #include "parquet/exception.h" -#include "parquet/variant/encoding.h" +#include "parquet/variant/format.h" -namespace parquet::variant::internal { +namespace parquet::variant { -namespace util = ::arrow::util; +class VariantPrimitiveView; +class VariantShortStringView; +class VariantObjectView; +class VariantArrayView; + +namespace internal { inline constexpr uint8_t kVariantVersion = 1; inline constexpr uint8_t kMetadataVersionMask = 0x0F; @@ -74,8 +80,9 @@ inline void ValidateDecimalScale(uint8_t scale) { } inline void ValidateUtf8(std::string_view value, std::string_view context) { - util::InitializeUTF8(); - if (!util::ValidateUTF8(reinterpret_cast(value.data()), value.size())) { + ::arrow::util::InitializeUTF8(); + if (!::arrow::util::ValidateUTF8(reinterpret_cast(value.data()), + value.size())) { throw ParquetInvalidOrCorruptedFileException("Invalid Variant encoding: ", context, " is not valid UTF-8"); } @@ -123,8 +130,9 @@ struct VariantDecimalPrimitiveTraits {}; using CType = C_TYPE; \ }; -VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal4, int32_t) -VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal8, int64_t) +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal4, ::arrow::Decimal32) +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal8, ::arrow::Decimal64) +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal16, ::arrow::Decimal128) #undef VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF @@ -136,10 +144,9 @@ template concept LengthPrefixedVariantPrimitive = type == VariantPrimitiveType::kBinary || type == VariantPrimitiveType::kString; -template -concept Decimal16VariantPrimitive = type == VariantPrimitiveType::kDecimal16; - template concept UuidVariantPrimitive = type == VariantPrimitiveType::kUuid; -} // namespace parquet::variant::internal +} // namespace internal + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/meson.build b/cpp/src/parquet/variant/meson.build index 7a3f74014ad5..89fdc8858353 100644 --- a/cpp/src/parquet/variant/meson.build +++ b/cpp/src/parquet/variant/meson.build @@ -19,9 +19,11 @@ exc = executable( 'parquet-variant-test', sources: [ 'builder_test.cc', - 'encoding_test.cc', + 'decoding_test.cc', + 'parquet_testing_test.cc', 'test_util_internal.cc', 'type_test.cc', + 'unshred_test.cc', 'validate_test.cc', ], dependencies: parquet_test_dep, @@ -29,6 +31,6 @@ exc = executable( test('parquet-variant-test', exc) install_headers( - ['builder.h', 'encoding.h', 'validate.h'], + ['builder.h', 'decoding.h', 'format.h', 'unshred.h', 'validate.h'], subdir: 'parquet/variant', ) diff --git a/cpp/src/parquet/variant/parquet_testing_test.cc b/cpp/src/parquet/variant/parquet_testing_test.cc new file mode 100644 index 000000000000..e2f1ee534299 --- /dev/null +++ b/cpp/src/parquet/variant/parquet_testing_test.cc @@ -0,0 +1,474 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/buffer.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension_type.h" +#include "arrow/io/file.h" +#include "arrow/io/memory.h" +#include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep +#include "arrow/table.h" +#include "arrow/testing/extension_type.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/util/io_util.h" +#include "parquet/arrow/reader.h" +#include "parquet/arrow/writer.h" +#include "parquet/exception.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/decoding.h" +#include "parquet/variant/unshred.h" +#include "parquet/variant/validate.h" + +#include +#include + +namespace parquet::variant { + +using ::arrow::internal::checked_pointer_cast; + +namespace { + +namespace rj = ::arrow::rapidjson; + +using ExpectedCases = std::map; + +constexpr int kMissingParquetFileCaseNumber = 3; + +const ExpectedCases kSkipUnshreddedCases = { + {43, "testPartiallyShreddedObjectMissingFieldConflict"}, + {125, "testPartiallyShreddedObjectFieldConflict"}, +}; + +const ExpectedCases kReadCompatOnlyCases = { + {41, "testArrayMissingValueColumn"}, + {43, "testPartiallyShreddedObjectMissingFieldConflict"}, + {84, "testShreddedObjectWithOptionalFieldStructs"}, + {125, "testPartiallyShreddedObjectFieldConflict"}, + {131, "testMissingValueColumn"}, + {132, "testShreddedObjectMissingFieldValueColumn"}, + {138, "testShreddedObjectMissingValueColumn"}, +}; + +const ExpectedCases kReadCompatErrorCases = { + {40, "testArrayWithElementValueTypedValueConflict"}, + {42, "testValueAndTypedValueConflict"}, + {87, "testNonObjectWithNonNullShreddedFields"}, + {128, "testEmptyPartiallyShreddedObjectConflict"}, +}; + +const ExpectedCases kUnsupportedTypedValueCases = { + {127, "testUnsignedInteger"}, + {137, "testFixedLengthByteArray"}, +}; + +std::optional ParquetTestingSiblingDir(std::string_view name) { + auto data = ::arrow::internal::GetEnvVar("PARQUET_TEST_DATA"); + if (!data.ok() || data->empty()) { + return std::nullopt; + } + auto dir = *data + "/../" + std::string(name); + PARQUET_ASSIGN_OR_THROW(auto dir_path, + ::arrow::internal::PlatformFilename::FromString(dir)); + PARQUET_ASSIGN_OR_THROW(auto exists, ::arrow::internal::FileExists(dir_path)); + if (!exists) { + return std::nullopt; + } + return dir; +} + +std::shared_ptr ReadFile(const std::string& path) { + PARQUET_ASSIGN_OR_THROW(auto file, ::arrow::io::ReadableFile::Open(path)); + PARQUET_ASSIGN_OR_THROW(auto size, file->GetSize()); + PARQUET_ASSIGN_OR_THROW(auto data, file->Read(size)); + return data; +} + +class TestVariantParquetTesting : public ::testing::Test { + protected: + using Case = std::pair; + + static EncodedVariantValue LoadEncodedVariantFile(const std::string& dir, + const std::string& name) { + auto metadata = ReadFile(dir + "/" + name + ".metadata"); + auto value = ReadFile(dir + "/" + name + ".value"); + return EncodedVariantValue{.metadata = std::move(metadata), + .value = std::move(value)}; + } + + static std::vector LoadEncodedVariantCases(const std::string& dir) { + constexpr std::string_view kMetadataSuffix = ".metadata"; + PARQUET_ASSIGN_OR_THROW(auto dir_path, + ::arrow::internal::PlatformFilename::FromString(dir)); + PARQUET_ASSIGN_OR_THROW(auto entries, ::arrow::internal::ListDir(dir_path)); + + std::vector cases; + for (const auto& entry : entries) { + auto filename = entry.ToString(); + if (filename.ends_with(kMetadataSuffix)) { + auto name = filename.substr(0, filename.size() - kMetadataSuffix.size()); + auto value = LoadEncodedVariantFile(dir, name); + cases.emplace_back(std::move(name), std::move(value)); + } + } + return cases; + } + + static std::shared_ptr<::arrow::Table> RoundTripVariantArray( + const std::shared_ptr<::arrow::extension::VariantArray>& array) { + auto table = ::arrow::Table::Make( + ::arrow::schema({::arrow::field("variant", array->type())}), {array}); + auto arrow_writer_properties = + ArrowWriterProperties::Builder().set_variant_validation_enabled(true)->build(); + PARQUET_ASSIGN_OR_THROW(auto sink, ::arrow::io::BufferOutputStream::Create()); + PARQUET_THROW_NOT_OK(parquet::arrow::WriteTable( + *table, ::arrow::default_memory_pool(), sink, /*chunk_size=*/table->num_rows(), + default_writer_properties(), std::move(arrow_writer_properties))); + PARQUET_ASSIGN_OR_THROW(auto buffer, sink->Finish()); + + std::optional<::arrow::ExtensionTypeGuard> guard; + if (::arrow::GetExtensionType( + std::string(::arrow::extension::kVariantExtensionName)) == nullptr) { + guard.emplace(array->type()); + } + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + reader_properties.set_variant_validation_enabled(true); + parquet::arrow::FileReaderBuilder reader_builder; + PARQUET_THROW_NOT_OK(reader_builder.Open( + std::make_shared<::arrow::io::BufferReader>(std::move(buffer)))); + reader_builder.properties(reader_properties); + PARQUET_ASSIGN_OR_THROW(auto reader, reader_builder.Build()); + PARQUET_ASSIGN_OR_THROW(auto read_table, reader->ReadTable()); + return read_table; + } + + void SetUp() override { + auto maybe_dir = ParquetTestingSiblingDir("variant"); + if (!maybe_dir.has_value()) { + GTEST_SKIP() << "PARQUET_TEST_DATA not set or variant folder missing"; + } + ASSERT_NO_THROW(cases_ = LoadEncodedVariantCases(*maybe_dir)); + } + + std::vector cases_; +}; + +class TestShreddedVariantParquetTesting : public ::testing::Test { + protected: + struct ShreddedVariantCase { + std::string test_name; + std::string parquet_file; + std::vector> expected_files; + std::optional error_message; + }; + + std::shared_ptr<::arrow::Table> ReadVariantTestingTable( + const std::string& parquet_file, + ArrowReaderProperties reader_properties = ArrowReaderProperties()) const { + auto registered_storage_type = ::arrow::struct_( + {::arrow::field("metadata", ::arrow::binary(), /*nullable=*/false), + ::arrow::field("value", ::arrow::binary(), /*nullable=*/false)}); + std::optional<::arrow::ExtensionTypeGuard> guard; + if (::arrow::GetExtensionType( + std::string(::arrow::extension::kVariantExtensionName)) == nullptr) { + guard.emplace(::arrow::extension::variant(registered_storage_type)); + } + + parquet::arrow::FileReaderBuilder builder; + PARQUET_THROW_NOT_OK( + builder.OpenFile(dir_ + "/" + parquet_file, /*memory_map=*/false)); + reader_properties.set_arrow_extensions_enabled(true); + builder.properties(reader_properties); + PARQUET_ASSIGN_OR_THROW(auto reader, builder.Build()); + PARQUET_ASSIGN_OR_THROW(auto table, reader->ReadTable()); + return table; + } + + EncodedVariantValue LoadEncodedVariantBinaryFile(const std::string& name) const { + auto encoded_buffer = ReadFile(dir_ + "/" + name); + size_t metadata_size; + auto metadata_view = VariantMetadataView::ParsePrefix( + std::string_view{*encoded_buffer}, &metadata_size); + auto metadata = + ::arrow::SliceBuffer(encoded_buffer, 0, static_cast(metadata_size)); + auto value = + ::arrow::SliceBuffer(encoded_buffer, static_cast(metadata_size)); + VariantValueView::Validate(std::string_view{*value}, metadata_view); + return EncodedVariantValue{.metadata = std::move(metadata), + .value = std::move(value)}; + } + + std::unordered_map LoadShreddedVariantCases() const { + auto data = ReadFile(dir_ + "/cases.json"); + const auto data_view = std::string_view{*data}; + rj::Document document; + document.Parse(data_view.data(), data_view.size()); + if (document.HasParseError()) { + throw ParquetException("Cannot parse cases.json: ", + rj::GetParseError_En(document.GetParseError())); + } + if (!document.IsArray()) { + throw ParquetException("Expected cases.json root array"); + } + + std::unordered_set case_numbers; + std::unordered_map cases; + cases.reserve(document.Size()); + for (const auto& value : document.GetArray()) { + if (!value.IsObject()) { + throw ParquetException("Expected cases.json entry object"); + } + auto case_number = value.FindMember("case_number"); + if (case_number == value.MemberEnd() || !case_number->value.IsInt()) { + throw ParquetException("Expected case_number integer"); + } + const int number = case_number->value.GetInt(); + if (!case_numbers.insert(number).second) { + throw ParquetException("Duplicate case_number: ", number); + } + + auto parquet_file = value.FindMember("parquet_file"); + if (parquet_file == value.MemberEnd()) { + if (number != kMissingParquetFileCaseNumber || value.MemberCount() != 1) { + throw ParquetException("Missing parquet_file for case ", number); + } + continue; + } + if (!parquet_file->value.IsString()) { + throw ParquetException("Expected parquet_file string for case ", number); + } + auto test = value.FindMember("test"); + if (test == value.MemberEnd() || !test->value.IsString()) { + throw ParquetException("Expected test string for case ", number); + } + + ShreddedVariantCase test_case; + test_case.test_name = test->value.GetString(); + test_case.parquet_file = parquet_file->value.GetString(); + + auto variant_file = value.FindMember("variant_file"); + auto variant_files = value.FindMember("variant_files"); + if (variant_file != value.MemberEnd() && variant_files != value.MemberEnd()) { + throw ParquetException("Both variant_file and variant_files are set for case ", + number); + } + if (variant_file != value.MemberEnd()) { + if (!variant_file->value.IsString()) { + throw ParquetException("Expected variant_file string for case ", number); + } + test_case.expected_files.emplace_back(variant_file->value.GetString()); + } else if (variant_files != value.MemberEnd()) { + if (!variant_files->value.IsArray()) { + throw ParquetException("Expected variant_files array for case ", number); + } + for (const auto& entry : variant_files->value.GetArray()) { + if (entry.IsNull()) { + test_case.expected_files.emplace_back(std::nullopt); + } else if (entry.IsString()) { + test_case.expected_files.emplace_back(entry.GetString()); + } else { + throw ParquetException("Expected variant_files string or null for case ", + number); + } + } + } + + auto error_message = value.FindMember("error_message"); + if (error_message != value.MemberEnd()) { + if (!error_message->value.IsString()) { + throw ParquetException("Expected error_message string for case ", number); + } + test_case.error_message = std::string(error_message->value.GetString()); + } else if (test_case.expected_files.empty()) { + throw ParquetException("Missing expected variant file for case ", number); + } + cases.emplace(number, std::move(test_case)); + } + return cases; + } + + const ShreddedVariantCase& FindCase(int case_number, std::string_view test_name) const { + auto it = cases_.find(case_number); + if (it == cases_.end()) { + throw ParquetException("Missing case_number: ", case_number); + } + if (it->second.test_name != test_name) { + throw ParquetException("Unexpected test name for case ", case_number, ": ", + it->second.test_name); + } + return it->second; + } + + void SetUp() override { + auto maybe_dir = ParquetTestingSiblingDir("shredded_variant"); + if (!maybe_dir.has_value()) { + GTEST_SKIP() << "PARQUET_TEST_DATA not set or shredded variant folder missing"; + } + dir_ = std::move(*maybe_dir); + ASSERT_NO_THROW(cases_ = LoadShreddedVariantCases()); + } + + std::string dir_; + std::unordered_map cases_; +}; + +} // namespace + +TEST_F(TestVariantParquetTesting, RoundTripsEncoded) { + VariantArrayBuilder builder; + for (const auto& [name, encoded] : cases_) { + SCOPED_TRACE(name); + builder.AppendEncoded(encoded); + } + auto input = builder.Finish(); + + std::shared_ptr<::arrow::Table> table; + ASSERT_NO_THROW(table = RoundTripVariantArray(input)); + ASSERT_NE(nullptr, table); + auto column = table->GetColumnByName("variant"); + ASSERT_NE(nullptr, column); + ASSERT_EQ(1, column->num_chunks()); + const auto& array = + checked_pointer_cast<::arrow::extension::VariantArray>(column->chunk(0)); + ASSERT_EQ(static_cast(cases_.size()), array->length()); + + for (size_t row = 0; row < cases_.size(); ++row) { + const auto& [name, expected] = cases_[row]; + SCOPED_TRACE(name); + auto actual = UnshredVariantRow(*array, static_cast(row)); + ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); + ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); + } +} + +TEST_F(TestShreddedVariantParquetTesting, UnshredsArrayWithMissingValueColumn) { + // Case 131 is the simplest single-row fixture without a top-level value column. + const auto& test_case = FindCase(131, "testMissingValueColumn"); + auto table = ReadVariantTestingTable(test_case.parquet_file); + auto column = table->GetColumnByName("var"); + ASSERT_NE(nullptr, column); + ASSERT_EQ(1, column->num_chunks()); + auto variant_array = + checked_pointer_cast<::arrow::extension::VariantArray>(column->chunk(0)); + ASSERT_NE(nullptr, variant_array); + + auto unshredded = UnshredVariantArray(*variant_array); + auto actual = UnshredVariantRow(*unshredded, /*row=*/0); + ASSERT_EQ(1, test_case.expected_files.size()); + ASSERT_TRUE(test_case.expected_files[0].has_value()); + auto expected = LoadEncodedVariantBinaryFile(*test_case.expected_files[0]); + ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); + ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); +} + +TEST_F(TestShreddedVariantParquetTesting, ReadsShreddedParquetTesting) { + size_t skipped_cases = 0; + for (const auto& [case_number, test_case] : cases_) { + if (kReadCompatErrorCases.contains(case_number) || + kUnsupportedTypedValueCases.contains(case_number)) { + continue; + } + ASSERT_FALSE(test_case.error_message.has_value()); + + SCOPED_TRACE(::testing::Message() + << "case " << case_number << " (" << test_case.test_name << ")"); + auto table = ReadVariantTestingTable(test_case.parquet_file); + auto column = table->GetColumnByName("var"); + ASSERT_NE(nullptr, column); + ASSERT_EQ(1, column->num_chunks()); + auto variant_array = + checked_pointer_cast<::arrow::extension::VariantArray>(column->chunk(0)); + ASSERT_NE(nullptr, variant_array); + + auto skip = kSkipUnshreddedCases.find(case_number); + if (skip != kSkipUnshreddedCases.end()) { + ASSERT_EQ(skip->second, test_case.test_name); + ++skipped_cases; + continue; + } + + ASSERT_EQ(static_cast(test_case.expected_files.size()), + variant_array->length()); + for (int64_t row = 0; row < variant_array->length(); ++row) { + SCOPED_TRACE(row); + if (!test_case.expected_files[row].has_value()) { + ASSERT_TRUE(variant_array->IsNull(row)); + continue; + } + + auto expected = LoadEncodedVariantBinaryFile(*test_case.expected_files[row]); + auto actual = UnshredVariantRow(*variant_array, row); + ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); + ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); + } + } + ASSERT_EQ(kSkipUnshreddedCases.size(), skipped_cases); +} + +TEST_F(TestShreddedVariantParquetTesting, ReadCompatOnlyShapes) { + ArrowReaderProperties reader_properties; + reader_properties.set_variant_validation_enabled(false); + for (const auto& [case_number, test_name] : kReadCompatOnlyCases) { + const auto& test_case = FindCase(case_number, test_name); + SCOPED_TRACE(case_number); + + auto table = ReadVariantTestingTable(test_case.parquet_file, reader_properties); + auto column = table->GetColumnByName("var"); + ASSERT_NE(nullptr, column); + ASSERT_THROW(ValidateVariants(*column), ParquetInvalidOrCorruptedFileException); + ASSERT_NO_THROW(ValidateVariants(*column)); + } +} + +TEST_F(TestShreddedVariantParquetTesting, RejectsUnsupportedTypedValues) { + for (const auto& [case_number, test_name] : kUnsupportedTypedValueCases) { + const auto& test_case = FindCase(case_number, test_name); + SCOPED_TRACE(::testing::Message() + << "case " << case_number << " (" << test_case.test_name << ")"); + ASSERT_TRUE(test_case.error_message.has_value()); + ASSERT_FALSE(test_case.error_message->empty()); + ASSERT_THROW(ReadVariantTestingTable(test_case.parquet_file), ParquetException); + } +} + +TEST_F(TestShreddedVariantParquetTesting, RejectsErrorsInReadCompat) { + ArrowReaderProperties reader_properties; + reader_properties.set_variant_validation_enabled(false); + for (const auto& [case_number, test_name] : kReadCompatErrorCases) { + const auto& test_case = FindCase(case_number, test_name); + SCOPED_TRACE(case_number); + ASSERT_TRUE(test_case.error_message.has_value()); + ASSERT_FALSE(test_case.error_message->empty()); + + auto table = ReadVariantTestingTable(test_case.parquet_file, reader_properties); + auto column = table->GetColumnByName("var"); + ASSERT_NE(nullptr, column); + ASSERT_THROW(ValidateVariants(*column), + ParquetInvalidOrCorruptedFileException); + } +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/slot_internal.cc b/cpp/src/parquet/variant/slot_internal.cc new file mode 100644 index 000000000000..56acae16e259 --- /dev/null +++ b/cpp/src/parquet/variant/slot_internal.cc @@ -0,0 +1,703 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/slot_internal.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/extension_type.h" +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/decimal.h" +#include "arrow/util/logging_internal.h" +#include "arrow/util/unreachable.h" +#include "parquet/exception.h" +#include "parquet/variant/append_table_internal.h" +#include "parquet/variant/array_internal.h" + +namespace parquet::variant::internal { + +namespace { + +using ::arrow::Array; +using ::arrow::DataType; +using ::arrow::ExtensionArray; +using ::arrow::StructArray; +using ::arrow::internal::checked_cast; + +#ifdef _MSC_VER +# define VARIANT_TARGET_DECL(...) (BuildTarget * target, ##__VA_ARGS__) +# define VARIANT_OBJECT_CALL_ARGS(...) destination.field_name, ##__VA_ARGS__ +#else +# define VARIANT_TARGET_DECL(...) \ + (BuildTarget * target __VA_OPT__(, ) __VA_ARGS__) // NOLINT +# define VARIANT_OBJECT_CALL_ARGS(...) \ + destination.field_name __VA_OPT__(, ) __VA_ARGS__ // NOLINT +#endif + +#define DEFINE_TARGET_APPEND(NAME, decl_args, ...) \ + void Append##NAME VARIANT_TARGET_DECL decl_args { \ + if (target == nullptr) { \ + return; \ + } \ + std::visit( \ + [&](auto& destination) { \ + using Target = std::decay_t; \ + auto& builder = destination.builder.get(); \ + if constexpr (std::is_same_v) { \ + builder.Append##NAME(VARIANT_OBJECT_CALL_ARGS(__VA_ARGS__)); \ + } else { \ + builder.Append##NAME(__VA_ARGS__); \ + } \ + }, \ + target->destination); \ + } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_TARGET_APPEND) +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_TARGET_APPEND) +DEFINE_TARGET_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_TARGET_APPEND + +#undef VARIANT_OBJECT_CALL_ARGS +#undef VARIANT_TARGET_DECL + +VariantObjectBuilder StartObject(BuildTarget& target) { + return std::visit( + [](auto& destination) -> VariantObjectBuilder { + using Target = std::decay_t; + auto& builder = destination.builder.get(); + if constexpr (std::is_same_v) { + return builder.StartObject(destination.field_name); + } else { + return builder.StartObject(); + } + }, + target.destination); +} + +VariantListBuilder StartList(BuildTarget& target) { + return std::visit( + [](auto& destination) -> VariantListBuilder { + using Target = std::decay_t; + auto& builder = destination.builder.get(); + if constexpr (std::is_same_v) { + return builder.StartList(destination.field_name); + } else { + return builder.StartList(); + } + }, + target.destination); +} + +[[noreturn]] void UnsupportedType(const DataType& type) { + throw ParquetInvalidOrCorruptedFileException("Illegal shredded value type: ", + type.ToString()); +} + +std::optional GetValueSlot(const std::shared_ptr& value_array, + int64_t row) { + if (value_array == nullptr || value_array->IsNull(row)) { + return std::nullopt; + } + return BinaryFieldView(*value_array, row); +} + +CompiledVariantRowPlan CompileVariantRowPlanImpl( + const std::shared_ptr& value_array, const std::shared_ptr& typed_array, + std::string_view path); + +CompiledVariantRowPlan::FieldGroup CompileFieldGroupPlan( + const std::shared_ptr& field_group_array, std::string_view path, + CompiledVariantRowPlan& parent) { + if (field_group_array->type_id() != ::arrow::Type::STRUCT) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + path, ": expected struct storage, got ", + field_group_array->type()->ToString()); + } + const auto& field_struct = checked_cast(*field_group_array); + const auto child_plan_index = parent.children.size(); + parent.children.push_back( + CompileVariantRowPlanImpl(field_struct.GetFieldByName("value"), + field_struct.GetFieldByName("typed_value"), path)); + return { + .array = field_group_array, + .child_plan_index = child_plan_index, + }; +} + +CompiledTypedScalarPlan CompileTypedScalarPlan( + const std::shared_ptr& typed_array) { + switch (typed_array->type_id()) { + case ::arrow::Type::BOOL: + return {.kind = TypedScalarKind::kBoolean}; + case ::arrow::Type::INT8: + return {.kind = TypedScalarKind::kInt8}; + case ::arrow::Type::INT16: + return {.kind = TypedScalarKind::kInt16}; + case ::arrow::Type::INT32: + return {.kind = TypedScalarKind::kInt32}; + case ::arrow::Type::INT64: + return {.kind = TypedScalarKind::kInt64}; + case ::arrow::Type::FLOAT: + return {.kind = TypedScalarKind::kFloat}; + case ::arrow::Type::DOUBLE: + return {.kind = TypedScalarKind::kDouble}; + case ::arrow::Type::BINARY: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::BINARY_VIEW: + return {.kind = TypedScalarKind::kBinary, .physical_type = typed_array->type_id()}; + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_STRING: + case ::arrow::Type::STRING_VIEW: + return {.kind = TypedScalarKind::kString, .physical_type = typed_array->type_id()}; + case ::arrow::Type::DATE32: + return {.kind = TypedScalarKind::kDate}; + case ::arrow::Type::TIME64: { + const auto& type = checked_cast(*typed_array->type()); + if (type.unit() != ::arrow::TimeUnit::MICRO) { + UnsupportedType(*typed_array->type()); + } + return {.kind = TypedScalarKind::kTimeNTZMicros}; + } + case ::arrow::Type::TIMESTAMP: { + const auto& type = + checked_cast(*typed_array->type()); + switch (type.unit()) { + case ::arrow::TimeUnit::MICRO: + return {.kind = TypedScalarKind::kTimestampMicros, + .adjusted_to_utc = !type.timezone().empty()}; + case ::arrow::TimeUnit::NANO: + return {.kind = TypedScalarKind::kTimestampNanos, + .adjusted_to_utc = !type.timezone().empty()}; + case ::arrow::TimeUnit::SECOND: + case ::arrow::TimeUnit::MILLI: + UnsupportedType(*typed_array->type()); + } + ::arrow::Unreachable("Unexpected timestamp unit"); + } + case ::arrow::Type::DECIMAL32: { + const auto& type = + checked_cast(*typed_array->type()); + return {.kind = TypedScalarKind::kDecimal4From32, + .scale = static_cast(type.scale())}; + } + case ::arrow::Type::DECIMAL64: { + const auto& type = + checked_cast(*typed_array->type()); + if (type.precision() <= 9) { + return {.kind = TypedScalarKind::kDecimal4From64, + .scale = static_cast(type.scale())}; + } + return {.kind = TypedScalarKind::kDecimal8From64, + .scale = static_cast(type.scale())}; + } + case ::arrow::Type::DECIMAL128: { + const auto& type = + checked_cast(*typed_array->type()); + if (type.precision() <= 9) { + return {.kind = TypedScalarKind::kDecimal4From128, + .scale = static_cast(type.scale())}; + } + if (type.precision() <= 18) { + return {.kind = TypedScalarKind::kDecimal8From128, + .scale = static_cast(type.scale())}; + } + return {.kind = TypedScalarKind::kDecimal16From128, + .scale = static_cast(type.scale())}; + } + case ::arrow::Type::FIXED_SIZE_BINARY: { + const auto& type = + checked_cast(*typed_array->type()); + if (type.byte_width() != 16) { + UnsupportedType(*typed_array->type()); + } + return {.kind = TypedScalarKind::kUuidFixed}; + } + case ::arrow::Type::EXTENSION: { + const auto& ext_type = + checked_cast(*typed_array->type()); + if (ext_type.extension_name() != "arrow.uuid") { + UnsupportedType(*typed_array->type()); + } + return {.kind = TypedScalarKind::kUuidExtension}; + } + default: + UnsupportedType(*typed_array->type()); + } +} + +CompiledVariantRowPlan CompileVariantRowPlanImpl( + const std::shared_ptr& value_array, const std::shared_ptr& typed_array, + std::string_view path) { + CompiledVariantRowPlan plan{ + .value_array = value_array, + .typed = std::nullopt, + .children = {}, + }; + if (typed_array == nullptr) { + return plan; + } + + switch (typed_array->type_id()) { + case ::arrow::Type::STRUCT: { + const auto& typed_struct = checked_cast(*typed_array); + CompiledVariantRowPlan::Object object; + object.fields.reserve(typed_struct.struct_type()->num_fields()); + object.field_names.reserve(typed_struct.struct_type()->num_fields()); + for (int i = 0; i < typed_struct.struct_type()->num_fields(); ++i) { + const auto& field_name = typed_struct.struct_type()->field(i)->name(); + if (!object.field_names.insert(field_name).second) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: duplicate shredded object field '", field_name, + "'"); + } + auto field_group = CompileFieldGroupPlan(typed_struct.field(i), field_name, plan); + object.fields.push_back( + {.field_name = field_name, .field_group = std::move(field_group)}); + } + plan.typed.emplace(CompiledVariantRowPlan::Typed{ + .array = typed_array, + .plan = std::move(object), + }); + return plan; + } + case ::arrow::Type::LIST: + case ::arrow::Type::LARGE_LIST: + case ::arrow::Type::LIST_VIEW: + case ::arrow::Type::LARGE_LIST_VIEW: + case ::arrow::Type::FIXED_SIZE_LIST: { + auto values = ValuesArray(*typed_array); + auto element = CompileFieldGroupPlan(values, path, plan); + plan.typed.emplace(CompiledVariantRowPlan::Typed{ + .array = typed_array, + .plan = CompiledVariantRowPlan::Array{.element = std::move(element)}, + }); + return plan; + } + default: + plan.typed.emplace(CompiledVariantRowPlan::Typed{ + .array = typed_array, + .plan = + CompiledVariantRowPlan::Primitive{ + .scalar = CompileTypedScalarPlan(typed_array), + }, + }); + return plan; + } +} + +void ValidateTypedFieldNames(const VariantMetadataView& metadata, + const CompiledVariantRowPlan::Object& object) { + for (const auto& field : object.fields) { + if (!metadata.FindString(field.field_name).has_value()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: shredded field '", field.field_name, + "' is not in metadata dictionary"); + } + } +} + +std::string_view GetStringScalarView(const CompiledTypedScalarPlan& plan, + const std::shared_ptr& typed_array, + int64_t row) { + switch (plan.physical_type) { + case ::arrow::Type::STRING: + return checked_cast(*typed_array).GetView(row); + case ::arrow::Type::LARGE_STRING: + return checked_cast(*typed_array).GetView(row); + case ::arrow::Type::STRING_VIEW: + return checked_cast(*typed_array).GetView(row); + default: + DCHECK(false); + return {}; + } +} + +void AppendTypedScalar(BuildTarget* target, const CompiledTypedScalarPlan& plan, + const std::shared_ptr& typed_array, int64_t row) { + if (target == nullptr) { + return; + } + switch (plan.kind) { + case TypedScalarKind::kBoolean: + AppendBoolean(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt8: + AppendInt8(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt16: + AppendInt16(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt32: + AppendInt32(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt64: + AppendInt64(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kFloat: + AppendFloat(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kDouble: + AppendDouble(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kBinary: + AppendBinary(target, BinaryFieldView(*typed_array, row)); + return; + case TypedScalarKind::kString: { + auto value = GetStringScalarView(plan, typed_array, row); + if (value.size() <= 0x3f) { + AppendShortString(target, value); + } else { + AppendString(target, value); + } + return; + } + case TypedScalarKind::kDate: + AppendDate(target, + checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kTimeNTZMicros: + AppendTimeNTZMicros( + target, checked_cast(*typed_array).Value(row)); + return; + case TypedScalarKind::kTimestampMicros: + AppendTimestampMicros( + target, checked_cast(*typed_array).Value(row), + plan.adjusted_to_utc); + return; + case TypedScalarKind::kTimestampNanos: + AppendTimestampNanos( + target, checked_cast(*typed_array).Value(row), + plan.adjusted_to_utc); + return; + case TypedScalarKind::kDecimal4From32: { + const auto value = ::arrow::Decimal32( + checked_cast(*typed_array).GetValue(row)); + AppendDecimal4(target, value, plan.scale); + return; + } + case TypedScalarKind::kDecimal4From64: { + const auto value = ::arrow::Decimal64( + checked_cast(*typed_array).GetValue(row)); + PARQUET_ASSIGN_OR_THROW(auto raw, value.ToInteger()); + AppendDecimal4(target, ::arrow::Decimal32(raw), plan.scale); + return; + } + case TypedScalarKind::kDecimal8From64: { + const auto value = ::arrow::Decimal64( + checked_cast(*typed_array).GetValue(row)); + AppendDecimal8(target, value, plan.scale); + return; + } + case TypedScalarKind::kDecimal4From128: { + const auto value = ::arrow::Decimal128( + checked_cast(*typed_array).GetValue(row)); + PARQUET_ASSIGN_OR_THROW(auto raw, value.ToInteger()); + AppendDecimal4(target, ::arrow::Decimal32(raw), plan.scale); + return; + } + case TypedScalarKind::kDecimal8From128: { + const auto value = ::arrow::Decimal128( + checked_cast(*typed_array).GetValue(row)); + PARQUET_ASSIGN_OR_THROW(auto raw, value.ToInteger()); + AppendDecimal8(target, ::arrow::Decimal64(raw), plan.scale); + return; + } + case TypedScalarKind::kDecimal16From128: { + const auto value = ::arrow::Decimal128( + checked_cast(*typed_array).GetValue(row)); + AppendDecimal16(target, value, plan.scale); + return; + } + case TypedScalarKind::kUuidFixed: { + auto value = + checked_cast(*typed_array).GetView(row); + AppendUuid(target, value); + return; + } + case TypedScalarKind::kUuidExtension: { + const auto& ext_array = checked_cast(*typed_array); + auto storage = ext_array.storage(); + auto value = + checked_cast(*storage).GetView(row); + AppendUuid(target, value); + return; + } + } +} + +template +void ProcessListElement(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::FieldGroup& field_group_plan, + int64_t row, std::string_view path, BuildTarget* target) { + const auto& field_array = *field_group_plan.array; + if (field_array.IsNull(row)) { + if constexpr (strict) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant field at ", path, ": field group must be required"); + } else { + AppendVariantNull(target); + return; + } + } + + DCHECK_LT(field_group_plan.child_plan_index, parent_plan.children.size()); + const auto& row_plan = parent_plan.children[field_group_plan.child_plan_index]; + ProcessSlot(metadata, row_plan, row, target, path); +} + +template +void ProcessTypedArraySlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::Typed& typed, + const CompiledVariantRowPlan::Array& array, int64_t row, + std::string_view path, BuildTarget* target) { + DCHECK(!typed.array->IsNull(row)); + + std::optional list_builder; + std::optional child_target; + BuildTarget* child_target_ptr = nullptr; + if (target != nullptr) { + list_builder.emplace(StartList(*target)); + child_target.emplace(BuildTarget{BuildTarget::ListElement{*list_builder}}); + child_target_ptr = &*child_target; + } + const auto [offset, length] = ValuesRangeAt(*typed.array, row); + for (int64_t i = 0; i < length; ++i) { + ProcessListElement(metadata, parent_plan, array.element, offset + i, path, + child_target_ptr); + } + if (list_builder.has_value()) { + list_builder->Finish(); + } +} + +template +bool TryProcessSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& plan, int64_t row, + std::string_view path, BuildTarget* target); + +template +void ProcessObjectField(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::ObjectField& field, int64_t row, + BuildTarget* target) { + const auto& field_array = *field.field_group.array; + if (field_array.IsNull(row)) { + if constexpr (strict) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + field.field_name, + ": field group must be required"); + } else { + return; + } + } + + DCHECK_LT(field.field_group.child_plan_index, parent_plan.children.size()); + const auto& row_plan = parent_plan.children[field.field_group.child_plan_index]; + if constexpr (strict) { + if (row_plan.value_array == nullptr) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + field.field_name, + ": missing value field"); + } + } + + if (target != nullptr) { + std::get(target->destination).field_name = field.field_name; + } + TryProcessSlot(metadata, row_plan, row, field.field_name, target); +} + +template +void ProcessObjectResidual(const VariantValueView& value_view, + const CompiledVariantRowPlan::Object& object, + BuildTarget* target) { + DCHECK_EQ(value_view.basic_type(), VariantBasicType::kObject); + const auto& object_view = std::get(value_view.data()); + for (const auto& field : object_view.fields()) { + if (object.field_names.contains(field.name)) { + if (strict || target != nullptr) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: value object contains shredded field: ", + field.name); + } + continue; + } + if (target == nullptr) { + continue; + } + BuildTarget residual_target = *target; + std::get(residual_target.destination).field_name = + field.name; + AppendEncodedValue(&residual_target, field.value); + } +} + +template +void ProcessTypedObjectSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::Typed& typed, + const CompiledVariantRowPlan::Object& object, int64_t row, + std::optional value, BuildTarget* target) { + DCHECK(!typed.array->IsNull(row)); + + std::optional object_value_view; + if (value.has_value()) { + if (VariantValueView::PeekBasicType(*value) != VariantBasicType::kObject) { + throw ParquetInvalidOrCorruptedFileException( + "Expected object in value field for partially shredded struct"); + } + if (strict || target != nullptr) { + object_value_view = VariantValueView::Make(*value, metadata); + } else { + VariantValueView::Validate(*value, metadata); + } + } + + std::optional object_builder; + std::optional child_target; + BuildTarget* child_target_ptr = nullptr; + if (target != nullptr) { + object_builder.emplace(StartObject(*target)); + child_target.emplace(BuildTarget{BuildTarget::ObjectField{ + .builder = *object_builder, + .field_name = {}, + }}); + child_target_ptr = &*child_target; + } + for (const auto& field : object.fields) { + ProcessObjectField(metadata, parent_plan, field, row, child_target_ptr); + } + if (object_value_view.has_value()) { + ProcessObjectResidual(*object_value_view, object, child_target_ptr); + } + if (object_builder.has_value()) { + object_builder->Finish(); + } +} + +template +bool TryProcessSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& plan, int64_t row, + std::string_view path, BuildTarget* target) { + const auto value = GetValueSlot(plan.value_array, row); + if (!plan.typed.has_value()) { + if (!value.has_value()) { + return false; + } + VariantValueView::Validate(*value, metadata); + AppendEncodedValue(target, *value); + return true; + } + + const auto& typed = *plan.typed; + return std::visit( + [&](const auto& typed_plan) -> bool { + using Plan = std::decay_t; + const bool typed_present = !typed.array->IsNull(row); + + if constexpr (std::is_same_v) { + ValidateTypedFieldNames(metadata, typed_plan); + } + + if (!typed_present) { + if (!value.has_value()) { + return false; + } + VariantValueView::Validate(*value, metadata); + if constexpr (std::is_same_v || + std::is_same_v) { + const auto basic_type = VariantValueView::PeekBasicType(*value); + if constexpr (std::is_same_v) { + if (basic_type == VariantBasicType::kArray) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: array value must be stored in " + "typed_value"); + } + } else if (basic_type == VariantBasicType::kObject) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: object value requires object typed_value"); + } + } + AppendEncodedValue(target, *value); + return true; + } + + if constexpr (!std::is_same_v) { + if (value.has_value()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded variant: both value and typed_value are non-null"); + } + } + + if constexpr (std::is_same_v) { + AppendTypedScalar(target, typed_plan.scalar, typed.array, row); + } else if constexpr (std::is_same_v) { + ProcessTypedArraySlot(metadata, plan, typed, typed_plan, row, path, + target); + } else { + ProcessTypedObjectSlot(metadata, plan, typed, typed_plan, row, value, + target); + } + return true; + }, + typed.plan); +} + +} // namespace + +CompiledVariantRowPlan CompileVariantRowPlan(const std::shared_ptr& value_array, + const std::shared_ptr& typed_array) { + return CompileVariantRowPlanImpl(value_array, typed_array, "root"); +} + +template +void ProcessSlot(const VariantMetadataView& metadata, const CompiledVariantRowPlan& plan, + int64_t row, BuildTarget* target, std::string_view path) { + if (TryProcessSlot(metadata, plan, row, path, target)) { + return; + } + + if constexpr (strict) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant at ", path, + ": value and typed_value are both null"); + } + AppendVariantNull(target); +} + +template void ProcessSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& plan, int64_t row, + BuildTarget* target, std::string_view path); +template void ProcessSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& plan, int64_t row, + BuildTarget* target, std::string_view path); + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/slot_internal.h b/cpp/src/parquet/variant/slot_internal.h new file mode 100644 index 000000000000..0808dfc1f148 --- /dev/null +++ b/cpp/src/parquet/variant/slot_internal.h @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/type_fwd.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/decoding.h" + +namespace parquet::variant::internal { + +enum class TypedScalarKind { + kBoolean, + kInt8, + kInt16, + kInt32, + kInt64, + kFloat, + kDouble, + kBinary, + kString, + kDate, + kTimeNTZMicros, + kTimestampMicros, + kTimestampNanos, + kDecimal4From32, + kDecimal4From64, + kDecimal8From64, + kDecimal4From128, + kDecimal8From128, + kDecimal16From128, + kUuidFixed, + kUuidExtension, +}; + +struct CompiledTypedScalarPlan { + TypedScalarKind kind; + ::arrow::Type::type physical_type = ::arrow::Type::NA; + uint8_t scale = 0; + bool adjusted_to_utc = false; +}; + +struct CompiledVariantRowPlan { + struct FieldGroup { + std::shared_ptr<::arrow::Array> array; + // Stores the child plan by index because the parent child vector may reallocate. + size_t child_plan_index = 0; + }; + + struct ObjectField { + // Borrows the field name owned by the typed array retained in the parent plan. + std::string_view field_name; + FieldGroup field_group; + }; + + struct Primitive { + CompiledTypedScalarPlan scalar; + }; + + struct Array { + FieldGroup element; + }; + + struct Object { + std::vector fields; + std::unordered_set field_names; + }; + + struct Typed { + std::shared_ptr<::arrow::Array> array; + std::variant plan; + }; + + std::shared_ptr<::arrow::Array> value_array; + std::optional typed; + std::vector children; +}; + +struct BuildTarget { + struct Row { + std::reference_wrapper builder; + }; + + struct ObjectField { + std::reference_wrapper builder; + std::string_view field_name; + }; + + struct ListElement { + std::reference_wrapper builder; + }; + + using Destination = std::variant; + Destination destination; +}; + +CompiledVariantRowPlan CompileVariantRowPlan( + const std::shared_ptr<::arrow::Array>& value_array, + const std::shared_ptr<::arrow::Array>& typed_array); + +template +void ProcessSlot(const VariantMetadataView& metadata, const CompiledVariantRowPlan& plan, + int64_t row, BuildTarget* target = nullptr, + std::string_view path = "root"); + +} // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/test_util_internal.cc b/cpp/src/parquet/variant/test_util_internal.cc index 064954f39af4..de95aabdb216 100644 --- a/cpp/src/parquet/variant/test_util_internal.cc +++ b/cpp/src/parquet/variant/test_util_internal.cc @@ -17,29 +17,20 @@ #include "parquet/variant/test_util_internal.h" -#include #include #include "arrow/array.h" // IWYU pragma: keep #include "arrow/array/builder_binary.h" -#include "arrow/array/builder_primitive.h" -#include "arrow/extension/uuid.h" #include "arrow/extension_type.h" #include "arrow/io/memory.h" #include "arrow/table.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" -#include "parquet/arrow/reader.h" #include "parquet/arrow/writer.h" #include "parquet/exception.h" namespace parquet::variant::internal { -VariantValueView MakeVariantValueView(const EncodedVariantValue& encoded) { - auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); - return VariantValueView::Make(std::string_view{*encoded.value}, metadata); -} - std::shared_ptr<::arrow::Table> VariantTable( const std::shared_ptr& variant_type, const std::vector>& storage_children, @@ -78,30 +69,6 @@ ::arrow::Status WriteVariantRecordBatch( return writer->Close(); } -std::optional ShreddedVariantTestingDir() { - const char* data = std::getenv("PARQUET_TEST_DATA"); - if (data == nullptr) { - return std::nullopt; - } - std::string path(data); - const auto pos = path.find_last_of("/\\"); - if (pos == std::string::npos) { - return std::nullopt; - } - return path.substr(0, pos) + "/shredded_variant"; -} - -Result> ReadVariantTestingTable(const std::string& path) { - ArrowReaderProperties reader_properties; - reader_properties.set_arrow_extensions_enabled(true); - - parquet::arrow::FileReaderBuilder builder; - RETURN_NOT_OK(builder.OpenFile(path, /*memory_map=*/false)); - builder.properties(reader_properties); - ARROW_ASSIGN_OR_RAISE(auto reader, builder.Build()); - return reader->ReadTable(); -} - std::shared_ptr EmptyVariantMetadata() { VariantBuilder builder; builder.AppendVariantNull(); @@ -119,78 +86,11 @@ std::shared_ptr BinaryArrayFromValues( const std::vector>& values) { ::arrow::BinaryBuilder builder; for (const auto& value : values) { - if (value.has_value()) { - ARROW_EXPECT_OK(builder.Append(*value)); - } else { - ARROW_EXPECT_OK(builder.AppendNull()); - } - } - std::shared_ptr out; - ARROW_EXPECT_OK(builder.Finish(&out)); - return out; -} - -std::shared_ptr BinaryViewArrayFromValues( - const std::vector>& values) { - ::arrow::BinaryViewBuilder builder; - for (const auto& value : values) { - if (value.has_value()) { - ARROW_EXPECT_OK(builder.Append(*value)); - } else { - ARROW_EXPECT_OK(builder.AppendNull()); - } - } - std::shared_ptr out; - ARROW_EXPECT_OK(builder.Finish(&out)); - return out; -} - -std::shared_ptr Int64ArrayFromValues( - const std::vector>& values) { - ::arrow::Int64Builder builder; - for (const auto& value : values) { - if (value.has_value()) { - ARROW_EXPECT_OK(builder.Append(*value)); - } else { - ARROW_EXPECT_OK(builder.AppendNull()); - } - } - std::shared_ptr out; - ARROW_EXPECT_OK(builder.Finish(&out)); - return out; -} - -std::shared_ptr Int32ArrayFromValues(const std::vector& values) { - ::arrow::Int32Builder builder; - ARROW_EXPECT_OK(builder.AppendValues(values)); - std::shared_ptr out; - ARROW_EXPECT_OK(builder.Finish(&out)); - return out; -} - -std::shared_ptr StringArrayFromValues( - const std::vector>& values) { - ::arrow::StringBuilder builder; - for (const auto& value : values) { - if (value.has_value()) { - ARROW_EXPECT_OK(builder.Append(*value)); - } else { - ARROW_EXPECT_OK(builder.AppendNull()); - } + ARROW_EXPECT_OK(builder.AppendOrNull(value)); } std::shared_ptr out; ARROW_EXPECT_OK(builder.Finish(&out)); return out; } -std::shared_ptr UuidArrayFromValues(const std::vector& values) { - ::arrow::FixedSizeBinaryBuilder builder(::arrow::fixed_size_binary(16)); - for (const auto& value : values) { - ARROW_EXPECT_OK(builder.Append(value)); - } - std::shared_ptr storage; - ARROW_EXPECT_OK(builder.Finish(&storage)); - return ::arrow::ExtensionType::WrapArray(::arrow::extension::uuid(), storage); -} - } // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/test_util_internal.h b/cpp/src/parquet/variant/test_util_internal.h index 30186c2f533c..9edab8403da5 100644 --- a/cpp/src/parquet/variant/test_util_internal.h +++ b/cpp/src/parquet/variant/test_util_internal.h @@ -20,7 +20,6 @@ #include #include #include -#include #include #include @@ -28,7 +27,6 @@ #include "arrow/type_fwd.h" #include "parquet/properties.h" #include "parquet/variant/builder.h" -#include "parquet/variant/encoding.h" namespace parquet::variant::internal { @@ -38,8 +36,6 @@ using ::arrow::DataType; using ::arrow::FieldVector; using ::arrow::Result; -VariantValueView MakeVariantValueView(const EncodedVariantValue& encoded); - std::shared_ptr<::arrow::Table> VariantTable( const std::shared_ptr& variant_type, const std::vector>& storage_children, @@ -56,10 +52,6 @@ ::arrow::Status WriteVariantRecordBatch( std::shared_ptr arrow_properties = default_arrow_writer_properties()); -std::optional ShreddedVariantTestingDir(); - -Result> ReadVariantTestingTable(const std::string& path); - std::shared_ptr EmptyVariantMetadata(); EncodedVariantValue Int8Variant(int8_t value); @@ -67,17 +59,4 @@ EncodedVariantValue Int8Variant(int8_t value); std::shared_ptr BinaryArrayFromValues( const std::vector>& values); -std::shared_ptr BinaryViewArrayFromValues( - const std::vector>& values); - -std::shared_ptr Int64ArrayFromValues( - const std::vector>& values); - -std::shared_ptr Int32ArrayFromValues(const std::vector& values); - -std::shared_ptr StringArrayFromValues( - const std::vector>& values); - -std::shared_ptr UuidArrayFromValues(const std::vector& values); - } // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/type_test.cc b/cpp/src/parquet/variant/type_test.cc index c6e93263ab83..d5b00802778c 100644 --- a/cpp/src/parquet/variant/type_test.cc +++ b/cpp/src/parquet/variant/type_test.cc @@ -52,32 +52,40 @@ using ::arrow::extension::kVariantExtensionName; using ::arrow::extension::VariantExtensionType; TEST(TestVariantType, Storage) { - auto unshredded = struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary(), /*nullable=*/false)}); - ASSERT_OK_AND_ASSIGN(auto type, VariantExtensionType::Make(unshredded)); - ASSERT_EQ( - kVariantExtensionName, - ::arrow::internal::checked_cast(*type).extension_name()); - - auto shredded = struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary()), field("typed_value", int64())}); - ASSERT_OK_AND_ASSIGN(auto shredded_type, VariantExtensionType::Make(shredded)); - auto variant_type = - ::arrow::internal::checked_pointer_cast(shredded_type); - ASSERT_EQ("metadata", variant_type->metadata()->name()); - ASSERT_EQ("value", variant_type->value()->name()); - ASSERT_EQ("typed_value", variant_type->typed_value()->name()); + { + auto unshredded = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN(auto type, VariantExtensionType::Make(unshredded)); + ASSERT_EQ( + kVariantExtensionName, + ::arrow::internal::checked_cast(*type).extension_name()); + } - auto typed_only = struct_( - {field("metadata", binary(), /*nullable=*/false), field("typed_value", int64())}); - ASSERT_OK(VariantExtensionType::Make(typed_only)); + { + auto shredded = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary()), field("typed_value", int64())}); + ASSERT_OK_AND_ASSIGN(auto shredded_type, VariantExtensionType::Make(shredded)); + auto variant_type = + ::arrow::internal::checked_pointer_cast(shredded_type); + ASSERT_EQ("metadata", variant_type->metadata()->name()); + ASSERT_EQ("value", variant_type->value()->name()); + ASSERT_EQ("typed_value", variant_type->typed_value()->name()); + } - auto flipped = - std::dynamic_pointer_cast(::arrow::extension::variant( - struct_({field("value", binary(), /*nullable=*/false), - field("metadata", binary(), /*nullable=*/false)}))); - ASSERT_EQ("metadata", flipped->metadata()->name()); - ASSERT_EQ("value", flipped->value()->name()); + { + auto typed_only = struct_( + {field("metadata", binary(), /*nullable=*/false), field("typed_value", int64())}); + ASSERT_RAISES(Invalid, VariantExtensionType::Make(typed_only)); + } + + { + auto flipped = + std::dynamic_pointer_cast(::arrow::extension::variant( + struct_({field("value", binary(), /*nullable=*/false), + field("metadata", binary(), /*nullable=*/false)}))); + ASSERT_EQ("metadata", flipped->metadata()->name()); + ASSERT_EQ("value", flipped->value()->name()); + } ASSERT_OK(VariantExtensionType::Make( struct_({field("metadata", binary_view(), /*nullable=*/false), @@ -164,4 +172,23 @@ TEST(TestVariantType, InvalidStorage) { } } +TEST(TestVariantType, ReadStorage) { + auto base_storage = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN(auto base_type, VariantExtensionType::Make(base_storage)); + auto variant_type = + ::arrow::internal::checked_pointer_cast(base_type); + + auto typed_only = struct_( + {field("metadata", binary(), /*nullable=*/false), field("typed_value", int64())}); + ASSERT_OK(variant_type->Deserialize(typed_only, /*serialized_data=*/"")); + + auto typed_only_field_group = struct_({field("typed_value", ::arrow::utf8())}); + auto shredded_object = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", + struct_({field("a", typed_only_field_group, /*nullable=*/false)}))}); + ASSERT_OK(variant_type->Deserialize(shredded_object, /*serialized_data=*/"")); +} + } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/unshred.cc b/cpp/src/parquet/variant/unshred.cc new file mode 100644 index 000000000000..a6a2a19b5168 --- /dev/null +++ b/cpp/src/parquet/variant/unshred.cc @@ -0,0 +1,141 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/unshred.h" + +#include "arrow/buffer.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/type.h" +#include "arrow/util/bit_block_counter.h" +#include "arrow/util/bitmap_ops.h" +#include "parquet/exception.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/decoding.h" +#include "parquet/variant/slot_internal.h" + +namespace parquet::variant { + +using ::arrow::binary_view; +using ::arrow::field; +using ::arrow::struct_; +using ::arrow::StructArray; +using ::arrow::extension::VariantArray; +using ::arrow::internal::checked_cast; +using ::arrow::internal::checked_pointer_cast; + +namespace { + +std::shared_ptr<::arrow::Buffer> ParentNullBitmap(const VariantArray& array, + ::arrow::MemoryPool* pool) { + if (array.null_bitmap() == nullptr) { + return nullptr; + } + if (array.offset() == 0) { + return array.null_bitmap(); + } + PARQUET_ASSIGN_OR_THROW(auto null_bitmap, + ::arrow::internal::CopyBitmap(pool, array.null_bitmap_data(), + array.offset(), array.length())); + return null_bitmap; +} + +void AppendUnshreddedRow(std::string_view metadata_value, + const internal::CompiledVariantRowPlan& row_plan, int64_t row, + VariantValueArrayBuilder& value_builder) { + auto metadata = VariantMetadataView::Make(metadata_value); + auto row_builder = value_builder.BindMetadata(metadata); + internal::BuildTarget target{internal::BuildTarget::Row{row_builder}}; + internal::ProcessSlot(metadata, row_plan, row, &target); + row_builder.Finish(); +} + +} // namespace + +EncodedVariantValue UnshredVariantRow(const VariantArray& array, int64_t row, + ::arrow::MemoryPool* pool) { + if (row < 0 || row >= array.length()) { + throw ParquetException("Variant row out of bounds: ", row); + } + if (array.IsNull(row)) { + throw ParquetException("Cannot unshred null Variant row"); + } + const auto& storage = checked_cast(*array.storage()); + auto metadata_array = storage.GetFieldByName("metadata"); + auto value_array = storage.GetFieldByName("value"); + auto typed_array = storage.GetFieldByName("typed_value"); + if (metadata_array->IsNull(row)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant extension storage: metadata is null"); + } + auto metadata_value = internal::BinaryFieldView(*metadata_array, row); + + if (typed_array == nullptr && value_array != nullptr) { + if (value_array->IsNull(row)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant extension storage: value is null"); + } + return EncodedVariantValue{ + .metadata = ::arrow::Buffer::FromString(std::string(metadata_value)), + .value = ::arrow::Buffer::FromString( + std::string(internal::BinaryFieldView(*value_array, row)))}; + } + + auto row_plan = internal::CompileVariantRowPlan(value_array, typed_array); + VariantValueArrayBuilder value_builder(pool); + AppendUnshreddedRow(metadata_value, row_plan, row, value_builder); + auto unshredded_value = value_builder.Finish(); + return EncodedVariantValue{ + .metadata = ::arrow::Buffer::FromString(std::string(metadata_value)), + .value = ::arrow::Buffer::FromString(std::string(unshredded_value->GetView(0)))}; +} + +std::shared_ptr<::arrow::extension::VariantArray> UnshredVariantArray( + const ::arrow::extension::VariantArray& array, ::arrow::MemoryPool* pool) { + const auto& storage = checked_cast(*array.storage()); + auto metadata_array = storage.GetFieldByName("metadata"); + auto value_array = storage.GetFieldByName("value"); + auto typed_array = storage.GetFieldByName("typed_value"); + + if (typed_array == nullptr && value_array != nullptr) { + return MakeVariantArrayFromStorage( + checked_pointer_cast(array.storage())); + } + + auto row_plan = internal::CompileVariantRowPlan(value_array, typed_array); + VariantValueArrayBuilder value_builder(pool); + ::arrow::internal::VisitBitBlocksVoid( + array.null_bitmap_data(), array.offset(), array.length(), + [&](int64_t row) { + if (metadata_array->IsNull(row)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant extension storage: metadata is null"); + } + auto metadata_value = internal::BinaryFieldView(*metadata_array, row); + AppendUnshreddedRow(metadata_value, row_plan, row, value_builder); + }, + [&] { value_builder.AppendNull(); }); + + std::shared_ptr<::arrow::Array> unshredded_value = value_builder.Finish(); + auto storage_type = + struct_({field("metadata", metadata_array->type(), /*nullable=*/false), + field("value", binary_view(), /*nullable=*/false)}); + return MakeVariantArrayFromChildren( + std::move(storage_type), {std::move(metadata_array), std::move(unshredded_value)}, + ParentNullBitmap(array, pool)); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/unshred.h b/cpp/src/parquet/variant/unshred.h new file mode 100644 index 000000000000..e6d7b617ac6f --- /dev/null +++ b/cpp/src/parquet/variant/unshred.h @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include "parquet/platform.h" +#include "parquet/variant/builder.h" + +namespace arrow::extension { +class VariantArray; +} // namespace arrow::extension + +namespace parquet::variant { + +/// Materialize a non-null Variant row into independently owned encoded buffers. +/// +/// The row index must be in bounds and the parent row must be non-null. When the input +/// is already unshredded, the encoded metadata and value are copied without validation. +/// Invalid or corrupted storage detected during materialization raises a Parquet +/// exception. +PARQUET_EXPORT +EncodedVariantValue UnshredVariantRow( + const ::arrow::extension::VariantArray& array, int64_t row, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + +/// Materialize a Variant array while preserving parent nulls. +/// +/// An already unshredded array may reuse its entire storage without validating the +/// encoded metadata or values. When materialization is required, the output may reuse +/// the input metadata buffers. Invalid or corrupted storage detected during +/// materialization raises a Parquet exception. +PARQUET_EXPORT +std::shared_ptr<::arrow::extension::VariantArray> UnshredVariantArray( + const ::arrow::extension::VariantArray& array, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/unshred_test.cc b/cpp/src/parquet/variant/unshred_test.cc new file mode 100644 index 000000000000..e8d462032afd --- /dev/null +++ b/cpp/src/parquet/variant/unshred_test.cc @@ -0,0 +1,355 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/unshred.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep + +#include "arrow/buffer.h" +#include "arrow/chunked_array.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/testing/builder.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/type.h" +#include "arrow/util/decimal.h" +#include "parquet/exception.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/test_util_internal.h" +#include "parquet/variant/validate.h" + +namespace parquet::variant { + +using ::arrow::binary; +using ::arrow::field; +using ::arrow::int64; +using ::arrow::struct_; +using ::arrow::internal::checked_pointer_cast; +using internal::BinaryArrayFromValues; + +namespace { + +void AssertUnshreddedRow(const ::arrow::extension::VariantArray& array, + const EncodedVariantValue& expected) { + auto actual = UnshredVariantRow(array, /*row=*/0); + ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); + ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); +} + +EncodedVariantValue ObjectWithInt64(std::string_view field_name, int64_t value) { + VariantBuilder builder; + auto object = builder.StartObject(); + object.AppendInt64(field_name, value); + object.Finish(); + return builder.Finish(); +} + +EncodedVariantValue EmptyObjectWithField(std::string_view field_name) { + VariantBuilder builder; + builder.AddFieldName(field_name); + auto object = builder.StartObject(); + object.Finish(); + return builder.Finish(); +} + +std::shared_ptr<::arrow::StructArray> MakeInt64FieldGroup( + const std::vector>& values, + std::string_view typed_values, const std::vector& is_valid = {}) { + auto field_group_type = + struct_({field("value", binary()), field("typed_value", int64())}); + std::shared_ptr<::arrow::Buffer> null_bitmap; + if (!is_valid.empty()) { + ::arrow::BitmapFromVector(is_valid, &null_bitmap); + } + PARQUET_ASSIGN_OR_THROW( + auto field_group, ::arrow::StructArray::Make( + {BinaryArrayFromValues(values), + ::arrow::ArrayFromJSON(int64(), std::string(typed_values))}, + field_group_type->fields(), std::move(null_bitmap))); + return field_group; +} + +std::shared_ptr<::arrow::extension::VariantArray> MakeSingleRowVariantArray( + std::string_view metadata, std::optional value, + const std::shared_ptr<::arrow::Array>& typed_value) { + auto storage_type = + struct_({field("metadata", binary(), /*nullable=*/false), field("value", binary()), + field("typed_value", typed_value->type())}); + return MakeVariantArrayFromChildren( + storage_type, + {BinaryArrayFromValues({metadata}), BinaryArrayFromValues({value}), typed_value}); +} + +} // namespace + +TEST(TestVariantUnshred, FastPathKeepsStorage) { + VariantArrayBuilder builder; + builder.AppendNull(); + builder.AppendInt32(42); + auto array = builder.Finish(); + + auto unshredded = UnshredVariantArray(*array); + const auto& input_storage = + checked_pointer_cast<::arrow::StructArray>(array->storage()); + const auto& output_storage = + checked_pointer_cast<::arrow::StructArray>(unshredded->storage()); + ASSERT_EQ(nullptr, output_storage->GetFieldByName("typed_value").get()); + + const auto& input_metadata = checked_pointer_cast<::arrow::BinaryViewArray>( + input_storage->GetFieldByName("metadata")); + const auto& input_value = checked_pointer_cast<::arrow::BinaryViewArray>( + input_storage->GetFieldByName("value")); + const auto& output_metadata = checked_pointer_cast<::arrow::BinaryViewArray>( + output_storage->GetFieldByName("metadata")); + const auto& output_value = checked_pointer_cast<::arrow::BinaryViewArray>( + output_storage->GetFieldByName("value")); + + ASSERT_EQ(input_metadata->data().get(), output_metadata->data().get()); + ASSERT_EQ(input_value->data().get(), output_value->data().get()); +} + +TEST(TestVariantUnshred, UnshredsValuesWithoutCopyingMetadata) { + VariantBuilder typed_only_row_builder; + typed_only_row_builder.AppendInt64(1); + auto typed_only_row = typed_only_row_builder.Finish(); + VariantBuilder residual_value_row_builder; + residual_value_row_builder.AppendInt8(7); + auto residual_value_row = residual_value_row_builder.Finish(); + + const auto metadata_value = std::string_view{*typed_only_row.metadata}; + ASSERT_EQ(metadata_value, std::string_view{*residual_value_row.metadata}); + + auto input_metadata = BinaryArrayFromValues({metadata_value, metadata_value}); + auto input_values = + BinaryArrayFromValues({std::nullopt, std::string_view{*residual_value_row.value}}); + auto input_typed_values = ArrayFromJSON(int64(), "[1, null]"); + auto input_storage_type = + struct_({field("metadata", binary(), false), field("value", binary()), + field("typed_value", int64())}); + + auto input_array = MakeVariantArrayFromChildren( + input_storage_type, {input_metadata, input_values, input_typed_values}); + auto unshredded = UnshredVariantArray(*input_array); + + const auto& output_storage = + checked_pointer_cast<::arrow::StructArray>(unshredded->storage()); + const auto& output_metadata = output_storage->GetFieldByName("metadata"); + ASSERT_EQ(input_metadata->data().get(), output_metadata->data().get()); + const auto& output_values = checked_pointer_cast<::arrow::BinaryViewArray>( + output_storage->GetFieldByName("value")); + ASSERT_EQ(2, output_values->length()); + ASSERT_EQ(std::string_view{*typed_only_row.value}, output_values->GetView(0)); + ASSERT_EQ(std::string_view{*residual_value_row.value}, output_values->GetView(1)); +} + +TEST(TestVariantUnshred, UnshredsDecimalTypedValues) { + auto assert_decimal = [](std::shared_ptr<::arrow::DataType> type, std::string_view json, + std::invocable auto&& append_expected) { + VariantBuilder expected_builder; + std::invoke(std::forward(append_expected), + expected_builder); + auto expected = expected_builder.Finish(); + auto typed_value = ::arrow::ArrayFromJSON(type, std::string(json)); + auto array = MakeSingleRowVariantArray(std::string_view{*expected.metadata}, + std::nullopt, typed_value); + + AssertUnshreddedRow(*array, expected); + }; + + assert_decimal(::arrow::decimal32(/*precision=*/9, /*scale=*/2), R"(["-1234567.89"])", + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(-123456789), /*scale=*/2); + }); + assert_decimal(::arrow::decimal64(/*precision=*/9, /*scale=*/0), R"(["123456789"])", + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(123456789), /*scale=*/0); + }); + assert_decimal(::arrow::decimal64(/*precision=*/18, /*scale=*/0), + R"(["-123456789012345678"])", [](VariantBuilder& builder) { + builder.AppendDecimal8(::arrow::Decimal64(-123456789012345678), + /*scale=*/0); + }); + assert_decimal(::arrow::decimal128(/*precision=*/9, /*scale=*/0), R"(["-123456789"])", + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(-123456789), /*scale=*/0); + }); + assert_decimal(::arrow::decimal128(/*precision=*/18, /*scale=*/0), + R"(["123456789012345678"])", [](VariantBuilder& builder) { + builder.AppendDecimal8(::arrow::Decimal64(123456789012345678), + /*scale=*/0); + }); + assert_decimal(::arrow::decimal128(/*precision=*/21, /*scale=*/0), + R"(["100000000000000000000"])", [](VariantBuilder& builder) { + builder.AppendDecimal16( + ::arrow::Decimal128(std::string("100000000000000000000")), + /*scale=*/0); + }); +} + +TEST(TestVariantUnshred, MissingTopLevelValue) { + VariantBuilder builder; + builder.AppendVariantNull(); + auto expected = builder.Finish(); + auto array = + MakeSingleRowVariantArray(std::string_view{*expected.metadata}, std::nullopt, + ::arrow::ArrayFromJSON(int64(), "[null]")); + ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); + + ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); + ValidateVariants(chunked); + AssertUnshreddedRow(*array, expected); +} + +TEST(TestVariantUnshred, MissingObjectField) { + auto expected = EmptyObjectWithField("a"); + auto field_group = MakeInt64FieldGroup({std::nullopt}, "[null]"); + auto typed_object_type = struct_({field("a", field_group->type(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN( + auto typed_object, + ::arrow::StructArray::Make({field_group}, typed_object_type->fields())); + auto array = MakeSingleRowVariantArray(std::string_view{*expected.metadata}, + std::nullopt, typed_object); + ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); + + ValidateVariants(chunked); + ValidateVariants(chunked); + AssertUnshreddedRow(*array, expected); +} + +TEST(TestVariantUnshred, NullObjectFieldGroup) { + auto expected = EmptyObjectWithField("a"); + auto field_group = MakeInt64FieldGroup({std::nullopt}, "[null]", {false}); + auto typed_object_type = struct_({field("a", field_group->type(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN( + auto typed_object, + ::arrow::StructArray::Make({field_group}, typed_object_type->fields())); + auto array = MakeSingleRowVariantArray(std::string_view{*expected.metadata}, + std::nullopt, typed_object); + ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); + + ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); + ValidateVariants(chunked); + AssertUnshreddedRow(*array, expected); +} + +TEST(TestVariantUnshred, MissingListElements) { + VariantBuilder builder; + auto list = builder.StartList(); + list.AppendVariantNull(); + list.AppendVariantNull(); + list.Finish(); + auto expected = builder.Finish(); + auto field_group = + MakeInt64FieldGroup({std::nullopt, std::nullopt}, "[null, null]", {true, false}); + auto typed_list_type = + ::arrow::list(field("element", field_group->type(), /*nullable=*/false)); + ASSERT_OK_AND_ASSIGN( + auto typed_list, + ::arrow::ListArray::FromArrays(typed_list_type, + *::arrow::ArrayFromJSON(::arrow::int32(), "[0, 2]"), + *field_group)); + auto array = MakeSingleRowVariantArray(std::string_view{*expected.metadata}, + std::nullopt, typed_list); + ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); + + ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); + ValidateVariants(chunked); + AssertUnshreddedRow(*array, expected); +} + +TEST(TestVariantUnshred, RejectsNullRow) { + VariantArrayBuilder builder; + builder.AppendNull(); + auto array = builder.Finish(); + + ASSERT_THROW(UnshredVariantRow(*array, /*row=*/0), ParquetException); +} + +TEST(TestVariantUnshred, PreservesSliceNullBitmap) { + VariantBuilder builder; + builder.AppendInt64(3); + auto encoded = builder.Finish(); + auto metadata = BinaryArrayFromValues({std::string_view{*encoded.metadata}, + std::string_view{*encoded.metadata}, + std::string_view{*encoded.metadata}}); + auto values = BinaryArrayFromValues({std::nullopt, std::nullopt, std::nullopt}); + auto typed = ArrayFromJSON(int64(), "[1, 2, 3]"); + std::shared_ptr<::arrow::Buffer> null_bitmap; + ::arrow::BitmapFromVector(std::vector{true, false, true}, &null_bitmap); + auto storage_type = struct_({field("metadata", binary(), false), + field("value", binary()), field("typed_value", int64())}); + + auto array = + MakeVariantArrayFromChildren(storage_type, {metadata, values, typed}, null_bitmap); + auto sliced = + checked_pointer_cast<::arrow::extension::VariantArray>(array->Slice(1, 2)); + auto unshredded = UnshredVariantArray(*sliced); + + ASSERT_EQ(2, unshredded->length()); + ASSERT_TRUE(unshredded->IsNull(0)); + ASSERT_FALSE(unshredded->IsNull(1)); + ASSERT_THROW(UnshredVariantRow(*unshredded, /*row=*/0), ParquetException); + auto actual = UnshredVariantRow(*unshredded, /*row=*/1); + ASSERT_EQ(std::string_view{*encoded.metadata}, std::string_view{*actual.metadata}); + ASSERT_EQ(std::string_view{*encoded.value}, std::string_view{*actual.value}); +} + +TEST(TestVariantUnshred, RejectsDuplicateResidualOnUnshred) { + auto encoded = ObjectWithInt64("a", 1); + auto field_group = MakeInt64FieldGroup({std::nullopt}, "[2]"); + auto typed_value_type = struct_({field("a", field_group->type(), /*nullable=*/false)}); + PARQUET_ASSIGN_OR_THROW( + auto typed_value, + ::arrow::StructArray::Make({field_group}, typed_value_type->fields())); + auto array = MakeSingleRowVariantArray(std::string_view{*encoded.metadata}, + std::string_view{*encoded.value}, typed_value); + ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); + + ValidateVariants(chunked); + ASSERT_THROW(UnshredVariantRow(*array, /*row=*/0), + ParquetInvalidOrCorruptedFileException); + ASSERT_THROW(UnshredVariantArray(*array), ParquetInvalidOrCorruptedFileException); +} + +TEST(TestVariantUnshred, DuplicateFields) { + auto encoded = ObjectWithInt64("a", 0); + auto first = MakeInt64FieldGroup({std::nullopt}, "[1]"); + auto second = MakeInt64FieldGroup({std::nullopt}, "[2]"); + auto typed_value_type = struct_({field("a", first->type(), /*nullable=*/false), + field("a", second->type(), /*nullable=*/false)}); + PARQUET_ASSIGN_OR_THROW( + auto typed_value, + ::arrow::StructArray::Make({first, second}, typed_value_type->fields())); + auto array = MakeSingleRowVariantArray(std::string_view{*encoded.metadata}, + std::nullopt, typed_value); + ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); + + ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); + ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); + ASSERT_THROW(UnshredVariantRow(*array, /*row=*/0), + ParquetInvalidOrCorruptedFileException); + ASSERT_THROW(UnshredVariantArray(*array), ParquetInvalidOrCorruptedFileException); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate.cc b/cpp/src/parquet/variant/validate.cc index e9bb5f5f47e1..a3ed1dac3f5e 100644 --- a/cpp/src/parquet/variant/validate.cc +++ b/cpp/src/parquet/variant/validate.cc @@ -19,10 +19,7 @@ #include #include -#include -#include #include -#include #include #include "arrow/array.h" // IWYU pragma: keep @@ -30,15 +27,15 @@ #include "arrow/chunked_array.h" #include "arrow/extension/parquet_variant.h" #include "arrow/extension_type.h" -#include "arrow/status.h" #include "arrow/type.h" -#include "arrow/util/bit_block_counter.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" #include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" #include "parquet/exception.h" -#include "parquet/variant/encoding.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/decoding.h" +#include "parquet/variant/slot_internal.h" namespace parquet::variant { @@ -46,46 +43,54 @@ namespace { using ::arrow::Array; using ::arrow::Buffer; -using ::arrow::DataType; +using ::arrow::ChunkedArray; using ::arrow::ExtensionArray; using ::arrow::ExtensionType; using ::arrow::MemoryPool; -using ::arrow::Status; +using ::arrow::StructArray; using ::arrow::extension::kVariantExtensionName; +using ::arrow::extension::VariantArray; using ::arrow::internal::checked_cast; -enum class VariantShreddedTypeKind { - None, - Primitive, - Object, - Array, -}; - struct VariantValidationPlan { std::shared_ptr array; std::vector children; }; -std::shared_ptr ValuesArray(const Array& array) { - switch (array.type_id()) { - case ::arrow::Type::LIST_VIEW: - return checked_cast(array).values(); - case ::arrow::Type::LARGE_LIST_VIEW: - return checked_cast(array).values(); - case ::arrow::Type::LIST: - return checked_cast(array).values(); - case ::arrow::Type::LARGE_LIST: - return checked_cast(array).values(); - case ::arrow::Type::FIXED_SIZE_LIST: - return checked_cast(array).values(); - case ::arrow::Type::MAP: - return checked_cast(array).values(); - default: - throw ParquetException("Expected list or map storage, got ", - array.type()->ToString()); +template +void ValidateVariantArrayRows(const VariantArray& array, + const std::shared_ptr& valid_rows, + MemoryPool* pool) { + const auto& storage = checked_cast(*array.storage()); + auto metadata_array = storage.GetFieldByName("metadata"); + auto value_array = storage.GetFieldByName("value"); + auto typed_array = storage.GetFieldByName("typed_value"); + + if constexpr (strict) { + if (value_array == nullptr) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant extension storage: missing top-level value field"); + } } + + auto row_plan = internal::CompileVariantRowPlan(value_array, typed_array); + std::string_view last_metadata_bytes; + std::optional last_metadata; + internal::VisitVisibleRows(valid_rows, array, [&](int64_t row) { + if (metadata_array->IsNull(row)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant extension storage: metadata is null"); + } + auto metadata_value = internal::BinaryFieldView(*metadata_array, row); + if (!last_metadata.has_value() || last_metadata_bytes != metadata_value) { + last_metadata = VariantMetadataView::Make(metadata_value); + last_metadata_bytes = metadata_value; + } + internal::ProcessSlot(*last_metadata, row_plan, row); + }); } +template std::optional BuildVariantValidationPlan( std::shared_ptr array) { switch (array->type_id()) { @@ -95,13 +100,13 @@ std::optional BuildVariantValidationPlan( if (ext_type.extension_name() == kVariantExtensionName) { return VariantValidationPlan{.array = std::move(array), .children = {}}; } - return BuildVariantValidationPlan(ext_array.storage()); + return BuildVariantValidationPlan(ext_array.storage()); } case ::arrow::Type::STRUCT: { - const auto& struct_array = checked_cast(*array); + const auto& struct_array = checked_cast(*array); std::vector children; for (auto field : struct_array.fields()) { - auto child_plan = BuildVariantValidationPlan(std::move(field)); + auto child_plan = BuildVariantValidationPlan(std::move(field)); if (child_plan.has_value()) { children.push_back(std::move(*child_plan)); } @@ -118,21 +123,20 @@ std::optional BuildVariantValidationPlan( case ::arrow::Type::LARGE_LIST: case ::arrow::Type::FIXED_SIZE_LIST: case ::arrow::Type::MAP: { - auto values = ValuesArray(*array); - auto child_plan = BuildVariantValidationPlan(std::move(values)); + auto values = internal::ValuesArray(*array); + auto child_plan = BuildVariantValidationPlan(std::move(values)); if (!child_plan.has_value()) { return std::nullopt; } - std::vector children; - children.push_back(std::move(*child_plan)); return VariantValidationPlan{.array = std::move(array), - .children = std::move(children)}; + .children = {std::move(*child_plan)}}; } default: { std::vector children; for (const auto& child_data : array->data()->child_data) { if (child_data != nullptr) { - auto child_plan = BuildVariantValidationPlan(::arrow::MakeArray(child_data)); + auto child_plan = + BuildVariantValidationPlan(::arrow::MakeArray(child_data)); if (child_plan.has_value()) { children.push_back(std::move(*child_plan)); } @@ -147,270 +151,14 @@ std::optional BuildVariantValidationPlan( } } -std::string_view BinaryFieldView(const Array& array, int64_t row) { - switch (array.type_id()) { - case ::arrow::Type::BINARY: - return checked_cast(array).GetView(row); - case ::arrow::Type::LARGE_BINARY: - return checked_cast(array).GetView(row); - case ::arrow::Type::BINARY_VIEW: - return checked_cast(array).GetView(row); - default: - throw ParquetInvalidOrCorruptedFileException("Expected binary Variant field, got ", - array.type()->ToString()); - } -} - -VariantShreddedTypeKind ShreddedTypeKind(const DataType& type) { - switch (type.id()) { - case ::arrow::Type::STRUCT: - return VariantShreddedTypeKind::Object; - case ::arrow::Type::LIST: - case ::arrow::Type::LARGE_LIST: - case ::arrow::Type::LIST_VIEW: - case ::arrow::Type::LARGE_LIST_VIEW: - case ::arrow::Type::FIXED_SIZE_LIST: - return VariantShreddedTypeKind::Array; - default: - return VariantShreddedTypeKind::Primitive; - } -} - -void ValidateVariantShredding( - const VariantMetadataView& metadata, std::string_view value, bool typed_value_present, - VariantShreddedTypeKind typed_value_kind, - std::span shredded_field_names = {}) { - if (typed_value_kind == VariantShreddedTypeKind::None) { - if (value.empty()) { - return; - } - VariantValueView::Validate(value, metadata); - return; - } - - if (typed_value_kind == VariantShreddedTypeKind::Object) { - for (const auto& name : shredded_field_names) { - if (!metadata.FindString(name).has_value()) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: shredded field '", name, - "' is not in metadata dictionary"); - } - } - } - - if (value.empty()) { - return; - } - auto value_view = VariantValueView::Make(value, metadata); - - switch (typed_value_kind) { - case VariantShreddedTypeKind::Primitive: - if (typed_value_present) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: value and primitive typed_value are both " - "non-null"); - } - break; - case VariantShreddedTypeKind::Array: - if (typed_value_present) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: value and array typed_value are both " - "non-null"); - } - if (value_view.basic_type() == VariantBasicType::kArray) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: array value must be stored in typed_value"); - } - break; - case VariantShreddedTypeKind::Object: - default: - if (value_view.basic_type() == VariantBasicType::kObject) { - if (!typed_value_present) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: object value requires object typed_value"); - } - for (const auto& name : shredded_field_names) { - if (std::get(value_view.data()).ContainsField(name)) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: value object contains shredded field: ", name); - } - } - } else if (typed_value_present) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid shredded Variant: partially shredded value must be an object"); - } - break; - } -} - -template -std::pair ValueOffsetAndLength(const Array& array, int64_t row) { - const auto& typed_array = checked_cast(array); - return {typed_array.value_offset(row), typed_array.value_length(row)}; -} - -std::pair ValuesRangeAt(const Array& array, int64_t row) { - switch (array.type_id()) { - case ::arrow::Type::LIST: - return ValueOffsetAndLength<::arrow::ListArray>(array, row); - case ::arrow::Type::LARGE_LIST: - return ValueOffsetAndLength<::arrow::LargeListArray>(array, row); - case ::arrow::Type::LIST_VIEW: - return ValueOffsetAndLength<::arrow::ListViewArray>(array, row); - case ::arrow::Type::LARGE_LIST_VIEW: - return ValueOffsetAndLength<::arrow::LargeListViewArray>(array, row); - case ::arrow::Type::MAP: - return ValueOffsetAndLength<::arrow::MapArray>(array, row); - case ::arrow::Type::FIXED_SIZE_LIST: - return ValueOffsetAndLength<::arrow::FixedSizeListArray>(array, row); - default: - throw ParquetException("Expected list or map storage, got ", - array.type()->ToString()); - } -} - -void ValidateShreddedVariantSlot(const VariantMetadataView& metadata, - const std::shared_ptr& value_array, - const std::shared_ptr& typed_array, int64_t row, - bool allow_missing, std::string_view path); - -void ValidateShreddedVariantField(const VariantMetadataView& metadata, - const Array& field_array, int64_t row, - bool allow_missing, std::string_view path) { - if (field_array.type_id() != ::arrow::Type::STRUCT) { - throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", - path, ": expected struct storage, got ", - field_array.type()->ToString()); - } - if (field_array.IsNull(row)) { - throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", - path, ": field group must be required"); - } - const auto& field_struct = checked_cast(field_array); - auto value_array = field_struct.GetFieldByName("value"); - auto typed_array = field_struct.GetFieldByName("typed_value"); - ValidateShreddedVariantSlot(metadata, value_array, typed_array, row, allow_missing, - path); -} - -void ValidateShreddedVariantSlot(const VariantMetadataView& metadata, - const std::shared_ptr& value_array, - const std::shared_ptr& typed_array, int64_t row, - bool allow_missing, std::string_view path) { - std::string_view value; - if (value_array != nullptr && !value_array->IsNull(row)) { - value = BinaryFieldView(*value_array, row); - if (value.empty()) { - VariantValueView::Validate(value, metadata); - } - } - - const bool typed_value_present = typed_array != nullptr && !typed_array->IsNull(row); - if (value.empty() && !typed_value_present) { - if (allow_missing) { - return; - } - throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant at ", path, - ": value and typed_value are both null"); - } - - const auto typed_kind = typed_array == nullptr ? VariantShreddedTypeKind::None - : ShreddedTypeKind(*typed_array->type()); - std::vector field_names; - if (typed_kind == VariantShreddedTypeKind::Object) { - field_names.reserve(typed_array->type()->num_fields()); - for (const auto& field : typed_array->type()->fields()) { - field_names.emplace_back(field->name()); - } - } - ValidateVariantShredding(metadata, value, typed_value_present, typed_kind, field_names); - - if (!typed_value_present) { - return; - } - - if (typed_kind == VariantShreddedTypeKind::Object) { - const auto& typed_struct = checked_cast(*typed_array); - for (int i = 0; i < typed_struct.struct_type()->num_fields(); ++i) { - ValidateShreddedVariantField(metadata, *typed_struct.field(i), row, - /*allow_missing=*/true, - typed_struct.struct_type()->field(i)->name()); - } - } else if (typed_kind == VariantShreddedTypeKind::Array) { - auto values = ValuesArray(*typed_array); - auto range = ValuesRangeAt(*typed_array, row); - const auto [offset, length] = range; - auto elements = values->Slice(offset, length); - for (int64_t i = 0; i < elements->length(); ++i) { - ValidateShreddedVariantField(metadata, *elements, i, - /*allow_missing=*/false, path); - } - } -} - -template -void VisitVisibleRows(const std::shared_ptr& valid_rows, const Array& array, - VisitVisible&& visit_visible) { - if (valid_rows == nullptr && !array.data()->MayHaveNulls()) { - for (int64_t row = 0; row < array.length(); ++row) { - visit_visible(row); - } - return; - } - PARQUET_THROW_NOT_OK(::arrow::internal::VisitTwoBitBlocks( - valid_rows != nullptr ? valid_rows->data() : nullptr, /*left_offset=*/0, - array.data()->MayHaveNulls() ? array.null_bitmap_data() : nullptr, array.offset(), - array.length(), - [&](int64_t row) { - visit_visible(row); - return Status::OK(); - }, - [] { return Status::OK(); })); -} - -void ValidateVariantExtensionArray(const ExtensionArray& array, - const std::shared_ptr& valid_rows) { - const auto& storage = checked_cast(*array.storage()); - auto metadata_array = storage.GetFieldByName("metadata"); - auto value_array = storage.GetFieldByName("value"); - auto typed_array = storage.GetFieldByName("typed_value"); - - std::string_view last_metadata_bytes; - std::optional last_metadata; - VisitVisibleRows(valid_rows, array, [&](int64_t row) { - if (metadata_array->IsNull(row)) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant extension storage: metadata is null"); - } - - auto metadata_value = BinaryFieldView(*metadata_array, row); - if (!last_metadata.has_value() || last_metadata_bytes != metadata_value) { - last_metadata = VariantMetadataView::Make(metadata_value); - last_metadata_bytes = metadata_value; - } - - if (typed_array == nullptr) { - if (value_array->IsNull(row)) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant extension storage: unshredded value is null"); - } - auto value = BinaryFieldView(*value_array, row); - VariantValueView::Validate(value, *last_metadata); - } else { - ValidateShreddedVariantSlot(*last_metadata, value_array, typed_array, row, - /*allow_missing=*/false, "variant"); - } - }); -} - +template void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, const std::shared_ptr& valid_rows) { switch (plan.array->type_id()) { - case ::arrow::Type::EXTENSION: { - ValidateVariantExtensionArray(checked_cast(*plan.array), - valid_rows); + case ::arrow::Type::EXTENSION: + ValidateVariantArrayRows(checked_cast(*plan.array), + valid_rows, pool); return; - } case ::arrow::Type::STRUCT: { const auto& struct_array = checked_cast(*plan.array); std::shared_ptr child_valid_rows = valid_rows; @@ -422,7 +170,7 @@ void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, struct_array.offset(), struct_array.length(), /*out_offset=*/0)); } for (const auto& child : plan.children) { - ValidateVariantPlan(child, pool, child_valid_rows); + ValidateVariantPlan(child, pool, child_valid_rows); } return; } @@ -433,36 +181,40 @@ void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, case ::arrow::Type::MAP: case ::arrow::Type::FIXED_SIZE_LIST: { DCHECK_EQ(plan.children.size(), 1); - auto values = ValuesArray(*plan.array); - std::shared_ptr values_valid_rows; - PARQUET_ASSIGN_OR_THROW(values_valid_rows, + auto values = internal::ValuesArray(*plan.array); + PARQUET_ASSIGN_OR_THROW(auto values_valid_rows, ::arrow::AllocateEmptyBitmap(values->length(), pool)); - VisitVisibleRows(valid_rows, *plan.array, [&](int64_t row) { - const auto [offset, length] = ValuesRangeAt(*plan.array, row); + internal::VisitVisibleRows(valid_rows, *plan.array, [&](int64_t row) { + const auto [offset, length] = internal::ValuesRangeAt(*plan.array, row); ::arrow::bit_util::SetBitsTo(values_valid_rows->mutable_data(), offset, length, true); }); - ValidateVariantPlan(plan.children[0], pool, values_valid_rows); + ValidateVariantPlan(plan.children[0], pool, values_valid_rows); return; } - default: { + default: for (const auto& child : plan.children) { - ValidateVariantPlan(child, pool, valid_rows); + ValidateVariantPlan(child, pool, valid_rows); } return; - } } } } // namespace +template void ValidateVariants(const ::arrow::ChunkedArray& data, MemoryPool* pool) { for (const auto& chunk : data.chunks()) { - auto plan = BuildVariantValidationPlan(chunk); + auto plan = BuildVariantValidationPlan(chunk); if (plan.has_value()) { - ValidateVariantPlan(*plan, pool, nullptr); + ValidateVariantPlan(*plan, pool, nullptr); } } } +template PARQUET_TEMPLATE_EXPORT void ValidateVariants( + const ::arrow::ChunkedArray& data, MemoryPool* pool); +template PARQUET_TEMPLATE_EXPORT void ValidateVariants( + const ::arrow::ChunkedArray& data, MemoryPool* pool); + } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate.h b/cpp/src/parquet/variant/validate.h index 9abbc7bfd3b6..ae3fbf58f263 100644 --- a/cpp/src/parquet/variant/validate.h +++ b/cpp/src/parquet/variant/validate.h @@ -17,18 +17,17 @@ #pragma once +#include "arrow/type_fwd.h" #include "parquet/platform.h" -namespace arrow { - -class ChunkedArray; -class MemoryPool; - -} // namespace arrow - namespace parquet::variant { -PARQUET_EXPORT -void ValidateVariants(const ::arrow::ChunkedArray& data, ::arrow::MemoryPool* pool); +/// Validate Variant arrays using strict writer rules or read-compatible rules. +/// The non-strict mode accepts invalid encodings retained in parquet-testing for +/// reader compatibility and must not be used to validate values for writing. +template +PARQUET_EXPORT void ValidateVariants( + const ::arrow::ChunkedArray& data, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/validate_test.cc b/cpp/src/parquet/variant/validate_test.cc index 5aa52f537392..44a264ccd1b3 100644 --- a/cpp/src/parquet/variant/validate_test.cc +++ b/cpp/src/parquet/variant/validate_test.cc @@ -17,11 +17,8 @@ #include "parquet/variant/validate.h" -#include #include -#include #include -#include #include "arrow/array.h" // IWYU pragma: keep #include "arrow/chunked_array.h" @@ -42,10 +39,7 @@ using ::arrow::binary; using ::arrow::field; using ::arrow::struct_; using internal::BinaryArrayFromValues; -using internal::Int32ArrayFromValues; using internal::Int8Variant; -using internal::ReadVariantTestingTable; -using internal::ShreddedVariantTestingDir; using internal::VariantTable; using internal::WriteVariantTable; @@ -64,67 +58,24 @@ TEST(TestVariantValidate, ListView) { ::arrow::StructArray::Make({metadata_array, value_array}, storage_type->fields())); auto variant_array = ::arrow::ExtensionType::WrapArray(variant_type, storage); - ASSERT_OK_AND_ASSIGN(auto valid_list, - ::arrow::ListViewArray::FromArrays( - *Int32ArrayFromValues({0}), *Int32ArrayFromValues({1}), - *variant_array, ::arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN( + auto valid_list, + ::arrow::ListViewArray::FromArrays(*::arrow::ArrayFromJSON(::arrow::int32(), "[0]"), + *::arrow::ArrayFromJSON(::arrow::int32(), "[1]"), + *variant_array)); ::arrow::ChunkedArray valid_data{valid_list}; - ValidateVariants(valid_data, ::arrow::default_memory_pool()); + ValidateVariants(valid_data); - ASSERT_OK_AND_ASSIGN(auto invalid_list, - ::arrow::ListViewArray::FromArrays( - *Int32ArrayFromValues({1}), *Int32ArrayFromValues({1}), - *variant_array, ::arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN( + auto invalid_list, + ::arrow::ListViewArray::FromArrays(*::arrow::ArrayFromJSON(::arrow::int32(), "[1]"), + *::arrow::ArrayFromJSON(::arrow::int32(), "[1]"), + *variant_array)); ::arrow::ChunkedArray invalid_data{invalid_list}; - ASSERT_THROW(ValidateVariants(invalid_data, ::arrow::default_memory_pool()), + ASSERT_THROW(ValidateVariants(invalid_data), ParquetInvalidOrCorruptedFileException); } -TEST(TestVariantValidate, ParquetTestingShredded) { - auto maybe_dir = ShreddedVariantTestingDir(); - if (!maybe_dir.has_value()) { - GTEST_SKIP() << "PARQUET_TEST_DATA not set"; - } - if (!std::filesystem::exists(*maybe_dir)) { - GTEST_SKIP() << *maybe_dir << " does not exist"; - } - - auto registered_storage_type = struct_({field("metadata", binary(), /*nullable=*/false), - field("value", binary(), /*nullable=*/false)}); - ::arrow::ExtensionTypeGuard guard(::arrow::extension::variant(registered_storage_type)); - - struct ShreddedCase { - std::string file_name; - bool has_top_level_value; - }; - const std::vector cases = { - {.file_name = "case-041.parquet", .has_top_level_value = false}, - {.file_name = "case-088.parquet", .has_top_level_value = true}, - {.file_name = "case-131.parquet", .has_top_level_value = false}, - {.file_name = "case-132.parquet", .has_top_level_value = true}, - {.file_name = "case-138.parquet", .has_top_level_value = false}, - }; - - for (const auto& test_case : cases) { - SCOPED_TRACE(test_case.file_name); - ASSERT_OK_AND_ASSIGN(auto table, - ReadVariantTestingTable(*maybe_dir + "/" + test_case.file_name)); - ASSERT_OK(table->ValidateFull()); - - auto field = table->schema()->GetFieldByName("var"); - ASSERT_NE(nullptr, field); - auto variant_type = - std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>( - field->type()); - ASSERT_NE(nullptr, variant_type); - ASSERT_EQ(test_case.has_top_level_value, variant_type->value() != nullptr); - - auto column = table->GetColumnByName("var"); - ASSERT_NE(nullptr, column); - ValidateVariants(*column, ::arrow::default_memory_pool()); - } -} - TEST(TestVariantValidate, DictionaryMetadata) { auto encoded = Int8Variant(42); @@ -154,7 +105,7 @@ TEST(TestVariantValidate, DictionaryMetadata) { ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); auto column = read_table->GetColumnByName("variant"); ASSERT_NE(nullptr, column); - ValidateVariants(*column, ::arrow::default_memory_pool()); + ValidateVariants(*column); } TEST(TestVariantValidate, ReadDictionaryOption) { @@ -186,7 +137,7 @@ TEST(TestVariantValidate, ReadDictionaryOption) { ASSERT_OK_AND_ASSIGN(auto read_table, reader->ReadTable()); auto column = read_table->GetColumnByName("variant"); ASSERT_NE(nullptr, column); - ValidateVariants(*column, ::arrow::default_memory_pool()); + ValidateVariants(*column); } } // namespace parquet::variant From 0f4f83ae823a8dddb652a60aa284646c8cb67b95 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Wed, 22 Jul 2026 20:56:52 +0800 Subject: [PATCH 08/10] support shredded --- cpp/src/arrow/extension/CMakeLists.txt | 6 +- cpp/src/arrow/extension/meson.build | 7 +- cpp/src/arrow/extension/parquet_variant.cc | 20 + cpp/src/arrow/extension/parquet_variant.h | 12 + .../arrow/extension/parquet_variant_test.cc | 85 ++ cpp/src/parquet/CMakeLists.txt | 1 + cpp/src/parquet/arrow/variant_test.cc | 131 +++ cpp/src/parquet/meson.build | 1 + cpp/src/parquet/variant/CMakeLists.txt | 1 + cpp/src/parquet/variant/array_internal.cc | 24 + cpp/src/parquet/variant/array_internal.h | 7 + cpp/src/parquet/variant/builder.cc | 11 +- cpp/src/parquet/variant/builder_test.cc | 21 +- cpp/src/parquet/variant/decoding.cc | 61 +- cpp/src/parquet/variant/decoding.h | 11 +- cpp/src/parquet/variant/decoding_test.cc | 20 +- cpp/src/parquet/variant/meson.build | 3 +- .../parquet/variant/parquet_testing_test.cc | 63 +- cpp/src/parquet/variant/shred.cc | 869 ++++++++++++++++++ cpp/src/parquet/variant/shred.h | 50 + cpp/src/parquet/variant/shred_test.cc | 788 ++++++++++++++++ cpp/src/parquet/variant/slot_internal.cc | 12 +- cpp/src/parquet/variant/test_util_internal.cc | 79 +- cpp/src/parquet/variant/test_util_internal.h | 26 +- cpp/src/parquet/variant/unshred.cc | 71 +- cpp/src/parquet/variant/unshred.h | 13 +- cpp/src/parquet/variant/unshred_test.cc | 85 +- cpp/src/parquet/variant/validate.cc | 24 +- 28 files changed, 2233 insertions(+), 269 deletions(-) create mode 100644 cpp/src/arrow/extension/parquet_variant_test.cc create mode 100644 cpp/src/parquet/variant/shred.cc create mode 100644 cpp/src/parquet/variant/shred.h create mode 100644 cpp/src/parquet/variant/shred_test.cc diff --git a/cpp/src/arrow/extension/CMakeLists.txt b/cpp/src/arrow/extension/CMakeLists.txt index ae52bc32a998..1128a03a35d3 100644 --- a/cpp/src/arrow/extension/CMakeLists.txt +++ b/cpp/src/arrow/extension/CMakeLists.txt @@ -15,7 +15,11 @@ # specific language governing permissions and limitations # under the License. -set(CANONICAL_EXTENSION_TESTS bool8_test.cc json_test.cc uuid_test.cc) +set(CANONICAL_EXTENSION_TESTS + bool8_test.cc + json_test.cc + parquet_variant_test.cc + uuid_test.cc) if(ARROW_JSON) list(APPEND CANONICAL_EXTENSION_TESTS tensor_extension_array_test.cc opaque_test.cc) diff --git a/cpp/src/arrow/extension/meson.build b/cpp/src/arrow/extension/meson.build index 84dafe4bbe32..23689eb34087 100644 --- a/cpp/src/arrow/extension/meson.build +++ b/cpp/src/arrow/extension/meson.build @@ -15,7 +15,12 @@ # specific language governing permissions and limitations # under the License. -canonical_extension_tests = ['bool8_test.cc', 'json_test.cc', 'uuid_test.cc'] +canonical_extension_tests = [ + 'bool8_test.cc', + 'json_test.cc', + 'parquet_variant_test.cc', + 'uuid_test.cc', +] if needs_json canonical_extension_tests += [ diff --git a/cpp/src/arrow/extension/parquet_variant.cc b/cpp/src/arrow/extension/parquet_variant.cc index 3d865fbe2ccc..f64c2bfe1207 100644 --- a/cpp/src/arrow/extension/parquet_variant.cc +++ b/cpp/src/arrow/extension/parquet_variant.cc @@ -20,6 +20,7 @@ #include #include +#include "arrow/array/array_nested.h" #include "arrow/extension_type.h" #include "arrow/result.h" #include "arrow/status.h" @@ -182,6 +183,25 @@ bool IsSupportedStorageTypeImpl(const std::shared_ptr& storage_type) { } // namespace +std::shared_ptr VariantArray::metadata() const { + return internal::checked_cast(*storage()) + .GetFieldByName("metadata"); +} + +std::shared_ptr VariantArray::value() const { + return internal::checked_cast(*storage()).GetFieldByName("value"); +} + +std::shared_ptr VariantArray::typed_value() const { + return internal::checked_cast(*storage()) + .GetFieldByName("typed_value"); +} + +bool VariantArray::is_shredded() const { + return internal::checked_cast(*type()).typed_value() != + nullptr; +} + VariantExtensionType::VariantExtensionType(const std::shared_ptr& storage_type) : ExtensionType(storage_type) { // IsSupportedStorageType should have been called already, asserting that diff --git a/cpp/src/arrow/extension/parquet_variant.h b/cpp/src/arrow/extension/parquet_variant.h index c74a794f6a01..75c26c205846 100644 --- a/cpp/src/arrow/extension/parquet_variant.h +++ b/cpp/src/arrow/extension/parquet_variant.h @@ -31,6 +31,18 @@ inline constexpr std::string_view kVariantExtensionName = "arrow.parquet.variant class ARROW_EXPORT VariantArray : public ExtensionArray { public: using ExtensionArray::ExtensionArray; + + /// \brief The metadata child array. + std::shared_ptr metadata() const; + + /// \brief The residual value child array, or null if it is absent from storage. + std::shared_ptr value() const; + + /// \brief The typed value child array, or null for an unshredded Variant array. + std::shared_ptr typed_value() const; + + /// \brief Whether the storage schema contains a typed value field. + bool is_shredded() const; }; /// EXPERIMENTAL: Variant is not yet fully supported. diff --git a/cpp/src/arrow/extension/parquet_variant_test.cc b/cpp/src/arrow/extension/parquet_variant_test.cc new file mode 100644 index 000000000000..a7fdecff69c9 --- /dev/null +++ b/cpp/src/arrow/extension/parquet_variant_test.cc @@ -0,0 +1,85 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/extension/parquet_variant.h" + +#include + +#include + +#include "arrow/array/array_nested.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" + +namespace arrow { + +TEST(TestVariantArray, Accessors) { + auto metadata = ArrayFromJSON(binary(), R"(["metadata"])"); + auto value = ArrayFromJSON(binary(), R"(["value"])"); + auto typed_value = ArrayFromJSON(int64(), "[1]"); + auto storage_type = struct_({field("typed_value", int64()), + field("metadata", binary(), /*nullable=*/false), + field("value", binary())}); + ASSERT_OK_AND_ASSIGN(auto storage, StructArray::Make({typed_value, metadata, value}, + storage_type->fields())); + ASSERT_OK_AND_ASSIGN(auto type, extension::VariantExtensionType::Make(storage_type)); + auto array = std::make_shared(type, storage); + + ASSERT_EQ(metadata->data(), array->metadata()->data()); + ASSERT_EQ(value->data(), array->value()->data()); + ASSERT_EQ(typed_value->data(), array->typed_value()->data()); + ASSERT_TRUE(array->is_shredded()); + + auto unshredded_type = struct_({field("value", binary(), /*nullable=*/false), + field("metadata", binary(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN(auto unshredded_storage, + StructArray::Make({value, metadata}, unshredded_type->fields())); + ASSERT_OK_AND_ASSIGN(auto unshredded_extension_type, + extension::VariantExtensionType::Make(unshredded_type)); + auto unshredded = std::make_shared(unshredded_extension_type, + unshredded_storage); + + ASSERT_EQ(metadata->data(), unshredded->metadata()->data()); + ASSERT_EQ(value->data(), unshredded->value()->data()); + ASSERT_EQ(nullptr, unshredded->typed_value()); + ASSERT_FALSE(unshredded->is_shredded()); +} + +TEST(TestVariantArray, CompatibleStorage) { + auto metadata = ArrayFromJSON(binary(), R"(["metadata"])"); + auto typed_value = ArrayFromJSON(int64(), "[1]"); + auto storage_type = struct_( + {field("typed_value", int64()), field("metadata", binary(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN( + auto storage, StructArray::Make({typed_value, metadata}, storage_type->fields())); + auto prototype_storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("value", binary(), /*nullable=*/false)}); + ASSERT_OK_AND_ASSIGN(auto prototype_type, + extension::VariantExtensionType::Make(prototype_storage_type)); + const auto& prototype = + internal::checked_cast(*prototype_type); + ASSERT_OK_AND_ASSIGN(auto type, prototype.Deserialize(storage_type, "")); + auto array = std::make_shared(type, storage); + + ASSERT_EQ(metadata->data(), array->metadata()->data()); + ASSERT_EQ(nullptr, array->value()); + ASSERT_EQ(typed_value->data(), array->typed_value()->data()); + ASSERT_TRUE(array->is_shredded()); +} + +} // namespace arrow diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index a2d6cfdd8f74..cf626f91af23 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -194,6 +194,7 @@ set(PARQUET_SRCS variant/array_internal.cc variant/builder.cc variant/decoding.cc + variant/shred.cc variant/slot_internal.cc variant/unshred.cc variant/validate.cc diff --git a/cpp/src/parquet/arrow/variant_test.cc b/cpp/src/parquet/arrow/variant_test.cc index 03dda1d12720..d471d2aa4f26 100644 --- a/cpp/src/parquet/arrow/variant_test.cc +++ b/cpp/src/parquet/arrow/variant_test.cc @@ -27,26 +27,55 @@ #include "parquet/arrow/reader.h" #include "parquet/arrow/writer.h" #include "parquet/exception.h" +#include "parquet/variant/array_internal.h" #include "parquet/variant/builder.h" +#include "parquet/variant/shred.h" #include "parquet/variant/test_util_internal.h" +#include "parquet/variant/unshred.h" #include +#include #include #include #include namespace parquet::arrow { +using ::arrow::Array; using ::arrow::binary; using ::arrow::field; using ::arrow::struct_; +using ::arrow::StructArray; +using ::arrow::extension::VariantArray; +using variant::internal::AssertUnshreddedValue; using variant::internal::BinaryArrayFromValues; using variant::internal::EmptyVariantMetadata; using variant::internal::Int8Variant; +using variant::internal::MakeInt64FieldGroup; +using variant::internal::RoundTripVariantArray; using variant::internal::VariantTable; using variant::internal::WriteVariantRecordBatch; using variant::internal::WriteVariantTable; +namespace { + +std::shared_ptr<::arrow::DataType> ExpectedShreddedStorageType( + const std::shared_ptr<::arrow::DataType>& binary_type, + std::string_view list_element_name) { + auto primitive_group_type = + struct_({field("value", binary_type), field("typed_value", ::arrow::int64())}); + auto list_type = ::arrow::list( + field(std::string(list_element_name), primitive_group_type, /*nullable=*/false)); + auto list_group_type = + struct_({field("value", binary_type), field("typed_value", list_type)}); + auto typed_type = struct_({field("count", primitive_group_type, false), + field("items", list_group_type, false)}); + return struct_({field("metadata", binary_type, false), field("value", binary_type), + field("typed_value", typed_type)}); +} + +} // namespace + TEST(TestVariantExtensionType, WriterValidatesUnshreddedVariantBytes) { auto encoded = Int8Variant(42); @@ -355,4 +384,106 @@ TEST(TestVariantExtensionType, WriterValidatesShreddedObjectConflicts) { ASSERT_OK(WriteVariantTable(missing_table)); } +TEST(TestVariantExtensionType, WriterRoundTripsManualShreddedVariant) { + variant::VariantBuilder expected_builder; + auto expected_object = expected_builder.StartObject(); + expected_object.AppendInt64("count", 2); + auto expected_items = expected_object.StartList("items"); + expected_items.AppendInt64(1); + expected_items.AppendInt64(2); + expected_items.Finish(); + expected_object.Finish(); + auto expected = expected_builder.Finish(); + + auto count_group = MakeInt64FieldGroup({std::nullopt}, "[2]"); + auto item_group = MakeInt64FieldGroup({std::nullopt, std::nullopt}, "[1, 2]"); + auto primitive_group_type = count_group->type(); + auto list_type = ::arrow::list(field("element", primitive_group_type, + /*nullable=*/false)); + ASSERT_OK_AND_ASSIGN( + auto items, + ::arrow::ListArray::FromArrays( + list_type, *::arrow::ArrayFromJSON(::arrow::int32(), "[0, 2]"), *item_group)); + auto list_group_type = + struct_({field("value", binary()), field("typed_value", list_type)}); + ASSERT_OK_AND_ASSIGN(auto items_group, + StructArray::Make({BinaryArrayFromValues({std::nullopt}), items}, + list_group_type->fields())); + + auto typed_type = struct_({field("count", primitive_group_type, false), + field("items", list_group_type, false)}); + ASSERT_OK_AND_ASSIGN( + auto typed, StructArray::Make({count_group, items_group}, typed_type->fields())); + auto storage_type = + struct_({field("metadata", binary(), false), field("value", binary()), + field("typed_value", typed_type)}); + auto input = variant::MakeVariantArrayFromChildren( + storage_type, {BinaryArrayFromValues({std::string_view{*expected.metadata}}), + BinaryArrayFromValues({std::nullopt}), typed}); + + ASSERT_OK_AND_ASSIGN(auto actual, RoundTripVariantArray(input)); + ASSERT_TRUE(actual->is_shredded()); + auto expected_storage_type = + ExpectedShreddedStorageType(binary(), /*list_element_name=*/"element"); + ASSERT_TRUE(actual->storage()->type()->Equals(expected_storage_type)) + << "Expected: " << expected_storage_type->ToString() + << "\nActual: " << actual->storage()->type()->ToString(); + ASSERT_TRUE(::arrow::extension::VariantExtensionType::IsSupportedStorageType( + actual->storage()->type())); + AssertUnshreddedValue(*actual, 0, expected); +} + +TEST(TestVariantExtensionType, WriterRoundTripsShredVariantArray) { + variant::VariantArrayBuilder input_builder; + input_builder.AppendNull(); + auto object = input_builder.StartObject(); + object.AppendInt64("count", 2); + auto items = object.StartList("items"); + items.AppendInt64(1); + items.AppendString("residual"); + items.Finish(); + object.AppendBoolean("extra", true); + object.Finish(); + + variant::VariantBuilder list_builder; + list_builder.AddFieldName("count"); + list_builder.AddFieldName("items"); + auto list = list_builder.StartList(); + list.AppendInt64(3); + list.Finish(); + input_builder.AppendEncoded(list_builder.Finish()); + auto input = input_builder.Finish(); + + auto target = struct_({field("count", ::arrow::int64()), + field("items", ::arrow::list(::arrow::int64()))}); + auto shredded = variant::ShredVariantArray(*input, target); + auto expected_shredded_storage_type = + ExpectedShreddedStorageType(::arrow::binary_view(), /*list_element_name=*/"item"); + ASSERT_TRUE(shredded->storage()->type()->Equals(expected_shredded_storage_type)) + << "Expected: " << expected_shredded_storage_type->ToString() + << "\nActual: " << shredded->storage()->type()->ToString(); + ASSERT_TRUE(shredded->is_shredded()); + + ASSERT_OK_AND_ASSIGN(auto actual, RoundTripVariantArray(shredded)); + ASSERT_EQ(input->length(), actual->length()); + auto expected_read_storage_type = + ExpectedShreddedStorageType(binary(), /*list_element_name=*/"element"); + ASSERT_TRUE(actual->storage()->type()->Equals(expected_read_storage_type)) + << "Expected: " << expected_read_storage_type->ToString() + << "\nActual: " << actual->storage()->type()->ToString(); + ASSERT_TRUE(actual->is_shredded()); + + auto unshredded = variant::UnshredVariantArray(*actual); + ASSERT_FALSE(unshredded->is_shredded()); + for (int64_t row = 0; row < actual->length(); ++row) { + ASSERT_EQ(input->IsNull(row), unshredded->IsNull(row)); + if (!input->IsNull(row)) { + ASSERT_EQ(variant::internal::BinaryFieldView(*input->metadata(), row), + variant::internal::BinaryFieldView(*unshredded->metadata(), row)); + ASSERT_EQ(variant::internal::BinaryFieldView(*input->value(), row), + variant::internal::BinaryFieldView(*unshredded->value(), row)); + } + } +} + } // namespace parquet::arrow diff --git a/cpp/src/parquet/meson.build b/cpp/src/parquet/meson.build index 6cbe589c5c67..9452f056438c 100644 --- a/cpp/src/parquet/meson.build +++ b/cpp/src/parquet/meson.build @@ -58,6 +58,7 @@ parquet_srcs = files( 'variant/array_internal.cc', 'variant/builder.cc', 'variant/decoding.cc', + 'variant/shred.cc', 'variant/slot_internal.cc', 'variant/unshred.cc', 'variant/validate.cc', diff --git a/cpp/src/parquet/variant/CMakeLists.txt b/cpp/src/parquet/variant/CMakeLists.txt index 036b5b744b27..c41bc3c6332f 100644 --- a/cpp/src/parquet/variant/CMakeLists.txt +++ b/cpp/src/parquet/variant/CMakeLists.txt @@ -20,6 +20,7 @@ add_parquet_test(variant-test builder_test.cc decoding_test.cc parquet_testing_test.cc + shred_test.cc test_util_internal.cc type_test.cc unshred_test.cc diff --git a/cpp/src/parquet/variant/array_internal.cc b/cpp/src/parquet/variant/array_internal.cc index c71fc31f4024..3a3b22ae1a45 100644 --- a/cpp/src/parquet/variant/array_internal.cc +++ b/cpp/src/parquet/variant/array_internal.cc @@ -19,6 +19,7 @@ #include "arrow/array.h" // IWYU pragma: keep #include "arrow/type.h" +#include "arrow/util/bitmap_ops.h" #include "arrow/util/checked_cast.h" #include "parquet/exception.h" @@ -87,4 +88,27 @@ std::string_view BinaryFieldView(const Array& array, int64_t row) { } } +std::shared_ptr<::arrow::Buffer> FinishNullBitmap( + ::arrow::TypedBufferBuilder& builder) { + if (builder.false_count() == 0) { + return nullptr; + } + PARQUET_ASSIGN_OR_THROW(auto bitmap, builder.Finish()); + return bitmap; +} + +std::shared_ptr<::arrow::Buffer> NullBitmapForOutput(const Array& array, + ::arrow::MemoryPool* pool) { + if (array.null_count() == 0) { + return nullptr; + } + if (array.offset() == 0) { + return array.null_bitmap(); + } + PARQUET_ASSIGN_OR_THROW(auto null_bitmap, + ::arrow::internal::CopyBitmap(pool, array.null_bitmap_data(), + array.offset(), array.length())); + return null_bitmap; +} + } // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/array_internal.h b/cpp/src/parquet/variant/array_internal.h index 26f80471f35a..87f3c1e23b9a 100644 --- a/cpp/src/parquet/variant/array_internal.h +++ b/cpp/src/parquet/variant/array_internal.h @@ -24,6 +24,7 @@ #include "arrow/array.h" // IWYU pragma: export #include "arrow/buffer.h" +#include "arrow/buffer_builder.h" #include "arrow/util/bit_block_counter.h" namespace parquet::variant::internal { @@ -34,6 +35,12 @@ std::pair ValuesRangeAt(const ::arrow::Array& array, int64_t r std::string_view BinaryFieldView(const ::arrow::Array& array, int64_t row); +std::shared_ptr<::arrow::Buffer> FinishNullBitmap( + ::arrow::TypedBufferBuilder& builder); + +std::shared_ptr<::arrow::Buffer> NullBitmapForOutput(const ::arrow::Array& array, + ::arrow::MemoryPool* pool); + template void VisitVisibleRows(const std::shared_ptr<::arrow::Buffer>& valid_rows, const ::arrow::Array& array, VisitVisible&& visit_visible) { diff --git a/cpp/src/parquet/variant/builder.cc b/cpp/src/parquet/variant/builder.cc index 29c516e59f11..7b58409a6cd2 100644 --- a/cpp/src/parquet/variant/builder.cc +++ b/cpp/src/parquet/variant/builder.cc @@ -37,6 +37,7 @@ #include "arrow/util/logging_internal.h" #include "parquet/exception.h" #include "parquet/variant/append_table_internal.h" +#include "parquet/variant/array_internal.h" #include "parquet/variant/binary_view_column_builder_internal.h" #include "parquet/variant/decoding.h" #include "parquet/variant/format_internal.h" @@ -938,10 +939,7 @@ struct VariantArrayBuilder::Impl { DCHECK_EQ(metadata_column.length(), validity_builder.length()); const int64_t null_count = validity_builder.false_count(); - std::shared_ptr null_bitmap; - if (null_count > 0) { - PARQUET_THROW_NOT_OK(validity_builder.Finish(&null_bitmap)); - } + auto null_bitmap = internal::FinishNullBitmap(validity_builder); auto metadata = metadata_column.Finish(); auto value = value_column.Finish(); @@ -1098,10 +1096,7 @@ struct VariantValueArrayBuilder::Impl { std::shared_ptr Finish() { const int64_t null_count = validity_builder.false_count(); - std::shared_ptr null_bitmap; - if (null_count > 0) { - PARQUET_THROW_NOT_OK(validity_builder.Finish(&null_bitmap)); - } + auto null_bitmap = internal::FinishNullBitmap(validity_builder); return value_column.Finish(std::move(null_bitmap), null_count); } diff --git a/cpp/src/parquet/variant/builder_test.cc b/cpp/src/parquet/variant/builder_test.cc index ec0312e9fb41..ef0c31a675d4 100644 --- a/cpp/src/parquet/variant/builder_test.cc +++ b/cpp/src/parquet/variant/builder_test.cc @@ -357,6 +357,7 @@ TEST(TestVariantBuilder, ArrayBuilder) { ASSERT_EQ(8, array->length()); ASSERT_TRUE(array->IsNull(0)); ASSERT_FALSE(array->IsNull(1)); + ASSERT_FALSE(array->is_shredded()); const auto& storage = ::arrow::internal::checked_cast(*array->storage()); @@ -376,12 +377,10 @@ TEST(TestVariantBuilder, Tags) { builder.AppendTimestampNanos(4, false); auto array = builder.Finish(); - const auto& storage = - ::arrow::internal::checked_cast(*array->storage()); const auto& metadata_values = - ::arrow::internal::checked_cast(*storage.field(0)); + ::arrow::internal::checked_cast(*array->metadata()); const auto& values = - ::arrow::internal::checked_cast(*storage.field(1)); + ::arrow::internal::checked_cast(*array->value()); const std::array expected = { VariantPrimitiveType::kBooleanTrue, VariantPrimitiveType::kBooleanFalse, VariantPrimitiveType::kTimestampMicros, VariantPrimitiveType::kTimestampNTZMicros, @@ -406,12 +405,10 @@ TEST(TestVariantBuilder, ArrayBuilderViews) { object.Finish(); auto array = builder.Finish(); - const auto& storage = - ::arrow::internal::checked_cast(*array->storage()); const auto& metadata = - ::arrow::internal::checked_cast(*storage.field(0)); + ::arrow::internal::checked_cast(*array->metadata()); const auto& values = - ::arrow::internal::checked_cast(*storage.field(1)); + ::arrow::internal::checked_cast(*array->value()); auto* metadata_views = metadata.data()->GetValues<::arrow::BinaryViewType::c_type>(1); auto* value_views = values.data()->GetValues<::arrow::BinaryViewType::c_type>(1); @@ -479,12 +476,10 @@ TEST(TestVariantBuilder, SharedMetadata) { std::move(storage_type), {std::move(metadata_values), std::move(values)}); ASSERT_EQ(2, array->length()); - const auto& storage = - ::arrow::internal::checked_cast(*array->storage()); const auto& metadata_array = - ::arrow::internal::checked_cast(*storage.field(0)); + ::arrow::internal::checked_cast(*array->metadata()); const auto& value_array = - ::arrow::internal::checked_cast(*storage.field(1)); + ::arrow::internal::checked_cast(*array->value()); ASSERT_EQ(std::string_view{*metadata_encoded.metadata}, metadata_array.GetView(0)); ASSERT_EQ(metadata_array.GetView(0), metadata_array.GetView(1)); ASSERT_EQ(metadata_encoded.metadata, metadata_array.data()->buffers[2]); @@ -591,6 +586,7 @@ TEST(TestVariantBuilder, FromStorage) { auto array = MakeVariantArrayFromStorage(storage); ASSERT_EQ(1, array->length()); ASSERT_TRUE(array->type()->Equals(variant(storage->type()))); + ASSERT_FALSE(array->is_shredded()); } TEST(TestVariantBuilder, FromShredded) { @@ -608,6 +604,7 @@ TEST(TestVariantBuilder, FromShredded) { auto array = MakeVariantArrayFromChildren(storage_type, {metadata, values, typed}); ASSERT_EQ(2, array->length()); ASSERT_TRUE(array->type()->Equals(variant(storage_type))); + ASSERT_TRUE(array->is_shredded()); } } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/decoding.cc b/cpp/src/parquet/variant/decoding.cc index e6e00eb69c88..f4734d04c878 100644 --- a/cpp/src/parquet/variant/decoding.cc +++ b/cpp/src/parquet/variant/decoding.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include "arrow/util/endian.h" @@ -61,6 +62,24 @@ void CheckAvailable(std::string_view data, size_t offset, size_t size, } } +size_t OffsetCount(uint32_t count) { + const auto converted_count = static_cast(count); + if (converted_count == std::numeric_limits::max()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: offset count overflow"); + } + return converted_count + 1; +} + +size_t OffsetBytes(size_t offset_count, uint8_t offset_size) { + if (offset_count > + std::numeric_limits::max() / static_cast(offset_size)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: offset table size overflow"); + } + return offset_count * static_cast(offset_size); +} + size_t PrimitivePayloadSize(std::string_view value, size_t offset, VariantPrimitiveType primitive) { switch (primitive) { @@ -141,12 +160,12 @@ size_t ParseArray(std::string_view value, const VariantMetadataView& metadata, CheckAvailable(value, offset, count_size, "array size"); const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); offset += count_size; + const size_t offset_count = OffsetCount(num_elements); - CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, - "array offsets"); + CheckAvailable(value, offset, OffsetBytes(offset_count, offset_size), "array offsets"); - std::vector offsets(num_elements + 1); - for (uint32_t i = 0; i <= num_elements; ++i) { + std::vector offsets(offset_count); + for (size_t i = 0; i < offset_count; ++i) { offsets[i] = ReadLittleEndian(value, offset, offset_size); offset += offset_size; } @@ -213,10 +232,11 @@ size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, offset += id_size; } - CheckAvailable(value, offset, (static_cast(num_elements) + 1) * offset_size, + const size_t offset_count = OffsetCount(num_elements); + CheckAvailable(value, offset, OffsetBytes(offset_count, offset_size), "object field offsets"); - std::vector field_offsets(num_elements + 1); - for (uint32_t i = 0; i <= num_elements; ++i) { + std::vector field_offsets(offset_count); + for (size_t i = 0; i < offset_count; ++i) { field_offsets[i] = ReadLittleEndian(value, offset, offset_size); offset += offset_size; } @@ -373,12 +393,12 @@ VariantMetadataView VariantMetadataView::ParsePrefix(std::string_view data, CheckAvailable(data, 1, view.offset_size_, "metadata dictionary size"); const uint32_t dictionary_size = ReadLittleEndian(data, 1, view.offset_size_); const size_t offsets_offset = 1 + view.offset_size_; - CheckAvailable(data, offsets_offset, - (static_cast(dictionary_size) + 1) * view.offset_size_, - "metadata dictionary offsets"); + const size_t offset_count = OffsetCount(dictionary_size); + const size_t offset_bytes = OffsetBytes(offset_count, view.offset_size_); + CheckAvailable(data, offsets_offset, offset_bytes, "metadata dictionary offsets"); - std::vector offsets(dictionary_size + 1); - for (uint32_t i = 0; i <= dictionary_size; ++i) { + std::vector offsets(offset_count); + for (size_t i = 0; i < offset_count; ++i) { offsets[i] = ReadLittleEndian(data, offsets_offset + i * view.offset_size_, view.offset_size_); } @@ -394,8 +414,7 @@ VariantMetadataView VariantMetadataView::ParsePrefix(std::string_view data, } } - const size_t bytes_offset = - offsets_offset + (static_cast(dictionary_size) + 1) * view.offset_size_; + const size_t bytes_offset = offsets_offset + offset_bytes; const size_t bytes_size = offsets[dictionary_size]; CheckAvailable(data, bytes_offset, bytes_size, "metadata dictionary bytes"); const size_t metadata_size = bytes_offset + bytes_size; @@ -463,31 +482,31 @@ std::optional VariantObjectView::GetField(std::string_view nam if (it == fields_.end() || it->name != name) { return std::nullopt; } - return VariantValueView::MakeShallow(it->value, metadata_.get()); + return VariantValueView::Make(it->value, metadata_.get()); } std::optional VariantObjectView::GetField(size_t index) const { if (index >= fields_.size()) { return std::nullopt; } - return VariantValueView::MakeShallow(fields_[index].value, metadata_.get()); + return VariantValueView::Make(fields_[index].value, metadata_.get()); } std::optional VariantArrayView::GetElement(size_t index) const { if (index >= elements_.size()) { return std::nullopt; } - return VariantValueView::MakeShallow(elements_[index], metadata_.get()); + return VariantValueView::Make(elements_[index], metadata_.get()); } VariantValueView VariantValueView::Make(std::string_view value, const VariantMetadataView& metadata) { - return ParseCompleteValue(value, metadata); + return ParseCompleteValue(value, metadata); } -VariantValueView VariantValueView::MakeShallow(std::string_view value, - const VariantMetadataView& metadata) { - return ParseCompleteValue(value, metadata); +VariantValueView VariantValueView::MakeWithValidate(std::string_view value, + const VariantMetadataView& metadata) { + return ParseCompleteValue(value, metadata); } VariantBasicType VariantValueView::PeekBasicType(std::string_view value) { diff --git a/cpp/src/parquet/variant/decoding.h b/cpp/src/parquet/variant/decoding.h index 6ca3f5d21bfc..f05c77f2827f 100644 --- a/cpp/src/parquet/variant/decoding.h +++ b/cpp/src/parquet/variant/decoding.h @@ -153,7 +153,7 @@ class PARQUET_EXPORT VariantValueView { using Data = std::variant; - /// Parse Variant value bytes without copying them. + /// Parse one Variant value node without copying it. /// /// The returned view and any nested object/list/string views are valid only while the /// input value bytes remain alive and unchanged. The metadata view object and its @@ -163,6 +163,12 @@ class PARQUET_EXPORT VariantValueView { static VariantValueView Make(std::string_view value, VariantMetadataView&& metadata) = delete; + /// Parse and recursively validate Variant value bytes without copying them. + static VariantValueView MakeWithValidate(std::string_view value, + const VariantMetadataView& metadata); + static VariantValueView MakeWithValidate(std::string_view value, + VariantMetadataView&& metadata) = delete; + /// Read the basic type from the value header without validating the remaining bytes. static VariantBasicType PeekBasicType(std::string_view value); @@ -182,9 +188,6 @@ class PARQUET_EXPORT VariantValueView { VariantValueView(std::string_view value, Data data) : value_(value), data_(std::move(data)) {} - static VariantValueView MakeShallow(std::string_view value, - const VariantMetadataView& metadata); - std::string_view value_; Data data_; }; diff --git a/cpp/src/parquet/variant/decoding_test.cc b/cpp/src/parquet/variant/decoding_test.cc index 25bb6e71bd4c..303096ae7193 100644 --- a/cpp/src/parquet/variant/decoding_test.cc +++ b/cpp/src/parquet/variant/decoding_test.cc @@ -40,6 +40,15 @@ concept CanMakeVariantValueView = (requires(std::string_view value, Metadata&& m static_assert(CanMakeVariantValueView); static_assert(!CanMakeVariantValueView); +template +concept CanMakeValidatedVariantValueView = + (requires(std::string_view value, Metadata&& metadata) { + VariantValueView::MakeWithValidate(value, std::forward(metadata)); + }); + +static_assert(CanMakeValidatedVariantValueView); +static_assert(!CanMakeValidatedVariantValueView); + TEST(TestVariantEncoding, EmptyMetadata) { VariantBuilder builder; builder.AppendVariantNull(); @@ -222,8 +231,15 @@ TEST(TestVariantEncoding, InvalidValue) { ASSERT_THROW(VariantValueView::Make("\x02\x01\x00\x00\x01\x00"sv, metadata), ParquetInvalidOrCorruptedFileException); // The outer and inner array offsets are valid, but the grandchild primitive tag is - // unknown. Make must recursively validate the entire value before returning. - ASSERT_THROW(VariantValueView::Make("\x03\x01\x00\x05\x03\x01\x00\x01\x54"sv, metadata), + // unknown. Shallow parsing accepts the root and rejects the child when accessed. + constexpr auto invalid_grandchild = "\x03\x01\x00\x05\x03\x01\x00\x01\x54"sv; + auto root = VariantValueView::Make(invalid_grandchild, metadata); + const auto& outer = std::get(root.data()); + auto inner_value = outer.GetElement(0); + ASSERT_TRUE(inner_value.has_value()); + const auto& inner = std::get(inner_value->data()); + ASSERT_THROW(inner.GetElement(0), ParquetInvalidOrCorruptedFileException); + ASSERT_THROW(VariantValueView::MakeWithValidate(invalid_grandchild, metadata), ParquetInvalidOrCorruptedFileException); } diff --git a/cpp/src/parquet/variant/meson.build b/cpp/src/parquet/variant/meson.build index 89fdc8858353..679c2a845e64 100644 --- a/cpp/src/parquet/variant/meson.build +++ b/cpp/src/parquet/variant/meson.build @@ -21,6 +21,7 @@ exc = executable( 'builder_test.cc', 'decoding_test.cc', 'parquet_testing_test.cc', + 'shred_test.cc', 'test_util_internal.cc', 'type_test.cc', 'unshred_test.cc', @@ -31,6 +32,6 @@ exc = executable( test('parquet-variant-test', exc) install_headers( - ['builder.h', 'decoding.h', 'format.h', 'unshred.h', 'validate.h'], + ['builder.h', 'decoding.h', 'format.h', 'shred.h', 'unshred.h', 'validate.h'], subdir: 'parquet/variant', ) diff --git a/cpp/src/parquet/variant/parquet_testing_test.cc b/cpp/src/parquet/variant/parquet_testing_test.cc index e2f1ee534299..9c18c0159d88 100644 --- a/cpp/src/parquet/variant/parquet_testing_test.cc +++ b/cpp/src/parquet/variant/parquet_testing_test.cc @@ -28,17 +28,16 @@ #include "arrow/extension/parquet_variant.h" #include "arrow/extension_type.h" #include "arrow/io/file.h" -#include "arrow/io/memory.h" #include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep #include "arrow/table.h" #include "arrow/testing/extension_type.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/io_util.h" #include "parquet/arrow/reader.h" -#include "parquet/arrow/writer.h" #include "parquet/exception.h" #include "parquet/variant/builder.h" #include "parquet/variant/decoding.h" +#include "parquet/variant/test_util_internal.h" #include "parquet/variant/unshred.h" #include "parquet/variant/validate.h" @@ -48,6 +47,8 @@ namespace parquet::variant { using ::arrow::internal::checked_pointer_cast; +using internal::AssertEncodedRow; +using internal::RoundTripVariantArray; namespace { @@ -136,35 +137,6 @@ class TestVariantParquetTesting : public ::testing::Test { return cases; } - static std::shared_ptr<::arrow::Table> RoundTripVariantArray( - const std::shared_ptr<::arrow::extension::VariantArray>& array) { - auto table = ::arrow::Table::Make( - ::arrow::schema({::arrow::field("variant", array->type())}), {array}); - auto arrow_writer_properties = - ArrowWriterProperties::Builder().set_variant_validation_enabled(true)->build(); - PARQUET_ASSIGN_OR_THROW(auto sink, ::arrow::io::BufferOutputStream::Create()); - PARQUET_THROW_NOT_OK(parquet::arrow::WriteTable( - *table, ::arrow::default_memory_pool(), sink, /*chunk_size=*/table->num_rows(), - default_writer_properties(), std::move(arrow_writer_properties))); - PARQUET_ASSIGN_OR_THROW(auto buffer, sink->Finish()); - - std::optional<::arrow::ExtensionTypeGuard> guard; - if (::arrow::GetExtensionType( - std::string(::arrow::extension::kVariantExtensionName)) == nullptr) { - guard.emplace(array->type()); - } - ArrowReaderProperties reader_properties; - reader_properties.set_arrow_extensions_enabled(true); - reader_properties.set_variant_validation_enabled(true); - parquet::arrow::FileReaderBuilder reader_builder; - PARQUET_THROW_NOT_OK(reader_builder.Open( - std::make_shared<::arrow::io::BufferReader>(std::move(buffer)))); - reader_builder.properties(reader_properties); - PARQUET_ASSIGN_OR_THROW(auto reader, reader_builder.Build()); - PARQUET_ASSIGN_OR_THROW(auto read_table, reader->ReadTable()); - return read_table; - } - void SetUp() override { auto maybe_dir = ParquetTestingSiblingDir("variant"); if (!maybe_dir.has_value()) { @@ -345,22 +317,14 @@ TEST_F(TestVariantParquetTesting, RoundTripsEncoded) { } auto input = builder.Finish(); - std::shared_ptr<::arrow::Table> table; - ASSERT_NO_THROW(table = RoundTripVariantArray(input)); - ASSERT_NE(nullptr, table); - auto column = table->GetColumnByName("variant"); - ASSERT_NE(nullptr, column); - ASSERT_EQ(1, column->num_chunks()); - const auto& array = - checked_pointer_cast<::arrow::extension::VariantArray>(column->chunk(0)); + ASSERT_OK_AND_ASSIGN(auto array, RoundTripVariantArray(input)); ASSERT_EQ(static_cast(cases_.size()), array->length()); + auto unshredded = UnshredVariantArray(*array); for (size_t row = 0; row < cases_.size(); ++row) { const auto& [name, expected] = cases_[row]; SCOPED_TRACE(name); - auto actual = UnshredVariantRow(*array, static_cast(row)); - ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); - ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); + AssertEncodedRow(*unshredded, static_cast(row), expected); } } @@ -376,12 +340,10 @@ TEST_F(TestShreddedVariantParquetTesting, UnshredsArrayWithMissingValueColumn) { ASSERT_NE(nullptr, variant_array); auto unshredded = UnshredVariantArray(*variant_array); - auto actual = UnshredVariantRow(*unshredded, /*row=*/0); ASSERT_EQ(1, test_case.expected_files.size()); ASSERT_TRUE(test_case.expected_files[0].has_value()); auto expected = LoadEncodedVariantBinaryFile(*test_case.expected_files[0]); - ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); - ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); + AssertEncodedRow(*unshredded, /*row=*/0, expected); } TEST_F(TestShreddedVariantParquetTesting, ReadsShreddedParquetTesting) { @@ -410,19 +372,18 @@ TEST_F(TestShreddedVariantParquetTesting, ReadsShreddedParquetTesting) { continue; } + auto unshredded = UnshredVariantArray(*variant_array); ASSERT_EQ(static_cast(test_case.expected_files.size()), - variant_array->length()); - for (int64_t row = 0; row < variant_array->length(); ++row) { + unshredded->length()); + for (int64_t row = 0; row < unshredded->length(); ++row) { SCOPED_TRACE(row); if (!test_case.expected_files[row].has_value()) { - ASSERT_TRUE(variant_array->IsNull(row)); + ASSERT_TRUE(unshredded->IsNull(row)); continue; } auto expected = LoadEncodedVariantBinaryFile(*test_case.expected_files[row]); - auto actual = UnshredVariantRow(*variant_array, row); - ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); - ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); + AssertEncodedRow(*unshredded, row, expected); } } ASSERT_EQ(kSkipUnshreddedCases.size(), skipped_cases); diff --git a/cpp/src/parquet/variant/shred.cc b/cpp/src/parquet/variant/shred.cc new file mode 100644 index 000000000000..4199252baa9d --- /dev/null +++ b/cpp/src/parquet/variant/shred.cc @@ -0,0 +1,869 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/shred.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/array/builder_base.h" +#include "arrow/array/builder_binary.h" +#include "arrow/buffer.h" +#include "arrow/buffer_builder.h" +#include "arrow/compute/cast.h" +#include "arrow/compute/exec.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension/uuid.h" +#include "arrow/extension_type.h" +#include "arrow/scalar.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/bitmap_ops.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" +#include "arrow/util/ubsan.h" +#include "arrow/util/unreachable.h" +#include "parquet/exception.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/decoding.h" + +namespace parquet::variant { + +using ::arrow::Array; +using ::arrow::ArrayBuilder; +using ::arrow::DataType; +using ::arrow::Field; +using ::arrow::MemoryPool; +using ::arrow::Scalar; +using ::arrow::StructArray; +using ::arrow::TimeUnit; +using ::arrow::extension::VariantArray; +using ::arrow::internal::checked_cast; + +namespace { + +struct ShreddedArrayParts { + std::shared_ptr value; + std::shared_ptr typed_value; + std::shared_ptr<::arrow::Buffer> null_bitmap; + int64_t null_count = 0; +}; + +std::shared_ptr FieldGroupType(const std::shared_ptr& typed_type) { + return ::arrow::struct_({::arrow::field("value", ::arrow::binary_view()), + ::arrow::field("typed_value", typed_type)}); +} + +std::shared_ptr MakeFieldGroup(ShreddedArrayParts parts) { + auto type = FieldGroupType(parts.typed_value->type()); + PARQUET_ASSIGN_OR_THROW( + auto out, + StructArray::Make({std::move(parts.value), std::move(parts.typed_value)}, + type->fields(), std::move(parts.null_bitmap), parts.null_count)); + return out; +} + +class ShredNode { + public: + ShredNode(std::shared_ptr typed_type, MemoryPool* pool) + : typed_type_(std::move(typed_type)), residual_(pool), group_validity_(pool) {} + virtual ~ShredNode() = default; + + const std::shared_ptr& typed_type() const { return typed_type_; } + std::shared_ptr field_group_type() const { + return FieldGroupType(typed_type_); + } + + virtual void AppendParentNull() = 0; + virtual void AppendMissing(const VariantMetadataView& metadata) = 0; + virtual void AppendValue(const VariantMetadataView& metadata, + const VariantValueView& value) = 0; + virtual ShreddedArrayParts Finish() = 0; + + protected: + void AppendGroupValidity(bool valid) { + PARQUET_THROW_NOT_OK(group_validity_.Append(valid)); + } + void AppendResidualNull() { residual_.AppendNull(); } + void AppendResidual(const VariantValueView& value) { + residual_.AppendEncodedValue(value.value()); + } + + ShreddedArrayParts FinishParts(std::shared_ptr typed_value) { + const auto null_count = group_validity_.false_count(); + return ShreddedArrayParts{.value = residual_.Finish(), + .typed_value = std::move(typed_value), + .null_bitmap = internal::FinishNullBitmap(group_validity_), + .null_count = null_count}; + } + + std::shared_ptr typed_type_; + VariantValueArrayBuilder residual_; + ::arrow::TypedBufferBuilder group_validity_; +}; + +template +T LoadLittleEndian(std::string_view bytes) { + return ::arrow::bit_util::FromLittleEndian( + ::arrow::util::SafeLoadAs(reinterpret_cast(bytes.data()))); +} + +bool TryAppendString(const DataType& storage_type, ArrayBuilder& builder, + std::string_view value) { + switch (storage_type.id()) { + case ::arrow::Type::STRING: + PARQUET_THROW_NOT_OK(checked_cast<::arrow::StringBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::LARGE_STRING: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::LargeStringBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::STRING_VIEW: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::StringViewBuilder&>(builder).Append(value)); + return true; + default: + return false; + } +} + +bool TryAppendBinary(const DataType& storage_type, ArrayBuilder& builder, + std::string_view value) { + switch (storage_type.id()) { + case ::arrow::Type::BINARY: + PARQUET_THROW_NOT_OK(checked_cast<::arrow::BinaryBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::LARGE_BINARY: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::LargeBinaryBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::BINARY_VIEW: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::BinaryViewBuilder&>(builder).Append(value)); + return true; + default: + return false; + } +} + +bool TryAppendPrimitiveDirect(const VariantValueView& value, const DataType& storage_type, + ArrayBuilder& builder) { + if (const auto* short_string = std::get_if(&value.data())) { + return TryAppendString(storage_type, builder, short_string->string()); + } + const auto* primitive = std::get_if(&value.data()); + if (primitive == nullptr) { + return false; + } + + const auto payload = primitive->payload(); + switch (primitive->type()) { + case VariantPrimitiveType::kBinary: + return TryAppendBinary(storage_type, builder, payload.substr(4)); + case VariantPrimitiveType::kString: + return TryAppendString(storage_type, builder, payload.substr(4)); + default: + return false; + } +} + +std::shared_ptr DecodeSourceScalar(const VariantValueView& value) { + if (const auto* short_string = std::get_if(&value.data())) { + return std::make_shared<::arrow::StringScalar>(std::string(short_string->string())); + } + const auto* primitive = std::get_if(&value.data()); + if (primitive == nullptr) { + return nullptr; + } + + const auto payload = primitive->payload(); + switch (primitive->type()) { + case VariantPrimitiveType::kNull: + return nullptr; + case VariantPrimitiveType::kBooleanTrue: + return std::make_shared<::arrow::BooleanScalar>(true); + case VariantPrimitiveType::kBooleanFalse: + return std::make_shared<::arrow::BooleanScalar>(false); + case VariantPrimitiveType::kInt8: + return std::make_shared<::arrow::Int8Scalar>(LoadLittleEndian(payload)); + case VariantPrimitiveType::kInt16: + return std::make_shared<::arrow::Int16Scalar>(LoadLittleEndian(payload)); + case VariantPrimitiveType::kInt32: + return std::make_shared<::arrow::Int32Scalar>(LoadLittleEndian(payload)); + case VariantPrimitiveType::kInt64: + return std::make_shared<::arrow::Int64Scalar>(LoadLittleEndian(payload)); + case VariantPrimitiveType::kFloat: { + const auto bits = LoadLittleEndian(payload); + return std::make_shared<::arrow::FloatScalar>(std::bit_cast(bits)); + } + case VariantPrimitiveType::kDouble: { + const auto bits = LoadLittleEndian(payload); + return std::make_shared<::arrow::DoubleScalar>(std::bit_cast(bits)); + } + case VariantPrimitiveType::kDecimal4: { + const auto scale = static_cast(payload[0]); + auto type = ::arrow::decimal32(/*precision=*/9, scale); + return std::make_shared<::arrow::Decimal32Scalar>( + ::arrow::Decimal32(LoadLittleEndian(payload.substr(1))), type); + } + case VariantPrimitiveType::kDecimal8: { + const auto scale = static_cast(payload[0]); + auto type = ::arrow::decimal64(/*precision=*/18, scale); + return std::make_shared<::arrow::Decimal64Scalar>( + ::arrow::Decimal64(LoadLittleEndian(payload.substr(1))), type); + } + case VariantPrimitiveType::kDecimal16: { + const auto scale = static_cast(payload[0]); + const auto low = LoadLittleEndian(payload.substr(1)); + const auto high = LoadLittleEndian(payload.substr(9)); + auto type = ::arrow::decimal128(/*precision=*/38, scale); + return std::make_shared<::arrow::Decimal128Scalar>(::arrow::Decimal128(high, low), + type); + } + case VariantPrimitiveType::kDate: + return std::make_shared<::arrow::Date32Scalar>(LoadLittleEndian(payload)); + case VariantPrimitiveType::kTimeNTZMicros: + return std::make_shared<::arrow::Time64Scalar>(LoadLittleEndian(payload), + TimeUnit::MICRO); + case VariantPrimitiveType::kTimestampMicros: + case VariantPrimitiveType::kTimestampNTZMicros: + case VariantPrimitiveType::kTimestampNanos: + case VariantPrimitiveType::kTimestampNTZNanos: { + const bool nanos = primitive->type() == VariantPrimitiveType::kTimestampNanos || + primitive->type() == VariantPrimitiveType::kTimestampNTZNanos; + const bool adjusted = primitive->type() == VariantPrimitiveType::kTimestampMicros || + primitive->type() == VariantPrimitiveType::kTimestampNanos; + const auto unit = nanos ? TimeUnit::NANO : TimeUnit::MICRO; + auto type = adjusted ? ::arrow::timestamp(unit, "UTC") : ::arrow::timestamp(unit); + return std::make_shared<::arrow::TimestampScalar>( + LoadLittleEndian(payload), type); + } + case VariantPrimitiveType::kBinary: + return std::make_shared<::arrow::BinaryScalar>(std::string(payload.substr(4))); + case VariantPrimitiveType::kString: + return std::make_shared<::arrow::StringScalar>(std::string(payload.substr(4))); + case VariantPrimitiveType::kUuid: + return std::make_shared<::arrow::FixedSizeBinaryScalar>( + ::arrow::Buffer::FromString(std::string(payload)), + ::arrow::fixed_size_binary(16)); + } + ::arrow::Unreachable("Unexpected Variant primitive type"); +} + +bool CanAttemptCast(const Scalar& source, const DataType& target) { + const auto source_id = source.type->id(); + switch (target.id()) { + case ::arrow::Type::BOOL: + return source_id == ::arrow::Type::BOOL || ::arrow::is_signed_integer(source_id) || + ::arrow::is_physical_floating(source_id) || + source_id == ::arrow::Type::STRING; + case ::arrow::Type::INT8: + case ::arrow::Type::INT16: + case ::arrow::Type::INT32: + case ::arrow::Type::INT64: + case ::arrow::Type::FLOAT: + case ::arrow::Type::DOUBLE: + return source_id == ::arrow::Type::BOOL || ::arrow::is_signed_integer(source_id) || + ::arrow::is_physical_floating(source_id) || ::arrow::is_decimal(source_id); + case ::arrow::Type::DECIMAL32: + case ::arrow::Type::DECIMAL64: + case ::arrow::Type::DECIMAL128: + return ::arrow::is_signed_integer(source_id) || + ::arrow::is_physical_floating(source_id) || ::arrow::is_decimal(source_id) || + source_id == ::arrow::Type::STRING; + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_STRING: + case ::arrow::Type::STRING_VIEW: + return source_id == ::arrow::Type::STRING; + case ::arrow::Type::BINARY: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::BINARY_VIEW: + return source_id == ::arrow::Type::BINARY; + case ::arrow::Type::DATE32: + return source_id == ::arrow::Type::DATE32; + case ::arrow::Type::TIME64: + return source.type->Equals(target); + case ::arrow::Type::FIXED_SIZE_BINARY: + return source.type->Equals(target); + case ::arrow::Type::TIMESTAMP: { + if (source_id != ::arrow::Type::TIMESTAMP) { + return false; + } + const auto& source_type = checked_cast(*source.type); + const auto& target_type = checked_cast(target); + return source_type.timezone().empty() == target_type.timezone().empty() && + (source_type.unit() == target_type.unit() || + (source_type.unit() == TimeUnit::MICRO && + target_type.unit() == TimeUnit::NANO)); + } + default: + return false; + } +} + +class PrimitiveShredNode : public ShredNode { + public: + PrimitiveShredNode(std::shared_ptr target, ::arrow::compute::ExecContext* ctx, + MemoryPool* pool) + : ShredNode(std::move(target), pool), ctx_(ctx) { + if (typed_type_->id() == ::arrow::Type::EXTENSION) { + storage_type_ = + checked_cast(*typed_type_).storage_type(); + } else { + storage_type_ = typed_type_; + } + PARQUET_ASSIGN_OR_THROW(builder_, ::arrow::MakeBuilder(storage_type_, pool)); + } + + void AppendParentNull() override { + AppendGroupValidity(false); + AppendResidualNull(); + PARQUET_THROW_NOT_OK(builder_->AppendNull()); + } + + void AppendMissing(const VariantMetadataView&) override { + AppendGroupValidity(true); + AppendResidualNull(); + PARQUET_THROW_NOT_OK(builder_->AppendNull()); + } + + void AppendValue(const VariantMetadataView&, const VariantValueView& value) override { + AppendGroupValidity(true); + if (TryAppendPrimitiveDirect(value, *storage_type_, *builder_)) { + AppendResidualNull(); + return; + } + const auto source = DecodeSourceScalar(value); + if (source != nullptr && CanAttemptCast(*source, *storage_type_)) { + auto casted = ::arrow::compute::Cast(::arrow::Datum(source), storage_type_, + ::arrow::compute::CastOptions::Safe(), ctx_); + if (casted.ok()) { + AppendResidualNull(); + PARQUET_THROW_NOT_OK(builder_->AppendScalar(*casted->scalar())); + return; + } + if (!casted.status().IsInvalid()) { + PARQUET_THROW_NOT_OK(casted.status()); + } + } + AppendResidual(value); + PARQUET_THROW_NOT_OK(builder_->AppendNull()); + } + + ShreddedArrayParts Finish() override { + std::shared_ptr typed; + PARQUET_THROW_NOT_OK(builder_->Finish(&typed)); + if (typed_type_->id() == ::arrow::Type::EXTENSION) { + typed = ::arrow::ExtensionType::WrapArray(typed_type_, std::move(typed)); + } + return FinishParts(std::move(typed)); + } + + private: + ::arrow::compute::ExecContext* ctx_; + std::shared_ptr storage_type_; + std::unique_ptr builder_; +}; + +std::unique_ptr CompileShredNode(const std::shared_ptr& target, + ::arrow::compute::ExecContext* ctx, + MemoryPool* pool); + +class ObjectShredNode : public ShredNode { + public: + ObjectShredNode(const ::arrow::StructType& type, ::arrow::compute::ExecContext* ctx, + MemoryPool* pool) + : ShredNode(nullptr, pool), typed_validity_(pool) { + if (type.num_fields() == 0) { + throw ParquetException("Variant shredding object target must not be empty"); + } + + std::vector> fields; + fields.reserve(type.num_fields()); + children_.reserve(type.num_fields()); + requested_field_names_.reserve(type.num_fields()); + for (const auto& field : type.fields()) { + auto child = CompileShredNode(field->type(), ctx, pool); + auto compiled_field = ::arrow::field(field->name(), child->field_group_type(), + /*nullable=*/false); + if (!requested_field_names_.insert(compiled_field->name()).second) { + throw ParquetException("Duplicate Variant shredding field: ", field->name()); + } + fields.push_back(std::move(compiled_field)); + children_.push_back(std::move(child)); + } + typed_type_ = ::arrow::struct_(fields); + } + + void AppendParentNull() override { + AppendGroupValidity(false); + AppendResidualNull(); + AppendTypedNull(); + } + + void AppendMissing(const VariantMetadataView& metadata) override { + ValidateMetadata(metadata); + AppendGroupValidity(true); + AppendResidualNull(); + AppendTypedNull(); + } + + void AppendValue(const VariantMetadataView& metadata, + const VariantValueView& value) override { + ValidateMetadata(metadata); + AppendGroupValidity(true); + const auto* object = std::get_if(&value.data()); + if (object == nullptr) { + AppendResidual(value); + AppendTypedNull(); + return; + } + + const auto& object_fields = object->fields(); + const auto first_residual = std::ranges::find_if( + object_fields, + [&](const auto& field) { return !requested_field_names_.contains(field.name); }); + + if (first_residual == object_fields.end()) { + AppendResidualNull(); + } else { + auto row = residual_.BindMetadata(metadata); + auto residual_object = row.StartObject(); + for (auto it = first_residual; it != object_fields.end(); ++it) { + if (!requested_field_names_.contains(it->name)) { + residual_object.AppendEncodedValue(it->name, it->value); + } + } + residual_object.Finish(); + row.Finish(); + } + + PARQUET_THROW_NOT_OK(typed_validity_.Append(true)); + const auto& fields = checked_cast(*typed_type_).fields(); + DCHECK_EQ(fields.size(), children_.size()); + for (size_t i = 0; i < children_.size(); ++i) { + auto child_value = object->GetField(fields[i]->name()); + if (child_value.has_value()) { + children_[i]->AppendValue(metadata, *child_value); + } else { + children_[i]->AppendMissing(metadata); + } + } + } + + ShreddedArrayParts Finish() override { + std::vector> fields; + fields.reserve(children_.size()); + for (auto& child : children_) { + fields.push_back(MakeFieldGroup(child->Finish())); + } + const auto null_count = typed_validity_.false_count(); + PARQUET_ASSIGN_OR_THROW( + auto typed, + StructArray::Make(std::move(fields), typed_type_->fields(), + internal::FinishNullBitmap(typed_validity_), null_count)); + return FinishParts(std::move(typed)); + } + + private: + void AppendTypedNull() { + PARQUET_THROW_NOT_OK(typed_validity_.Append(false)); + for (auto& child : children_) { + child->AppendParentNull(); + } + } + + void ValidateMetadata(const VariantMetadataView& metadata) { + if (last_validated_metadata_.has_value() && + *last_validated_metadata_ == metadata.metadata()) { + return; + } + const auto& fields = checked_cast(*typed_type_).fields(); + for (const auto& field : fields) { + if (!metadata.FindString(field->name()).has_value()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: shredded field '", field->name(), + "' is not in metadata dictionary"); + } + } + last_validated_metadata_ = metadata.metadata(); + } + + std::unordered_set requested_field_names_; + std::vector> children_; + ::arrow::TypedBufferBuilder typed_validity_; + std::optional last_validated_metadata_; +}; + +template +void AppendIndex(::arrow::TypedBufferBuilder* builder, int64_t value) { + if (value > std::numeric_limits::max()) { + throw ParquetException("Variant array has too many elements"); + } + PARQUET_THROW_NOT_OK(builder->Append(static_cast(value))); +} + +template +std::shared_ptr<::arrow::Buffer> FinishBuffer( + ::arrow::TypedBufferBuilder* builder) { + PARQUET_ASSIGN_OR_THROW(auto buffer, builder->Finish()); + return buffer; +} + +template TypeClass> +class ListLayout; + +template TypeClass> + requires ::arrow::is_var_length_list_type::value +class ListLayout { + public: + using ArrayType = typename ::arrow::TypeTraits::ArrayType; + using Offset = typename TypeClass::offset_type; + + ListLayout(const TypeClass&, MemoryPool* pool) : offsets_(pool) { + AppendIndex(&offsets_, 0); + } + + std::shared_ptr MakeType(std::shared_ptr element_field) const { + return std::make_shared(std::move(element_field)); + } + + void AppendPresent(int64_t, int64_t end) { AppendIndex(&offsets_, end); } + void AppendNull(int64_t child_length) { AppendIndex(&offsets_, child_length); } + + std::shared_ptr FinishArray(const std::shared_ptr& type, + int64_t length, std::shared_ptr values, + std::shared_ptr<::arrow::Buffer> null_bitmap, + int64_t null_count) { + return std::make_shared(type, length, FinishBuffer(&offsets_), + std::move(values), std::move(null_bitmap), + null_count); + } + + private: + ::arrow::TypedBufferBuilder offsets_; +}; + +template TypeClass> + requires ::arrow::is_list_view_type::value +class ListLayout { + public: + using ArrayType = typename ::arrow::TypeTraits::ArrayType; + using Offset = typename TypeClass::offset_type; + + ListLayout(const TypeClass&, MemoryPool* pool) : offsets_(pool), sizes_(pool) {} + + std::shared_ptr MakeType(std::shared_ptr element_field) const { + return std::make_shared(std::move(element_field)); + } + + void AppendPresent(int64_t start, int64_t end) { + AppendIndex(&offsets_, start); + AppendIndex(&sizes_, end - start); + } + + void AppendNull(int64_t child_length) { + AppendIndex(&offsets_, child_length); + AppendIndex(&sizes_, 0); + } + + std::shared_ptr FinishArray(const std::shared_ptr& type, + int64_t length, std::shared_ptr values, + std::shared_ptr<::arrow::Buffer> null_bitmap, + int64_t null_count) { + return std::make_shared(type, length, FinishBuffer(&offsets_), + FinishBuffer(&sizes_), std::move(values), + std::move(null_bitmap), null_count); + } + + private: + ::arrow::TypedBufferBuilder offsets_; + ::arrow::TypedBufferBuilder sizes_; +}; + +template <> +class ListLayout<::arrow::FixedSizeListType> { + public: + using ArrayType = typename ::arrow::TypeTraits<::arrow::FixedSizeListType>::ArrayType; + + ListLayout(const ::arrow::FixedSizeListType& type, MemoryPool*) + : list_size_(type.list_size()) { + if (list_size_ < 0) { + throw ParquetException("Invalid fixed-size list size: ", list_size_); + } + } + + std::shared_ptr MakeType(std::shared_ptr element_field) const { + return std::make_shared<::arrow::FixedSizeListType>(std::move(element_field), + list_size_); + } + + void ValidateLength(size_t length) const { + if (std::cmp_not_equal(length, list_size_)) { + throw ParquetException("Expected fixed-size list of size ", list_size_, ", got ", + length); + } + } + + int32_t list_size() const { return list_size_; } + void AppendPresent(int64_t, int64_t) {} + void AppendNull(int64_t) {} + + std::shared_ptr FinishArray(const std::shared_ptr& type, + int64_t length, std::shared_ptr values, + std::shared_ptr<::arrow::Buffer> null_bitmap, + int64_t null_count) { + return std::make_shared(type, length, std::move(values), + std::move(null_bitmap), null_count); + } + + private: + int32_t list_size_; +}; + +template TypeClass> +class ListShredNode : public ShredNode { + public: + ListShredNode(const TypeClass& target, ::arrow::compute::ExecContext* ctx, + MemoryPool* pool) + : ShredNode(nullptr, pool), typed_validity_(pool), layout_(target, pool) { + child_ = CompileShredNode(target.value_type(), ctx, pool); + const auto& value_field = target.value_field(); + auto element_field = arrow::field(value_field->name(), child_->field_group_type(), + /*nullable=*/false, value_field->metadata()); + typed_type_ = layout_.MakeType(std::move(element_field)); + } + + void AppendParentNull() override { + AppendGroupValidity(false); + AppendResidualNull(); + AppendTypedNull(); + } + + void AppendMissing(const VariantMetadataView&) override { + AppendGroupValidity(true); + AppendResidualNull(); + AppendTypedNull(); + } + + void AppendValue(const VariantMetadataView& metadata, + const VariantValueView& value) override { + const auto* array = std::get_if(&value.data()); + if constexpr (std::same_as) { + if (array != nullptr) { + layout_.ValidateLength(array->elements().size()); + } + } + AppendGroupValidity(true); + if (array == nullptr) { + AppendResidual(value); + AppendTypedNull(); + return; + } + + AppendResidualNull(); + PARQUET_THROW_NOT_OK(typed_validity_.Append(true)); + const auto start = child_length_; + for (size_t i = 0; i < array->elements().size(); ++i) { + auto element = array->GetElement(i); + DCHECK(element.has_value()); + child_->AppendValue(metadata, *element); + ++child_length_; + } + layout_.AppendPresent(start, child_length_); + } + + ShreddedArrayParts Finish() override { + auto values = MakeFieldGroup(child_->Finish()); + const auto length = typed_validity_.length(); + const auto null_count = typed_validity_.false_count(); + auto null_bitmap = internal::FinishNullBitmap(typed_validity_); + auto typed = layout_.FinishArray(typed_type_, length, std::move(values), + std::move(null_bitmap), null_count); + return FinishParts(std::move(typed)); + } + + private: + void AppendTypedNull() { + PARQUET_THROW_NOT_OK(typed_validity_.Append(false)); + if constexpr (std::same_as) { + for (int32_t i = 0; i < layout_.list_size(); ++i) { + child_->AppendParentNull(); + ++child_length_; + } + } + layout_.AppendNull(child_length_); + } + + std::unique_ptr child_; + ::arrow::TypedBufferBuilder typed_validity_; + ListLayout layout_; + int64_t child_length_ = 0; +}; + +void ValidatePrimitiveTarget(const std::shared_ptr& target) { + switch (target->id()) { + case ::arrow::Type::BOOL: + case ::arrow::Type::INT8: + case ::arrow::Type::INT16: + case ::arrow::Type::INT32: + case ::arrow::Type::INT64: + case ::arrow::Type::FLOAT: + case ::arrow::Type::DOUBLE: + case ::arrow::Type::DATE32: + case ::arrow::Type::BINARY: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::BINARY_VIEW: + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_STRING: + case ::arrow::Type::STRING_VIEW: + return; + case ::arrow::Type::DECIMAL32: + case ::arrow::Type::DECIMAL64: + case ::arrow::Type::DECIMAL128: { + const auto& decimal = checked_cast(*target); + if (decimal.scale() < 0 || decimal.scale() > decimal.precision()) { + throw ParquetException("Invalid Variant shredding decimal type: ", + target->ToString()); + } + return; + } + case ::arrow::Type::TIME64: + if (checked_cast(*target).unit() == TimeUnit::MICRO) { + return; + } + break; + case ::arrow::Type::TIMESTAMP: { + const auto unit = checked_cast(*target).unit(); + if (unit == TimeUnit::MICRO || unit == TimeUnit::NANO) { + return; + } + break; + } + case ::arrow::Type::FIXED_SIZE_BINARY: + if (checked_cast(*target).byte_width() == 16) { + return; + } + break; + case ::arrow::Type::EXTENSION: { + const auto& extension = checked_cast(*target); + if (extension.extension_name() == "arrow.uuid" && + ::arrow::extension::UuidType::IsSupportedStorageType( + extension.storage_type())) { + return; + } + break; + } + default: + break; + } + throw ParquetException("Unsupported Variant shredding type: ", target->ToString()); +} + +std::unique_ptr CompileShredNode(const std::shared_ptr& target, + ::arrow::compute::ExecContext* ctx, + MemoryPool* pool) { + if (target == nullptr) { + throw ParquetException("Variant shredding target type must not be null"); + } + switch (target->id()) { + case ::arrow::Type::STRUCT: + return std::make_unique( + checked_cast(*target), ctx, pool); + case ::arrow::Type::LIST: + return std::make_unique>( + checked_cast(*target), ctx, pool); + case ::arrow::Type::LARGE_LIST: + return std::make_unique>( + checked_cast(*target), ctx, pool); + case ::arrow::Type::LIST_VIEW: + return std::make_unique>( + checked_cast(*target), ctx, pool); + case ::arrow::Type::LARGE_LIST_VIEW: + return std::make_unique>( + checked_cast(*target), ctx, pool); + case ::arrow::Type::FIXED_SIZE_LIST: + return std::make_unique>( + checked_cast(*target), ctx, pool); + default: + ValidatePrimitiveTarget(target); + return std::make_unique(target, ctx, pool); + } +} + +} // namespace + +std::shared_ptr ShredVariantArray( + const VariantArray& array, const std::shared_ptr& typed_value_type, + MemoryPool* pool) { + if (array.is_shredded()) { + throw ParquetException("Cannot shred an already shredded Variant array"); + } + auto metadata_array = array.metadata(); + auto value_array = array.value(); + DCHECK_NE(metadata_array, nullptr); + DCHECK_NE(value_array, nullptr); + + ::arrow::compute::ExecContext ctx(pool); + auto root = CompileShredNode(typed_value_type, &ctx, pool); + auto output_type = ::arrow::struct_( + {::arrow::field("metadata", metadata_array->type(), /*nullable=*/false), + ::arrow::field("value", ::arrow::binary_view()), + ::arrow::field("typed_value", root->typed_type())}); + DCHECK(::arrow::extension::VariantExtensionType::IsSupportedStorageType(output_type)); + + std::string_view last_metadata_bytes; + std::optional last_metadata; + for (int64_t row = 0; row < array.length(); ++row) { + if (array.IsNull(row)) { + root->AppendParentNull(); + continue; + } + if (metadata_array->IsNull(row) || value_array->IsNull(row)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid unshredded Variant: visible metadata or value is null"); + } + const auto metadata_bytes = internal::BinaryFieldView(*metadata_array, row); + if (!last_metadata.has_value() || last_metadata_bytes != metadata_bytes) { + last_metadata = VariantMetadataView::Make(metadata_bytes); + last_metadata_bytes = metadata_bytes; + } + auto value = VariantValueView::Make(internal::BinaryFieldView(*value_array, row), + *last_metadata); + root->AppendValue(*last_metadata, value); + } + + auto parts = root->Finish(); + return MakeVariantArrayFromChildren( + std::move(output_type), + {std::move(metadata_array), std::move(parts.value), std::move(parts.typed_value)}, + internal::NullBitmapForOutput(array, pool)); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/shred.h b/cpp/src/parquet/variant/shred.h new file mode 100644 index 000000000000..34b710bffafe --- /dev/null +++ b/cpp/src/parquet/variant/shred.h @@ -0,0 +1,50 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include "arrow/type_fwd.h" +#include "parquet/platform.h" + +namespace arrow::extension { +class VariantArray; +} // namespace arrow::extension + +namespace parquet::variant { + +/// Shred an unshredded Variant array using a logical Arrow target type. +/// +/// Supported targets are Variant primitives, non-empty structs, and list, large-list, +/// list-view, large-list-view, or fixed-size-list types. The logical target is +/// recursively compiled into the required Variant field groups. Values that cannot be +/// safely converted to the requested primitive type remain encoded in the residual +/// value column. Successful conversions are materialized as the target type and may +/// therefore change the encoded Variant representation after unshredding. +/// +/// Input metadata and parent SQL nulls are preserved. The metadata child may be reused, +/// but the output always has a deterministic shredded schema containing value and +/// typed_value, including for empty and all-parent-null inputs. The input must contain +/// metadata and value fields and must not already contain typed_value. Invalid targets, +/// invalid visible Variant data, and failures that cannot be represented as residual +/// values raise a Parquet exception. +PARQUET_EXPORT +std::shared_ptr<::arrow::extension::VariantArray> ShredVariantArray( + const ::arrow::extension::VariantArray& array, + const std::shared_ptr<::arrow::DataType>& typed_value_type, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/shred_test.cc b/cpp/src/parquet/variant/shred_test.cc new file mode 100644 index 000000000000..6853b97baa78 --- /dev/null +++ b/cpp/src/parquet/variant/shred_test.cc @@ -0,0 +1,788 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "parquet/variant/shred.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/chunked_array.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension/uuid.h" +#include "arrow/testing/builder.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/decimal.h" +#include "arrow/util/key_value_metadata.h" +#include "parquet/exception.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/decoding.h" +#include "parquet/variant/test_util_internal.h" +#include "parquet/variant/validate.h" + +namespace parquet::variant { + +namespace { + +using ::arrow::Array; +using ::arrow::BaseListType; +using ::arrow::BinaryViewArray; +using ::arrow::ChunkedArray; +using ::arrow::DataType; +using ::arrow::Field; +using ::arrow::StructArray; +using ::arrow::Type; +using ::arrow::extension::VariantArray; +using ::arrow::internal::checked_cast; +using ::arrow::internal::checked_pointer_cast; +using internal::AssertUnshreddedValue; +using internal::BinaryArrayFromValues; + +std::shared_ptr SingleValue(const EncodedVariantValue& encoded) { + VariantArrayBuilder builder; + builder.AppendEncoded(encoded); + return builder.Finish(); +} + +template AppendInput, + std::invocable AppendExpected> +void AssertConverted(const std::shared_ptr& target, AppendInput append_input, + AppendExpected append_expected) { + VariantBuilder input_builder; + std::invoke(append_input, input_builder); + auto input = input_builder.Finish(); + auto shredded = ShredVariantArray(*SingleValue(input), target); + + ASSERT_TRUE(shredded->value()->IsNull(0)); + ASSERT_FALSE(shredded->typed_value()->IsNull(0)); + ASSERT_TRUE(shredded->is_shredded()); + ValidateVariants(ChunkedArray{shredded}); + + VariantBuilder expected_builder; + std::invoke(append_expected, expected_builder); + AssertUnshreddedValue(*shredded, 0, expected_builder.Finish()); +} + +void AssertFallback(const EncodedVariantValue& input, + const std::shared_ptr& target) { + auto shredded = ShredVariantArray(*SingleValue(input), target); + const auto& residual = checked_cast(*shredded->value()); + ASSERT_EQ(std::string_view{*input.value}, residual.GetView(0)); + ASSERT_TRUE(shredded->typed_value()->IsNull(0)); + ValidateVariants(ChunkedArray{shredded}); + AssertUnshreddedValue(*shredded, 0, input); +} + +template + requires ::arrow::is_var_length_list_type::value +void AssertListOffsets( + const Array& array, + std::initializer_list expected_offsets) { + using ArrayType = typename ::arrow::TypeTraits::ArrayType; + const auto& list = checked_cast(array); + ASSERT_EQ(2, list.data()->buffers.size()); + ASSERT_EQ(expected_offsets.size(), static_cast(list.length() + 1)); + int64_t index = 0; + for (const auto expected : expected_offsets) { + ASSERT_EQ(expected, list.raw_value_offsets()[index++]); + } +} + +template + requires ::arrow::is_list_view_type::value +void AssertListViewDimensions( + const Array& array, + std::initializer_list expected_offsets, + std::initializer_list expected_sizes) { + using ArrayType = typename ::arrow::TypeTraits::ArrayType; + const auto& list = checked_cast(array); + ASSERT_EQ(3, list.data()->buffers.size()); + ASSERT_EQ(expected_offsets.size(), static_cast(list.length())); + ASSERT_EQ(expected_sizes.size(), static_cast(list.length())); + int64_t index = 0; + for (const auto expected : expected_offsets) { + ASSERT_EQ(expected, list.raw_value_offsets()[index++]); + } + index = 0; + for (const auto expected : expected_sizes) { + ASSERT_EQ(expected, list.raw_value_sizes()[index++]); + } +} + +std::vector> ListTargets( + const std::shared_ptr& element) { + return {::arrow::list(element), ::arrow::large_list(element), + ::arrow::list_view(element), ::arrow::large_list_view(element), + ::arrow::fixed_size_list(element, 3)}; +} + +} // namespace + +TEST(TestVariantShred, TargetTypes) { + VariantArrayBuilder input_builder; + input_builder.AppendNull(); + auto input = input_builder.Finish(); + + std::vector> targets = { + ::arrow::boolean(), + ::arrow::int8(), + ::arrow::int16(), + ::arrow::int32(), + ::arrow::int64(), + ::arrow::float32(), + ::arrow::float64(), + ::arrow::decimal32(9, 2), + ::arrow::decimal64(18, 2), + ::arrow::decimal128(38, 2), + ::arrow::date32(), + ::arrow::time64(::arrow::TimeUnit::MICRO), + ::arrow::timestamp(::arrow::TimeUnit::MICRO), + ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC"), + ::arrow::binary(), + ::arrow::large_binary(), + ::arrow::binary_view(), + ::arrow::utf8(), + ::arrow::large_utf8(), + ::arrow::utf8_view(), + ::arrow::fixed_size_binary(16), + ::arrow::extension::uuid(), + ::arrow::struct_({::arrow::field("a", ::arrow::int64())}), + }; + auto element = ::arrow::field("item", ::arrow::int64()); + auto lists = ListTargets(element); + targets.insert(targets.end(), lists.begin(), lists.end()); + + for (const auto& target : targets) { + auto shredded = ShredVariantArray(*input, target); + ASSERT_EQ(1, shredded->length()) << target->ToString(); + ASSERT_TRUE(shredded->IsNull(0)) << target->ToString(); + ASSERT_TRUE(shredded->is_shredded()) << target->ToString(); + ASSERT_TRUE(shredded->metadata()->type()->Equals(input->metadata()->type())); + ValidateVariants(ChunkedArray{shredded}); + } +} + +TEST(TestVariantShred, InvalidTargets) { + VariantArrayBuilder builder; + auto input = builder.Finish(); + std::vector> targets = { + nullptr, + ::arrow::null(), + ::arrow::uint8(), + ::arrow::uint64(), + ::arrow::decimal32(4, -1), + ::arrow::decimal32(4, 5), + ::arrow::time32(::arrow::TimeUnit::MILLI), + ::arrow::time64(::arrow::TimeUnit::NANO), + ::arrow::timestamp(::arrow::TimeUnit::MILLI), + ::arrow::fixed_size_binary(15), + ::arrow::struct_({}), + ::arrow::struct_( + {::arrow::field("a", ::arrow::int64()), ::arrow::field("a", ::arrow::utf8())}), + ::arrow::fixed_size_list(::arrow::int64(), -1), + ::arrow::map(::arrow::utf8(), ::arrow::int64()), + ::arrow::duration(::arrow::TimeUnit::MICRO), + input->type(), + }; + + for (const auto& target : targets) { + ASSERT_THROW(ShredVariantArray(*input, target), ParquetException) + << (target == nullptr ? "null" : target->ToString()); + } +} + +TEST(TestVariantShred, AlreadyShredded) { + VariantArrayBuilder builder; + builder.AppendInt8(1); + auto shredded = ShredVariantArray(*builder.Finish(), ::arrow::int64()); + ASSERT_THROW(ShredVariantArray(*shredded, ::arrow::int64()), ParquetException); +} + +TEST(TestVariantShred, Primitive) { + AssertConverted( + ::arrow::boolean(), [](VariantBuilder& builder) { builder.AppendBoolean(true); }, + [](VariantBuilder& builder) { builder.AppendBoolean(true); }); + AssertConverted( + ::arrow::int32(), [](VariantBuilder& builder) { builder.AppendInt32(7); }, + [](VariantBuilder& builder) { builder.AppendInt32(7); }); + AssertConverted( + ::arrow::float32(), [](VariantBuilder& builder) { builder.AppendFloat(9); }, + [](VariantBuilder& builder) { builder.AppendFloat(9); }); + AssertConverted( + ::arrow::boolean(), [](VariantBuilder& builder) { builder.AppendInt16(1); }, + [](VariantBuilder& builder) { builder.AppendBoolean(true); }); + AssertConverted( + ::arrow::int8(), [](VariantBuilder& builder) { builder.AppendBoolean(true); }, + [](VariantBuilder& builder) { builder.AppendInt8(1); }); + AssertConverted( + ::arrow::int16(), [](VariantBuilder& builder) { builder.AppendInt8(7); }, + [](VariantBuilder& builder) { builder.AppendInt16(7); }); + AssertConverted( + ::arrow::int32(), [](VariantBuilder& builder) { builder.AppendInt16(7); }, + [](VariantBuilder& builder) { builder.AppendInt32(7); }); + AssertConverted( + ::arrow::int64(), [](VariantBuilder& builder) { builder.AppendInt8(7); }, + [](VariantBuilder& builder) { builder.AppendInt64(7); }); + AssertConverted( + ::arrow::float32(), [](VariantBuilder& builder) { builder.AppendInt16(9); }, + [](VariantBuilder& builder) { builder.AppendFloat(9); }); + AssertConverted( + ::arrow::float64(), [](VariantBuilder& builder) { builder.AppendInt16(9); }, + [](VariantBuilder& builder) { builder.AppendDouble(9); }); + AssertConverted( + ::arrow::boolean(), [](VariantBuilder& builder) { builder.AppendString("true"); }, + [](VariantBuilder& builder) { builder.AppendBoolean(true); }); + AssertConverted( + ::arrow::utf8_view(), + [](VariantBuilder& builder) { builder.AppendShortString("hello"); }, + [](VariantBuilder& builder) { builder.AppendShortString("hello"); }); + const std::string long_string(80, 'x'); + AssertConverted( + ::arrow::utf8(), + [&long_string](VariantBuilder& builder) { builder.AppendString(long_string); }, + [&long_string](VariantBuilder& builder) { builder.AppendString(long_string); }); + AssertConverted( + ::arrow::large_utf8(), + [&long_string](VariantBuilder& builder) { builder.AppendString(long_string); }, + [&long_string](VariantBuilder& builder) { builder.AppendString(long_string); }); + AssertConverted( + ::arrow::binary(), [](VariantBuilder& builder) { builder.AppendBinary("abc"); }, + [](VariantBuilder& builder) { builder.AppendBinary("abc"); }); + AssertConverted( + ::arrow::binary_view(), + [](VariantBuilder& builder) { builder.AppendBinary("abc"); }, + [](VariantBuilder& builder) { builder.AppendBinary("abc"); }); + AssertConverted( + ::arrow::large_binary(), + [](VariantBuilder& builder) { builder.AppendBinary("abc"); }, + [](VariantBuilder& builder) { builder.AppendBinary("abc"); }); + AssertConverted( + ::arrow::date32(), [](VariantBuilder& builder) { builder.AppendDate(123); }, + [](VariantBuilder& builder) { builder.AppendDate(123); }); + AssertConverted( + ::arrow::time64(::arrow::TimeUnit::MICRO), + [](VariantBuilder& builder) { builder.AppendTimeNTZMicros(123); }, + [](VariantBuilder& builder) { builder.AppendTimeNTZMicros(123); }); +} + +TEST(TestVariantShred, MixedPrimitive) { + VariantArrayBuilder builder; + builder.AppendInt8(1); + builder.AppendInt16(2); + builder.AppendInt8(3); + builder.AppendDouble(4); + auto shredded = ShredVariantArray(*builder.Finish(), ::arrow::int64()); + for (int64_t row = 0; row < shredded->length(); ++row) { + VariantBuilder expected; + expected.AppendInt64(row + 1); + AssertUnshreddedValue(*shredded, row, expected.Finish()); + } + ValidateVariants(ChunkedArray{shredded}); +} + +TEST(TestVariantShred, PrimitiveFallback) { + VariantBuilder overflow; + overflow.AppendInt64(128); + AssertFallback(overflow.Finish(), ::arrow::int8()); + + VariantBuilder fractional; + fractional.AppendDouble(1.5); + AssertFallback(fractional.Finish(), ::arrow::int64()); + + VariantBuilder nan; + nan.AppendDouble(std::numeric_limits::quiet_NaN()); + AssertFallback(nan.Finish(), ::arrow::int64()); + + VariantBuilder infinity; + infinity.AppendDouble(std::numeric_limits::infinity()); + AssertFallback(infinity.Finish(), ::arrow::int64()); + + VariantBuilder malformed; + malformed.AppendString("not-a-decimal"); + AssertFallback(malformed.Finish(), ::arrow::decimal64(18, 2)); + + VariantBuilder wrong_family; + wrong_family.AppendString("42"); + AssertFallback(wrong_family.Finish(), ::arrow::int64()); + + VariantBuilder null_value; + null_value.AppendVariantNull(); + AssertFallback(null_value.Finish(), ::arrow::int64()); + + VariantBuilder decimal_boolean; + decimal_boolean.AppendDecimal4(::arrow::Decimal32(1), /*scale=*/0); + AssertFallback(decimal_boolean.Finish(), ::arrow::boolean()); + + VariantBuilder boolean_decimal; + boolean_decimal.AppendBoolean(true); + AssertFallback(boolean_decimal.Finish(), ::arrow::decimal32(9, 0)); +} + +TEST(TestVariantShred, Decimal) { + AssertConverted( + ::arrow::decimal32(9, 2), + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(1234), /*scale=*/2); + }, + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(1234), /*scale=*/2); + }); + AssertConverted( + ::arrow::decimal32(9, 2), [](VariantBuilder& builder) { builder.AppendInt16(12); }, + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(1200), /*scale=*/2); + }); + AssertConverted( + ::arrow::decimal64(18, 3), + [](VariantBuilder& builder) { + builder.AppendDecimal4(::arrow::Decimal32(1234), /*scale=*/2); + }, + [](VariantBuilder& builder) { + builder.AppendDecimal8(::arrow::Decimal64(12340), /*scale=*/3); + }); + AssertConverted( + ::arrow::decimal128(38, 2), + [](VariantBuilder& builder) { builder.AppendString("12.34"); }, + [](VariantBuilder& builder) { + builder.AppendDecimal16(::arrow::Decimal128(1234), /*scale=*/2); + }); +} + +TEST(TestVariantShred, DecimalScales) { + VariantArrayBuilder builder; + builder.AppendDecimal4(::arrow::Decimal32(12), /*scale=*/1); + builder.AppendDecimal4(::arrow::Decimal32(123), /*scale=*/2); + builder.AppendDecimal4(::arrow::Decimal32(34), /*scale=*/1); + auto shredded = ShredVariantArray(*builder.Finish(), ::arrow::decimal64(18, 3)); + constexpr std::array expected_values = {1200, 1230, 3400}; + for (int64_t row = 0; row < shredded->length(); ++row) { + VariantBuilder expected; + expected.AppendDecimal8(::arrow::Decimal64(expected_values[row]), /*scale=*/3); + AssertUnshreddedValue(*shredded, row, expected.Finish()); + } + ValidateVariants(ChunkedArray{shredded}); +} + +TEST(TestVariantShred, DecimalFallback) { + VariantBuilder precision_overflow; + precision_overflow.AppendDecimal8(::arrow::Decimal64(1234), /*scale=*/0); + AssertFallback(precision_overflow.Finish(), ::arrow::decimal32(3, 0)); +} + +TEST(TestVariantShred, Temporal) { + AssertConverted( + ::arrow::timestamp(::arrow::TimeUnit::MICRO), + [](VariantBuilder& builder) { builder.AppendTimestampMicros(12, false); }, + [](VariantBuilder& builder) { builder.AppendTimestampMicros(12, false); }); + AssertConverted( + ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC"), + [](VariantBuilder& builder) { builder.AppendTimestampMicros(12, true); }, + [](VariantBuilder& builder) { builder.AppendTimestampNanos(12000, true); }); + + VariantBuilder wrong_timezone; + wrong_timezone.AppendTimestampMicros(12, false); + AssertFallback(wrong_timezone.Finish(), + ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC")); + + VariantBuilder nanos_to_micros; + nanos_to_micros.AppendTimestampNanos(12000, true); + AssertFallback(nanos_to_micros.Finish(), + ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC")); + + VariantBuilder overflow; + overflow.AppendTimestampMicros(std::numeric_limits::max() / 1000 + 1, true); + AssertFallback(overflow.Finish(), ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC")); +} + +TEST(TestVariantShred, Uuid) { + constexpr std::string_view uuid = "0123456789abcdef"; + AssertConverted( + ::arrow::fixed_size_binary(16), + [uuid](VariantBuilder& builder) { builder.AppendUuid(uuid); }, + [uuid](VariantBuilder& builder) { builder.AppendUuid(uuid); }); + AssertConverted( + ::arrow::extension::uuid(), + [uuid](VariantBuilder& builder) { builder.AppendUuid(uuid); }, + [uuid](VariantBuilder& builder) { builder.AppendUuid(uuid); }); +} + +TEST(TestVariantShred, Object) { + VariantArrayBuilder builder; + auto object = builder.StartObject(); + object.AppendInt8("a", 1); + object.AppendString("b", "x"); + object.Finish(); + auto input = builder.Finish(); + + auto target = ::arrow::struct_( + {::arrow::field("a", ::arrow::int64(), /*nullable=*/true, + ::arrow::key_value_metadata({{"ignored", "metadata"}}))}); + auto shredded = ShredVariantArray(*input, target); + const auto& typed = checked_cast(*shredded->typed_value()); + ASSERT_TRUE(typed.type()->field(0)->type()->Equals( + ::arrow::struct_({::arrow::field("value", ::arrow::binary_view()), + ::arrow::field("typed_value", ::arrow::int64())}))); + ASSERT_FALSE(typed.type()->field(0)->nullable()); + ASSERT_EQ(nullptr, typed.type()->field(0)->metadata()); + const auto& field_group = checked_cast(*typed.field(0)); + ASSERT_TRUE(field_group.field(0)->IsNull(0)); + ASSERT_FALSE(field_group.field(1)->IsNull(0)); + + VariantBuilder expected_builder; + auto expected_object = expected_builder.StartObject(); + expected_object.AppendInt64("a", 1); + expected_object.AppendString("b", "x"); + expected_object.Finish(); + ValidateVariants(ChunkedArray{shredded}); + AssertUnshreddedValue(*shredded, 0, expected_builder.Finish()); +} + +TEST(TestVariantShred, ObjectResidual) { + VariantBuilder input_builder; + auto input_object = input_builder.StartObject(); + input_object.AppendString("a", "bad"); + input_object.AppendInt8("extra", 2); + input_object.Finish(); + auto input = input_builder.Finish(); + auto target = ::arrow::struct_({::arrow::field("a", ::arrow::int64())}); + auto shredded = ShredVariantArray(*SingleValue(input), target); + + const auto& typed = checked_cast(*shredded->typed_value()); + const auto& field_group = checked_cast(*typed.field(0)); + ASSERT_FALSE(field_group.field(0)->IsNull(0)); + ASSERT_TRUE(field_group.field(1)->IsNull(0)); + ASSERT_FALSE(shredded->value()->IsNull(0)); + ValidateVariants(ChunkedArray{shredded}); + AssertUnshreddedValue(*shredded, 0, input); + + VariantBuilder missing_builder; + missing_builder.AddFieldName("a"); + auto missing_object = missing_builder.StartObject(); + missing_object.AppendInt8("extra", 2); + missing_object.Finish(); + auto missing = ShredVariantArray(*SingleValue(missing_builder.Finish()), target); + const auto& missing_typed = checked_cast(*missing->typed_value()); + const auto& missing_group = checked_cast(*missing_typed.field(0)); + ASSERT_FALSE(missing_group.IsNull(0)); + ASSERT_TRUE(missing_group.field(0)->IsNull(0)); + ASSERT_TRUE(missing_group.field(1)->IsNull(0)); + + VariantArrayBuilder null_builder; + auto null_object = null_builder.StartObject(); + null_object.AppendVariantNull("a"); + null_object.Finish(); + auto explicit_null = ShredVariantArray(*null_builder.Finish(), target); + const auto& null_typed = + checked_cast(*explicit_null->typed_value()); + const auto& null_group = checked_cast(*null_typed.field(0)); + ASSERT_FALSE(null_group.field(0)->IsNull(0)); + ASSERT_TRUE(null_group.field(1)->IsNull(0)); + VariantBuilder null_value_builder; + null_value_builder.AppendVariantNull(); + ASSERT_EQ(std::string_view{*null_value_builder.Finish().value}, + checked_cast(*null_group.field(0)).GetView(0)); + + VariantBuilder non_object_builder; + non_object_builder.AddFieldName("a"); + non_object_builder.AppendInt8(1); + AssertFallback(non_object_builder.Finish(), target); +} + +TEST(TestVariantShred, MetadataNames) { + VariantArrayBuilder builder; + builder.AppendInt8(1); + auto target = ::arrow::struct_({::arrow::field("missing", ::arrow::int64())}); + ASSERT_THROW(ShredVariantArray(*builder.Finish(), target), + ParquetInvalidOrCorruptedFileException); + + VariantBuilder nested_builder; + nested_builder.AddFieldName("a"); + auto object = nested_builder.StartObject(); + object.Finish(); + auto nested_target = ::arrow::struct_( + {::arrow::field("a", ::arrow::struct_({::arrow::field("b", ::arrow::int64())}))}); + ASSERT_THROW(ShredVariantArray(*SingleValue(nested_builder.Finish()), nested_target), + ParquetInvalidOrCorruptedFileException); +} + +TEST(TestVariantShred, ListKinds) { + VariantArrayBuilder builder; + auto list = builder.StartList(); + list.AppendInt8(1); + list.AppendString("bad"); + list.AppendVariantNull(); + list.Finish(); + builder.AppendInt8(7); + builder.AppendNull(); + auto input = builder.Finish(); + + auto element_metadata = ::arrow::key_value_metadata({{"key", "value"}}); + auto element = + ::arrow::field("element", ::arrow::int64(), /*nullable=*/true, element_metadata); + for (const auto& target : ListTargets(element)) { + auto shredded = ShredVariantArray(*input, target); + const auto& typed = *shredded->typed_value(); + ASSERT_EQ(target->id(), typed.type_id()); + const auto& typed_list_type = checked_cast(*typed.type()); + ASSERT_EQ("element", typed_list_type.value_field()->name()); + ASSERT_FALSE(typed_list_type.value_field()->nullable()); + ASSERT_TRUE(typed_list_type.value_field()->metadata()->Equals(*element_metadata)); + ASSERT_FALSE(typed.IsNull(0)); + ASSERT_TRUE(typed.IsNull(1)); + ASSERT_TRUE(typed.IsNull(2)); + + switch (target->id()) { + case Type::LIST: + AssertListOffsets<::arrow::ListType>(typed, {0, 3, 3, 3}); + break; + case Type::LARGE_LIST: + AssertListOffsets<::arrow::LargeListType>(typed, {0, 3, 3, 3}); + break; + case Type::LIST_VIEW: + AssertListViewDimensions<::arrow::ListViewType>(typed, {0, 3, 3}, {3, 0, 0}); + break; + case Type::LARGE_LIST_VIEW: + AssertListViewDimensions<::arrow::LargeListViewType>(typed, {0, 3, 3}, {3, 0, 0}); + break; + case Type::FIXED_SIZE_LIST: + ASSERT_EQ(1, typed.data()->buffers.size()); + ASSERT_EQ(9, internal::ValuesArray(typed)->length()); + break; + default: + FAIL() << "Expected list-like target"; + } + + const auto& values = checked_cast(*internal::ValuesArray(typed)); + const auto& residual = checked_cast(*values.field(0)); + ASSERT_TRUE(residual.IsNull(0)); + ASSERT_FALSE(residual.IsNull(1)); + ASSERT_FALSE(residual.IsNull(2)); + ASSERT_FALSE(values.field(1)->IsNull(0)); + ASSERT_TRUE(values.field(1)->IsNull(1)); + ASSERT_TRUE(values.field(1)->IsNull(2)); + VariantBuilder null_value_builder; + null_value_builder.AppendVariantNull(); + ASSERT_EQ(std::string_view{*null_value_builder.Finish().value}, residual.GetView(2)); + + const auto& root_residual = checked_cast(*shredded->value()); + VariantBuilder scalar_builder; + scalar_builder.AppendInt8(7); + auto scalar = scalar_builder.Finish(); + ASSERT_EQ(std::string_view{*scalar.value}, root_residual.GetView(1)); + ASSERT_TRUE(shredded->IsNull(2)); + + VariantBuilder expected_builder; + auto expected_list = expected_builder.StartList(); + expected_list.AppendInt64(1); + expected_list.AppendString("bad"); + expected_list.AppendVariantNull(); + expected_list.Finish(); + ValidateVariants(ChunkedArray{shredded}); + AssertUnshreddedValue(*shredded, 0, expected_builder.Finish()); + AssertUnshreddedValue(*shredded, 1, scalar); + } +} + +TEST(TestVariantShred, ListResidual) { + VariantBuilder scalar; + scalar.AppendInt8(1); + auto scalar_input = scalar.Finish(); + auto target = ::arrow::list(::arrow::int64()); + AssertFallback(scalar_input, target); + + VariantArrayBuilder builder; + auto list = builder.StartList(); + list.AppendInt8(1); + list.AppendString("bad"); + list.Finish(); + auto shredded = ShredVariantArray(*builder.Finish(), target); + const auto& typed = *shredded->typed_value(); + const auto& values = checked_cast(*internal::ValuesArray(typed)); + const auto& residual = checked_cast(*values.field(0)); + + VariantBuilder bad_element; + bad_element.AppendString("bad"); + ASSERT_EQ(std::string_view{*bad_element.Finish().value}, residual.GetView(1)); +} + +TEST(TestVariantShred, Nested) { + VariantArrayBuilder builder; + auto root = builder.StartObject(); + auto items = root.StartList("items"); + auto item = items.StartObject(); + item.AppendInt8("x", 3); + item.Finish(); + items.Finish(); + root.Finish(); + auto input = builder.Finish(); + + auto target = ::arrow::struct_({::arrow::field( + "items", + ::arrow::list(::arrow::struct_({::arrow::field("x", ::arrow::int64())})))}); + auto shredded = ShredVariantArray(*input, target); + + VariantBuilder expected_builder; + auto expected_root = expected_builder.StartObject(); + auto expected_items = expected_root.StartList("items"); + auto expected_item = expected_items.StartObject(); + expected_item.AppendInt64("x", 3); + expected_item.Finish(); + expected_items.Finish(); + expected_root.Finish(); + ValidateVariants(ChunkedArray{shredded}); + AssertUnshreddedValue(*shredded, 0, expected_builder.Finish()); +} + +TEST(TestVariantShred, FixedSize) { + VariantArrayBuilder builder; + auto list = builder.StartList(); + list.AppendInt8(1); + list.AppendInt8(2); + list.Finish(); + ASSERT_THROW( + ShredVariantArray(*builder.Finish(), + ::arrow::fixed_size_list(::arrow::field("item", ::arrow::int64()), + /*list_size=*/3)), + ParquetException); +} + +TEST(TestVariantShred, NullsAndSlices) { + VariantArrayBuilder builder; + builder.AppendInt8(0); + builder.AppendNull(); + builder.AppendVariantNull(); + builder.AppendInt16(2); + auto full = builder.Finish(); + auto input = checked_pointer_cast(full->Slice(1, 3)); + auto shredded = ShredVariantArray(*input, ::arrow::int64()); + + ASSERT_EQ(3, shredded->length()); + ASSERT_TRUE(shredded->IsNull(0)); + ASSERT_FALSE(shredded->IsNull(1)); + ASSERT_FALSE(shredded->IsNull(2)); + ASSERT_EQ(0, shredded->offset()); + ASSERT_EQ(input->metadata()->data().get(), shredded->metadata()->data().get()); + + VariantBuilder null_expected; + null_expected.AppendVariantNull(); + AssertUnshreddedValue(*shredded, 1, null_expected.Finish()); + VariantBuilder int_expected; + int_expected.AppendInt64(2); + AssertUnshreddedValue(*shredded, 2, int_expected.Finish()); + ValidateVariants(ChunkedArray{shredded}); +} + +TEST(TestVariantShred, PerRowMetadata) { + VariantArrayBuilder builder; + auto first = builder.StartObject(); + first.AppendInt8("a", 1); + first.AppendString("first", "x"); + first.Finish(); + auto second = builder.StartObject(); + second.AppendInt16("a", 2); + second.AppendBoolean("second", true); + second.Finish(); + auto input = builder.Finish(); + auto shredded = ShredVariantArray( + *input, ::arrow::struct_({::arrow::field("a", ::arrow::int64())})); + + ASSERT_EQ(input->metadata()->data().get(), shredded->metadata()->data().get()); + const auto& input_metadata = checked_cast(*input->metadata()); + const auto& output_metadata = + checked_cast(*shredded->metadata()); + for (int64_t row = 0; row < input->length(); ++row) { + ASSERT_EQ(input_metadata.GetView(row), output_metadata.GetView(row)); + } + + VariantBuilder first_expected_builder; + auto first_expected = first_expected_builder.StartObject(); + first_expected.AppendInt64("a", 1); + first_expected.AppendString("first", "x"); + first_expected.Finish(); + AssertUnshreddedValue(*shredded, 0, first_expected_builder.Finish()); + VariantBuilder second_expected_builder; + auto second_expected = second_expected_builder.StartObject(); + second_expected.AppendInt64("a", 2); + second_expected.AppendBoolean("second", true); + second_expected.Finish(); + AssertUnshreddedValue(*shredded, 1, second_expected_builder.Finish()); + for (int64_t row = 0; row < input->length(); ++row) { + VariantMetadataView::Make(output_metadata.GetView(row)); + } + ValidateVariants(ChunkedArray{shredded}); +} + +TEST(TestVariantShred, Empty) { + VariantArrayBuilder empty_builder; + auto empty = ShredVariantArray( + *empty_builder.Finish(), ::arrow::struct_({::arrow::field("a", ::arrow::int64())})); + ASSERT_EQ(0, empty->length()); + ASSERT_TRUE(empty->is_shredded()); + ValidateVariants(ChunkedArray{empty}); +} + +TEST(TestVariantShred, AllNull) { + VariantArrayBuilder null_builder; + null_builder.AppendNull(); + null_builder.AppendNull(); + auto input = null_builder.Finish(); + auto shredded = ShredVariantArray(*input, ::arrow::list(::arrow::int64())); + ASSERT_EQ(2, shredded->length()); + ASSERT_TRUE(shredded->IsNull(0)); + ASSERT_TRUE(shredded->IsNull(1)); + ASSERT_TRUE(shredded->is_shredded()); + ASSERT_EQ(input->metadata()->data().get(), shredded->metadata()->data().get()); + ASSERT_NE(input->storage().get(), shredded->storage().get()); + ValidateVariants(ChunkedArray{shredded}); +} + +TEST(TestVariantShred, HiddenNullBytes) { + VariantBuilder valid_builder; + valid_builder.AppendInt8(1); + auto valid = valid_builder.Finish(); + auto metadata = BinaryArrayFromValues( + {std::string_view("\xff", 1), std::string_view{*valid.metadata}}); + auto values = BinaryArrayFromValues( + {std::string_view("\xff", 1), std::string_view{*valid.value}}); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), /*nullable=*/false), + ::arrow::field("value", ::arrow::binary(), /*nullable=*/false)}); + std::shared_ptr<::arrow::Buffer> null_bitmap; + ::arrow::BitmapFromVector(std::vector{false, true}, &null_bitmap); + auto hidden = + MakeVariantArrayFromChildren(storage_type, {metadata, values}, null_bitmap); + + auto shredded = ShredVariantArray(*hidden, ::arrow::int64()); + ASSERT_TRUE(shredded->IsNull(0)); + ASSERT_FALSE(shredded->IsNull(1)); + ValidateVariants(ChunkedArray{shredded}); + + auto visible = MakeVariantArrayFromChildren(storage_type, {metadata, values}); + ASSERT_THROW(ShredVariantArray(*visible, ::arrow::int64()), + ParquetInvalidOrCorruptedFileException); +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/slot_internal.cc b/cpp/src/parquet/variant/slot_internal.cc index 56acae16e259..abd334dabe8f 100644 --- a/cpp/src/parquet/variant/slot_internal.cc +++ b/cpp/src/parquet/variant/slot_internal.cc @@ -575,8 +575,10 @@ void ProcessTypedObjectSlot(const VariantMetadataView& metadata, throw ParquetInvalidOrCorruptedFileException( "Expected object in value field for partially shredded struct"); } - if (strict || target != nullptr) { + if (target != nullptr) { object_value_view = VariantValueView::Make(*value, metadata); + } else if constexpr (strict) { + object_value_view = VariantValueView::MakeWithValidate(*value, metadata); } else { VariantValueView::Validate(*value, metadata); } @@ -613,7 +615,9 @@ bool TryProcessSlot(const VariantMetadataView& metadata, if (!value.has_value()) { return false; } - VariantValueView::Validate(*value, metadata); + if (target == nullptr) { + VariantValueView::Validate(*value, metadata); + } AppendEncodedValue(target, *value); return true; } @@ -632,7 +636,9 @@ bool TryProcessSlot(const VariantMetadataView& metadata, if (!value.has_value()) { return false; } - VariantValueView::Validate(*value, metadata); + if (target == nullptr) { + VariantValueView::Validate(*value, metadata); + } if constexpr (std::is_same_v || std::is_same_v) { const auto basic_type = VariantValueView::PeekBasicType(*value); diff --git a/cpp/src/parquet/variant/test_util_internal.cc b/cpp/src/parquet/variant/test_util_internal.cc index de95aabdb216..a612e6023cb5 100644 --- a/cpp/src/parquet/variant/test_util_internal.cc +++ b/cpp/src/parquet/variant/test_util_internal.cc @@ -17,17 +17,26 @@ #include "parquet/variant/test_util_internal.h" +#include +#include #include #include "arrow/array.h" // IWYU pragma: keep #include "arrow/array/builder_binary.h" +#include "arrow/extension/parquet_variant.h" #include "arrow/extension_type.h" #include "arrow/io/memory.h" #include "arrow/table.h" +#include "arrow/testing/builder.h" +#include "arrow/testing/extension_type.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "parquet/arrow/reader.h" #include "parquet/arrow/writer.h" #include "parquet/exception.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/unshred.h" namespace parquet::variant::internal { @@ -55,9 +64,8 @@ Result> WriteVariantTable( return sink->Finish(); } -::arrow::Status WriteVariantRecordBatch( - const std::shared_ptr<::arrow::Table>& table, - std::shared_ptr arrow_properties) { +Status WriteVariantRecordBatch(const std::shared_ptr<::arrow::Table>& table, + std::shared_ptr arrow_properties) { ARROW_ASSIGN_OR_RAISE(auto sink, ::arrow::io::BufferOutputStream::Create( 1024, ::arrow::default_memory_pool())); ARROW_ASSIGN_OR_RAISE(auto writer, @@ -69,6 +77,39 @@ ::arrow::Status WriteVariantRecordBatch( return writer->Close(); } +Result> RoundTripVariantArray( + const std::shared_ptr<::arrow::extension::VariantArray>& array) { + auto table = ::arrow::Table::Make( + ::arrow::schema({::arrow::field("variant", array->type())}), {array}); + auto arrow_properties = + ArrowWriterProperties::Builder().set_variant_validation_enabled(true)->build(); + ARROW_ASSIGN_OR_RAISE(auto buffer, WriteVariantTable(table, default_writer_properties(), + arrow_properties)); + + std::optional<::arrow::ExtensionTypeGuard> guard; + if (::arrow::GetExtensionType(std::string(::arrow::extension::kVariantExtensionName)) == + nullptr) { + guard.emplace(array->type()); + } + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + reader_properties.set_variant_validation_enabled(true); + parquet::arrow::FileReaderBuilder builder; + RETURN_NOT_OK( + builder.Open(std::make_shared<::arrow::io::BufferReader>(std::move(buffer)))); + builder.properties(reader_properties); + ARROW_ASSIGN_OR_RAISE(auto reader, builder.Build()); + ARROW_ASSIGN_OR_RAISE(auto read_table, reader->ReadTable()); + RETURN_NOT_OK(read_table->ValidateFull()); + + auto column = read_table->GetColumnByName("variant"); + if (column == nullptr || column->num_chunks() != 1) { + return Status::Invalid("Expected one Variant output chunk"); + } + return ::arrow::internal::checked_pointer_cast<::arrow::extension::VariantArray>( + column->chunk(0)); +} + std::shared_ptr EmptyVariantMetadata() { VariantBuilder builder; builder.AppendVariantNull(); @@ -93,4 +134,36 @@ std::shared_ptr BinaryArrayFromValues( return out; } +std::shared_ptr<::arrow::StructArray> MakeInt64FieldGroup( + const std::vector>& values, + std::string_view typed_values, const std::vector& is_valid) { + auto field_group_type = + ::arrow::struct_({::arrow::field("value", ::arrow::binary()), + ::arrow::field("typed_value", ::arrow::int64())}); + std::shared_ptr<::arrow::Buffer> null_bitmap; + if (!is_valid.empty()) { + ::arrow::BitmapFromVector(is_valid, &null_bitmap); + } + PARQUET_ASSIGN_OR_THROW( + auto field_group, + ::arrow::StructArray::Make( + {BinaryArrayFromValues(values), + ::arrow::ArrayFromJSON(::arrow::int64(), std::string(typed_values))}, + field_group_type->fields(), std::move(null_bitmap))); + return field_group; +} + +void AssertEncodedRow(const ::arrow::extension::VariantArray& array, int64_t row, + const EncodedVariantValue& expected) { + ASSERT_EQ(std::string_view{*expected.metadata}, + BinaryFieldView(*array.metadata(), row)); + ASSERT_EQ(std::string_view{*expected.value}, BinaryFieldView(*array.value(), row)); +} + +void AssertUnshreddedValue(const ::arrow::extension::VariantArray& array, int64_t row, + const EncodedVariantValue& expected) { + auto unshredded = UnshredVariantArray(array); + AssertEncodedRow(*unshredded, row, expected); +} + } // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/test_util_internal.h b/cpp/src/parquet/variant/test_util_internal.h index 9edab8403da5..62fca0f518e3 100644 --- a/cpp/src/parquet/variant/test_util_internal.h +++ b/cpp/src/parquet/variant/test_util_internal.h @@ -24,10 +24,15 @@ #include #include "arrow/result.h" +#include "arrow/status.h" #include "arrow/type_fwd.h" #include "parquet/properties.h" #include "parquet/variant/builder.h" +namespace arrow::extension { +class VariantArray; +} // namespace arrow::extension + namespace parquet::variant::internal { using ::arrow::Array; @@ -35,6 +40,7 @@ using ::arrow::Buffer; using ::arrow::DataType; using ::arrow::FieldVector; using ::arrow::Result; +using ::arrow::Status; std::shared_ptr<::arrow::Table> VariantTable( const std::shared_ptr& variant_type, @@ -47,10 +53,12 @@ Result> WriteVariantTable( std::shared_ptr arrow_properties = default_arrow_writer_properties()); -::arrow::Status WriteVariantRecordBatch( - const std::shared_ptr<::arrow::Table>& table, - std::shared_ptr arrow_properties = - default_arrow_writer_properties()); +Status WriteVariantRecordBatch(const std::shared_ptr<::arrow::Table>& table, + std::shared_ptr arrow_properties = + default_arrow_writer_properties()); + +Result> RoundTripVariantArray( + const std::shared_ptr<::arrow::extension::VariantArray>& array); std::shared_ptr EmptyVariantMetadata(); @@ -59,4 +67,14 @@ EncodedVariantValue Int8Variant(int8_t value); std::shared_ptr BinaryArrayFromValues( const std::vector>& values); +std::shared_ptr<::arrow::StructArray> MakeInt64FieldGroup( + const std::vector>& values, + std::string_view typed_values, const std::vector& is_valid = {}); + +void AssertEncodedRow(const ::arrow::extension::VariantArray& array, int64_t row, + const EncodedVariantValue& expected); + +void AssertUnshreddedValue(const ::arrow::extension::VariantArray& array, int64_t row, + const EncodedVariantValue& expected); + } // namespace parquet::variant::internal diff --git a/cpp/src/parquet/variant/unshred.cc b/cpp/src/parquet/variant/unshred.cc index a6a2a19b5168..2c5e2c103cb6 100644 --- a/cpp/src/parquet/variant/unshred.cc +++ b/cpp/src/parquet/variant/unshred.cc @@ -17,11 +17,9 @@ #include "parquet/variant/unshred.h" -#include "arrow/buffer.h" #include "arrow/extension/parquet_variant.h" #include "arrow/type.h" #include "arrow/util/bit_block_counter.h" -#include "arrow/util/bitmap_ops.h" #include "parquet/exception.h" #include "parquet/variant/array_internal.h" #include "parquet/variant/decoding.h" @@ -32,27 +30,10 @@ namespace parquet::variant { using ::arrow::binary_view; using ::arrow::field; using ::arrow::struct_; -using ::arrow::StructArray; using ::arrow::extension::VariantArray; -using ::arrow::internal::checked_cast; -using ::arrow::internal::checked_pointer_cast; namespace { -std::shared_ptr<::arrow::Buffer> ParentNullBitmap(const VariantArray& array, - ::arrow::MemoryPool* pool) { - if (array.null_bitmap() == nullptr) { - return nullptr; - } - if (array.offset() == 0) { - return array.null_bitmap(); - } - PARQUET_ASSIGN_OR_THROW(auto null_bitmap, - ::arrow::internal::CopyBitmap(pool, array.null_bitmap_data(), - array.offset(), array.length())); - return null_bitmap; -} - void AppendUnshreddedRow(std::string_view metadata_value, const internal::CompiledVariantRowPlan& row_plan, int64_t row, VariantValueArrayBuilder& value_builder) { @@ -65,54 +46,14 @@ void AppendUnshreddedRow(std::string_view metadata_value, } // namespace -EncodedVariantValue UnshredVariantRow(const VariantArray& array, int64_t row, - ::arrow::MemoryPool* pool) { - if (row < 0 || row >= array.length()) { - throw ParquetException("Variant row out of bounds: ", row); - } - if (array.IsNull(row)) { - throw ParquetException("Cannot unshred null Variant row"); - } - const auto& storage = checked_cast(*array.storage()); - auto metadata_array = storage.GetFieldByName("metadata"); - auto value_array = storage.GetFieldByName("value"); - auto typed_array = storage.GetFieldByName("typed_value"); - if (metadata_array->IsNull(row)) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant extension storage: metadata is null"); - } - auto metadata_value = internal::BinaryFieldView(*metadata_array, row); - - if (typed_array == nullptr && value_array != nullptr) { - if (value_array->IsNull(row)) { - throw ParquetInvalidOrCorruptedFileException( - "Invalid Variant extension storage: value is null"); - } - return EncodedVariantValue{ - .metadata = ::arrow::Buffer::FromString(std::string(metadata_value)), - .value = ::arrow::Buffer::FromString( - std::string(internal::BinaryFieldView(*value_array, row)))}; - } - - auto row_plan = internal::CompileVariantRowPlan(value_array, typed_array); - VariantValueArrayBuilder value_builder(pool); - AppendUnshreddedRow(metadata_value, row_plan, row, value_builder); - auto unshredded_value = value_builder.Finish(); - return EncodedVariantValue{ - .metadata = ::arrow::Buffer::FromString(std::string(metadata_value)), - .value = ::arrow::Buffer::FromString(std::string(unshredded_value->GetView(0)))}; -} - std::shared_ptr<::arrow::extension::VariantArray> UnshredVariantArray( const ::arrow::extension::VariantArray& array, ::arrow::MemoryPool* pool) { - const auto& storage = checked_cast(*array.storage()); - auto metadata_array = storage.GetFieldByName("metadata"); - auto value_array = storage.GetFieldByName("value"); - auto typed_array = storage.GetFieldByName("typed_value"); + auto metadata_array = array.metadata(); + auto value_array = array.value(); + auto typed_array = array.typed_value(); - if (typed_array == nullptr && value_array != nullptr) { - return MakeVariantArrayFromStorage( - checked_pointer_cast(array.storage())); + if (!array.is_shredded()) { + return std::make_shared(array.data()); } auto row_plan = internal::CompileVariantRowPlan(value_array, typed_array); @@ -135,7 +76,7 @@ std::shared_ptr<::arrow::extension::VariantArray> UnshredVariantArray( field("value", binary_view(), /*nullable=*/false)}); return MakeVariantArrayFromChildren( std::move(storage_type), {std::move(metadata_array), std::move(unshredded_value)}, - ParentNullBitmap(array, pool)); + internal::NullBitmapForOutput(array, pool)); } } // namespace parquet::variant diff --git a/cpp/src/parquet/variant/unshred.h b/cpp/src/parquet/variant/unshred.h index e6d7b617ac6f..cccdbe31c4e2 100644 --- a/cpp/src/parquet/variant/unshred.h +++ b/cpp/src/parquet/variant/unshred.h @@ -17,8 +17,8 @@ #pragma once +#include "arrow/type_fwd.h" #include "parquet/platform.h" -#include "parquet/variant/builder.h" namespace arrow::extension { class VariantArray; @@ -26,17 +26,6 @@ class VariantArray; namespace parquet::variant { -/// Materialize a non-null Variant row into independently owned encoded buffers. -/// -/// The row index must be in bounds and the parent row must be non-null. When the input -/// is already unshredded, the encoded metadata and value are copied without validation. -/// Invalid or corrupted storage detected during materialization raises a Parquet -/// exception. -PARQUET_EXPORT -EncodedVariantValue UnshredVariantRow( - const ::arrow::extension::VariantArray& array, int64_t row, - ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); - /// Materialize a Variant array while preserving parent nulls. /// /// An already unshredded array may reuse its entire storage without validating the diff --git a/cpp/src/parquet/variant/unshred_test.cc b/cpp/src/parquet/variant/unshred_test.cc index e8d462032afd..e1944c5833eb 100644 --- a/cpp/src/parquet/variant/unshred_test.cc +++ b/cpp/src/parquet/variant/unshred_test.cc @@ -46,17 +46,13 @@ using ::arrow::field; using ::arrow::int64; using ::arrow::struct_; using ::arrow::internal::checked_pointer_cast; +using internal::AssertEncodedRow; +using internal::AssertUnshreddedValue; using internal::BinaryArrayFromValues; +using internal::MakeInt64FieldGroup; namespace { -void AssertUnshreddedRow(const ::arrow::extension::VariantArray& array, - const EncodedVariantValue& expected) { - auto actual = UnshredVariantRow(array, /*row=*/0); - ASSERT_EQ(std::string_view{*expected.metadata}, std::string_view{*actual.metadata}); - ASSERT_EQ(std::string_view{*expected.value}, std::string_view{*actual.value}); -} - EncodedVariantValue ObjectWithInt64(std::string_view field_name, int64_t value) { VariantBuilder builder; auto object = builder.StartObject(); @@ -73,23 +69,6 @@ EncodedVariantValue EmptyObjectWithField(std::string_view field_name) { return builder.Finish(); } -std::shared_ptr<::arrow::StructArray> MakeInt64FieldGroup( - const std::vector>& values, - std::string_view typed_values, const std::vector& is_valid = {}) { - auto field_group_type = - struct_({field("value", binary()), field("typed_value", int64())}); - std::shared_ptr<::arrow::Buffer> null_bitmap; - if (!is_valid.empty()) { - ::arrow::BitmapFromVector(is_valid, &null_bitmap); - } - PARQUET_ASSIGN_OR_THROW( - auto field_group, ::arrow::StructArray::Make( - {BinaryArrayFromValues(values), - ::arrow::ArrayFromJSON(int64(), std::string(typed_values))}, - field_group_type->fields(), std::move(null_bitmap))); - return field_group; -} - std::shared_ptr<::arrow::extension::VariantArray> MakeSingleRowVariantArray( std::string_view metadata, std::optional value, const std::shared_ptr<::arrow::Array>& typed_value) { @@ -110,23 +89,10 @@ TEST(TestVariantUnshred, FastPathKeepsStorage) { auto array = builder.Finish(); auto unshredded = UnshredVariantArray(*array); - const auto& input_storage = - checked_pointer_cast<::arrow::StructArray>(array->storage()); - const auto& output_storage = - checked_pointer_cast<::arrow::StructArray>(unshredded->storage()); - ASSERT_EQ(nullptr, output_storage->GetFieldByName("typed_value").get()); - - const auto& input_metadata = checked_pointer_cast<::arrow::BinaryViewArray>( - input_storage->GetFieldByName("metadata")); - const auto& input_value = checked_pointer_cast<::arrow::BinaryViewArray>( - input_storage->GetFieldByName("value")); - const auto& output_metadata = checked_pointer_cast<::arrow::BinaryViewArray>( - output_storage->GetFieldByName("metadata")); - const auto& output_value = checked_pointer_cast<::arrow::BinaryViewArray>( - output_storage->GetFieldByName("value")); - - ASSERT_EQ(input_metadata->data().get(), output_metadata->data().get()); - ASSERT_EQ(input_value->data().get(), output_value->data().get()); + ASSERT_FALSE(array->is_shredded()); + ASSERT_FALSE(unshredded->is_shredded()); + ASSERT_EQ(array->metadata()->data().get(), unshredded->metadata()->data().get()); + ASSERT_EQ(array->value()->data().get(), unshredded->value()->data().get()); } TEST(TestVariantUnshred, UnshredsValuesWithoutCopyingMetadata) { @@ -152,12 +118,12 @@ TEST(TestVariantUnshred, UnshredsValuesWithoutCopyingMetadata) { input_storage_type, {input_metadata, input_values, input_typed_values}); auto unshredded = UnshredVariantArray(*input_array); - const auto& output_storage = - checked_pointer_cast<::arrow::StructArray>(unshredded->storage()); - const auto& output_metadata = output_storage->GetFieldByName("metadata"); + ASSERT_TRUE(input_array->is_shredded()); + ASSERT_FALSE(unshredded->is_shredded()); + const auto& output_metadata = unshredded->metadata(); ASSERT_EQ(input_metadata->data().get(), output_metadata->data().get()); - const auto& output_values = checked_pointer_cast<::arrow::BinaryViewArray>( - output_storage->GetFieldByName("value")); + const auto& output_values = + checked_pointer_cast<::arrow::BinaryViewArray>(unshredded->value()); ASSERT_EQ(2, output_values->length()); ASSERT_EQ(std::string_view{*typed_only_row.value}, output_values->GetView(0)); ASSERT_EQ(std::string_view{*residual_value_row.value}, output_values->GetView(1)); @@ -174,7 +140,7 @@ TEST(TestVariantUnshred, UnshredsDecimalTypedValues) { auto array = MakeSingleRowVariantArray(std::string_view{*expected.metadata}, std::nullopt, typed_value); - AssertUnshreddedRow(*array, expected); + AssertUnshreddedValue(*array, /*row=*/0, expected); }; assert_decimal(::arrow::decimal32(/*precision=*/9, /*scale=*/2), R"(["-1234567.89"])", @@ -218,7 +184,7 @@ TEST(TestVariantUnshred, MissingTopLevelValue) { ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); ValidateVariants(chunked); - AssertUnshreddedRow(*array, expected); + AssertUnshreddedValue(*array, /*row=*/0, expected); } TEST(TestVariantUnshred, MissingObjectField) { @@ -234,7 +200,7 @@ TEST(TestVariantUnshred, MissingObjectField) { ValidateVariants(chunked); ValidateVariants(chunked); - AssertUnshreddedRow(*array, expected); + AssertUnshreddedValue(*array, /*row=*/0, expected); } TEST(TestVariantUnshred, NullObjectFieldGroup) { @@ -250,7 +216,7 @@ TEST(TestVariantUnshred, NullObjectFieldGroup) { ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); ValidateVariants(chunked); - AssertUnshreddedRow(*array, expected); + AssertUnshreddedValue(*array, /*row=*/0, expected); } TEST(TestVariantUnshred, MissingListElements) { @@ -275,15 +241,7 @@ TEST(TestVariantUnshred, MissingListElements) { ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); ValidateVariants(chunked); - AssertUnshreddedRow(*array, expected); -} - -TEST(TestVariantUnshred, RejectsNullRow) { - VariantArrayBuilder builder; - builder.AppendNull(); - auto array = builder.Finish(); - - ASSERT_THROW(UnshredVariantRow(*array, /*row=*/0), ParquetException); + AssertUnshreddedValue(*array, /*row=*/0, expected); } TEST(TestVariantUnshred, PreservesSliceNullBitmap) { @@ -309,10 +267,7 @@ TEST(TestVariantUnshred, PreservesSliceNullBitmap) { ASSERT_EQ(2, unshredded->length()); ASSERT_TRUE(unshredded->IsNull(0)); ASSERT_FALSE(unshredded->IsNull(1)); - ASSERT_THROW(UnshredVariantRow(*unshredded, /*row=*/0), ParquetException); - auto actual = UnshredVariantRow(*unshredded, /*row=*/1); - ASSERT_EQ(std::string_view{*encoded.metadata}, std::string_view{*actual.metadata}); - ASSERT_EQ(std::string_view{*encoded.value}, std::string_view{*actual.value}); + AssertEncodedRow(*unshredded, /*row=*/1, encoded); } TEST(TestVariantUnshred, RejectsDuplicateResidualOnUnshred) { @@ -327,8 +282,6 @@ TEST(TestVariantUnshred, RejectsDuplicateResidualOnUnshred) { ::arrow::ChunkedArray chunked(::arrow::ArrayVector{array}); ValidateVariants(chunked); - ASSERT_THROW(UnshredVariantRow(*array, /*row=*/0), - ParquetInvalidOrCorruptedFileException); ASSERT_THROW(UnshredVariantArray(*array), ParquetInvalidOrCorruptedFileException); } @@ -347,8 +300,6 @@ TEST(TestVariantUnshred, DuplicateFields) { ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); ASSERT_THROW(ValidateVariants(chunked), ParquetInvalidOrCorruptedFileException); - ASSERT_THROW(UnshredVariantRow(*array, /*row=*/0), - ParquetInvalidOrCorruptedFileException); ASSERT_THROW(UnshredVariantArray(*array), ParquetInvalidOrCorruptedFileException); } diff --git a/cpp/src/parquet/variant/validate.cc b/cpp/src/parquet/variant/validate.cc index a3ed1dac3f5e..e07269937813 100644 --- a/cpp/src/parquet/variant/validate.cc +++ b/cpp/src/parquet/variant/validate.cc @@ -59,12 +59,10 @@ struct VariantValidationPlan { template void ValidateVariantArrayRows(const VariantArray& array, - const std::shared_ptr& valid_rows, - MemoryPool* pool) { - const auto& storage = checked_cast(*array.storage()); - auto metadata_array = storage.GetFieldByName("metadata"); - auto value_array = storage.GetFieldByName("value"); - auto typed_array = storage.GetFieldByName("typed_value"); + const std::shared_ptr& valid_rows) { + auto metadata_array = array.metadata(); + auto value_array = array.value(); + auto typed_array = array.typed_value(); if constexpr (strict) { if (value_array == nullptr) { @@ -90,7 +88,6 @@ void ValidateVariantArrayRows(const VariantArray& array, }); } -template std::optional BuildVariantValidationPlan( std::shared_ptr array) { switch (array->type_id()) { @@ -100,13 +97,13 @@ std::optional BuildVariantValidationPlan( if (ext_type.extension_name() == kVariantExtensionName) { return VariantValidationPlan{.array = std::move(array), .children = {}}; } - return BuildVariantValidationPlan(ext_array.storage()); + return BuildVariantValidationPlan(ext_array.storage()); } case ::arrow::Type::STRUCT: { const auto& struct_array = checked_cast(*array); std::vector children; for (auto field : struct_array.fields()) { - auto child_plan = BuildVariantValidationPlan(std::move(field)); + auto child_plan = BuildVariantValidationPlan(std::move(field)); if (child_plan.has_value()) { children.push_back(std::move(*child_plan)); } @@ -124,7 +121,7 @@ std::optional BuildVariantValidationPlan( case ::arrow::Type::FIXED_SIZE_LIST: case ::arrow::Type::MAP: { auto values = internal::ValuesArray(*array); - auto child_plan = BuildVariantValidationPlan(std::move(values)); + auto child_plan = BuildVariantValidationPlan(std::move(values)); if (!child_plan.has_value()) { return std::nullopt; } @@ -135,8 +132,7 @@ std::optional BuildVariantValidationPlan( std::vector children; for (const auto& child_data : array->data()->child_data) { if (child_data != nullptr) { - auto child_plan = - BuildVariantValidationPlan(::arrow::MakeArray(child_data)); + auto child_plan = BuildVariantValidationPlan(::arrow::MakeArray(child_data)); if (child_plan.has_value()) { children.push_back(std::move(*child_plan)); } @@ -157,7 +153,7 @@ void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, switch (plan.array->type_id()) { case ::arrow::Type::EXTENSION: ValidateVariantArrayRows(checked_cast(*plan.array), - valid_rows, pool); + valid_rows); return; case ::arrow::Type::STRUCT: { const auto& struct_array = checked_cast(*plan.array); @@ -205,7 +201,7 @@ void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, template void ValidateVariants(const ::arrow::ChunkedArray& data, MemoryPool* pool) { for (const auto& chunk : data.chunks()) { - auto plan = BuildVariantValidationPlan(chunk); + auto plan = BuildVariantValidationPlan(chunk); if (plan.has_value()) { ValidateVariantPlan(*plan, pool, nullptr); } From 670e2f481d97aad512c0221ff0dfadc7d3d66b30 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Wed, 22 Jul 2026 21:05:43 +0800 Subject: [PATCH 09/10] try /Zc:preprocessor in MSVC --- cpp/src/parquet/CMakeLists.txt | 5 +++++ cpp/src/parquet/meson.build | 21 +++++++++++++++++++-- cpp/src/parquet/variant/builder.cc | 8 ++------ cpp/src/parquet/variant/slot_internal.cc | 13 ++++--------- 4 files changed, 30 insertions(+), 17 deletions(-) diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index cf626f91af23..48eb8bac706c 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -200,6 +200,11 @@ set(PARQUET_SRCS variant/validate.cc xxhasher.cc) +if(MSVC) + set_source_files_properties(variant/builder.cc variant/slot_internal.cc + PROPERTIES COMPILE_OPTIONS /Zc:preprocessor) +endif() + if(ARROW_HAVE_RUNTIME_AVX2) # AVX2 is used as a proxy for BMI2. list(APPEND diff --git a/cpp/src/parquet/meson.build b/cpp/src/parquet/meson.build index 9452f056438c..aab7db634b4d 100644 --- a/cpp/src/parquet/meson.build +++ b/cpp/src/parquet/meson.build @@ -56,15 +56,18 @@ parquet_srcs = files( 'stream_writer.cc', 'types.cc', 'variant/array_internal.cc', - 'variant/builder.cc', 'variant/decoding.cc', 'variant/shred.cc', - 'variant/slot_internal.cc', 'variant/unshred.cc', 'variant/validate.cc', 'xxhasher.cc', ) +parquet_variant_va_opt_srcs = files( + 'variant/builder.cc', + 'variant/slot_internal.cc', +) + thrift_dep = dependency('thrift', allow_fallback: false, required: false) if not thrift_dep.found() cmake = import('cmake') @@ -124,11 +127,25 @@ else parquet_srcs += files('encryption/encryption_internal_nossl.cc') endif +parquet_link_whole = [] +if cpp_compiler.get_id() == 'msvc' + parquet_variant_va_opt_lib = static_library( + 'parquet-variant-va-opt', + parquet_variant_va_opt_srcs, + dependencies: parquet_deps, + cpp_args: ['/Zc:preprocessor'], + ) + parquet_link_whole += [parquet_variant_va_opt_lib] +else + parquet_srcs += parquet_variant_va_opt_srcs +endif + parquet_lib = library( 'arrow-parquet', sources: parquet_srcs, dependencies: parquet_deps, gnu_symbol_visibility: 'inlineshidden', + link_whole: parquet_link_whole, ) parquet_dep = declare_dependency(link_with: parquet_lib) diff --git a/cpp/src/parquet/variant/builder.cc b/cpp/src/parquet/variant/builder.cc index 7b58409a6cd2..2adf712c67dc 100644 --- a/cpp/src/parquet/variant/builder.cc +++ b/cpp/src/parquet/variant/builder.cc @@ -753,12 +753,8 @@ VariantObjectBuilder::VariantObjectBuilder(VariantObjectBuilder&&) noexcept = de VariantObjectBuilder& VariantObjectBuilder::operator=(VariantObjectBuilder&&) noexcept = default; -#ifdef _MSC_VER -# define VARIANT_OBJECT_DECL(...) (std::string_view field_name, ##__VA_ARGS__) -#else -# define VARIANT_OBJECT_DECL(...) \ - (std::string_view field_name __VA_OPT__(, ) __VA_ARGS__) // NOLINT -#endif +#define VARIANT_OBJECT_DECL(...) \ + (std::string_view field_name __VA_OPT__(, ) __VA_ARGS__) // NOLINT #define DEFINE_OBJECT_DIRECT_APPEND(NAME, decl_args, ...) \ void VariantObjectBuilder::Append##NAME VARIANT_OBJECT_DECL decl_args { \ diff --git a/cpp/src/parquet/variant/slot_internal.cc b/cpp/src/parquet/variant/slot_internal.cc index abd334dabe8f..e32f63e8e908 100644 --- a/cpp/src/parquet/variant/slot_internal.cc +++ b/cpp/src/parquet/variant/slot_internal.cc @@ -45,15 +45,10 @@ using ::arrow::ExtensionArray; using ::arrow::StructArray; using ::arrow::internal::checked_cast; -#ifdef _MSC_VER -# define VARIANT_TARGET_DECL(...) (BuildTarget * target, ##__VA_ARGS__) -# define VARIANT_OBJECT_CALL_ARGS(...) destination.field_name, ##__VA_ARGS__ -#else -# define VARIANT_TARGET_DECL(...) \ - (BuildTarget * target __VA_OPT__(, ) __VA_ARGS__) // NOLINT -# define VARIANT_OBJECT_CALL_ARGS(...) \ - destination.field_name __VA_OPT__(, ) __VA_ARGS__ // NOLINT -#endif +#define VARIANT_TARGET_DECL(...) \ + (BuildTarget * target __VA_OPT__(, ) __VA_ARGS__) // NOLINT +#define VARIANT_OBJECT_CALL_ARGS(...) \ + destination.field_name __VA_OPT__(, ) __VA_ARGS__ // NOLINT #define DEFINE_TARGET_APPEND(NAME, decl_args, ...) \ void Append##NAME VARIANT_TARGET_DECL decl_args { \ From c343dedf62faf694dc0cda9cdc2c0b9b9863b201 Mon Sep 17 00:00:00 2001 From: Zehua Zou Date: Wed, 22 Jul 2026 21:11:36 +0800 Subject: [PATCH 10/10] address ci --- cpp/src/arrow/extension/CMakeLists.txt | 7 ++---- cpp/src/parquet/variant/array_internal.cc | 28 +++++++++++------------ cpp/src/parquet/variant/array_internal.h | 7 ++++-- cpp/src/parquet/variant/meson.build | 9 +++++++- cpp/src/parquet/variant/shred.cc | 13 ++++------- 5 files changed, 33 insertions(+), 31 deletions(-) diff --git a/cpp/src/arrow/extension/CMakeLists.txt b/cpp/src/arrow/extension/CMakeLists.txt index 1128a03a35d3..fbb0daf69613 100644 --- a/cpp/src/arrow/extension/CMakeLists.txt +++ b/cpp/src/arrow/extension/CMakeLists.txt @@ -15,11 +15,8 @@ # specific language governing permissions and limitations # under the License. -set(CANONICAL_EXTENSION_TESTS - bool8_test.cc - json_test.cc - parquet_variant_test.cc - uuid_test.cc) +set(CANONICAL_EXTENSION_TESTS bool8_test.cc json_test.cc parquet_variant_test.cc + uuid_test.cc) if(ARROW_JSON) list(APPEND CANONICAL_EXTENSION_TESTS tensor_extension_array_test.cc opaque_test.cc) diff --git a/cpp/src/parquet/variant/array_internal.cc b/cpp/src/parquet/variant/array_internal.cc index 3a3b22ae1a45..d2a007fe10b8 100644 --- a/cpp/src/parquet/variant/array_internal.cc +++ b/cpp/src/parquet/variant/array_internal.cc @@ -28,6 +28,20 @@ namespace parquet::variant::internal { using ::arrow::Array; using ::arrow::internal::checked_cast; +std::string_view BinaryFieldView(const Array& array, int64_t row) { + switch (array.type_id()) { + case ::arrow::Type::BINARY: + return checked_cast(array).GetView(row); + case ::arrow::Type::LARGE_BINARY: + return checked_cast(array).GetView(row); + case ::arrow::Type::BINARY_VIEW: + return checked_cast(array).GetView(row); + default: + throw ParquetInvalidOrCorruptedFileException("Expected binary Variant field, got ", + array.type()->ToString()); + } +} + std::shared_ptr ValuesArray(const Array& array) { switch (array.type_id()) { case ::arrow::Type::LIST_VIEW: @@ -74,20 +88,6 @@ std::pair ValuesRangeAt(const Array& array, int64_t row) { } } -std::string_view BinaryFieldView(const Array& array, int64_t row) { - switch (array.type_id()) { - case ::arrow::Type::BINARY: - return checked_cast(array).GetView(row); - case ::arrow::Type::LARGE_BINARY: - return checked_cast(array).GetView(row); - case ::arrow::Type::BINARY_VIEW: - return checked_cast(array).GetView(row); - default: - throw ParquetInvalidOrCorruptedFileException("Expected binary Variant field, got ", - array.type()->ToString()); - } -} - std::shared_ptr<::arrow::Buffer> FinishNullBitmap( ::arrow::TypedBufferBuilder& builder) { if (builder.false_count() == 0) { diff --git a/cpp/src/parquet/variant/array_internal.h b/cpp/src/parquet/variant/array_internal.h index 87f3c1e23b9a..937f69d667c1 100644 --- a/cpp/src/parquet/variant/array_internal.h +++ b/cpp/src/parquet/variant/array_internal.h @@ -26,15 +26,18 @@ #include "arrow/buffer.h" #include "arrow/buffer_builder.h" #include "arrow/util/bit_block_counter.h" +#include "parquet/platform.h" namespace parquet::variant::internal { +PARQUET_EXPORT +std::string_view BinaryFieldView(const ::arrow::Array& array, int64_t row); + +PARQUET_EXPORT std::shared_ptr<::arrow::Array> ValuesArray(const ::arrow::Array& array); std::pair ValuesRangeAt(const ::arrow::Array& array, int64_t row); -std::string_view BinaryFieldView(const ::arrow::Array& array, int64_t row); - std::shared_ptr<::arrow::Buffer> FinishNullBitmap( ::arrow::TypedBufferBuilder& builder); diff --git a/cpp/src/parquet/variant/meson.build b/cpp/src/parquet/variant/meson.build index 679c2a845e64..cecf81f4333f 100644 --- a/cpp/src/parquet/variant/meson.build +++ b/cpp/src/parquet/variant/meson.build @@ -32,6 +32,13 @@ exc = executable( test('parquet-variant-test', exc) install_headers( - ['builder.h', 'decoding.h', 'format.h', 'shred.h', 'unshred.h', 'validate.h'], + [ + 'builder.h', + 'decoding.h', + 'format.h', + 'shred.h', + 'unshred.h', + 'validate.h', + ], subdir: 'parquet/variant', ) diff --git a/cpp/src/parquet/variant/shred.cc b/cpp/src/parquet/variant/shred.cc index 4199252baa9d..0364c2db1f2f 100644 --- a/cpp/src/parquet/variant/shred.cc +++ b/cpp/src/parquet/variant/shred.cc @@ -17,7 +17,6 @@ #include "parquet/variant/shred.h" -#include #include #include #include @@ -218,14 +217,10 @@ std::shared_ptr DecodeSourceScalar(const VariantValueView& value) { return std::make_shared<::arrow::Int32Scalar>(LoadLittleEndian(payload)); case VariantPrimitiveType::kInt64: return std::make_shared<::arrow::Int64Scalar>(LoadLittleEndian(payload)); - case VariantPrimitiveType::kFloat: { - const auto bits = LoadLittleEndian(payload); - return std::make_shared<::arrow::FloatScalar>(std::bit_cast(bits)); - } - case VariantPrimitiveType::kDouble: { - const auto bits = LoadLittleEndian(payload); - return std::make_shared<::arrow::DoubleScalar>(std::bit_cast(bits)); - } + case VariantPrimitiveType::kFloat: + return std::make_shared<::arrow::FloatScalar>(LoadLittleEndian(payload)); + case VariantPrimitiveType::kDouble: + return std::make_shared<::arrow::DoubleScalar>(LoadLittleEndian(payload)); case VariantPrimitiveType::kDecimal4: { const auto scale = static_cast(payload[0]); auto type = ::arrow::decimal32(/*precision=*/9, scale);