From 33d7ea3fc80dbc3108e114dc42950c396654e46c Mon Sep 17 00:00:00 2001 From: Zhao Junwang Date: Thu, 15 Oct 2026 00:15:47 +0800 Subject: [PATCH] feat: add Iceberg v3 variant Parquet support Map variant schemas to Arrow and Parquet, read unshredded and shredded variants, and support explicit shredding targets when writing Parquet. Preserve nested projections and null semantics, validate variant values, and add coverage for round trips, residuals, filtering, and invalid input. Pin Arrow PR #50252 to 0f16412bb1c39611aeed7393eba09e637f4d86fa with archive SHA256 verification. Require its variant APIs by default, force the pinned dependency in CI, and document the build options. --- .github/workflows/cpp-linter.yml | 2 + .github/workflows/sanitizer_test.yml | 2 + CMakeLists.txt | 4 + ci/scripts/build_iceberg.sh | 4 + .../IcebergThirdpartyToolchain.cmake | 70 +- mkdocs/docs/getting-started.md | 9 +- mkdocs/docs/library.md | 37 + src/iceberg/CMakeLists.txt | 1 + src/iceberg/arrow/arrow_register.cc | 14 + src/iceberg/arrow/arrow_register.h | 2 +- src/iceberg/file_writer.h | 6 + src/iceberg/parquet/parquet_data_util.cc | 12 +- src/iceberg/parquet/parquet_reader.cc | 26 +- src/iceberg/parquet/parquet_reader.h | 2 +- src/iceberg/parquet/parquet_schema_util.cc | 151 +++- .../parquet/parquet_schema_util_internal.h | 9 +- src/iceberg/parquet/parquet_variant_util.cc | 690 +++++++++++++++ .../parquet/parquet_variant_util_internal.h | 65 ++ src/iceberg/parquet/parquet_writer.cc | 35 +- src/iceberg/parquet/parquet_writer.h | 2 +- src/iceberg/schema_internal.cc | 213 ++++- src/iceberg/schema_internal.h | 2 + src/iceberg/test/CMakeLists.txt | 4 + src/iceberg/test/arrow_test.cc | 143 +++- src/iceberg/test/parquet_variant_test.cc | 800 ++++++++++++++++++ src/iceberg/test/schema_test.cc | 7 + 26 files changed, 2215 insertions(+), 97 deletions(-) create mode 100644 src/iceberg/parquet/parquet_variant_util.cc create mode 100644 src/iceberg/parquet/parquet_variant_util_internal.h create mode 100644 src/iceberg/test/parquet_variant_test.cc diff --git a/.github/workflows/cpp-linter.yml b/.github/workflows/cpp-linter.yml index be59a610d5..6f598b02bb 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 b2d1e011a8..f51b6ca107 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 7a9c3b2ff1..b267ae9e59 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 640b772692..a9db77ce23 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 f94ad70318..0b00da848f 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 ab0741fd9a..5daa4f4dd5 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 6b769425ad..ac4c55ed12 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 c751cd7a6e..6924145889 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 4fedc914cc..1eeeeb8ab4 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 7a67be5168..a74362475f 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 4b1045a409..6579e402f4 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 73597cbfb8..e4a93119d2 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 549a3d3aed..a85382b089 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 3a8e57cf3a..f3b6563866 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 2a230fe4cc..abcc10ab9a 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 6612de0569..f80df30be1 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 0000000000..afe6ddcfb3 --- /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 0000000000..fba5a2d7d4 --- /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 08c3941383..4ba735cec7 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 562ffb8013..543fd5757c 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 0a61898843..14d3f73140 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 fd9ec62661..013b46757d 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 8622a1a67f..50a0857323 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 c66814809f..ae1c2c0ccb 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 0000000000..70a65130e5 --- /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 90cfdb13c9..a1b86fe8d5 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(