diff --git a/.github/workflows/cpp-linter.yml b/.github/workflows/cpp-linter.yml index be59a610d..6f598b02b 100644 --- a/.github/workflows/cpp-linter.yml +++ b/.github/workflows/cpp-linter.yml @@ -72,6 +72,8 @@ jobs: run: | mkdir build && cd build cmake .. -G Ninja -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \ + -DICEBERG_REQUIRE_ARROW_VARIANT=ON \ + -DICEBERG_USE_BUNDLED_ARROW=ON \ -DCMAKE_C_COMPILER_LAUNCHER=sccache \ -DCMAKE_CXX_COMPILER_LAUNCHER=sccache \ -DICEBERG_BUILD_SQL_CATALOG=ON \ diff --git a/.github/workflows/sanitizer_test.yml b/.github/workflows/sanitizer_test.yml index b2d1e011a..f51b6ca10 100644 --- a/.github/workflows/sanitizer_test.yml +++ b/.github/workflows/sanitizer_test.yml @@ -63,6 +63,8 @@ jobs: run: | mkdir build && cd build cmake .. -G Ninja -DCMAKE_BUILD_TYPE=Debug -DICEBERG_ENABLE_ASAN=ON -DICEBERG_ENABLE_UBSAN=ON \ + -DICEBERG_REQUIRE_ARROW_VARIANT=ON \ + -DICEBERG_USE_BUNDLED_ARROW=ON \ -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache cmake --build . --verbose - name: Save sccache diff --git a/CMakeLists.txt b/CMakeLists.txt index 7a9c3b2ff..b267ae9e5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -53,6 +53,10 @@ option(ICEBERG_BUILD_BENCHMARKS "Build benchmarks" OFF) option(ICEBERG_BUILD_BUNDLE "Build the battery included library" ON) option(ICEBERG_ENCRYPTION "Build table encryption support (Parquet modular encryption, AES GCM streams)" OFF) +option(ICEBERG_USE_BUNDLED_ARROW + "Use the pinned Arrow source instead of an installed Arrow" OFF) +option(ICEBERG_REQUIRE_ARROW_VARIANT + "Require Arrow variant APIs when building the bundle" ON) option(ICEBERG_BUILD_REST "Build rest catalog client" ON) option(ICEBERG_BUILD_REST_INTEGRATION_TESTS "Build rest catalog integration tests" OFF) option(ICEBERG_BUILD_HIVE "Build hive (HMS) catalog client" OFF) diff --git a/ci/scripts/build_iceberg.sh b/ci/scripts/build_iceberg.sh index 640b77269..a9db77ce2 100755 --- a/ci/scripts/build_iceberg.sh +++ b/ci/scripts/build_iceberg.sh @@ -168,6 +168,10 @@ if [[ -n "${ICEBERG_EXTRA_CMAKE_ARGS:-}" ]]; then CMAKE_ARGS+=("${EXTRA_CMAKE_ARGS[@]}") fi +# CI must compile and run the complete variant tests, even if a runner happens +# to discover an older system Arrow installation. +CMAKE_ARGS+=("-DICEBERG_USE_BUNDLED_ARROW=ON" "-DICEBERG_REQUIRE_ARROW_VARIANT=ON") + cmake "${CMAKE_ARGS[@]}" ${source_dir} cmake --build . --target install diff --git a/cmake_modules/IcebergThirdpartyToolchain.cmake b/cmake_modules/IcebergThirdpartyToolchain.cmake index f94ad7031..0b00da848 100644 --- a/cmake_modules/IcebergThirdpartyToolchain.cmake +++ b/cmake_modules/IcebergThirdpartyToolchain.cmake @@ -89,22 +89,21 @@ endfunction() # ICEBERG_CPR_URL - cpr tarball URL # # Example usage: -# export ICEBERG_ARROW_URL="https://your-mirror.com/apache-arrow-25.0.0.tar.gz" +# export ICEBERG_ARROW_URL="https://your-mirror.com/arrow-0f16412bb1c39611aeed7393eba09e637f4d86fa.tar.gz" # cmake -S . -B build # -set(ICEBERG_ARROW_BUILD_VERSION "25.0.0") +# Temporary pin for apache/arrow#50252. Replace with an Apache Arrow release +# containing the variant APIs once that work has merged and been released. +set(ICEBERG_ARROW_BUILD_COMMIT "0f16412bb1c39611aeed7393eba09e637f4d86fa") set(ICEBERG_ARROW_BUILD_SHA256_CHECKSUM - "12afc2dc8137bdd4a68876cec939f664c9d55cfc7b75f55b45163ebb4e344d81") + "4e184655c31b4b28cf7d0cd1034f9b3891d5d861520d427a90726b28e52e72f9") if(DEFINED ENV{ICEBERG_ARROW_URL}) set(ARROW_SOURCE_URL "$ENV{ICEBERG_ARROW_URL}") else() set(ARROW_SOURCE_URL - "https://www.apache.org/dyn/closer.lua?action=download&filename=/arrow/arrow-${ICEBERG_ARROW_BUILD_VERSION}/apache-arrow-${ICEBERG_ARROW_BUILD_VERSION}.tar.gz" - "https://downloads.apache.org/arrow/arrow-${ICEBERG_ARROW_BUILD_VERSION}/apache-arrow-${ICEBERG_ARROW_BUILD_VERSION}.tar.gz" - "https://archive.apache.org/dist/arrow/arrow-${ICEBERG_ARROW_BUILD_VERSION}/apache-arrow-${ICEBERG_ARROW_BUILD_VERSION}.tar.gz" - ) + "https://codeload.github.com/HuaHuaY/arrow/tar.gz/${ICEBERG_ARROW_BUILD_COMMIT}") endif() set(ICEBERG_NANOARROW_BUILD_VERSION "0.9.0") @@ -180,12 +179,67 @@ macro(resolve_openssl_for_encryption) find_package(OpenSSL REQUIRED) endmacro() +function(iceberg_check_arrow_variant) + if(NOT ICEBERG_REQUIRE_ARROW_VARIANT) + return() + endif() + + include(CheckCXXSourceCompiles) + # This also works before bundled Arrow has been built: compile the API probe + # without linking against its libraries. + set(CMAKE_TRY_COMPILE_TARGET_TYPE STATIC_LIBRARY) + set(CMAKE_REQUIRED_INCLUDES) + foreach(dependency Arrow Parquet) + if(dependency STREQUAL "Arrow") + set(component arrow) + else() + set(component parquet) + endif() + if(TARGET ${dependency}::${component}_static) + set(dependency_target ${dependency}::${component}_static) + else() + set(dependency_target ${dependency}::${component}_shared) + endif() + get_target_property(dependency_includes ${dependency_target} + INTERFACE_INCLUDE_DIRECTORIES) + list(APPEND CMAKE_REQUIRED_INCLUDES ${dependency_includes}) + endforeach() + # Recheck when an existing build switches Arrow installations or source pins. + unset(ICEBERG_ARROW_HAS_VARIANT CACHE) + check_cxx_source_compiles("#include + #include + #include + #include + #include + #include + #include + using Shred = decltype(&parquet::variant::ShredVariantArray); + using Unshred = decltype(&parquet::variant::UnshredVariantArray); + using Validate = decltype(&parquet::variant::ValidateVariants); + using ReaderValidation = decltype(&parquet::ArrowReaderProperties::set_variant_validation_enabled); + using WriterValidation = decltype(&parquet::ArrowWriterProperties::Builder::set_variant_validation_enabled); + using IsShredded = decltype(&arrow::extension::VariantArray::is_shredded); + int main() { return 0; }" + ICEBERG_ARROW_HAS_VARIANT) + if(NOT ICEBERG_ARROW_HAS_VARIANT) + message(FATAL_ERROR "Arrow/Parquet lacks the variant APIs required by ICEBERG_REQUIRE_ARROW_VARIANT. " + "Use the pinned bundled Arrow (ICEBERG_USE_BUNDLED_ARROW=ON) " + "or a compatible Arrow installation. Set ICEBERG_REQUIRE_ARROW_VARIANT=OFF " + "only for builds that do not need variant I/O.") + endif() +endfunction() + function(resolve_arrow_dependency) if(ICEBERG_ENCRYPTION) resolve_openssl_for_encryption() endif() prepare_fetchcontent() + if(ICEBERG_USE_BUNDLED_ARROW) + # This override is scoped to Arrow; other dependencies can still use + # installed packages, including the Windows CI runner's vcpkg packages. + set(FETCHCONTENT_TRY_FIND_PACKAGE_MODE NEVER) + endif() # Prevent Arrow from injecting -Werror into CMAKE_CXX_FLAGS_DEBUG via # arrow_add_werror_if_debug(). PRODUCTION level only adds standard warnings. @@ -340,6 +394,8 @@ function(resolve_arrow_dependency) list(APPEND ICEBERG_SYSTEM_DEPENDENCIES OpenSSL) endif() + iceberg_check_arrow_variant() + set(ICEBERG_SYSTEM_DEPENDENCIES ${ICEBERG_SYSTEM_DEPENDENCIES} PARENT_SCOPE) diff --git a/mkdocs/docs/getting-started.md b/mkdocs/docs/getting-started.md index ab0741fd9..5daa4f4dd 100644 --- a/mkdocs/docs/getting-started.md +++ b/mkdocs/docs/getting-started.md @@ -51,7 +51,7 @@ cmake --install build ### Bundle Library (with vendored dependencies) ```bash -cmake -S . -B build -G Ninja -DCMAKE_INSTALL_PREFIX=/path/to/install -DICEBERG_BUILD_BUNDLE=ON +cmake -S . -B build -G Ninja -DCMAKE_INSTALL_PREFIX=/path/to/install -DICEBERG_BUILD_BUNDLE=ON -DFETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER cmake --build build cmake --install build ``` @@ -72,6 +72,8 @@ cmake --install build | `ICEBERG_BUILD_SHARED` | `OFF` | Build shared library | | `ICEBERG_BUILD_TESTS` | `ON` | Build tests | | `ICEBERG_BUILD_BUNDLE` | `ON` | Build the battery-included library | +| `ICEBERG_USE_BUNDLED_ARROW` | `OFF` | Force the pinned Arrow source; CI enables this | +| `ICEBERG_REQUIRE_ARROW_VARIANT` | `ON` | Require Arrow variant APIs in the bundle; CI enforces this | | `ICEBERG_BUILD_REST` | `ON` | Build REST catalog client | | `ICEBERG_BUILD_REST_INTEGRATION_TESTS` | `OFF` | Build REST catalog integration tests | | `ICEBERG_BUILD_HIVE` | `OFF` | Build Hive (HMS) catalog client | @@ -159,6 +161,9 @@ If you experience network issues when downloading dependencies, you can override Example: ```bash -export ICEBERG_ARROW_URL="https://your-mirror.com/apache-arrow-22.0.0.tar.gz" +export ICEBERG_ARROW_URL="https://your-mirror.com/arrow-0f16412bb1c39611aeed7393eba09e637f4d86fa.tar.gz" cmake -S . -B build -G Ninja ``` + +The Arrow mirror must serve the same pinned source archive: its SHA256 is checked +against the value in `cmake_modules/IcebergThirdpartyToolchain.cmake`. diff --git a/mkdocs/docs/library.md b/mkdocs/docs/library.md index 6b769425a..ac4c55ed1 100644 --- a/mkdocs/docs/library.md +++ b/mkdocs/docs/library.md @@ -42,6 +42,43 @@ are supplied by the `iceberg_bundle` library. This library combines `iceberg_data` with the Arrow C++, Avro, Parquet, and Arrow filesystem integrations. It gives applications one link-time dependency. +Variant values use the `arrow.parquet.variant` extension with binary `metadata` +and `value` storage. Call `iceberg::arrow::RegisterAll()` before importing these +schemas into Arrow C++. Parquet readers and writers also register the extension +when opened. + +Parquet variant I/O requires the Arrow variant builders, validation, and shredding +utilities introduced by [Arrow PR #50252](https://github.com/apache/arrow/pull/50252). +The bundled dependency temporarily pins Arrow revision +`0f16412bb1c39611aeed7393eba09e637f4d86fa` from that PR and verifies the source +archive's SHA256. It will be replaced by an Apache Arrow release containing these +APIs. Use `ICEBERG_USE_BUNDLED_ARROW=ON` to force this Arrow source, or provide a +compatible Arrow installation through `CMAKE_PREFIX_PATH`. + +`ICEBERG_REQUIRE_ARROW_VARIANT` defaults to `ON`: configuration fails when the +selected Arrow/Parquet lacks the required APIs. CI explicitly uses the pinned +Arrow source and keeps this option enabled; the variant tests fail to compile +if their APIs are missing. To use an older Arrow installation for other +functionality, explicitly set this option +to `OFF`; Parquet variant I/O then returns `NotSupported`. + +Variants are written unshredded by default. To shred a field, supply a logical +target keyed by its Iceberg field ID: + +```cpp +iceberg::WriterOptions options; +// Set the path, schema, and FileIO as usual. +options.variant_shredding_types.emplace(7, iceberg::int64()); +``` + +Targets can also contain structs and lists. Their field IDs are ignored; they +describe the requested layout inside the variant. The layout stays fixed for +the file, and incompatible values remain in residual storage. Input batches +always use unshredded variants. Readers reconstruct the same unshredded schema +from either file layout. Validation stays enabled, and a null column value is +distinct from an encoded variant null. Avro variant data and variant bounds +metrics are not supported. + ### `iceberg_rest`, `iceberg_hive`, `iceberg_sql_catalog` These optional catalog libraries cover REST, Hive Metastore, and SQL databases. diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index c751cd7a6..692414588 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -323,6 +323,7 @@ if(ICEBERG_BUILD_BUNDLE) parquet/parquet_reader.cc parquet/parquet_register.cc parquet/parquet_schema_util.cc + parquet/parquet_variant_util.cc parquet/parquet_writer.cc parquet/row_number_reader.cc) diff --git a/src/iceberg/arrow/arrow_register.cc b/src/iceberg/arrow/arrow_register.cc index 4fedc914c..1eeeeb8ab 100644 --- a/src/iceberg/arrow/arrow_register.cc +++ b/src/iceberg/arrow/arrow_register.cc @@ -23,6 +23,10 @@ #include #include +#include +#include +#include + #include "iceberg/arrow/arrow_io_util.h" #include "iceberg/arrow/s3/s3_properties.h" #include "iceberg/file_io_registry.h" @@ -57,6 +61,16 @@ void RegisterS3FileIO() { void EnsureArrowFileIOsRegistered() { static std::once_flag flag; std::call_once(flag, []() { + // Unlike UUID, Arrow does not register its variant extension by default. + // Register it so C Data imports reconstruct variant arrays from Iceberg's + // canonical metadata/value storage. + if (!::arrow::GetExtensionType("arrow.parquet.variant")) { + auto variant = ::arrow::extension::variant( + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary(), false)})); + (void)::arrow::RegisterExtensionType( + std::static_pointer_cast<::arrow::ExtensionType>(variant)); + } RegisterLocalFileIO(); RegisterS3FileIO(); }); diff --git a/src/iceberg/arrow/arrow_register.h b/src/iceberg/arrow/arrow_register.h index 7a67be516..a74362475 100644 --- a/src/iceberg/arrow/arrow_register.h +++ b/src/iceberg/arrow/arrow_register.h @@ -26,7 +26,7 @@ namespace iceberg::arrow { -/// \brief Register Arrow FileIOs and Arrow-backed C Data utilities. +/// \brief Register Arrow FileIOs, the variant extension, and C Data utilities. /// /// This operation is idempotent and safe to call multiple times. ICEBERG_BUNDLE_EXPORT void RegisterAll(); diff --git a/src/iceberg/file_writer.h b/src/iceberg/file_writer.h index 4b1045a40..6579e402f 100644 --- a/src/iceberg/file_writer.h +++ b/src/iceberg/file_writer.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include "iceberg/arrow_c_data.h" @@ -85,6 +86,11 @@ struct ICEBERG_EXPORT WriterOptions { std::shared_ptr metrics_config = MetricsConfig::Default(); /// \brief Format-specific or implementation-specific properties. WriterProperties properties; + /// \brief Optional logical shredding targets for Parquet variant fields, by field ID. + /// Targets may be supported primitive, struct, or list types. Their field IDs are + /// ignored. Unconfigured fields are written unshredded. Input arrays always use the + /// unshredded variant representation; the target layout stays fixed for this file. + std::unordered_map> variant_shredding_types; }; /// \brief Base writer class to write data from different file formats. diff --git a/src/iceberg/parquet/parquet_data_util.cc b/src/iceberg/parquet/parquet_data_util.cc index 73597cbfb..e4a93119d 100644 --- a/src/iceberg/parquet/parquet_data_util.cc +++ b/src/iceberg/parquet/parquet_data_util.cc @@ -30,7 +30,9 @@ #include "iceberg/arrow/literal_util_internal.h" #include "iceberg/metadata_columns.h" #include "iceberg/parquet/parquet_data_util_internal.h" +#include "iceberg/parquet/parquet_variant_util_internal.h" #include "iceberg/schema.h" +#include "iceberg/schema_internal.h" #include "iceberg/schema_util.h" #include "iceberg/type.h" #include "iceberg/util/checked_cast.h" @@ -117,9 +119,8 @@ Result> ProjectNestedArray( Result> MakeNullArray( const std::shared_ptr<::arrow::DataType>& type, int64_t length, ::arrow::MemoryPool* pool) { - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto builder, ::arrow::MakeBuilder(type, pool)); - ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendNulls(length)); - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto array, builder->Finish()); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto array, + ::arrow::MakeArrayOfNull(type, length, pool)); return array; } @@ -127,6 +128,11 @@ Result> ProjectPrimitiveArray( const std::shared_ptr<::arrow::Array>& array, const std::shared_ptr<::arrow::DataType>& output_arrow_type, ::arrow::MemoryPool* pool) { + if (output_arrow_type->id() == ::arrow::Type::EXTENSION && + static_cast(*output_arrow_type).extension_name() == + kArrowVariantExtensionName) { + return ProjectVariantArray(array, output_arrow_type, pool); + } if (array->type()->Equals(output_arrow_type)) { return array; } diff --git a/src/iceberg/parquet/parquet_reader.cc b/src/iceberg/parquet/parquet_reader.cc index 549a3d3ae..a85382b08 100644 --- a/src/iceberg/parquet/parquet_reader.cc +++ b/src/iceberg/parquet/parquet_reader.cc @@ -35,15 +35,18 @@ #include #include #include +#include #include #include #include "iceberg/arrow/arrow_io_internal.h" +#include "iceberg/arrow/arrow_register.h" #include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow/metadata_column_util_internal.h" #include "iceberg/parquet/parquet_data_util_internal.h" #include "iceberg/parquet/parquet_register.h" #include "iceberg/parquet/parquet_schema_util_internal.h" +#include "iceberg/parquet/parquet_variant_util_internal.h" #include "iceberg/parquet/row_number_reader_internal.h" #include "iceberg/result.h" #include "iceberg/schema_internal.h" @@ -277,9 +280,11 @@ class ParquetReader::Impl { public: // Open the Parquet reader with the given options Status Open(const ReaderOptions& options) { + arrow::RegisterAll(); if (options.projection == nullptr) { return InvalidArgument("Projected schema is required by Parquet reader"); } + ICEBERG_RETURN_UNEXPECTED(CheckVariantSupport(*options.projection)); read_schema_ = options.projection; @@ -289,6 +294,11 @@ class ParquetReader::Impl { arrow_reader_properties.set_batch_size( options.properties.Get(ReaderProperties::kBatchSize)); arrow_reader_properties.set_arrow_extensions_enabled(true); +#if ICEBERG_HAS_ARROW_VARIANT + // Arrow's per-column validator cannot see null ancestors of required variants. + // Validate the complete batch with its parent validity in Next() instead. + arrow_reader_properties.set_variant_validation_enabled(false); +#endif use_large_list_ = options.properties.Get(ReaderProperties::kArrowUseLargeList); if (use_large_list_) { arrow_reader_properties.set_list_type(::arrow::Type::LARGE_LIST); @@ -323,6 +333,7 @@ class ParquetReader::Impl { if (!batch) { return std::nullopt; } + ICEBERG_ASSIGN_OR_RAISE(batch, NormalizeVariantReadBatch(batch, pool_)); std::shared_ptr<::arrow::Int64Array> row_positions; if (context_->row_number_reader_) { @@ -496,9 +507,16 @@ class ParquetReader::Impl { std::unique_ptr context_; }; +ParquetReader::ParquetReader() = default; ParquetReader::~ParquetReader() = default; -Result> ParquetReader::Next() { return impl_->Next(); } +Result> ParquetReader::Next() { + try { + return impl_->Next(); + } catch (const ::parquet::ParquetException& error) { + return Invalid("Cannot read Parquet batch: {}", error.what()); + } +} Result ParquetReader::Schema() { return impl_->Schema(); } @@ -508,7 +526,11 @@ Result> ParquetReader::Metadata() { Status ParquetReader::Open(const ReaderOptions& options) { impl_ = std::make_unique(); - return impl_->Open(options); + try { + return impl_->Open(options); + } catch (const ::parquet::ParquetException& error) { + return InvalidSchema("Cannot open Parquet reader: {}", error.what()); + } } Status ParquetReader::Close() { return impl_->Close(); } diff --git a/src/iceberg/parquet/parquet_reader.h b/src/iceberg/parquet/parquet_reader.h index 3a8e57cf3..f3b656386 100644 --- a/src/iceberg/parquet/parquet_reader.h +++ b/src/iceberg/parquet/parquet_reader.h @@ -30,7 +30,7 @@ namespace iceberg::parquet { /// \brief A reader that reads ArrowArray from Parquet files. class ICEBERG_BUNDLE_EXPORT ParquetReader : public Reader { public: - ParquetReader() = default; + ParquetReader(); ~ParquetReader() override; diff --git a/src/iceberg/parquet/parquet_schema_util.cc b/src/iceberg/parquet/parquet_schema_util.cc index 2a230fe4c..abcc10ab9 100644 --- a/src/iceberg/parquet/parquet_schema_util.cc +++ b/src/iceberg/parquet/parquet_schema_util.cc @@ -17,13 +17,16 @@ * under the License. */ +#include #include #include #include #include #include +#include #include +#include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/constants.h" #include "iceberg/metadata_columns.h" #include "iceberg/name_mapping.h" @@ -196,7 +199,9 @@ Result<::parquet::schema::NodePtr> ToParquetPrimitive( physical_type, type_length, field_id); } -Result<::parquet::schema::NodePtr> ToParquetNode(const SchemaField& field) { +Result<::parquet::schema::NodePtr> ToParquetNode( + const SchemaField& field, + const std::shared_ptr<::arrow::DataType>& storage_type = nullptr) { const auto repetition = field.optional() ? ::parquet::Repetition::OPTIONAL : ::parquet::Repetition::REQUIRED; const auto& type = *field.type(); @@ -213,8 +218,12 @@ Result<::parquet::schema::NodePtr> ToParquetNode(const SchemaField& field) { } ::parquet::schema::NodeVector children; children.reserve(struct_type.fields().size()); - for (const auto& child_field : struct_type.fields()) { - ICEBERG_ASSIGN_OR_RAISE(auto child, ToParquetNode(child_field)); + for (size_t i = 0; i < struct_type.fields().size(); ++i) { + const auto& child_field = struct_type.fields()[i]; + ICEBERG_ASSIGN_OR_RAISE( + auto child, + ToParquetNode(child_field, + storage_type ? storage_type->field(i)->type() : nullptr)); children.push_back(std::move(child)); } return ::parquet::schema::GroupNode::Make( @@ -223,7 +232,10 @@ Result<::parquet::schema::NodePtr> ToParquetNode(const SchemaField& field) { } case TypeId::kList: { const auto& list_type = internal::checked_cast(type); - ICEBERG_ASSIGN_OR_RAISE(auto element, ToParquetNode(list_type.element())); + ICEBERG_ASSIGN_OR_RAISE( + auto element, + ToParquetNode(list_type.element(), + storage_type ? storage_type->field(0)->type() : nullptr)); auto repeated = ::parquet::schema::GroupNode::Make( "list", ::parquet::Repetition::REPEATED, {std::move(element)}); return ::parquet::schema::GroupNode::Make( @@ -232,8 +244,15 @@ Result<::parquet::schema::NodePtr> ToParquetNode(const SchemaField& field) { } case TypeId::kMap: { const auto& map_type = internal::checked_cast(type); - ICEBERG_ASSIGN_OR_RAISE(auto key, ToParquetNode(map_type.key())); - ICEBERG_ASSIGN_OR_RAISE(auto value, ToParquetNode(map_type.value())); + const auto* arrow_map = + storage_type ? static_cast(storage_type.get()) + : nullptr; + ICEBERG_ASSIGN_OR_RAISE( + auto key, + ToParquetNode(map_type.key(), arrow_map ? arrow_map->key_type() : nullptr)); + ICEBERG_ASSIGN_OR_RAISE( + auto value, + ToParquetNode(map_type.value(), arrow_map ? arrow_map->item_type() : nullptr)); auto repeated = ::parquet::schema::GroupNode::Make("key_value", ::parquet::Repetition::REPEATED, {std::move(key), std::move(value)}); @@ -241,8 +260,29 @@ Result<::parquet::schema::NodePtr> ToParquetNode(const SchemaField& field) { std::string(field.name()), repetition, {std::move(repeated)}, ::parquet::LogicalType::Map(), field.field_id()); } - case TypeId::kVariant: - return NotSupported("Cannot write Iceberg variant type to Parquet"); + case TypeId::kVariant: { + if (storage_type) { + auto arrow_field = + ::arrow::field(std::string(field.name()), storage_type, field.optional()) + ->WithMetadata( + ::arrow::key_value_metadata({std::string(kParquetFieldIdKey)}, + {std::to_string(field.field_id())})); + auto arrow_schema = ::arrow::schema({std::move(arrow_field)}); + std::shared_ptr<::parquet::SchemaDescriptor> descriptor; + ICEBERG_ARROW_RETURN_NOT_OK(::parquet::arrow::ToParquetSchema( + arrow_schema.get(), *::parquet::default_writer_properties(), &descriptor)); + return static_cast( + descriptor->schema_root().get()) + ->field(0); + } + auto metadata = ::parquet::schema::PrimitiveNode::Make( + "metadata", ::parquet::Repetition::REQUIRED, ::parquet::Type::BYTE_ARRAY); + auto value = ::parquet::schema::PrimitiveNode::Make( + "value", ::parquet::Repetition::REQUIRED, ::parquet::Type::BYTE_ARRAY); + return ::parquet::schema::GroupNode::Make( + std::string(field.name()), repetition, {std::move(metadata), std::move(value)}, + ::parquet::LogicalType::Variant(), field.field_id()); + } default: return InvalidSchema("Expected nested Iceberg type, got {}", type); } @@ -256,24 +296,38 @@ bool HasSelectedColumn(const FieldProjection& projection) { if (projection.attributes) { const auto& attributes = internal::checked_cast(*projection.attributes); - if (attributes.column_id) { + if (attributes.column_id || !attributes.column_ids.empty()) { return true; } } return std::ranges::any_of(projection.children, HasSelectedColumn); } -std::optional FirstColumnIndex( - const ::parquet::arrow::SchemaField& parquet_field) { - if (parquet_field.column_index >= 0) { - return parquet_field.column_index; +void CollectVariantColumns(const ::parquet::arrow::SchemaField& field, + std::vector& column_ids) { + if (field.column_index >= 0) column_ids.push_back(field.column_index); + for (const auto& child : field.children) CollectVariantColumns(child, column_ids); +} + +std::vector FirstAnchorColumns(const ::parquet::arrow::SchemaField& field) { + if (field.field->type()->id() == ::arrow::Type::EXTENSION && + static_cast(*field.field->type()).extension_name() == + kArrowVariantExtensionName) { + std::vector columns; + CollectVariantColumns(field, columns); + return columns; } - for (const auto& child : parquet_field.children) { - if (auto column_index = FirstColumnIndex(child)) { - return column_index; - } + if (field.column_index >= 0) return {field.column_index}; + for (const auto& child : field.children) { + auto columns = FirstAnchorColumns(child); + if (!columns.empty()) return columns; } - return std::nullopt; + return {}; +} + +std::shared_ptr AnchorAttributes(std::vector columns) { + if (columns.size() == 1) return std::make_shared(columns[0]); + return std::make_shared(std::move(columns)); } // Pick a physical column as an anchor when all children are null-projected, @@ -285,9 +339,9 @@ void SelectAnchorColumnIfEmpty( return; } for (const auto& parquet_field : parquet_fields) { - if (auto column_index = FirstColumnIndex(parquet_field)) { - projection->attributes = - std::make_shared(column_index.value()); + auto columns = FirstAnchorColumns(parquet_field); + if (!columns.empty()) { + projection->attributes = AnchorAttributes(std::move(columns)); return; } } @@ -476,8 +530,12 @@ Status ValidateParquetTypeCompatibility( case TypeId::kUnknown: return {}; case TypeId::kVariant: - return NotSupported("Reading Iceberg type {} from Parquet is not supported", - expected_type); + if (arrow_type->id() == ::arrow::Type::EXTENSION && + static_cast(*arrow_type).extension_name() == + kArrowVariantExtensionName) { + return {}; + } + break; case TypeId::kStruct: if (arrow_type->id() == ::arrow::Type::STRUCT) { return {}; @@ -547,7 +605,33 @@ class ProjectionBuilder { ValidateParquetTypeCompatibility(expected_type, parquet_field)); } - if (expected_type.is_nested()) { + if (expected_type.is_variant()) { + std::vector column_ids; + CollectVariantColumns(parquet_field, column_ids); + if (column_ids.empty()) { + return InvalidSchema("Variant field has no physical columns"); + } + if (descr_ == nullptr) + return InvalidSchema("Missing variant Parquet schema metadata"); + for (auto column_id : column_ids) { + if (column_id < 0 || column_id >= descr_->num_columns()) { + return InvalidSchema("Variant physical column index is out of range"); + } + auto node = descr_->Column(column_id)->schema_node()->parent(); + while (node != nullptr && !node->logical_type()->is_variant()) + node = node->parent(); + // ProjectStruct has already matched the manifest field ID, including IDs + // supplied by name mapping. The physical node may legitimately have no ID. + if (node == nullptr || + (node->field_id() >= 0 && node->field_id() != expected_field.field_id())) { + return InvalidSchema( + "Iceberg variant requires a Parquet VARIANT group with the matching field " + "ID"); + } + } + projection.attributes = + std::make_shared(std::move(column_ids)); + } else if (expected_type.is_nested()) { ICEBERG_ASSIGN_OR_RAISE(projection, ProjectNested(expected_type, parquet_field.children)); } else { @@ -682,9 +766,9 @@ class ProjectionBuilder { ICEBERG_ASSIGN_OR_RAISE(auto sub_projection, ProjectField(sub_field, sub_node, i)); if (sub_projection.kind == FieldProjection::Kind::kNull && !HasSelectedColumn(sub_projection)) { - if (auto column_index = FirstColumnIndex(sub_node)) { - sub_projection.attributes = - std::make_shared(column_index.value()); + auto columns = FirstAnchorColumns(sub_node); + if (!columns.empty()) { + sub_projection.attributes = AnchorAttributes(std::move(columns)); } } result.children.emplace_back(std::move(sub_projection)); @@ -733,6 +817,8 @@ void CollectColumnIds(const FieldProjection& field_projection, if (attributes.column_id) { column_ids->push_back(attributes.column_id.value()); } + column_ids->insert(column_ids->end(), attributes.column_ids.begin(), + attributes.column_ids.end()); } for (const auto& child : field_projection.children) { CollectColumnIds(child, column_ids); @@ -831,11 +917,16 @@ void ApplyNameMapping(const NameMapping& mapping, } Result> ToParquetSchema( - const Schema& iceberg_schema) { + const Schema& iceberg_schema, + const std::shared_ptr<::arrow::Schema>& storage_schema) { ::parquet::schema::NodeVector fields; fields.reserve(iceberg_schema.fields().size()); - for (const auto& field : iceberg_schema.fields()) { - ICEBERG_ASSIGN_OR_RAISE(auto parquet_field, ToParquetNode(field)); + for (size_t i = 0; i < iceberg_schema.fields().size(); ++i) { + const auto& field = iceberg_schema.fields()[i]; + ICEBERG_ASSIGN_OR_RAISE( + auto parquet_field, + ToParquetNode(field, + storage_schema ? storage_schema->field(i)->type() : nullptr)); fields.push_back(std::move(parquet_field)); } diff --git a/src/iceberg/parquet/parquet_schema_util_internal.h b/src/iceberg/parquet/parquet_schema_util_internal.h index 6612de056..f80df30be 100644 --- a/src/iceberg/parquet/parquet_schema_util_internal.h +++ b/src/iceberg/parquet/parquet_schema_util_internal.h @@ -20,6 +20,8 @@ #pragma once #include +#include +#include #include #include @@ -33,10 +35,14 @@ namespace iceberg::parquet { /// \brief Parquet specific attributes for the field. struct ParquetExtraAttributes : public FieldProjection::ExtraAttributes { explicit ParquetExtraAttributes(int32_t column_id) : column_id(column_id) {} + explicit ParquetExtraAttributes(std::vector column_ids) + : column_ids(std::move(column_ids)) {} ~ParquetExtraAttributes() override = default; /// \brief The column id of projected Parquet column. std::optional column_id; + /// \brief All physical leaves of a projected variant group. + std::vector column_ids; }; /// \brief Project an Iceberg Schema onto a Parquet Schema. @@ -76,6 +82,7 @@ void ApplyNameMapping(const NameMapping& mapping, /// \brief Convert an Iceberg schema to a Parquet schema. Result> ToParquetSchema( - const Schema& iceberg_schema); + const Schema& iceberg_schema, + const std::shared_ptr<::arrow::Schema>& storage_schema = nullptr); } // namespace iceberg::parquet diff --git a/src/iceberg/parquet/parquet_variant_util.cc b/src/iceberg/parquet/parquet_variant_util.cc new file mode 100644 index 000000000..afe6ddcfb --- /dev/null +++ b/src/iceberg/parquet/parquet_variant_util.cc @@ -0,0 +1,690 @@ +/* + * 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 +#include +#include +#include + +#include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/parquet/parquet_variant_util_internal.h" +#include "iceberg/schema.h" +#include "iceberg/schema_internal.h" +#include "iceberg/type.h" +#include "iceberg/util/formatter.h" // IWYU pragma: keep +#include "iceberg/util/macros.h" + +#if ICEBERG_HAS_ARROW_VARIANT + +# include +# include +# include +# include +# include +#endif + +namespace iceberg::parquet { + +namespace { + +[[maybe_unused]] unexpected MissingVariantSupport() { + return NotSupported( + "Parquet variant I/O requires Arrow's variant validation and " + "shredding utilities (apache/arrow#50252)"); +} + +bool HasVariant(const Type& type) { + if (type.is_variant()) return true; + if (type.is_nested()) { + for (const auto& field : static_cast(type).fields()) { + if (HasVariant(*field.type())) return true; + } + } + return false; +} + +#if ICEBERG_HAS_ARROW_VARIANT + +using ::parquet::variant::VariantMetadataView; +using ::parquet::variant::VariantValueView; + +Result> WithValidity( + const std::shared_ptr<::arrow::Array>& array, const std::vector& valid, + ::arrow::MemoryPool* pool) { + auto data = array->data()->Copy(); + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto bitmap, ::arrow::AllocateEmptyBitmap(array->offset() + array->length(), pool)); + int64_t null_count = 0; + for (int64_t i = 0; i < array->length(); ++i) { + ::arrow::bit_util::SetBitTo(bitmap->mutable_data(), array->offset() + i, valid[i]); + null_count += !valid[i]; + } + data->buffers[0] = std::move(bitmap); + data->null_count = null_count; + return ::arrow::MakeArray(std::move(data)); +} + +bool HasArrowVariant(const ::arrow::DataType& type) { + if (type.id() == ::arrow::Type::EXTENSION) { + return static_cast(type).extension_name() == + kArrowVariantExtensionName; + } + return std::ranges::any_of( + type.fields(), [](const auto& field) { return HasArrowVariant(*field->type()); }); +} + +Result> MaskVariantParents( + const std::shared_ptr<::arrow::Array>& array, const std::vector* visible, + ::arrow::MemoryPool* pool) { + if (!HasArrowVariant(*array->type())) return array; + if (array->type_id() == ::arrow::Type::EXTENSION) { + if (visible == nullptr || + std::ranges::all_of(*visible, [](uint8_t v) { return v; })) { + return array; + } + auto valid = *visible; + for (int64_t i = 0; i < array->length(); ++i) valid[i] &= array->IsValid(i); + return WithValidity(array, valid, pool); + } + auto data = array->data()->Copy(); + if (array->type_id() == ::arrow::Type::STRUCT) { + for (auto& child_data : data->child_data) { + auto child = ::arrow::MakeArray(child_data); + if (!HasArrowVariant(*child->type())) continue; + std::vector child_visible(child->length(), 0); + for (int64_t i = 0; i < array->length(); ++i) { + child_visible[array->offset() + i] = + array->IsValid(i) && (visible == nullptr || (*visible)[i]); + } + ICEBERG_ASSIGN_OR_RAISE(auto masked, + MaskVariantParents(child, &child_visible, pool)); + child_data = masked->data(); + } + } else if (array->type_id() == ::arrow::Type::LIST || + array->type_id() == ::arrow::Type::LARGE_LIST || + array->type_id() == ::arrow::Type::MAP) { + auto child = ::arrow::MakeArray(data->child_data[0]); + std::vector child_visible(child->length(), 0); + auto mark_visible = [&](const auto& list) { + for (int64_t i = 0; i < array->length(); ++i) { + if (array->IsValid(i) && (visible == nullptr || (*visible)[i])) { + std::fill(child_visible.begin() + list.value_offset(i), + child_visible.begin() + list.value_offset(i + 1), 1); + } + } + }; + if (array->type_id() == ::arrow::Type::LARGE_LIST) { + mark_visible(static_cast(*array)); + } else { + mark_visible(static_cast(*array)); + } + ICEBERG_ASSIGN_OR_RAISE(auto masked, MaskVariantParents(child, &child_visible, pool)); + data->child_data[0] = masked->data(); + } + return ::arrow::MakeArray(std::move(data)); +} + +// Arrow's shredder builds binary-view columns, whose slicing is not yet +// supported by the Parquet writer. Use binary storage for the file schema. +std::shared_ptr<::arrow::DataType> MaterializeStorageType( + const std::shared_ptr<::arrow::DataType>& type) { + if (type->id() == ::arrow::Type::BINARY_VIEW) return ::arrow::binary(); + if (type->id() == ::arrow::Type::STRUCT) { + ::arrow::FieldVector fields; + for (const auto& field : type->fields()) { + fields.push_back(field->WithType(MaterializeStorageType(field->type()))); + } + return ::arrow::struct_(std::move(fields)); + } + if (type->id() == ::arrow::Type::LIST) { + return ::arrow::list( + type->field(0)->WithType(MaterializeStorageType(type->field(0)->type()))); + } + return type; +} + +Result> CastVariantStorage( + const ::arrow::extension::VariantArray& array, + const std::shared_ptr<::arrow::DataType>& output_type) { + const auto& extension = static_cast(*output_type); + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto storage, ::arrow::compute::Cast(array.storage(), extension.storage_type())); + return ::arrow::ExtensionType::WrapArray(output_type, storage.make_array()); +} + +// Shredding targets describe values, not Iceberg fields. Remove Iceberg IDs and +// other field metadata before handing their logical layout to Arrow. +std::shared_ptr<::arrow::DataType> StripTargetMetadata( + const std::shared_ptr<::arrow::DataType>& type) { + if (type->id() == ::arrow::Type::STRUCT) { + ::arrow::FieldVector fields; + for (const auto& field : type->fields()) { + fields.push_back(::arrow::field(field->name(), StripTargetMetadata(field->type()), + field->nullable())); + } + return ::arrow::struct_(std::move(fields)); + } + if (type->id() == ::arrow::Type::LIST) { + const auto& field = type->field(0); + return ::arrow::list(::arrow::field(field->name(), StripTargetMetadata(field->type()), + field->nullable())); + } + return type; +} + +Status CheckTarget(const Type& type) { + if (type.is_struct() || type.is_list()) { + const auto& nested = static_cast(type); + for (const auto& field : nested.fields()) { + if (!field.type()) { + return InvalidArgument("Variant shredding target contains a null type"); + } + ICEBERG_RETURN_UNEXPECTED(CheckTarget(*field.type())); + } + return {}; + } + if (!type.is_primitive() || type.type_id() == TypeId::kUnknown || + type.type_id() == TypeId::kGeometry || type.type_id() == TypeId::kGeography || + type.type_id() == TypeId::kFixed) { + return InvalidArgument("Unsupported variant shredding target: {}", type); + } + return {}; +} + +Result> TargetToArrow( + const std::shared_ptr& type) { + ICEBERG_RETURN_UNEXPECTED(CheckTarget(*type)); + Schema schema({SchemaField::MakeOptional(0, "target", type)}); + ArrowSchema c_schema; + ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(schema, &c_schema)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema, ::arrow::ImportSchema(&c_schema)); + return StripTargetMetadata(arrow_schema->field(0)->type()); +} + +void AddTargetNames(const Type& type, ::parquet::variant::VariantBuilder& builder) { + if (type.is_nested()) { + for (const auto& field : static_cast(type).fields()) { + if (type.is_struct()) builder.AddFieldName(field.name()); + AddTargetNames(*field.type(), builder); + } + } +} + +std::shared_ptr<::arrow::extension::VariantArray> PrepareVariant( + const ::arrow::extension::VariantArray& array, const Type* target, + const std::vector* visible, ::arrow::MemoryPool* pool) { + const auto& storage = static_cast(*array.storage()); + const auto& metadata = static_cast(*storage.field(0)); + const auto& values = static_cast(*storage.field(1)); + ::parquet::variant::VariantArrayBuilder output(pool); + for (int64_t i = 0; i < array.length(); ++i) { + if (array.IsNull(i) || (visible != nullptr && !(*visible)[i])) { + // Required storage children use empty bytes for a null outer variant. + output.AppendNull(); + continue; + } + if (metadata.IsNull(i) || values.IsNull(i)) { + throw ::parquet::ParquetException("Non-null variant has null metadata or value"); + } + if (target == nullptr) { + output.AppendEncoded({::arrow::Buffer::FromString(std::string(metadata.GetView(i))), + ::arrow::Buffer::FromString(std::string(values.GetView(i)))}); + continue; + } + auto dictionary = ::parquet::variant::VariantMetadataView::Make(metadata.GetView(i)); + ::parquet::variant::VariantBuilder row(pool); + // Preserve existing dictionary IDs so encoded object references stay valid. + for (uint32_t id = 0; id < dictionary.dictionary_size(); ++id) { + row.AddFieldName(dictionary.string(id)); + } + AddTargetNames(*target, row); + row.AppendEncodedValue(values.GetView(i)); + output.AppendEncoded(row.Finish()); + } + return output.Finish(); +} + +bool MatchesShreddingType(const VariantValueView& value, const Type& target) { + using Primitive = ::parquet::variant::VariantPrimitiveType; + if (std::holds_alternative<::parquet::variant::VariantShortStringView>(value.data())) { + return target.type_id() == TypeId::kString; + } + const auto* primitive = + std::get_if<::parquet::variant::VariantPrimitiveView>(&value.data()); + if (primitive == nullptr) return false; + const auto source = primitive->type(); + switch (target.type_id()) { + case TypeId::kBoolean: + return source == Primitive::kBooleanTrue || source == Primitive::kBooleanFalse; + case TypeId::kInt: + case TypeId::kLong: + return source >= Primitive::kInt8 && source <= Primitive::kInt64; + case TypeId::kFloat: + return source == Primitive::kFloat; + case TypeId::kDouble: + return source == Primitive::kDouble; + case TypeId::kDecimal: + return (source == Primitive::kDecimal4 || source == Primitive::kDecimal8 || + source == Primitive::kDecimal16) && + static_cast(primitive->payload()[0]) == + static_cast(target).scale(); + case TypeId::kString: + return source == Primitive::kString; + case TypeId::kBinary: + return source == Primitive::kBinary; + case TypeId::kDate: + return source == Primitive::kDate; + case TypeId::kTime: + return source == Primitive::kTimeNTZMicros; + case TypeId::kTimestamp: + return source == Primitive::kTimestampNTZMicros; + case TypeId::kTimestampTz: + return source == Primitive::kTimestampMicros; + case TypeId::kTimestampNs: + return source == Primitive::kTimestampNTZNanos; + case TypeId::kTimestampTzNs: + return source == Primitive::kTimestampNanos; + case TypeId::kUuid: + return source == Primitive::kUuid; + default: + return false; + } +} + +using VariantSlots = std::vector>; + +// The pinned Arrow shredder casts between unrelated Variant types. Restore such +// values to their residual slots, including slots nested inside objects and lists. +Result> PreserveShreddedTypes( + const Type& target, const std::shared_ptr<::arrow::StructArray>& group, + const VariantSlots& sources, ::arrow::MemoryPool* pool) { + auto typed = group->GetFieldByName("typed_value"); + auto data = group->data()->Copy(); + const auto& group_type = static_cast(*group->type()); + if (target.is_struct()) { + const auto& fields = static_cast(target).fields(); + const auto& structure = static_cast(*typed); + auto typed_data = typed->data()->Copy(); + for (size_t i = 0; i < fields.size(); ++i) { + VariantSlots children(sources.size()); + for (size_t row = 0; row < sources.size(); ++row) { + if (!sources[row] || typed->IsNull(row)) continue; + const auto* object = + std::get_if<::parquet::variant::VariantObjectView>(&sources[row]->data()); + if (object != nullptr) children[row] = object->GetField(fields[i].name()); + } + ICEBERG_ASSIGN_OR_RAISE( + auto child, + PreserveShreddedTypes( + *fields[i].type(), + std::static_pointer_cast<::arrow::StructArray>(structure.field(i)), + children, pool)); + typed_data->child_data[i] = child->data(); + } + typed = ::arrow::MakeArray(std::move(typed_data)); + } else if (target.is_list()) { + const auto& list = static_cast(*typed); + VariantSlots children(list.values()->length()); + for (size_t row = 0; row < sources.size(); ++row) { + if (!sources[row] || typed->IsNull(row)) continue; + const auto* values = + std::get_if<::parquet::variant::VariantArrayView>(&sources[row]->data()); + if (values == nullptr) continue; + for (int64_t i = 0; i < list.value_length(row); ++i) { + children[list.value_offset(row) + i] = values->GetElement(i); + } + } + ICEBERG_ASSIGN_OR_RAISE( + auto child, PreserveShreddedTypes( + *static_cast(target).element().type(), + std::static_pointer_cast<::arrow::StructArray>(list.values()), + children, pool)); + auto typed_data = typed->data()->Copy(); + typed_data->child_data[0] = child->data(); + typed = ::arrow::MakeArray(std::move(typed_data)); + } else { + std::vector valid(sources.size()); + bool needs_restore = false; + for (size_t i = 0; i < sources.size(); ++i) { + valid[i] = typed->IsValid(i); + if (valid[i] && sources[i] && !MatchesShreddingType(*sources[i], target)) { + valid[i] = 0; + needs_restore = true; + } + } + if (!needs_restore) return group; + const auto& residual = + static_cast(*group->GetFieldByName("value")); + ::arrow::BinaryBuilder builder(pool); + for (size_t i = 0; i < sources.size(); ++i) { + if (typed->IsValid(i) && !valid[i]) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(sources[i]->value())); + } else if (residual.IsNull(i)) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.AppendNull()); + } else { + ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(residual.GetView(i))); + } + } + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto values, builder.Finish()); + data->child_data[group_type.GetFieldIndex("value")] = values->data(); + ICEBERG_ASSIGN_OR_RAISE(typed, WithValidity(typed, valid, pool)); + } + data->child_data[group_type.GetFieldIndex("typed_value")] = typed->data(); + return ::arrow::MakeArray(std::move(data)); +} + +Result> ShredPreservingTypes( + const ::arrow::extension::VariantArray& source, + const ::arrow::extension::VariantArray& prepared, const Type& target, + const std::shared_ptr<::arrow::DataType>& arrow_target, + const std::shared_ptr<::arrow::DataType>& output_type, ::arrow::MemoryPool* pool) { + auto shredded = ::parquet::variant::ShredVariantArray(prepared, arrow_target, pool); + ICEBERG_ASSIGN_OR_RAISE(auto materialized, CastVariantStorage(*shredded, output_type)); + const auto& metadata = static_cast(*source.metadata()); + const auto& values = static_cast(*source.value()); + // Object/list views borrow these metadata objects; reserve before creating views. + std::vector dictionaries; + dictionaries.reserve(source.length()); + VariantSlots sources(source.length()); + for (int64_t i = 0; i < source.length(); ++i) { + if (prepared.IsNull(i)) continue; + dictionaries.push_back(VariantMetadataView::Make(metadata.GetView(i))); + sources[i] = VariantValueView::Make(values.GetView(i), dictionaries.back()); + } + ICEBERG_ASSIGN_OR_RAISE( + auto storage, + PreserveShreddedTypes( + target, + std::static_pointer_cast<::arrow::StructArray>( + static_cast(*materialized).storage()), + sources, pool)); + return ::arrow::ExtensionType::WrapArray(output_type, storage); +} + +#endif + +Result> MakeWriteType( + const Type& type, int32_t field_id, const std::shared_ptr<::arrow::DataType>& input, + const std::unordered_map>& targets, + ::arrow::MemoryPool* pool) { + if (type.is_variant()) { + auto target = targets.find(field_id); + if (target == targets.end()) { + return input; + } +#if ICEBERG_HAS_ARROW_VARIANT + ICEBERG_ASSIGN_OR_RAISE(auto target_type, TargetToArrow(target->second)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto empty, ::arrow::MakeArrayOfNull(input, 0, pool)); + auto prepared = + PrepareVariant(static_cast(*empty), + target->second.get(), nullptr, pool); + auto shredded = ::parquet::variant::ShredVariantArray(*prepared, target_type, pool); + return ::arrow::extension::variant(MaterializeStorageType( + static_cast(*shredded->type()).storage_type())); +#else + return MissingVariantSupport(); +#endif + } + if (type.is_struct()) { + ::arrow::FieldVector fields; + const auto& iceberg_fields = static_cast(type).fields(); + for (size_t i = 0; i < iceberg_fields.size(); ++i) { + const auto& field = iceberg_fields[i]; + ICEBERG_ASSIGN_OR_RAISE(auto child, + MakeWriteType(*field.type(), field.field_id(), + input->field(i)->type(), targets, pool)); + fields.push_back(input->field(i)->WithType(std::move(child))); + } + return ::arrow::struct_(std::move(fields)); + } + if (type.is_list()) { + const auto& element = static_cast(type).element(); + ICEBERG_ASSIGN_OR_RAISE(auto child, + MakeWriteType(*element.type(), element.field_id(), + input->field(0)->type(), targets, pool)); + return ::arrow::list(input->field(0)->WithType(std::move(child))); + } + if (type.is_map()) { + const auto& map = static_cast(type); + const auto& arrow_map = static_cast(*input); + ICEBERG_ASSIGN_OR_RAISE(auto key, + MakeWriteType(*map.key().type(), map.key().field_id(), + arrow_map.key_type(), targets, pool)); + ICEBERG_ASSIGN_OR_RAISE(auto value, + MakeWriteType(*map.value().type(), map.value().field_id(), + arrow_map.item_type(), targets, pool)); + return std::make_shared<::arrow::MapType>( + arrow_map.key_field()->WithType(std::move(key)), + arrow_map.item_field()->WithType(std::move(value)), arrow_map.keys_sorted()); + } + return input; +} + +Result> ShredArray( + const Type& type, int32_t field_id, bool optional, + const std::shared_ptr<::arrow::Array>& array, + const std::shared_ptr<::arrow::DataType>& output_type, + const std::unordered_map>& targets, + const std::vector* visible, ::arrow::MemoryPool* pool) { + if (!HasVariant(type)) { + return array; + } + auto data = array->data()->Copy(); + if (type.is_variant()) { + if (!optional) { + for (int64_t i = 0; i < array->length(); ++i) { + if (array->IsNull(i) && (visible == nullptr || (*visible)[i])) { + return InvalidArgument("Required variant field {} contains a null", field_id); + } + } + } +#if ICEBERG_HAS_ARROW_VARIANT + auto target = targets.find(field_id); + const auto& storage = static_cast( + *static_cast(*array).storage()); + if (target == targets.end() && storage.field(0)->null_count() == 0 && + storage.field(1)->null_count() == 0 && + (visible == nullptr || + std::ranges::all_of(*visible, [](uint8_t valid) { return valid; }))) { + return array; + } + auto prepared = PrepareVariant( + static_cast(*array), + target == targets.end() ? nullptr : target->second.get(), visible, pool); + if (target == targets.end()) { + return ProjectVariantArray(prepared, output_type, pool); + } + ICEBERG_ASSIGN_OR_RAISE(auto arrow_target, TargetToArrow(target->second)); + return ShredPreservingTypes( + static_cast(*array), *prepared, + *target->second, arrow_target, output_type, pool); +#else + return MissingVariantSupport(); +#endif + } + std::vector child_visible; + if (type.is_struct()) { + const auto& fields = static_cast(type).fields(); + for (size_t child_index = 0; child_index < fields.size(); ++child_index) { + auto child = ::arrow::MakeArray(data->child_data[child_index]); + child_visible.assign(child->length(), 0); + for (int64_t i = 0; i < array->length(); ++i) { + child_visible[array->offset() + i] = + array->IsValid(i) && (visible == nullptr || (*visible)[i]); + } + ICEBERG_ASSIGN_OR_RAISE( + auto transformed, + ShredArray(*fields[child_index].type(), fields[child_index].field_id(), + fields[child_index].optional(), child, + output_type->field(child_index)->type(), targets, &child_visible, + pool)); + data->child_data[child_index] = transformed->data(); + } + } else if (type.is_list() || type.is_map()) { + const auto& list = static_cast(*array); + auto child = ::arrow::MakeArray(data->child_data[0]); + child_visible.assign(child->length(), 0); + for (int64_t i = 0; i < array->length(); ++i) { + if (array->IsValid(i) && (visible == nullptr || (*visible)[i])) { + std::fill(child_visible.begin() + list.value_offset(i), + child_visible.begin() + list.value_offset(i + 1), 1); + } + } + if (type.is_list()) { + ICEBERG_ASSIGN_OR_RAISE( + auto transformed, + ShredArray(*static_cast(type).element().type(), + static_cast(type).element().field_id(), + static_cast(type).element().optional(), child, + output_type->field(0)->type(), targets, &child_visible, pool)); + data->child_data[0] = transformed->data(); + } else { + const auto& map = static_cast(type); + StructType entries({map.key(), map.value()}); + ICEBERG_ASSIGN_OR_RAISE( + auto transformed, + ShredArray(entries, -1, true, child, output_type->field(0)->type(), targets, + &child_visible, pool)); + data->child_data[0] = transformed->data(); + } + } + data->type = output_type; + return ::arrow::MakeArray(std::move(data)); +} + +} // namespace + +Status CheckVariantSupport(const Type& type) { + if (type.is_variant()) { +#if ICEBERG_HAS_ARROW_VARIANT + return {}; +#else + return MissingVariantSupport(); +#endif + } + if (type.is_nested()) { + for (const auto& field : static_cast(type).fields()) { + ICEBERG_RETURN_UNEXPECTED(CheckVariantSupport(*field.type())); + } + } + return {}; +} + +Result> NormalizeVariantReadBatch( + const std::shared_ptr<::arrow::RecordBatch>& batch, ::arrow::MemoryPool* pool) { +#if ICEBERG_HAS_ARROW_VARIANT + if (!std::ranges::any_of(batch->schema()->fields(), [](const auto& field) { + return HasArrowVariant(*field->type()); + })) { + return batch; + } + try { + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto array, batch->ToStructArray()); + ICEBERG_ASSIGN_OR_RAISE(auto masked, MaskVariantParents(array, nullptr, pool)); + ::parquet::variant::ValidateVariants(::arrow::ChunkedArray(masked), pool); + return ::arrow::RecordBatch::Make( + batch->schema(), batch->num_rows(), + static_cast(*masked).fields()); + } catch (const ::parquet::ParquetException& error) { + return Invalid("Invalid Parquet variant: {}", error.what()); + } +#else + return batch; +#endif +} + +Result> ProjectVariantArray( + const std::shared_ptr<::arrow::Array>& array, + const std::shared_ptr<::arrow::DataType>& output_type, ::arrow::MemoryPool* pool) { +#if ICEBERG_HAS_ARROW_VARIANT + try { + const auto& variant = static_cast(*array); + auto unshredded = ::parquet::variant::UnshredVariantArray(variant, pool); + // Arrow may return binary-view storage. The C Data interface uses canonical + // binary storage so that the schema stays the same across files and batches. + return CastVariantStorage(*unshredded, output_type); + } catch (const ::parquet::ParquetException& error) { + return Invalid("Invalid Parquet variant: {}", error.what()); + } +#else + return MissingVariantSupport(); +#endif +} + +Result> MakeVariantWriteSchema( + const Schema& schema, const std::shared_ptr<::arrow::Schema>& input_schema, + const std::unordered_map>& targets, + ::arrow::MemoryPool* pool) { + ICEBERG_RETURN_UNEXPECTED(CheckVariantSupport(schema)); + for (const auto& [field_id, target] : targets) { + ICEBERG_ASSIGN_OR_RAISE(auto field, schema.FindFieldById(field_id)); + if (!field || !field->get().type()->is_variant()) { + return InvalidArgument("Variant shredding field id {} does not identify a variant", + field_id); + } + if (!target) { + return InvalidArgument("Variant shredding target for field id {} is null", + field_id); + } + } + try { + ICEBERG_ASSIGN_OR_RAISE( + auto type, MakeWriteType(schema, -1, ::arrow::struct_(input_schema->fields()), + targets, pool)); + return ::arrow::schema(type->fields(), input_schema->metadata()); + } catch (const ::parquet::ParquetException& error) { + return InvalidArgument("Invalid variant shredding target: {}", error.what()); + } +} + +Result> ShredVariantBatch( + const Schema& schema, const std::shared_ptr<::arrow::RecordBatch>& batch, + const std::shared_ptr<::arrow::Schema>& output_schema, + const std::unordered_map>& targets, + ::arrow::MemoryPool* pool) { + try { + if (!HasVariant(schema)) return batch; + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto array, batch->ToStructArray()); +#if ICEBERG_HAS_ARROW_VARIANT + if (!targets.empty()) { + ::parquet::variant::ValidateVariants(::arrow::ChunkedArray(array), pool); + } +#endif + ICEBERG_ASSIGN_OR_RAISE( + auto output, + ShredArray(schema, -1, false, array, ::arrow::struct_(output_schema->fields()), + targets, nullptr, pool)); + const auto& structure = static_cast(*output); + return ::arrow::RecordBatch::Make(output_schema, batch->num_rows(), + structure.fields()); + } catch (const ::parquet::ParquetException& error) { + return Invalid("Cannot shred variant batch: {}", error.what()); + } +} + +} // namespace iceberg::parquet diff --git a/src/iceberg/parquet/parquet_variant_util_internal.h b/src/iceberg/parquet/parquet_variant_util_internal.h new file mode 100644 index 000000000..fba5a2d7d --- /dev/null +++ b/src/iceberg/parquet/parquet_variant_util_internal.h @@ -0,0 +1,65 @@ +/* + * 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 "iceberg/result.h" +#include "iceberg/type_fwd.h" + +// Arrow's variant utilities are not available in Arrow 25. Keep ordinary Parquet +// I/O available with that release while enabling variant I/O with the new APIs. +#if __has_include() && \ + __has_include() && \ + __has_include() && \ + __has_include() && \ + __has_include() +# define ICEBERG_HAS_ARROW_VARIANT 1 +#else +# define ICEBERG_HAS_ARROW_VARIANT 0 +#endif + +namespace iceberg::parquet { + +Status CheckVariantSupport(const Type& type); + +// Propagate ancestor nulls to variant arrays before validating complete batches. +Result> NormalizeVariantReadBatch( + const std::shared_ptr<::arrow::RecordBatch>& batch, ::arrow::MemoryPool* pool); + +Result> ProjectVariantArray( + const std::shared_ptr<::arrow::Array>& array, + const std::shared_ptr<::arrow::DataType>& output_type, ::arrow::MemoryPool* pool); + +Result> MakeVariantWriteSchema( + const Schema& schema, const std::shared_ptr<::arrow::Schema>& input_schema, + const std::unordered_map>& targets, + ::arrow::MemoryPool* pool); + +Result> ShredVariantBatch( + const Schema& schema, const std::shared_ptr<::arrow::RecordBatch>& batch, + const std::shared_ptr<::arrow::Schema>& output_schema, + const std::unordered_map>& targets, + ::arrow::MemoryPool* pool); + +} // namespace iceberg::parquet diff --git a/src/iceberg/parquet/parquet_writer.cc b/src/iceberg/parquet/parquet_writer.cc index 08c394138..4ba735cec 100644 --- a/src/iceberg/parquet/parquet_writer.cc +++ b/src/iceberg/parquet/parquet_writer.cc @@ -34,13 +34,16 @@ #include #include #include +#include #include #include #include "iceberg/arrow/arrow_io_internal.h" +#include "iceberg/arrow/arrow_register.h" #include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/parquet/parquet_metrics_internal.h" #include "iceberg/parquet/parquet_schema_util_internal.h" +#include "iceberg/parquet/parquet_variant_util_internal.h" #include "iceberg/schema_internal.h" #include "iceberg/type.h" #include "iceberg/util/macros.h" @@ -238,7 +241,10 @@ Result> ParseCodecLevel(const WriterProperties& propertie class ParquetWriter::Impl { public: Status Open(const WriterOptions& options) { + arrow::RegisterAll(); + ICEBERG_PRECHECK(options.schema != nullptr, "Writer schema must not be null"); schema_ = options.schema; + variant_shredding_types_ = options.variant_shredding_types; ICEBERG_ASSIGN_OR_RAISE(auto compression, ParseCompression(options.properties)); ICEBERG_ASSIGN_OR_RAISE(auto compression_level, ParseCodecLevel(options.properties)); @@ -258,9 +264,12 @@ class ParquetWriter::Impl { ArrowSchema c_schema; ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*schema_, &c_schema)); - ICEBERG_ARROW_ASSIGN_OR_RETURN(arrow_schema_, ::arrow::ImportSchema(&c_schema)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(input_arrow_schema_, ::arrow::ImportSchema(&c_schema)); + ICEBERG_ASSIGN_OR_RAISE(arrow_schema_, + MakeVariantWriteSchema(*schema_, input_arrow_schema_, + variant_shredding_types_, pool_)); - ICEBERG_ASSIGN_OR_RAISE(parquet_schema_, ToParquetSchema(*schema_)); + ICEBERG_ASSIGN_OR_RAISE(parquet_schema_, ToParquetSchema(*schema_, arrow_schema_)); auto schema_node = std::static_pointer_cast<::parquet::schema::GroupNode>( parquet_schema_->schema_root()); @@ -281,8 +290,9 @@ class ParquetWriter::Impl { } Status Write(ArrowArray* array) { - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, - ::arrow::ImportRecordBatch(array, arrow_schema_)); + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto batch, ::arrow::ImportRecordBatch(array, input_arrow_schema_)); + ICEBERG_ARROW_RETURN_NOT_OK(batch->ValidateFull()); ICEBERG_ARROW_ASSIGN_OR_RETURN(auto struct_array, batch->ToStructArray()); FieldMetricsCollector field_metrics_collector(field_metrics_, *metrics_config_, @@ -290,7 +300,13 @@ class ParquetWriter::Impl { ICEBERG_RETURN_UNEXPECTED( field_metrics_collector.VisitStruct(*schema_, *struct_array)); - ICEBERG_ARROW_RETURN_NOT_OK(writer_->WriteRecordBatch(*batch)); + ICEBERG_ASSIGN_OR_RAISE(batch, ShredVariantBatch(*schema_, batch, arrow_schema_, + variant_shredding_types_, pool_)); + try { + ICEBERG_ARROW_RETURN_NOT_OK(writer_->WriteRecordBatch(*batch)); + } catch (const ::parquet::ParquetException& error) { + return Invalid("Cannot write Parquet batch: {}", error.what()); + } return {}; } @@ -346,6 +362,8 @@ class ParquetWriter::Impl { std::shared_ptr<::arrow::Schema> arrow_schema_; // Parquet schema descriptor generated from the Arrow schema. std::shared_ptr<::parquet::SchemaDescriptor> parquet_schema_; + std::shared_ptr<::arrow::Schema> input_arrow_schema_; + std::unordered_map> variant_shredding_types_; // Metrics config for collecting metrics during write. std::shared_ptr metrics_config_; // The output stream to write Parquet file. @@ -362,11 +380,16 @@ class ParquetWriter::Impl { std::unordered_map field_metrics_; }; +ParquetWriter::ParquetWriter() = default; ParquetWriter::~ParquetWriter() = default; Status ParquetWriter::Open(const WriterOptions& options) { impl_ = std::make_unique(); - return impl_->Open(options); + try { + return impl_->Open(options); + } catch (const ::parquet::ParquetException& error) { + return InvalidSchema("Cannot open Parquet writer: {}", error.what()); + } } Status ParquetWriter::Write(ArrowArray* array) { return impl_->Write(array); } diff --git a/src/iceberg/parquet/parquet_writer.h b/src/iceberg/parquet/parquet_writer.h index 562ffb801..543fd5757 100644 --- a/src/iceberg/parquet/parquet_writer.h +++ b/src/iceberg/parquet/parquet_writer.h @@ -30,7 +30,7 @@ namespace iceberg::parquet { /// \brief A writer that writes ArrowArray to Parquet files. class ICEBERG_BUNDLE_EXPORT ParquetWriter : public Writer { public: - ParquetWriter() = default; + ParquetWriter(); ~ParquetWriter() override; diff --git a/src/iceberg/schema_internal.cc b/src/iceberg/schema_internal.cc index 0a6189884..14d3f7314 100644 --- a/src/iceberg/schema_internal.cc +++ b/src/iceberg/schema_internal.cc @@ -39,31 +39,6 @@ constexpr const char* kArrowExtensionName = "ARROW:extension:name"; constexpr const char* kArrowExtensionMetadata = "ARROW:extension:metadata"; constexpr int32_t kUnknownFieldId = -1; -Status CheckArrowCompatible(const Type& type) { - switch (type.type_id()) { - case TypeId::kVariant: - return NotSupported("Iceberg type {} is not supported by Arrow conversion", - type.ToString()); - case TypeId::kStruct: - for (const auto& field : static_cast(type).fields()) { - ICEBERG_RETURN_UNEXPECTED(CheckArrowCompatible(*field.type())); - } - break; - case TypeId::kList: - ICEBERG_RETURN_UNEXPECTED( - CheckArrowCompatible(*static_cast(type).element().type())); - break; - case TypeId::kMap: { - const auto& map_type = static_cast(type); - ICEBERG_RETURN_UNEXPECTED(CheckArrowCompatible(*map_type.key().type())); - ICEBERG_RETURN_UNEXPECTED(CheckArrowCompatible(*map_type.value().type())); - } break; - default: - break; - } - return {}; -} - // Convert an Iceberg type to Arrow schema. Return value is Nanoarrow error code. ArrowErrorCode ToArrowSchema(const Type& type, bool optional, std::string_view name, std::optional field_id, ArrowSchema* schema) { @@ -182,8 +157,15 @@ ArrowErrorCode ToArrowSchema(const Type& type, bool optional, std::string_view n NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_NA)); break; case TypeId::kVariant: - ArrowBufferReset(&metadata_buffer); - return EINVAL; + NANOARROW_RETURN_NOT_OK(ArrowSchemaSetTypeStruct(schema, 2)); + NANOARROW_RETURN_NOT_OK( + ToArrowSchema(*binary(), false, "metadata", std::nullopt, schema->children[0])); + NANOARROW_RETURN_NOT_OK( + ToArrowSchema(*binary(), false, "value", std::nullopt, schema->children[1])); + NANOARROW_RETURN_NOT_OK( + ArrowMetadataBuilderAppend(&metadata_buffer, ArrowCharView(kArrowExtensionName), + ArrowCharView(kArrowVariantExtensionName))); + break; } if (!name.empty()) { @@ -209,8 +191,6 @@ Status ToArrowSchema(const Schema& schema, ArrowSchema* out) { return InvalidArgument("Output Arrow schema cannot be null"); } - ICEBERG_RETURN_UNEXPECTED(CheckArrowCompatible(schema)); - ArrowSchemaInit(out); internal::ArrowSchemaGuard schema_guard(out); @@ -234,9 +214,10 @@ Result GetFieldId(const ArrowSchema& schema) { ArrowStringView field_id_key{.data = kParquetFieldIdKey.data(), .size_bytes = kParquetFieldIdKey.size()}; - ArrowStringView field_id_value; + ArrowStringView field_id_value{}; if (ArrowMetadataGetValue(schema.metadata, field_id_key, &field_id_value) != - NANOARROW_OK) { + NANOARROW_OK || + field_id_value.data == nullptr) { return kUnknownFieldId; } @@ -249,6 +230,124 @@ Result GetFieldId(const ArrowSchema& schema) { return field_id_result.value(); } +Status ValidateVariantTypedValue(const ArrowSchema& schema); + +bool IsVariantBinary(const ArrowSchemaView& view) { + return view.type == NANOARROW_TYPE_BINARY || view.type == NANOARROW_TYPE_LARGE_BINARY || + view.type == NANOARROW_TYPE_BINARY_VIEW; +} + +Status ValidateVariantFieldGroup(const ArrowSchema& schema) { + ArrowSchemaView view; + ArrowError error{}; + if (ArrowSchemaViewInit(&view, &schema, &error) != NANOARROW_OK || + view.storage_type != NANOARROW_TYPE_STRUCT || schema.dictionary != nullptr || + view.extension_name.size_bytes != 0) { + return InvalidSchema("Variant field group must use struct storage"); + } + bool has_value = false; + bool has_typed_value = false; + for (int64_t i = 0; i < schema.n_children; ++i) { + const auto& child = *schema.children[i]; + std::string_view name = child.name == nullptr ? "" : child.name; + if (name == "typed_value" && !has_typed_value) { + ICEBERG_RETURN_UNEXPECTED(ValidateVariantTypedValue(child)); + has_typed_value = true; + } else if (name == "value" && !has_value) { + ArrowSchemaView child_view; + if (ArrowSchemaViewInit(&child_view, &child, &error) != NANOARROW_OK || + !IsVariantBinary(child_view) || child.dictionary != nullptr || + (child.flags & ARROW_FLAG_NULLABLE) == 0) { + return InvalidSchema("Variant field group value must be optional binary"); + } + has_value = true; + } else { + return InvalidSchema("Unexpected or duplicate variant field group child '{}'", + name); + } + } + if (!has_value && !has_typed_value) { + return InvalidSchema("Variant field group requires value or typed_value"); + } + return {}; +} + +Status ValidateVariantTypedValue(const ArrowSchema& schema) { + // Match Arrow's reader-compatible layouts without linking Arrow C++ into core. + ArrowSchemaView view; + ArrowError error{}; + if (ArrowSchemaViewInit(&view, &schema, &error) != NANOARROW_OK || + schema.dictionary != nullptr || (schema.flags & ARROW_FLAG_NULLABLE) == 0) { + return InvalidSchema("Variant typed_value must have a supported optional type"); + } + const std::string_view extension_name( + view.extension_name.data == nullptr ? "" : view.extension_name.data, + view.extension_name.size_bytes); + if (!extension_name.empty()) { + if (extension_name == kArrowUuidExtensionName && + view.storage_type == NANOARROW_TYPE_FIXED_SIZE_BINARY && view.fixed_size == 16 && + view.extension_metadata.size_bytes == 0) { + return {}; + } + return InvalidSchema("Unsupported variant typed_value extension '{}'", + extension_name); + } + switch (view.type) { + case NANOARROW_TYPE_STRUCT: + if (schema.n_children == 0) { + return InvalidSchema("Variant typed_value object must not be empty"); + } + for (int64_t i = 0; i < schema.n_children; ++i) { + ICEBERG_RETURN_UNEXPECTED(ValidateVariantFieldGroup(*schema.children[i])); + } + return {}; + case NANOARROW_TYPE_LIST: + case NANOARROW_TYPE_LARGE_LIST: + case NANOARROW_TYPE_LIST_VIEW: + case NANOARROW_TYPE_LARGE_LIST_VIEW: + case NANOARROW_TYPE_FIXED_SIZE_LIST: + if ((schema.children[0]->flags & ARROW_FLAG_NULLABLE) != 0) { + return InvalidSchema("Variant typed_value list element must be required"); + } + return ValidateVariantFieldGroup(*schema.children[0]); + case NANOARROW_TYPE_BOOL: + case NANOARROW_TYPE_INT8: + case NANOARROW_TYPE_INT16: + case NANOARROW_TYPE_INT32: + case NANOARROW_TYPE_INT64: + case NANOARROW_TYPE_FLOAT: + case NANOARROW_TYPE_DOUBLE: + case NANOARROW_TYPE_DATE32: + case NANOARROW_TYPE_BINARY: + case NANOARROW_TYPE_LARGE_BINARY: + case NANOARROW_TYPE_BINARY_VIEW: + case NANOARROW_TYPE_STRING: + case NANOARROW_TYPE_LARGE_STRING: + case NANOARROW_TYPE_STRING_VIEW: + return {}; + case NANOARROW_TYPE_DECIMAL32: + case NANOARROW_TYPE_DECIMAL64: + case NANOARROW_TYPE_DECIMAL128: + if (view.decimal_scale >= 0 && view.decimal_scale <= view.decimal_precision) { + return {}; + } + break; + case NANOARROW_TYPE_TIME64: + if (view.time_unit == NANOARROW_TIME_UNIT_MICRO) return {}; + break; + case NANOARROW_TYPE_TIMESTAMP: + if (view.time_unit == NANOARROW_TIME_UNIT_MICRO || + view.time_unit == NANOARROW_TIME_UNIT_NANO) { + return {}; + } + break; + default: + break; + } + return InvalidSchema("Unsupported variant typed_value type '{}'", + ArrowTypeString(view.type)); +} + Result> FromArrowSchema(const ArrowSchema& schema) { auto to_schema_field = [](const ArrowSchema& schema) -> Result> { @@ -273,6 +372,58 @@ Result> FromArrowSchema(const ArrowSchema& schema) { arrow_error.message); } + const std::string_view extension_name( + schema_view.extension_name.data == nullptr ? "" : schema_view.extension_name.data, + schema_view.extension_name.size_bytes); + if (extension_name == kArrowVariantExtensionName) { + if (schema_view.extension_metadata.size_bytes != 0) { + return InvalidSchema("Variant extension metadata must be empty"); + } + if (schema_view.storage_type != NANOARROW_TYPE_STRUCT || schema.n_children < 2 || + schema.n_children > 3) { + return InvalidSchema( + "Variant extension requires struct storage with metadata and value or " + "typed_value"); + } + bool has_metadata = false; + bool has_value = false; + bool has_typed_value = false; + bool value_nullable = true; + for (int64_t i = 0; i < schema.n_children; ++i) { + const auto& child = *schema.children[i]; + std::string_view name = child.name == nullptr ? "" : child.name; + if (name == "typed_value" && !has_typed_value) { + ICEBERG_RETURN_UNEXPECTED(ValidateVariantTypedValue(child)); + has_typed_value = true; + continue; + } + ArrowSchemaView child_view; + if (ArrowSchemaViewInit(&child_view, &child, &arrow_error) != NANOARROW_OK || + !IsVariantBinary(child_view) || child.dictionary != nullptr) { + return InvalidSchema("Variant storage field '{}' must be binary", name); + } + if (name == "metadata" && !has_metadata) { + if ((child.flags & ARROW_FLAG_NULLABLE) != 0) { + return InvalidSchema("Variant metadata must be required"); + } + has_metadata = true; + } else if (name == "value" && !has_value) { + has_value = true; + value_nullable = (child.flags & ARROW_FLAG_NULLABLE) != 0; + } else { + return InvalidSchema("Unexpected or duplicate variant storage field '{}'", name); + } + } + if (!has_metadata || (!has_value && !has_typed_value)) { + return InvalidSchema("Variant storage requires metadata and value or typed_value"); + } + if (has_typed_value != value_nullable) { + return InvalidSchema( + "Variant value must be required when unshredded and optional when shredded"); + } + return variant(); + } + switch (schema_view.type) { case NANOARROW_TYPE_STRUCT: { std::vector fields; diff --git a/src/iceberg/schema_internal.h b/src/iceberg/schema_internal.h index fd9ec6266..013b46757 100644 --- a/src/iceberg/schema_internal.h +++ b/src/iceberg/schema_internal.h @@ -32,6 +32,8 @@ namespace iceberg { /// \brief Canonical Arrow extension name used for Iceberg UUID values. inline constexpr const char* kArrowUuidExtensionName = "arrow.uuid"; +/// \brief Canonical Arrow extension name used for Iceberg variant values. +inline constexpr const char* kArrowVariantExtensionName = "arrow.parquet.variant"; /// \brief Convert an Iceberg schema to an Arrow schema. /// diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 8622a1a67..50a085732 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -223,8 +223,12 @@ if(ICEBERG_BUILD_BUNDLE) parquet_data_test.cc parquet_metrics_test.cc parquet_schema_test.cc + parquet_variant_test.cc parquet_test.cc row_number_reader_test.cc) + target_compile_definitions(parquet_test + PRIVATE ICEBERG_REQUIRE_ARROW_VARIANT=$ + ) if(ICEBERG_ENCRYPTION) target_sources(parquet_test PRIVATE parquet_encryption_test.cc) endif() diff --git a/src/iceberg/test/arrow_test.cc b/src/iceberg/test/arrow_test.cc index c66814809..ae1c2c0cc 100644 --- a/src/iceberg/test/arrow_test.cc +++ b/src/iceberg/test/arrow_test.cc @@ -30,6 +30,7 @@ #include #include +#include #include #include #include @@ -38,6 +39,7 @@ #include #include +#include "iceberg/arrow/arrow_register.h" #include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow/metadata_column_util_internal.h" #include "iceberg/constants.h" @@ -139,22 +141,54 @@ INSTANTIATE_TEST_SUITE_P( ToArrowSchemaParam{.iceberg_type = iceberg::unknown(), .arrow_type = ::arrow::null()})); -TEST(ToArrowSchemaTest, UnsupportedV3Types) { - const std::vector> unsupported_types = {iceberg::variant()}; - - for (const auto& unsupported_type : unsupported_types) { - Schema schema( - {SchemaField::MakeOptional(/*field_id=*/1, "unsupported", unsupported_type)}, - /*schema_id=*/0); - ArrowSchema arrow_schema{}; - ASSERT_THAT(ToArrowSchema(schema, &arrow_schema), - HasErrorMessage("is not supported by Arrow conversion")); - EXPECT_EQ(arrow_schema.release, nullptr); +TEST(ToArrowSchemaTest, VariantRoundtrip) { + arrow::RegisterAll(); + auto list_type = std::make_shared(3, variant(), true); + Schema schema({SchemaField::MakeOptional(1, "payload", variant()), + SchemaField::MakeRequired(2, "events", list_type)}); + ArrowSchema c_schema; + ASSERT_THAT(ToArrowSchema(schema, &c_schema), IsOk()); + auto imported = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + ASSERT_EQ(imported->field(0)->type()->id(), ::arrow::Type::EXTENSION); + const auto& extension = + static_cast(*imported->field(0)->type()); + EXPECT_EQ(extension.extension_name(), kArrowVariantExtensionName); + const auto& storage = extension.storage_type(); + ASSERT_EQ(storage->num_fields(), 2); + for (const auto& child : storage->fields()) { + EXPECT_EQ(child->type()->id(), ::arrow::Type::BINARY); + EXPECT_FALSE(child->nullable()); + EXPECT_TRUE(child->metadata() == nullptr || + !child->metadata()->Contains(kParquetFieldIdKey)); + } + ASSERT_TRUE(::arrow::ExportSchema(*imported, &c_schema).ok()); + auto roundtrip = FromArrowSchema(c_schema, 0); + c_schema.release(&c_schema); + ASSERT_THAT(roundtrip, IsOk()); + EXPECT_EQ(*roundtrip.value(), schema); +} + +TEST(FromArrowSchemaTest, RejectMalformedVariantStorage) { + auto metadata = + ::arrow::key_value_metadata({"ARROW:extension:name"}, {kArrowVariantExtensionName}); + for (const auto& storage : + {::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false)}), + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary()), + ::arrow::field("value", ::arrow::binary(), false)}), + ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::int32(), false)})}) { + auto schema = + ::arrow::schema({::arrow::field("payload", storage)->WithMetadata(metadata)}); + ArrowSchema c_schema; + ASSERT_TRUE(::arrow::ExportSchema(*schema, &c_schema).ok()); + auto result = FromArrowSchema(c_schema, 0); + c_schema.release(&c_schema); + EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); } } TEST(ToArrowSchemaTest, ReleasesSchemaOnConversionFailure) { - // fixed(0) passes CheckArrowCompatible but nanoarrow rejects its non-positive width. + // nanoarrow rejects fixed(0) after initializing the output schema. Schema schema({SchemaField::MakeOptional(/*field_id=*/1, "invalid_fixed", fixed(0))}, /*schema_id=*/0); ArrowSchema arrow_schema{}; @@ -163,6 +197,91 @@ TEST(ToArrowSchemaTest, ReleasesSchemaOnConversionFailure) { EXPECT_EQ(arrow_schema.release, nullptr); } +TEST(FromArrowSchemaTest, RejectMalformedVariantTypedValue) { + auto metadata = ::arrow::key_value_metadata( + {"ARROW:extension:name", std::string(kParquetFieldIdKey)}, + {kArrowVariantExtensionName, "1"}); + auto value = ::arrow::field("value", ::arrow::binary()); + for (const auto& typed : + {::arrow::field("typed_value", ::arrow::int32(), false), + ::arrow::field("typed_value", ::arrow::map(::arrow::utf8(), ::arrow::int32())), + ::arrow::field("typed_value", ::arrow::struct_({})), + ::arrow::field("typed_value", + ::arrow::struct_({::arrow::field("x", ::arrow::int32())})), + ::arrow::field("typed_value", ::arrow::list(::arrow::field( + "element", ::arrow::struct_({value})))), + ::arrow::field( + "typed_value", + ::arrow::struct_({::arrow::field( + "x", + ::arrow::struct_({::arrow::field("value", ::arrow::binary(), false)}), + false)})), + ::arrow::field("typed_value", + ::arrow::struct_({::arrow::field( + "x", ::arrow::struct_({value, value}), false)})), + ::arrow::field( + "typed_value", + ::arrow::struct_({::arrow::field( + "x", ::arrow::struct_({::arrow::field("typed_value", ::arrow::uint64())}), + false)})), + ::arrow::field("typed_value", ::arrow::timestamp(::arrow::TimeUnit::MILLI)), + ::arrow::field("typed_value", ::arrow::decimal128(10, -1))}) { + SCOPED_TRACE(typed->ToString()); + auto storage = ::arrow::struct_( + {::arrow::field("metadata", ::arrow::binary(), false), value, typed}); + auto schema = + ::arrow::schema({::arrow::field("payload", storage)->WithMetadata(metadata)}); + ArrowSchema c_schema; + ASSERT_TRUE(::arrow::ExportSchema(*schema, &c_schema).ok()); + auto result = FromArrowSchema(c_schema, 0); + c_schema.release(&c_schema); + EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); + } +} + +TEST(FromArrowSchemaTest, AcceptShreddedVariantTypedValue) { + auto metadata = ::arrow::key_value_metadata( + {"ARROW:extension:name", std::string(kParquetFieldIdKey)}, + {kArrowVariantExtensionName, "1"}); + auto value = ::arrow::field("value", ::arrow::binary()); + auto group = ::arrow::struct_({value, ::arrow::field("typed_value", ::arrow::int64())}); + for (const auto& type : + {::arrow::int64(), ::arrow::extension::uuid(), ::arrow::decimal128(30, 2), + ::arrow::struct_({::arrow::field("count", group, false)}), + ::arrow::list(::arrow::field("element", group, false)), + ::arrow::large_list(::arrow::field("element", group, false))}) { + for (bool typed_only : {false, true}) { + ::arrow::FieldVector fields{::arrow::field("metadata", ::arrow::binary(), false)}; + if (!typed_only) fields.push_back(value); + fields.push_back(::arrow::field("typed_value", type)); + auto schema = ::arrow::schema( + {::arrow::field("payload", ::arrow::struct_(fields))->WithMetadata(metadata)}); + ArrowSchema c_schema; + ASSERT_TRUE(::arrow::ExportSchema(*schema, &c_schema).ok()); + auto result = FromArrowSchema(c_schema, 0); + c_schema.release(&c_schema); + ASSERT_THAT(result, IsOk()); + EXPECT_TRUE(result.value()->fields()[0].type()->is_variant()); + } + } +} + +TEST(FromArrowSchemaTest, VariantWithoutFieldId) { + auto metadata = + ::arrow::key_value_metadata({"ARROW:extension:name"}, {kArrowVariantExtensionName}); + auto storage = ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false), + ::arrow::field("value", ::arrow::binary(), false)}); + auto schema = + ::arrow::schema({::arrow::field("payload", storage)->WithMetadata(metadata)}); + ArrowSchema c_schema; + ASSERT_TRUE(::arrow::ExportSchema(*schema, &c_schema).ok()); + auto result = FromArrowSchema(c_schema, 0); + c_schema.release(&c_schema); + ASSERT_THAT(result, IsOk()); + EXPECT_TRUE(result.value()->fields()[0].type()->is_variant()); + EXPECT_EQ(result.value()->fields()[0].field_id(), -1); +} + namespace { void CheckArrowField(const ::arrow::Field& field, ::arrow::Type::type type_id, diff --git a/src/iceberg/test/parquet_variant_test.cc b/src/iceberg/test/parquet_variant_test.cc new file mode 100644 index 000000000..70a65130e --- /dev/null +++ b/src/iceberg/test/parquet_variant_test.cc @@ -0,0 +1,800 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "iceberg/arrow/arrow_io_internal.h" +#include "iceberg/arrow/arrow_register.h" +#include "iceberg/arrow_c_data_util_internal.h" +#include "iceberg/file_reader.h" +#include "iceberg/file_writer.h" +#include "iceberg/name_mapping.h" +#include "iceberg/parquet/parquet_reader.h" +#include "iceberg/parquet/parquet_schema_util_internal.h" +#include "iceberg/parquet/parquet_variant_util_internal.h" +#include "iceberg/parquet/parquet_writer.h" +#include "iceberg/schema.h" +#include "iceberg/schema_internal.h" +#include "iceberg/test/matchers.h" +#include "iceberg/test/std_io.h" +#include "iceberg/test/temp_file_test_base.h" +#include "iceberg/type.h" + +#if ICEBERG_REQUIRE_ARROW_VARIANT +static_assert(ICEBERG_HAS_ARROW_VARIANT, + "Variant tests require Arrow's variant APIs; they must not be skipped"); +#endif + +#if ICEBERG_HAS_ARROW_VARIANT +# include +# include +#endif + +namespace iceberg::parquet { + +TEST(ParquetVariantSchemaTest, StorageChildrenHaveNoFieldIds) { + Schema schema({SchemaField::MakeRequired(1, "payload", variant())}); + auto result = ToParquetSchema(schema); + ASSERT_THAT(result, IsOk()); + auto group = static_cast( + result.value()->schema_root().get()); + auto field = static_cast(group->field(0).get()); + EXPECT_TRUE(field->logical_type()->is_variant()); + EXPECT_EQ(field->field_id(), 1); + ASSERT_EQ(field->field_count(), 2); + for (int i = 0; i < 2; ++i) { + EXPECT_EQ(field->field(i)->field_id(), -1); + EXPECT_EQ(field->field(i)->repetition(), ::parquet::Repetition::REQUIRED); + } +} + +#if !ICEBERG_HAS_ARROW_VARIANT + +TEST(ParquetVariantTest, MissingArrowUtilitiesFailsBeforeCreatingFile) { + auto schema = std::make_shared( + std::vector{SchemaField::MakeOptional(1, "payload", variant())}); + auto io = std::make_shared(); + ParquetWriter writer; + EXPECT_THAT(writer.Open({.path = "unused-variant.parquet", .schema = schema, .io = io}), + HasErrorMessage("requires Arrow's variant validation")); + EXPECT_FALSE(std::filesystem::exists("unused-variant.parquet")); +} + +#else + +namespace { + +using ::parquet::variant::VariantArrayBuilder; +using ::parquet::variant::VariantPrimitiveType; +using ::parquet::variant::VariantValueView; + +std::shared_ptr< ::arrow::Schema> ArrowSchemaFor(const Schema& schema) { + ArrowSchema c_schema; + arrow::RegisterAll(); + if (!ToArrowSchema(schema, &c_schema).has_value()) { + return nullptr; + } + return ::arrow::ImportSchema(&c_schema).ValueOrDie(); +} + +std::shared_ptr< ::arrow::Array> CanonicalVariant( + const std::shared_ptr< ::arrow::Array>& array, + const std::shared_ptr< ::arrow::DataType>& type) { + auto extension = std::static_pointer_cast< ::arrow::ExtensionArray>(array); + auto storage_type = + std::static_pointer_cast< ::arrow::ExtensionType>(type)->storage_type(); + auto storage = ::arrow::compute::Cast(extension->storage(), storage_type).ValueOrDie(); + return ::arrow::ExtensionType::WrapArray(type, storage.make_array()); +} + +// Compare logical objects and arrays by value, including integer widths and +// short-string versus long-string encodings, rather than by serialized bytes. +nlohmann::json SemanticValue(const VariantValueView& value) { + return std::visit( + [&](const auto& view) -> nlohmann::json { + using View = std::decay_t; + if constexpr (std::is_same_v) { + auto object = nlohmann::json::object(); + for (const auto& field : view.fields()) { + object[field.name] = SemanticValue(*view.GetField(field.name)); + } + return object; + } else if constexpr (std::is_same_v) { + auto array = nlohmann::json::array(); + for (size_t i = 0; i < view.elements().size(); ++i) { + array.push_back(SemanticValue(*view.GetElement(i))); + } + return array; + } else if constexpr (std::is_same_v) { + return std::string(view.string()); + } else { + auto type = view.type(); + if (type == VariantPrimitiveType::kNull) return nullptr; + if (type == VariantPrimitiveType::kBooleanTrue) return true; + if (type == VariantPrimitiveType::kBooleanFalse) return false; + if (type >= VariantPrimitiveType::kInt8 && + type <= VariantPrimitiveType::kInt64) { + auto payload = view.payload(); + uint64_t integer = 0; + for (size_t i = 0; i < payload.size(); ++i) { + integer |= static_cast(static_cast(payload[i])) + << (8 * i); + } + if (payload.size() < 8 && (static_cast(payload.back()) & 0x80)) { + integer |= ~uint64_t{0} << (payload.size() * 8); + } + return std::bit_cast(integer); + } + if (type == VariantPrimitiveType::kString) { + // The primitive string payload starts with a four-byte length. + return std::string(view.payload().substr(4)); + } + return nlohmann::json::array( + {static_cast(type), std::string(view.payload())}); + } + }, + value.data()); +} + +void ExpectVariantsEqual(const ::arrow::Array& expected, const ::arrow::Array& actual) { + ASSERT_EQ(expected.length(), actual.length()); + const auto& expected_storage = static_cast( + *static_cast(expected).storage()); + const auto& actual_storage = static_cast( + *static_cast(actual).storage()); + const auto& expected_metadata = + static_cast(*expected_storage.field(0)); + const auto& actual_metadata = + static_cast(*actual_storage.field(0)); + const auto& expected_values = + static_cast(*expected_storage.field(1)); + const auto& actual_values = + static_cast(*actual_storage.field(1)); + for (int64_t i = 0; i < expected.length(); ++i) { + ASSERT_EQ(expected.IsNull(i), actual.IsNull(i)) << i; + if (expected.IsNull(i)) continue; + auto em = ::parquet::variant::VariantMetadataView::Make(expected_metadata.GetView(i)); + auto am = ::parquet::variant::VariantMetadataView::Make(actual_metadata.GetView(i)); + EXPECT_EQ( + SemanticValue(VariantValueView::MakeWithValidate(expected_values.GetView(i), em)), + SemanticValue(VariantValueView::MakeWithValidate(actual_values.GetView(i), am))) + << i; + } +} + +} // namespace + +class ParquetVariantTest : public TempFileTestBase { + protected: + std::shared_ptr io_ = std::make_shared(); + std::shared_ptr schema_ = std::make_shared( + std::vector{SchemaField::MakeOptional(1, "payload", variant())}); + + std::shared_ptr< ::arrow::Array> Roundtrip( + const std::shared_ptr< ::arrow::Array>& input, + std::shared_ptr target = nullptr, int batches = 1, int32_t field_id = 1, + std::shared_ptr read_schema = nullptr) { + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + WriterOptions options{.path = path, .schema = schema_, .io = io_}; + options.properties.Set(WriterProperties::kParquetCompression, + std::string("uncompressed")); + if (target) options.variant_shredding_types.emplace(field_id, std::move(target)); + ParquetWriter writer; + auto opened = writer.Open(options); + EXPECT_THAT(opened, IsOk()); + if (!opened) return nullptr; + auto arrow_schema = ArrowSchemaFor(*schema_); + for (int i = 0; i < batches; ++i) { + auto batch = ::arrow::RecordBatch::Make(arrow_schema, input->length(), {input}); + ArrowArray c_array; + EXPECT_TRUE(::arrow::ExportRecordBatch(*batch, &c_array).ok()); + EXPECT_THAT(writer.Write(&c_array), IsOk()); + } + EXPECT_THAT(writer.Close(), IsOk()); + auto metrics = writer.metrics(); + EXPECT_THAT(metrics, IsOk()); + if (metrics) { + EXPECT_FALSE(metrics->lower_bounds.contains(1)); + EXPECT_FALSE(metrics->upper_bounds.contains(1)); + } + auto stream = arrow::OpenArrowInputStream(io_, path, std::nullopt).value(); + auto file = ::parquet::ParquetFileReader::Open(stream); + const auto& descr = file->metadata()->schema(); + auto root = + static_cast(descr->schema_root().get()); + if (schema_->fields()[0].type()->is_variant()) { + EXPECT_TRUE(root->field(0)->logical_type()->is_variant()); + } + for (int i = 0; i < descr->num_columns(); ++i) { + auto node = descr->Column(i)->schema_node().get(); + for (auto parent = node->parent(); parent != nullptr; parent = parent->parent()) { + if (parent->logical_type()->is_variant()) { + EXPECT_EQ(node->field_id(), -1); + } + } + } + ::parquet::arrow::SchemaManifest manifest; + auto properties = ::parquet::default_arrow_reader_properties(); + properties.set_arrow_extensions_enabled(true); + EXPECT_TRUE( + ::parquet::arrow::SchemaManifest::Make(descr, nullptr, properties, &manifest) + .ok()); + auto projection = Project(*schema_, manifest); + EXPECT_THAT(projection, IsOk()); + if (projection) + EXPECT_EQ(SelectedColumnIndices(*projection).size(), descr->num_columns()); + file->Close(); + + ParquetReader reader; + if (!read_schema) read_schema = schema_; + EXPECT_THAT(reader.Open({.path = path, .io = io_, .projection = read_schema}), + IsOk()); + auto c_schema = reader.Schema().value(); + auto output_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + EXPECT_TRUE(output_schema->Equals(*ArrowSchemaFor(*read_schema))); + std::vector > chunks; + while (true) { + auto next = reader.Next(); + EXPECT_THAT(next, IsOk()); + if (!next || !next->has_value()) break; + auto c_array = next->value(); + auto batch = ::arrow::ImportRecordBatch(&c_array, output_schema).ValueOrDie(); + chunks.push_back(batch->column(0)); + } + EXPECT_THAT(reader.Close(), IsOk()); + if (chunks.empty()) + return ::arrow::MakeArrayOfNull(arrow_schema->field(0)->type(), 0).ValueOrDie(); + return ::arrow::Concatenate(chunks).ValueOrDie(); + } +}; + +TEST_F(ParquetVariantTest, UnshreddedValues) { + VariantArrayBuilder builder; + builder.AppendNull(); + builder.AppendVariantNull(); + builder.AppendInt8(-42); + builder.AppendString("hello"); + builder.AppendDate(12345); + builder.AppendTimestampNanos(123456789, true); + builder.AppendBinary(std::string("\0\xff", 2)); + builder.AppendDecimal16(::arrow::Decimal128(12345), 2); + auto object = builder.StartObject(); + object.AppendString("name", "event"); + auto list = object.StartList("items"); + list.AppendBoolean(true); + list.AppendVariantNull(); + list.AppendInt64(9); + list.Finish(); + object.Finish(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto output = Roundtrip(input); + ASSERT_NE(output, nullptr); + ExpectVariantsEqual(*input, *output); +} + +TEST_F(ParquetVariantTest, PrimitiveShreddingAndResidualValues) { + VariantArrayBuilder builder; + builder.AppendInt8(-42); + builder.AppendInt64(100); + builder.AppendString("residual"); + builder.AppendNull(); + builder.AppendVariantNull(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto output = Roundtrip(input, int64(), 2); + ASSERT_NE(output, nullptr); + ASSERT_EQ(output->length(), input->length() * 2); + ExpectVariantsEqual(*input, *output->Slice(0, input->length())); + ExpectVariantsEqual(*input, *output->Slice(input->length())); +} + +TEST_F(ParquetVariantTest, ObjectAndListShredding) { + VariantArrayBuilder builder; + auto object = builder.StartObject(); + object.AppendInt8("count", 7); + object.AppendString("other", "residual field"); + auto list = object.StartList("items"); + list.AppendInt32(3); + list.AppendString("mixed"); + list.AppendVariantNull(); + list.Finish(); + object.Finish(); + auto missing = builder.StartObject(); + missing.AppendVariantNull("count"); + missing.Finish(); + builder.AppendInt64(11); + builder.AppendNull(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto target = std::make_shared(std::vector{ + SchemaField::MakeOptional(91, "count", int64()), + SchemaField::MakeOptional(92, "items", + std::make_shared(93, int64(), true))}); + auto output = Roundtrip(input, target); + ASSERT_NE(output, nullptr); + ExpectVariantsEqual(*input, *output); +} + +TEST_F(ParquetVariantTest, ShreddingPreservesPrimitiveTypes) { + VariantArrayBuilder builder; + builder.AppendBoolean(true); + builder.AppendBoolean(false); + builder.AppendInt64(42); + builder.AppendFloat(1.25F); + builder.AppendDouble(1.25); + builder.AppendString("true"); + builder.AppendString("12.34"); + builder.AppendDecimal16(::arrow::Decimal128(1234), 2); + builder.AppendDecimal16(::arrow::Decimal128(1230), 2); + builder.AppendTimestampMicros(123456, false); + builder.AppendTimestampNanos(123456000, false); + builder.AppendNull(); + builder.AppendVariantNull(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + for (const auto& target : std::vector >{ + boolean(), int32(), int64(), float32(), float64(), decimal(30, 2), + decimal(30, 1), timestamp_ns()}) { + SCOPED_TRACE(target->ToString()); + auto output = Roundtrip(input, target); + ASSERT_NE(output, nullptr); + ExpectVariantsEqual(*input, *output); + } +} + +TEST_F(ParquetVariantTest, NestedShreddingPreservesPrimitiveTypes) { + VariantArrayBuilder builder; + auto object = builder.StartObject(); + object.AppendBoolean("count", true); + auto list = object.StartList("items"); + list.AppendInt8(7); + list.AppendBoolean(false); + list.AppendDouble(2.0); + list.Finish(); + object.Finish(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto target = std::make_shared(std::vector{ + SchemaField::MakeOptional(91, "count", int64()), + SchemaField::MakeOptional(92, "items", + std::make_shared(93, int64(), true))}); + auto output = Roundtrip(input, target); + ASSERT_NE(output, nullptr); + ExpectVariantsEqual(*input, *output); +} + +TEST_F(ParquetVariantTest, EmptyAndAllNullBatches) { + auto type = ArrowSchemaFor(*schema_)->field(0)->type(); + for (int64_t length : {0, 3}) { + auto input = ::arrow::MakeArrayOfNull(type, length).ValueOrDie(); + auto output = Roundtrip(input, int64()); + ASSERT_NE(output, nullptr); + ExpectVariantsEqual(*input, *output); + } +} + +TEST_F(ParquetVariantTest, MissingOptionalVariantUsesCompleteAnchor) { + VariantArrayBuilder builder; + builder.AppendInt64(1); + builder.AppendNull(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto projected = std::make_shared( + std::vector{SchemaField::MakeOptional(5, "added_later", variant())}); + auto output = Roundtrip(input, int64(), 1, 1, projected); + ASSERT_NE(output, nullptr); + EXPECT_EQ(output->length(), 2); + EXPECT_EQ(output->null_count(), 2); +} + +TEST_F(ParquetVariantTest, NullStructParent) { + VariantArrayBuilder builder; + builder.AppendInt64(7); + builder.AppendInt64(11); + builder.AppendVariantNull(); + auto values = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + schema_ = std::make_shared(std::vector{SchemaField::MakeOptional( + 1, "payload", + std::make_shared( + std::vector{SchemaField::MakeOptional(2, "nested", variant())}))}); + auto type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto bitmap = ::arrow::AllocateEmptyBitmap(3).ValueOrDie(); + bitmap->mutable_data()[0] = 0x05; + auto input = + ::arrow::StructArray::Make({values}, type->fields(), bitmap, 1).ValueOrDie(); + auto output = Roundtrip(input, int64(), 1, 2); + ASSERT_NE(output, nullptr); + EXPECT_TRUE(output->IsNull(1)); + const auto& child = static_cast(*output).field(0); + ExpectVariantsEqual(*values->Slice(0, 1), *child->Slice(0, 1)); + ExpectVariantsEqual(*values->Slice(2, 1), *child->Slice(2, 1)); +} + +TEST_F(ParquetVariantTest, RequiredVariantRejectsSqlNull) { + schema_ = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "payload", variant())}); + auto type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto input = ::arrow::MakeArrayOfNull(type, 1).ValueOrDie(); + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + ParquetWriter writer; + WriterProperties properties; + properties.Set(WriterProperties::kParquetCompression, std::string("uncompressed")); + ASSERT_THAT(writer.Open({.path = path, + .schema = schema_, + .io = io_, + .properties = std::move(properties)}), + IsOk()); + auto batch = ::arrow::RecordBatch::Make(ArrowSchemaFor(*schema_), 1, {input}); + ArrowArray c_array; + ASSERT_TRUE(::arrow::ExportRecordBatch(*batch, &c_array).ok()); + EXPECT_FALSE(writer.Write(&c_array).has_value()); + EXPECT_THAT(writer.Close(), IsOk()); +} + +TEST_F(ParquetVariantTest, RequiredVariantUnderNullStructParent) { + VariantArrayBuilder builder; + builder.AppendInt64(7); + builder.AppendInt64(11); + builder.AppendInt64(13); + auto values = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + schema_ = std::make_shared(std::vector{SchemaField::MakeOptional( + 1, "payload", + std::make_shared( + std::vector{SchemaField::MakeRequired(2, "nested", variant())}))}); + auto type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto bitmap = ::arrow::AllocateEmptyBitmap(3).ValueOrDie(); + bitmap->mutable_data()[0] = 0x05; + auto input = + ::arrow::StructArray::Make({values}, type->fields(), bitmap, 1).ValueOrDie(); + for (const auto& target : std::vector >{nullptr, int64()}) { + auto output = Roundtrip(input->Slice(1, 2), target, 1, 2); + ASSERT_NE(output, nullptr); + ASSERT_EQ(output->length(), 2); + EXPECT_TRUE(output->IsNull(0)); + const auto& child = static_cast(*output).field(0); + ExpectVariantsEqual(*values->Slice(2, 1), *child->Slice(1, 1)); + } +} + +TEST_F(ParquetVariantTest, RequiredVariantInLargeListWithNullElements) { + VariantArrayBuilder builder; + builder.AppendInt64(7); + builder.AppendInt64(11); + builder.AppendInt64(13); + auto values = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto element = std::make_shared( + std::vector{SchemaField::MakeRequired(3, "nested", variant())}); + schema_ = std::make_shared(std::vector{SchemaField::MakeOptional( + 1, "payload", std::make_shared(2, element, true))}); + auto arrow_schema = ArrowSchemaFor(*schema_); + auto bitmap = ::arrow::AllocateEmptyBitmap(3).ValueOrDie(); + bitmap->mutable_data()[0] = 0x02; + auto entries = + ::arrow::StructArray::Make( + {values}, arrow_schema->field(0)->type()->field(0)->type()->fields(), bitmap, 2) + .ValueOrDie(); + auto offsets = + ::arrow::json::ArrayFromJSONString(::arrow::int32(), "[0, 2, 3]").ValueOrDie(); + auto input = + ::arrow::ListArray::FromArrays(arrow_schema->field(0)->type(), *offsets, *entries) + .ValueOrDie(); + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + ParquetWriter writer; + WriterOptions options{.path = path, .schema = schema_, .io = io_}; + options.properties.Set(WriterProperties::kParquetCompression, + std::string("uncompressed")); + options.variant_shredding_types.emplace(3, int64()); + ASSERT_THAT(writer.Open(options), IsOk()); + auto batch = ::arrow::RecordBatch::Make(arrow_schema, 2, {input}); + ArrowArray c_array; + ASSERT_TRUE(::arrow::ExportRecordBatch(*batch, &c_array).ok()); + ASSERT_THAT(writer.Write(&c_array), IsOk()); + ASSERT_THAT(writer.Close(), IsOk()); + + ReaderOptions read_options{.path = path, .io = io_, .projection = schema_}; + read_options.properties.Set(ReaderProperties::kArrowUseLargeList, true); + read_options.properties.Set(ReaderProperties::kBatchSize, int64_t{1}); + ParquetReader reader; + ASSERT_THAT(reader.Open(read_options), IsOk()); + auto c_schema = reader.Schema().value(); + auto output_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + ASSERT_EQ(output_schema->field(0)->type()->id(), ::arrow::Type::LARGE_LIST); + for (int i = 0; i < 2; ++i) { + auto next = reader.Next(); + ASSERT_THAT(next, IsOk()); + ASSERT_TRUE(next->has_value()); + auto output_array = next->value(); + auto output = ::arrow::ImportRecordBatch(&output_array, output_schema).ValueOrDie(); + auto list = std::static_pointer_cast< ::arrow::LargeListArray>(output->column(0)); + auto elements = std::static_pointer_cast< ::arrow::StructArray>(list->value_slice(0)); + ASSERT_EQ(elements->length(), i == 0 ? 2 : 1); + EXPECT_TRUE(elements->IsNull(0)); + if (i == 0) { + ExpectVariantsEqual(*values->Slice(1, 1), *elements->field(0)->Slice(1, 1)); + } + } + EXPECT_FALSE(reader.Next().value().has_value()); + EXPECT_THAT(reader.Close(), IsOk()); +} + +TEST_F(ParquetVariantTest, RequiredVariantInMapWithNullValues) { + VariantArrayBuilder builder; + builder.AppendInt64(7); + builder.AppendInt64(11); + auto values = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + auto item = std::make_shared( + std::vector{SchemaField::MakeRequired(4, "nested", variant())}); + schema_ = std::make_shared(std::vector{SchemaField::MakeOptional( + 1, "payload", + std::make_shared(SchemaField::MakeRequired(2, "key", string()), + SchemaField::MakeOptional(3, "value", item)))}); + auto map_type = std::static_pointer_cast< ::arrow::MapType>( + ArrowSchemaFor(*schema_)->field(0)->type()); + auto bitmap = ::arrow::AllocateEmptyBitmap(2).ValueOrDie(); + bitmap->mutable_data()[0] = 0x02; + auto items = + ::arrow::StructArray::Make({values}, map_type->item_type()->fields(), bitmap, 1) + .ValueOrDie(); + auto offsets = + ::arrow::json::ArrayFromJSONString(::arrow::int32(), "[0, 2]").ValueOrDie(); + auto keys = + ::arrow::json::ArrayFromJSONString(::arrow::utf8(), "[\"a\", \"b\"]").ValueOrDie(); + auto input = ::arrow::MapArray::FromArrays(map_type, offsets, keys, items).ValueOrDie(); + auto output = Roundtrip(input, int64(), 1, 4); + ASSERT_NE(output, nullptr); + ASSERT_EQ(output->length(), 1); + auto output_items = std::static_pointer_cast< ::arrow::StructArray>( + static_cast(*output).items()); + EXPECT_TRUE(output_items->IsNull(0)); + ExpectVariantsEqual(*values->Slice(1, 1), *output_items->field(0)->Slice(1, 1)); +} + +TEST_F(ParquetVariantTest, NestedVariantsAndSlicedArrays) { + VariantArrayBuilder builder; + builder.AppendInt8(7); + builder.AppendString("residual"); + builder.AppendNull(); + builder.AppendVariantNull(); + auto canonical_type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto values = CanonicalVariant(builder.Finish(), canonical_type); + + // A sliced struct retains an offset into its unsliced child arrays. + auto structure = std::make_shared( + std::vector{SchemaField::MakeOptional(2, "nested", variant())}); + schema_ = std::make_shared( + std::vector{SchemaField::MakeOptional(1, "payload", structure)}); + auto structure_type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto input_struct = + ::arrow::StructArray::Make({values}, structure_type->fields()).ValueOrDie(); + auto output_struct = Roundtrip(input_struct->Slice(1, 3), int64(), 1, 2); + ASSERT_NE(output_struct, nullptr); + ExpectVariantsEqual(*values->Slice(1, 3), + *static_cast(*output_struct).field(0)); + + auto offsets = + ::arrow::json::ArrayFromJSONString(::arrow::int32(), "[0, 2, 2, 4]").ValueOrDie(); + schema_ = std::make_shared(std::vector{SchemaField::MakeOptional( + 1, "payload", std::make_shared(2, variant(), true))}); + auto list_type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto input_list = + ::arrow::ListArray::FromArrays(list_type, *offsets, *values).ValueOrDie(); + auto output_list = Roundtrip(input_list, int64(), 1, 2); + ASSERT_NE(output_list, nullptr); + ExpectVariantsEqual(*values, + *static_cast(*output_list).values()); + + schema_ = std::make_shared(std::vector{SchemaField::MakeOptional( + 1, "payload", + std::make_shared(SchemaField::MakeRequired(2, "key", string()), + SchemaField::MakeOptional(3, "value", variant())))}); + auto keys = + ::arrow::json::ArrayFromJSONString(::arrow::utf8(), "[\"a\",\"b\",\"c\",\"d\"]") + .ValueOrDie(); + auto map_type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto input_map = + ::arrow::MapArray::FromArrays(map_type, offsets, keys, values).ValueOrDie(); + auto output_map = Roundtrip(input_map, int64(), 1, 3); + ASSERT_NE(output_map, nullptr); + ExpectVariantsEqual(*values, + *static_cast(*output_map).items()); +} + +TEST_F(ParquetVariantTest, CorruptValuesReturnStatus) { + auto type = ArrowSchemaFor(*schema_)->field(0)->type(); + auto metadata = + ::arrow::json::ArrayFromJSONString(::arrow::binary(), "[\"invalid\"]").ValueOrDie(); + auto values = + ::arrow::json::ArrayFromJSONString(::arrow::binary(), "[\"invalid\"]").ValueOrDie(); + auto storage_type = static_cast(*type).storage_type(); + auto storage = + ::arrow::StructArray::Make({metadata, values}, storage_type->fields()).ValueOrDie(); + auto array = ::arrow::ExtensionType::WrapArray(type, storage); + ParquetWriter writer; + WriterProperties writer_properties; + writer_properties.Set(WriterProperties::kParquetCompression, + std::string("uncompressed")); + ASSERT_THAT(writer.Open({.path = CreateNewTempFilePathWithSuffix(".parquet"), + .schema = schema_, + .io = io_, + .properties = std::move(writer_properties)}), + IsOk()); + ArrowArray c_array; + auto batch = ::arrow::RecordBatch::Make(ArrowSchemaFor(*schema_), 1, {array}); + ASSERT_TRUE(::arrow::ExportRecordBatch(*batch, &c_array).ok()); + EXPECT_FALSE(writer.Write(&c_array).has_value()); + EXPECT_EQ(c_array.release, nullptr); + EXPECT_THAT(writer.Close(), IsOk()); + + // Write a malformed file directly through Arrow with validation disabled, + // then check that Iceberg's reader validates it and returns a Status. + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + auto stream = arrow::OpenArrowOutputStream(io_, path).value(); + auto properties = ::parquet::ArrowWriterProperties::Builder() + .set_variant_validation_enabled(false) + ->build(); + auto unchecked = ::parquet::arrow::FileWriter::Open( + *batch->schema(), ::arrow::default_memory_pool(), stream, + ::parquet::default_writer_properties(), properties) + .ValueOrDie(); + ASSERT_TRUE(unchecked->WriteRecordBatch(*batch).ok()); + ASSERT_TRUE(unchecked->Close().ok()); + ASSERT_TRUE(stream->Close().ok()); + ParquetReader reader; + ASSERT_THAT(reader.Open({.path = path, .io = io_, .projection = schema_}), IsOk()); + EXPECT_FALSE(reader.Next().has_value()); + EXPECT_THAT(reader.Close(), IsOk()); +} + +TEST_F(ParquetVariantTest, ReadTypedOnlyExternalLayout) { + VariantArrayBuilder builder; + builder.AppendInt64(42); + builder.AppendInt64(-7); + builder.AppendNull(); + auto input = + CanonicalVariant(builder.Finish(), ArrowSchemaFor(*schema_)->field(0)->type()); + const auto& storage = static_cast( + *static_cast(*input).storage()); + const auto& metadata = static_cast(*storage.field(0)); + for (int32_t field_id : {1, -1}) { + SCOPED_TRACE(field_id); + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + auto stream = arrow::OpenArrowOutputStream(io_, path).value(); + auto metadata_node = ::parquet::schema::PrimitiveNode::Make( + "metadata", ::parquet::Repetition::REQUIRED, ::parquet::Type::BYTE_ARRAY); + auto typed_node = ::parquet::schema::PrimitiveNode::Make( + "typed_value", ::parquet::Repetition::OPTIONAL, ::parquet::Type::INT64); + auto variant_node = ::parquet::schema::GroupNode::Make( + "payload", ::parquet::Repetition::OPTIONAL, {metadata_node, typed_node}, + ::parquet::LogicalType::Variant(), field_id); + auto root = ::parquet::schema::GroupNode::Make( + "schema", ::parquet::Repetition::REQUIRED, {variant_node}); + auto writer = ::parquet::ParquetFileWriter::Open( + stream, std::static_pointer_cast< ::parquet::schema::GroupNode>(root)); + auto row_group = writer->AppendRowGroup(); + int16_t metadata_defs[] = {1, 1, 0}; + int16_t typed_defs[] = {2, 2, 0}; + ::parquet::ByteArray encoded_metadata[] = { + ::parquet::ByteArray( + static_cast(metadata.GetView(0).size()), + reinterpret_cast(metadata.GetView(0).data())), + ::parquet::ByteArray( + static_cast(metadata.GetView(1).size()), + reinterpret_cast(metadata.GetView(1).data()))}; + int64_t integers[] = {42, -7}; + auto metadata_writer = + static_cast< ::parquet::ByteArrayWriter*>(row_group->NextColumn()); + EXPECT_EQ(metadata_writer->WriteBatch(3, metadata_defs, nullptr, encoded_metadata), + 2); + metadata_writer->Close(); + auto typed_writer = static_cast< ::parquet::Int64Writer*>(row_group->NextColumn()); + EXPECT_EQ(typed_writer->WriteBatch(3, typed_defs, nullptr, integers), 2); + typed_writer->Close(); + row_group->Close(); + writer->Close(); + ASSERT_TRUE(stream->Close().ok()); + + ParquetReader reader; + auto mapping = CreateMapping(*schema_).value(); + ASSERT_THAT(reader.Open({.path = path, + .io = io_, + .projection = schema_, + .name_mapping = std::move(mapping)}), + IsOk()); + auto next = reader.Next(); + ASSERT_THAT(next, IsOk()); + ASSERT_TRUE(next->has_value()); + auto c_array = next->value(); + auto batch = + ::arrow::ImportRecordBatch(&c_array, ArrowSchemaFor(*schema_)).ValueOrDie(); + ExpectVariantsEqual(*input, *batch->column(0)); + EXPECT_THAT(reader.Close(), IsOk()); + } +} + +TEST_F(ParquetVariantTest, BatchFilteringPreservesExtensionAndNulls) { + VariantArrayBuilder builder; + builder.AppendInt64(42); + builder.AppendNull(); + builder.AppendVariantNull(); + auto arrow_schema = ArrowSchemaFor(*schema_); + auto input = CanonicalVariant(builder.Finish(), arrow_schema->field(0)->type()); + int32_t rows[] = {2, 0, 1}; + for (auto function : {ProjectionContext::ProjectBatchFunction{nullptr}, + ProjectionContext::ResolveProjectBatchFunction()}) { + auto projection = ProjectionContext::Make(*schema_, *schema_, function); + ASSERT_THAT(projection, IsOk()); + auto batch = ::arrow::RecordBatch::Make(arrow_schema, 3, {input}); + ArrowArray c_array; + ASSERT_TRUE(::arrow::ExportRecordBatch(*batch, &c_array).ok()); + auto projected = ProjectBatch(&c_array, rows, *projection); + ASSERT_THAT(projected, IsOk()); + auto output_c_array = projected.value(); + auto output = ::arrow::ImportRecordBatch(&output_c_array, arrow_schema).ValueOrDie(); + ExpectVariantsEqual(*input->Slice(2, 1), *output->column(0)->Slice(0, 1)); + ExpectVariantsEqual(*input->Slice(0, 2), *output->column(0)->Slice(1, 2)); + } +} + +TEST_F(ParquetVariantTest, InvalidTargetsFailBeforeCreatingFile) { + for (auto targets : + {std::unordered_map >{{999, int64()}}, + std::unordered_map >{{1, nullptr}}, + std::unordered_map >{{1, variant()}}, + std::unordered_map >{{1, unknown()}}}) { + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + ParquetWriter writer; + EXPECT_THAT(writer.Open({.path = path, + .schema = schema_, + .io = io_, + .variant_shredding_types = targets}), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_FALSE(std::filesystem::exists(path)); + } +} + +#endif + +} // namespace iceberg::parquet diff --git a/src/iceberg/test/schema_test.cc b/src/iceberg/test/schema_test.cc index 90cfdb13c..a1b86fe8d 100644 --- a/src/iceberg/test/schema_test.cc +++ b/src/iceberg/test/schema_test.cc @@ -99,6 +99,13 @@ TEST(SchemaTest, Equality) { } TEST(SchemaTest, ValidateRejectsV3TypesBeforeFormatV3) { + iceberg::Schema variant_schema( + {iceberg::SchemaField(1, "payload", iceberg::variant(), true)}); + for (int32_t version : {1, 2}) { + EXPECT_THAT(variant_schema.Validate(version), + iceberg::IsError(iceberg::ErrorKind::kInvalidSchema)); + } + EXPECT_THAT(variant_schema.Validate(3), iceberg::IsOk()); iceberg::Schema timestamp_ns_schema( {iceberg::SchemaField(1, "timestamp_ns", iceberg::timestamp_ns(), false)}); iceberg::Schema timestamptz_ns_schema(