diff --git a/cpp/src/arrow/extension/CMakeLists.txt b/cpp/src/arrow/extension/CMakeLists.txt index ae52bc32a998..fbb0daf69613 100644 --- a/cpp/src/arrow/extension/CMakeLists.txt +++ b/cpp/src/arrow/extension/CMakeLists.txt @@ -15,7 +15,8 @@ # 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 95aa5a0eb68e..f64c2bfe1207 100644 --- a/cpp/src/arrow/extension/parquet_variant.cc +++ b/cpp/src/arrow/extension/parquet_variant.cc @@ -17,28 +17,203 @@ #include "arrow/extension/parquet_variant.h" +#include #include +#include "arrow/array/array_nested.h" #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 { +namespace { + +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; + } +} + +template +bool IsSupportedTypedValue(const std::shared_ptr& field); + +template +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; + } + } + + 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; + } + + switch (field->type()->id()) { + case Type::STRUCT: + return field->type()->num_fields() > 0 && + 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: { + auto& inner_field = field->type()->field(0); + return !inner_field->nullable() && + IsVariantFieldGroup(inner_field->type()); + } + default: + return IsSupportedPrimitiveTypedValue(field->type()); + } +} + +template +bool IsSupportedStorageTypeImpl(const std::shared_ptr& storage_type) { + 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; + } + } + + if (metadata == nullptr) { + return false; + } + + 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; + } +} + +} // 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) { - // 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,7 +224,14 @@ bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const { Result> VariantExtensionType::Deserialize( std::shared_ptr storage_type, const std::string& serialized) const { - return VariantExtensionType::Make(std::move(storage_type)); + 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 std::make_shared(std::move(storage_type)); } std::string VariantExtensionType::Serialize() const { return ""; } @@ -57,50 +239,14 @@ 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; -} -} // 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(); - } - } - } - - return false; + return IsSupportedStorageTypeImpl(storage_type); } Result> VariantExtensionType::Make( @@ -113,9 +259,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..75c26c205846 100644 --- a/cpp/src/arrow/extension/parquet_variant.h +++ b/cpp/src/arrow/extension/parquet_variant.h @@ -18,15 +18,31 @@ #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; + + /// \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. @@ -43,13 +59,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 +97,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/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 f8a42b5b96bf..48eb8bac706c 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,20 @@ set(PARQUET_SRCS statistics.cc stream_reader.cc stream_writer.cc - types.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) + +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. @@ -363,6 +375,7 @@ add_subdirectory(api) add_subdirectory(arrow) add_subdirectory(encryption) add_subdirectory(geospatial) +add_subdirectory(variant) arrow_install_all_headers("parquet") @@ -409,7 +422,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..b96790ecd059 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,84 @@ 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); + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), /*nullable=*/false), + ::arrow::field("value", ::arrow::binary(), /*nullable=*/false)}); - // 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{}); + ASSERT_NO_FATAL_FAILURE( + CheckParquetVariantSchema("variant_unshredded", {metadata, value}, storage_type, + /*check_serialized_arrow_schema=*/true)); +} - // 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); +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)); +} - auto arrow_schema = ::arrow::schema({::arrow::field( - "variant_unshredded", register_extension ? variant_extension : arrow_variant)}); +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_OK(ConvertSchema(parquet_fields, /*metadata=*/nullptr, props)); + { + 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)); } - for (bool register_extension : {true, false}) { - ::arrow::ExtensionTypeGuard guard(register_extension - ? ::arrow::DataTypeVector{variant_extension} - : ::arrow::DataTypeVector{}); - - // 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)}); + 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}); - 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)); - } + 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) { @@ -1534,6 +1597,241 @@ 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, ParquetVariantShreddedWithoutValue) { + auto storage_type = + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("typed_value", ::arrow::int64())}); + 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)})); +} + +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/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/schema.cc b/cpp/src/parquet/arrow/schema.cc index bc4de6c39b5e..1c5296206c25 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,40 @@ 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) { + // `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"); + } - NodePtr value_node; - RETURN_NOT_OK( - FieldToNode("value", type->value(), properties, arrow_properties, &value_node)); + std::vector children; + children.reserve(type->storage_type()->num_fields()); - *out = GroupNode::Make(name, RepetitionFromNullable(nullable), - {std::move(metadata_node), std::move(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(children), LogicalType::Variant(), field_id); return Status::OK(); @@ -280,8 +534,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 +542,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 +694,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 +725,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 +776,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 +796,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 +849,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 +1439,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..d471d2aa4f26 100644 --- a/cpp/src/parquet/arrow/variant_test.cc +++ b/cpp/src/parquet/arrow/variant_test.cc @@ -15,57 +15,474 @@ // 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/builder.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/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); + + 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) { + 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) { + 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 = ::arrow::ArrayFromJSON(::arrow::utf8(), R"(["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)); + 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) { + 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); + 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)); +} + +TEST(TestVariantExtensionType, WriterSkipsNullParents) { + 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, ::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)); +} + +TEST(TestVariantExtensionType, WriterValidatesShreddedPrimitiveConflicts) { + 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 = ::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)); + + 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 = ::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 storage_type = struct_({field("metadata", binary(), /*nullable=*/false), + field("typed_value", ::arrow::int64())}); + ASSERT_RAISES(Invalid, ::arrow::extension::VariantExtensionType::Make(storage_type)); +} + +TEST(TestVariantExtensionType, ReadsDictionaryEncodedMetadata) { + 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) { + 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) { + 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 = ::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)); +} + +TEST(TestVariantExtensionType, WriterValidatesShreddedObjectConflicts) { + variant::VariantBuilder object_builder; + 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())}); + 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}), + ::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())); + + 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)); + + 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())); + auto missing_table = + VariantTable(variant_type, {metadata_array, valid_value_array, missing_typed_array}, + storage_type->fields()); + 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()); -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)); + 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)); + } } } diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index e0fbe308219c..ae04ceea6702 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,10 @@ 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()) { + PARQUET_CATCH_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 +453,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 +475,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()) { + PARQUET_CATCH_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 +542,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..aab7db634b4d 100644 --- a/cpp/src/parquet/meson.build +++ b/cpp/src/parquet/meson.build @@ -55,9 +55,19 @@ parquet_srcs = files( 'stream_reader.cc', 'stream_writer.cc', 'types.cc', + 'variant/array_internal.cc', + 'variant/decoding.cc', + 'variant/shred.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') @@ -117,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) @@ -130,6 +154,7 @@ subdir('api') subdir('arrow') subdir('encryption') subdir('geospatial') +subdir('variant') install_headers( [ @@ -229,6 +254,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..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) {} @@ -1259,15 +1260,24 @@ 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; } 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. @@ -1306,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_; }; @@ -1332,7 +1343,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 +1448,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 +1482,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 +1521,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 +1540,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 +1553,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..c41bc3c6332f --- /dev/null +++ b/cpp/src/parquet/variant/CMakeLists.txt @@ -0,0 +1,29 @@ +# 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 + decoding_test.cc + parquet_testing_test.cc + shred_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..d2a007fe10b8 --- /dev/null +++ b/cpp/src/parquet/variant/array_internal.cc @@ -0,0 +1,114 @@ +// 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/bitmap_ops.h" +#include "arrow/util/checked_cast.h" +#include "parquet/exception.h" + +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: + 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::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 new file mode 100644 index 000000000000..937f69d667c1 --- /dev/null +++ b/cpp/src/parquet/variant/array_internal.h @@ -0,0 +1,63 @@ +// 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/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::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) { + 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 new file mode 100644 index 000000000000..2adf712c67dc --- /dev/null +++ b/cpp/src/parquet/variant/builder.cc @@ -0,0 +1,1238 @@ +// 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/array/array_binary.h" +#include "arrow/buffer_builder.h" +#include "arrow/builder.h" // IWYU pragma: keep +#include "arrow/type.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/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" + +namespace parquet::variant { + +using ::arrow::binary_view; +using ::arrow::BinaryViewArray; +using ::arrow::BufferBuilder; +using ::arrow::ExtensionType; +using ::arrow::field; +using ::arrow::struct_; +using ::arrow::StructType; +using ::arrow::Type; +using ::arrow::TypedBufferBuilder; +using ::arrow::extension::VariantExtensionType; +using internal::BinaryViewColumnBuilder; + +namespace bit_util = ::arrow::bit_util; + +namespace { + +void AppendLittleEndian(BufferBuilder& out, uint32_t value, uint8_t width) { + DCHECK_LE(width, sizeof(uint32_t)); + const auto little_endian = bit_util::ToLittleEndian(value); + PARQUET_THROW_NOT_OK(out.Append(&little_endian, width)); +} + +template + requires(std::is_arithmetic_v) +void AppendFixedLittleEndian(BufferBuilder& out, T value) { + const auto little_endian = bit_util::ToLittleEndian(value); + PARQUET_THROW_NOT_OK(out.Append(&little_endian, sizeof(T))); +} + +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) { + if (value <= std::numeric_limits::max()) { + return 1; + } + if (value <= std::numeric_limits::max()) { + return 2; + } + if (value <= 0xFFFFFFU) { + return 3; + } + return 4; +} + +class MetadataBuilder { + public: + 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; +}; + +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 TryUpsert(std::string_view name) override { + auto it = field_ids_.find(name); + if (it != field_ids_.end()) { + return it->second; + } + if (field_names_.size() >= std::numeric_limits::max()) { + throw ParquetException("Variant metadata dictionary is too large"); + } + 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 override { + DCHECK_LT(id, field_names_.size()); + return field_names_[id]; + } + + size_t size() const override { return field_names_.size(); } + + 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 override { + 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()) { + throw ParquetException("Variant metadata dictionary is too large"); + } + + const uint8_t offset_size = + WidthForValue(std::max(field_names_.size(), bytes_size)); + 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_; + const auto header = + static_cast(internal::kVariantVersion | + (sorted_strings ? internal::kMetadataSortedStringsMask : 0) | + ((offset_size - 1) << 6)); + PARQUET_THROW_NOT_OK(out.Append(&header, sizeof(header))); + AppendLittleEndian(out, static_cast(field_names_.size()), offset_size); + + uint32_t offset = 0; + AppendLittleEndian(out, offset, offset_size); + for (const auto& string : field_names_) { + offset += static_cast(string.size()); + AppendLittleEndian(out, offset, offset_size); + } + for (const auto& string : field_names_) { + PARQUET_THROW_NOT_OK(out.Append(string)); + } + } + + private: + 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; + } + + 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_; +}; + +class VariantValueWriter { + public: + explicit VariantValueWriter(BufferBuilder& out) : out_(out) {} + + template + requires internal::HeaderOnlyVariantPrimitive + void Append() { + AppendPrimitiveHeader(); + } + + template + requires internal::FixedVariantPrimitive + void Append(internal::VariantFixedPrimitiveTraits::CType value) { + AppendPrimitiveHeader(); + if constexpr (sizeof(value) == 1) { + const auto byte = static_cast(value); + PARQUET_THROW_NOT_OK(out_.Append(&byte, sizeof(byte))); + } else { + AppendFixedLittleEndian(out_, value); + } + } + + template + requires internal::DecimalVariantPrimitive + void Append(internal::VariantDecimalPrimitiveTraits::CType value, uint8_t scale) { + internal::ValidateDecimalScale(scale); + AppendPrimitiveHeader(); + PARQUET_THROW_NOT_OK(out_.Append(&scale, sizeof(scale))); + const auto bytes = ToLittleEndianBytes(value); + PARQUET_THROW_NOT_OK(out_.Append(bytes.data(), static_cast(bytes.size()))); + } + + template + requires internal::LengthPrefixedVariantPrimitive + void Append(std::string_view value) { + if (value.size() > std::numeric_limits::max()) { + throw ParquetException("Variant ", + type == VariantPrimitiveType::kBinary ? "binary" : "string", + " value is too large"); + } + if constexpr (type == VariantPrimitiveType::kString) { + internal::ValidateUtf8(value, "primitive string value"); + } + AppendPrimitiveHeader(); + AppendLittleEndian(out_, static_cast(value.size()), 4); + PARQUET_THROW_NOT_OK(out_.Append(value)); + } + + template + requires internal::UuidVariantPrimitive + void Append(std::string_view big_endian_bytes) { + if (big_endian_bytes.size() != 16) { + throw ParquetException("Variant UUID values must be 16 bytes"); + } + AppendPrimitiveHeader(); + 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"); + } + internal::ValidateUtf8(value, "short string value"); + const auto header = static_cast( + (value.size() << 2) | static_cast(VariantBasicType::kShortString)); + PARQUET_THROW_NOT_OK(out_.Append(&header, sizeof(header))); + 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() { + const auto header = + static_cast((static_cast(type) << 2) | + static_cast(VariantBasicType::kPrimitive)); + PARQUET_THROW_NOT_OK(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 { + 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; + std::unique_ptr metadata; + std::vector frames; + bool root_has_value = false; +}; + +void CheckRootWritable(const VariantBuildState& state) { + if (!state.frames.empty()) { + throw ParquetException("VariantBuilder has an active container"); + } + if (state.root_has_value) { + throw ParquetException("VariantBuilder already has a root value"); + } +} + +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"); + } + const auto& frame = state.frames.back(); + if (frame.finished || frame.kind != kind) { + throw ParquetException("Variant nested builder has invalid state"); + } +} + +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); + } +} + +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); +} + +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"); + } + + 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 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) { + AppendLittleEndian(state.value, field.field_id, id_size); + } + for (const auto& field : fields) { + AppendLittleEndian(state.value, field.offset, offset_size); + } + AppendLittleEndian(state.value, values_size, offset_size); + DCHECK_EQ(state.value.length(), frame.value_start + header_size); + state.value.UnsafeAdvance(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 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) { + AppendLittleEndian(state.value, offset, offset_size); + } + 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; + +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; + DCHECK_GE(values_size, 0); + if (std::cmp_greater(values_size, std::numeric_limits::max())) { + throw ParquetException("Variant container values are too large"); + } + + 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(); + if (!is_root) { + return; + } + + state.root_has_value = true; + if (callback) { + callback(state); + } +} + +template +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); + throw; + } + state.root_has_value = true; +} + +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(); + + const auto field_id = state.metadata->TryUpsert(field_name); + if (!frame.object_field_ids.insert(field_id).second) { + state.metadata->Truncate(metadata_size); + throw ParquetException("Duplicate Variant object field: ", field_name); + } + const auto offset = state.value.length() - frame.value_start; + DCHECK_GE(offset, 0); + 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); + try { + std::invoke(std::forward(write), writer); + } catch (...) { + state.value.Rewind(value_size); + state.metadata->Truncate(metadata_size); + TruncateFrameEntries(state.frames[frame_index], entry_count); + throw; + } +} + +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(); + + const auto offset = state.value.length() - frame.value_start; + DCHECK_GE(offset, 0); + 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); + try { + std::invoke(std::forward(write), writer); + } catch (...) { + state.value.Rewind(value_size); + TruncateFrameEntries(state.frames[frame_index], entry_count); + throw; + } +} + +} // namespace + +namespace internal { + +struct NestedVariantBuilderImpl { + NestedVariantBuilderImpl(VariantBuildState& state, size_t frame_index, + RootFinishCallback callback) + : state(state), frame_index(frame_index), callback(std::move(callback)) {} + + ~NestedVariantBuilderImpl() { + if (!active || frame_index >= state.frames.size()) { + return; + } + 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); + } + state.frames.resize(frame_index); + } + + VariantBuildState& state; + size_t frame_index = 0; + RootFinishCallback callback; + bool active = true; +}; + +} // namespace internal + +struct VariantBuilder::Impl { + explicit Impl(MemoryPool* pool, bool validate) + : pool(pool), + value_builder(pool), + validate(validate), + state(value_builder, std::make_unique()) {} + + void Reset() { + value_builder.Reset(); + state.Reset(std::make_unique()); + } + + MemoryPool* pool; + BufferBuilder value_builder; + bool validate; + VariantBuildState state; +}; + +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); +} + +uint32_t VariantBuilder::AddFieldName(std::string_view name) { + return impl_->state.metadata->TryUpsert(name); +} + +#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 + +#define DEFINE_SPECIAL_APPEND(NAME, decl_args, ...) \ + void VariantBuilder::Append##NAME decl_args { \ + AppendValueWith(impl_->state, [&](VariantValueWriter& writer) { \ + writer.Append##NAME(__VA_ARGS__); \ + }); \ + } +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_SPECIAL_APPEND) +DEFINE_SPECIAL_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_SPECIAL_APPEND + +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() { + return VariantObjectBuilder(StartContainer(impl_->state)); +} + +VariantListBuilder VariantBuilder::StartList() { + return VariantListBuilder(StartContainer(impl_->state)); +} + +EncodedVariantValue VariantBuilder::Finish() { + if (!impl_->state.frames.empty()) { + throw ParquetException("Cannot finish VariantBuilder with active containers"); + } + if (!impl_->state.root_has_value) { + throw ParquetException("Cannot finish empty VariantBuilder"); + } + + 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; +} + +void VariantBuilder::Reset() { impl_->Reset(); } + +VariantObjectBuilder::VariantObjectBuilder( + std::unique_ptr impl) + : impl_(std::move(impl)) {} +VariantObjectBuilder::~VariantObjectBuilder() = default; +VariantObjectBuilder::VariantObjectBuilder(VariantObjectBuilder&&) noexcept = default; +VariantObjectBuilder& VariantObjectBuilder::operator=(VariantObjectBuilder&&) noexcept = + default; + +#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 { \ + 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) { + return VariantObjectBuilder( + ObjectStartContainer(*impl_, field_name)); +} + +VariantListBuilder VariantObjectBuilder::StartList(std::string_view field_name) { + return VariantListBuilder( + ObjectStartContainer(*impl_, field_name)); +} + +void VariantObjectBuilder::Finish() { + FinishFrame(impl_->state, impl_->frame_index, + impl_->callback); + impl_->active = false; +} + +VariantListBuilder::VariantListBuilder( + std::unique_ptr impl) + : impl_(std::move(impl)) {} +VariantListBuilder::~VariantListBuilder() = default; +VariantListBuilder::VariantListBuilder(VariantListBuilder&&) noexcept = default; +VariantListBuilder& VariantListBuilder::operator=(VariantListBuilder&&) noexcept = + default; + +#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__); \ + }); \ + } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_LIST_DIRECT_APPEND) +#undef DEFINE_LIST_DIRECT_APPEND + +#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__); }); \ + } +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_LIST_SPECIAL_APPEND) +DEFINE_LIST_SPECIAL_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_LIST_SPECIAL_APPEND + +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() { + return VariantObjectBuilder(ListStartContainer(*impl_)); +} + +VariantListBuilder VariantListBuilder::StartList() { + return VariantListBuilder(ListStartContainer(*impl_)); +} + +void VariantListBuilder::Finish() { + 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), + state(value_column.bytes_builder(), std::make_unique()) { + } + + 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; + } + } + + 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) { + 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_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()); + 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(); + auto null_bitmap = internal::FinishNullBitmap(validity_builder); + + 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)}); + PARQUET_ASSIGN_OR_THROW(auto storage, + StructArray::Make({metadata, value}, storage_type->fields(), + std::move(null_bitmap), null_count)); + 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) { + 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) + : impl_(std::make_unique(pool)) {} +VariantArrayBuilder::~VariantArrayBuilder() = default; +VariantArrayBuilder::VariantArrayBuilder(VariantArrayBuilder&&) noexcept = default; +VariantArrayBuilder& VariantArrayBuilder::operator=(VariantArrayBuilder&&) noexcept = + default; + +void VariantArrayBuilder::AppendNull() { impl_->AppendNull(); } + +#define DEFINE_ARRAY_DIRECT_APPEND(NAME, decl_args, ...) \ + void VariantArrayBuilder::Append##NAME decl_args { \ + impl_->AppendValue([&](VariantValueWriter& writer) { \ + writer.Append(__VA_ARGS__); \ + }); \ + } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_ARRAY_DIRECT_APPEND) +#undef DEFINE_ARRAY_DIRECT_APPEND + +#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 + +void VariantArrayBuilder::AppendEncoded(const EncodedVariantValue& value) { + impl_->AppendEncoded(value); +} + +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)); +} + +VariantObjectBuilder VariantArrayBuilder::StartObject() { + return VariantObjectBuilder(ArrayStartContainer(*impl_)); +} + +VariantListBuilder VariantArrayBuilder::StartList() { + return VariantListBuilder(ArrayStartContainer(*impl_)); +} + +std::shared_ptr VariantArrayBuilder::Finish() { + auto out = impl_->Finish(); + Reset(); + return out; +} + +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; + } + } + + 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; + } + } + + 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"); + } + + 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(); + auto null_bitmap = internal::FinishNullBitmap(validity_builder); + 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); +} + +VariantValueRowBuilder VariantValueArrayBuilder::BindMetadata( + const VariantMetadataView& metadata) { + return {*this, metadata}; +} + +std::shared_ptr VariantValueArrayBuilder::Finish() { + auto out = impl_->Finish(); + Reset(); + return out; +} + +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"); + } + PARQUET_ASSIGN_OR_THROW(auto type, VariantExtensionType::Make(storage->type())); + auto array = ExtensionType::WrapArray(type, std::move(storage)); + return std::static_pointer_cast(array); +} + +std::shared_ptr MakeVariantArrayFromChildren( + std::shared_ptr storage_type, std::vector> children, + std::shared_ptr null_bitmap) { + if (storage_type->id() != Type::STRUCT) { + 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())) { + 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) { + throw ParquetException("Variant storage child ", i, " is null"); + } + if (!children[i]->type()->Equals(struct_type.field(i)->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) { + throw ParquetException("Variant storage child lengths must match"); + } + } + + PARQUET_ASSIGN_OR_THROW(auto storage, + StructArray::Make(std::move(children), struct_type.fields(), + std::move(null_bitmap))); + return MakeVariantArrayFromStorage(std::move(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..380a5eafa2c6 --- /dev/null +++ b/cpp/src/parquet/variant/builder.h @@ -0,0 +1,331 @@ +// 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/type_fwd.h" +#include "parquet/platform.h" + +namespace parquet::variant { + +using ::arrow::Array; +using ::arrow::BinaryViewArray; +using ::arrow::Buffer; +using ::arrow::DataType; +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 PARQUET_EXPORT VariantBuilder { + public: + explicit VariantBuilder(MemoryPool* pool = ::arrow::default_memory_pool(), + bool validate = true); + ~VariantBuilder(); + VariantBuilder(const VariantBuilder&) = delete; + VariantBuilder& operator=(const VariantBuilder&) = delete; + VariantBuilder(VariantBuilder&&) noexcept; + VariantBuilder& operator=(VariantBuilder&&) noexcept; + + 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(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(); + + EncodedVariantValue 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; + + 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, 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); + + /// Commit this nested object into its parent builder. Destroying an unfinished + /// nested builder rolls back the object contents written through this builder. + void Finish(); + + private: + friend class VariantBuilder; + friend class VariantListBuilder; + friend class VariantArrayBuilder; + friend class VariantValueRowBuilder; + + 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; + + 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); + + 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. + void Finish(); + + private: + friend class VariantBuilder; + friend class VariantObjectBuilder; + friend class VariantArrayBuilder; + friend class VariantValueRowBuilder; + + 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; + + 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(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 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(); + + std::shared_ptr 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; + + 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); + +PARQUET_EXPORT +std::shared_ptr 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..ef0c31a675d4 --- /dev/null +++ b/cpp/src/parquet/variant/builder_test.cc @@ -0,0 +1,610 @@ +// 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/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 { + +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; + +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) { + 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, + VariantPrimitiveType 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 + +TEST(TestVariantBuilder, Object) { + VariantBuilder builder; + auto object = builder.StartObject(); + object.AppendVariantNull("c"); + object.AppendVariantNull("b"); + object.AppendVariantNull("a"); + object.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::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; + 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& array = std::get(view.data()); + ASSERT_EQ(3, array.elements().size()); + + auto element = array.GetElement(1); + ASSERT_TRUE(element.has_value()); + ASSERT_EQ(VariantPrimitiveType::kInt32, + 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) { + VariantBuilder builder; + 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_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) { + const std::string uuid(16, '\1'); + VariantBuilder builder; + auto object = builder.StartObject(); + object.AppendInt8("int8", 1); + object.AppendInt64("int64", 2); + object.AppendDouble("double", 3); + object.AppendDecimal4("decimal", ::arrow::Decimal32(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", 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) { + const std::string uuid(16, '\2'); + VariantBuilder builder; + auto list = builder.StartList(); + list.AppendInt8(1); + list.AppendInt64(2); + list.AppendDouble(3); + list.AppendDecimal4(::arrow::Decimal32(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); + 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, ObjectRollback) { + VariantBuilder builder; + auto object = builder.StartObject(); + { + auto child = object.StartObject("value"); + child.AppendString("drop", "x"); + } + object.AppendInt32("value", 1); + object.Finish(); + + auto encoded = builder.Finish(); + 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("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) { + VariantBuilder builder; + 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); + builder.AppendString(std::string(128, 'x')); + auto encoded = builder.Finish(); + ASSERT_GT(pool.total_bytes_allocated(), 0); + + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + 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(); + 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)); + ASSERT_FALSE(array->is_shredded()); + + 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_VIEW, storage.field(0)->type_id()); + 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& metadata_values = + ::arrow::internal::checked_cast(*array->metadata()); + const auto& values = + ::arrow::internal::checked_cast(*array->value()); + 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"); + 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& metadata = + ::arrow::internal::checked_cast(*array->metadata()); + const auto& values = + ::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); + 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, 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& metadata_array = + ::arrow::internal::checked_cast(*array->metadata()); + const auto& value_array = + ::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]); + + 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); + 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)})); + + 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) { + VariantBuilder value; + value.AppendVariantNull(); + 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())}); + + 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 new file mode 100644 index 000000000000..f4734d04c878 --- /dev/null +++ b/cpp/src/parquet/variant/decoding.cc @@ -0,0 +1,526 @@ +// 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/decoding.h" + +#include +#include +#include +#include + +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" +#include "parquet/exception.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; + std::memcpy(&value, data.data() + offset, width); + return bit_util::FromLittleEndian(value); +} + +void CheckAvailable(std::string_view data, size_t offset, size_t size, + std::string_view context) { + if (offset > data.size() || data.size() - offset < size) { + throw ParquetInvalidOrCorruptedFileException("Invalid Variant encoding: truncated ", + context); + } +} + +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) { + case VariantPrimitiveType::kNull: + case VariantPrimitiveType::kBooleanTrue: + case VariantPrimitiveType::kBooleanFalse: + return 0; + case VariantPrimitiveType::kInt8: + return 1; + case VariantPrimitiveType::kInt16: + return 2; + case VariantPrimitiveType::kInt32: + case VariantPrimitiveType::kDate: + case VariantPrimitiveType::kFloat: + return 4; + case VariantPrimitiveType::kInt64: + case VariantPrimitiveType::kDouble: + case VariantPrimitiveType::kTimestampMicros: + case VariantPrimitiveType::kTimestampNTZMicros: + case VariantPrimitiveType::kTimeNTZMicros: + case VariantPrimitiveType::kTimestampNanos: + case VariantPrimitiveType::kTimestampNTZNanos: + return 8; + case VariantPrimitiveType::kDecimal4: + return 5; + case VariantPrimitiveType::kDecimal8: + return 9; + case VariantPrimitiveType::kDecimal16: + return 17; + case VariantPrimitiveType::kUuid: + return 16; + case VariantPrimitiveType::kBinary: + case VariantPrimitiveType::kString: { + CheckAvailable(value, offset, 4, "variable-length size"); + const uint32_t length = ReadLittleEndian(value, offset, 4); + return 4 + static_cast(length); + } + } + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown primitive type"); +} + +size_t ParsePrimitive(std::string_view value, size_t offset, + VariantPrimitiveType primitive) { + if (!internal::IsKnownVariantPrimitive(primitive)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown primitive type ", static_cast(primitive)); + } + + 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]); + internal::ValidateDecimalScale(scale); + } + + if (primitive == VariantPrimitiveType::kString) { + const uint32_t length = ReadLittleEndian(value, offset, 4); + internal::ValidateUtf8(value.substr(offset + 4, length), "primitive string value"); + } + + 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); + const bool is_large = (header & 0x04) != 0; + const size_t count_size = is_large ? 4 : 1; + + size_t offset = 1; + 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, OffsetBytes(offset_count, offset_size), "array offsets"); + + std::vector offsets(offset_count); + for (size_t i = 0; i < offset_count; ++i) { + offsets[i] = ReadLittleEndian(value, offset, offset_size); + offset += offset_size; + } + + if (offsets[0] != 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]) { + 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]; + CheckAvailable(value, values_start, total_value_size, "array values"); + + 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"); + } + } + } + + const size_t consumed = values_start + total_value_size; + if (out != nullptr) { + 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); + 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; + CheckAvailable(value, offset, count_size, "object size"); + const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); + offset += count_size; + + 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; + } + + const size_t offset_count = OffsetCount(num_elements); + CheckAvailable(value, offset, OffsetBytes(offset_count, offset_size), + "object field offsets"); + 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; + } + + const size_t values_start = offset; + const size_t total_value_size = field_offsets[num_elements]; + if (num_elements == 0 && total_value_size != 0) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: empty object must have zero value size"); + } + CheckAvailable(value, values_start, total_value_size, "object values"); + + std::vector object_fields; + if (out != nullptr) { + object_fields.reserve(num_elements); + } + + 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()); + } + 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"); + } + } + + const auto field_offset = field_offsets[i]; + if (field_offset >= total_value_size) { + throw ParquetInvalidOrCorruptedFileException( + "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()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object field offsets must be unique"); + } + if (value_offsets.front() != 0) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object values have leading data"); + } + + 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"); + } + } + } + + 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)}); + } + } + + const size_t consumed = values_start + total_value_size; + if (out != nullptr) { + *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) { + const auto basic_type = VariantValueView::PeekBasicType(value); + const auto header = static_cast(static_cast(value[0]) >> 2); + + switch (basic_type) { + case VariantBasicType::kPrimitive: { + const auto primitive = static_cast(header); + const size_t payload_size = ParsePrimitive(value, 1, primitive); + const size_t consumed = 1 + payload_size; + if (out != nullptr) { + *out = ViewAccess::Make( + value.substr(0, consumed), ViewAccess::Make( + primitive, value.substr(1, payload_size))); + } + return consumed; + } + case VariantBasicType::kShortString: { + 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 = ViewAccess::Make( + value.substr(0, consumed), + ViewAccess::Make(value.substr(1, header))); + } + return consumed; + } + case VariantBasicType::kObject: + return ParseObject(value, metadata, header, out); + case VariantBasicType::kArray: + 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::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( + "Invalid Variant metadata: expected version 1, got ", static_cast(version)); + } + + VariantMetadataView view; + view.sorted_strings_ = (header & internal::kMetadataSortedStringsMask) != 0; + view.offset_size_ = static_cast(((header >> 6) & 0x03) + 1); + + 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_; + 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(offset_count); + for (size_t i = 0; i < offset_count; ++i) { + offsets[i] = + ReadLittleEndian(data, offsets_offset + i * view.offset_size_, view.offset_size_); + } + + if (offsets[0] != 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]) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant metadata: dictionary offsets must be monotonic"); + } + } + + 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; + 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 = + 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( + "Invalid Variant metadata: sorted dictionary strings must be unique and " + "lexicographically sorted"); + } + view.strings_.push_back(string); + } + + 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]; +} + +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; +} + +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::Make(it->value, metadata_.get()); +} + +std::optional VariantObjectView::GetField(size_t index) const { + if (index >= fields_.size()) { + return std::nullopt; + } + 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::Make(elements_[index], metadata_.get()); +} + +VariantValueView VariantValueView::Make(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) { + 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); + if (consumed != value.size()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: trailing bytes after value"); + } +} + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/decoding.h b/cpp/src/parquet/variant/decoding.h new file mode 100644 index 000000000000..f05c77f2827f --- /dev/null +++ b/cpp/src/parquet/variant/decoding.h @@ -0,0 +1,195 @@ +// 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 "parquet/platform.h" +#include "parquet/variant/format.h" + +namespace parquet::variant { + +namespace internal { + +class ViewAccess; + +} + +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 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_; } + /// Number of strings in the metadata dictionary used by object field ids. + uint32_t dictionary_size() const { return static_cast(strings_.size()); } + + /// 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: + std::string_view metadata_; + bool sorted_strings_ = false; + uint8_t offset_size_ = 0; + std::vector strings_; +}; + +class PARQUET_EXPORT VariantPrimitiveView { + public: + 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: + 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: + static constexpr VariantBasicType kBasicType = VariantBasicType::kObject; + + const std::vector& fields() const { return fields_; } + 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: + 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 { + public: + using Data = std::variant; + + /// 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 + /// 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; + + /// 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); + + static void Validate(std::string_view value, const VariantMetadataView& metadata); + + std::string_view value() const { return value_; } + 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)) {} + + std::string_view value_; + Data data_; +}; + +} // namespace parquet::variant diff --git a/cpp/src/parquet/variant/decoding_test.cc b/cpp/src/parquet/variant/decoding_test.cc new file mode 100644 index 000000000000..303096ae7193 --- /dev/null +++ b/cpp/src/parquet/variant/decoding_test.cc @@ -0,0 +1,266 @@ +// 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/decoding.h" +#include "parquet/variant/builder.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); + +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(); + auto encoded = builder.Finish(); + 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 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 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()); + }; + + check([](VariantBuilder& b) { b.AppendVariantNull(); }, VariantPrimitiveType::kNull); + check([](VariantBuilder& b) { b.AppendBoolean(true); }, + VariantPrimitiveType::kBooleanTrue); + 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(::arrow::Decimal32(123), 2); }, + VariantPrimitiveType::kDecimal4); + check([](VariantBuilder& b) { b.AppendDecimal8(::arrow::Decimal64(123), 2); }, + VariantPrimitiveType::kDecimal8); + 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); + check([](VariantBuilder& b) { b.AppendTimeNTZMicros(1); }, + VariantPrimitiveType::kTimeNTZMicros); + check([](VariantBuilder& b) { b.AppendTimestampMicros(1, true); }, + VariantPrimitiveType::kTimestampMicros); + check([](VariantBuilder& b) { b.AppendTimestampMicros(1, false); }, + VariantPrimitiveType::kTimestampNTZMicros); + check([](VariantBuilder& b) { b.AppendTimestampNanos(1, true); }, + VariantPrimitiveType::kTimestampNanos); + check([](VariantBuilder& b) { b.AppendTimestampNanos(1, false); }, + VariantPrimitiveType::kTimestampNTZNanos); + check( + [&](VariantBuilder& b) { + b.AppendUuid(std::string_view(uuid.data(), uuid.size())); + }, + VariantPrimitiveType::kUuid); +} + +TEST(TestVariantEncoding, ShortString) { + VariantBuilder builder; + builder.AppendShortString("abc"); + 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::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("\x02\x00\x00"sv), + ParquetInvalidOrCorruptedFileException); + // Dictionary offsets are truncated. + ASSERT_THROW(VariantMetadataView::Make("\x01\x01\x00"sv), + ParquetInvalidOrCorruptedFileException); + // Dictionary string payload is not valid UTF-8. + ASSERT_THROW(VariantMetadataView::Make("\x01\x01\x00\x01\xff"sv), + ParquetInvalidOrCorruptedFileException); +} + +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 + 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("\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("\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) { + VariantBuilder builder; + builder.AppendVariantNull(); + auto encoded = builder.Finish(); + auto metadata = VariantMetadataView::Make(std::string_view{*encoded.metadata}); + + VariantValueView::Validate(std::string_view{*encoded.value}, metadata); + std::string invalid_value(std::string_view{*encoded.value}); + invalid_value.push_back('\0'); + ASSERT_THROW(VariantValueView::Validate(invalid_value, metadata), + 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(); + 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_THROW(VariantValueView::Make("\x54"sv, metadata), + ParquetInvalidOrCorruptedFileException); + // Short string payload is not valid UTF-8. + ASSERT_THROW(VariantValueView::Make("\x05\xff"sv, metadata), + ParquetInvalidOrCorruptedFileException); + // Null primitive has trailing bytes after the encoded value. + ASSERT_THROW(VariantValueView::Make("\x00\x00"sv, metadata), + ParquetInvalidOrCorruptedFileException); + // Object references field id 0, but the metadata dictionary is empty. + 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. 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); +} + +TEST(TestVariantEncoding, InvalidObject) { + VariantBuilder builder; + 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 + // - 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 auto duplicate_offsets = "\x02\x02\x00\x01\x00\x00\x01\x00"sv; + ASSERT_THROW(VariantValueView::Validate(duplicate_offsets, metadata), + ParquetInvalidOrCorruptedFileException); +} + +} // namespace parquet::variant 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/format_internal.h b/cpp/src/parquet/variant/format_internal.h new file mode 100644 index 000000000000..ea74fda27117 --- /dev/null +++ b/cpp/src/parquet/variant/format_internal.h @@ -0,0 +1,152 @@ +// 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/type_fwd.h" +#include "arrow/util/utf8.h" +#include "parquet/exception.h" +#include "parquet/variant/format.h" + +namespace parquet::variant { + +class VariantPrimitiveView; +class VariantShortStringView; +class VariantObjectView; +class VariantArrayView; + +namespace internal { + +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 void ValidateDecimalScale(uint8_t scale) { + if (scale > 38) { + throw ParquetInvalidOrCorruptedFileException("Invalid Variant decimal scale ", scale, + " exceeds 38"); + } +} + +inline void ValidateUtf8(std::string_view value, std::string_view context) { + ::arrow::util::InitializeUTF8(); + if (!::arrow::util::ValidateUTF8(reinterpret_cast(value.data()), + value.size())) { + throw ParquetInvalidOrCorruptedFileException("Invalid Variant encoding: ", context, + " is not valid UTF-8"); + } +} + +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, ::arrow::Decimal32) +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal8, ::arrow::Decimal64) +VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF(Decimal16, ::arrow::Decimal128) + +#undef VARIANT_DECIMAL_PRIMITIVE_TRAITS_DEF + +template +concept DecimalVariantPrimitive = + requires { typename VariantDecimalPrimitiveTraits::CType; }; + +template +concept LengthPrefixedVariantPrimitive = + type == VariantPrimitiveType::kBinary || type == VariantPrimitiveType::kString; + +template +concept UuidVariantPrimitive = type == VariantPrimitiveType::kUuid; + +} // namespace internal + +} // 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..cecf81f4333f --- /dev/null +++ b/cpp/src/parquet/variant/meson.build @@ -0,0 +1,44 @@ +# 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', + 'decoding_test.cc', + 'parquet_testing_test.cc', + 'shred_test.cc', + 'test_util_internal.cc', + 'type_test.cc', + 'unshred_test.cc', + 'validate_test.cc', + ], + dependencies: parquet_test_dep, +) +test('parquet-variant-test', exc) + +install_headers( + [ + '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 new file mode 100644 index 000000000000..9c18c0159d88 --- /dev/null +++ b/cpp/src/parquet/variant/parquet_testing_test.cc @@ -0,0 +1,435 @@ +// 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/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/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" + +#include +#include + +namespace parquet::variant { + +using ::arrow::internal::checked_pointer_cast; +using internal::AssertEncodedRow; +using internal::RoundTripVariantArray; + +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; + } + + 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(); + + 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); + AssertEncodedRow(*unshredded, static_cast(row), expected); + } +} + +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); + 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]); + AssertEncodedRow(*unshredded, /*row=*/0, expected); +} + +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; + } + + auto unshredded = UnshredVariantArray(*variant_array); + ASSERT_EQ(static_cast(test_case.expected_files.size()), + unshredded->length()); + for (int64_t row = 0; row < unshredded->length(); ++row) { + SCOPED_TRACE(row); + if (!test_case.expected_files[row].has_value()) { + ASSERT_TRUE(unshredded->IsNull(row)); + continue; + } + + auto expected = LoadEncodedVariantBinaryFile(*test_case.expected_files[row]); + AssertEncodedRow(*unshredded, row, expected); + } + } + 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/shred.cc b/cpp/src/parquet/variant/shred.cc new file mode 100644 index 000000000000..0364c2db1f2f --- /dev/null +++ b/cpp/src/parquet/variant/shred.cc @@ -0,0 +1,864 @@ +// 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 "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: + 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); + 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 new file mode 100644 index 000000000000..e32f63e8e908 --- /dev/null +++ b/cpp/src/parquet/variant/slot_internal.cc @@ -0,0 +1,704 @@ +// 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; + +#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 { \ + 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 (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); + } + } + + 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; + } + if (target == nullptr) { + 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; + } + if (target == nullptr) { + 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 new file mode 100644 index 000000000000..a612e6023cb5 --- /dev/null +++ b/cpp/src/parquet/variant/test_util_internal.cc @@ -0,0 +1,169 @@ +// 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 + +#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 { + +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(); +} + +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(); +} + +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(); + auto encoded = builder.Finish(); + return std::move(encoded.metadata); +} + +EncodedVariantValue Int8Variant(int8_t value) { + VariantBuilder builder; + builder.AppendInt8(value); + return builder.Finish(); +} + +std::shared_ptr BinaryArrayFromValues( + const std::vector>& values) { + ::arrow::BinaryBuilder builder; + for (const auto& value : values) { + ARROW_EXPECT_OK(builder.AppendOrNull(value)); + } + std::shared_ptr out; + ARROW_EXPECT_OK(builder.Finish(&out)); + 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 new file mode 100644 index 000000000000..62fca0f518e3 --- /dev/null +++ b/cpp/src/parquet/variant/test_util_internal.h @@ -0,0 +1,80 @@ +// 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 "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; +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, + 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()); + +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(); + +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/type_test.cc b/cpp/src/parquet/variant/type_test.cc new file mode 100644 index 000000000000..d5b00802778c --- /dev/null +++ b/cpp/src/parquet/variant/type_test.cc @@ -0,0 +1,194 @@ +// 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_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), + 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))); + } +} + +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..2c5e2c103cb6 --- /dev/null +++ b/cpp/src/parquet/variant/unshred.cc @@ -0,0 +1,82 @@ +// 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/extension/parquet_variant.h" +#include "arrow/type.h" +#include "arrow/util/bit_block_counter.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::extension::VariantArray; + +namespace { + +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 + +std::shared_ptr<::arrow::extension::VariantArray> UnshredVariantArray( + const ::arrow::extension::VariantArray& array, ::arrow::MemoryPool* pool) { + auto metadata_array = array.metadata(); + auto value_array = array.value(); + auto typed_array = array.typed_value(); + + if (!array.is_shredded()) { + return std::make_shared(array.data()); + } + + 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)}, + internal::NullBitmapForOutput(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..cccdbe31c4e2 --- /dev/null +++ b/cpp/src/parquet/variant/unshred.h @@ -0,0 +1,40 @@ +// 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 { + +/// 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..e1944c5833eb --- /dev/null +++ b/cpp/src/parquet/variant/unshred_test.cc @@ -0,0 +1,306 @@ +// 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::AssertEncodedRow; +using internal::AssertUnshreddedValue; +using internal::BinaryArrayFromValues; +using internal::MakeInt64FieldGroup; + +namespace { + +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::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); + 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) { + 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); + + 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>(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)); +} + +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); + + AssertUnshreddedValue(*array, /*row=*/0, 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); + AssertUnshreddedValue(*array, /*row=*/0, 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); + AssertUnshreddedValue(*array, /*row=*/0, 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); + AssertUnshreddedValue(*array, /*row=*/0, 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); + AssertUnshreddedValue(*array, /*row=*/0, expected); +} + +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)); + AssertEncodedRow(*unshredded, /*row=*/1, encoded); +} + +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(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(UnshredVariantArray(*array), ParquetInvalidOrCorruptedFileException); +} + +} // 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..e07269937813 --- /dev/null +++ b/cpp/src/parquet/variant/validate.cc @@ -0,0 +1,216 @@ +// 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 "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/type.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/array_internal.h" +#include "parquet/variant/decoding.h" +#include "parquet/variant/slot_internal.h" + +namespace parquet::variant { + +namespace { + +using ::arrow::Array; +using ::arrow::Buffer; +using ::arrow::ChunkedArray; +using ::arrow::ExtensionArray; +using ::arrow::ExtensionType; +using ::arrow::MemoryPool; +using ::arrow::StructArray; +using ::arrow::extension::kVariantExtensionName; +using ::arrow::extension::VariantArray; +using ::arrow::internal::checked_cast; + +struct VariantValidationPlan { + std::shared_ptr array; + std::vector children; +}; + +template +void ValidateVariantArrayRows(const VariantArray& array, + 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) { + 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); + }); +} + +std::optional 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()) { + 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: { + auto values = internal::ValuesArray(*array); + auto child_plan = BuildVariantValidationPlan(std::move(values)); + if (!child_plan.has_value()) { + return std::nullopt; + } + return VariantValidationPlan{.array = std::move(array), + .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)); + 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)}; + } + } +} + +template +void ValidateVariantPlan(const VariantValidationPlan& plan, MemoryPool* pool, + const std::shared_ptr& valid_rows) { + switch (plan.array->type_id()) { + case ::arrow::Type::EXTENSION: + ValidateVariantArrayRows(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()) { + 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) { + ValidateVariantPlan(child, pool, child_valid_rows); + } + return; + } + 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); + auto values = internal::ValuesArray(*plan.array); + PARQUET_ASSIGN_OR_THROW(auto values_valid_rows, + ::arrow::AllocateEmptyBitmap(values->length(), pool)); + 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); + return; + } + default: + for (const auto& child : plan.children) { + 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); + if (plan.has_value()) { + 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 new file mode 100644 index 000000000000..ae3fbf58f263 --- /dev/null +++ b/cpp/src/parquet/variant/validate.h @@ -0,0 +1,33 @@ +// 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 parquet::variant { + +/// 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 new file mode 100644 index 000000000000..44a264ccd1b3 --- /dev/null +++ b/cpp/src/parquet/variant/validate_test.cc @@ -0,0 +1,143 @@ +// 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 "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/exception.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::Int8Variant; +using internal::VariantTable; +using internal::WriteVariantTable; + +TEST(TestVariantValidate, ListView) { + 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(*::arrow::ArrayFromJSON(::arrow::int32(), "[0]"), + *::arrow::ArrayFromJSON(::arrow::int32(), "[1]"), + *variant_array)); + ::arrow::ChunkedArray valid_data{valid_list}; + ValidateVariants(valid_data); + + 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), + ParquetInvalidOrCorruptedFileException); +} + +TEST(TestVariantValidate, DictionaryMetadata) { + 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); + ValidateVariants(*column); +} + +TEST(TestVariantValidate, ReadDictionaryOption) { + 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); + ValidateVariants(*column); +} + +} // namespace parquet::variant