From 3184397686b1f31f5e1eb436c6b08498344f001c Mon Sep 17 00:00:00 2001 From: Nicholas Jiang Date: Thu, 9 Jul 2026 20:59:43 +0800 Subject: [PATCH] feat(variant): support variant data type Add the VARIANT data type for storing and efficiently querying semi-structured data (e.g. JSON), following the parquet-format VariantEncoding.md / VariantShredding.md specifications. - Type system: VARIANT is parsed from / serialized to "VARIANT" and is represented in Arrow as struct marked with paimon.extension.type=paimon.type.variant; field-id assignment and schema validation treat variant structs as leaves. - Public API: paimon::Variant builds variants from JSON, renders them back to JSON, and extracts values by path (e.g. $.a[0]) cast to scalar targets or, via VariantGetArrow, to nested STRUCT / LIST / MAP / VARIANT targets. - Parquet: a variant column is written as an unannotated group of REQUIRED BINARY value (field id 0) and metadata (field id 1); ORC/Avro report NotImplemented for VARIANT columns. - Shredding writes: a generic write-plan framework shared with MAP shared-shredding rewrites variant columns into typed parquet sub-columns, either from the configured variant.shreddingSchema or inferred per file from sampled rows when variant.inferShreddingSchema is enabled (tuned by the variant.shredding.* options). - Shredding reads: shredded files are detected structurally and reassembled back into variants transparently; a variant column can also be read as a variant-access projection that extracts specific paths while pruning the scan to metadata and the requested typed_value sub-columns. Co-Authored-By: Claude Fable 5 --- docs/source/user_guide/data_types.rst | 38 + include/paimon/data/variant.h | 197 ++++++ include/paimon/defs.h | 23 + src/paimon/CMakeLists.txt | 28 + .../map_shared_shredding_batch_converter.h | 7 +- ...ap_shared_shredding_write_plan_factory.cpp | 69 ++ .../map_shared_shredding_write_plan_factory.h | 67 ++ .../shredding/shredding_batch_converter.h | 44 ++ .../data/shredding/shredding_file_reader.cpp | 132 ++++ .../data/shredding/shredding_file_reader.h | 65 ++ .../data/shredding/shredding_read_plan.h | 50 ++ .../shredding_write_plan_factories.cpp | 46 ++ .../shredding_write_plan_factories.h | 47 ++ .../shredding/shredding_write_plan_factory.h | 69 ++ .../common/data/variant/generic_variant.cpp | 422 ++++++++++++ .../common/data/variant/generic_variant.h | 144 ++++ .../data/variant/generic_variant_test.cpp | 348 ++++++++++ .../infer_variant_shredding_schema.cpp | 361 ++++++++++ .../variant/infer_variant_shredding_schema.h | 56 ++ .../infer_variant_shredding_schema_test.cpp | 195 ++++++ src/paimon/common/data/variant/variant.cpp | 187 +++++ .../data/variant/variant_access_utils.cpp | 179 +++++ .../data/variant/variant_access_utils.h | 77 +++ .../data/variant/variant_binary_util.cpp | 589 ++++++++++++++++ .../common/data/variant/variant_binary_util.h | 198 ++++++ .../common/data/variant/variant_builder.cpp | 652 ++++++++++++++++++ .../common/data/variant/variant_builder.h | 144 ++++ src/paimon/common/data/variant/variant_defs.h | 161 +++++ .../common/data/variant/variant_get.cpp | 407 +++++++++++ src/paimon/common/data/variant/variant_get.h | 83 +++ .../common/data/variant/variant_get_test.cpp | 309 +++++++++ .../data/variant/variant_json_utils.cpp | 294 ++++++++ .../common/data/variant/variant_json_utils.h | 59 ++ .../data/variant/variant_path_segment.cpp | 115 +++ .../data/variant/variant_path_segment.h | 57 ++ .../data/variant/variant_reassembler.cpp | 260 +++++++ .../common/data/variant/variant_reassembler.h | 68 ++ .../common/data/variant/variant_schema.h | 95 +++ .../variant_shredding_batch_converter.cpp | 106 +++ .../variant_shredding_batch_converter.h | 62 ++ .../variant_shredding_read_plan_factory.cpp | 341 +++++++++ .../variant_shredding_read_plan_factory.h | 51 ++ .../data/variant/variant_shredding_test.cpp | 248 +++++++ .../data/variant/variant_shredding_utils.cpp | 291 ++++++++ .../data/variant/variant_shredding_utils.h | 61 ++ .../variant/variant_shredding_write_plan.cpp | 81 +++ .../variant/variant_shredding_write_plan.h | 81 +++ .../variant_shredding_write_plan_factory.cpp | 144 ++++ .../variant_shredding_write_plan_factory.h | 72 ++ ...iant_shredding_write_plan_factory_test.cpp | 132 ++++ .../data/variant/variant_shredding_writer.cpp | 468 +++++++++++++ .../data/variant/variant_shredding_writer.h | 103 +++ .../common/data/variant/variant_test.cpp | 107 +++ .../data/variant/variant_type_utils.cpp | 125 ++++ .../common/data/variant/variant_type_utils.h | 75 ++ .../data/variant/variant_type_utils_test.cpp | 102 +++ src/paimon/common/defs.cpp | 9 + src/paimon/common/types/data_type.cpp | 15 +- .../common/types/data_type_json_parser.cpp | 54 +- .../types/data_type_json_parser_test.cpp | 12 + src/paimon/common/types/data_type_test.cpp | 14 + src/paimon/common/utils/field_type_utils.h | 15 + src/paimon/core/append/append_only_writer.cpp | 6 +- src/paimon/core/core_options.cpp | 36 + src/paimon/core/core_options.h | 9 + .../core/io/infer_shredding_file_writer.h | 175 +++++ .../io/infer_shredding_file_writer_test.cpp | 212 ++++++ src/paimon/core/io/rolling_file_writer.h | 4 +- ...edding_append_data_file_writer_factory.cpp | 47 +- ...hredding_append_data_file_writer_factory.h | 12 +- ...ing_key_value_data_file_writer_factory.cpp | 50 +- ...dding_key_value_data_file_writer_factory.h | 12 +- src/paimon/core/io/single_file_writer.h | 6 +- .../compact/merge_tree_compact_rewriter.cpp | 6 +- .../core/mergetree/merge_tree_writer.cpp | 8 +- .../core/operation/abstract_split_read.cpp | 31 + .../core/operation/abstract_split_read.h | 7 + .../append_only_file_store_write.cpp | 6 +- .../core/operation/internal_read_context.cpp | 10 + .../core/postpone/postpone_bucket_writer.cpp | 8 +- .../core/schema/arrow_schema_validator.cpp | 17 + .../schema/arrow_schema_validator_test.cpp | 44 ++ src/paimon/core/schema/schema_validation.cpp | 7 + src/paimon/core/schema/table_schema.cpp | 8 +- src/paimon/core/schema/table_schema_test.cpp | 15 + .../core/utils/nested_projection_utils.cpp | 13 + .../format/avro/avro_schema_converter.cpp | 4 + .../avro/avro_schema_converter_test.cpp | 8 + src/paimon/format/orc/orc_format_writer.cpp | 4 + .../format/orc/orc_format_writer_test.cpp | 18 + src/paimon/format/parquet/CMakeLists.txt | 1 + .../format/parquet/variant_parquet_test.cpp | 568 +++++++++++++++ src/paimon/testing/utils/test_helper.h | 35 + src/paimon/testing/utils/variant_test_data.h | 64 ++ test/inte/CMakeLists.txt | 7 + test/inte/variant_table_inte_test.cpp | 233 +++++++ 96 files changed, 10727 insertions(+), 75 deletions(-) create mode 100644 include/paimon/data/variant.h create mode 100644 src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.cpp create mode 100644 src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.h create mode 100644 src/paimon/common/data/shredding/shredding_batch_converter.h create mode 100644 src/paimon/common/data/shredding/shredding_file_reader.cpp create mode 100644 src/paimon/common/data/shredding/shredding_file_reader.h create mode 100644 src/paimon/common/data/shredding/shredding_read_plan.h create mode 100644 src/paimon/common/data/shredding/shredding_write_plan_factories.cpp create mode 100644 src/paimon/common/data/shredding/shredding_write_plan_factories.h create mode 100644 src/paimon/common/data/shredding/shredding_write_plan_factory.h create mode 100644 src/paimon/common/data/variant/generic_variant.cpp create mode 100644 src/paimon/common/data/variant/generic_variant.h create mode 100644 src/paimon/common/data/variant/generic_variant_test.cpp create mode 100644 src/paimon/common/data/variant/infer_variant_shredding_schema.cpp create mode 100644 src/paimon/common/data/variant/infer_variant_shredding_schema.h create mode 100644 src/paimon/common/data/variant/infer_variant_shredding_schema_test.cpp create mode 100644 src/paimon/common/data/variant/variant.cpp create mode 100644 src/paimon/common/data/variant/variant_access_utils.cpp create mode 100644 src/paimon/common/data/variant/variant_access_utils.h create mode 100644 src/paimon/common/data/variant/variant_binary_util.cpp create mode 100644 src/paimon/common/data/variant/variant_binary_util.h create mode 100644 src/paimon/common/data/variant/variant_builder.cpp create mode 100644 src/paimon/common/data/variant/variant_builder.h create mode 100644 src/paimon/common/data/variant/variant_defs.h create mode 100644 src/paimon/common/data/variant/variant_get.cpp create mode 100644 src/paimon/common/data/variant/variant_get.h create mode 100644 src/paimon/common/data/variant/variant_get_test.cpp create mode 100644 src/paimon/common/data/variant/variant_json_utils.cpp create mode 100644 src/paimon/common/data/variant/variant_json_utils.h create mode 100644 src/paimon/common/data/variant/variant_path_segment.cpp create mode 100644 src/paimon/common/data/variant/variant_path_segment.h create mode 100644 src/paimon/common/data/variant/variant_reassembler.cpp create mode 100644 src/paimon/common/data/variant/variant_reassembler.h create mode 100644 src/paimon/common/data/variant/variant_schema.h create mode 100644 src/paimon/common/data/variant/variant_shredding_batch_converter.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_batch_converter.h create mode 100644 src/paimon/common/data/variant/variant_shredding_read_plan_factory.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_read_plan_factory.h create mode 100644 src/paimon/common/data/variant/variant_shredding_test.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_utils.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_utils.h create mode 100644 src/paimon/common/data/variant/variant_shredding_write_plan.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_write_plan.h create mode 100644 src/paimon/common/data/variant/variant_shredding_write_plan_factory.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_write_plan_factory.h create mode 100644 src/paimon/common/data/variant/variant_shredding_write_plan_factory_test.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_writer.cpp create mode 100644 src/paimon/common/data/variant/variant_shredding_writer.h create mode 100644 src/paimon/common/data/variant/variant_test.cpp create mode 100644 src/paimon/common/data/variant/variant_type_utils.cpp create mode 100644 src/paimon/common/data/variant/variant_type_utils.h create mode 100644 src/paimon/common/data/variant/variant_type_utils_test.cpp create mode 100644 src/paimon/core/io/infer_shredding_file_writer.h create mode 100644 src/paimon/core/io/infer_shredding_file_writer_test.cpp create mode 100644 src/paimon/format/parquet/variant_parquet_test.cpp create mode 100644 src/paimon/testing/utils/variant_test_data.h create mode 100644 test/inte/variant_table_inte_test.cpp diff --git a/docs/source/user_guide/data_types.rst b/docs/source/user_guide/data_types.rst index e5690841c..74339163e 100644 --- a/docs/source/user_guide/data_types.rst +++ b/docs/source/user_guide/data_types.rst @@ -217,3 +217,41 @@ and `Arrow DataTypes `` where n is the unique name of a field, t is the logical type of a field, d is the description of a field. + + * - ``VARIANT`` + - Struct + - Data type of semi-structured data (e.g. JSON). A variant value contains one of: + a primitive (e.g. integer, string), an array of variant values, or an object + mapping string keys to variant values. + + Variant values are encoded with two binaries following the parquet-format + `Variant Binary Encoding `_ + specification (compatible with the Java Paimon / Spark implementation). In + C++ Paimon, a variant field is represented in an Arrow schema as + ``Struct{value: Binary NOT NULL, metadata: Binary NOT NULL}`` marked with + Paimon-specific field metadata; use ``paimon::Variant::ArrowField`` to + construct such a field, and ``paimon::Variant`` (``FromJson``/``ToJson``/ + ``VariantGet``) to build and inspect values. + + Only the parquet file format supports VARIANT columns. VARIANT cannot be + used as a primary key, partition key or bucket key, and no predicate + pushdown applies to it. + + When writing, variant columns can optionally be *shredded* into typed + parquet columns per the parquet-format + `Variant Shredding `_ + specification by setting ``variant.shreddingSchema`` to a ROW type JSON + whose fields map variant column names to their shredding types. + Alternatively, setting ``variant.inferShreddingSchema`` to ``true`` + infers a shredding schema per file from the first written rows + (tuned by ``variant.shredding.maxSchemaWidth``, + ``variant.shredding.maxSchemaDepth``, + ``variant.shredding.minFieldCardinalityRatio`` and + ``variant.shredding.maxInferBufferRow``). Readers reassemble shredded + columns transparently. + + When reading, instead of the full variant, specific paths can be + extracted by replacing the variant column in the read schema with a + projection built by ``paimon::VariantAccessBuilder`` (e.g. ``$.a.b`` as + BIGINT). For shredded files only the required typed sub-columns are + read. diff --git a/include/paimon/data/variant.h b/include/paimon/data/variant.h new file mode 100644 index 000000000..e4a6fae6e --- /dev/null +++ b/include/paimon/data/variant.h @@ -0,0 +1,197 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/result.h" +#include "paimon/visibility.h" + +struct ArrowArray; +struct ArrowSchema; + +namespace paimon { + +/// Arguments controlling how a variant value is cast to a target type in `Variant::VariantGet`. +struct PAIMON_EXPORT VariantCastArgs { + /// Whether an invalid cast fails the call (true) or yields SQL NULL (false). + bool fail_on_error = true; + /// The time zone used when rendering TIMESTAMP values. Supported forms are `UTC`/`Z`/`GMT`, + /// fixed offsets such as `+08:00`, and IANA region ids such as `Asia/Shanghai`. + std::string zone_id = "UTC"; +}; + +/// A Variant represents a type that contains one of: 1) Primitive: A type and corresponding +/// value (e.g. INT, STRING); 2) Array: An ordered list of Variant values; 3) Object: An +/// unordered collection of string/Variant pairs (i.e. key/value pairs). An object may not +/// contain duplicate keys. +/// +/// A Variant is encoded with 2 binaries: the value and the metadata, following the parquet +/// Variant Binary Encoding specification (compatible with the Java / Spark implementation). The +/// encoding allows representation of semi-structured data (e.g. JSON) in a form that can be +/// efficiently queried by path. +/// +/// In an Arrow schema, a Variant field is represented as +/// `struct` marked with Paimon-specific field +/// metadata; use `Variant::ArrowField` to construct such a field. +class PAIMON_EXPORT Variant { + public: + ~Variant(); + + /// Parses a JSON string as a Variant (duplicate object keys are rejected). + /// + /// @param json The JSON document text. + /// @param pool The memory pool used for the variant buffers. + /// @return A result containing the created variant or an error. + static Result> FromJson(const std::string& json, + const std::shared_ptr& pool); + + /// Creates a Variant from already-encoded value and metadata binaries (copied into `pool`). + /// + /// @param value The variant value binary. + /// @param value_length The length of the value binary. + /// @param metadata The variant metadata binary. + /// @param metadata_length The length of the metadata binary. + /// @param pool The memory pool used for the variant buffers. + /// @return A result containing the created variant or an error. + static Result> Create(const char* value, uint64_t value_length, + const char* metadata, uint64_t metadata_length, + const std::shared_ptr& pool); + + /// The variant value binary. The view remains valid as long as this variant exists. + std::string_view Value() const; + + /// The variant metadata binary. The view remains valid as long as this variant exists. + std::string_view Metadata() const; + + /// The size of the variant in bytes (value size + metadata size). + int64_t SizeInBytes() const; + + /// Stringifies the variant in JSON format. + /// + /// @param zone_id The time zone used when rendering TIMESTAMP values. + /// @return A result containing the JSON text or an error. + Result ToJson(const std::string& zone_id = "UTC") const; + + /// Extracts a sub-variant value according to a path which starts with a `$`, e.g. `$.key`, + /// `$['key']`, `$["key"]`, `$.array[0]`, and casts the value to the target type. + /// + /// @param path The extraction path. + /// @param target_type The target Arrow type (C data interface, consumed by the call); only + /// scalar types are supported. Use `VariantGetArrow` for nested targets. + /// @param cast_args Cast behavior arguments. + /// @return A result containing the extracted literal, or nullopt for SQL NULL (unmatched + /// path, variant null, or an invalid cast with `fail_on_error == false`). + Result> VariantGet(const std::string& path, + struct ArrowSchema* target_type, + const VariantCastArgs& cast_args) const; + + /// Extracts a sub-variant value according to a path which starts with a `$` and casts the + /// value to the (possibly nested) target field type. + /// + /// In addition to scalar types, the target may be a STRUCT (cast from a variant object, + /// children matched by field name; unmatched children are null), a MAP with string keys + /// (cast from a variant object), a LIST (cast from a variant array), or a variant-marked + /// field created by `Variant::ArrowField` (the sub-variant is deeply re-encoded). + /// + /// @param path The extraction path. + /// @param target_field The target Arrow field (C data interface, consumed by the call). + /// @param cast_args Cast behavior arguments. + /// @return A result containing a length-1 Arrow array of the target type whose single slot + /// holds the result; a null slot represents SQL NULL (unmatched path, variant null, + /// or an invalid cast with `fail_on_error == false`). + Result> VariantGetArrow( + const std::string& path, struct ArrowSchema* target_field, + const VariantCastArgs& cast_args) const; + + /// Extracts a sub-variant value according to a path which starts with a `$` and renders it + /// as JSON text. + /// + /// @param path The extraction path. + /// @param zone_id The time zone used when rendering TIMESTAMP values. + /// @return A result containing the JSON text, or nullopt if the path does not match. + Result> VariantGetJson(const std::string& path, + const std::string& zone_id = "UTC") const; + + /// Creates an Arrow field definition for the Variant type. + /// + /// This function constructs an Arrow Field (internally + /// `struct`) and exports it to the C data + /// interface structure `::ArrowSchema`. It automatically injects Paimon-specific metadata to + /// identify the field as a VARIANT. + /// + /// @param field_name The name of the Arrow field. + /// @param nullable Whether the field is nullable. + /// @param metadata A map of key-value metadata to be attached to the field. + /// @return A result containing a unique pointer to the generated `::ArrowSchema` or an error. + static Result> ArrowField( + const std::string& field_name, bool nullable = true, + std::unordered_map metadata = {}); + + private: + class Impl; + + explicit Variant(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +/// Builds a variant-access projection field: a struct field that replaces a VARIANT column in +/// the read schema so that, instead of the full variant, only the described paths are extracted +/// at read time (reading only the required shredded sub-columns from shredded files). +/// +/// Example: read `$.age` as INT64 and `$.city` as STRING from variant column `v`: +/// +/// VariantAccessBuilder builder; +/// builder.AddField(age_type, "$.age"); +/// builder.AddField(city_type, "$.city"); +/// auto field = builder.Build("v"); // use in ReadContextBuilder::SetReadSchema +/// +/// The resulting struct has one child per added field, named by its position ("0", "1", ...). +class PAIMON_EXPORT VariantAccessBuilder { + public: + VariantAccessBuilder(); + ~VariantAccessBuilder(); + + /// Adds an extracted field. + /// + /// @param target_type The Arrow type the extracted value is cast to (C data interface, + /// consumed by the call). Nested targets follow `Variant::VariantGetArrow` semantics. + /// @param path The extraction path, e.g. `$.a.b` or `$.array[0]`. + /// @param fail_on_error Whether an invalid cast fails the read (true) or yields SQL NULL. + /// @param zone_id The time zone used when rendering TIMESTAMP values. + Status AddField(struct ArrowSchema* target_type, const std::string& path, + bool fail_on_error = true, const std::string& zone_id = "UTC"); + + /// Builds the projection field named `field_name`, to be used in the read schema in place + /// of the variant column with the same name. + Result> Build(const std::string& field_name) const; + + private: + class Impl; + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/include/paimon/defs.h b/include/paimon/defs.h index bdce64a5c..7a5b4385a 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -46,6 +46,7 @@ enum class FieldType { MAP = 14, STRUCT = 15, BLOB = 16, + VARIANT = 17, UNKNOWN = 128, }; @@ -391,6 +392,28 @@ struct PAIMON_EXPORT Options { /// Only effective when map.storage-layout = shared-shredding. static const char MAP_SHARED_SHREDDING_COLUMN_PLACEMENT_POLICY[]; + /// "variant.shreddingSchema" - The Variant shredding schema for writing: a ROW type JSON + /// whose fields map variant column names to their shredding types. No default value. + static const char VARIANT_SHREDDING_SCHEMA[]; + /// "parquet.variant.shreddingSchema" - Fallback key of "variant.shreddingSchema". + static const char PARQUET_VARIANT_SHREDDING_SCHEMA[]; + /// "variant.inferShreddingSchema" - Whether to automatically infer the shredding schema when + /// writing Variant columns. Default value is "false". + static const char VARIANT_INFER_SHREDDING_SCHEMA[]; + /// "variant.shredding.maxSchemaWidth" - Maximum number of shredded fields allowed in an + /// inferred schema. Default value is 300. + static const char VARIANT_SHREDDING_MAX_SCHEMA_WIDTH[]; + /// "variant.shredding.maxSchemaDepth" - Maximum traversal depth in Variant values during + /// schema inference. Default value is 50. + static const char VARIANT_SHREDDING_MAX_SCHEMA_DEPTH[]; + /// "variant.shredding.minFieldCardinalityRatio" - Minimum fraction of rows that must contain + /// a field for it to be shredded. Fields below this threshold stay in the un-shredded + /// Variant binary. Default value is 0.1. + static const char VARIANT_SHREDDING_MIN_FIELD_CARDINALITY_RATIO[]; + /// "variant.shredding.maxInferBufferRow" - Maximum number of rows to buffer for schema + /// inference. Default value is 4096. + static const char VARIANT_SHREDDING_MAX_INFER_BUFFER_ROW[]; + /// "blob-as-descriptor" - Read blob field using blob descriptor rather than blob /// bytes. Default value is "false". static const char BLOB_AS_DESCRIPTOR[]; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index b044badea..5cb15a9d4 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -38,6 +38,23 @@ set(PAIMON_COMMON_SRCS common/data/serializer/row_compacted_serializer.cpp common/data/serializer/binary_serializer_utils.cpp common/data/timestamp.cpp + common/data/variant/generic_variant.cpp + common/data/variant/infer_variant_shredding_schema.cpp + common/data/variant/variant.cpp + common/data/variant/variant_binary_util.cpp + common/data/variant/variant_builder.cpp + common/data/variant/variant_get.cpp + common/data/variant/variant_json_utils.cpp + common/data/variant/variant_path_segment.cpp + common/data/variant/variant_reassembler.cpp + common/data/variant/variant_shredding_batch_converter.cpp + common/data/variant/variant_access_utils.cpp + common/data/variant/variant_shredding_read_plan_factory.cpp + common/data/variant/variant_shredding_utils.cpp + common/data/variant/variant_shredding_write_plan.cpp + common/data/variant/variant_shredding_write_plan_factory.cpp + common/data/variant/variant_shredding_writer.cpp + common/data/variant/variant_type_utils.cpp common/defs.cpp common/executor/executor.cpp common/factories/singleton.cpp @@ -139,11 +156,14 @@ set(PAIMON_COMMON_SRCS common/utils/crc32c.cpp common/utils/decimal_utils.cpp common/data/shredding/map_shared_shredding_utils.cpp + common/data/shredding/map_shared_shredding_write_plan_factory.cpp + common/data/shredding/shredding_write_plan_factories.cpp common/data/shredding/map_shared_shredding_context.cpp common/data/shredding/map_shared_shredding_batch_converter.cpp common/data/shredding/map_shared_shredding_column_allocator.cpp common/data/shredding/lru_map_shared_shredding_column_allocator.cpp common/data/shredding/map_shared_shredding_file_reader.cpp + common/data/shredding/shredding_file_reader.cpp common/utils/delta_varint_compressor.cpp common/utils/fields_comparator.cpp common/utils/path_util.cpp @@ -428,6 +448,13 @@ if(PAIMON_BUILD_TESTS) common/data/blob_descriptor_test.cpp common/data/blob_view_struct_test.cpp common/data/blob_utils_test.cpp + common/data/variant/generic_variant_test.cpp + common/data/variant/infer_variant_shredding_schema_test.cpp + common/data/variant/variant_shredding_write_plan_factory_test.cpp + common/data/variant/variant_get_test.cpp + common/data/variant/variant_shredding_test.cpp + common/data/variant/variant_test.cpp + common/data/variant/variant_type_utils_test.cpp common/executor/default_executor_test.cpp common/format/column_stats_test.cpp common/fs/external_path_provider_test.cpp @@ -620,6 +647,7 @@ if(PAIMON_BUILD_TESTS) core/index/index_file_meta_serializer_test.cpp core/index/index_file_handler_test.cpp core/io/compact_increment_test.cpp + core/io/infer_shredding_file_writer_test.cpp core/io/concat_key_value_record_reader_test.cpp core/io/data_file_meta_serializer_test.cpp core/io/data_file_path_factory_test.cpp diff --git a/src/paimon/common/data/shredding/map_shared_shredding_batch_converter.h b/src/paimon/common/data/shredding/map_shared_shredding_batch_converter.h index 6a697e94f..7349d8b39 100644 --- a/src/paimon/common/data/shredding/map_shared_shredding_batch_converter.h +++ b/src/paimon/common/data/shredding/map_shared_shredding_batch_converter.h @@ -27,6 +27,7 @@ #include "paimon/common/data/shredding/map_shared_shredding_column_allocator.h" #include "paimon/common/data/shredding/map_shared_shredding_field_dict.h" #include "paimon/common/data/shredding/map_shredding_defs.h" +#include "paimon/common/data/shredding/shredding_batch_converter.h" #include "paimon/memory/memory_pool.h" #include "paimon/result.h" #include "paimon/status.h" @@ -44,7 +45,7 @@ class MapSharedShreddingContext; /// /// Non-shared-shredding columns are passed through unchanged. /// Each shared-shredding column has its own FieldDict and ColumnAllocator. -class MapSharedShreddingBatchConverter { +class MapSharedShreddingBatchConverter : public ShreddingBatchConverter { public: /// Creates a converter for one file write cycle. /// Computes per-file K from context, builds physical schema, and constructs the converter. @@ -59,12 +60,12 @@ class MapSharedShreddingBatchConverter { const std::shared_ptr& pool); /// Returns the physical schema produced for this converter. - const std::shared_ptr& GetPhysicalSchema() const; + const std::shared_ptr& GetPhysicalSchema() const override; /// Converts a logical batch to a physical batch. /// @param logical_batch Input ArrowArray (C ABI) with logical schema. Consumed on success. /// @return Owned physical ArrowArray (C ABI) with physical schema. - Result> Convert(ArrowArray* logical_batch); + Result> Convert(ArrowArray* logical_batch) override; /// Builds MapSharedShreddingFieldMeta for one shredding column (by field name). /// Called at file close to serialize metadata. diff --git a/src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.cpp b/src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.cpp new file mode 100644 index 000000000..39ad2cec1 --- /dev/null +++ b/src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.cpp @@ -0,0 +1,69 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/map_shared_shredding_write_plan_factory.h" + +#include + +#include "paimon/common/data/shredding/map_shared_shredding_batch_converter.h" +#include "paimon/common/data/shredding/map_shared_shredding_context.h" +#include "paimon/common/data/shredding/map_shared_shredding_utils.h" +#include "paimon/common/data/shredding/map_shredding_defs.h" + +namespace paimon { + +MapSharedShreddingWritePlanFactory::MapSharedShreddingWritePlanFactory( + const CoreOptions& options, const std::shared_ptr& write_schema, + const std::shared_ptr& context, + const std::shared_ptr& pool) + : options_(options), write_schema_(write_schema), context_(context), pool_(pool) {} + +bool MapSharedShreddingWritePlanFactory::ShouldCreateWritePlan() const { + return context_ != nullptr; +} + +bool MapSharedShreddingWritePlanFactory::ShouldInferWritePlan() const { + return false; +} + +int32_t MapSharedShreddingWritePlanFactory::InferBufferRowCount() const { + return 0; +} + +Result> +MapSharedShreddingWritePlanFactory::CreateConverter( + const std::string& file_format_identifier, + const std::vector>& sample_batches) const { + if (context_ == nullptr) { + return Status::Invalid("Shared-shredding write plan requires a shredding context."); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr converter, + MapSharedShreddingBatchConverter::Create(write_schema_, context_, options_, pool_)); + return std::shared_ptr(std::move(converter)); +} + +ShreddingWritePlanFactory::MetadataFinalizer +MapSharedShreddingWritePlanFactory::CreateMetadataFinalizer( + const std::shared_ptr& converter) const { + // The converter is created by CreateConverter above; the concrete type is guaranteed. + auto map_converter = std::static_pointer_cast(converter); + return MapSharedShreddingUtils::BuildMetadataFinalizer( + map_converter, MapSharedShreddingDefine::kDefaultDictCompression, context_, + map_converter->GetPhysicalSchema()); +} + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.h b/src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.h new file mode 100644 index 000000000..84ce7a31e --- /dev/null +++ b/src/paimon/common/data/shredding/map_shared_shredding_write_plan_factory.h @@ -0,0 +1,67 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/shredding_write_plan_factory.h" +#include "paimon/core/core_options.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class Schema; +} // namespace arrow + +namespace paimon { + +class MapSharedShreddingContext; + +/// Creates MAP shared-shredding batch converters driven by the cross-file adaptive-K context. +/// The write plan is never inferred from samples; per-file field metadata is persisted into the +/// file footer by the metadata finalizer. +class MapSharedShreddingWritePlanFactory : public ShreddingWritePlanFactory { + public: + MapSharedShreddingWritePlanFactory(const CoreOptions& options, + const std::shared_ptr& write_schema, + const std::shared_ptr& context, + const std::shared_ptr& pool); + + bool ShouldCreateWritePlan() const override; + + bool ShouldInferWritePlan() const override; + + int32_t InferBufferRowCount() const override; + + Result> CreateConverter( + const std::string& file_format_identifier, + const std::vector>& sample_batches) const override; + + MetadataFinalizer CreateMetadataFinalizer( + const std::shared_ptr& converter) const override; + + private: + CoreOptions options_; + std::shared_ptr write_schema_; + std::shared_ptr context_; + std::shared_ptr pool_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_batch_converter.h b/src/paimon/common/data/shredding/shredding_batch_converter.h new file mode 100644 index 000000000..7057bdc05 --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_batch_converter.h @@ -0,0 +1,44 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "arrow/c/abi.h" +#include "paimon/result.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +/// Converts logical write batches into a physical (shredded) file layout. Implemented by both +/// the VARIANT shredding and the MAP shared-shredding batch converters so that shredded file +/// writers can be composed generically. +class ShreddingBatchConverter { + public: + virtual ~ShreddingBatchConverter() = default; + + /// The physical file schema that converted batches conform to. + virtual const std::shared_ptr& GetPhysicalSchema() const = 0; + + /// Converts a logical batch into the physical layout. Consumes `logical_batch`. + virtual Result> Convert(::ArrowArray* logical_batch) = 0; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_file_reader.cpp b/src/paimon/common/data/shredding/shredding_file_reader.cpp new file mode 100644 index 000000000..b7b6d03ef --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_file_reader.cpp @@ -0,0 +1,132 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/shredding_file_reader.h" + +#include +#include + +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" + +namespace paimon { + +ShreddingFileReader::ShreddingFileReader( + std::unique_ptr&& reader, + std::map>&& plans, + const std::shared_ptr& pool) + : arrow_pool_(GetArrowPool(pool)), reader_(std::move(reader)), plans_(std::move(plans)) {} + +Result> ShreddingFileReader::GetFileSchema() const { + return reader_->GetFileSchema(); +} + +Status ShreddingFileReader::SetReadSchema(::ArrowSchema* read_schema, + const std::shared_ptr& predicate, + const std::optional& selection_bitmap) { + if (!read_schema) { + return Status::Invalid("invalid read schema in ShreddingFileReader, cannot be null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_read_schema, + arrow::ImportSchema(read_schema)); + arrow::FieldVector resolved_fields = logical_read_schema->fields(); + bool any_resolved = false; + for (auto& resolved_field : resolved_fields) { + auto it = plans_.find(resolved_field->name()); + if (it == plans_.end()) { + continue; + } + // Push the physical (possibly pruned) subtree down so the inner reader materializes it. + resolved_field = it->second->PhysicalField(); + any_resolved = true; + } + if (!any_resolved) { + return Status::Invalid("no planned shredded columns exist in the read schema"); + } + auto resolved_schema = arrow::schema(resolved_fields, logical_read_schema->metadata()); + auto c_resolved_schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*resolved_schema, c_resolved_schema.get())); + return reader_->SetReadSchema(c_resolved_schema.get(), predicate, selection_bitmap); +} + +Result ShreddingFileReader::NextBatch() { + return Status::Invalid( + "paimon inner reader ShreddingFileReader should use NextBatchWithBitmap"); +} + +Result ShreddingFileReader::NextBatchWithBitmap() { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader_->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + return batch_with_bitmap; + } + + auto& [batch, bitmap] = batch_with_bitmap; + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (arrow_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("cannot cast batch to StructArray in ShreddingFileReader"); + } + auto struct_array = std::static_pointer_cast(arrow_array); + + arrow::ArrayVector resolved_arrays = struct_array->fields(); + arrow::FieldVector resolved_fields = struct_array->struct_type()->fields(); + for (int32_t field_idx = 0; field_idx < struct_array->num_fields(); ++field_idx) { + const auto& physical_field = struct_array->struct_type()->field(field_idx); + auto it = plans_.find(physical_field->name()); + if (it == plans_.end()) { + continue; + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr logical_array, + it->second->Assemble(struct_array->field(field_idx), arrow_pool_.get())); + resolved_arrays[field_idx] = logical_array; + resolved_fields[field_idx] = it->second->LogicalField(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr new_struct_array, + arrow::StructArray::Make(resolved_arrays, resolved_fields)); + auto new_c_array = std::make_unique(); + auto new_c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*new_struct_array, new_c_array.get(), new_c_schema.get())); + batch = std::make_pair(std::move(new_c_array), std::move(new_c_schema)); + return batch_with_bitmap; +} + +std::shared_ptr ShreddingFileReader::GetReaderMetrics() const { + return reader_->GetReaderMetrics(); +} + +void ShreddingFileReader::Close() { + reader_->Close(); +} + +Result ShreddingFileReader::GetPreviousBatchFileRowId(uint64_t batch_row_id) const { + return reader_->GetPreviousBatchFileRowId(batch_row_id); +} + +Result ShreddingFileReader::GetNumberOfRows() const { + return reader_->GetNumberOfRows(); +} + +bool ShreddingFileReader::SupportPreciseBitmapSelection() const { + return reader_->SupportPreciseBitmapSelection(); +} + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_file_reader.h b/src/paimon/common/data/shredding/shredding_file_reader.h new file mode 100644 index 000000000..19eade2ed --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_file_reader.h @@ -0,0 +1,65 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "paimon/common/data/shredding/shredding_read_plan.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/reader/file_batch_reader.h" + +namespace paimon { + +/// A file batch reader wrapper that pushes per-column physical (shredded, possibly pruned) +/// subtrees down to the inner reader and assembles the physical batches back into the logical +/// columns as planned by `ShreddingColumnReadPlan`s. +class ShreddingFileReader : public FileBatchReader { + public: + ShreddingFileReader(std::unique_ptr&& reader, + std::map>&& plans, + const std::shared_ptr& pool); + + Result> GetFileSchema() const override; + + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) override; + + Result NextBatch() override; + + Result NextBatchWithBitmap() override; + + std::shared_ptr GetReaderMetrics() const override; + + void Close() override; + + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; + + Result GetNumberOfRows() const override; + + bool SupportPreciseBitmapSelection() const override; + + private: + std::shared_ptr arrow_pool_; + std::unique_ptr reader_; + std::map> plans_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_read_plan.h b/src/paimon/common/data/shredding/shredding_read_plan.h new file mode 100644 index 000000000..afbaa7ffe --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_read_plan.h @@ -0,0 +1,50 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class Field; +class MemoryPool; +} // namespace arrow + +namespace paimon { + +/// A per-column read plan translating between a column's logical shape and the physical +/// (shredded, possibly pruned) shape stored in one file: the physical field is pushed down to +/// the format reader and the physical batches are assembled back into logical arrays. +class ShreddingColumnReadPlan { + public: + virtual ~ShreddingColumnReadPlan() = default; + + /// The logical field restored on the assembled output. + virtual const std::shared_ptr& LogicalField() const = 0; + + /// The physical file field (possibly a pruned subtree) to read from the file. + virtual const std::shared_ptr& PhysicalField() const = 0; + + /// Assembles the physical column array back into the logical column array. + virtual Result> Assemble( + const std::shared_ptr& physical, arrow::MemoryPool* pool) const = 0; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_write_plan_factories.cpp b/src/paimon/common/data/shredding/shredding_write_plan_factories.cpp new file mode 100644 index 000000000..74909afdc --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_write_plan_factories.cpp @@ -0,0 +1,46 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/shredding_write_plan_factories.h" + +#include "paimon/common/data/shredding/map_shared_shredding_write_plan_factory.h" +#include "paimon/common/data/variant/variant_shredding_write_plan_factory.h" +#include "paimon/core/core_options.h" + +namespace paimon { + +std::shared_ptr ShreddingWritePlanFactories::SelectActive( + const CoreOptions& options, const std::shared_ptr& write_schema, + const std::shared_ptr& shredding_context, + const std::shared_ptr& pool) { + // MAP shared-shredding is active exactly when a context exists; constructing its factory + // copies the options, so skip it otherwise. + if (shredding_context != nullptr) { + auto map_factory = std::make_shared( + options, write_schema, shredding_context, pool); + if (map_factory->ShouldCreateWritePlan()) { + return map_factory; + } + } + auto variant_factory = + std::make_shared(options, write_schema, pool); + if (variant_factory->ShouldCreateWritePlan()) { + return variant_factory; + } + return std::shared_ptr(nullptr); +} + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_write_plan_factories.h b/src/paimon/common/data/shredding/shredding_write_plan_factories.h new file mode 100644 index 000000000..675f5a803 --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_write_plan_factories.h @@ -0,0 +1,47 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/shredding_write_plan_factory.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class CoreOptions; +class MapSharedShreddingContext; + +/// Composes the known shredding write-plan factories (MAP shared-shredding and VARIANT +/// shredding) and selects the one active for a write schema. +class ShreddingWritePlanFactories { + public: + /// Returns the single active write-plan factory for the write, or nullptr when no shredding + /// applies. MAP shared-shredding takes precedence over VARIANT shredding, preserving the + /// selection order of the writer call sites. + static std::shared_ptr SelectActive( + const CoreOptions& options, const std::shared_ptr& write_schema, + const std::shared_ptr& shredding_context, + const std::shared_ptr& pool); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/shredding/shredding_write_plan_factory.h b/src/paimon/common/data/shredding/shredding_write_plan_factory.h new file mode 100644 index 000000000..3b44c6498 --- /dev/null +++ b/src/paimon/common/data/shredding/shredding_write_plan_factory.h @@ -0,0 +1,69 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/common/data/shredding/shredding_batch_converter.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class Schema; +} // namespace arrow + +namespace paimon { + +/// Decides whether write batches must be rewritten into a physical (shredded) layout before +/// they reach the format writer, and creates the per-file batch converter -- either immediately +/// from configuration or, for inference, from sampled logical batches buffered by the writer. +class ShreddingWritePlanFactory { + public: + /// Finalizes per-file shredding metadata; the returned schema (or nullptr) is persisted into + /// the file footer right before the file is finished. + using MetadataFinalizer = std::function>()>; + + virtual ~ShreddingWritePlanFactory() = default; + + /// Whether a write plan (immediate or inferred) applies to the write schema. + virtual bool ShouldCreateWritePlan() const = 0; + + /// Whether the write plan must be inferred from sampled rows instead of configuration. + virtual bool ShouldInferWritePlan() const = 0; + + /// The number of rows buffered per file to sample the inferred write plan from. + virtual int32_t InferBufferRowCount() const = 0; + + /// Creates the per-file batch converter. `sample_batches` holds the logical batches sampled + /// for inference and is empty when the plan comes from configuration. Returns nullptr when + /// no conversion is useful for this file (the file is written with the logical schema). + virtual Result> CreateConverter( + const std::string& file_format_identifier, + const std::vector>& sample_batches) const = 0; + + /// The per-file metadata finalizer persisted into the file footer, or nullptr when the + /// physical schema is self-describing (as it is for VARIANT shredding). + virtual MetadataFinalizer CreateMetadataFinalizer( + const std::shared_ptr& converter) const { + return nullptr; + } +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/generic_variant.cpp b/src/paimon/common/data/variant/generic_variant.cpp new file mode 100644 index 000000000..dd77756ab --- /dev/null +++ b/src/paimon/common/data/variant/generic_variant.cpp @@ -0,0 +1,422 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#include "paimon/common/data/variant/generic_variant.h" + +#include +#include +#include + +#include "arrow/util/base64.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/data/variant/variant_json_utils.h" + +namespace paimon { + +namespace { + +Status ToJsonImpl(std::string_view value, std::string_view metadata, int32_t pos, + const std::string& zone_id, std::string* out) { + PAIMON_ASSIGN_OR_RAISE(VariantValueType type, VariantBinaryUtil::GetType(value, pos)); + switch (type) { + case VariantValueType::kObject: { + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ObjectInfo info, + VariantBinaryUtil::GetObjectInfo(value, pos)); + out->push_back('{'); + for (int32_t i = 0; i < info.num_elements; ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t id, + VariantBinaryUtil::ReadUnsigned( + value, info.id_start + info.id_size * i, info.id_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, + VariantBinaryUtil::ReadUnsigned(value, info.offset_start + info.offset_size * i, + info.offset_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t element_pos, + VariantBinaryUtil::CheckedElementPos(info.data_start, offset, value.size())); + if (i != 0) { + out->push_back(','); + } + PAIMON_ASSIGN_OR_RAISE(std::string_view key, + VariantBinaryUtil::GetMetadataKey(metadata, id)); + VariantJsonUtils::AppendEscapedJson(key, out); + out->push_back(':'); + PAIMON_RETURN_NOT_OK(ToJsonImpl(value, metadata, element_pos, zone_id, out)); + } + out->push_back('}'); + return Status::OK(); + } + case VariantValueType::kArray: { + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ArrayInfo info, + VariantBinaryUtil::GetArrayInfo(value, pos)); + out->push_back('['); + for (int32_t i = 0; i < info.num_elements; ++i) { + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, + VariantBinaryUtil::ReadUnsigned(value, info.offset_start + info.offset_size * i, + info.offset_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t element_pos, + VariantBinaryUtil::CheckedElementPos(info.data_start, offset, value.size())); + if (i != 0) { + out->push_back(','); + } + PAIMON_RETURN_NOT_OK(ToJsonImpl(value, metadata, element_pos, zone_id, out)); + } + out->push_back(']'); + return Status::OK(); + } + case VariantValueType::kNull: + out->append("null"); + return Status::OK(); + case VariantValueType::kBoolean: { + PAIMON_ASSIGN_OR_RAISE(bool b, VariantBinaryUtil::GetBoolean(value, pos)); + out->append(b ? "true" : "false"); + return Status::OK(); + } + case VariantValueType::kLong: { + PAIMON_ASSIGN_OR_RAISE(int64_t l, VariantBinaryUtil::GetLong(value, pos)); + out->append(std::to_string(l)); + return Status::OK(); + } + case VariantValueType::kString: { + PAIMON_ASSIGN_OR_RAISE(std::string_view s, VariantBinaryUtil::GetString(value, pos)); + VariantJsonUtils::AppendEscapedJson(s, out); + return Status::OK(); + } + case VariantValueType::kDouble: { + PAIMON_ASSIGN_OR_RAISE(double d, VariantBinaryUtil::GetDouble(value, pos)); + std::string repr = VariantJsonUtils::JavaDoubleToString(d); + if (std::isfinite(d)) { + out->append(repr); + } else { + out->push_back('"'); + out->append(repr); + out->push_back('"'); + } + return Status::OK(); + } + case VariantValueType::kDecimal: { + PAIMON_ASSIGN_OR_RAISE(VariantDecimal d, VariantBinaryUtil::GetDecimal(value, pos)); + out->append(d.ToPlainString()); + return Status::OK(); + } + case VariantValueType::kDate: { + PAIMON_ASSIGN_OR_RAISE(int64_t days, VariantBinaryUtil::GetLong(value, pos)); + out->push_back('"'); + out->append(VariantJsonUtils::DateToString(static_cast(days))); + out->push_back('"'); + return Status::OK(); + } + case VariantValueType::kTimestamp: { + PAIMON_ASSIGN_OR_RAISE(int64_t micros, VariantBinaryUtil::GetLong(value, pos)); + PAIMON_ASSIGN_OR_RAISE(int32_t offset_seconds, + VariantJsonUtils::GetZoneOffsetSeconds(zone_id, micros)); + out->push_back('"'); + out->append(VariantJsonUtils::TimestampToString(micros, offset_seconds, true)); + out->push_back('"'); + return Status::OK(); + } + case VariantValueType::kTimestampNtz: { + PAIMON_ASSIGN_OR_RAISE(int64_t micros, VariantBinaryUtil::GetLong(value, pos)); + out->push_back('"'); + out->append(VariantJsonUtils::TimestampToString(micros, 0, false)); + out->push_back('"'); + return Status::OK(); + } + case VariantValueType::kFloat: { + PAIMON_ASSIGN_OR_RAISE(float f, VariantBinaryUtil::GetFloat(value, pos)); + std::string repr = VariantJsonUtils::JavaFloatToString(f); + if (std::isfinite(f)) { + out->append(repr); + } else { + out->push_back('"'); + out->append(repr); + out->push_back('"'); + } + return Status::OK(); + } + case VariantValueType::kBinary: { + PAIMON_ASSIGN_OR_RAISE(std::string_view binary, + VariantBinaryUtil::GetBinary(value, pos)); + out->push_back('"'); + out->append(arrow::util::base64_encode(binary)); + out->push_back('"'); + return Status::OK(); + } + case VariantValueType::kUuid: { + PAIMON_ASSIGN_OR_RAISE(std::string_view uuid, VariantBinaryUtil::GetUuid(value, pos)); + out->push_back('"'); + out->append(VariantBinaryUtil::UuidToString(uuid)); + out->push_back('"'); + return Status::OK(); + } + } + return VariantBinaryUtil::MalformedVariant(); +} + +} // namespace + +GenericVariant::GenericVariant(std::shared_ptr value, std::shared_ptr metadata, + int32_t pos) + : value_(std::move(value)), metadata_(std::move(metadata)), pos_(pos) {} + +Result> GenericVariant::Create(std::shared_ptr value, + std::shared_ptr metadata) { + if (!value || !metadata) { + return Status::Invalid("variant value and metadata must not be null"); + } + // There is currently only one allowed version. + if (metadata->size() < 1 || (static_cast((*metadata)[0]) & + VariantDefs::kVersionMask) != VariantDefs::kVersion) { + return VariantBinaryUtil::MalformedVariant(); + } + // Don't attempt to use a Variant larger than 128 MiB. We'll never produce one, and it risks + // memory instability. + if (metadata->size() > static_cast(VariantDefs::kSizeLimit) || + value->size() > static_cast(VariantDefs::kSizeLimit)) { + return VariantBinaryUtil::VariantConstructorSizeLimit(); + } + return std::shared_ptr( + new GenericVariant(std::move(value), std::move(metadata), 0)); +} + +Result> GenericVariant::Create( + std::string_view value, std::string_view metadata, const std::shared_ptr& pool) { + // Reject over-limit inputs before allocating and copying them. + if (value.size() > static_cast(VariantDefs::kSizeLimit) || + metadata.size() > static_cast(VariantDefs::kSizeLimit)) { + return VariantBinaryUtil::VariantConstructorSizeLimit(); + } + std::shared_ptr value_bytes = Bytes::AllocateBytes(value.size(), pool.get()); + if (!value.empty()) { + std::memcpy(value_bytes->data(), value.data(), value.size()); + } + std::shared_ptr metadata_bytes = Bytes::AllocateBytes(metadata.size(), pool.get()); + if (!metadata.empty()) { + // An empty (malformed) metadata view may carry a null data pointer; the size check + // keeps memcpy away from it. + std::memcpy(metadata_bytes->data(), metadata.data(), metadata.size()); + } + return Create(std::move(value_bytes), std::move(metadata_bytes)); +} + +Result> GenericVariant::FromJson( + std::string_view json, const std::shared_ptr& pool) { + return VariantBuilder::ParseJson(json, /*allow_duplicate_keys=*/false, pool); +} + +Result GenericVariant::Value() const { + std::string_view raw = RawValue(); + if (pos_ == 0) { + return raw; + } + PAIMON_ASSIGN_OR_RAISE(int32_t size, VariantBinaryUtil::ValueSize(raw, pos_)); + PAIMON_RETURN_NOT_OK( + VariantBinaryUtil::CheckIndex(pos_ + size - 1, static_cast(raw.size()))); + return raw.substr(pos_, size); +} + +std::string_view GenericVariant::RawValue() const { + return {value_->data(), value_->size()}; +} + +std::string_view GenericVariant::Metadata() const { + return {metadata_->data(), metadata_->size()}; +} + +int64_t GenericVariant::SizeInBytes() const { + return static_cast(value_->size()) + static_cast(metadata_->size()); +} + +Result GenericVariant::ToJson(const std::string& zone_id) const { + std::string result; + PAIMON_RETURN_NOT_OK(ToJsonImpl(RawValue(), Metadata(), pos_, zone_id, &result)); + return result; +} + +Result GenericVariant::GetType() const { + return VariantBinaryUtil::GetType(RawValue(), pos_); +} + +Result GenericVariant::GetTypeInfo() const { + return VariantBinaryUtil::GetTypeInfo(RawValue(), pos_); +} + +Result GenericVariant::GetBoolean() const { + return VariantBinaryUtil::GetBoolean(RawValue(), pos_); +} + +Result GenericVariant::GetLong() const { + return VariantBinaryUtil::GetLong(RawValue(), pos_); +} + +Result GenericVariant::GetDouble() const { + return VariantBinaryUtil::GetDouble(RawValue(), pos_); +} + +Result GenericVariant::GetDecimal() const { + return VariantBinaryUtil::GetDecimal(RawValue(), pos_); +} + +Result GenericVariant::GetFloat() const { + return VariantBinaryUtil::GetFloat(RawValue(), pos_); +} + +Result GenericVariant::GetBinary() const { + return VariantBinaryUtil::GetBinary(RawValue(), pos_); +} + +Result GenericVariant::GetString() const { + return VariantBinaryUtil::GetString(RawValue(), pos_); +} + +Result GenericVariant::GetUuid() const { + return VariantBinaryUtil::GetUuid(RawValue(), pos_); +} + +Result GenericVariant::ObjectSize() const { + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ObjectInfo info, + VariantBinaryUtil::GetObjectInfo(RawValue(), pos_)); + return info.num_elements; +} + +std::shared_ptr GenericVariant::SubVariant(int32_t pos) const { + return std::shared_ptr(new GenericVariant(value_, metadata_, pos)); +} + +Result> GenericVariant::GetFieldByKey(std::string_view key) const { + std::string_view raw = RawValue(); + std::string_view metadata = Metadata(); + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ObjectInfo info, + VariantBinaryUtil::GetObjectInfo(raw, pos_)); + // Use linear search for a short list. Switch to binary search when the length reaches + // `kBinarySearchThreshold`. + if (info.num_elements < VariantDefs::kBinarySearchThreshold) { + for (int32_t i = 0; i < info.num_elements; ++i) { + PAIMON_ASSIGN_OR_RAISE( + int32_t id, VariantBinaryUtil::ReadUnsigned(raw, info.id_start + info.id_size * i, + info.id_size)); + PAIMON_ASSIGN_OR_RAISE(std::string_view field_key, + VariantBinaryUtil::GetMetadataKey(metadata, id)); + if (field_key == key) { + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, + VariantBinaryUtil::ReadUnsigned(raw, info.offset_start + info.offset_size * i, + info.offset_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t element_pos, + VariantBinaryUtil::CheckedElementPos(info.data_start, offset, raw.size())); + return SubVariant(element_pos); + } + } + } else { + int32_t low = 0; + int32_t high = info.num_elements - 1; + while (low <= high) { + // Use an unsigned shift to compute the middle of `low` and `high`, which properly + // handles the case where `low + high` overflows int32. + auto mid = static_cast( + (static_cast(low) + static_cast(high)) >> 1); + PAIMON_ASSIGN_OR_RAISE( + int32_t id, VariantBinaryUtil::ReadUnsigned(raw, info.id_start + info.id_size * mid, + info.id_size)); + PAIMON_ASSIGN_OR_RAISE(std::string_view field_key, + VariantBinaryUtil::GetMetadataKey(metadata, id)); + int32_t cmp = field_key.compare(key); + if (cmp < 0) { + low = mid + 1; + } else if (cmp > 0) { + high = mid - 1; + } else { + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, + VariantBinaryUtil::ReadUnsigned(raw, info.offset_start + info.offset_size * mid, + info.offset_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t element_pos, + VariantBinaryUtil::CheckedElementPos(info.data_start, offset, raw.size())); + return SubVariant(element_pos); + } + } + } + return std::shared_ptr(nullptr); +} + +Result> GenericVariant::GetFieldAtIndex( + int32_t index) const { + std::string_view raw = RawValue(); + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ObjectInfo info, + VariantBinaryUtil::GetObjectInfo(raw, pos_)); + if (index < 0 || index >= info.num_elements) { + return std::optional(std::nullopt); + } + PAIMON_ASSIGN_OR_RAISE( + int32_t id, + VariantBinaryUtil::ReadUnsigned(raw, info.id_start + info.id_size * index, info.id_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, VariantBinaryUtil::ReadUnsigned( + raw, info.offset_start + info.offset_size * index, info.offset_size)); + PAIMON_ASSIGN_OR_RAISE(std::string_view key, VariantBinaryUtil::GetMetadataKey(Metadata(), id)); + ObjectField field; + field.key = std::string(key); + PAIMON_ASSIGN_OR_RAISE(int32_t field_pos, VariantBinaryUtil::CheckedElementPos( + info.data_start, offset, raw.size())); + field.value = SubVariant(field_pos); + return std::optional(std::move(field)); +} + +Result GenericVariant::GetDictionaryIdAtIndex(int32_t index) const { + std::string_view raw = RawValue(); + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ObjectInfo info, + VariantBinaryUtil::GetObjectInfo(raw, pos_)); + if (index < 0 || index >= info.num_elements) { + return VariantBinaryUtil::MalformedVariant(); + } + return VariantBinaryUtil::ReadUnsigned(raw, info.id_start + info.id_size * index, info.id_size); +} + +Result GenericVariant::ArraySize() const { + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ArrayInfo info, + VariantBinaryUtil::GetArrayInfo(RawValue(), pos_)); + return info.num_elements; +} + +Result> GenericVariant::GetElementAtIndex(int32_t index) const { + std::string_view raw = RawValue(); + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ArrayInfo info, + VariantBinaryUtil::GetArrayInfo(raw, pos_)); + if (index < 0 || index >= info.num_elements) { + return std::shared_ptr(nullptr); + } + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, VariantBinaryUtil::ReadUnsigned( + raw, info.offset_start + info.offset_size * index, info.offset_size)); + PAIMON_ASSIGN_OR_RAISE(int32_t element_pos, VariantBinaryUtil::CheckedElementPos( + info.data_start, offset, raw.size())); + return SubVariant(element_pos); +} + +bool GenericVariant::operator==(const GenericVariant& other) const { + return pos_ == other.pos_ && RawValue() == other.RawValue() && Metadata() == other.Metadata(); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/generic_variant.h b/src/paimon/common/data/variant/generic_variant.h new file mode 100644 index 000000000..8e2ecc128 --- /dev/null +++ b/src/paimon/common/data/variant/generic_variant.h @@ -0,0 +1,144 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/common/data/variant/variant_binary_util.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace paimon { + +/// A Variant represents a type that contains one of: 1) Primitive: A type and corresponding value +/// (e.g. INT, STRING); 2) Array: An ordered list of Variant values; 3) Object: An unordered +/// collection of string/Variant pairs (i.e. key/value pairs). An object may not contain duplicate +/// keys. +/// +/// A Variant is encoded with 2 binaries: the value and the metadata. The Variant Binary Encoding +/// allows representation of semi-structured data (e.g. JSON) in a form that can be efficiently +/// queried by path. The design is intended to allow efficient access to nested data even in the +/// presence of very wide or deep structures. +class GenericVariant { + public: + /// A field of a variant object. + struct ObjectField { + std::string key; + std::shared_ptr value; + }; + + /// Creates a variant taking ownership of the given buffers. Fails with `MALFORMED_VARIANT` + /// if the metadata version is unsupported, or `VARIANT_CONSTRUCTOR_SIZE_LIMIT` if either + /// buffer exceeds the 128MiB size limit. + static Result> Create(std::shared_ptr value, + std::shared_ptr metadata); + + /// Creates a variant by copying the given buffers into `pool`. + static Result> Create(std::string_view value, + std::string_view metadata, + const std::shared_ptr& pool); + + /// Parses a JSON string as a variant (duplicate object keys are rejected). + static Result> FromJson( + std::string_view json, const std::shared_ptr& pool); + + /// The variant value binary. For a sub-variant (`Pos() != 0`), the view covers only the + /// sub-variant slice of the underlying buffer. + Result Value() const; + + /// The whole underlying value buffer, regardless of `Pos()`. + std::string_view RawValue() const; + + /// The variant metadata binary. + std::string_view Metadata() const; + + /// The variant value doesn't use the whole value binary, but starts from the `Pos()` index + /// and spans a size of `VariantBinaryUtil::ValueSize`. This design avoids frequent copies of + /// the value binary when reading a sub-variant in an array/object element. + int32_t Pos() const { + return pos_; + } + + /// The size of the variant in bytes (value size + metadata size). + int64_t SizeInBytes() const; + + /// Stringifies the variant in JSON format. `zone_id` controls the rendering of TIMESTAMP + /// values; supported forms are "UTC"/"Z"/"GMT" and fixed offsets such as "+08:00". + Result ToJson(const std::string& zone_id = "UTC") const; + + /// The value type of the variant. + Result GetType() const; + + /// The type info bits of the variant value header. + Result GetTypeInfo() const; + + Result GetBoolean() const; + Result GetLong() const; + Result GetDouble() const; + Result GetDecimal() const; + Result GetFloat() const; + Result GetBinary() const; + Result GetString() const; + /// The 16-byte big-endian UUID value. + Result GetUuid() const; + + /// The number of object fields in the variant. It is only legal to call it when `GetType()` + /// is `kObject`. + Result ObjectSize() const; + + /// Finds the field value whose key is equal to `key`. Returns nullptr if the key is not + /// found. It is only legal to call it when `GetType()` is `kObject`. The returned sub-variant + /// shares the underlying buffers with this variant. + Result> GetFieldByKey(std::string_view key) const; + + /// Gets the object field at the `index` slot. Returns nullopt if `index` is out of the bound + /// of `[0, ObjectSize())`. It is only legal to call it when `GetType()` is `kObject`. + Result> GetFieldAtIndex(int32_t index) const; + + /// Gets the metadata dictionary id for the object field at the `index` slot. It is only + /// legal to call it when `GetType()` is `kObject`. + Result GetDictionaryIdAtIndex(int32_t index) const; + + /// The number of array elements in the variant. It is only legal to call it when `GetType()` + /// is `kArray`. + Result ArraySize() const; + + /// Gets the array element at the `index` slot. Returns nullptr if `index` is out of the + /// bound of `[0, ArraySize())`. It is only legal to call it when `GetType()` is `kArray`. + Result> GetElementAtIndex(int32_t index) const; + + bool operator==(const GenericVariant& other) const; + + private: + GenericVariant(std::shared_ptr value, std::shared_ptr metadata, int32_t pos); + + std::shared_ptr SubVariant(int32_t pos) const; + + std::shared_ptr value_; + std::shared_ptr metadata_; + int32_t pos_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/generic_variant_test.cpp b/src/paimon/common/data/variant/generic_variant_test.cpp new file mode 100644 index 000000000..84e542a34 --- /dev/null +++ b/src/paimon/common/data/variant/generic_variant_test.cpp @@ -0,0 +1,348 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/generic_variant.h" + +#include +#include + +#include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class GenericVariantTest : public ::testing::Test { + public: + static std::string ToHex(std::string_view data) { + constexpr char kHexDigits[] = "0123456789abcdef"; + std::string result; + result.reserve(data.size() * 2); + for (char c : data) { + auto byte = static_cast(c); + result.push_back(kHexDigits[byte >> 4]); + result.push_back(kHexDigits[byte & 0xF]); + } + return result; + } + + std::shared_ptr FromJson(const std::string& json) { + auto result = GenericVariant::FromJson(json, pool_); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return result.value(); + } + + // Asserts that parsing `json` produces exactly the value/metadata binaries produced by the + // Java implementation (`GenericVariantBuilder`), and that rendering back to JSON matches the + // Java `toJson` output. + void CheckGolden(const std::string& json, const std::string& expected_value_hex, + const std::string& expected_metadata_hex, + const std::string& expected_to_json) { + auto variant = FromJson(json); + ASSERT_OK_AND_ASSIGN(std::string_view value, variant->Value()); + ASSERT_EQ(ToHex(value), expected_value_hex) << "value bytes mismatch for: " << json; + ASSERT_EQ(ToHex(variant->Metadata()), expected_metadata_hex) + << "metadata bytes mismatch for: " << json; + ASSERT_OK_AND_ASSIGN(std::string to_json, variant->ToJson()); + ASSERT_EQ(to_json, expected_to_json); + } + + protected: + std::shared_ptr pool_ = GetDefaultPool(); +}; + +// Golden binaries generated by the Java implementation (org.apache.paimon.data.variant); these +// pin the cross-implementation byte compatibility of the variant encoding. +TEST_F(GenericVariantTest, GoldenPrimitives) { + CheckGolden("null", "00", "010000", "null"); + CheckGolden("true", "04", "010000", "true"); + CheckGolden("false", "08", "010000", "false"); + CheckGolden("1", "0c01", "010000", "1"); + CheckGolden("-1", "0cff", "010000", "-1"); + CheckGolden("300", "102c01", "010000", "300"); + CheckGolden("100000", "14a0860100", "010000", "100000"); + CheckGolden("12345678901234", "18f22fce733a0b0000", "010000", "12345678901234"); + CheckGolden("1e40", "1ca55cc3f129633d48", "010000", "1.0E40"); + CheckGolden("1.0123456789012345678901234567890123456789", "1c240bf2619132f03f", "010000", + "1.0123456789012346"); + CheckGolden("2.5e-3", "1c7b14ae47e17a643f", "010000", "0.0025"); + CheckGolden("100.99", "200273270000", "010000", "100.99"); + CheckGolden("-0.5", "2001fbffffff", "010000", "-0.5"); + CheckGolden("0.0", "200100000000", "010000", "0"); + CheckGolden("12345678.90123", "2405cb04fb711f010000", "010000", "12345678.90123"); + CheckGolden("1234567890123456789.0123456789", "280a1581396eb1c9be46321be42700000000", "010000", + "1234567890123456789.0123456789"); + // An integer that overflows int64 is parsed as an exact decimal. + CheckGolden("123456789012345678901234567890", "2800d20a3f4eeee073c3f60fe98e01000000", "010000", + "123456789012345678901234567890"); +} + +TEST_F(GenericVariantTest, GoldenStrings) { + CheckGolden("\"Hello, World!\"", "3548656c6c6f2c20576f726c6421", "010000", "\"Hello, World!\""); + CheckGolden("\"\"", "01", "010000", "\"\""); + CheckGolden( + "\"This is a long string that definitely exceeds the sixty-three byte short string limit " + "...!\"", + "405a000000546869732069732061206c6f6e6720737472696e67207468617420646566696e6974656c792065" + "786365656473207468652073697874792d746872656520627974652073686f727420737472696e67206c696d" + "6974202e2e2e21", + "010000", + "\"This is a long string that definitely exceeds the sixty-three byte short string limit " + "...!\""); +} + +TEST_F(GenericVariantTest, GoldenContainers) { + CheckGolden("{}", "020000", "010000", "{}"); + CheckGolden("[]", "030000", "010000", "[]"); + CheckGolden(R"({"a": 1, "b": "hello"})", "020200010002080c011568656c6c6f", "01020001026162", + R"({"a":1,"b":"hello"})"); + CheckGolden(R"([1, "two", 3.5, null, true, {"k":[]}])", + "03060002060c0d0e160c010d74776f20012300000000040201000003030000", "010100016b", + R"([1,"two",3.5,null,true,{"k":[]}])"); +} + +TEST_F(GenericVariantTest, GoldenNested) { + CheckGolden( + "{\"object\":{\"name\":\"Apache Paimon\",\"age\":2,\"address\":{\"street\":\"Main " + "St\",\"city\":\"Hangzhou\"}},\"array\":[1,2,3,4,5],\"string\":\"Hello, " + "World!\",\"long\":12345678901234,\"double\":1." + "0123456789012345678901234567890123456789,\"decimal\":100.99,\"boolean1\":true," + "\"boolean2\":false,\"nullField\":null}", + "0209060b0c0a09080d000731696a635a516b00436c0203030201100e0028354170616368652050" + "61696d6f6e0c02020205040800111d4d61696e2053742148616e677a686f75030500020406080a" + "0c010c020c030c040c053548656c6c6f2c20576f726c642118f22fce733a0b00001c240bf26191" + "32f03f200273270000040800", + "010e00060a0d141a1e23292d333a424a536f626a6563746e616d6561676561646472657373737472656574" + "636974796172726179737472696e676c6f6e67646f75626c65646563696d616c626f6f6c65616e31626f6f" + "6c65616e326e756c6c4669656c64", + "{\"array\":[1,2,3,4,5],\"boolean1\":true,\"boolean2\":false,\"decimal\":100.99," + "\"double\":1.0123456789012346,\"long\":12345678901234,\"nullField\":null,\"object\":{" + "\"address\":{\"city\":\"Hangzhou\",\"street\":\"Main St\"},\"age\":2,\"name\":\"Apache " + "Paimon\"},\"string\":\"Hello, World!\"}"); +} + +TEST_F(GenericVariantTest, GoldenUnicodeEscape) { + CheckGolden(R"({"\u4e2d\u6587": "\u4f60\u597d\n\t\"quoted\""})", + "020100001141e4bda0e5a5bd0a092271756f74656422", "01010006e4b8ade69687", + "{\"中文\":\"你好\\n\\t\\\"quoted\\\"\"}"); +} + +TEST_F(GenericVariantTest, TypedAccessors) { + auto variant = FromJson(R"({"a": 1, "b": "hello"})"); + ASSERT_OK_AND_ASSIGN(VariantValueType type, variant->GetType()); + ASSERT_EQ(type, VariantValueType::kObject); + ASSERT_OK_AND_ASSIGN(int32_t object_size, variant->ObjectSize()); + ASSERT_EQ(object_size, 2); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr a, variant->GetFieldByKey("a")); + ASSERT_NE(a, nullptr); + ASSERT_OK_AND_ASSIGN(int64_t a_value, a->GetLong()); + ASSERT_EQ(a_value, 1); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr b, variant->GetFieldByKey("b")); + ASSERT_NE(b, nullptr); + ASSERT_OK_AND_ASSIGN(std::string_view b_value, b->GetString()); + ASSERT_EQ(b_value, "hello"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr missing, variant->GetFieldByKey("c")); + ASSERT_EQ(missing, nullptr); + + ASSERT_OK_AND_ASSIGN(auto field0, variant->GetFieldAtIndex(0)); + ASSERT_TRUE(field0.has_value()); + ASSERT_EQ(field0->key, "a"); + ASSERT_OK_AND_ASSIGN(auto field_oob, variant->GetFieldAtIndex(2)); + ASSERT_FALSE(field_oob.has_value()); + + auto array = FromJson("[10, 20, 30]"); + ASSERT_OK_AND_ASSIGN(int32_t array_size, array->ArraySize()); + ASSERT_EQ(array_size, 3); + ASSERT_OK_AND_ASSIGN(std::shared_ptr elem, array->GetElementAtIndex(1)); + ASSERT_NE(elem, nullptr); + ASSERT_OK_AND_ASSIGN(int64_t elem_value, elem->GetLong()); + ASSERT_EQ(elem_value, 20); + ASSERT_OK_AND_ASSIGN(std::shared_ptr elem_oob, array->GetElementAtIndex(3)); + ASSERT_EQ(elem_oob, nullptr); +} + +TEST_F(GenericVariantTest, ObjectBinarySearch) { + // More fields than kBinarySearchThreshold exercises the binary-search lookup. + std::string json = "{"; + for (int32_t i = 0; i < 40; ++i) { + if (i != 0) { + json += ","; + } + json += "\"key" + std::to_string(i) + "\":" + std::to_string(i); + } + json += "}"; + auto variant = FromJson(json); + for (int32_t i = 0; i < 40; ++i) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr field, + variant->GetFieldByKey("key" + std::to_string(i))); + ASSERT_NE(field, nullptr); + ASSERT_OK_AND_ASSIGN(int64_t value, field->GetLong()); + ASSERT_EQ(value, i); + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr missing, variant->GetFieldByKey("key40")); + ASSERT_EQ(missing, nullptr); +} + +TEST_F(GenericVariantTest, DuplicateKeys) { + ASSERT_NOK(GenericVariant::FromJson("{\"a\": 1, \"a\": 2}", pool_)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr variant, + VariantBuilder::ParseJson("{\"a\": 1, \"a\": 2}", /*allow_duplicate_keys=*/true, pool_)); + ASSERT_OK_AND_ASSIGN(std::string to_json, variant->ToJson()); + ASSERT_EQ(to_json, "{\"a\":2}"); +} + +TEST_F(GenericVariantTest, CorruptedLayoutRejected) { + // Headers claiming a near-INT32_MAX element count must be rejected by the 64-bit layout + // bound instead of overflowing the 32-bit offset arithmetic. + std::string metadata; + metadata.push_back(static_cast(0x01)); + metadata.push_back(static_cast(0x00)); + { + // Object, large size, 4-byte ids and offsets, num_elements = INT32_MAX. + std::string value; + value.push_back(static_cast((0x1F << 2) | 0x02)); + value.append({static_cast(0xFF), static_cast(0xFF), static_cast(0xFF), + static_cast(0x7F)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, + GenericVariant::Create(value, metadata, pool_)); + ASSERT_NOK(variant->ToJson()); + } + { + // Array, large size, 4-byte offsets, num_elements = INT32_MAX. + std::string value; + value.push_back(static_cast((0x07 << 2) | 0x03)); + value.append({static_cast(0xFF), static_cast(0xFF), static_cast(0xFF), + static_cast(0x7F)}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, + GenericVariant::Create(value, metadata, pool_)); + ASSERT_NOK(variant->ToJson()); + } +} + +TEST_F(GenericVariantTest, OverLimitInputRejectedBeforeAllocation) { + // The size check must run before the buffers are allocated/copied; the fake-length view is + // never dereferenced. + char byte = 0; + std::string_view huge(&byte, static_cast(VariantDefs::kSizeLimit) + 1); + ASSERT_NOK(GenericVariant::Create(huge, std::string_view(&byte, 1), pool_)); + ASSERT_NOK(GenericVariant::Create(std::string_view(&byte, 1), huge, pool_)); +} + +TEST_F(GenericVariantTest, MalformedInput) { + ASSERT_NOK(GenericVariant::FromJson("", pool_)); + ASSERT_NOK(GenericVariant::FromJson("{", pool_)); + ASSERT_NOK(GenericVariant::FromJson("{\"a\":}", pool_)); + ASSERT_NOK(GenericVariant::FromJson("NaN", pool_)); + + // Unsupported metadata version. + std::string bad_metadata = std::string("\x02\x00\x00", 3); + std::string value = std::string("\x00", 1); + ASSERT_NOK(GenericVariant::Create(value, bad_metadata, pool_)); + // Empty metadata. + ASSERT_NOK(GenericVariant::Create(value, std::string(), pool_)); +} + +TEST_F(GenericVariantTest, SizeInBytesAndViews) { + auto variant = FromJson(R"({"a": 1, "b": "hello"})"); + ASSERT_EQ(variant->SizeInBytes(), + static_cast(variant->RawValue().size() + variant->Metadata().size())); + // A sub-variant shares buffers with its parent and reports a positive position. + ASSERT_OK_AND_ASSIGN(std::shared_ptr b, variant->GetFieldByKey("b")); + ASSERT_GT(b->Pos(), 0); + ASSERT_OK_AND_ASSIGN(std::string_view b_slice, b->Value()); + ASSERT_EQ(ToHex(b_slice), "1568656c6c6f"); +} + +TEST_F(GenericVariantTest, AppendVariantRebuild) { + // Rebuilding a sub-variant through a fresh builder produces a self-contained variant. + auto variant = FromJson(R"({"outer": {"x": [1, 2], "y": "z"}})"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr outer, variant->GetFieldByKey("outer")); + VariantBuilder builder(/*allow_duplicate_keys=*/false); + ASSERT_OK(builder.AppendVariant(*outer)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rebuilt, builder.Build(pool_)); + ASSERT_EQ(rebuilt->Pos(), 0); + ASSERT_OK_AND_ASSIGN(std::string to_json, rebuilt->ToJson()); + ASSERT_EQ(to_json, "{\"x\":[1,2],\"y\":\"z\"}"); +} + +TEST_F(GenericVariantTest, TimestampAndSpecialTypesToJson) { + // JSON can't produce date/timestamp/binary/uuid variants; build them directly. + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendDate(19737)); // 2024-01-15 + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "\"2024-01-15\""); + } + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendTimestamp(1705312496123456LL)); + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "\"2024-01-15 09:54:56.123456+00:00\""); + ASSERT_OK_AND_ASSIGN(std::string json_shanghai, v->ToJson("Asia/Shanghai")); + ASSERT_EQ(json_shanghai, "\"2024-01-15 17:54:56.123456+08:00\""); + ASSERT_OK_AND_ASSIGN(std::string json_offset, v->ToJson("+08:00")); + ASSERT_EQ(json_offset, "\"2024-01-15 17:54:56.123456+08:00\""); + ASSERT_NOK(v->ToJson("Not/AZone")); + } + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendTimestampNtz(1705312496000000LL)); + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "\"2024-01-15 09:54:56\""); + } + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendBinary(std::string_view("\x01\x02\x03", 3))); + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "\"AQID\""); + } + { + VariantBuilder builder(false); + std::string uuid_bytes = + std::string("\x12\x3e\x45\x67\xe8\x9b\x12\xd3\xa4\x56\x42\x66\x14\x17\x40\x00", 16); + ASSERT_OK(builder.AppendUuid(uuid_bytes)); + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "\"123e4567-e89b-12d3-a456-426614174000\""); + } + { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendFloat(1.5f)); + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "1.5"); + } +} + +TEST_F(GenericVariantTest, NonFiniteDoubleToJson) { + VariantBuilder builder(false); + ASSERT_OK(builder.AppendDouble(std::numeric_limits::infinity())); + ASSERT_OK_AND_ASSIGN(auto v, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, v->ToJson()); + ASSERT_EQ(json, "\"Infinity\""); +} + +} // namespace paimon::test diff --git a/src/paimon/common/data/variant/infer_variant_shredding_schema.cpp b/src/paimon/common/data/variant/infer_variant_shredding_schema.cpp new file mode 100644 index 000000000..44f04470a --- /dev/null +++ b/src/paimon/common/data/variant/infer_variant_shredding_schema.cpp @@ -0,0 +1,361 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/infer_variant_shredding_schema.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "paimon/common/data/variant/variant_binary_util.h" +#include "paimon/common/data/variant/variant_defs.h" + +namespace paimon { + +namespace { + +constexpr int32_t kMaxRowFieldSize = 1000; + +// The inference type lattice. A scalar node holds an arrow type (`arrow::null()` is the untyped +// VARIANT sentinel); object nodes track per-field occurrence counts so that rare fields can be +// dropped in the final schema. +struct SimpleSchema { + struct Field { + std::string name; + std::shared_ptr schema; + int64_t count; + }; + + bool is_object = false; + bool is_array = false; + std::vector fields; + std::shared_ptr element; + std::shared_ptr scalar; + + static std::shared_ptr Variant() { + auto schema = std::make_shared(); + schema->scalar = arrow::null(); + return schema; + } + + static std::shared_ptr Scalar(std::shared_ptr type) { + auto schema = std::make_shared(); + schema->scalar = std::move(type); + return schema; + } +}; + +// Tracks the total number of shredded fields remaining, shared across the whole schema. +struct MaxFields { + int32_t remaining; +}; + +std::shared_ptr MergeSchema(const std::shared_ptr& s1, + const std::shared_ptr& s2); + +// Merges two decimals with possibly different scales. +std::shared_ptr MergeDecimal(const arrow::Decimal128Type& d1, + const arrow::Decimal128Type& d2) { + int32_t scale = std::max(d1.scale(), d2.scale()); + int32_t range = std::max(d1.precision() - d1.scale(), d2.precision() - d2.scale()); + if (range + scale > VariantDefs::kMaxDecimal16Precision) { + // Decimal cannot support precision > 38. + return SimpleSchema::Variant(); + } + return SimpleSchema::Scalar(arrow::decimal128(range + scale, scale)); +} + +std::shared_ptr MergeDecimalWithLong(const arrow::Decimal128Type& d) { + if (d.scale() == 0 && d.precision() <= 18) { + return SimpleSchema::Scalar(arrow::int64()); + } + // A long can always fit in a decimal(19, 0). + auto long_decimal = std::static_pointer_cast(arrow::decimal128(19, 0)); + return MergeDecimal(d, *long_decimal); +} + +std::shared_ptr MergeObjects(const std::shared_ptr& s1, + const std::shared_ptr& s2) { + auto result = std::make_shared(); + result->is_object = true; + size_t f1_idx = 0; + size_t f2_idx = 0; + while (f1_idx < s1->fields.size() && f2_idx < s2->fields.size() && + result->fields.size() < kMaxRowFieldSize) { + const auto& field1 = s1->fields[f1_idx]; + const auto& field2 = s2->fields[f2_idx]; + int32_t comp = field1.name.compare(field2.name); + if (comp == 0) { + result->fields.push_back(SimpleSchema::Field{field1.name, + MergeSchema(field1.schema, field2.schema), + field1.count + field2.count}); + ++f1_idx; + ++f2_idx; + } else if (comp < 0) { + result->fields.push_back(field1); + ++f1_idx; + } else { + result->fields.push_back(field2); + ++f2_idx; + } + } + while (f1_idx < s1->fields.size() && result->fields.size() < kMaxRowFieldSize) { + result->fields.push_back(s1->fields[f1_idx++]); + } + while (f2_idx < s2->fields.size() && result->fields.size() < kMaxRowFieldSize) { + result->fields.push_back(s2->fields[f2_idx++]); + } + return result; +} + +std::shared_ptr MergeSchema(const std::shared_ptr& s1, + const std::shared_ptr& s2) { + // Allow null (missing) to merge into any typed schema. + if (s1 == nullptr) { + return s2; + } + if (s2 == nullptr) { + return s1; + } + if (s1->is_object && s2->is_object) { + return MergeObjects(s1, s2); + } + if (s1->is_array && s2->is_array) { + auto result = std::make_shared(); + result->is_array = true; + result->element = MergeSchema(s1->element, s2->element); + if (result->element == nullptr) { + result->element = SimpleSchema::Variant(); + } + return result; + } + if (s1->scalar != nullptr && s2->scalar != nullptr) { + bool s1_decimal = s1->scalar->id() == arrow::Type::DECIMAL128; + bool s2_decimal = s2->scalar->id() == arrow::Type::DECIMAL128; + bool s1_long = s1->scalar->id() == arrow::Type::INT64; + bool s2_long = s2->scalar->id() == arrow::Type::INT64; + if (s1_decimal && s2_decimal) { + return MergeDecimal(static_cast(*s1->scalar), + static_cast(*s2->scalar)); + } + if (s1_decimal && s2_long) { + return MergeDecimalWithLong(static_cast(*s1->scalar)); + } + if (s1_long && s2_decimal) { + return MergeDecimalWithLong(static_cast(*s2->scalar)); + } + if (s1->scalar->Equals(*s2->scalar)) { + return s1; + } + } + return SimpleSchema::Variant(); +} + +// Returns an appropriate schema for shredding a variant value. Unlike a generic schema-of +// expression, the merged types stay consistent with what shredding allows (e.g. an integer and a +// double merge to VARIANT, not double). +Result> SchemaOf(const GenericVariant& variant, int32_t max_depth) { + PAIMON_ASSIGN_OR_RAISE(VariantValueType type, variant.GetType()); + switch (type) { + case VariantValueType::kObject: { + if (max_depth <= 0) { + return SimpleSchema::Variant(); + } + PAIMON_ASSIGN_OR_RAISE(int32_t size, variant.ObjectSize()); + auto result = std::make_shared(); + result->is_object = true; + result->fields.reserve(size); + for (int32_t i = 0; i < size; ++i) { + PAIMON_ASSIGN_OR_RAISE(auto field, variant.GetFieldAtIndex(i)); + if (!field.has_value()) { + return VariantBinaryUtil::MalformedVariant(); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr field_schema, + SchemaOf(*field->value, max_depth - 1)); + if (field_schema == nullptr) { + field_schema = SimpleSchema::Variant(); + } + result->fields.push_back(SimpleSchema::Field{field->key, field_schema, 1}); + } + // According to the variant spec, object fields must be sorted alphabetically. + for (size_t i = 1; i < result->fields.size(); ++i) { + if (result->fields[i - 1].name.compare(result->fields[i].name) >= 0) { + return Status::Invalid("Variant object fields must be sorted alphabetically"); + } + } + return result; + } + case VariantValueType::kArray: { + if (max_depth <= 0) { + return SimpleSchema::Variant(); + } + PAIMON_ASSIGN_OR_RAISE(int32_t size, variant.ArraySize()); + std::shared_ptr element_type; + for (int32_t i = 0; i < size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element, + variant.GetElementAtIndex(i)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element_schema, + SchemaOf(*element, max_depth - 1)); + element_type = MergeSchema(element_type, element_schema); + } + auto result = std::make_shared(); + result->is_array = true; + result->element = element_type == nullptr ? SimpleSchema::Variant() : element_type; + return result; + } + case VariantValueType::kNull: + return std::shared_ptr(nullptr); + case VariantValueType::kBoolean: + return SimpleSchema::Scalar(arrow::boolean()); + case VariantValueType::kLong: { + // Compute the smallest decimal that can contain this value. + PAIMON_ASSIGN_OR_RAISE(int64_t value, variant.GetLong()); + VariantDecimal decimal{value, 0}; + int32_t precision = decimal.Precision(); + if (precision <= 18) { + return SimpleSchema::Scalar(arrow::decimal128(precision, 0)); + } + return SimpleSchema::Scalar(arrow::int64()); + } + case VariantValueType::kString: + return SimpleSchema::Scalar(arrow::utf8()); + case VariantValueType::kDouble: + return SimpleSchema::Scalar(arrow::float64()); + case VariantValueType::kDecimal: { + PAIMON_ASSIGN_OR_RAISE(VariantDecimal decimal, variant.GetDecimal()); + if (decimal.scale < 0) { + // GetDecimal strips trailing zeros and can return a negative scale (100.00 -> + // 1E+2), which neither the shredded parquet type nor reassembly's AppendDecimal + // can represent; scale the value back up. Bounded by construction: the encoded + // decimal had at most 38 digits before the point. + for (; decimal.scale < 0; ++decimal.scale) { + decimal.unscaled *= 10; + } + } + int32_t precision = decimal.Precision(); + int32_t scale = decimal.scale; + // Ensure precision is at least scale (and at least 1) to be valid. + if (precision < scale) { + precision = scale; + } + if (precision == 0) { + precision = 1; + } + return SimpleSchema::Scalar(arrow::decimal128(precision, scale)); + } + case VariantValueType::kDate: + case VariantValueType::kTimestamp: + case VariantValueType::kTimestampNtz: + // The shredding schema builder rejects temporal leaf types (as the Java one does), + // so inferring them would abort the whole write; keep such values in the untyped + // column instead. + return SimpleSchema::Scalar(arrow::null()); + case VariantValueType::kFloat: + return SimpleSchema::Scalar(arrow::float32()); + case VariantValueType::kBinary: + return SimpleSchema::Scalar(arrow::binary()); + default: + return SimpleSchema::Variant(); + } +} + +// Finalizes the inferred schema: 1) widen integer types to int64, 2) replace empty objects with +// VARIANT, 3) limit the total number of shredded fields in the schema. +std::shared_ptr FinalizeSimpleSchema(const std::shared_ptr& schema, + int64_t min_cardinality, + MaxFields* max_fields) { + // Every field uses a value column. + --max_fields->remaining; + if (max_fields->remaining <= 0) { + return arrow::null(); + } + if (schema == nullptr || + (schema->scalar != nullptr && schema->scalar->id() == arrow::Type::NA)) { + return arrow::null(); + } + if (schema->is_object) { + arrow::FieldVector new_fields; + for (const auto& field : schema->fields) { + if (field.count >= min_cardinality && max_fields->remaining > 0) { + auto new_type = FinalizeSimpleSchema(field.schema, min_cardinality, max_fields); + new_fields.push_back(arrow::field(field.name, new_type)); + } + } + if (!new_fields.empty()) { + return arrow::struct_(new_fields); + } + return arrow::null(); + } + if (schema->is_array) { + auto new_element = FinalizeSimpleSchema(schema->element, min_cardinality, max_fields); + return arrow::list(new_element); + } + switch (schema->scalar->id()) { + case arrow::Type::INT8: + case arrow::Type::INT16: + case arrow::Type::INT32: + case arrow::Type::INT64: + --max_fields->remaining; + return arrow::int64(); + case arrow::Type::DECIMAL128: { + const auto& decimal_type = static_cast(*schema->scalar); + --max_fields->remaining; + if (decimal_type.precision() <= 18 && decimal_type.scale() == 0) { + return arrow::int64(); + } + if (decimal_type.precision() <= 18) { + return arrow::decimal128(18, decimal_type.scale()); + } + return arrow::decimal128(VariantDefs::kMaxDecimal16Precision, decimal_type.scale()); + } + default: + // All other scalar types use typed_value. + --max_fields->remaining; + return schema->scalar; + } +} + +} // namespace + +Result> InferVariantShreddingSchema::InferColumnShreddingType( + const std::vector>& samples) const { + int64_t num_non_null_values = 0; + std::shared_ptr simple_schema; + for (const auto& sample : samples) { + if (sample == nullptr) { + continue; + } + ++num_non_null_values; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr row_schema, + SchemaOf(*sample, max_schema_depth_)); + simple_schema = MergeSchema(simple_schema, row_schema); + } + // Don't infer a schema for fields that appear in less than min_field_cardinality_ratio of + // the rows. + auto min_cardinality = static_cast( + std::ceil(static_cast(num_non_null_values) * min_field_cardinality_ratio_)); + MaxFields max_fields{max_schema_width_}; + std::shared_ptr finalized = + FinalizeSimpleSchema(simple_schema, min_cardinality, &max_fields); + if (finalized->id() == arrow::Type::NA) { + // The whole column stays unshredded. + return std::shared_ptr(nullptr); + } + return finalized; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/infer_variant_shredding_schema.h b/src/paimon/common/data/variant/infer_variant_shredding_schema.h new file mode 100644 index 000000000..860cd59d0 --- /dev/null +++ b/src/paimon/common/data/variant/infer_variant_shredding_schema.h @@ -0,0 +1,56 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/generic_variant.h" +#include "paimon/result.h" + +namespace arrow { +class DataType; +} // namespace arrow + +namespace paimon { + +/// Infers a shredding type for a variant column from sampled values (mirroring the Java +/// `InferVariantShreddingSchema`). Rare fields (below the cardinality ratio) stay in the +/// un-shredded variant binary, integer types widen to int64, and the total number of shredded +/// fields is limited. +class InferVariantShreddingSchema { + public: + InferVariantShreddingSchema(int32_t max_schema_width, int32_t max_schema_depth, + double min_field_cardinality_ratio) + : max_schema_width_(max_schema_width), + max_schema_depth_(max_schema_depth), + min_field_cardinality_ratio_(min_field_cardinality_ratio) {} + + /// Infers the shredding type of one variant column from its sampled non-null values, e.g. + /// `struct{a: int64, b: string}`. `arrow::null()` leaves denote untyped variant sub-values. + /// Returns nullptr when no useful shredding schema was found (the column should stay + /// unshredded). + Result> InferColumnShreddingType( + const std::vector>& samples) const; + + private: + int32_t max_schema_width_; + int32_t max_schema_depth_; + double min_field_cardinality_ratio_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/infer_variant_shredding_schema_test.cpp b/src/paimon/common/data/variant/infer_variant_shredding_schema_test.cpp new file mode 100644 index 000000000..22a23d2f9 --- /dev/null +++ b/src/paimon/common/data/variant/infer_variant_shredding_schema_test.cpp @@ -0,0 +1,195 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/infer_variant_shredding_schema.h" + +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_shredding_utils.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class InferVariantShreddingSchemaTest : public ::testing::Test { + public: + std::vector> Samples(const std::vector& jsons) { + std::vector> samples; + for (const char* json : jsons) { + if (json == nullptr) { + samples.push_back(nullptr); + continue; + } + auto variant = GenericVariant::FromJson(json, pool_); + EXPECT_TRUE(variant.ok()) << variant.status().ToString(); + samples.push_back(variant.value()); + } + return samples; + } + + protected: + std::shared_ptr pool_ = GetDefaultPool(); + InferVariantShreddingSchema infer_{/*max_schema_width=*/300, /*max_schema_depth=*/50, + /*min_field_cardinality_ratio=*/0.1}; +}; + +TEST_F(InferVariantShreddingSchemaTest, InferObjectSchema) { + auto samples = Samples({ + R"({"age": 35, "city": "Hangzhou"})", + R"({"age": 20, "city": "Beijing", "tags": [1, 2]})", + nullptr, + R"({"age": 120000000000, "city": "Shanghai"})", + }); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + infer_.InferColumnShreddingType(samples)); + ASSERT_NE(inferred, nullptr); + // Integers widen to int64, strings stay, arrays of small ints infer as list. + auto expected = + arrow::struct_({arrow::field("age", arrow::int64()), arrow::field("city", arrow::utf8()), + arrow::field("tags", arrow::list(arrow::int64()))}); + ASSERT_TRUE(inferred->Equals(*expected)) << inferred->ToString(); + ASSERT_OK(VariantShreddingUtils::VariantShreddingSchema(inferred)); +} + +TEST_F(InferVariantShreddingSchemaTest, MixedTypesFallToVariant) { + auto samples = Samples({ + R"({"x": 1, "y": 1.5e0})", + R"({"x": "string now", "y": 2.5e0})", + }); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + infer_.InferColumnShreddingType(samples)); + ASSERT_NE(inferred, nullptr); + // x saw both int and string: untyped variant leaf; y stays double (exponent notation + // parses as double, plain decimals parse as DECIMAL). + auto expected = + arrow::struct_({arrow::field("x", arrow::null()), arrow::field("y", arrow::float64())}); + ASSERT_TRUE(inferred->Equals(*expected)) << inferred->ToString(); + ASSERT_OK(VariantShreddingUtils::VariantShreddingSchema(inferred)); +} + +TEST_F(InferVariantShreddingSchemaTest, RareFieldsDropped) { + std::vector jsons; + for (int i = 0; i < 19; ++i) { + jsons.push_back("{\"common\": 1}"); + } + jsons.push_back(R"({"common": 2, "rare": true})"); + auto samples = Samples(jsons); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + infer_.InferColumnShreddingType(samples)); + ASSERT_NE(inferred, nullptr); + // "rare" appears in 1/20 rows (< 0.1 ratio): dropped from the typed schema. + auto expected = arrow::struct_({arrow::field("common", arrow::int64())}); + ASSERT_TRUE(inferred->Equals(*expected)) << inferred->ToString(); +} + +TEST_F(InferVariantShreddingSchemaTest, DecimalMerging) { + auto samples = Samples({ + "{\"d\": 100.99}", + "{\"d\": 1.5}", + "{\"d\": 42}", + }); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + infer_.InferColumnShreddingType(samples)); + ASSERT_NE(inferred, nullptr); + // Decimals merge to a widened decimal (scale 2, enough integer digits), capped at 18 digits. + auto expected = arrow::struct_({arrow::field("d", arrow::decimal128(18, 2))}); + ASSERT_TRUE(inferred->Equals(*expected)) << inferred->ToString(); +} + +TEST_F(InferVariantShreddingSchemaTest, NoUsefulSchema) { + auto scalar_samples = Samples({"1", "2"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr scalar_inferred, + infer_.InferColumnShreddingType(scalar_samples)); + ASSERT_NE(scalar_inferred, nullptr); + ASSERT_TRUE(scalar_inferred->Equals(*arrow::int64())) << scalar_inferred->ToString(); + + // Conflicting top-level types stay unshredded. + auto mixed_samples = Samples({"1", "\"a string\""}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr mixed_inferred, + infer_.InferColumnShreddingType(mixed_samples)); + ASSERT_EQ(mixed_inferred, nullptr); + + // All-null columns stay unshredded. + auto null_samples = Samples({nullptr, nullptr}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr null_inferred, + infer_.InferColumnShreddingType(null_samples)); + ASSERT_EQ(null_inferred, nullptr); +} + +TEST_F(InferVariantShreddingSchemaTest, MaxSchemaWidthLimit) { + InferVariantShreddingSchema narrow_infer{/*max_schema_width=*/3, /*max_schema_depth=*/50, + /*min_field_cardinality_ratio=*/0.1}; + auto samples = Samples({R"({"a": 1, "b": 2, "c": 3, "d": 4, "e": 5})"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + narrow_infer.InferColumnShreddingType(samples)); + if (inferred != nullptr) { + ASSERT_EQ(inferred->id(), arrow::Type::STRUCT); + ASSERT_LT(inferred->num_fields(), 5); + } +} + +TEST_F(InferVariantShreddingSchemaTest, MaxSchemaDepthLimit) { + InferVariantShreddingSchema shallow_infer{/*max_schema_width=*/300, /*max_schema_depth=*/1, + /*min_field_cardinality_ratio=*/0.1}; + auto samples = Samples({R"({"outer": {"inner": 1}})"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + shallow_infer.InferColumnShreddingType(samples)); + ASSERT_NE(inferred, nullptr); + // Depth 1: the nested object stays an untyped variant leaf. + auto expected = arrow::struct_({arrow::field("outer", arrow::null())}); + ASSERT_TRUE(inferred->Equals(*expected)) << inferred->ToString(); +} + +TEST_F(InferVariantShreddingSchemaTest, TrailingZeroDecimalNormalized) { + // GetDecimal strips 100.00 to 1E+2 (scale -2); the inferred type must carry a non-negative + // scale or reassembling the shredded file would be rejected. After normalization the value + // is an integral decimal, which finalization widens to int64. + auto samples = Samples({"100.00"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr inferred, + infer_.InferColumnShreddingType(samples)); + ASSERT_NE(inferred, nullptr); + ASSERT_TRUE(inferred->Equals(*arrow::int64())) << inferred->ToString(); + + auto mixed = Samples({"100.00", "1.5"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr mixed_inferred, + infer_.InferColumnShreddingType(mixed)); + ASSERT_NE(mixed_inferred, nullptr); + ASSERT_TRUE(mixed_inferred->Equals(*arrow::decimal128(18, 1))) << mixed_inferred->ToString(); +} + +TEST_F(InferVariantShreddingSchemaTest, TemporalValuesStayUnshredded) { + // The shredding schema builder rejects temporal leaf types, so inferring them would abort + // the write; date/timestamp samples must leave the column unshredded. + VariantBuilder date_builder(/*allow_duplicate_keys=*/false); + ASSERT_OK(date_builder.AppendDate(19000)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr date_variant, date_builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr date_inferred, + infer_.InferColumnShreddingType({date_variant})); + ASSERT_EQ(date_inferred, nullptr); + + VariantBuilder ts_builder(/*allow_duplicate_keys=*/false); + ASSERT_OK(ts_builder.AppendTimestamp(1700000000000000)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr ts_variant, ts_builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr ts_inferred, + infer_.InferColumnShreddingType({ts_variant})); + ASSERT_EQ(ts_inferred, nullptr); +} + +} // namespace paimon::test diff --git a/src/paimon/common/data/variant/variant.cpp b/src/paimon/common/data/variant/variant.cpp new file mode 100644 index 000000000..b3b43d42a --- /dev/null +++ b/src/paimon/common/data/variant/variant.cpp @@ -0,0 +1,187 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/data/variant.h" + +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_access_utils.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/data/variant/variant_get.h" +#include "paimon/common/data/variant/variant_path_segment.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" + +namespace paimon { + +class Variant::Impl { + public: + Impl(std::shared_ptr variant, std::shared_ptr pool) + : variant_(std::move(variant)), pool_(std::move(pool)), arrow_pool_(GetArrowPool(pool_)) {} + + const std::shared_ptr& GetVariant() const { + return variant_; + } + + const std::shared_ptr& GetPool() const { + return pool_; + } + + arrow::MemoryPool* GetArrowMemoryPool() const { + return arrow_pool_.get(); + } + + private: + std::shared_ptr variant_; + std::shared_ptr pool_; + std::unique_ptr arrow_pool_; +}; + +Variant::Variant(std::unique_ptr&& impl) : impl_(std::move(impl)) {} +Variant::~Variant() = default; + +Result> Variant::FromJson(const std::string& json, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + GenericVariant::FromJson(json, pool)); + auto impl = std::make_unique(std::move(variant), pool); + return std::unique_ptr(new Variant(std::move(impl))); +} + +Result> Variant::Create(const char* value, uint64_t value_length, + const char* metadata, uint64_t metadata_length, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr variant, + GenericVariant::Create(std::string_view(value, value_length), + std::string_view(metadata, metadata_length), pool)); + auto impl = std::make_unique(std::move(variant), pool); + return std::unique_ptr(new Variant(std::move(impl))); +} + +std::string_view Variant::Value() const { + return impl_->GetVariant()->RawValue(); +} + +std::string_view Variant::Metadata() const { + return impl_->GetVariant()->Metadata(); +} + +int64_t Variant::SizeInBytes() const { + return impl_->GetVariant()->SizeInBytes(); +} + +Result Variant::ToJson(const std::string& zone_id) const { + return impl_->GetVariant()->ToJson(zone_id); +} + +Result> Variant::VariantGet(const std::string& path, + struct ArrowSchema* target_type, + const VariantCastArgs& cast_args) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr target_field, + arrow::ImportField(target_type)); + return VariantGetExecutor::Get(impl_->GetVariant(), path, target_field->type(), cast_args); +} + +Result> Variant::VariantGetArrow( + const std::string& path, struct ArrowSchema* target_field, + const VariantCastArgs& cast_args) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr field, + arrow::ImportField(target_field)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr array, + VariantGetExecutor::GetAsArrow(impl_->GetVariant(), path, field, cast_args, + impl_->GetPool(), impl_->GetArrowMemoryPool())); + auto result = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, result.get())); + return result; +} + +Result> Variant::VariantGetJson(const std::string& path, + const std::string& zone_id) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr extracted, + VariantGetExecutor::ExtractByPath(impl_->GetVariant(), path)); + if (extracted == nullptr) { + return std::optional(std::nullopt); + } + PAIMON_ASSIGN_OR_RAISE(std::string json, extracted->ToJson(zone_id)); + return std::optional(std::move(json)); +} + +Result> Variant::ArrowField( + const std::string& field_name, bool nullable, + std::unordered_map metadata) { + auto variant_field = VariantTypeUtils::ToArrowField(field_name, nullable, std::move(metadata)); + auto field = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportField(*variant_field, field.get())); + return field; +} + +class VariantAccessBuilder::Impl { + public: + arrow::FieldVector fields; +}; + +VariantAccessBuilder::VariantAccessBuilder() : impl_(std::make_unique()) {} +VariantAccessBuilder::~VariantAccessBuilder() = default; + +Status VariantAccessBuilder::AddField(struct ArrowSchema* target_type, const std::string& path, + bool fail_on_error, const std::string& zone_id) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr target, + arrow::ImportField(target_type)); + // Validate the path eagerly so mistakes fail at build time, not at read time. + PAIMON_RETURN_NOT_OK(VariantPathSegment::Parse(path).status()); + // Keep the target field's own metadata (e.g. the variant extension marker of a + // `Variant::ArrowField` target, which drives the deep re-encode cast) and add the access + // description to it. + std::vector keys = {DataField::DESCRIPTION}; + std::vector values = { + VariantAccessUtils::BuildVariantMetadata(path, fail_on_error, zone_id)}; + if (target->metadata() != nullptr) { + for (int64_t i = 0; i < target->metadata()->size(); ++i) { + if (target->metadata()->key(i) == DataField::DESCRIPTION) { + continue; + } + keys.push_back(target->metadata()->key(i)); + values.push_back(target->metadata()->value(i)); + } + } + impl_->fields.push_back(arrow::field(std::to_string(impl_->fields.size()), target->type(), + /*nullable=*/true, + arrow::KeyValueMetadata::Make(keys, values))); + return Status::OK(); +} + +Result> VariantAccessBuilder::Build( + const std::string& field_name) const { + if (impl_->fields.empty()) { + return Status::Invalid("a variant-access projection needs at least one field"); + } + auto access_field = + arrow::field(field_name, arrow::struct_(impl_->fields), /*nullable=*/true, + arrow::KeyValueMetadata::Make({VariantDefs::kExtensionTypeKey}, + {VariantDefs::kExtensionTypeValue})); + auto field = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportField(*access_field, field.get())); + return field; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_access_utils.cpp b/src/paimon/common/data/variant/variant_access_utils.cpp new file mode 100644 index 000000000..e8decfc36 --- /dev/null +++ b/src/paimon/common/data/variant/variant_access_utils.cpp @@ -0,0 +1,179 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_access_utils.h" + +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/types/data_field.h" + +namespace paimon { + +namespace { + +Result GetDescription(const std::shared_ptr& field) { + if (field->metadata() == nullptr) { + return std::string(); + } + auto result = field->metadata()->Get(DataField::DESCRIPTION); + if (!result.ok()) { + return std::string(); + } + return result.ValueOrDie(); +} + +// Parses the description body from the right: the last two delimited tokens are `failOnError` +// and `timeZoneId` (neither contains the delimiter), and everything before them is the path, +// which may itself contain the delimiter inside object keys (e.g. `$['a;b']`). +std::vector SplitDescription(const std::string& description) { + std::string body = description.substr(sizeof(VariantAccessUtils::kMetadataKey) - 1); + size_t tz_sep = body.rfind(VariantAccessUtils::kDelimiter); + if (tz_sep == std::string::npos) { + return {body}; + } + size_t fail_sep = + tz_sep == 0 ? std::string::npos : body.rfind(VariantAccessUtils::kDelimiter, tz_sep - 1); + if (fail_sep == std::string::npos) { + return {body.substr(0, tz_sep), body.substr(tz_sep + 1)}; + } + return {body.substr(0, fail_sep), body.substr(fail_sep + 1, tz_sep - fail_sep - 1), + body.substr(tz_sep + 1)}; +} + +bool HasAccessDescription(const std::shared_ptr& field) { + auto description = GetDescription(field); + return description.ok() && description.value().rfind(VariantAccessUtils::kMetadataKey, 0) == 0; +} + +} // namespace + +constexpr char VariantAccessUtils::kMetadataKey[]; +constexpr char VariantAccessUtils::kDelimiter; + +std::string VariantAccessUtils::BuildVariantMetadata(const std::string& path, bool fail_on_error, + const std::string& zone_id) { + return fmt::format("{}{}{}{}{}{}", kMetadataKey, path, kDelimiter, + fail_on_error ? "true" : "false", kDelimiter, zone_id); +} + +bool VariantAccessUtils::IsVariantAccessType(const std::shared_ptr& type) { + if (type == nullptr || type->id() != arrow::Type::STRUCT || type->num_fields() == 0) { + return false; + } + for (const auto& child : type->fields()) { + if (!HasAccessDescription(child)) { + return false; + } + } + return true; +} + +Result> VariantAccessUtils::ParseAccessSpecs( + const std::shared_ptr& access_field) { + if (!IsVariantAccessType(access_field->type())) { + return Status::Invalid( + fmt::format("field '{}' is not a variant-access projection", access_field->name())); + } + std::vector specs; + specs.reserve(access_field->type()->num_fields()); + for (const auto& child : access_field->type()->fields()) { + PAIMON_ASSIGN_OR_RAISE(std::string description, GetDescription(child)); + std::vector parts = SplitDescription(description); + if (parts.size() != 3) { + return Status::Invalid( + fmt::format("malformed variant access description '{}' on field '{}'", description, + child->name())); + } + VariantAccessSpec spec; + spec.path = parts[0]; + PAIMON_ASSIGN_OR_RAISE(spec.segments, VariantPathSegment::Parse(spec.path)); + spec.cast_args.fail_on_error = parts[1] == "true"; + spec.cast_args.zone_id = parts[2]; + spec.target_field = child; + specs.push_back(std::move(spec)); + } + return specs; +} + +Result> VariantAccessUtils::ClipShreddedFileField( + const std::vector& specs, const std::shared_ptr& file_field) { + if (file_field->type()->id() != arrow::Type::STRUCT) { + return Status::Invalid( + fmt::format("variant file field '{}' is not a struct", file_field->name())); + } + const auto& file_struct = static_cast(*file_field->type()); + std::shared_ptr typed_value = + file_struct.GetFieldByName(VariantDefs::kTypedValueFieldName); + if (typed_value == nullptr) { + // The file stores the column unshredded; there is nothing to prune. + return file_field; + } + + bool can_clip = true; + std::set fields_to_read; + for (const auto& spec : specs) { + if (spec.segments.empty()) { + // A root path needs the whole variant. + can_clip = false; + break; + } + if (spec.segments[0].kind == VariantPathSegment::Kind::kObjectExtraction) { + // Only top-level object keys are pruned; nested paths still narrow to their + // top-level key. + fields_to_read.insert(spec.segments[0].key); + } else { + can_clip = false; + break; + } + } + if (!can_clip) { + return file_field; + } + + std::shared_ptr metadata_field = + file_struct.GetFieldByName(VariantDefs::kMetadataFieldName); + std::shared_ptr value_field = + file_struct.GetFieldByName(VariantDefs::kValueFieldName); + if (metadata_field == nullptr) { + return Status::Invalid( + fmt::format("shredded variant field '{}' misses metadata", file_field->name())); + } + + arrow::FieldVector typed_fields; + if (typed_value->type()->id() == arrow::Type::STRUCT) { + for (const auto& typed_child : typed_value->type()->fields()) { + if (fields_to_read.erase(typed_child->name()) > 0) { + typed_fields.push_back(typed_child); + } + } + } + + arrow::FieldVector clipped_fields = {metadata_field}; + if (!fields_to_read.empty() && value_field != nullptr) { + // Some requested key is not shredded; keep `value` for the binary fallback. + clipped_fields.push_back(value_field); + } + if (!typed_fields.empty()) { + clipped_fields.push_back(typed_value->WithType(arrow::struct_(typed_fields))); + } + return file_field->WithType(arrow::struct_(clipped_fields)); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_access_utils.h b/src/paimon/common/data/variant/variant_access_utils.h new file mode 100644 index 000000000..6032d5a7a --- /dev/null +++ b/src/paimon/common/data/variant/variant_access_utils.h @@ -0,0 +1,77 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_path_segment.h" +#include "paimon/data/variant.h" +#include "paimon/result.h" + +namespace arrow { +class DataType; +class Field; +} // namespace arrow + +namespace paimon { + +/// One extracted field of a variant-access projection: the extraction path, cast behavior, and +/// the target field the extracted value is cast to. +struct VariantAccessSpec { + std::string path; + std::vector segments; + VariantCastArgs cast_args; + std::shared_ptr target_field; +}; + +/// Utilities for variant-access projections: a variant column read as a struct whose children +/// each carry a `__VARIANT_METADATA;;` description (mirroring the +/// Java `VariantMetadataUtils`). Such a projection extracts the described paths from the variant +/// column at read time, reading only the required shredded sub-columns. +class VariantAccessUtils { + public: + static constexpr char kMetadataKey[] = "__VARIANT_METADATA"; + static constexpr char kDelimiter = ';'; + + VariantAccessUtils() = delete; + ~VariantAccessUtils() = delete; + + /// Builds the description string encoding one access spec. + static std::string BuildVariantMetadata(const std::string& path, bool fail_on_error, + const std::string& zone_id); + + /// Whether `type` is a variant-access projection: a struct with at least one field whose + /// children all carry a `__VARIANT_METADATA` description. + static bool IsVariantAccessType(const std::shared_ptr& type); + + /// Parses the access specs of a variant-access projection field. + static Result> ParseAccessSpecs( + const std::shared_ptr& access_field); + + /// Prunes a shredded file field down to the sub-columns required by the access specs: + /// `metadata` is always kept, `typed_value` is narrowed to the requested top-level keys, and + /// `value` is kept only when some requested key is not shredded. Returns `file_field` + /// unchanged when the file is unshredded or the paths cannot be pruned (root or array-first + /// paths). + static Result> ClipShreddedFileField( + const std::vector& specs, + const std::shared_ptr& file_field); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_binary_util.cpp b/src/paimon/common/data/variant/variant_binary_util.cpp new file mode 100644 index 000000000..fd4a9c022 --- /dev/null +++ b/src/paimon/common/data/variant/variant_binary_util.cpp @@ -0,0 +1,589 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#include "paimon/common/data/variant/variant_binary_util.h" + +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_defs.h" + +namespace paimon { + +int32_t VariantDecimal::Precision() const { + __uint128_t abs_value = + unscaled < 0 ? -static_cast<__uint128_t>(unscaled) : static_cast<__uint128_t>(unscaled); + int32_t digits = 1; + while (abs_value >= 10) { + abs_value /= 10; + ++digits; + } + return digits; +} + +VariantDecimal VariantDecimal::StripTrailingZeros() const { + VariantDecimal result = *this; + if (result.unscaled == 0) { + result.scale = 0; + return result; + } + while (result.unscaled % 10 == 0) { + result.unscaled /= 10; + --result.scale; + } + return result; +} + +std::string VariantDecimal::ToPlainString() const { + __uint128_t abs_value = + unscaled < 0 ? -static_cast<__uint128_t>(unscaled) : static_cast<__uint128_t>(unscaled); + std::string digits; + if (abs_value == 0) { + digits = "0"; + } else { + while (abs_value > 0) { + digits.push_back(static_cast('0' + static_cast(abs_value % 10))); + abs_value /= 10; + } + std::reverse(digits.begin(), digits.end()); + } + std::string result; + if (unscaled < 0) { + result.push_back('-'); + } + if (scale <= 0) { + result.append(digits); + result.append(static_cast(-scale), '0'); + } else if (static_cast(scale) < digits.size()) { + result.append(digits, 0, digits.size() - scale); + result.push_back('.'); + result.append(digits, digits.size() - scale, scale); + } else { + result.append("0."); + result.append(static_cast(scale) - digits.size(), '0'); + result.append(digits); + } + return result; +} + +Status VariantBinaryUtil::MalformedVariant() { + return Status::Invalid("MALFORMED_VARIANT"); +} + +Status VariantBinaryUtil::UnknownPrimitiveTypeInVariant(int32_t id) { + return Status::Invalid(fmt::format("UNKNOWN_PRIMITIVE_TYPE_IN_VARIANT, id: {}", id)); +} + +Status VariantBinaryUtil::VariantConstructorSizeLimit() { + return Status::Invalid("VARIANT_CONSTRUCTOR_SIZE_LIMIT"); +} + +Status VariantBinaryUtil::UnexpectedType(VariantValueType type) { + static constexpr const char* kTypeNames[] = { + "OBJECT", "ARRAY", "NULL", "BOOLEAN", "LONG", "STRING", "DOUBLE", + "DECIMAL", "DATE", "TIMESTAMP", "TIMESTAMP_NTZ", "FLOAT", "BINARY", "UUID"}; + return Status::Invalid( + fmt::format("Expect type to be {}", kTypeNames[static_cast(type)])); +} + +Status VariantBinaryUtil::CheckIndex(int32_t pos, int32_t length) { + if (pos < 0 || pos >= length) { + return MalformedVariant(); + } + return Status::OK(); +} + +void VariantBinaryUtil::WriteLong(uint8_t* bytes, int32_t pos, int64_t value, int32_t num_bytes) { + for (int32_t i = 0; i < num_bytes; ++i) { + bytes[pos + i] = static_cast((static_cast(value) >> (8 * i)) & 0xFF); + } +} + +Result VariantBinaryUtil::ReadLong(std::string_view bytes, int32_t pos, + int32_t num_bytes) { + auto length = static_cast(bytes.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(pos, length)); + PAIMON_RETURN_NOT_OK(CheckIndex(pos + num_bytes - 1, length)); + uint64_t result = 0; + // All bytes except the most significant byte should be unsign-extended and shifted. The most + // significant byte should be sign-extended. + for (int32_t i = 0; i < num_bytes - 1; ++i) { + uint64_t unsigned_byte_value = static_cast(bytes[pos + i]); + result |= unsigned_byte_value << (8 * i); + } + int64_t signed_byte_value = static_cast(bytes[pos + num_bytes - 1]); + result |= static_cast(signed_byte_value) << (8 * (num_bytes - 1)); + return static_cast(result); +} + +Result VariantBinaryUtil::ReadUnsigned(std::string_view bytes, int32_t pos, + int32_t num_bytes) { + auto length = static_cast(bytes.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(pos, length)); + PAIMON_RETURN_NOT_OK(CheckIndex(pos + num_bytes - 1, length)); + int64_t result = 0; + // Similar to the `ReadLong` loop, but all bytes should be unsign-extended. + for (int32_t i = 0; i < num_bytes; ++i) { + int64_t unsigned_byte_value = static_cast(bytes[pos + i]); + result |= unsigned_byte_value << (8 * i); + } + if (result < 0 || result > std::numeric_limits::max()) { + return MalformedVariant(); + } + return static_cast(result); +} + +uint8_t VariantBinaryUtil::PrimitiveHeader(int32_t type) { + return static_cast(type << 2 | VariantDefs::kPrimitive); +} + +uint8_t VariantBinaryUtil::ShortStrHeader(int32_t size) { + return static_cast(size << 2 | VariantDefs::kShortStr); +} + +uint8_t VariantBinaryUtil::ObjectHeader(bool large_size, int32_t id_size, int32_t offset_size) { + return static_cast(((large_size ? 1 : 0) << (VariantDefs::kBasicTypeBits + 4)) | + ((id_size - 1) << (VariantDefs::kBasicTypeBits + 2)) | + ((offset_size - 1) << VariantDefs::kBasicTypeBits) | + VariantDefs::kObject); +} + +uint8_t VariantBinaryUtil::ArrayHeader(bool large_size, int32_t offset_size) { + return static_cast(((large_size ? 1 : 0) << (VariantDefs::kBasicTypeBits + 2)) | + ((offset_size - 1) << VariantDefs::kBasicTypeBits) | + VariantDefs::kArray); +} + +Result VariantBinaryUtil::GetTypeInfo(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + return (static_cast(value[pos]) >> VariantDefs::kBasicTypeBits) & + VariantDefs::kTypeInfoMask; +} + +Result VariantBinaryUtil::GetType(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + switch (basic_type) { + case VariantDefs::kShortStr: + return VariantValueType::kString; + case VariantDefs::kObject: + return VariantValueType::kObject; + case VariantDefs::kArray: + return VariantValueType::kArray; + default: + switch (type_info) { + case VariantDefs::kNull: + return VariantValueType::kNull; + case VariantDefs::kTrue: + case VariantDefs::kFalse: + return VariantValueType::kBoolean; + case VariantDefs::kInt1: + case VariantDefs::kInt2: + case VariantDefs::kInt4: + case VariantDefs::kInt8: + return VariantValueType::kLong; + case VariantDefs::kDouble: + return VariantValueType::kDouble; + case VariantDefs::kDecimal4: + case VariantDefs::kDecimal8: + case VariantDefs::kDecimal16: + return VariantValueType::kDecimal; + case VariantDefs::kDate: + return VariantValueType::kDate; + case VariantDefs::kTimestamp: + return VariantValueType::kTimestamp; + case VariantDefs::kTimestampNtz: + return VariantValueType::kTimestampNtz; + case VariantDefs::kFloat: + return VariantValueType::kFloat; + case VariantDefs::kBinary: + return VariantValueType::kBinary; + case VariantDefs::kLongStr: + return VariantValueType::kString; + case VariantDefs::kUuid: + return VariantValueType::kUuid; + default: + return UnknownPrimitiveTypeInVariant(type_info); + } + } +} + +Result VariantBinaryUtil::ValueSize(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + switch (basic_type) { + case VariantDefs::kShortStr: + return 1 + type_info; + case VariantDefs::kObject: { + PAIMON_ASSIGN_OR_RAISE(ObjectInfo info, GetObjectInfo(value, pos)); + PAIMON_ASSIGN_OR_RAISE( + int32_t data_size, + ReadUnsigned(value, info.offset_start + info.num_elements * info.offset_size, + info.offset_size)); + return info.data_start - pos + data_size; + } + case VariantDefs::kArray: { + PAIMON_ASSIGN_OR_RAISE(ArrayInfo info, GetArrayInfo(value, pos)); + PAIMON_ASSIGN_OR_RAISE( + int32_t data_size, + ReadUnsigned(value, info.offset_start + info.num_elements * info.offset_size, + info.offset_size)); + return info.data_start - pos + data_size; + } + default: + switch (type_info) { + case VariantDefs::kNull: + case VariantDefs::kTrue: + case VariantDefs::kFalse: + return 1; + case VariantDefs::kInt1: + return 2; + case VariantDefs::kInt2: + return 3; + case VariantDefs::kInt4: + case VariantDefs::kDate: + case VariantDefs::kFloat: + return 5; + case VariantDefs::kInt8: + case VariantDefs::kDouble: + case VariantDefs::kTimestamp: + case VariantDefs::kTimestampNtz: + return 9; + case VariantDefs::kDecimal4: + return 6; + case VariantDefs::kDecimal8: + return 10; + case VariantDefs::kDecimal16: + return 18; + case VariantDefs::kBinary: + case VariantDefs::kLongStr: { + PAIMON_ASSIGN_OR_RAISE(int32_t data_size, + ReadUnsigned(value, pos + 1, VariantDefs::kU32Size)); + return 1 + VariantDefs::kU32Size + data_size; + } + case VariantDefs::kUuid: + return 17; + default: + return UnknownPrimitiveTypeInVariant(type_info); + } + } +} + +Result VariantBinaryUtil::GetBoolean(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kPrimitive || + (type_info != VariantDefs::kTrue && type_info != VariantDefs::kFalse)) { + return UnexpectedType(VariantValueType::kBoolean); + } + return type_info == VariantDefs::kTrue; +} + +Result VariantBinaryUtil::GetLong(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + constexpr const char* kExceptionMessage = "Expect type to be LONG/DATE/TIMESTAMP/TIMESTAMP_NTZ"; + if (basic_type != VariantDefs::kPrimitive) { + return Status::Invalid(kExceptionMessage); + } + switch (type_info) { + case VariantDefs::kInt1: + return ReadLong(value, pos + 1, 1); + case VariantDefs::kInt2: + return ReadLong(value, pos + 1, 2); + case VariantDefs::kInt4: + case VariantDefs::kDate: + return ReadLong(value, pos + 1, 4); + case VariantDefs::kInt8: + case VariantDefs::kTimestamp: + case VariantDefs::kTimestampNtz: + return ReadLong(value, pos + 1, 8); + default: + return Status::Invalid(kExceptionMessage); + } +} + +Result VariantBinaryUtil::GetDouble(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kPrimitive || type_info != VariantDefs::kDouble) { + return UnexpectedType(VariantValueType::kDouble); + } + PAIMON_ASSIGN_OR_RAISE(int64_t bits, ReadLong(value, pos + 1, 8)); + double result; + memcpy(&result, &bits, sizeof(result)); + return result; +} + +namespace { +// Checks whether the precision and scale of the decimal are within the limit. +Status CheckDecimal(const VariantDecimal& d, int32_t max_precision) { + if (d.Precision() > max_precision || d.scale > max_precision) { + return VariantBinaryUtil::MalformedVariant(); + } + return Status::OK(); +} +} // namespace + +Result VariantBinaryUtil::GetDecimalWithOriginalScale(std::string_view value, + int32_t pos) { + auto length = static_cast(value.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(pos, length)); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kPrimitive) { + return UnexpectedType(VariantValueType::kDecimal); + } + PAIMON_RETURN_NOT_OK(CheckIndex(pos + 1, length)); + // Interpret the scale byte as unsigned. If it is a negative byte, the unsigned value must be + // greater than `kMaxDecimal16Precision` and will trigger an error in `CheckDecimal`. + int32_t scale = static_cast(value[pos + 1]); + VariantDecimal result; + result.scale = scale; + switch (type_info) { + case VariantDefs::kDecimal4: { + PAIMON_ASSIGN_OR_RAISE(int64_t unscaled, ReadLong(value, pos + 2, 4)); + result.unscaled = unscaled; + PAIMON_RETURN_NOT_OK(CheckDecimal(result, VariantDefs::kMaxDecimal4Precision)); + break; + } + case VariantDefs::kDecimal8: { + PAIMON_ASSIGN_OR_RAISE(int64_t unscaled, ReadLong(value, pos + 2, 8)); + result.unscaled = unscaled; + PAIMON_RETURN_NOT_OK(CheckDecimal(result, VariantDefs::kMaxDecimal8Precision)); + break; + } + case VariantDefs::kDecimal16: { + PAIMON_RETURN_NOT_OK(CheckIndex(pos + 17, length)); + __uint128_t unscaled = 0; + for (int32_t i = 0; i < 16; ++i) { + unscaled |= static_cast<__uint128_t>(static_cast(value[pos + 2 + i])) + << (8 * i); + } + result.unscaled = static_cast<__int128_t>(unscaled); + PAIMON_RETURN_NOT_OK(CheckDecimal(result, VariantDefs::kMaxDecimal16Precision)); + break; + } + default: + return UnexpectedType(VariantValueType::kDecimal); + } + return result; +} + +Result VariantBinaryUtil::GetDecimal(std::string_view value, int32_t pos) { + PAIMON_ASSIGN_OR_RAISE(VariantDecimal decimal, GetDecimalWithOriginalScale(value, pos)); + return decimal.StripTrailingZeros(); +} + +Result VariantBinaryUtil::GetFloat(std::string_view value, int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kPrimitive || type_info != VariantDefs::kFloat) { + return UnexpectedType(VariantValueType::kFloat); + } + PAIMON_ASSIGN_OR_RAISE(int64_t bits, ReadLong(value, pos + 1, 4)); + auto int_bits = static_cast(bits); + float result; + memcpy(&result, &int_bits, sizeof(result)); + return result; +} + +Result VariantBinaryUtil::GetBinary(std::string_view value, int32_t pos) { + auto length = static_cast(value.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(pos, length)); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kPrimitive || type_info != VariantDefs::kBinary) { + return UnexpectedType(VariantValueType::kBinary); + } + int32_t start = pos + 1 + VariantDefs::kU32Size; + PAIMON_ASSIGN_OR_RAISE(int32_t data_size, ReadUnsigned(value, pos + 1, VariantDefs::kU32Size)); + if (data_size > 0) { + PAIMON_RETURN_NOT_OK(CheckIndex(start + data_size - 1, length)); + } + return value.substr(start, data_size); +} + +Result VariantBinaryUtil::GetString(std::string_view value, int32_t pos) { + auto length = static_cast(value.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(pos, length)); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type == VariantDefs::kShortStr || + (basic_type == VariantDefs::kPrimitive && type_info == VariantDefs::kLongStr)) { + int32_t start; + int32_t str_size; + if (basic_type == VariantDefs::kShortStr) { + start = pos + 1; + str_size = type_info; + } else { + start = pos + 1 + VariantDefs::kU32Size; + PAIMON_ASSIGN_OR_RAISE(str_size, ReadUnsigned(value, pos + 1, VariantDefs::kU32Size)); + } + if (str_size > 0) { + PAIMON_RETURN_NOT_OK(CheckIndex(start + str_size - 1, length)); + } + return value.substr(start, str_size); + } + return UnexpectedType(VariantValueType::kString); +} + +Result VariantBinaryUtil::GetUuid(std::string_view value, int32_t pos) { + auto length = static_cast(value.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(pos, length)); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kPrimitive || type_info != VariantDefs::kUuid) { + return UnexpectedType(VariantValueType::kUuid); + } + int32_t start = pos + 1; + PAIMON_RETURN_NOT_OK(CheckIndex(start + 15, length)); + return value.substr(start, 16); +} + +std::string VariantBinaryUtil::UuidToString(std::string_view uuid_bytes) { + constexpr char kHexDigits[] = "0123456789abcdef"; + std::string result; + result.reserve(36); + for (size_t i = 0; i < 16; ++i) { + if (i == 4 || i == 6 || i == 8 || i == 10) { + result.push_back('-'); + } + auto byte = static_cast(uuid_bytes[i]); + result.push_back(kHexDigits[byte >> 4]); + result.push_back(kHexDigits[byte & 0xF]); + } + return result; +} + +Result VariantBinaryUtil::GetObjectInfo(std::string_view value, + int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kObject) { + return UnexpectedType(VariantValueType::kObject); + } + // Refer to the comment of the `VariantDefs::kObject` constant for the details of the object + // header encoding. Suppose `type_info` has a bit representation of 0_b4_b3b2_b1b0, the + // following line extracts b4 to determine whether the object uses a 1/4-byte size. + bool large_size = ((type_info >> 4) & 0x1) != 0; + int32_t size_bytes = large_size ? VariantDefs::kU32Size : 1; + ObjectInfo info; + PAIMON_ASSIGN_OR_RAISE(info.num_elements, ReadUnsigned(value, pos + 1, size_bytes)); + // Extracts b3b2 to determine the integer size of the field id list. + info.id_size = ((type_info >> 2) & 0x3) + 1; + // Extracts b1b0 to determine the integer size of the offset list. + info.offset_size = (type_info & 0x3) + 1; + info.id_start = pos + 1 + size_bytes; + // A corrupted variant can claim a near-INT32_MAX element count; compute the layout in + // 64-bit and bound it by the buffer before 32-bit arithmetic could overflow. + int64_t offset_start = static_cast(info.id_start) + + static_cast(info.num_elements) * info.id_size; + int64_t data_start = + offset_start + (static_cast(info.num_elements) + 1) * info.offset_size; + if (data_start > static_cast(value.size())) { + return MalformedVariant(); + } + info.offset_start = static_cast(offset_start); + info.data_start = static_cast(data_start); + return info; +} + +Result VariantBinaryUtil::GetArrayInfo(std::string_view value, + int32_t pos) { + PAIMON_RETURN_NOT_OK(CheckIndex(pos, static_cast(value.size()))); + auto header = static_cast(value[pos]); + int32_t basic_type = header & VariantDefs::kBasicTypeMask; + int32_t type_info = (header >> VariantDefs::kBasicTypeBits) & VariantDefs::kTypeInfoMask; + if (basic_type != VariantDefs::kArray) { + return UnexpectedType(VariantValueType::kArray); + } + // Suppose `type_info` has a bit representation of 000_b2_b1b0, the following line extracts b2 + // to determine whether the array uses a 1/4-byte size. + bool large_size = ((type_info >> 2) & 0x1) != 0; + int32_t size_bytes = large_size ? VariantDefs::kU32Size : 1; + ArrayInfo info; + PAIMON_ASSIGN_OR_RAISE(info.num_elements, ReadUnsigned(value, pos + 1, size_bytes)); + // Extracts b1b0 to determine the integer size of the offset list. + info.offset_size = (type_info & 0x3) + 1; + info.offset_start = pos + 1 + size_bytes; + // See GetObjectInfo: bound the 64-bit layout by the buffer before 32-bit overflow. + int64_t data_start = static_cast(info.offset_start) + + (static_cast(info.num_elements) + 1) * info.offset_size; + if (data_start > static_cast(value.size())) { + return MalformedVariant(); + } + info.data_start = static_cast(data_start); + return info; +} + +Result VariantBinaryUtil::GetMetadataKey(std::string_view metadata, int32_t id) { + auto length = static_cast(metadata.size()); + PAIMON_RETURN_NOT_OK(CheckIndex(0, length)); + // Extracts the highest 2 bits in the metadata header to determine the integer size of the + // offset list. + int32_t offset_size = ((static_cast(metadata[0]) >> 6) & 0x3) + 1; + PAIMON_ASSIGN_OR_RAISE(int32_t dict_size, ReadUnsigned(metadata, 1, offset_size)); + if (id >= dict_size) { + return MalformedVariant(); + } + // There are a header byte, a `dict_size` with `offset_size` bytes, and `(dict_size + 1)` + // offsets before the string data. + // Bound the offset-table layout in 64-bit before 32-bit arithmetic could overflow on a + // corrupted dictionary size. + int64_t string_start64 = 1 + (static_cast(dict_size) + 2) * offset_size; + if (string_start64 > static_cast(length)) { + return MalformedVariant(); + } + auto string_start = static_cast(string_start64); + PAIMON_ASSIGN_OR_RAISE(int32_t offset, + ReadUnsigned(metadata, 1 + (id + 1) * offset_size, offset_size)); + PAIMON_ASSIGN_OR_RAISE(int32_t next_offset, + ReadUnsigned(metadata, 1 + (id + 2) * offset_size, offset_size)); + if (offset > next_offset) { + return MalformedVariant(); + } + if (static_cast(string_start) + next_offset > static_cast(length)) { + return MalformedVariant(); + } + return metadata.substr(string_start + offset, next_offset - offset); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_binary_util.h b/src/paimon/common/data/variant/variant_binary_util.h new file mode 100644 index 000000000..6856d4473 --- /dev/null +++ b/src/paimon/common/data/variant/variant_binary_util.h @@ -0,0 +1,198 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#pragma once + +#include +#include +#include + +#include "paimon/result.h" + +namespace paimon { + +/// The value type of a variant value. It is determined by the header byte but not a 1:1 mapping +/// (for example, INT1/2/4/8 all map to `kLong`). +enum class VariantValueType { + kObject, + kArray, + kNull, + kBoolean, + kLong, + kString, + kDouble, + kDecimal, + kDate, + kTimestamp, + kTimestampNtz, + kFloat, + kBinary, + kUuid, +}; + +/// An arbitrary-precision (up to 38 digits) decimal value decoded from a variant binary. The +/// numeric value is `unscaled * 10^(-scale)`. Unlike `paimon::Decimal`, `scale` may become +/// negative after `StripTrailingZeros` (mirroring `java.math.BigDecimal`). +struct VariantDecimal { + __int128_t unscaled = 0; + int32_t scale = 0; + + /// Number of decimal digits in the unscaled value (a value of zero has precision 1). + int32_t Precision() const; + + /// Removes trailing zero digits from the unscaled value, increasing `10^-scale` accordingly. + VariantDecimal StripTrailingZeros() const; + + /// Plain (non-scientific) string representation, e.g. `-12.340` or `100`. + std::string ToPlainString() const; + + bool operator==(const VariantDecimal& other) const { + return unscaled == other.unscaled && scale == other.scale; + } +}; + +/// Static functions for manipulating variant binaries. See `VariantDefs` for the binary format. +class VariantBinaryUtil { + public: + VariantBinaryUtil() = delete; + ~VariantBinaryUtil() = delete; + + /// The decoded layout of a variant object value. + struct ObjectInfo { + /// Number of object fields. + int32_t num_elements; + /// The integer size of the field id list. + int32_t id_size; + /// The integer size of the offset list. + int32_t offset_size; + /// The starting index of the field id list in the variant value. + int32_t id_start; + /// The starting index of the offset list in the variant value. + int32_t offset_start; + /// The starting index of field data in the variant value. + int32_t data_start; + }; + + /// The decoded layout of a variant array value. + struct ArrayInfo { + /// Number of array elements. + int32_t num_elements; + /// The integer size of the offset list. + int32_t offset_size; + /// The starting index of the offset list in the variant value. + int32_t offset_start; + /// The starting index of element data in the variant value. + int32_t data_start; + }; + + static Status MalformedVariant(); + static Status UnknownPrimitiveTypeInVariant(int32_t id); + static Status VariantConstructorSizeLimit(); + static Status UnexpectedType(VariantValueType type); + + /// Checks the validity of an index `pos` in a buffer of `length` bytes. Returns + /// `MALFORMED_VARIANT` if it is out of bound. + static Status CheckIndex(int32_t pos, int32_t length); + + /// Writes the least significant `num_bytes` bytes in `value` into + /// `bytes[pos, pos + num_bytes)` in little endian. + static void WriteLong(uint8_t* bytes, int32_t pos, int64_t value, int32_t num_bytes); + + /// Reads a little-endian signed long value from `bytes[pos, pos + num_bytes)`. + static Result ReadLong(std::string_view bytes, int32_t pos, int32_t num_bytes); + + /// Reads a little-endian unsigned int value from `bytes[pos, pos + num_bytes)`. The value + /// must fit into a non-negative int32. + static Result ReadUnsigned(std::string_view bytes, int32_t pos, int32_t num_bytes); + + /// Adds a buffer-relative element `offset` to `base` in 64-bit and validates the result + /// stays inside `buffer_size`, guarding the 32-bit addition against corrupted offsets. + static Result CheckedElementPos(int32_t base, int32_t offset, size_t buffer_size) { + int64_t pos = static_cast(base) + offset; + if (pos >= static_cast(buffer_size)) { + return MalformedVariant(); + } + return static_cast(pos); + } + + static uint8_t PrimitiveHeader(int32_t type); + static uint8_t ShortStrHeader(int32_t size); + static uint8_t ObjectHeader(bool large_size, int32_t id_size, int32_t offset_size); + static uint8_t ArrayHeader(bool large_size, int32_t offset_size); + + /// Gets the type info bits from the variant value `value[pos...]`. + static Result GetTypeInfo(std::string_view value, int32_t pos); + + /// Gets the value type of the variant value `value[pos...]`. It is only legal to call `Get*` + /// if `GetType` returns the corresponding type (for example, it is only legal to call + /// `GetLong` if `GetType` returns `kLong`). + static Result GetType(std::string_view value, int32_t pos); + + /// Computes the size in bytes of the variant value `value[pos...]`. `value.size() - pos` is + /// an upper bound of the size, but the actual size can be smaller. + static Result ValueSize(std::string_view value, int32_t pos); + + static Result GetBoolean(std::string_view value, int32_t pos); + + /// Gets a long value from the variant value `value[pos...]`. It is only legal to call it if + /// `GetType` returns one of `kLong/kDate/kTimestamp/kTimestampNtz`. If the type is `kDate`, + /// the return value is guaranteed to fit into an int32 and represents the number of days from + /// the Unix epoch. If the type is `kTimestamp/kTimestampNtz`, the return value represents the + /// number of microseconds from the Unix epoch. + static Result GetLong(std::string_view value, int32_t pos); + + static Result GetDouble(std::string_view value, int32_t pos); + + /// Gets a decimal value from the variant value `value[pos...]`, keeping the stored scale. + static Result GetDecimalWithOriginalScale(std::string_view value, int32_t pos); + + /// Gets a decimal value from the variant value `value[pos...]` with trailing zeros stripped. + static Result GetDecimal(std::string_view value, int32_t pos); + + static Result GetFloat(std::string_view value, int32_t pos); + + /// Gets a binary value from the variant value `value[pos...]`. The returned view aliases + /// `value` and remains valid only as long as the underlying buffer. + static Result GetBinary(std::string_view value, int32_t pos); + + /// Gets a string value from the variant value `value[pos...]`. The returned view aliases + /// `value` and remains valid only as long as the underlying buffer. + static Result GetString(std::string_view value, int32_t pos); + + /// Gets a UUID value (16 bytes, big-endian) from the variant value `value[pos...]`. The + /// returned view aliases `value`. + static Result GetUuid(std::string_view value, int32_t pos); + + /// Formats a 16-byte big-endian UUID as the canonical lower-case string, e.g. + /// `123e4567-e89b-12d3-a456-426614174000`. + static std::string UuidToString(std::string_view uuid_bytes); + + /// Decodes the layout of the variant object value `value[pos...]`. + static Result GetObjectInfo(std::string_view value, int32_t pos); + + /// Decodes the layout of the variant array value `value[pos...]`. + static Result GetArrayInfo(std::string_view value, int32_t pos); + + /// Gets the key at `id` in the variant metadata. An out-of-bound `id` is considered a + /// malformed variant because it is read from the corresponding variant value. + static Result GetMetadataKey(std::string_view metadata, int32_t id); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_builder.cpp b/src/paimon/common/data/variant/variant_builder.cpp new file mode 100644 index 000000000..5aa44213e --- /dev/null +++ b/src/paimon/common/data/variant/variant_builder.cpp @@ -0,0 +1,652 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#include "paimon/common/data/variant/variant_builder.h" + +#include +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "rapidjson/error/en.h" +#include "rapidjson/memorystream.h" +#include "rapidjson/reader.h" + +namespace paimon { + +namespace { + +// A rapidjson SAX handler that feeds parsed JSON events into a `VariantBuilder`, mirroring +// `GenericVariantBuilder.buildJson` in the Java implementation. Numbers are delivered as raw +// text (`kParseNumbersAsStringsFlag`) so that exact decimal semantics match Jackson's. +class JsonToVariantHandler + : public rapidjson::BaseReaderHandler, JsonToVariantHandler> { + public: + explicit JsonToVariantHandler(VariantBuilder* builder) : builder_(builder) {} + + bool Null() { + return Ok(BeforeValue()) && Ok(builder_->AppendNull()); + } + + bool Bool(bool b) { + return Ok(BeforeValue()) && Ok(builder_->AppendBoolean(b)); + } + + bool RawNumber(const char* str, rapidjson::SizeType length, bool /*copy*/) { + return Ok(BeforeValue()) && Ok(AppendNumber(std::string_view(str, length))); + } + + bool String(const char* str, rapidjson::SizeType length, bool /*copy*/) { + return Ok(BeforeValue()) && Ok(builder_->AppendString(std::string_view(str, length))); + } + + bool StartObject() { + if (!Ok(BeforeValue())) { + return false; + } + contexts_.emplace_back(Context{true, builder_->GetWritePos(), {}, {}, {}}); + return true; + } + + bool Key(const char* str, rapidjson::SizeType length, bool /*copy*/) { + contexts_.back().pending_key.assign(str, length); + return true; + } + + bool EndObject(rapidjson::SizeType /*member_count*/) { + Context context = std::move(contexts_.back()); + contexts_.pop_back(); + return Ok(builder_->FinishWritingObject(context.start, &context.fields)); + } + + bool StartArray() { + if (!Ok(BeforeValue())) { + return false; + } + contexts_.emplace_back(Context{false, builder_->GetWritePos(), {}, {}, {}}); + return true; + } + + bool EndArray(rapidjson::SizeType /*element_count*/) { + Context context = std::move(contexts_.back()); + contexts_.pop_back(); + return Ok(builder_->FinishWritingArray(context.start, context.offsets)); + } + + const Status& status() const { + return status_; + } + + private: + struct Context { + bool is_object; + int32_t start; + std::vector fields; + std::vector offsets; + std::string pending_key; + }; + + bool Ok(const Status& status) { + if (!status.ok()) { + status_ = status; + return false; + } + return true; + } + + // Records the offset of the value that is about to be appended in the enclosing container. + Status BeforeValue() { + if (contexts_.empty()) { + return Status::OK(); + } + Context& top = contexts_.back(); + int32_t offset = builder_->GetWritePos() - top.start; + if (top.is_object) { + int32_t id = builder_->AddKey(top.pending_key); + top.fields.emplace_back(top.pending_key, id, offset); + } else { + top.offsets.push_back(offset); + } + return Status::OK(); + } + + // Mirrors the Java number handling: integers that fit in a long are appended as long; + // everything else is first tried as an exact decimal and falls back to double. + Status AppendNumber(std::string_view text) { + bool integral = true; + for (char c : text) { + if (c != '-' && !(c >= '0' && c <= '9')) { + integral = false; + break; + } + } + if (integral) { + int64_t long_value = 0; + auto [ptr, ec] = std::from_chars(text.data(), text.data() + text.size(), long_value); + if (ec == std::errc() && ptr == text.data() + text.size()) { + return builder_->AppendLong(long_value); + } + } + PAIMON_ASSIGN_OR_RAISE(bool appended, TryAppendDecimal(text)); + if (appended) { + return Status::OK(); + } + char* end = nullptr; + std::string text_copy(text); + double double_value = std::strtod(text_copy.c_str(), &end); + if (end != text_copy.c_str() + text_copy.size()) { + return Status::Invalid(fmt::format("Invalid JSON number: {}", text)); + } + return builder_->AppendDouble(double_value); + } + + // Tries to append a JSON number as an exact decimal. Returns whether it succeeded. The input + // must only use the decimal format (an integer value with an optional '.' in it) and must not + // use scientific notation. It also must fit into the precision limitation of decimal types. + Result TryAppendDecimal(std::string_view text) { + for (char c : text) { + if (c != '-' && c != '.' && !(c >= '0' && c <= '9')) { + return false; + } + } + bool negative = false; + size_t i = 0; + if (i < text.size() && text[i] == '-') { + negative = true; + ++i; + } + __int128_t unscaled = 0; + int32_t scale = 0; + int32_t significant_digits = 0; + bool seen_point = false; + bool seen_nonzero = false; + for (; i < text.size(); ++i) { + char c = text[i]; + if (c == '.') { + seen_point = true; + continue; + } + if (seen_point) { + ++scale; + } + if (c != '0' || seen_nonzero) { + seen_nonzero = true; + ++significant_digits; + } + if (significant_digits > VariantDefs::kMaxDecimal16Precision) { + return false; + } + unscaled = unscaled * 10 + (c - '0'); + } + if (scale > VariantDefs::kMaxDecimal16Precision) { + return false; + } + VariantDecimal decimal; + decimal.unscaled = negative ? -unscaled : unscaled; + decimal.scale = scale; + PAIMON_RETURN_NOT_OK(builder_->AppendDecimal(decimal)); + return true; + } + + VariantBuilder* builder_; + std::vector contexts_; + Status status_; +}; + +} // namespace + +Result> VariantBuilder::ParseJson( + std::string_view json, bool allow_duplicate_keys, const std::shared_ptr& pool) { + VariantBuilder builder(allow_duplicate_keys); + JsonToVariantHandler handler(&builder); + rapidjson::Reader reader; + rapidjson::MemoryStream stream(json.data(), json.size()); + rapidjson::ParseResult result = + reader.Parse(stream, handler); + if (!result) { + if (!handler.status().ok()) { + return handler.status(); + } + return Status::Invalid(fmt::format("Failed to parse JSON: {} (at offset {})", + rapidjson::GetParseError_En(result.Code()), + result.Offset())); + } + return builder.Build(pool); +} + +Result> VariantBuilder::Build( + const std::shared_ptr& pool) { + auto num_keys = static_cast(dictionary_keys_.size()); + // Use int64 to avoid overflow in accumulating lengths. + int64_t dictionary_string_size = 0; + for (const std::string& key : dictionary_keys_) { + dictionary_string_size += static_cast(key.size()); + } + // Determine the number of bytes required per offset entry. The largest offset is the + // one-past-the-end value, which is the total string size. It's very unlikely that the number + // of keys could be larger, but incorporate that into the calculation in case of pathological + // data. + int64_t max_size = std::max(dictionary_string_size, static_cast(num_keys)); + if (max_size > VariantDefs::kSizeLimit) { + return Status::Invalid("VARIANT_SIZE_LIMIT"); + } + int32_t offset_size = GetIntegerSize(static_cast(max_size)); + + int32_t offset_start = 1 + offset_size; + int32_t string_start = offset_start + (num_keys + 1) * offset_size; + int64_t metadata_size = string_start + dictionary_string_size; + if (metadata_size > VariantDefs::kSizeLimit) { + return Status::Invalid("VARIANT_SIZE_LIMIT"); + } + + std::shared_ptr metadata = + Bytes::AllocateBytes(static_cast(metadata_size), pool.get()); + auto* metadata_data = reinterpret_cast(metadata->data()); + int32_t header_byte = VariantDefs::kVersion | ((offset_size - 1) << 6); + VariantBinaryUtil::WriteLong(metadata_data, 0, header_byte, 1); + VariantBinaryUtil::WriteLong(metadata_data, 1, num_keys, offset_size); + int32_t current_offset = 0; + for (int32_t i = 0; i < num_keys; ++i) { + VariantBinaryUtil::WriteLong(metadata_data, offset_start + i * offset_size, current_offset, + offset_size); + const std::string& key = dictionary_keys_[i]; + memcpy(metadata_data + string_start + current_offset, key.data(), key.size()); + current_offset += static_cast(key.size()); + } + VariantBinaryUtil::WriteLong(metadata_data, offset_start + num_keys * offset_size, + current_offset, offset_size); + + std::shared_ptr value = + Bytes::AllocateBytes(static_cast(write_pos_), pool.get()); + memcpy(value->data(), write_buffer_.data(), static_cast(write_pos_)); + return GenericVariant::Create(std::move(value), std::move(metadata)); +} + +Status VariantBuilder::AppendString(std::string_view str) { + bool long_str = static_cast(str.size()) > VariantDefs::kMaxShortStrSize; + PAIMON_RETURN_NOT_OK(CheckCapacity((long_str ? 1 + VariantDefs::kU32Size : 1) + + static_cast(str.size()))); + if (long_str) { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kLongStr); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, + static_cast(str.size()), VariantDefs::kU32Size); + write_pos_ += VariantDefs::kU32Size; + } else { + write_buffer_[write_pos_++] = + VariantBinaryUtil::ShortStrHeader(static_cast(str.size())); + } + memcpy(write_buffer_.data() + write_pos_, str.data(), str.size()); + write_pos_ += static_cast(str.size()); + return Status::OK(); +} + +Status VariantBuilder::AppendNull() { + PAIMON_RETURN_NOT_OK(CheckCapacity(1)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kNull); + return Status::OK(); +} + +Status VariantBuilder::AppendBoolean(bool b) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1)); + write_buffer_[write_pos_++] = + VariantBinaryUtil::PrimitiveHeader(b ? VariantDefs::kTrue : VariantDefs::kFalse); + return Status::OK(); +} + +Status VariantBuilder::AppendLong(int64_t l) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 8)); + if (l == static_cast(l)) { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kInt1); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, l, 1); + write_pos_ += 1; + } else if (l == static_cast(l)) { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kInt2); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, l, 2); + write_pos_ += 2; + } else if (l == static_cast(l)) { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kInt4); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, l, 4); + write_pos_ += 4; + } else { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kInt8); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, l, 8); + write_pos_ += 8; + } + return Status::OK(); +} + +Status VariantBuilder::AppendDouble(double d) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 8)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDouble); + int64_t bits; + memcpy(&bits, &d, sizeof(bits)); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, bits, 8); + write_pos_ += 8; + return Status::OK(); +} + +Status VariantBuilder::AppendDecimal(const VariantDecimal& d) { + PAIMON_RETURN_NOT_OK(CheckCapacity(2 + 16)); + int32_t precision = d.Precision(); + if (d.scale < 0 || d.scale > VariantDefs::kMaxDecimal16Precision || + precision > VariantDefs::kMaxDecimal16Precision) { + return Status::Invalid( + fmt::format("Decimal precision {} and scale {} must fit into the variant decimal " + "limit {}", + precision, d.scale, VariantDefs::kMaxDecimal16Precision)); + } + if (d.scale <= VariantDefs::kMaxDecimal4Precision && + precision <= VariantDefs::kMaxDecimal4Precision) { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDecimal4); + write_buffer_[write_pos_++] = static_cast(d.scale); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, + static_cast(d.unscaled), 4); + write_pos_ += 4; + } else if (d.scale <= VariantDefs::kMaxDecimal8Precision && + precision <= VariantDefs::kMaxDecimal8Precision) { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDecimal8); + write_buffer_[write_pos_++] = static_cast(d.scale); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, + static_cast(d.unscaled), 8); + write_pos_ += 8; + } else { + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDecimal16); + write_buffer_[write_pos_++] = static_cast(d.scale); + auto bits = static_cast<__uint128_t>(d.unscaled); + for (int32_t i = 0; i < 16; ++i) { + write_buffer_[write_pos_ + i] = static_cast((bits >> (8 * i)) & 0xFF); + } + write_pos_ += 16; + } + return Status::OK(); +} + +Status VariantBuilder::AppendDate(int32_t days_since_epoch) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 4)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDate); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, days_since_epoch, 4); + write_pos_ += 4; + return Status::OK(); +} + +Status VariantBuilder::AppendTimestamp(int64_t micros_since_epoch) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 8)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kTimestamp); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, micros_since_epoch, 8); + write_pos_ += 8; + return Status::OK(); +} + +Status VariantBuilder::AppendTimestampNtz(int64_t micros_since_epoch) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 8)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kTimestampNtz); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, micros_since_epoch, 8); + write_pos_ += 8; + return Status::OK(); +} + +Status VariantBuilder::AppendFloat(float f) { + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 4)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kFloat); + int32_t bits; + memcpy(&bits, &f, sizeof(bits)); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, bits, 4); + write_pos_ += 4; + return Status::OK(); +} + +Status VariantBuilder::AppendBinary(std::string_view binary) { + PAIMON_RETURN_NOT_OK( + CheckCapacity(1 + VariantDefs::kU32Size + static_cast(binary.size()))); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kBinary); + VariantBinaryUtil::WriteLong(write_buffer_.data(), write_pos_, + static_cast(binary.size()), VariantDefs::kU32Size); + write_pos_ += VariantDefs::kU32Size; + memcpy(write_buffer_.data() + write_pos_, binary.data(), binary.size()); + write_pos_ += static_cast(binary.size()); + return Status::OK(); +} + +Status VariantBuilder::AppendUuid(std::string_view uuid_bytes) { + if (uuid_bytes.size() != 16) { + return Status::Invalid("UUID must be 16 bytes"); + } + PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 16)); + write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kUuid); + // UUID is stored big-endian, so don't use WriteLong. + memcpy(write_buffer_.data() + write_pos_, uuid_bytes.data(), 16); + write_pos_ += 16; + return Status::OK(); +} + +int32_t VariantBuilder::AddKey(std::string_view key) { + auto it = dictionary_.find(std::string(key)); + if (it != dictionary_.end()) { + return it->second; + } + auto id = static_cast(dictionary_keys_.size()); + dictionary_.emplace(std::string(key), id); + dictionary_keys_.emplace_back(key); + return id; +} + +Status VariantBuilder::FinishWritingObject(int32_t start, std::vector* fields) { + auto size = static_cast(fields->size()); + std::sort(fields->begin(), fields->end(), + [](const FieldEntry& a, const FieldEntry& b) { return a.key < b.key; }); + int32_t max_id = size == 0 ? 0 : (*fields)[0].id; + if (allow_duplicate_keys_) { + int32_t distinct_pos = 0; + // Maintain a list of distinct keys in-place. + for (int32_t i = 1; i < size; ++i) { + max_id = std::max(max_id, (*fields)[i].id); + if ((*fields)[i].id == (*fields)[i - 1].id) { + // Found a duplicate key. Keep the field with a greater offset. + if ((*fields)[distinct_pos].offset < (*fields)[i].offset) { + (*fields)[distinct_pos].offset = (*fields)[i].offset; + } + } else { + // Found a distinct key. Add the field to the list. + ++distinct_pos; + (*fields)[distinct_pos] = (*fields)[i]; + } + } + if (distinct_pos + 1 < size) { + size = distinct_pos + 1; + fields->erase(fields->begin() + size, fields->end()); + // Sort the fields by offsets so that we can move the value data of each field to the + // new offset without overwriting the fields after it. + std::sort(fields->begin(), fields->end(), + [](const FieldEntry& a, const FieldEntry& b) { return a.offset < b.offset; }); + int32_t current_offset = 0; + for (int32_t i = 0; i < size; ++i) { + int32_t old_offset = (*fields)[i].offset; + PAIMON_ASSIGN_OR_RAISE( + int32_t field_size, + VariantBinaryUtil::ValueSize( + std::string_view(reinterpret_cast(write_buffer_.data()), + static_cast(write_pos_)), + start + old_offset)); + memmove(write_buffer_.data() + start + current_offset, + write_buffer_.data() + start + old_offset, static_cast(field_size)); + (*fields)[i].offset = current_offset; + current_offset += field_size; + } + write_pos_ = start + current_offset; + // Change back to the sort order by field keys to meet the variant spec. + std::sort(fields->begin(), fields->end(), + [](const FieldEntry& a, const FieldEntry& b) { return a.key < b.key; }); + } + } else { + for (int32_t i = 1; i < size; ++i) { + max_id = std::max(max_id, (*fields)[i].id); + if ((*fields)[i].key == (*fields)[i - 1].key) { + return Status::Invalid("VARIANT_DUPLICATE_KEY"); + } + } + } + int32_t data_size = write_pos_ - start; + bool large_size = size > VariantDefs::kU8Max; + int32_t size_bytes = large_size ? VariantDefs::kU32Size : 1; + int32_t id_size = GetIntegerSize(max_id); + int32_t offset_size = GetIntegerSize(data_size); + // The space for the header byte, object size, id list, and offset list. + int32_t header_size = 1 + size_bytes + size * id_size + (size + 1) * offset_size; + PAIMON_RETURN_NOT_OK(CheckCapacity(header_size)); + // Shift the just-written field data to make room for the object header section. + memmove(write_buffer_.data() + start + header_size, write_buffer_.data() + start, + static_cast(data_size)); + write_pos_ += header_size; + write_buffer_[start] = VariantBinaryUtil::ObjectHeader(large_size, id_size, offset_size); + VariantBinaryUtil::WriteLong(write_buffer_.data(), start + 1, size, size_bytes); + int32_t id_start = start + 1 + size_bytes; + int32_t offset_start = id_start + size * id_size; + for (int32_t i = 0; i < size; ++i) { + VariantBinaryUtil::WriteLong(write_buffer_.data(), id_start + i * id_size, (*fields)[i].id, + id_size); + VariantBinaryUtil::WriteLong(write_buffer_.data(), offset_start + i * offset_size, + (*fields)[i].offset, offset_size); + } + VariantBinaryUtil::WriteLong(write_buffer_.data(), offset_start + size * offset_size, data_size, + offset_size); + return Status::OK(); +} + +Status VariantBuilder::FinishWritingArray(int32_t start, const std::vector& offsets) { + int32_t data_size = write_pos_ - start; + auto size = static_cast(offsets.size()); + bool large_size = size > VariantDefs::kU8Max; + int32_t size_bytes = large_size ? VariantDefs::kU32Size : 1; + int32_t offset_size = GetIntegerSize(data_size); + // The space for the header byte, array size, and offset list. + int32_t header_size = 1 + size_bytes + (size + 1) * offset_size; + PAIMON_RETURN_NOT_OK(CheckCapacity(header_size)); + // Shift the just-written element data to make room for the header section. + memmove(write_buffer_.data() + start + header_size, write_buffer_.data() + start, + static_cast(data_size)); + write_pos_ += header_size; + write_buffer_[start] = VariantBinaryUtil::ArrayHeader(large_size, offset_size); + VariantBinaryUtil::WriteLong(write_buffer_.data(), start + 1, size, size_bytes); + int32_t offset_start = start + 1 + size_bytes; + for (int32_t i = 0; i < size; ++i) { + VariantBinaryUtil::WriteLong(write_buffer_.data(), offset_start + i * offset_size, + offsets[i], offset_size); + } + VariantBinaryUtil::WriteLong(write_buffer_.data(), offset_start + size * offset_size, data_size, + offset_size); + return Status::OK(); +} + +Status VariantBuilder::AppendVariant(const GenericVariant& v) { + return AppendVariantImpl(v.RawValue(), v.Metadata(), v.Pos()); +} + +Status VariantBuilder::AppendVariantImpl(std::string_view value, std::string_view metadata, + int32_t pos) { + PAIMON_RETURN_NOT_OK(VariantBinaryUtil::CheckIndex(pos, static_cast(value.size()))); + int32_t basic_type = static_cast(value[pos]) & VariantDefs::kBasicTypeMask; + switch (basic_type) { + case VariantDefs::kObject: { + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ObjectInfo info, + VariantBinaryUtil::GetObjectInfo(value, pos)); + std::vector fields; + fields.reserve(info.num_elements); + int32_t start = write_pos_; + for (int32_t i = 0; i < info.num_elements; ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t id, + VariantBinaryUtil::ReadUnsigned( + value, info.id_start + info.id_size * i, info.id_size)); + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, + VariantBinaryUtil::ReadUnsigned(value, info.offset_start + info.offset_size * i, + info.offset_size)); + int32_t element_pos = info.data_start + offset; + PAIMON_ASSIGN_OR_RAISE(std::string_view key, + VariantBinaryUtil::GetMetadataKey(metadata, id)); + int32_t new_id = AddKey(key); + fields.emplace_back(std::string(key), new_id, write_pos_ - start); + PAIMON_RETURN_NOT_OK(AppendVariantImpl(value, metadata, element_pos)); + } + return FinishWritingObject(start, &fields); + } + case VariantDefs::kArray: { + PAIMON_ASSIGN_OR_RAISE(VariantBinaryUtil::ArrayInfo info, + VariantBinaryUtil::GetArrayInfo(value, pos)); + std::vector offsets; + offsets.reserve(info.num_elements); + int32_t start = write_pos_; + for (int32_t i = 0; i < info.num_elements; ++i) { + PAIMON_ASSIGN_OR_RAISE( + int32_t offset, + VariantBinaryUtil::ReadUnsigned(value, info.offset_start + info.offset_size * i, + info.offset_size)); + int32_t element_pos = info.data_start + offset; + offsets.push_back(write_pos_ - start); + PAIMON_RETURN_NOT_OK(AppendVariantImpl(value, metadata, element_pos)); + } + return FinishWritingArray(start, offsets); + } + default: + return ShallowAppendVariant(value, pos); + } +} + +Status VariantBuilder::ShallowAppendVariant(std::string_view value, int32_t pos) { + PAIMON_ASSIGN_OR_RAISE(int32_t size, VariantBinaryUtil::ValueSize(value, pos)); + PAIMON_RETURN_NOT_OK( + VariantBinaryUtil::CheckIndex(pos + size - 1, static_cast(value.size()))); + PAIMON_RETURN_NOT_OK(CheckCapacity(size)); + memcpy(write_buffer_.data() + write_pos_, value.data() + pos, static_cast(size)); + write_pos_ += size; + return Status::OK(); +} + +Status VariantBuilder::CheckCapacity(int32_t additional) { + int32_t required = write_pos_ + additional; + if (required > static_cast(write_buffer_.size())) { + // Allocate a new buffer with a capacity of the next power of 2 of `required`. + auto new_capacity = static_cast(write_buffer_.size()); + while (new_capacity < required) { + new_capacity *= 2; + if (new_capacity > VariantDefs::kSizeLimit) { + return Status::Invalid("VARIANT_SIZE_LIMIT"); + } + } + write_buffer_.resize(static_cast(new_capacity)); + } + return Status::OK(); +} + +int32_t VariantBuilder::GetIntegerSize(int32_t value) { + if (value <= VariantDefs::kU8Max) { + return 1; + } + if (value <= VariantDefs::kU16Max) { + return 2; + } + if (value <= VariantDefs::kU24Max) { + return 3; + } + return 4; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_builder.h b/src/paimon/common/data/variant/variant_builder.h new file mode 100644 index 000000000..af24b4b89 --- /dev/null +++ b/src/paimon/common/data/variant/variant_builder.h @@ -0,0 +1,144 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_binary_util.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace paimon { + +/// Builds variant value and metadata binaries, either by parsing JSON values or by appending +/// values directly. +class VariantBuilder { + public: + /// Temporarily stores the information of a field. All fields of a JSON object are collected, + /// sorted by their keys, and then the variant object is built in sorted order. + struct FieldEntry { + std::string key; + int32_t id; + int32_t offset; + + FieldEntry(std::string key, int32_t id, int32_t offset) + : key(std::move(key)), id(id), offset(offset) {} + }; + + explicit VariantBuilder(bool allow_duplicate_keys) + : allow_duplicate_keys_(allow_duplicate_keys) {} + + /// Parses a JSON string as a variant value. + /// + /// When `allow_duplicate_keys` is true, the last occurrence of a duplicate object key wins; + /// otherwise duplicate keys make the parse fail. + static Result> ParseJson( + std::string_view json, bool allow_duplicate_keys, const std::shared_ptr& pool); + + /// Builds the variant metadata from the collected dictionary keys and returns the variant + /// result. + Result> Build(const std::shared_ptr& pool); + + /// The variant value written so far, without metadata. Used in shredding to produce a final + /// value where all shredded values refer to a common metadata. + std::string_view ValueWithoutMetadata() const { + return {reinterpret_cast(write_buffer_.data()), + static_cast(write_pos_)}; + } + + Status AppendString(std::string_view str); + Status AppendNull(); + Status AppendBoolean(bool b); + /// Appends a long value. The actual used integer type depends on the value range. + Status AppendLong(int64_t l); + Status AppendDouble(double d); + /// Appends a decimal value. Its precision and scale must fit into `kMaxDecimal16Precision`. + Status AppendDecimal(const VariantDecimal& d); + Status AppendDate(int32_t days_since_epoch); + Status AppendTimestamp(int64_t micros_since_epoch); + Status AppendTimestampNtz(int64_t micros_since_epoch); + Status AppendFloat(float f); + Status AppendBinary(std::string_view binary); + /// Appends a UUID value (16 bytes, big-endian). + Status AppendUuid(std::string_view uuid_bytes); + + /// Adds a key to the variant dictionary and returns its id. If the key already exists, the + /// dictionary is not modified. + int32_t AddKey(std::string_view key); + + /// The current write position of the variant builder. It is used together with + /// `FinishWritingObject` or `FinishWritingArray`. + int32_t GetWritePos() const { + return write_pos_; + } + + /// Finishes writing a variant object after all of its fields have already been written. The + /// process is as follows: + /// 1. The caller calls `GetWritePos` before writing any fields to obtain the `start` + /// parameter. + /// 2. The caller appends all the object fields to the builder. In the meantime, it should + /// maintain the `fields` parameter. Before appending each field, it should append an entry + /// to `fields` to record the offset of the field, computed as `GetWritePos() - start`. + /// 3. The caller calls `FinishWritingObject` to finish writing a variant object. + /// + /// This function sorts the fields by key. If there are duplicate field keys: + /// - when `allow_duplicate_keys` is true, the field with the greatest offset value (the last + /// appended one) is kept; + /// - otherwise, the call fails. + Status FinishWritingObject(int32_t start, std::vector* fields); + + /// Finishes writing a variant array after all of its elements have already been written. The + /// process is similar to that of `FinishWritingObject`. + Status FinishWritingArray(int32_t start, const std::vector& offsets); + + /// Appends a variant value. The keys of the input variant are inserted into the current + /// variant dictionary and the value is rebuilt with new field ids. For scalar values, the + /// binary slice is copied directly. + Status AppendVariant(const GenericVariant& v); + + /// Appends the variant value without rewriting or creating any metadata. This is used when + /// building an object during shredding, where there is a fixed pre-existing metadata that all + /// shredded values refer to. + Status ShallowAppendVariant(std::string_view value, int32_t pos); + + private: + Status CheckCapacity(int32_t additional); + Status AppendVariantImpl(std::string_view value, std::string_view metadata, int32_t pos); + static int32_t GetIntegerSize(int32_t value); + + // The write buffer in building the variant value. Its first `write_pos_` bytes have been + // written. + std::vector write_buffer_ = std::vector(128); + int32_t write_pos_ = 0; + // Maps keys to a monotonically increasing id. + std::unordered_map dictionary_; + // Stores all keys in `dictionary_` in the order of id. + std::vector dictionary_keys_; + const bool allow_duplicate_keys_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_defs.h b/src/paimon/common/data/variant/variant_defs.h new file mode 100644 index 000000000..9ead20ddd --- /dev/null +++ b/src/paimon/common/data/variant/variant_defs.h @@ -0,0 +1,161 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace paimon { + +/// Constants of the Paimon Variant type and the Variant Binary Encoding, which follows the +/// parquet-format VariantEncoding.md specification (compatible with the Java / Spark +/// implementation). +/// +/// A variant is made up of 2 binaries: value and metadata. A variant value consists of a one-byte +/// header and a number of content bytes (can be zero). The header byte is divided into upper 6 +/// bits (called "type info") and lower 2 bits (called "basic type"). +/// +/// The variant metadata includes a version id and a dictionary of distinct strings +/// (case-sensitive). Its binary format is: +/// - Version: 1-byte unsigned integer. The only acceptable value is 1 currently. +/// - Dictionary size: `offset_size`-byte little-endian unsigned integer. The number of keys in the +/// dictionary. +/// - Offsets: (size + 1) * `offset_size`-byte little-endian unsigned integers. `offsets[i]` +/// represents the starting position of string i, counting starting from the address of +/// `offsets[0]`. Strings must be stored contiguously, so we don't need to store the string size, +/// instead, we compute it with `offset[i + 1] - offset[i]`. +/// - UTF-8 string data. +/// +/// A Variant field uses `struct` as its +/// underlying physical storage in Apache Arrow Schema, and is marked as the Paimon Variant +/// extension type by attaching specific **KeyValueMetadata** on the outer field. +class VariantDefs { + public: + VariantDefs() = delete; + ~VariantDefs() = delete; + + /// Metadata key identifying a Paimon extension type field (shared with BLOB). + static constexpr char kExtensionTypeKey[] = "paimon.extension.type"; + /// Metadata value identifying a Paimon Variant extension type field. + static constexpr char kExtensionTypeValue[] = "paimon.type.variant"; + + /// Name of the binary child field holding the variant value. + static constexpr char kValueFieldName[] = "value"; + /// Name of the binary child field holding the variant metadata. + static constexpr char kMetadataFieldName[] = "metadata"; + /// Name of the typed child field of a shredded variant (parquet VariantShredding.md). + static constexpr char kTypedValueFieldName[] = "typed_value"; + /// Paimon field id of the `value` child field. + static constexpr int32_t kValueFieldId = 0; + /// Paimon field id of the `metadata` child field. + static constexpr int32_t kMetadataFieldId = 1; + + static constexpr int32_t kBasicTypeBits = 2; + static constexpr int32_t kBasicTypeMask = 0x3; + static constexpr int32_t kTypeInfoMask = 0x3F; + /// The inclusive maximum value of the type info value. It is the size limit of `kShortStr`. + static constexpr int32_t kMaxShortStrSize = 0x3F; + + /// Primitive value. The type info value must be one of the primitive type values below. + static constexpr int32_t kPrimitive = 0; + /// Short string value. The type info value is the string size, which must be in + /// `[0, kMaxShortStrSize]`. The string content bytes directly follow the header byte. + static constexpr int32_t kShortStr = 1; + /// Object value. The content contains a size, a list of field ids, a list of field offsets, + /// and the actual field data. The length of the id list is `size`, while the length of the + /// offset list is `size + 1`, where the last offset represents the total size of the field + /// data. The fields in an object must be sorted by the field name in alphabetical order. + /// Duplicate field names in one object are not allowed. + /// The type info is 0_b4_b3b2_b1b0 (MSB is 0), where: + /// - b4 specifies the type of size. When it is 0/1, `size` is a little-endian 1/4-byte + /// unsigned integer. + /// - b3b2/b1b0 specifies the integer type of id and offset. When the 2 bits are 0/1/2, the + /// list contains 1/2/3-byte little-endian unsigned integers. + static constexpr int32_t kObject = 2; + /// Array value. The content contains a size, a list of field offsets, and the actual element + /// data. It is similar to an object without the id list. The type info is 000_b2_b1b0: + /// - b2 specifies the type of size. + /// - b1b0 specifies the integer type of offset. + static constexpr int32_t kArray = 3; + + /// JSON null value. Empty content. + static constexpr int32_t kNull = 0; + /// True value. Empty content. + static constexpr int32_t kTrue = 1; + /// False value. Empty content. + static constexpr int32_t kFalse = 2; + /// 1-byte little-endian signed integer. + static constexpr int32_t kInt1 = 3; + /// 2-byte little-endian signed integer. + static constexpr int32_t kInt2 = 4; + /// 4-byte little-endian signed integer. + static constexpr int32_t kInt4 = 5; + /// 8-byte little-endian signed integer. + static constexpr int32_t kInt8 = 6; + /// 8-byte IEEE double. + static constexpr int32_t kDouble = 7; + /// 4-byte decimal. Content is 1-byte scale + 4-byte little-endian signed integer. + static constexpr int32_t kDecimal4 = 8; + /// 8-byte decimal. Content is 1-byte scale + 8-byte little-endian signed integer. + static constexpr int32_t kDecimal8 = 9; + /// 16-byte decimal. Content is 1-byte scale + 16-byte little-endian signed integer. + static constexpr int32_t kDecimal16 = 10; + /// Date value. Content is 4-byte little-endian signed integer that represents the number of + /// days from the Unix epoch. + static constexpr int32_t kDate = 11; + /// Timestamp value. Content is 8-byte little-endian signed integer that represents the number + /// of microseconds elapsed since the Unix epoch, 1970-01-01 00:00:00 UTC. It is displayed to + /// users in their local time zones and may be displayed differently depending on the + /// execution environment. + static constexpr int32_t kTimestamp = 12; + /// Timestamp_ntz value. It has the same content as `kTimestamp` but should always be + /// interpreted as if the local time zone is UTC. + static constexpr int32_t kTimestampNtz = 13; + /// 4-byte IEEE float. + static constexpr int32_t kFloat = 14; + /// Binary value. The content is (4-byte little-endian unsigned integer representing the + /// binary size) + (size bytes of binary content). + static constexpr int32_t kBinary = 15; + /// Long string value. The content is (4-byte little-endian unsigned integer representing the + /// string size) + (size bytes of string content). + static constexpr int32_t kLongStr = 16; + /// UUID, 16-byte big-endian. + static constexpr int32_t kUuid = 20; + + /// The only acceptable variant version. It is stored in the lower 4 bits of the first + /// metadata byte. + static constexpr uint8_t kVersion = 1; + static constexpr uint8_t kVersionMask = 0x0F; + + static constexpr int32_t kU8Max = 0xFF; + static constexpr int32_t kU16Max = 0xFFFF; + static constexpr int32_t kU24Max = 0xFFFFFF; + static constexpr int32_t kU24Size = 3; + static constexpr int32_t kU32Size = 4; + + /// Both variant value and variant metadata need to be no longer than 128MiB. + static constexpr int32_t kSizeLimit = 128 * 1024 * 1024; + + static constexpr int32_t kMaxDecimal4Precision = 9; + static constexpr int32_t kMaxDecimal8Precision = 18; + static constexpr int32_t kMaxDecimal16Precision = 38; + + /// Object field lookup switches from linear search to binary search when the object size + /// reaches this threshold. + static constexpr int32_t kBinarySearchThreshold = 32; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_get.cpp b/src/paimon/common/data/variant/variant_get.cpp new file mode 100644 index 000000000..82258412d --- /dev/null +++ b/src/paimon/common/data/variant/variant_get.cpp @@ -0,0 +1,407 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_get.h" + +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_binary_util.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_json_utils.h" +#include "paimon/common/data/variant/variant_path_segment.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/field_type_utils.h" +#include "paimon/core/casting/cast_executor_factory.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" + +namespace paimon { + +namespace { + +Result> InvalidCast(const std::shared_ptr& variant, + const std::shared_ptr& target_type, + const VariantCastArgs& cast_args) { + if (cast_args.fail_on_error) { + PAIMON_ASSIGN_OR_RAISE(std::string json, variant->ToJson(cast_args.zone_id)); + return Status::Invalid(fmt::format("Invalid cast {} to {}", json, target_type->ToString())); + } + return std::optional(std::nullopt); +} + +Result> CastVariant(const std::shared_ptr& variant, + const std::shared_ptr& target_type, + const VariantCastArgs& cast_args) { + switch (target_type->id()) { + case arrow::Type::type::STRUCT: + case arrow::Type::type::LIST: + case arrow::Type::type::MAP: + return Status::NotImplemented(fmt::format( + "variant_get to nested type {} is not supported", target_type->ToString())); + default: + break; + } + + PAIMON_ASSIGN_OR_RAISE(VariantValueType variant_type, variant->GetType()); + if (variant_type == VariantValueType::kNull) { + return std::optional(std::nullopt); + } + + bool target_is_string = target_type->id() == arrow::Type::type::STRING; + if (variant_type == VariantValueType::kUuid) { + // There's no UUID type in Paimon. We only allow it to be cast to string. + if (target_is_string) { + PAIMON_ASSIGN_OR_RAISE(std::string_view uuid, variant->GetUuid()); + std::string uuid_str = VariantBinaryUtil::UuidToString(uuid); + return std::optional( + Literal(FieldType::STRING, uuid_str.data(), uuid_str.size())); + } + return InvalidCast(variant, target_type, cast_args); + } + + std::optional input; + std::shared_ptr input_type; + switch (variant_type) { + case VariantValueType::kObject: + case VariantValueType::kArray: { + if (target_is_string) { + PAIMON_ASSIGN_OR_RAISE(std::string json, variant->ToJson(cast_args.zone_id)); + return std::optional(Literal(FieldType::STRING, json.data(), json.size())); + } + return InvalidCast(variant, target_type, cast_args); + } + case VariantValueType::kBoolean: { + PAIMON_ASSIGN_OR_RAISE(bool value, variant->GetBoolean()); + input = Literal(value); + input_type = arrow::boolean(); + break; + } + case VariantValueType::kLong: { + PAIMON_ASSIGN_OR_RAISE(int64_t value, variant->GetLong()); + input = Literal(value); + input_type = arrow::int64(); + break; + } + case VariantValueType::kString: { + PAIMON_ASSIGN_OR_RAISE(std::string_view value, variant->GetString()); + input = Literal(FieldType::STRING, value.data(), value.size()); + input_type = arrow::utf8(); + break; + } + case VariantValueType::kDouble: { + PAIMON_ASSIGN_OR_RAISE(double value, variant->GetDouble()); + input = Literal(value); + input_type = arrow::float64(); + break; + } + case VariantValueType::kDecimal: { + PAIMON_ASSIGN_OR_RAISE(VariantDecimal value, variant->GetDecimal()); + if (value.scale < 0) { + // `paimon::Decimal` requires a non-negative scale; scale the value back up. + for (; value.scale < 0; ++value.scale) { + value.unscaled *= 10; + } + } + int32_t precision = std::max(value.Precision(), value.scale); + int32_t scale = value.scale; + input = Literal(Decimal(precision, scale, value.unscaled)); + input_type = arrow::decimal128(precision, scale); + break; + } + case VariantValueType::kDate: { + PAIMON_ASSIGN_OR_RAISE(int64_t value, variant->GetLong()); + input = Literal(FieldType::DATE, static_cast(value)); + input_type = arrow::date32(); + break; + } + case VariantValueType::kFloat: { + PAIMON_ASSIGN_OR_RAISE(float value, variant->GetFloat()); + input = Literal(value); + input_type = arrow::float32(); + break; + } + case VariantValueType::kBinary: { + PAIMON_ASSIGN_OR_RAISE(std::string_view value, variant->GetBinary()); + input = Literal(FieldType::BINARY, value.data(), value.size()); + input_type = arrow::binary(); + break; + } + case VariantValueType::kTimestamp: + case VariantValueType::kTimestampNtz: { + PAIMON_ASSIGN_OR_RAISE(int64_t micros, variant->GetLong()); + // Floor the division so negative epochs keep a non-negative sub-millisecond part. + int64_t millis = micros / 1000; + auto sub_micros = static_cast(micros % 1000); + if (sub_micros < 0) { + millis -= 1; + sub_micros += 1000; + } + input = Literal(Timestamp::FromEpochMillis(millis, sub_micros * 1000)); + input_type = variant_type == VariantValueType::kTimestamp + ? arrow::timestamp(arrow::TimeUnit::MICRO, "UTC") + : arrow::timestamp(arrow::TimeUnit::MICRO); + break; + } + default: + return Status::Invalid(fmt::format("Unsupported variant type in variant_get: {}", + static_cast(variant_type))); + } + + if (input_type->Equals(*target_type)) { + return input; + } + + PAIMON_ASSIGN_OR_RAISE(FieldType input_field_type, + FieldTypeUtils::ConvertToFieldType(input_type->id())); + PAIMON_ASSIGN_OR_RAISE(FieldType target_field_type, + FieldTypeUtils::ConvertToFieldType(target_type->id())); + std::shared_ptr executor = + CastExecutorFactory::GetCastExecutorFactory()->GetCastExecutor(input_field_type, + target_field_type); + if (executor == nullptr) { + return InvalidCast(variant, target_type, cast_args); + } + Result cast_result = executor->Cast(*input, target_type); + if (!cast_result.ok()) { + return InvalidCast(variant, target_type, cast_args); + } + return std::optional(std::move(cast_result).value()); +} + +Status AppendLiteralToBuilder(const Literal& literal, + const std::shared_ptr& target_type, + arrow::ArrayBuilder* builder) { + switch (target_type->id()) { + case arrow::Type::type::BOOL: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::INT8: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::INT16: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::INT32: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::INT64: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::FLOAT: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::DOUBLE: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::STRING: + return ToPaimonStatus(static_cast(builder)->Append( + literal.GetValue())); + case arrow::Type::type::BINARY: + return ToPaimonStatus(static_cast(builder)->Append( + literal.GetValue())); + case arrow::Type::type::DATE32: + return ToPaimonStatus( + static_cast(builder)->Append(literal.GetValue())); + case arrow::Type::type::TIMESTAMP: { + auto timestamp = literal.GetValue(); + const auto& timestamp_type = static_cast(*target_type); + int64_t value; + switch (timestamp_type.unit()) { + case arrow::TimeUnit::SECOND: + value = timestamp.GetMillisecond() / 1000; + break; + case arrow::TimeUnit::MILLI: + value = timestamp.GetMillisecond(); + break; + case arrow::TimeUnit::MICRO: + value = timestamp.ToMicrosecond(); + break; + case arrow::TimeUnit::NANO: + value = timestamp.ToNanosecond(); + break; + default: + return Status::Invalid("Unsupported timestamp unit"); + } + return ToPaimonStatus(static_cast(builder)->Append(value)); + } + case arrow::Type::type::DECIMAL128: { + auto decimal = literal.GetValue(); + arrow::Decimal128 value(static_cast(decimal.HighBits()), decimal.LowBits()); + return ToPaimonStatus(static_cast(builder)->Append(value)); + } + default: + return Status::Invalid( + fmt::format("Unsupported variant_get target type: {}", target_type->ToString())); + } +} + +} // namespace + +Status VariantGetExecutor::CastToBuilder(const std::shared_ptr& variant, + const std::shared_ptr& target_field, + arrow::ArrayBuilder* builder, + const VariantCastArgs& cast_args, + const std::shared_ptr& pool) { + if (variant == nullptr) { + return ToPaimonStatus(builder->AppendNull()); + } + + auto invalid_cast = [&]() -> Status { + if (cast_args.fail_on_error) { + PAIMON_ASSIGN_OR_RAISE(std::string json, variant->ToJson(cast_args.zone_id)); + return Status::Invalid( + fmt::format("Invalid cast {} to {}", json, target_field->type()->ToString())); + } + return ToPaimonStatus(builder->AppendNull()); + }; + + if (VariantTypeUtils::IsVariantField(target_field)) { + VariantBuilder variant_builder(/*allow_duplicate_keys=*/false); + PAIMON_RETURN_NOT_OK(variant_builder.AppendVariant(*variant)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr copied, variant_builder.Build(pool)); + auto* struct_builder = static_cast(builder); + PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Append()); + PAIMON_ASSIGN_OR_RAISE(std::string_view value, copied->Value()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(struct_builder->field_builder(0))->Append(value)); + return ToPaimonStatus(static_cast(struct_builder->field_builder(1)) + ->Append(copied->Metadata())); + } + + PAIMON_ASSIGN_OR_RAISE(VariantValueType variant_type, variant->GetType()); + if (variant_type == VariantValueType::kNull) { + return ToPaimonStatus(builder->AppendNull()); + } + + switch (target_field->type()->id()) { + case arrow::Type::type::STRUCT: { + if (variant_type != VariantValueType::kObject) { + return invalid_cast(); + } + auto* struct_builder = static_cast(builder); + PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Append()); + const auto& struct_type = static_cast(*target_field->type()); + for (int i = 0; i < struct_type.num_fields(); ++i) { + const std::shared_ptr& child_field = struct_type.field(i); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr child, + variant->GetFieldByKey(child_field->name())); + PAIMON_RETURN_NOT_OK(CastToBuilder( + child, child_field, struct_builder->field_builder(i), cast_args, pool)); + } + return Status::OK(); + } + case arrow::Type::type::MAP: { + const auto& map_type = static_cast(*target_field->type()); + if (map_type.key_type()->id() != arrow::Type::type::STRING || + variant_type != VariantValueType::kObject) { + return invalid_cast(); + } + auto* map_builder = static_cast(builder); + PAIMON_RETURN_NOT_OK_FROM_ARROW(map_builder->Append()); + PAIMON_ASSIGN_OR_RAISE(int32_t object_size, variant->ObjectSize()); + for (int32_t i = 0; i < object_size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::optional field, + variant->GetFieldAtIndex(i)); + if (!field.has_value()) { + return Status::Invalid(fmt::format("Malformed variant object at index {}", i)); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(map_builder->key_builder()) + ->Append(field->key)); + PAIMON_RETURN_NOT_OK(CastToBuilder(field->value, map_type.item_field(), + map_builder->item_builder(), cast_args, pool)); + } + return Status::OK(); + } + case arrow::Type::type::LIST: { + if (variant_type != VariantValueType::kArray) { + return invalid_cast(); + } + auto* list_builder = static_cast(builder); + PAIMON_RETURN_NOT_OK_FROM_ARROW(list_builder->Append()); + const auto& list_type = static_cast(*target_field->type()); + PAIMON_ASSIGN_OR_RAISE(int32_t array_size, variant->ArraySize()); + for (int32_t i = 0; i < array_size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element, + variant->GetElementAtIndex(i)); + PAIMON_RETURN_NOT_OK(CastToBuilder(element, list_type.value_field(), + list_builder->value_builder(), cast_args, pool)); + } + return Status::OK(); + } + default: { + PAIMON_ASSIGN_OR_RAISE(std::optional literal, + CastVariant(variant, target_field->type(), cast_args)); + if (!literal.has_value()) { + return ToPaimonStatus(builder->AppendNull()); + } + return AppendLiteralToBuilder(*literal, target_field->type(), builder); + } + } +} + +Result> VariantGetExecutor::GetAsArrow( + const std::shared_ptr& variant, const std::string& path, + const std::shared_ptr& target_field, const VariantCastArgs& cast_args, + const std::shared_ptr& pool, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr extracted, ExtractByPath(variant, path)); + std::unique_ptr builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder(arrow_pool, target_field->type(), &builder)); + PAIMON_RETURN_NOT_OK(CastToBuilder(extracted, target_field, builder.get(), cast_args, pool)); + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Finish(&array)); + return array; +} + +Result> VariantGetExecutor::ExtractByPath( + const std::shared_ptr& variant, const std::string& path) { + PAIMON_ASSIGN_OR_RAISE(std::vector segments, + VariantPathSegment::Parse(path)); + std::shared_ptr current = variant; + for (const VariantPathSegment& segment : segments) { + PAIMON_ASSIGN_OR_RAISE(VariantValueType type, current->GetType()); + if (segment.kind == VariantPathSegment::Kind::kObjectExtraction && + type == VariantValueType::kObject) { + PAIMON_ASSIGN_OR_RAISE(current, current->GetFieldByKey(segment.key)); + } else if (segment.kind == VariantPathSegment::Kind::kArrayExtraction && + type == VariantValueType::kArray) { + PAIMON_ASSIGN_OR_RAISE(current, current->GetElementAtIndex(segment.index)); + } else { + return std::shared_ptr(nullptr); + } + if (current == nullptr) { + return std::shared_ptr(nullptr); + } + } + return current; +} + +Result> VariantGetExecutor::Get( + const std::shared_ptr& variant, const std::string& path, + const std::shared_ptr& target_type, const VariantCastArgs& cast_args) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr extracted, ExtractByPath(variant, path)); + if (extracted == nullptr) { + return std::optional(std::nullopt); + } + return CastVariant(extracted, target_type, cast_args); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_get.h b/src/paimon/common/data/variant/variant_get.h new file mode 100644 index 000000000..eb2953a23 --- /dev/null +++ b/src/paimon/common/data/variant/variant_get.h @@ -0,0 +1,83 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/generic_variant.h" +#include "paimon/data/variant.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class ArrayBuilder; +class DataType; +class Field; +class MemoryPool; +} // namespace arrow + +namespace paimon { + +/// Implements `variant_get` semantics: extracting a sub-variant by a JSONPath-like path and +/// casting it to a target type. +class VariantGetExecutor { + public: + VariantGetExecutor() = delete; + ~VariantGetExecutor() = delete; + + /// Extracts a sub-variant value according to a path which starts with a `$`, e.g. `$.key`, + /// `$['key']`, `$["key"]`, `$.array[0]`. Returns nullptr if the path does not match the + /// variant structure. + static Result> ExtractByPath( + const std::shared_ptr& variant, const std::string& path); + + /// Extracts a sub-variant by `path` and casts it to `target_type`. Returns nullopt for SQL + /// NULL (unmatched path, variant null, or an invalid cast with `fail_on_error == false`). + /// + /// Nested target types (ROW/ARRAY/MAP/VARIANT) are not supported by the `Literal` result + /// type; use `CastToBuilder` / `GetAsArrow` for structured results. + static Result> Get(const std::shared_ptr& variant, + const std::string& path, + const std::shared_ptr& target_type, + const VariantCastArgs& cast_args); + + /// Casts `variant` to the type of `target_field` and appends the result to `builder`. + /// + /// Supported targets beyond scalars: a variant-marked struct field (the variant is deeply + /// re-encoded), STRUCT (from a variant object, children matched by name; unmatched children + /// are null), MAP with string keys (from a variant object), and LIST (from a variant array). + /// A nullptr `variant`, a variant null, or an invalid cast with `fail_on_error == false` + /// appends null. + static Status CastToBuilder(const std::shared_ptr& variant, + const std::shared_ptr& target_field, + arrow::ArrayBuilder* builder, const VariantCastArgs& cast_args, + const std::shared_ptr& pool); + + /// Extracts a sub-variant by `path`, casts it to the (possibly nested) type of + /// `target_field`, and returns a length-1 arrow array holding the result (a null slot + /// represents SQL NULL). `arrow_pool` must outlive the returned array. + static Result> GetAsArrow( + const std::shared_ptr& variant, const std::string& path, + const std::shared_ptr& target_field, const VariantCastArgs& cast_args, + const std::shared_ptr& pool, arrow::MemoryPool* arrow_pool); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_get_test.cpp b/src/paimon/common/data/variant/variant_get_test.cpp new file mode 100644 index 000000000..51a6ec8a9 --- /dev/null +++ b/src/paimon/common/data/variant/variant_get_test.cpp @@ -0,0 +1,309 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_get.h" + +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_path_segment.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class VariantGetTest : public ::testing::Test { + public: + void SetUp() override { + // The same document as the Java GenericVariantTest#testVariantGet. + std::string json = + "{\n" + " \"object\": {\n" + " \"name\": \"Apache Paimon\",\n" + " \"age\": 2,\n" + " \"address\": {\n" + " \"street\": \"Main St\",\n" + " \"city\": \"Hangzhou\"\n" + " }\n" + " },\n" + " \"array\": [1, 2, 3, 4, 5],\n" + " \"string\": \"Hello, World!\",\n" + " \"long\": 12345678901234,\n" + " \"double\": 1.0123456789012345678901234567890123456789,\n" + " \"decimal\": 100.99,\n" + " \"boolean1\": true,\n" + " \"boolean2\": false,\n" + " \"nullField\": null\n" + "}\n"; + ASSERT_OK_AND_ASSIGN(variant_, GenericVariant::FromJson(json, pool_)); + cast_args_.fail_on_error = false; + } + + std::optional Get(const std::string& path, + const std::shared_ptr& target) { + auto result = VariantGetExecutor::Get(variant_, path, target, cast_args_); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return result.value(); + } + + std::string GetString(const std::string& path, const std::shared_ptr& target) { + std::optional literal = Get(path, target); + EXPECT_TRUE(literal.has_value()); + return literal->ToString(); + } + + Result> GetAsArrow( + const std::string& path, const std::shared_ptr& target_field) { + return VariantGetExecutor::GetAsArrow(variant_, path, target_field, cast_args_, pool_, + arrow::default_memory_pool()); + } + + protected: + std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr variant_; + VariantCastArgs cast_args_; +}; + +TEST_F(VariantGetTest, PathSegmentParse) { + ASSERT_OK_AND_ASSIGN(auto segments, + VariantPathSegment::Parse("$[\"object\"]['address'].city[3]")); + ASSERT_EQ(segments.size(), 4); + ASSERT_EQ(segments[0].kind, VariantPathSegment::Kind::kObjectExtraction); + ASSERT_EQ(segments[0].key, "object"); + ASSERT_EQ(segments[1].kind, VariantPathSegment::Kind::kObjectExtraction); + ASSERT_EQ(segments[1].key, "address"); + ASSERT_EQ(segments[2].kind, VariantPathSegment::Kind::kObjectExtraction); + ASSERT_EQ(segments[2].key, "city"); + ASSERT_EQ(segments[3].kind, VariantPathSegment::Kind::kArrayExtraction); + ASSERT_EQ(segments[3].index, 3); + + ASSERT_OK_AND_ASSIGN(auto root_only, VariantPathSegment::Parse("$")); + ASSERT_TRUE(root_only.empty()); + + // Java parity: the root `$` is located anywhere in the path (Java uses Matcher#find), so a + // prefix before the first `$` is tolerated and ignored. + ASSERT_OK_AND_ASSIGN(auto prefixed, VariantPathSegment::Parse("abc$.x")); + ASSERT_EQ(prefixed.size(), 1); + ASSERT_EQ(prefixed[0].kind, VariantPathSegment::Kind::kObjectExtraction); + ASSERT_EQ(prefixed[0].key, "x"); + + ASSERT_NOK(VariantPathSegment::Parse("")); + ASSERT_NOK(VariantPathSegment::Parse("no_root")); + ASSERT_NOK(VariantPathSegment::Parse("$.")); + ASSERT_NOK(VariantPathSegment::Parse("$[abc]")); + ASSERT_NOK(VariantPathSegment::Parse("$['unterminated]")); +} + +TEST_F(VariantGetTest, ScalarTargets) { + ASSERT_EQ(GetString("$.string", arrow::utf8()), "Hello, World!"); + auto long_value = Get("$.long", arrow::int64()); + ASSERT_TRUE(long_value.has_value()); + ASSERT_EQ(long_value->GetValue(), 12345678901234LL); + ASSERT_EQ(GetString("$.long", arrow::utf8()), "12345678901234"); + auto double_value = Get("$.double", arrow::float64()); + ASSERT_TRUE(double_value.has_value()); + ASSERT_DOUBLE_EQ(double_value->GetValue(), 1.0123456789012346); + auto decimal_value = Get("$.decimal", arrow::decimal128(5, 2)); + ASSERT_TRUE(decimal_value.has_value()); + ASSERT_EQ(decimal_value->GetValue().ToUnscaledLong(), 10099); + ASSERT_EQ(GetString("$.decimal", arrow::utf8()), "100.99"); + auto bool1 = Get("$.boolean1", arrow::boolean()); + ASSERT_TRUE(bool1.has_value()); + ASSERT_TRUE(bool1->GetValue()); + auto bool2 = Get("$.boolean2", arrow::boolean()); + ASSERT_TRUE(bool2.has_value()); + ASSERT_FALSE(bool2->GetValue()); + // Variant null maps to SQL NULL. + ASSERT_FALSE(Get("$.nullField", arrow::boolean()).has_value()); + auto elem = Get("$.array[3]", arrow::int64()); + ASSERT_TRUE(elem.has_value()); + ASSERT_EQ(elem->GetValue(), 4); +} + +TEST_F(VariantGetTest, ContainerToJsonString) { + ASSERT_EQ(GetString("$.object", arrow::utf8()), + "{\"address\":{\"city\":\"Hangzhou\",\"street\":\"Main St\"},\"age\":2,\"name\":" + "\"Apache Paimon\"}"); + ASSERT_EQ(GetString("$.object.name", arrow::utf8()), "Apache Paimon"); + ASSERT_EQ(GetString("$.object.address.street", arrow::utf8()), "Main St"); + ASSERT_EQ(GetString("$[\"object\"]['address'].city", arrow::utf8()), "Hangzhou"); + ASSERT_EQ(GetString("$.array", arrow::utf8()), "[1,2,3,4,5]"); +} + +TEST_F(VariantGetTest, UnmatchedPathAndInvalidCast) { + // A path that does not exist yields SQL NULL. + ASSERT_FALSE(Get("$.missing", arrow::utf8()).has_value()); + ASSERT_FALSE(Get("$.string[0]", arrow::utf8()).has_value()); + ASSERT_FALSE(Get("$.array.key", arrow::utf8()).has_value()); + // An invalid cast yields SQL NULL when fail_on_error is false. + ASSERT_FALSE(Get("$.object", arrow::int64()).has_value()); + // ... and an error when fail_on_error is true. + cast_args_.fail_on_error = true; + auto result = VariantGetExecutor::Get(variant_, "$.object", arrow::int64(), cast_args_); + ASSERT_NOK(result); +} + +TEST_F(VariantGetTest, NestedTargetsNotImplemented) { + auto result = + VariantGetExecutor::Get(variant_, "$.array", arrow::list(arrow::int32()), cast_args_); + ASSERT_TRUE(result.status().IsNotImplemented()); +} + +TEST_F(VariantGetTest, BinaryAndTimestampSources) { + // Java-encoded data can carry binary/timestamp scalars that JSON parsing never produces. + { + VariantBuilder builder(/*allow_duplicate_keys=*/false); + ASSERT_OK(builder.AppendBinary(std::string_view("\x01\x02\x03", 3))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN(std::optional literal, + VariantGetExecutor::Get(variant, "$", arrow::binary(), cast_args_)); + ASSERT_TRUE(literal.has_value()); + ASSERT_EQ(literal->GetValue(), std::string("\x01\x02\x03", 3)); + } + { + VariantBuilder builder(/*allow_duplicate_keys=*/false); + ASSERT_OK(builder.AppendTimestamp(1700000000123456)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN( + std::optional literal, + VariantGetExecutor::Get(variant, "$", arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), + cast_args_)); + ASSERT_TRUE(literal.has_value()); + ASSERT_EQ(literal->GetValue().ToMicrosecond(), 1700000000123456); + } + { + VariantBuilder builder(/*allow_duplicate_keys=*/false); + ASSERT_OK(builder.AppendTimestampNtz(-1001)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, builder.Build(pool_)); + ASSERT_OK_AND_ASSIGN( + std::optional literal, + VariantGetExecutor::Get(variant, "$", arrow::timestamp(arrow::TimeUnit::MICRO), + cast_args_)); + ASSERT_TRUE(literal.has_value()); + // Negative epochs must floor, not truncate toward zero. + ASSERT_EQ(literal->GetValue().ToMicrosecond(), -1001); + } +} + +TEST_F(VariantGetTest, ExtractByPath) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr address, + VariantGetExecutor::ExtractByPath(variant_, "$.object.address")); + ASSERT_NE(address, nullptr); + ASSERT_OK_AND_ASSIGN(std::string json, address->ToJson()); + ASSERT_EQ(json, "{\"city\":\"Hangzhou\",\"street\":\"Main St\"}"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr missing, + VariantGetExecutor::ExtractByPath(variant_, "$.object.missing")); + ASSERT_EQ(missing, nullptr); +} + +TEST_F(VariantGetTest, CastToStructTarget) { + auto target = arrow::field( + "r", arrow::struct_( + {arrow::field("name", arrow::utf8()), arrow::field("age", arrow::int64()), + arrow::field("address", arrow::struct_({arrow::field("city", arrow::utf8())})), + arrow::field("missing", arrow::utf8())})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, GetAsArrow("$.object", target)); + ASSERT_EQ(array->length(), 1); + const auto& row = static_cast(*array); + ASSERT_FALSE(row.IsNull(0)); + ASSERT_EQ(static_cast(*row.field(0)).GetString(0), "Apache Paimon"); + ASSERT_EQ(static_cast(*row.field(1)).Value(0), 2); + const auto& address = static_cast(*row.field(2)); + ASSERT_EQ(static_cast(*address.field(0)).GetString(0), "Hangzhou"); + // A target field absent from the variant object is null. + ASSERT_TRUE(row.field(3)->IsNull(0)); +} + +TEST_F(VariantGetTest, CastToStructFromNonObject) { + auto target = arrow::field("r", arrow::struct_({arrow::field("a", arrow::int64())})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, GetAsArrow("$.array", target)); + ASSERT_TRUE(array->IsNull(0)); + cast_args_.fail_on_error = true; + auto result = VariantGetExecutor::GetAsArrow(variant_, "$.array", target, cast_args_, pool_, + arrow::default_memory_pool()); + ASSERT_NOK(result); +} + +TEST_F(VariantGetTest, CastToMapTarget) { + auto target = arrow::field("m", arrow::map(arrow::utf8(), arrow::utf8())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, + GetAsArrow("$.object.address", target)); + const auto& map = static_cast(*array); + ASSERT_EQ(map.value_length(0), 2); + const auto& keys = static_cast(*map.keys()); + const auto& items = static_cast(*map.items()); + ASSERT_EQ(keys.GetString(0), "city"); + ASSERT_EQ(items.GetString(0), "Hangzhou"); + ASSERT_EQ(keys.GetString(1), "street"); + ASSERT_EQ(items.GetString(1), "Main St"); + // A map with a non-string key type is an invalid cast. + auto bad_target = arrow::field("m", arrow::map(arrow::int64(), arrow::utf8())); + ASSERT_OK_AND_ASSIGN(array, GetAsArrow("$.object.address", bad_target)); + ASSERT_TRUE(array->IsNull(0)); +} + +TEST_F(VariantGetTest, CastToListTarget) { + auto target = arrow::field("l", arrow::list(arrow::int64())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, GetAsArrow("$.array", target)); + const auto& list = static_cast(*array); + ASSERT_EQ(list.value_length(0), 5); + const auto& values = static_cast(*list.values()); + for (int64_t i = 0; i < 5; ++i) { + ASSERT_EQ(values.Value(i), i + 1); + } + auto list_of_structs = + arrow::field("l", arrow::list(arrow::struct_({arrow::field("a", arrow::int64())}))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rows, + GenericVariant::FromJson("[{\"a\": 1}, {\"a\": 2}]", pool_)); + ASSERT_OK_AND_ASSIGN( + array, VariantGetExecutor::GetAsArrow(rows, "$", list_of_structs, cast_args_, pool_, + arrow::default_memory_pool())); + const auto& struct_list = static_cast(*array); + ASSERT_EQ(struct_list.value_length(0), 2); + const auto& elements = static_cast(*struct_list.values()); + ASSERT_EQ(static_cast(*elements.field(0)).Value(0), 1); + ASSERT_EQ(static_cast(*elements.field(0)).Value(1), 2); +} + +TEST_F(VariantGetTest, CastToVariantTarget) { + auto target = VariantTypeUtils::ToArrowField("v", /*nullable=*/true, {}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, + GetAsArrow("$.object.address", target)); + const auto& row = static_cast(*array); + ASSERT_FALSE(row.IsNull(0)); + std::string_view value = static_cast(*row.field(0)).GetView(0); + std::string_view metadata = static_cast(*row.field(1)).GetView(0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr copied, + GenericVariant::Create(value, metadata, pool_)); + ASSERT_OK_AND_ASSIGN(std::string json, copied->ToJson()); + ASSERT_EQ(json, "{\"city\":\"Hangzhou\",\"street\":\"Main St\"}"); +} + +TEST_F(VariantGetTest, NestedTargetNullSemantics) { + auto target = arrow::field("r", arrow::struct_({arrow::field("a", arrow::int64())})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr array, GetAsArrow("$.missing", target)); + ASSERT_TRUE(array->IsNull(0)); + // A variant null yields SQL NULL even with fail_on_error == true. + cast_args_.fail_on_error = true; + ASSERT_OK_AND_ASSIGN(array, GetAsArrow("$.nullField", target)); + ASSERT_TRUE(array->IsNull(0)); +} + +} // namespace paimon::test diff --git a/src/paimon/common/data/variant/variant_json_utils.cpp b/src/paimon/common/data/variant/variant_json_utils.cpp new file mode 100644 index 000000000..2d95b190b --- /dev/null +++ b/src/paimon/common/data/variant/variant_json_utils.cpp @@ -0,0 +1,294 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_json_utils.h" + +#include +#include +#include +#include + +#include "arrow/vendored/datetime.h" +#include "fmt/format.h" + +namespace paimon { + +namespace { + +constexpr int64_t kMicrosPerSecond = 1000000LL; +constexpr int64_t kMicrosPerDay = 86400LL * kMicrosPerSecond; + +// Formats the shortest round-trip decimal digits + decimal exponent into the textual form of +// `java.lang.Double.toString` / `java.lang.Float.toString`: plain notation when +// `1e-3 <= |value| < 1e7`, computerized scientific notation otherwise. +std::string JavaFloatingToString(std::string_view digits, int32_t exp, bool negative) { + std::string result; + if (negative) { + result.push_back('-'); + } + if (exp >= -3 && exp < 7) { + if (exp >= 0) { + size_t int_digits = static_cast(exp) + 1; + if (digits.size() > int_digits) { + result.append(digits, 0, int_digits); + result.push_back('.'); + result.append(digits.substr(int_digits)); + } else { + result.append(digits); + result.append(int_digits - digits.size(), '0'); + result.append(".0"); + } + } else { + result.append("0."); + result.append(static_cast(-exp) - 1, '0'); + result.append(digits); + } + } else { + result.push_back(digits[0]); + result.push_back('.'); + if (digits.size() > 1) { + result.append(digits.substr(1)); + } else { + result.push_back('0'); + } + result.push_back('E'); + result.append(std::to_string(exp)); + } + return result; +} + +template +std::string FloatingToJavaString(T value) { + if (std::isnan(value)) { + return "NaN"; + } + if (std::isinf(value)) { + return value > 0 ? "Infinity" : "-Infinity"; + } + if (value == 0) { + return std::signbit(value) ? "-0.0" : "0.0"; + } + // The shortest round-trip representation in scientific form, e.g. `1.234e+05`. gcc 8's + // lacks floating-point to_chars, so probe increasing precision until the value + // round-trips; dropping a trailing zero digit is an exact rounding, so the first precision + // that round-trips carries no trailing zeros. + constexpr int kMaxFractionDigits = std::is_same_v ? 8 : 16; + char buf[64]; + int len = 0; + for (int precision = 0; precision <= kMaxFractionDigits; ++precision) { + len = std::snprintf(buf, sizeof(buf), "%.*e", precision, static_cast(value)); + T parsed; + if constexpr (std::is_same_v) { + parsed = std::strtof(buf, nullptr); + } else { + parsed = std::strtod(buf, nullptr); + } + if (parsed == value) { + break; + } + } + std::string_view repr(buf, static_cast(len)); + bool negative = repr[0] == '-'; + if (negative) { + repr.remove_prefix(1); + } + size_t e_pos = repr.find('e'); + std::string_view mantissa = repr.substr(0, e_pos); + size_t exp_start = e_pos + 1; + // `std::from_chars` does not accept a leading '+'. + if (repr[exp_start] == '+') { + ++exp_start; + } + int32_t exp = 0; + std::from_chars(repr.data() + exp_start, repr.data() + repr.size(), exp); + std::string digits; + digits.push_back(mantissa[0]); + if (mantissa.size() > 2) { + digits.append(mantissa.substr(2)); + } + return JavaFloatingToString(digits, exp, negative); +} + +// Appends a year like `java.time.LocalDate.toString`: absolute value padded to at least 4 +// digits; years above 9999 get a `+` prefix. +void AppendJavaYear(int64_t year, std::string* out) { + if (year < 0) { + out->push_back('-'); + year = -year; + out->append(fmt::format("{:04}", year)); + } else { + if (year > 9999) { + out->push_back('+'); + } + out->append(fmt::format("{:04}", year)); + } +} + +void AppendDate(int64_t days_since_epoch, std::string* out) { + using arrow_vendored::date::days; + using arrow_vendored::date::sys_days; + using arrow_vendored::date::year_month_day; + year_month_day ymd{sys_days{days{days_since_epoch}}}; + AppendJavaYear(static_cast(ymd.year()), out); + out->append(fmt::format("-{:02}-{:02}", static_cast(ymd.month()), + static_cast(ymd.day()))); +} + +int64_t FloorDiv(int64_t x, int64_t y) { + int64_t quotient = x / y; + if ((x % y != 0) && ((x < 0) != (y < 0))) { + --quotient; + } + return quotient; +} + +} // namespace + +void VariantJsonUtils::AppendEscapedJson(std::string_view str, std::string* out) { + out->push_back('"'); + for (char c : str) { + switch (c) { + case '"': + out->append("\\\""); + break; + case '\\': + out->append("\\\\"); + break; + case '\b': + out->append("\\b"); + break; + case '\f': + out->append("\\f"); + break; + case '\n': + out->append("\\n"); + break; + case '\r': + out->append("\\r"); + break; + case '\t': + out->append("\\t"); + break; + default: + if (static_cast(c) < 0x20) { + out->append(fmt::format("\\u{:04x}", static_cast(c))); + } else { + out->push_back(c); + } + } + } + out->push_back('"'); +} + +std::string VariantJsonUtils::JavaDoubleToString(double value) { + return FloatingToJavaString(value); +} + +std::string VariantJsonUtils::JavaFloatToString(float value) { + return FloatingToJavaString(value); +} + +std::string VariantJsonUtils::DateToString(int32_t days_since_epoch) { + std::string result; + AppendDate(days_since_epoch, &result); + return result; +} + +std::string VariantJsonUtils::TimestampToString(int64_t micros_since_epoch, int32_t offset_seconds, + bool with_offset) { + int64_t local_micros = micros_since_epoch + static_cast(offset_seconds) * 1000000; + int64_t days = FloorDiv(local_micros, kMicrosPerDay); + int64_t micros_of_day = local_micros - days * kMicrosPerDay; + std::string result; + AppendDate(days, &result); + int64_t seconds_of_day = micros_of_day / kMicrosPerSecond; + int64_t micros_of_second = micros_of_day % kMicrosPerSecond; + result.append(fmt::format(" {:02}:{:02}:{:02}", seconds_of_day / 3600, + (seconds_of_day / 60) % 60, seconds_of_day % 60)); + if (micros_of_second != 0) { + std::string fraction = fmt::format("{:06}", micros_of_second); + while (fraction.back() == '0') { + fraction.pop_back(); + } + result.push_back('.'); + result.append(fraction); + } + if (with_offset) { + int32_t abs_offset = offset_seconds >= 0 ? offset_seconds : -offset_seconds; + result.append(fmt::format("{}{:02}:{:02}", offset_seconds >= 0 ? '+' : '-', + abs_offset / 3600, (abs_offset / 60) % 60)); + } + return result; +} + +Result VariantJsonUtils::GetZoneOffsetSeconds(const std::string& zone_id, + int64_t micros_since_epoch) { + std::string_view id = zone_id; + if (id == "Z" || id == "UTC" || id == "GMT" || id == "UT") { + return 0; + } + // `UTC+08:00` style ids: strip the prefix and parse the remaining fixed offset. + if (id.size() > 3 && (id.substr(0, 3) == "UTC" || id.substr(0, 3) == "GMT")) { + id.remove_prefix(3); + } else if (id.size() > 2 && id.substr(0, 2) == "UT" && (id[2] == '+' || id[2] == '-')) { + id.remove_prefix(2); + } + if (!id.empty() && (id[0] == '+' || id[0] == '-')) { + bool negative = id[0] == '-'; + id.remove_prefix(1); + // Accepted forms: H, HH, HH:MM, HHMM, HH:MM:SS, HHMMSS. + std::string digits; + for (char c : id) { + if (c >= '0' && c <= '9') { + digits.push_back(c); + } else if (c != ':') { + return Status::Invalid(fmt::format("Invalid zone offset: {}", zone_id)); + } + } + int32_t hours = 0; + int32_t minutes = 0; + int32_t seconds = 0; + if (digits.size() == 1 || digits.size() == 2) { + hours = std::stoi(digits); + } else if (digits.size() == 4) { + hours = std::stoi(digits.substr(0, 2)); + minutes = std::stoi(digits.substr(2, 2)); + } else if (digits.size() == 6) { + hours = std::stoi(digits.substr(0, 2)); + minutes = std::stoi(digits.substr(2, 2)); + seconds = std::stoi(digits.substr(4, 2)); + } else { + return Status::Invalid(fmt::format("Invalid zone offset: {}", zone_id)); + } + if (hours > 18 || minutes > 59 || seconds > 59) { + return Status::Invalid(fmt::format("Invalid zone offset: {}", zone_id)); + } + int32_t total = hours * 3600 + minutes * 60 + seconds; + return negative ? -total : total; + } + // IANA region id, resolved at the given instant (honoring DST). + try { + const auto* zone = arrow_vendored::date::locate_zone(zone_id); + std::chrono::time_point tp{ + std::chrono::microseconds(micros_since_epoch)}; + auto info = zone->get_info(tp); + return static_cast(info.offset.count()); + } catch (const std::exception& e) { + return Status::Invalid(fmt::format("Invalid zone id: {}, {}", zone_id, e.what())); + } +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_json_utils.h b/src/paimon/common/data/variant/variant_json_utils.h new file mode 100644 index 000000000..de7b4111f --- /dev/null +++ b/src/paimon/common/data/variant/variant_json_utils.h @@ -0,0 +1,59 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/result.h" + +namespace paimon { + +/// Helpers for rendering variant values as JSON text, matching the output of the Java +/// implementation (`GenericVariant.toJsonImpl`) character-for-character. +class VariantJsonUtils { + public: + VariantJsonUtils() = delete; + ~VariantJsonUtils() = delete; + + /// Appends `str` as a quoted JSON string with escaping into `out`. + static void AppendEscapedJson(std::string_view str, std::string* out); + + /// Formats a double like `java.lang.Double.toString`, e.g. `1.0`, `-0.001`, `1.0E7`. + static std::string JavaDoubleToString(double value); + + /// Formats a float like `java.lang.Float.toString`. + static std::string JavaFloatToString(float value); + + /// Formats days-since-epoch like `java.time.LocalDate.toString`, e.g. `2024-01-15`. + static std::string DateToString(int32_t days_since_epoch); + + /// Formats microseconds-since-epoch at the given offset like the Java formatter + /// `yyyy-MM-dd HH:mm:ss[.fraction]` (fraction with trailing zeros trimmed), optionally + /// followed by a `+HH:MM` offset suffix. + static std::string TimestampToString(int64_t micros_since_epoch, int32_t offset_seconds, + bool with_offset); + + /// Resolves the UTC offset (in seconds) of `zone_id` at the given instant. Supports fixed + /// offsets (`+08:00`, `-05:30`, optionally prefixed with `UTC`/`GMT`), `Z`, `UTC`, `GMT`, + /// and IANA region ids such as `Asia/Shanghai`. + static Result GetZoneOffsetSeconds(const std::string& zone_id, + int64_t micros_since_epoch); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_path_segment.cpp b/src/paimon/common/data/variant/variant_path_segment.cpp new file mode 100644 index 000000000..61885e6d9 --- /dev/null +++ b/src/paimon/common/data/variant/variant_path_segment.cpp @@ -0,0 +1,115 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_path_segment.h" + +#include +#include + +#include "fmt/format.h" + +namespace paimon { + +namespace { + +// Matches `[123]` at the beginning of `remaining` and consumes it. +bool TryParseIndex(std::string_view* remaining, int32_t* index) { + std::string_view s = *remaining; + if (s.empty() || s[0] != '[') { + return false; + } + size_t i = 1; + int64_t value = 0; + size_t digits = 0; + while (i < s.size() && s[i] >= '0' && s[i] <= '9') { + value = value * 10 + (s[i] - '0'); + if (value > std::numeric_limits::max()) { + return false; + } + ++i; + ++digits; + } + if (digits == 0 || i >= s.size() || s[i] != ']') { + return false; + } + *index = static_cast(value); + remaining->remove_prefix(i + 1); + return true; +} + +// Matches `['key']` or `["key"]` at the beginning of `remaining` and consumes it. +bool TryParseQuotedKey(std::string_view* remaining, std::string* key) { + std::string_view s = *remaining; + if (s.size() < 4 || s[0] != '[' || (s[1] != '\'' && s[1] != '"')) { + return false; + } + char quote = s[1]; + size_t end = s.find(quote, 2); + if (end == std::string_view::npos || end == 2 || end + 1 >= s.size() || s[end + 1] != ']') { + return false; + } + key->assign(s.substr(2, end - 2)); + remaining->remove_prefix(end + 2); + return true; +} + +// Matches `.key` (one or more characters that are neither `.` nor `[`) at the beginning of +// `remaining` and consumes it. +bool TryParseDotKey(std::string_view* remaining, std::string* key) { + std::string_view s = *remaining; + if (s.size() < 2 || s[0] != '.') { + return false; + } + size_t end = 1; + while (end < s.size() && s[end] != '.' && s[end] != '[') { + ++end; + } + if (end == 1) { + return false; + } + key->assign(s.substr(1, end - 1)); + remaining->remove_prefix(end); + return true; +} + +} // namespace + +Result> VariantPathSegment::Parse(const std::string& path) { + // Mirrors the Java parser: the root `$` is located with a find, so segments are parsed + // after the FIRST `$` and any prefix before it is ignored. + size_t root = path.find('$'); + if (path.empty() || root == std::string::npos) { + return Status::Invalid(fmt::format("Invalid path: {}", path)); + } + std::string_view remaining = std::string_view(path).substr(root + 1); + std::vector segments; + while (!remaining.empty()) { + int32_t index = 0; + if (TryParseIndex(&remaining, &index)) { + segments.push_back(ArrayExtraction(index)); + continue; + } + std::string key; + if (TryParseDotKey(&remaining, &key) || TryParseQuotedKey(&remaining, &key)) { + segments.push_back(ObjectExtraction(std::move(key))); + continue; + } + return Status::Invalid(fmt::format("Invalid path: {}", path)); + } + return segments; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_path_segment.h b/src/paimon/common/data/variant/variant_path_segment.h new file mode 100644 index 000000000..48bdbb1bb --- /dev/null +++ b/src/paimon/common/data/variant/variant_path_segment.h @@ -0,0 +1,57 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/result.h" + +namespace paimon { + +/// A path segment for variant get, representing either an object key access or an array index +/// access. +struct VariantPathSegment { + enum class Kind { kObjectExtraction, kArrayExtraction }; + + Kind kind; + /// The object key when `kind` is `kObjectExtraction`. + std::string key; + /// The array index when `kind` is `kArrayExtraction`. + int32_t index = 0; + + static VariantPathSegment ObjectExtraction(std::string key) { + VariantPathSegment segment; + segment.kind = Kind::kObjectExtraction; + segment.key = std::move(key); + return segment; + } + + static VariantPathSegment ArrayExtraction(int32_t index) { + VariantPathSegment segment; + segment.kind = Kind::kArrayExtraction; + segment.index = index; + return segment; + } + + /// Parses a path starting with `$`. Supported segments after the root are `.key`, `['key']`, + /// `["key"]` and `[index]`, e.g. `$.user.addresses[0]['city']`. + static Result> Parse(const std::string& path); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_reassembler.cpp b/src/paimon/common/data/variant/variant_reassembler.cpp new file mode 100644 index 000000000..83dd360af --- /dev/null +++ b/src/paimon/common/data/variant/variant_reassembler.cpp @@ -0,0 +1,260 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#include "paimon/common/data/variant/variant_reassembler.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon { + +namespace { + +// A row view over a shredded struct array, the Arrow analog of the Java `ShreddedRow`. +struct Cursor { + const arrow::StructArray* array; + int64_t row; + + bool IsNullAt(int32_t field) const { + return array->field(field)->IsNull(row); + } + + std::string_view GetBinary(int32_t field) const { + return static_cast(*array->field(field)).GetView(row); + } +}; + +Status Rebuild(const Cursor& cursor, std::string_view metadata, const VariantSchema& schema, + const std::shared_ptr& pool, VariantBuilder* builder); + +Status RebuildTypedScalar(const Cursor& cursor, const VariantSchema& schema, + VariantBuilder* builder) { + int32_t typed_idx = schema.typed_idx; + const arrow::Array& typed_array = *cursor.array->field(typed_idx); + int64_t row = cursor.row; + const VariantSchema::ScalarType& scalar = schema.scalar_schema.value(); + switch (scalar.kind) { + case VariantSchema::ScalarKind::kString: + return builder->AppendString( + static_cast(typed_array).GetView(row)); + case VariantSchema::ScalarKind::kByte: + return builder->AppendLong( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kShort: + return builder->AppendLong( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kInt: + return builder->AppendLong( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kLong: + return builder->AppendLong( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kFloat: + return builder->AppendFloat( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kDouble: + return builder->AppendDouble( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kBoolean: + return builder->AppendBoolean( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kBinary: + return builder->AppendBinary( + static_cast(typed_array).GetView(row)); + case VariantSchema::ScalarKind::kDecimal: { + const auto& decimal_array = static_cast(typed_array); + arrow::Decimal128 value(decimal_array.GetValue(row)); + VariantDecimal decimal; + decimal.unscaled = (static_cast<__int128_t>(value.high_bits()) << 64) | + static_cast<__int128_t>(static_cast<__uint128_t>(value.low_bits())); + decimal.scale = scalar.scale; + return builder->AppendDecimal(decimal); + } + case VariantSchema::ScalarKind::kDate: + return builder->AppendDate( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kTimestampLtz: + return builder->AppendTimestamp( + static_cast(typed_array).Value(row)); + case VariantSchema::ScalarKind::kTimestampNtz: + return builder->AppendTimestampNtz( + static_cast(typed_array).Value(row)); + default: + return Status::NotImplemented("unsupported variant scalar kind in reassembly"); + } +} + +// Rebuilds a variant value from the shredded data according to the reconstruction algorithm in +// the parquet-format VariantShredding.md specification, appending the result to `builder`. +Status Rebuild(const Cursor& cursor, std::string_view metadata, const VariantSchema& schema, + const std::shared_ptr& pool, VariantBuilder* builder) { + int32_t typed_idx = schema.typed_idx; + int32_t variant_idx = schema.variant_idx; + if (typed_idx >= 0 && !cursor.IsNullAt(typed_idx)) { + if (schema.scalar_schema.has_value()) { + return RebuildTypedScalar(cursor, schema, builder); + } else if (schema.array_schema != nullptr) { + const auto& list_array = + static_cast(*cursor.array->field(typed_idx)); + const auto& element_array = + static_cast(*list_array.values()); + int64_t element_start = list_array.value_offset(cursor.row); + int64_t element_end = list_array.value_offset(cursor.row + 1); + int32_t start = builder->GetWritePos(); + std::vector offsets; + offsets.reserve(element_end - element_start); + for (int64_t i = element_start; i < element_end; ++i) { + offsets.push_back(builder->GetWritePos() - start); + PAIMON_RETURN_NOT_OK(Rebuild(Cursor{&element_array, i}, metadata, + *schema.array_schema, pool, builder)); + } + return builder->FinishWritingArray(start, offsets); + } else { + const auto& object_array = + static_cast(*cursor.array->field(typed_idx)); + Cursor object_cursor{&object_array, cursor.row}; + std::vector fields; + int32_t start = builder->GetWritePos(); + for (size_t field_idx = 0; field_idx < schema.object_schema.size(); ++field_idx) { + // Shredded fields must not be null. + if (object_cursor.IsNullAt(static_cast(field_idx))) { + return VariantBinaryUtil::MalformedVariant(); + } + const std::string& field_name = schema.object_schema[field_idx].name; + const VariantSchema& field_schema = *schema.object_schema[field_idx].schema; + const auto& field_array = static_cast( + *object_array.field(static_cast(field_idx))); + Cursor field_cursor{&field_array, cursor.row}; + // If the field doesn't have a non-null `typed_value` or `value`, it is missing. + if ((field_schema.typed_idx >= 0 && + !field_cursor.IsNullAt(field_schema.typed_idx)) || + (field_schema.variant_idx >= 0 && + !field_cursor.IsNullAt(field_schema.variant_idx))) { + int32_t id = builder->AddKey(field_name); + fields.emplace_back(field_name, id, builder->GetWritePos() - start); + PAIMON_RETURN_NOT_OK( + Rebuild(field_cursor, metadata, field_schema, pool, builder)); + } + } + if (variant_idx >= 0 && !cursor.IsNullAt(variant_idx)) { + // Add the leftover fields in the variant binary. + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr leftover, + GenericVariant::Create(cursor.GetBinary(variant_idx), metadata, pool)); + PAIMON_ASSIGN_OR_RAISE(VariantValueType leftover_type, leftover->GetType()); + if (leftover_type != VariantValueType::kObject) { + return VariantBinaryUtil::MalformedVariant(); + } + PAIMON_ASSIGN_OR_RAISE(int32_t leftover_size, leftover->ObjectSize()); + for (int32_t i = 0; i < leftover_size; ++i) { + PAIMON_ASSIGN_OR_RAISE(auto field, leftover->GetFieldAtIndex(i)); + if (!field.has_value()) { + return VariantBinaryUtil::MalformedVariant(); + } + // `value` must not contain any shredded field. + if (schema.object_schema_map.count(field->key) > 0) { + return VariantBinaryUtil::MalformedVariant(); + } + int32_t id = builder->AddKey(field->key); + fields.emplace_back(field->key, id, builder->GetWritePos() - start); + PAIMON_RETURN_NOT_OK(builder->AppendVariant(*field->value)); + } + } + return builder->FinishWritingObject(start, &fields); + } + } else if (variant_idx >= 0 && !cursor.IsNullAt(variant_idx)) { + // `typed_value` doesn't exist or is null. Read from `value`. + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr variant, + GenericVariant::Create(cursor.GetBinary(variant_idx), metadata, pool)); + return builder->AppendVariant(*variant); + } else { + // The variant is missing in a context where it must be present; the data is invalid. + return VariantBinaryUtil::MalformedVariant(); + } +} + +} // namespace + +Status VariantReassembler::RebuildValue(const arrow::StructArray& shredded, int64_t row, + std::string_view metadata, const VariantSchema& schema, + const std::shared_ptr& pool, + VariantBuilder* builder) { + return Rebuild(Cursor{&shredded, row}, metadata, schema, pool, builder); +} + +Result> VariantReassembler::AssembleVariantArray( + const std::shared_ptr& shredded, + const std::shared_ptr& schema, arrow::MemoryPool* pool) { + if (schema->top_level_metadata_idx < 0) { + return VariantBinaryUtil::MalformedVariant(); + } + std::shared_ptr paimon_pool = GetDefaultPool(); + auto output_type = VariantTypeUtils::UnshreddedStructType(); + std::unique_ptr output_builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder(pool, output_type, &output_builder)); + auto* struct_builder = static_cast(output_builder.get()); + auto* value_builder = static_cast(struct_builder->field_builder(0)); + auto* metadata_builder = static_cast(struct_builder->field_builder(1)); + + bool unshredded = schema->IsUnshredded(); + for (int64_t row = 0; row < shredded->length(); ++row) { + if (shredded->IsNull(row)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->AppendNull()); + continue; + } + Cursor cursor{shredded.get(), row}; + if (cursor.IsNullAt(schema->top_level_metadata_idx)) { + return VariantBinaryUtil::MalformedVariant(); + } + std::string_view metadata = cursor.GetBinary(schema->top_level_metadata_idx); + PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Append()); + if (unshredded) { + // Rebuilding is unnecessary for unshredded variants. + if (cursor.IsNullAt(schema->variant_idx)) { + return VariantBinaryUtil::MalformedVariant(); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW( + value_builder->Append(cursor.GetBinary(schema->variant_idx))); + PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata_builder->Append(metadata)); + } else { + VariantBuilder builder(/*allow_duplicate_keys=*/false); + PAIMON_RETURN_NOT_OK(Rebuild(cursor, metadata, *schema, paimon_pool, &builder)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + builder.Build(paimon_pool)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Append(variant->RawValue())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata_builder->Append(variant->Metadata())); + } + } + std::shared_ptr result; + PAIMON_RETURN_NOT_OK_FROM_ARROW(output_builder->Finish(&result)); + return result; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_reassembler.h b/src/paimon/common/data/variant/variant_reassembler.h new file mode 100644 index 000000000..bb65db659 --- /dev/null +++ b/src/paimon/common/data/variant/variant_reassembler.h @@ -0,0 +1,68 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#pragma once + +#include +#include + +#include "paimon/common/data/variant/variant_schema.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class MemoryPool; +class StructArray; +} // namespace arrow + +namespace paimon { + +class VariantBuilder; + +/// Reassembles shredded variant columns back into the unshredded +/// `struct` representation, implementing the reconstruction +/// algorithm of the parquet-format VariantShredding.md specification (mirroring the Java +/// `ShreddingUtils.rebuild`). +class VariantReassembler { + public: + VariantReassembler() = delete; + ~VariantReassembler() = delete; + + /// Reassembles a shredded variant column into a `struct` array. + /// + /// @param shredded The physical shredded array read from the file. + /// @param schema The shredding schema of the column + /// (`VariantShreddingUtils::BuildVariantSchema` of the file type). + /// @param pool The Arrow memory pool used for the output. + /// @return The unshredded variant array (`VariantTypeUtils::UnshreddedStructType`). + static Result> AssembleVariantArray( + const std::shared_ptr& shredded, + const std::shared_ptr& schema, arrow::MemoryPool* pool); + + /// Rebuilds the variant value at `row` of a shredded (sub-)struct into `builder`, following + /// the same reconstruction algorithm. `schema` describes `shredded`, which may be any level + /// of the shredded tree; `metadata` is the column's top-level metadata binary. + static Status RebuildValue(const arrow::StructArray& shredded, int64_t row, + std::string_view metadata, const VariantSchema& schema, + const std::shared_ptr& pool, VariantBuilder* builder); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_schema.h b/src/paimon/common/data/variant/variant_schema.h new file mode 100644 index 000000000..e6a2d2192 --- /dev/null +++ b/src/paimon/common/data/variant/variant_schema.h @@ -0,0 +1,95 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace paimon { + +/// Defines a valid shredding schema, as described in the parquet-format VariantShredding.md +/// specification. A shredding schema contains a `value` and an optional `typed_value` field. If a +/// `typed_value` is an array or struct, it recursively contains its own shredding schema for +/// elements and fields, respectively. The schema also contains a `metadata` field at the top +/// level, but not in recursively shredded fields. +class VariantSchema { + public: + enum class ScalarKind { + kBoolean, + kByte, + kShort, + kInt, + kLong, + kFloat, + kDouble, + kString, + kBinary, + kDecimal, + kDate, + kTimestampLtz, + kTimestampNtz, + kUuid, + }; + + struct ScalarType { + ScalarKind kind; + // Only meaningful when `kind` is `kDecimal`. + int32_t precision = 0; + int32_t scale = 0; + }; + + /// Represents one field of an object in the shredding schema. + struct ObjectField { + std::string name; + std::shared_ptr schema; + }; + + /// The index of the typed_value, value, and metadata fields in the schema, respectively. If a + /// given field is not in the schema, its value must be set to -1 to indicate that it is + /// invalid. The indices of valid fields are contiguous and start from 0. + int32_t typed_idx = -1; + int32_t variant_idx = -1; + /// Must be non-negative in the top-level schema, and -1 at all other nesting levels. + int32_t top_level_metadata_idx = -1; + /// The number of fields in the schema, i.e. a value between 1 and 3, depending on which of + /// value, typed_value and metadata are present. + int32_t num_fields = 0; + + /// Exactly one of the following describes typed_value (or none if there is no typed_value). + std::optional scalar_schema; + bool has_object_schema = false; + std::vector object_schema; + /// Fast lookup of object fields by name; values are indices into `object_schema`. + std::unordered_map object_schema_map; + std::shared_ptr array_schema; + + /// Whether the variant column is unshredded. The user is not required to do anything special, + /// but can have certain optimizations for unshredded variants. + bool IsUnshredded() const { + return top_level_metadata_idx >= 0 && variant_idx >= 0 && typed_idx < 0; + } +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_batch_converter.cpp b/src/paimon/common/data/variant/variant_shredding_batch_converter.cpp new file mode 100644 index 000000000..9e70927b8 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_batch_converter.cpp @@ -0,0 +1,106 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_shredding_batch_converter.h" + +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_shredding_writer.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" + +namespace paimon { + +VariantShreddingBatchConverter::VariantShreddingBatchConverter( + const std::shared_ptr& plan, const std::shared_ptr& pool) + : plan_(plan), pool_(pool), arrow_pool_(GetArrowPool(pool)) {} + +Result> VariantShreddingBatchConverter::Create( + const std::shared_ptr& plan, + const std::shared_ptr& pool) { + if (!plan) { + return Status::Invalid("variant shredding batch converter requires a write plan"); + } + return std::shared_ptr( + new VariantShreddingBatchConverter(plan, pool)); +} + +const std::shared_ptr& VariantShreddingBatchConverter::GetPhysicalSchema() const { + return plan_->PhysicalSchema(); +} + +Result> VariantShreddingBatchConverter::Convert( + ArrowArray* logical_batch) { + auto logical_struct_type = arrow::struct_(plan_->LogicalSchema()->fields()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_array, + arrow::ImportArray(logical_batch, logical_struct_type)); + const auto& logical_struct = std::static_pointer_cast(logical_array); + + arrow::ArrayVector physical_arrays = logical_struct->fields(); + const auto& physical_fields = plan_->PhysicalSchema()->fields(); + for (int32_t i = 0; i < logical_struct->num_fields(); ++i) { + const std::string& field_name = logical_struct->struct_type()->field(i)->name(); + auto schema_it = plan_->ColumnSchemas().find(field_name); + if (schema_it == plan_->ColumnSchemas().end()) { + continue; + } + const auto& physical_type = plan_->ColumnPhysicalTypes().at(field_name); + if (logical_struct->field(i)->type_id() != arrow::Type::STRUCT) { + return Status::Invalid(fmt::format( + "variant column {} is not a struct column", field_name)); + } + auto variant_column = + std::static_pointer_cast(logical_struct->field(i)); + if (variant_column->num_fields() != 2) { + return Status::Invalid(fmt::format( + "variant column {} is not a struct column", field_name)); + } + const auto& value_column = + std::static_pointer_cast(variant_column->field(0)); + const auto& metadata_column = + std::static_pointer_cast(variant_column->field(1)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + VariantShreddedColumnWriter::Create(schema_it->second, physical_type, + arrow_pool_.get())); + for (int64_t row = 0; row < variant_column->length(); ++row) { + if (variant_column->IsNull(row)) { + PAIMON_RETURN_NOT_OK(writer->AppendNull()); + continue; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + GenericVariant::Create(value_column->GetView(row), + metadata_column->GetView(row), pool_)); + PAIMON_RETURN_NOT_OK(writer->Append(*variant)); + } + PAIMON_ASSIGN_OR_RAISE(physical_arrays[i], writer->Finish()); + } + + arrow::FieldVector physical_struct_fields(physical_fields.begin(), physical_fields.end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr physical_struct, + arrow::StructArray::Make(physical_arrays, physical_struct_fields)); + auto result = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*physical_struct, result.get())); + return result; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_batch_converter.h b/src/paimon/common/data/variant/variant_shredding_batch_converter.h new file mode 100644 index 000000000..a39c0e866 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_batch_converter.h @@ -0,0 +1,62 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/shredding_batch_converter.h" +#include "paimon/common/data/variant/variant_shredding_write_plan.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +struct ArrowArray; + +namespace arrow { +class MemoryPool; +class Schema; +} // namespace arrow + +namespace paimon { + +/// Converts logical batches containing VARIANT columns into physical batches where each planned +/// variant column is replaced by its shredded struct representation. Unplanned columns are +/// passed through unchanged. +class VariantShreddingBatchConverter : public ShreddingBatchConverter { + public: + static Result> Create( + const std::shared_ptr& plan, + const std::shared_ptr& pool); + + /// The physical schema produced by this converter. + const std::shared_ptr& GetPhysicalSchema() const override; + + /// Converts a logical batch to a physical batch. + /// @param logical_batch Input ArrowArray (C ABI) with the logical schema. Consumed on + /// success. + /// @return Owned physical ArrowArray (C ABI) with the physical schema. + Result> Convert(ArrowArray* logical_batch) override; + + private: + VariantShreddingBatchConverter(const std::shared_ptr& plan, + const std::shared_ptr& pool); + + std::shared_ptr plan_; + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_read_plan_factory.cpp b/src/paimon/common/data/variant/variant_shredding_read_plan_factory.cpp new file mode 100644 index 000000000..910916e33 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_read_plan_factory.cpp @@ -0,0 +1,341 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_shredding_read_plan_factory.h" + +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_access_utils.h" +#include "paimon/common/data/variant/variant_binary_util.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_get.h" +#include "paimon/common/data/variant/variant_reassembler.h" +#include "paimon/common/data/variant/variant_schema.h" +#include "paimon/common/data/variant/variant_shredding_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" + +namespace paimon { + +namespace { + +/// Reassembles the full variant of a shredded file column back into +/// `struct` (a plain VARIANT read). +class FullVariantColumnReadPlan : public ShreddingColumnReadPlan { + public: + FullVariantColumnReadPlan(std::shared_ptr logical_field, + std::shared_ptr physical_field, + std::shared_ptr schema) + : logical_field_(std::move(logical_field)), + physical_field_(std::move(physical_field)), + schema_(std::move(schema)) {} + + const std::shared_ptr& LogicalField() const override { + return logical_field_; + } + + const std::shared_ptr& PhysicalField() const override { + return physical_field_; + } + + Result> Assemble(const std::shared_ptr& physical, + arrow::MemoryPool* pool) const override { + if (physical->type_id() != arrow::Type::STRUCT) { + return Status::Invalid(fmt::format("cannot cast shredded variant field {} to a struct", + physical_field_->name())); + } + auto physical_struct = std::static_pointer_cast(physical); + return VariantReassembler::AssembleVariantArray(physical_struct, schema_, pool); + } + + private: + std::shared_ptr logical_field_; + std::shared_ptr physical_field_; + std::shared_ptr schema_; +}; + +/// One access path segment resolved against the shredded schema of one file. +struct ResolvedSegment { + VariantPathSegment raw; + bool is_object = false; + // The `typed_value` index at this level, or -1 when the path leaves the shredded schema + // here and continues inside the `value` binary. + int32_t typed_idx = -1; + // The object field index inside the typed object, or the array element index. + int32_t extraction_idx = -1; +}; + +struct ResolvedSpec { + VariantAccessSpec spec; + std::vector segments; +}; + +ResolvedSpec ResolveSpec(const VariantAccessSpec& spec, const VariantSchema* root) { + ResolvedSpec resolved; + resolved.spec = spec; + const VariantSchema* schema = root; + for (const auto& segment : spec.segments) { + ResolvedSegment r; + r.raw = segment; + if (segment.kind == VariantPathSegment::Kind::kObjectExtraction) { + r.is_object = true; + if (schema != nullptr && !schema->object_schema.empty()) { + auto it = schema->object_schema_map.find(segment.key); + if (it != schema->object_schema_map.end()) { + r.typed_idx = schema->typed_idx; + r.extraction_idx = it->second; + schema = schema->object_schema[it->second].schema.get(); + } else { + schema = nullptr; + } + } else { + schema = nullptr; + } + } else { + if (schema != nullptr && schema->array_schema != nullptr) { + r.typed_idx = schema->typed_idx; + r.extraction_idx = segment.index; + schema = schema->array_schema.get(); + } else { + schema = nullptr; + } + } + resolved.segments.push_back(std::move(r)); + } + return resolved; +} + +/// Extracts the paths described by a variant-access projection, reading typed sub-columns +/// directly and falling back to the `value` binary where the path is not shredded. +class VariantAccessColumnReadPlan : public ShreddingColumnReadPlan { + public: + VariantAccessColumnReadPlan(std::shared_ptr logical_field, + std::shared_ptr physical_field, + std::shared_ptr schema, + std::vector specs, std::shared_ptr pool) + : logical_field_(std::move(logical_field)), + physical_field_(std::move(physical_field)), + schema_(std::move(schema)), + specs_(std::move(specs)), + pool_(std::move(pool)) {} + + const std::shared_ptr& LogicalField() const override { + return logical_field_; + } + + const std::shared_ptr& PhysicalField() const override { + return physical_field_; + } + + Result> Assemble(const std::shared_ptr& physical, + arrow::MemoryPool* pool) const override { + if (physical->type_id() != arrow::Type::STRUCT) { + return Status::Invalid(fmt::format("cannot cast shredded variant field {} to a struct", + physical_field_->name())); + } + const auto& physical_struct = static_cast(*physical); + std::unique_ptr builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder(pool, logical_field_->type(), &builder)); + auto* struct_builder = static_cast(builder.get()); + for (int64_t row = 0; row < physical_struct.length(); ++row) { + if (physical_struct.IsNull(row)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->AppendNull()); + continue; + } + if (physical_struct.field(schema_->top_level_metadata_idx)->IsNull(row)) { + return VariantBinaryUtil::MalformedVariant(); + } + std::string_view metadata = static_cast( + *physical_struct.field(schema_->top_level_metadata_idx)) + .GetView(row); + PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Append()); + for (size_t i = 0; i < specs_.size(); ++i) { + PAIMON_RETURN_NOT_OK( + ExtractField(physical_struct, row, metadata, specs_[i], + struct_builder->field_builder(static_cast(i)))); + } + } + std::shared_ptr result; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Finish(&result)); + return result; + } + + private: + Status ExtractField(const arrow::StructArray& root, int64_t root_row, std::string_view metadata, + const ResolvedSpec& resolved, arrow::ArrayBuilder* builder) const { + const arrow::StructArray* current = &root; + int64_t row = root_row; + const VariantSchema* schema = schema_.get(); + size_t segment_idx = 0; + while (segment_idx < resolved.segments.size()) { + const ResolvedSegment& segment = resolved.segments[segment_idx]; + if (segment.typed_idx < 0) { + // The path leaves the shredded schema here; walk the remaining raw path inside + // the `value` binary. + return ExtractFromBinary(*current, row, metadata, resolved, segment_idx, schema, + builder); + } + if (current->field(segment.typed_idx)->IsNull(row)) { + return ToPaimonStatus(builder->AppendNull()); + } + if (segment.is_object) { + const auto& object_array = + static_cast(*current->field(segment.typed_idx)); + const auto& field_array = static_cast( + *object_array.field(segment.extraction_idx)); + if (field_array.IsNull(row)) { + // Shredded object fields must not be null. + return VariantBinaryUtil::MalformedVariant(); + } + schema = schema->object_schema[segment.extraction_idx].schema.get(); + current = &field_array; + // A field is missing when neither its typed_value nor its value is present. + bool typed_present = + schema->typed_idx >= 0 && !current->field(schema->typed_idx)->IsNull(row); + bool variant_present = + schema->variant_idx >= 0 && !current->field(schema->variant_idx)->IsNull(row); + if (!typed_present && !variant_present) { + return ToPaimonStatus(builder->AppendNull()); + } + } else { + const auto& list_array = + static_cast(*current->field(segment.typed_idx)); + if (segment.extraction_idx >= list_array.value_length(row)) { + return ToPaimonStatus(builder->AppendNull()); + } + int64_t element_row = list_array.value_offset(row) + segment.extraction_idx; + const auto& element_array = + static_cast(*list_array.values()); + if (element_array.IsNull(element_row)) { + // Shredded array elements must not be null. + return VariantBinaryUtil::MalformedVariant(); + } + schema = schema->array_schema.get(); + current = &element_array; + row = element_row; + } + ++segment_idx; + } + + // The terminal position: rebuild the (sub-)variant and cast it to the target type. + if (schema->typed_idx >= 0 && !current->field(schema->typed_idx)->IsNull(row)) { + VariantBuilder variant_builder(/*allow_duplicate_keys=*/false); + PAIMON_RETURN_NOT_OK(VariantReassembler::RebuildValue(*current, row, metadata, *schema, + pool_, &variant_builder)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + variant_builder.Build(pool_)); + return VariantGetExecutor::CastToBuilder(variant, resolved.spec.target_field, builder, + resolved.spec.cast_args, pool_); + } + if (schema->variant_idx >= 0 && !current->field(schema->variant_idx)->IsNull(row)) { + std::string_view value = + static_cast(*current->field(schema->variant_idx)) + .GetView(row); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + GenericVariant::Create(value, metadata, pool_)); + return VariantGetExecutor::CastToBuilder(variant, resolved.spec.target_field, builder, + resolved.spec.cast_args, pool_); + } + return VariantBinaryUtil::MalformedVariant(); + } + + Status ExtractFromBinary(const arrow::StructArray& current, int64_t row, + std::string_view metadata, const ResolvedSpec& resolved, + size_t segment_idx, const VariantSchema* schema, + arrow::ArrayBuilder* builder) const { + if (schema->variant_idx < 0 || current.field(schema->variant_idx)->IsNull(row)) { + return ToPaimonStatus(builder->AppendNull()); + } + std::string_view value = + static_cast(*current.field(schema->variant_idx)) + .GetView(row); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + GenericVariant::Create(value, metadata, pool_)); + for (; segment_idx < resolved.segments.size() && variant != nullptr; ++segment_idx) { + const VariantPathSegment& raw = resolved.segments[segment_idx].raw; + PAIMON_ASSIGN_OR_RAISE(VariantValueType type, variant->GetType()); + if (raw.kind == VariantPathSegment::Kind::kObjectExtraction && + type == VariantValueType::kObject) { + PAIMON_ASSIGN_OR_RAISE(variant, variant->GetFieldByKey(raw.key)); + } else if (raw.kind == VariantPathSegment::Kind::kArrayExtraction && + type == VariantValueType::kArray) { + PAIMON_ASSIGN_OR_RAISE(variant, variant->GetElementAtIndex(raw.index)); + } else { + variant = nullptr; + } + } + return VariantGetExecutor::CastToBuilder(variant, resolved.spec.target_field, builder, + resolved.spec.cast_args, pool_); + } + + std::shared_ptr logical_field_; + std::shared_ptr physical_field_; + std::shared_ptr schema_; + std::vector specs_; + std::shared_ptr pool_; +}; + +} // namespace + +Result>> +VariantShreddingReadPlanFactory::CreateReadPlans(const std::shared_ptr& read_schema, + const std::shared_ptr& file_schema, + const std::shared_ptr& pool) { + std::map> plans; + for (const auto& read_field : read_schema->fields()) { + if (!VariantTypeUtils::IsVariantField(read_field) && + !VariantAccessUtils::IsVariantAccessType(read_field->type())) { + continue; + } + auto file_field = file_schema->GetFieldByName(read_field->name()); + if (file_field == nullptr) { + // The column is absent in the file (schema evolution); it is filled with nulls + // downstream. + continue; + } + bool file_shredded = VariantShreddingUtils::IsShreddedFileType(file_field->type()); + if (VariantAccessUtils::IsVariantAccessType(read_field->type())) { + PAIMON_ASSIGN_OR_RAISE(std::vector specs, + VariantAccessUtils::ParseAccessSpecs(read_field)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr physical_field, + VariantAccessUtils::ClipShreddedFileField(specs, file_field)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr schema, + VariantShreddingUtils::BuildVariantSchema(physical_field->type())); + std::vector resolved; + resolved.reserve(specs.size()); + for (const auto& spec : specs) { + resolved.push_back(ResolveSpec(spec, schema.get())); + } + plans.emplace(read_field->name(), std::make_shared( + read_field, physical_field, std::move(schema), + std::move(resolved), pool)); + } else if (file_shredded) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr schema, + VariantShreddingUtils::BuildVariantSchema(file_field->type())); + plans.emplace(read_field->name(), std::make_shared( + read_field, file_field, std::move(schema))); + } + // A plain VARIANT read of an unshredded file needs no plan. + } + return plans; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_read_plan_factory.h b/src/paimon/common/data/variant/variant_shredding_read_plan_factory.h new file mode 100644 index 000000000..993296d0d --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_read_plan_factory.h @@ -0,0 +1,51 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/shredding/shredding_read_plan.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +/// Builds per-column read plans for VARIANT columns: +/// - a plain VARIANT read of a shredded file reassembles the full variant; +/// - a variant-access projection (struct whose children carry `__VARIANT_METADATA` +/// descriptions) extracts the described paths, reading only the required shredded +/// sub-columns from a shredded file (or the binary from an unshredded file). +class VariantShreddingReadPlanFactory { + public: + VariantShreddingReadPlanFactory() = delete; + ~VariantShreddingReadPlanFactory() = delete; + + /// Creates the per-column read plans for the variant columns of `read_schema` against + /// `file_schema`; the map is empty when no plan applies (no shredded file column and no + /// variant-access projection). + static Result>> CreateReadPlans( + const std::shared_ptr& read_schema, + const std::shared_ptr& file_schema, const std::shared_ptr& pool); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_test.cpp b/src/paimon/common/data/variant/variant_shredding_test.cpp new file mode 100644 index 000000000..22214fd19 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_test.cpp @@ -0,0 +1,248 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/data/variant/variant_reassembler.h" +#include "paimon/common/data/variant/variant_schema.h" +#include "paimon/common/data/variant/variant_shredding_utils.h" +#include "paimon/common/data/variant/variant_shredding_writer.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class VariantShreddingTest : public ::testing::Test { + public: + // Shreds the given JSON documents (nullptr = null variant) with the given shredding type, + // asserts the reassembled variants render back to the same JSON, and returns the shredded + // array for structural checks. + std::shared_ptr RoundTrip( + const std::shared_ptr& shredding_type, + const std::vector& jsons) { + auto physical_result = VariantShreddingUtils::VariantShreddingSchema(shredding_type); + EXPECT_TRUE(physical_result.ok()) << physical_result.status().ToString(); + auto physical = physical_result.value(); + auto schema_result = VariantShreddingUtils::BuildVariantSchema(physical); + EXPECT_TRUE(schema_result.ok()) << schema_result.status().ToString(); + auto schema = schema_result.value(); + + auto writer_result = + VariantShreddedColumnWriter::Create(schema, physical, arrow::default_memory_pool()); + EXPECT_TRUE(writer_result.ok()) << writer_result.status().ToString(); + auto& writer = writer_result.value(); + std::vector expected_jsons; + for (const char* json : jsons) { + if (json == nullptr) { + EXPECT_OK(writer->AppendNull()); + expected_jsons.emplace_back(); + continue; + } + auto variant = GenericVariant::FromJson(json, pool_); + EXPECT_TRUE(variant.ok()) << variant.status().ToString(); + auto to_json = variant.value()->ToJson(); + EXPECT_TRUE(to_json.ok()); + expected_jsons.push_back(to_json.value()); + EXPECT_OK(writer->Append(*variant.value())); + } + auto shredded_result = writer->Finish(); + EXPECT_TRUE(shredded_result.ok()) << shredded_result.status().ToString(); + auto shredded = std::static_pointer_cast(shredded_result.value()); + + auto assembled_result = VariantReassembler::AssembleVariantArray( + shredded, schema, arrow::default_memory_pool()); + EXPECT_TRUE(assembled_result.ok()) << assembled_result.status().ToString(); + auto assembled = std::static_pointer_cast(assembled_result.value()); + EXPECT_EQ(assembled->length(), static_cast(jsons.size())); + auto value_column = std::static_pointer_cast(assembled->field(0)); + auto metadata_column = std::static_pointer_cast(assembled->field(1)); + for (size_t i = 0; i < jsons.size(); ++i) { + SCOPED_TRACE("row " + std::to_string(i)); + if (jsons[i] == nullptr) { + EXPECT_TRUE(assembled->IsNull(i)); + continue; + } + EXPECT_FALSE(assembled->IsNull(i)); + auto variant = GenericVariant::Create(value_column->GetView(i), + metadata_column->GetView(i), pool_); + EXPECT_TRUE(variant.ok()) << variant.status().ToString(); + auto actual_json = variant.value()->ToJson(); + EXPECT_TRUE(actual_json.ok()) << actual_json.status().ToString(); + EXPECT_EQ(actual_json.value(), expected_jsons[i]); + } + return shredded; + } + + protected: + std::shared_ptr pool_ = GetDefaultPool(); +}; + +TEST_F(VariantShreddingTest, ShreddingSchemaShape) { + auto shredding_type = + arrow::struct_({arrow::field("a", arrow::int32()), arrow::field("b", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr physical, + VariantShreddingUtils::VariantShreddingSchema(shredding_type)); + // struct{metadata: binary not null, value: binary, typed_value: struct{a: + // struct{value, typed_value} not null, b: ... not null}} + auto expected = arrow::struct_( + {arrow::field("metadata", arrow::binary(), false), + arrow::field("value", arrow::binary(), true), + arrow::field( + "typed_value", + arrow::struct_( + {arrow::field("a", + arrow::struct_({arrow::field("value", arrow::binary(), true), + arrow::field("typed_value", arrow::int32(), true)}), + false), + arrow::field("b", + arrow::struct_({arrow::field("value", arrow::binary(), true), + arrow::field("typed_value", arrow::utf8(), true)}), + false)}), + true)}); + ASSERT_TRUE(physical->Equals(*expected)) << physical->ToString(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr schema, + VariantShreddingUtils::BuildVariantSchema(physical)); + ASSERT_EQ(schema->top_level_metadata_idx, 0); + ASSERT_EQ(schema->variant_idx, 1); + ASSERT_EQ(schema->typed_idx, 2); + ASSERT_TRUE(schema->has_object_schema); + ASSERT_EQ(schema->object_schema.size(), 2); + ASSERT_FALSE(schema->IsUnshredded()); + ASSERT_TRUE(VariantShreddingUtils::IsShreddedFileType(physical)); + ASSERT_FALSE(VariantShreddingUtils::IsShreddedFileType( + arrow::struct_({arrow::field("value", arrow::binary(), false), + arrow::field("metadata", arrow::binary(), false)}))); + + // Invalid shredding types are rejected. + ASSERT_NOK(VariantShreddingUtils::VariantShreddingSchema(arrow::date32())); + ASSERT_NOK( + VariantShreddingUtils::VariantShreddingSchema(arrow::map(arrow::utf8(), arrow::int32()))); +} + +TEST_F(VariantShreddingTest, ShredObject) { + // Mirrors the Java GenericVariantTest#testShredding scenarios. + auto variant_json = R"({"a": 1, "b": "hello"})"; + // Happy path: all fields shredded, no residual value. + { + auto shredded = RoundTrip( + arrow::struct_({arrow::field("a", arrow::int32()), arrow::field("b", arrow::utf8())}), + {variant_json}); + ASSERT_TRUE(shredded->field(1)->IsNull(0)); // top-level value (residual) is null + ASSERT_FALSE(shredded->field(2)->IsNull(0)); // typed_value is present + } + // Missing field "c" in the data: present in schema, both children null. + { + auto shredded = RoundTrip( + arrow::struct_({arrow::field("a", arrow::int32()), arrow::field("c", arrow::utf8()), + arrow::field("b", arrow::utf8())}), + {variant_json}); + auto typed = std::static_pointer_cast(shredded->field(2)); + auto c_group = std::static_pointer_cast(typed->field(1)); + ASSERT_FALSE(c_group->IsNull(0)); + ASSERT_TRUE(c_group->field(0)->IsNull(0)); + ASSERT_TRUE(c_group->field(1)->IsNull(0)); + } + // "a" is not present in the shredding schema: it goes to the residual value. + { + auto shredded = RoundTrip( + arrow::struct_({arrow::field("b", arrow::utf8()), arrow::field("c", arrow::utf8())}), + {variant_json}); + auto value_column = std::static_pointer_cast(shredded->field(1)); + ASSERT_FALSE(value_column->IsNull(0)); + // The residual must equal the standalone encoding of {"a": 1}. + ASSERT_OK_AND_ASSIGN(std::shared_ptr residual_expected, + GenericVariant::FromJson("{\"a\": 1}", pool_)); + ASSERT_OK_AND_ASSIGN(std::string_view residual_value, residual_expected->Value()); + ASSERT_EQ(value_column->GetView(0), residual_value); + } +} + +TEST_F(VariantShreddingTest, ShredAllTypes) { + // Mirrors the Java GenericVariantTest#testShreddingAllTypes. + const char* json = + "{\n" + " \"c1\": \"Hello, World!\",\n" + " \"c2\": 12345678901234,\n" + " \"c3\": 1.0123456789012345678901234567890123456789,\n" + " \"c4\": 100.99,\n" + " \"c5\": true,\n" + " \"c6\": null,\n" + " \"c7\": {\"street\" : \"Main St\",\"city\" : \"Hangzhou\"},\n" + " \"c8\": [1, 2]\n" + "}\n"; + auto shredding_type = arrow::struct_( + {arrow::field("c1", arrow::utf8()), arrow::field("c2", arrow::int64()), + arrow::field("c3", arrow::float64()), arrow::field("c4", arrow::decimal128(5, 2)), + arrow::field("c5", arrow::boolean()), arrow::field("c6", arrow::utf8()), + arrow::field("c7", arrow::struct_({arrow::field("street", arrow::utf8()), + arrow::field("city", arrow::utf8())})), + arrow::field("c8", arrow::list(arrow::int32()))}); + auto shredded = RoundTrip(shredding_type, {json, nullptr, json}); + + // c6 is a variant null: it stays in the field's value column ("00"), typed_value is null. + auto typed = std::static_pointer_cast(shredded->field(2)); + auto c6_group = std::static_pointer_cast(typed->field(5)); + ASSERT_FALSE(c6_group->IsNull(0)); + auto c6_value = std::static_pointer_cast(c6_group->field(0)); + ASSERT_FALSE(c6_value->IsNull(0)); + ASSERT_EQ(c6_value->GetView(0), std::string_view("\x00", 1)); + ASSERT_TRUE(c6_group->field(1)->IsNull(0)); + + // Nothing was left over at the top level. + ASSERT_TRUE(shredded->field(1)->IsNull(0)); + + // No shredding at all: everything stays in the top-level value. + auto no_match_type = arrow::struct_({arrow::field("other", arrow::utf8())}); + auto unshredded = RoundTrip(no_match_type, {json}); + auto value_column = std::static_pointer_cast(unshredded->field(1)); + ASSERT_FALSE(value_column->IsNull(0)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected, + GenericVariant::FromJson(json, pool_)); + ASSERT_OK_AND_ASSIGN(std::string_view expected_value, expected->Value()); + ASSERT_EQ(value_column->GetView(0), expected_value); +} + +TEST_F(VariantShreddingTest, ShredScalarsAndMismatches) { + // Top-level scalar shredding with type mismatches falling back to the value column. + auto long_type = arrow::struct_({arrow::field("x", arrow::int64())}); + RoundTrip(long_type, + {"{\"x\": 5}", R"({"x": "not a number"})", "{\"x\": 3.25}", "{\"x\": [1]}", "{}"}); + // Decimal rescale: 100.99 fits decimal(9, 4) exactly (allowNumericScaleChanges). + auto decimal_type = arrow::struct_({arrow::field("x", arrow::decimal128(9, 4))}); + RoundTrip(decimal_type, + {"{\"x\": 100.99}", "{\"x\": 42}", "{\"x\": 0.123456789}", "{\"x\": 100.00}"}); + // A 38-digit unscaled value cannot rescale to scale 1 without overflowing the 128-bit + // decimal; it must fall back to the value column instead of being written corrupted. + auto wide_decimal_type = arrow::struct_({arrow::field("x", arrow::decimal128(38, 1))}); + RoundTrip(wide_decimal_type, + {"{\"x\": 99999999999999999999999999999999999999}", "{\"x\": 1.5}"}); + // Integer target from decimal that is numerically integral. + auto int_type = arrow::struct_({arrow::field("x", arrow::int32())}); + RoundTrip(int_type, {"{\"x\": 5.0}", "{\"x\": 5.5}", "{\"x\": 123456789012345678}"}); + auto array_type = arrow::struct_( + {arrow::field("arr", arrow::list(arrow::struct_({arrow::field("k", arrow::utf8())})))}); + RoundTrip(array_type, {R"({"arr": [{"k": "v1"}, {"k": "v2", "extra": 1}, {"other": 2}]})", + "{\"arr\": [1, 2]}", R"({"arr": {"k": "v"}})"}); +} + +} // namespace paimon::test diff --git a/src/paimon/common/data/variant/variant_shredding_utils.cpp b/src/paimon/common/data/variant/variant_shredding_utils.cpp new file mode 100644 index 000000000..1cef5dbc2 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_utils.cpp @@ -0,0 +1,291 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#include "paimon/common/data/variant/variant_shredding_utils.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/data/variant/variant_type_utils.h" + +namespace paimon { + +namespace { + +Status InvalidVariantShreddingSchema(const std::shared_ptr& type) { + return Status::Invalid( + fmt::format("Invalid variant shredding schema: {}", type ? type->ToString() : "null")); +} + +// Mirrors the Java `PaimonShreddingUtils.variantShreddingSchema(dataType, isTopLevel, +// isObjectField)`. +Result> VariantShreddingSchemaImpl( + const std::shared_ptr& data_type, bool is_top_level, bool is_object_field) { + arrow::FieldVector fields; + if (is_top_level) { + fields.push_back(arrow::field(VariantDefs::kMetadataFieldName, arrow::binary(), + /*nullable=*/false)); + } + switch (data_type->id()) { + case arrow::Type::LIST: { + const auto& list_type = std::static_pointer_cast(data_type); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element_type, + VariantShreddingSchemaImpl(list_type->value_type(), + /*is_top_level=*/false, + /*is_object_field=*/false)); + fields.push_back( + arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true)); + fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName, + arrow::list(element_type), /*nullable=*/true)); + break; + } + case arrow::Type::STRUCT: { + // The field name level is always non-nullable: Variant null values are represented in + // the "value" column as "00", and missing values are represented by setting both + // "value" and "typed_value" to null. + const auto& struct_type = std::static_pointer_cast(data_type); + arrow::FieldVector shredded_fields; + for (const auto& field : struct_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr field_type, + VariantShreddingSchemaImpl(field->type(), + /*is_top_level=*/false, + /*is_object_field=*/true)); + shredded_fields.push_back( + arrow::field(field->name(), field_type, /*nullable=*/false)); + } + fields.push_back( + arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true)); + fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName, + arrow::struct_(shredded_fields), /*nullable=*/true)); + break; + } + case arrow::Type::NA: { + // `arrow::null()` denotes an untyped VARIANT leaf in shredding types. It doesn't + // need a typed column. If there is no typed column, value is required for array + // elements or top-level fields, but optional for objects (where a null represents a + // missing field). + fields.push_back(arrow::field(VariantDefs::kValueFieldName, arrow::binary(), + /*nullable=*/is_object_field)); + break; + } + case arrow::Type::STRING: + case arrow::Type::BOOL: + case arrow::Type::BINARY: + case arrow::Type::DECIMAL128: + case arrow::Type::INT8: + case arrow::Type::INT16: + case arrow::Type::INT32: + case arrow::Type::INT64: + case arrow::Type::FLOAT: + case arrow::Type::DOUBLE: { + fields.push_back( + arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true)); + fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName, data_type, + /*nullable=*/true)); + break; + } + default: + return InvalidVariantShreddingSchema(data_type); + } + return arrow::struct_(fields); +} + +Result> BuildVariantSchemaImpl( + const std::shared_ptr& type, bool top_level) { + if (type->id() != arrow::Type::STRUCT) { + return InvalidVariantShreddingSchema(type); + } + const auto& struct_type = std::static_pointer_cast(type); + // The struct must not be empty or contain duplicate field names. The latter is enforced in + // the loop below. + if (struct_type->num_fields() == 0) { + return InvalidVariantShreddingSchema(type); + } + + auto schema = std::make_shared(); + schema->num_fields = struct_type->num_fields(); + + for (int32_t i = 0; i < struct_type->num_fields(); ++i) { + const auto& field = struct_type->field(i); + const auto& field_type = field->type(); + if (field->name() == VariantDefs::kTypedValueFieldName) { + if (schema->typed_idx != -1) { + return InvalidVariantShreddingSchema(type); + } + schema->typed_idx = i; + switch (field_type->id()) { + case arrow::Type::STRUCT: { + const auto& object_type = + std::static_pointer_cast(field_type); + schema->has_object_schema = true; + schema->object_schema.reserve(object_type->num_fields()); + for (int32_t index = 0; index < object_type->num_fields(); ++index) { + const auto& object_field = object_type->field(index); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr field_schema, + BuildVariantSchemaImpl(object_field->type(), /*top_level=*/false)); + schema->object_schema.push_back( + VariantSchema::ObjectField{object_field->name(), field_schema}); + auto [it, inserted] = + schema->object_schema_map.emplace(object_field->name(), index); + if (!inserted) { + return InvalidVariantShreddingSchema(type); + } + } + break; + } + case arrow::Type::LIST: { + const auto& list_type = std::static_pointer_cast(field_type); + PAIMON_ASSIGN_OR_RAISE( + schema->array_schema, + BuildVariantSchemaImpl(list_type->value_type(), /*top_level=*/false)); + break; + } + case arrow::Type::BOOL: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kBoolean}; + break; + case arrow::Type::INT8: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kByte}; + break; + case arrow::Type::INT16: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kShort}; + break; + case arrow::Type::INT32: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kInt}; + break; + case arrow::Type::INT64: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kLong}; + break; + case arrow::Type::FLOAT: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kFloat}; + break; + case arrow::Type::DOUBLE: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kDouble}; + break; + case arrow::Type::STRING: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kString}; + break; + case arrow::Type::BINARY: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kBinary}; + break; + case arrow::Type::DATE32: + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kDate}; + break; + case arrow::Type::DECIMAL128: { + const auto& decimal_type = + std::static_pointer_cast(field_type); + schema->scalar_schema = + VariantSchema::ScalarType{VariantSchema::ScalarKind::kDecimal, + decimal_type->precision(), decimal_type->scale()}; + break; + } + default: + return InvalidVariantShreddingSchema(type); + } + } else if (field->name() == VariantDefs::kValueFieldName) { + if (schema->variant_idx != -1 || field_type->id() != arrow::Type::BINARY) { + return InvalidVariantShreddingSchema(type); + } + schema->variant_idx = i; + } else if (field->name() == VariantDefs::kMetadataFieldName) { + if (schema->top_level_metadata_idx != -1 || field_type->id() != arrow::Type::BINARY) { + return InvalidVariantShreddingSchema(type); + } + schema->top_level_metadata_idx = i; + } else { + return InvalidVariantShreddingSchema(type); + } + } + + if (top_level != (schema->top_level_metadata_idx >= 0)) { + return InvalidVariantShreddingSchema(type); + } + return schema; +} + +} // namespace + +Result> VariantShreddingUtils::VariantShreddingSchema( + const std::shared_ptr& shredding_type) { + return VariantShreddingSchemaImpl(shredding_type, /*is_top_level=*/true, + /*is_object_field=*/false); +} + +Result> VariantShreddingUtils::BuildVariantSchema( + const std::shared_ptr& struct_type) { + return BuildVariantSchemaImpl(struct_type, /*top_level=*/true); +} + +Result> VariantShreddingUtils::ScalarSchemaToArrowType( + const VariantSchema::ScalarType& scalar) { + switch (scalar.kind) { + case VariantSchema::ScalarKind::kBoolean: + return arrow::boolean(); + case VariantSchema::ScalarKind::kByte: + return arrow::int8(); + case VariantSchema::ScalarKind::kShort: + return arrow::int16(); + case VariantSchema::ScalarKind::kInt: + return arrow::int32(); + case VariantSchema::ScalarKind::kLong: + return arrow::int64(); + case VariantSchema::ScalarKind::kFloat: + return arrow::float32(); + case VariantSchema::ScalarKind::kDouble: + return arrow::float64(); + case VariantSchema::ScalarKind::kString: + return arrow::utf8(); + case VariantSchema::ScalarKind::kBinary: + return arrow::binary(); + case VariantSchema::ScalarKind::kDecimal: + return arrow::decimal128(scalar.precision, scalar.scale); + case VariantSchema::ScalarKind::kDate: + return arrow::date32(); + default: + return Status::NotImplemented(fmt::format("Unsupported variant scalar kind: {}", + static_cast(scalar.kind))); + } +} + +bool VariantShreddingUtils::IsShreddedFileType( + const std::shared_ptr& file_variant_type) { + if (!file_variant_type || file_variant_type->id() != arrow::Type::STRUCT) { + return false; + } + const auto& struct_type = std::static_pointer_cast(file_variant_type); + return struct_type->GetFieldByName(VariantDefs::kTypedValueFieldName) != nullptr; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_utils.h b/src/paimon/common/data/variant/variant_shredding_utils.h new file mode 100644 index 000000000..4953a12b4 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_utils.h @@ -0,0 +1,61 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_schema.h" +#include "paimon/result.h" + +namespace arrow { +class DataType; +} // namespace arrow + +namespace paimon { + +/// Utils for converting between shredding schemas (`VariantSchema`) and their physical Arrow +/// representation, mirroring the Java `PaimonShreddingUtils` schema functions. +class VariantShreddingUtils { + public: + VariantShreddingUtils() = delete; + ~VariantShreddingUtils() = delete; + + /// Given an expected schema of a Variant value, returns a suitable physical schema for + /// shredding, by inserting appropriate intermediate value/typed_value fields at each level. + /// For example, to represent the JSON `{"a": 1, "b": "hello"}`, the schema + /// `struct{a: int32, b: string}` could be passed into this function, and it would return the + /// shredding schema: `struct{metadata: binary, value: binary, typed_value: struct{a: + /// struct{value: binary, typed_value: int32}, b: struct{value: binary, typed_value: + /// string}}}`. + static Result> VariantShreddingSchema( + const std::shared_ptr& shredding_type); + + /// Builds a `VariantSchema` from the physical shredded struct type (the inverse of + /// `VariantShreddingSchema`), validating field names and types. + static Result> BuildVariantSchema( + const std::shared_ptr& struct_type); + + /// The Arrow type of the typed_value column for a scalar shredding schema. + static Result> ScalarSchemaToArrowType( + const VariantSchema::ScalarType& scalar); + + /// Whether the physical struct type of a variant field in a data file is shredded (contains + /// a `typed_value` child). + static bool IsShreddedFileType(const std::shared_ptr& file_variant_type); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_write_plan.cpp b/src/paimon/common/data/variant/variant_shredding_write_plan.cpp new file mode 100644 index 000000000..d195d790e --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_write_plan.cpp @@ -0,0 +1,81 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_shredding_write_plan.h" + +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_shredding_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/types/data_type_json_parser.h" +#include "rapidjson/document.h" + +namespace paimon { + +Result> VariantShreddingWritePlan::Create( + const std::shared_ptr& logical_schema, + const std::map>& column_shredding_types) { + auto plan = std::shared_ptr(new VariantShreddingWritePlan()); + plan->logical_schema_ = logical_schema; + arrow::FieldVector physical_fields; + physical_fields.reserve(logical_schema->num_fields()); + for (const auto& field : logical_schema->fields()) { + auto it = column_shredding_types.find(field->name()); + if (it == column_shredding_types.end() || !VariantTypeUtils::IsVariantField(field)) { + physical_fields.push_back(field); + continue; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr physical_type, + VariantShreddingUtils::VariantShreddingSchema(it->second)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant_schema, + VariantShreddingUtils::BuildVariantSchema(physical_type)); + plan->column_schemas_.emplace(field->name(), std::move(variant_schema)); + plan->column_physical_types_.emplace(field->name(), physical_type); + physical_fields.push_back(field->WithType(physical_type)); + } + if (plan->column_schemas_.empty()) { + return Status::Invalid( + "variant shredding write plan matches no variant column in the write schema"); + } + plan->physical_schema_ = arrow::schema(physical_fields, logical_schema->metadata()); + return plan; +} + +Result> VariantShreddingWritePlan::FromConfiguredSchema( + const std::shared_ptr& logical_schema, + const std::string& configured_schema_json) { + rapidjson::Document doc; + doc.Parse(configured_schema_json.c_str()); + if (doc.HasParseError()) { + return Status::Invalid(fmt::format("failed to parse variant shredding schema json: {}", + configured_schema_json)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr configured_field, + DataTypeJsonParser::ParseType("shredding_schema", doc)); + if (configured_field->type()->id() != arrow::Type::STRUCT) { + return Status::Invalid("variant shredding schema must be a ROW type"); + } + std::map> column_shredding_types; + for (const auto& column_field : configured_field->type()->fields()) { + column_shredding_types.emplace(column_field->name(), column_field->type()); + } + return Create(logical_schema, column_shredding_types); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_write_plan.h b/src/paimon/common/data/variant/variant_shredding_write_plan.h new file mode 100644 index 000000000..6c90d17ab --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_write_plan.h @@ -0,0 +1,81 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_schema.h" +#include "paimon/result.h" + +namespace arrow { +class DataType; +class Schema; +} // namespace arrow + +namespace paimon { + +/// A physical write plan for variant shredding: maps the logical write schema (with variant +/// columns) to the physical schema where planned variant columns are replaced by their shredded +/// struct representation. +class VariantShreddingWritePlan { + public: + /// Creates a plan shredding the given top-level variant columns. + /// + /// @param logical_schema The logical write schema. + /// @param column_shredding_types The shredding type per variant column name, e.g. + /// `{"v": struct{a: int32, b: string}}`. Names that are not top-level variant columns + /// of `logical_schema` are ignored. + static Result> Create( + const std::shared_ptr& logical_schema, + const std::map>& column_shredding_types); + + /// Creates a plan from the `variant.shreddingSchema` option value: a ROW type JSON whose + /// fields map variant column names to their shredding types. + static Result> FromConfiguredSchema( + const std::shared_ptr& logical_schema, + const std::string& configured_schema_json); + + const std::shared_ptr& LogicalSchema() const { + return logical_schema_; + } + + const std::shared_ptr& PhysicalSchema() const { + return physical_schema_; + } + + /// The shredding schema per planned variant column name. + const std::map>& ColumnSchemas() const { + return column_schemas_; + } + + /// The physical shredded struct type per planned variant column name. + const std::map>& ColumnPhysicalTypes() const { + return column_physical_types_; + } + + private: + VariantShreddingWritePlan() = default; + + std::shared_ptr logical_schema_; + std::shared_ptr physical_schema_; + std::map> column_schemas_; + std::map> column_physical_types_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_write_plan_factory.cpp b/src/paimon/common/data/variant/variant_shredding_write_plan_factory.cpp new file mode 100644 index 000000000..4a95e60e7 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_write_plan_factory.cpp @@ -0,0 +1,144 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_shredding_write_plan_factory.h" + +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/infer_variant_shredding_schema.h" +#include "paimon/common/data/variant/variant_shredding_batch_converter.h" +#include "paimon/common/data/variant/variant_shredding_write_plan.h" +#include "paimon/common/data/variant/variant_type_utils.h" + +namespace paimon { + +VariantShreddingWritePlanFactory::VariantShreddingWritePlanFactory( + const CoreOptions& options, const std::shared_ptr& write_schema, + const std::shared_ptr& pool) + : write_schema_(write_schema), pool_(pool) { + configured_schema_ = options.GetVariantShreddingSchema(); + auto resolve = [this](auto&& result, auto* target) { + if (result.ok()) { + *target = result.value(); + } else if (config_status_.ok()) { + config_status_ = result.status(); + } + }; + resolve(options.VariantInferShreddingSchemaEnabled(), &infer_enabled_); + resolve(options.GetVariantShreddingMaxSchemaWidth(), &max_schema_width_); + resolve(options.GetVariantShreddingMaxSchemaDepth(), &max_schema_depth_); + resolve(options.GetVariantShreddingMinFieldCardinalityRatio(), &min_field_cardinality_ratio_); + resolve(options.GetVariantShreddingMaxInferBufferRow(), &max_infer_buffer_row_); +} + +bool VariantShreddingWritePlanFactory::ShouldCreateWritePlan() const { + return ContainsVariantField() && + (HasConfiguredShreddingSchema() || infer_enabled_ || !config_status_.ok()); +} + +bool VariantShreddingWritePlanFactory::ShouldInferWritePlan() const { + return ContainsVariantField() && !HasConfiguredShreddingSchema() && infer_enabled_ && + config_status_.ok(); +} + +int32_t VariantShreddingWritePlanFactory::InferBufferRowCount() const { + return max_infer_buffer_row_; +} + +bool VariantShreddingWritePlanFactory::HasConfiguredShreddingSchema() const { + return configured_schema_.has_value(); +} + +bool VariantShreddingWritePlanFactory::ContainsVariantField() const { + for (const auto& field : write_schema_->fields()) { + if (VariantTypeUtils::IsVariantField(field)) { + return true; + } + } + return false; +} + +Result> VariantShreddingWritePlanFactory::CreateConverter( + const std::string& file_format_identifier, + const std::vector>& sample_batches) const { + PAIMON_RETURN_NOT_OK(config_status_); + if (file_format_identifier != "parquet") { + return Status::NotImplemented( + fmt::format("variant shredding is only supported by the parquet file format, got {}", + file_format_identifier)); + } + + std::shared_ptr plan; + if (HasConfiguredShreddingSchema()) { + PAIMON_ASSIGN_OR_RAISE(plan, VariantShreddingWritePlan::FromConfiguredSchema( + write_schema_, configured_schema_.value())); + } else { + InferVariantShreddingSchema inferrer(max_schema_width_, max_schema_depth_, + min_field_cardinality_ratio_); + std::map> column_shredding_types; + for (int i = 0; i < write_schema_->num_fields(); ++i) { + const std::shared_ptr& field = write_schema_->field(i); + if (!VariantTypeUtils::IsVariantField(field)) { + continue; + } + std::vector> samples; + for (const auto& sample_batch : sample_batches) { + const auto& struct_array = static_cast(*sample_batch); + std::shared_ptr column = struct_array.field(i); + if (column == nullptr) { + return Status::Invalid( + fmt::format("sample batch misses the variant column '{}'", field->name())); + } + const auto& variant_array = static_cast(*column); + const auto& value_array = + static_cast(*variant_array.field(0)); + const auto& metadata_array = + static_cast(*variant_array.field(1)); + for (int64_t row = 0; row < variant_array.length(); ++row) { + if (variant_array.IsNull(row)) { + continue; + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr variant, + GenericVariant::Create(std::string_view(value_array.GetView(row)), + std::string_view(metadata_array.GetView(row)), + pool_)); + samples.push_back(std::move(variant)); + } + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr shredding_type, + inferrer.InferColumnShreddingType(samples)); + if (shredding_type != nullptr) { + column_shredding_types.emplace(field->name(), std::move(shredding_type)); + } + } + if (column_shredding_types.empty()) { + // No useful shredding schema was found; write the file unshredded. + return std::shared_ptr(nullptr); + } + PAIMON_ASSIGN_OR_RAISE( + plan, VariantShreddingWritePlan::Create(write_schema_, column_shredding_types)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr converter, + VariantShreddingBatchConverter::Create(plan, pool_)); + return std::shared_ptr(std::move(converter)); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_write_plan_factory.h b/src/paimon/common/data/variant/variant_shredding_write_plan_factory.h new file mode 100644 index 000000000..98dee8646 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_write_plan_factory.h @@ -0,0 +1,72 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/common/data/shredding/shredding_write_plan_factory.h" +#include "paimon/core/core_options.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class Schema; +} // namespace arrow + +namespace paimon { + +/// Creates VARIANT shredding batch converters, either from the configured +/// `variant.shreddingSchema` or, when `variant.inferShreddingSchema` is enabled, by inferring a +/// shredding schema from sampled rows buffered per file. +class VariantShreddingWritePlanFactory : public ShreddingWritePlanFactory { + public: + VariantShreddingWritePlanFactory(const CoreOptions& options, + const std::shared_ptr& write_schema, + const std::shared_ptr& pool); + + bool ShouldCreateWritePlan() const override; + + bool ShouldInferWritePlan() const override; + + int32_t InferBufferRowCount() const override; + + Result> CreateConverter( + const std::string& file_format_identifier, + const std::vector>& sample_batches) const override; + + private: + bool HasConfiguredShreddingSchema() const; + bool ContainsVariantField() const; + + std::shared_ptr write_schema_; + std::shared_ptr pool_; + + // Option values resolved at construction; a parse failure is deferred to CreateConverter. + std::optional configured_schema_; + bool infer_enabled_ = false; + int32_t max_schema_width_ = 0; + int32_t max_schema_depth_ = 0; + double min_field_cardinality_ratio_ = 0.0; + int32_t max_infer_buffer_row_ = 0; + Status config_status_ = Status::OK(); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_write_plan_factory_test.cpp b/src/paimon/common/data/variant/variant_shredding_write_plan_factory_test.cpp new file mode 100644 index 000000000..281605578 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_write_plan_factory_test.cpp @@ -0,0 +1,132 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_shredding_write_plan_factory.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/core_options.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/testing/utils/variant_test_data.h" + +namespace paimon::test { + +class VariantShreddingWritePlanFactoryTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + std::vector fields = {DataField(1, arrow::field("id", arrow::int32())), + DataField(2, VariantTypeUtils::ToArrowField("v"))}; + schema_ = DataField::ConvertDataFieldsToArrowSchema(fields); + } + + Result MakeOptions(std::map options) { + // Keep the manifest format resolvable in test binaries without the avro plugin. + options.emplace("manifest.format", "parquet"); + return CoreOptions::FromMap(options); + } + + std::shared_ptr BuildBatch(const std::vector& jsons) { + auto result = BuildVariantBatch(schema_->field(0), schema_->field(1), jsons, pool_); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return std::move(result).value(); + } + + protected: + std::shared_ptr pool_; + std::shared_ptr schema_; +}; + +TEST_F(VariantShreddingWritePlanFactoryTest, InactiveWithoutOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, MakeOptions({})); + VariantShreddingWritePlanFactory factory(options, schema_, pool_); + ASSERT_FALSE(factory.ShouldCreateWritePlan()); + ASSERT_FALSE(factory.ShouldInferWritePlan()); +} + +TEST_F(VariantShreddingWritePlanFactoryTest, ConfiguredSchema) { + const char* shredding_schema_json = R"({ + "type": "ROW", + "fields": [ { + "id": 0, + "name": "v", + "type": { + "type": "ROW", + "fields": [ + {"id": 1, "name": "age", "type": "INT"}, + {"id": 2, "name": "city", "type": "STRING"} + ] + } + } ] + })"; + ASSERT_OK_AND_ASSIGN(CoreOptions options, + MakeOptions({{"variant.shreddingSchema", shredding_schema_json}})); + VariantShreddingWritePlanFactory factory(options, schema_, pool_); + ASSERT_TRUE(factory.ShouldCreateWritePlan()); + ASSERT_FALSE(factory.ShouldInferWritePlan()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr converter, + factory.CreateConverter("parquet", {})); + ASSERT_NE(converter, nullptr); + auto variant_field = converter->GetPhysicalSchema()->GetFieldByName("v"); + ASSERT_NE(variant_field, nullptr); + const auto& physical_type = static_cast(*variant_field->type()); + ASSERT_NE(physical_type.GetFieldByName("typed_value"), nullptr); + // Variant shredding only supports the parquet format. + ASSERT_TRUE(factory.CreateConverter("orc", {}).status().IsNotImplemented()); +} + +TEST_F(VariantShreddingWritePlanFactoryTest, InferredSchema) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + MakeOptions({{"variant.inferShreddingSchema", "true"}})); + VariantShreddingWritePlanFactory factory(options, schema_, pool_); + ASSERT_TRUE(factory.ShouldCreateWritePlan()); + ASSERT_TRUE(factory.ShouldInferWritePlan()); + ASSERT_EQ(factory.InferBufferRowCount(), 4096); + + std::vector> samples = { + BuildBatch({R"({"age": 35, "city": "Chicago"})", R"({"age": 25, "city": "Hangzhou"})"}), + BuildBatch({R"({"age": 18, "city": "Beijing"})", nullptr})}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr converter, + factory.CreateConverter("parquet", samples)); + ASSERT_NE(converter, nullptr); + auto variant_field = converter->GetPhysicalSchema()->GetFieldByName("v"); + ASSERT_NE(variant_field, nullptr); + const auto& physical_type = static_cast(*variant_field->type()); + auto typed_value = physical_type.GetFieldByName("typed_value"); + ASSERT_NE(typed_value, nullptr); + const auto& typed_struct = static_cast(*typed_value->type()); + ASSERT_NE(typed_struct.GetFieldByName("age"), nullptr); + ASSERT_NE(typed_struct.GetFieldByName("city"), nullptr); +} + +TEST_F(VariantShreddingWritePlanFactoryTest, InferredSchemaWithoutSamples) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + MakeOptions({{"variant.inferShreddingSchema", "true"}})); + VariantShreddingWritePlanFactory factory(options, schema_, pool_); + // With no useful samples the file stays unshredded. + ASSERT_OK_AND_ASSIGN(std::shared_ptr converter, + factory.CreateConverter("parquet", {})); + ASSERT_EQ(converter, nullptr); +} + +} // namespace paimon::test diff --git a/src/paimon/common/data/variant/variant_shredding_writer.cpp b/src/paimon/common/data/variant/variant_shredding_writer.cpp new file mode 100644 index 000000000..4a14f7e5f --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_writer.cpp @@ -0,0 +1,468 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#include "paimon/common/data/variant/variant_shredding_writer.h" + +#include +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_builder.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/utils/arrow/status_utils.h" + +namespace paimon { + +namespace { + +// Rescales `unscaled` by `10^power`, failing when the result would exceed the 38-digit +// limit. The precision check must precede each multiplication: a 38-digit value times 10 +// overflows the signed 128-bit representation before a post-check could reject it. +Result<__int128_t> ScaleUpUnscaled(__int128_t unscaled, int32_t power) { + __int128_t result = unscaled; + for (int32_t i = 0; i < power; ++i) { + VariantDecimal probe{result, 0}; + if (probe.Precision() >= VariantDefs::kMaxDecimal16Precision) { + return Status::Invalid("decimal overflow while rescaling"); + } + result *= 10; + } + return result; +} + +Status AppendDecimalTo(arrow::ArrayBuilder* builder, __int128_t unscaled) { + auto* decimal_builder = static_cast(builder); + arrow::Decimal128 value(static_cast(unscaled >> 64), + static_cast(static_cast<__uint128_t>(unscaled))); + PAIMON_RETURN_NOT_OK_FROM_ARROW(decimal_builder->Append(value)); + return Status::OK(); +} + +} // namespace + +VariantShreddedColumnWriter::VariantShreddedColumnWriter( + const std::shared_ptr& schema, + std::unique_ptr&& root_builder) + : schema_(schema), root_builder_(std::move(root_builder)) {} + +Result> VariantShreddedColumnWriter::Create( + const std::shared_ptr& schema, + const std::shared_ptr& physical_type, arrow::MemoryPool* pool) { + if (!schema || schema->top_level_metadata_idx < 0) { + return Status::Invalid("variant shredding schema must contain a top-level metadata field"); + } + std::unique_ptr root_builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder(pool, physical_type, &root_builder)); + if (root_builder->type()->id() != arrow::Type::STRUCT) { + return Status::Invalid("variant shredded physical type must be a struct"); + } + auto writer = std::unique_ptr( + new VariantShreddedColumnWriter(schema, std::move(root_builder))); + PAIMON_RETURN_NOT_OK(BuildNode( + schema, static_cast(writer->root_builder_.get()), &writer->root_)); + return writer; +} + +Status VariantShreddedColumnWriter::BuildNode(const std::shared_ptr& schema, + arrow::StructBuilder* group, Node* node) { + node->schema = schema.get(); + node->group = group; + if (group->num_children() != schema->num_fields) { + return Status::Invalid( + fmt::format("variant shredded builder has {} children but the schema has {} fields", + group->num_children(), schema->num_fields)); + } + if (schema->top_level_metadata_idx >= 0) { + node->metadata = static_cast( + group->field_builder(schema->top_level_metadata_idx)); + } + if (schema->variant_idx >= 0) { + node->value = static_cast(group->field_builder(schema->variant_idx)); + } + if (schema->typed_idx >= 0) { + arrow::ArrayBuilder* typed_builder = group->field_builder(schema->typed_idx); + if (schema->has_object_schema) { + node->typed_object = static_cast(typed_builder); + node->object_children.resize(schema->object_schema.size()); + for (size_t i = 0; i < schema->object_schema.size(); ++i) { + auto* child_group = static_cast( + node->typed_object->field_builder(static_cast(i))); + PAIMON_RETURN_NOT_OK(BuildNode(schema->object_schema[i].schema, child_group, + &node->object_children[i])); + } + } else if (schema->array_schema) { + node->typed_list = static_cast(typed_builder); + node->array_element = std::make_unique(); + auto* element_group = + static_cast(node->typed_list->value_builder()); + PAIMON_RETURN_NOT_OK( + BuildNode(schema->array_schema, element_group, node->array_element.get())); + } else if (schema->scalar_schema) { + node->typed_scalar = typed_builder; + } else { + return Status::Invalid("variant shredding schema typed_value has no schema"); + } + } + return Status::OK(); +} + +Status VariantShreddedColumnWriter::Append(const GenericVariant& variant) { + return AppendVariantNode(&root_, variant); +} + +Status VariantShreddedColumnWriter::AppendNull() { + PAIMON_RETURN_NOT_OK_FROM_ARROW(root_.group->AppendNull()); + return Status::OK(); +} + +Result> VariantShreddedColumnWriter::Finish() { + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(root_builder_->Finish(&array)); + return array; +} + +Status VariantShreddedColumnWriter::AppendVariantNode(Node* node, const GenericVariant& variant) { + const VariantSchema& schema = *node->schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->group->Append()); + if (schema.top_level_metadata_idx >= 0) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->metadata->Append(variant.Metadata())); + } + PAIMON_ASSIGN_OR_RAISE(VariantValueType variant_type, variant.GetType()); + if (schema.array_schema != nullptr && variant_type == VariantValueType::kArray) { + // The array element is always a struct containing untyped and typed fields. + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_list->Append()); + PAIMON_ASSIGN_OR_RAISE(int32_t size, variant.ArraySize()); + for (int32_t i = 0; i < size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr element, + variant.GetElementAtIndex(i)); + PAIMON_RETURN_NOT_OK(AppendVariantNode(node->array_element.get(), *element)); + } + if (node->value != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->AppendNull()); + } + } else if (schema.has_object_schema && variant_type == VariantValueType::kObject) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_object->Append()); + // Collect any field that exists in the variant, but not in the shredding schema, into a + // residual variant object that shares the top-level metadata. + VariantBuilder residual(/*allow_duplicate_keys=*/false); + std::vector field_entries; + std::vector matched(node->object_children.size(), false); + int32_t start = residual.GetWritePos(); + PAIMON_ASSIGN_OR_RAISE(int32_t object_size, variant.ObjectSize()); + for (int32_t i = 0; i < object_size; ++i) { + PAIMON_ASSIGN_OR_RAISE(auto field, variant.GetFieldAtIndex(i)); + if (!field.has_value()) { + return VariantBinaryUtil::MalformedVariant(); + } + auto it = schema.object_schema_map.find(field->key); + if (it != schema.object_schema_map.end()) { + PAIMON_RETURN_NOT_OK( + AppendVariantNode(&node->object_children[it->second], *field->value)); + matched[it->second] = true; + } else { + // The field is not shredded. Put it in the untyped value column. The shallow + // append is needed for correctness, since the metadata ids must stay unchanged. + PAIMON_ASSIGN_OR_RAISE(int32_t id, variant.GetDictionaryIdAtIndex(i)); + field_entries.emplace_back(field->key, id, residual.GetWritePos() - start); + PAIMON_RETURN_NOT_OK( + residual.ShallowAppendVariant(field->value->RawValue(), field->value->Pos())); + } + } + // Set missing fields to non-null with all fields set to null. + for (size_t i = 0; i < matched.size(); ++i) { + if (!matched[i]) { + PAIMON_RETURN_NOT_OK(AppendMissingNode(&node->object_children[i])); + } + } + if (residual.GetWritePos() != start) { + PAIMON_RETURN_NOT_OK(residual.FinishWritingObject(start, &field_entries)); + if (node->value == nullptr) { + return Status::Invalid( + "variant shredding schema has no value column for residual fields"); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->Append(residual.ValueWithoutMetadata())); + } else if (node->value != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->AppendNull()); + } + } else if (schema.scalar_schema.has_value()) { + bool shredded = false; + PAIMON_RETURN_NOT_OK(TryTypedShred(node, variant, variant_type, &shredded)); + if (shredded) { + if (node->value != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->AppendNull()); + } + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_scalar->AppendNull()); + if (node->value == nullptr) { + return Status::Invalid( + "variant shredding schema has no value column for untyped values"); + } + PAIMON_ASSIGN_OR_RAISE(std::string_view value, variant.Value()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->Append(value)); + } + } else { + if (node->typed_list != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_list->AppendNull()); + } else if (node->typed_object != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_object->AppendNull()); + } else if (node->typed_scalar != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_scalar->AppendNull()); + } + if (node->value == nullptr) { + return Status::Invalid( + "variant shredding schema has no value column for untyped values"); + } + PAIMON_ASSIGN_OR_RAISE(std::string_view value, variant.Value()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->Append(value)); + } + return Status::OK(); +} + +Status VariantShreddedColumnWriter::AppendMissingNode(Node* node) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->group->Append()); + if (node->metadata != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->metadata->AppendNull()); + } + if (node->value != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->value->AppendNull()); + } + if (node->typed_list != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_list->AppendNull()); + } else if (node->typed_object != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_object->AppendNull()); + } else if (node->typed_scalar != nullptr) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(node->typed_scalar->AppendNull()); + } + return Status::OK(); +} + +Status VariantShreddedColumnWriter::TryTypedShred(Node* node, const GenericVariant& variant, + VariantValueType variant_type, bool* shredded) { + const VariantSchema::ScalarType& target = node->schema->scalar_schema.value(); + *shredded = false; + switch (variant_type) { + case VariantValueType::kLong: { + PAIMON_ASSIGN_OR_RAISE(int64_t value, variant.GetLong()); + switch (target.kind) { + // Check that the target type can hold the actual value. + case VariantSchema::ScalarKind::kByte: + if (value == static_cast(value)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(value))); + *shredded = true; + } + break; + case VariantSchema::ScalarKind::kShort: + if (value == static_cast(value)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(value))); + *shredded = true; + } + break; + case VariantSchema::ScalarKind::kInt: + if (value == static_cast(value)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(value))); + *shredded = true; + } + break; + case VariantSchema::ScalarKind::kLong: + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar)->Append(value)); + *shredded = true; + break; + case VariantSchema::ScalarKind::kDecimal: { + // If the integer can fit in the given decimal precision, allow it. + auto scaled = ScaleUpUnscaled(value, target.scale); + if (scaled.ok()) { + VariantDecimal probe{scaled.value(), target.scale}; + if (probe.Precision() <= target.precision) { + PAIMON_RETURN_NOT_OK( + AppendDecimalTo(node->typed_scalar, scaled.value())); + *shredded = true; + } + } + break; + } + default: + break; + } + break; + } + case VariantValueType::kDecimal: { + if (target.kind == VariantSchema::ScalarKind::kDecimal) { + // Use the original scale so that scale information is retained. + PAIMON_ASSIGN_OR_RAISE(VariantDecimal value, + VariantBinaryUtil::GetDecimalWithOriginalScale( + variant.RawValue(), variant.Pos())); + if (value.Precision() <= target.precision && value.scale == target.scale) { + PAIMON_RETURN_NOT_OK(AppendDecimalTo(node->typed_scalar, value.unscaled)); + *shredded = true; + break; + } + // Convert to the target scale, and see if it fits without losing information. + int32_t scale_diff = target.scale - value.scale; + __int128_t rescaled = value.unscaled; + bool exact = true; + if (scale_diff >= 0) { + auto scaled = ScaleUpUnscaled(value.unscaled, scale_diff); + if (scaled.ok()) { + rescaled = scaled.value(); + } else { + exact = false; + } + } else { + for (int32_t i = 0; i < -scale_diff && exact; ++i) { + if (rescaled % 10 != 0) { + exact = false; + } else { + rescaled /= 10; + } + } + } + if (exact) { + VariantDecimal probe{rescaled, target.scale}; + if (probe.Precision() <= target.precision) { + PAIMON_RETURN_NOT_OK(AppendDecimalTo(node->typed_scalar, rescaled)); + *shredded = true; + } + } + } else if (target.kind == VariantSchema::ScalarKind::kByte || + target.kind == VariantSchema::ScalarKind::kShort || + target.kind == VariantSchema::ScalarKind::kInt || + target.kind == VariantSchema::ScalarKind::kLong) { + // Check if the decimal happens to be an integer. + PAIMON_ASSIGN_OR_RAISE(VariantDecimal value, variant.GetDecimal()); + if (value.scale > 0) { + break; + } + auto scaled = ScaleUpUnscaled(value.unscaled, -value.scale); + if (!scaled.ok()) { + break; + } + __int128_t integral = scaled.value(); + if (integral != static_cast(integral)) { + break; + } + auto long_value = static_cast(integral); + switch (target.kind) { + case VariantSchema::ScalarKind::kByte: + if (long_value == static_cast(long_value)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(long_value))); + *shredded = true; + } + break; + case VariantSchema::ScalarKind::kShort: + if (long_value == static_cast(long_value)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(long_value))); + *shredded = true; + } + break; + case VariantSchema::ScalarKind::kInt: + if (long_value == static_cast(long_value)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(long_value))); + *shredded = true; + } + break; + default: + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(long_value)); + *shredded = true; + break; + } + } + break; + } + case VariantValueType::kBoolean: { + if (target.kind == VariantSchema::ScalarKind::kBoolean) { + PAIMON_ASSIGN_OR_RAISE(bool value, variant.GetBoolean()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar)->Append(value)); + *shredded = true; + } + break; + } + case VariantValueType::kString: { + if (target.kind == VariantSchema::ScalarKind::kString) { + PAIMON_ASSIGN_OR_RAISE(std::string_view value, variant.GetString()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar)->Append(value)); + *shredded = true; + } + break; + } + case VariantValueType::kDouble: { + if (target.kind == VariantSchema::ScalarKind::kDouble) { + PAIMON_ASSIGN_OR_RAISE(double value, variant.GetDouble()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar)->Append(value)); + *shredded = true; + } + break; + } + case VariantValueType::kFloat: { + if (target.kind == VariantSchema::ScalarKind::kFloat) { + PAIMON_ASSIGN_OR_RAISE(float value, variant.GetFloat()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar)->Append(value)); + *shredded = true; + } + break; + } + case VariantValueType::kDate: { + if (target.kind == VariantSchema::ScalarKind::kDate) { + PAIMON_ASSIGN_OR_RAISE(int64_t value, variant.GetLong()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar) + ->Append(static_cast(value))); + *shredded = true; + } + break; + } + case VariantValueType::kBinary: { + if (target.kind == VariantSchema::ScalarKind::kBinary) { + PAIMON_ASSIGN_OR_RAISE(std::string_view value, variant.GetBinary()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + static_cast(node->typed_scalar)->Append(value)); + *shredded = true; + } + break; + } + default: + // TIMESTAMP/TIMESTAMP_NTZ/UUID typed columns are not producible by the configured + // shredding schema; the value stays in the untyped column. + break; + } + return Status::OK(); +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_shredding_writer.h b/src/paimon/common/data/variant/variant_shredding_writer.h new file mode 100644 index 000000000..9ce0ae9c3 --- /dev/null +++ b/src/paimon/common/data/variant/variant_shredding_writer.h @@ -0,0 +1,103 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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. + */ + +/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed + * by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE + * file distributed with this work for additional information regarding copyright ownership. */ + +#pragma once + +#include +#include + +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_schema.h" +#include "paimon/result.h" + +namespace arrow { +class Array; +class ArrayBuilder; +class BinaryBuilder; +class DataType; +class ListBuilder; +class MemoryPool; +class StructBuilder; +} // namespace arrow + +namespace paimon { + +/// Shreds variant values of one column into a physical shredded Arrow array, implementing the +/// `castShredded` algorithm of the parquet-format VariantShredding.md specification (mirroring +/// the Java `VariantShreddingWriter`). Decimals and integers are allowed to shred to numerically +/// equivalent values of a different scale (`allowNumericScaleChanges` in Java is always true). +class VariantShreddedColumnWriter { + public: + /// Creates a writer for one variant column. + /// + /// @param schema The shredding schema of the column. + /// @param physical_type The physical shredded struct type + /// (`VariantShreddingUtils::VariantShreddingSchema` output). + /// @param pool The Arrow memory pool used by the builders. + static Result> Create( + const std::shared_ptr& schema, + const std::shared_ptr& physical_type, arrow::MemoryPool* pool); + + /// Shreds one variant value and appends the result row. + Status Append(const GenericVariant& variant); + + /// Appends a null variant row. + Status AppendNull(); + + /// Finishes and returns the shredded array of all appended rows. + Result> Finish(); + + private: + /// Builder handles of one shredding schema node (a `metadata`/`value`/`typed_value` group). + struct Node { + const VariantSchema* schema = nullptr; + arrow::StructBuilder* group = nullptr; + arrow::BinaryBuilder* metadata = nullptr; + arrow::BinaryBuilder* value = nullptr; + // Exactly one of the following is set when `schema->typed_idx >= 0`. + arrow::ArrayBuilder* typed_scalar = nullptr; + arrow::ListBuilder* typed_list = nullptr; + arrow::StructBuilder* typed_object = nullptr; + std::vector object_children; + std::unique_ptr array_element; + }; + + VariantShreddedColumnWriter(const std::shared_ptr& schema, + std::unique_ptr&& root_builder); + + static Status BuildNode(const std::shared_ptr& schema, + arrow::StructBuilder* group, Node* node); + + Status AppendVariantNode(Node* node, const GenericVariant& variant); + + /// Appends a missing object field: the group is present with all its children set to null. + Status AppendMissingNode(Node* node); + + /// Tries to append the variant as a typed scalar. Sets `*shredded` to whether it succeeded + /// (on failure nothing is appended). + Status TryTypedShred(Node* node, const GenericVariant& variant, VariantValueType variant_type, + bool* shredded); + + std::shared_ptr schema_; + std::unique_ptr root_builder_; + Node root_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_test.cpp b/src/paimon/common/data/variant/variant_test.cpp new file mode 100644 index 000000000..4eac6f572 --- /dev/null +++ b/src/paimon/common/data/variant/variant_test.cpp @@ -0,0 +1,107 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/data/variant.h" + +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class VariantPublicApiTest : public ::testing::Test { + protected: + std::shared_ptr pool_ = GetDefaultPool(); +}; + +TEST_F(VariantPublicApiTest, FromJsonAndAccessors) { + ASSERT_OK_AND_ASSIGN(auto variant, + Variant::FromJson("{\"age\": 35, \"city\": \"Hangzhou\"}", pool_)); + ASSERT_GT(variant->Value().size(), 0); + ASSERT_GT(variant->Metadata().size(), 0); + ASSERT_EQ(variant->SizeInBytes(), + static_cast(variant->Value().size() + variant->Metadata().size())); + ASSERT_OK_AND_ASSIGN(std::string json, variant->ToJson()); + ASSERT_EQ(json, "{\"age\":35,\"city\":\"Hangzhou\"}"); + + ASSERT_OK_AND_ASSIGN( + auto rebuilt, + Variant::Create(variant->Value().data(), variant->Value().size(), + variant->Metadata().data(), variant->Metadata().size(), pool_)); + ASSERT_OK_AND_ASSIGN(std::string rebuilt_json, rebuilt->ToJson()); + ASSERT_EQ(rebuilt_json, json); +} + +TEST_F(VariantPublicApiTest, VariantGet) { + ASSERT_OK_AND_ASSIGN(auto variant, + Variant::FromJson("{\"age\": 35, \"city\": \"Hangzhou\"}", pool_)); + VariantCastArgs cast_args; + cast_args.fail_on_error = false; + { + auto target = std::make_unique(); + ASSERT_TRUE(arrow::ExportField(arrow::Field("t", arrow::int64()), target.get()).ok()); + ASSERT_OK_AND_ASSIGN(std::optional literal, + variant->VariantGet("$.age", target.get(), cast_args)); + ASSERT_TRUE(literal.has_value()); + ASSERT_EQ(literal->GetValue(), 35); + } + { + auto target = std::make_unique(); + ASSERT_TRUE(arrow::ExportField(arrow::Field("t", arrow::utf8()), target.get()).ok()); + ASSERT_OK_AND_ASSIGN(std::optional literal, + variant->VariantGet("$.missing", target.get(), cast_args)); + ASSERT_FALSE(literal.has_value()); + } + ASSERT_OK_AND_ASSIGN(std::optional sub_json, variant->VariantGetJson("$")); + ASSERT_TRUE(sub_json.has_value()); + ASSERT_EQ(*sub_json, "{\"age\":35,\"city\":\"Hangzhou\"}"); + ASSERT_OK_AND_ASSIGN(std::optional missing_json, + variant->VariantGetJson("$.missing")); + ASSERT_FALSE(missing_json.has_value()); +} + +TEST_F(VariantPublicApiTest, ArrowField) { + ASSERT_OK_AND_ASSIGN(auto c_field, Variant::ArrowField("v", /*nullable=*/true)); + auto imported = arrow::ImportField(c_field.get()); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + std::shared_ptr field = imported.ValueOrDie(); + ASSERT_TRUE(VariantTypeUtils::IsVariantField(field)); + ASSERT_TRUE(field->nullable()); + ASSERT_TRUE(field->type()->Equals(VariantTypeUtils::UnshreddedStructType())); +} + +TEST_F(VariantPublicApiTest, VariantGetArrow) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr variant, + Variant::FromJson("{\"user\": {\"name\": \"Paimon\"}}", pool_)); + auto target_type = arrow::struct_({arrow::field("name", arrow::utf8())}); + auto target = std::make_unique(); + ASSERT_TRUE(arrow::ExportField(arrow::Field("t", target_type), target.get()).ok()); + VariantCastArgs cast_args; + ASSERT_OK_AND_ASSIGN(std::unique_ptr c_array, + variant->VariantGetArrow("$.user", target.get(), cast_args)); + auto imported = arrow::ImportArray(c_array.get(), target_type); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + std::shared_ptr array = imported.ValueOrDie(); + ASSERT_EQ(array->length(), 1); + const auto& row = static_cast(*array); + ASSERT_EQ(static_cast(*row.field(0)).GetString(0), "Paimon"); +} + +} // namespace paimon::test diff --git a/src/paimon/common/data/variant/variant_type_utils.cpp b/src/paimon/common/data/variant/variant_type_utils.cpp new file mode 100644 index 000000000..7c19e6646 --- /dev/null +++ b/src/paimon/common/data/variant/variant_type_utils.cpp @@ -0,0 +1,125 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_type_utils.h" + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/types/data_field.h" + +namespace paimon { + +bool VariantTypeUtils::IsVariantField(const std::shared_ptr& field) { + if (field->type()->id() != arrow::Type::STRUCT) { + return false; + } + if (!field->HasMetadata()) { + return false; + } + return IsVariantMetadata(field->metadata()); +} + +bool VariantTypeUtils::IsUnshreddedVariantType(const std::shared_ptr& type) { + if (type == nullptr || type->id() != arrow::Type::STRUCT || type->num_fields() != 2) { + return false; + } + const auto& value_field = type->field(0); + const auto& metadata_field = type->field(1); + return value_field->name() == VariantDefs::kValueFieldName && + value_field->type()->id() == arrow::Type::BINARY && !value_field->nullable() && + metadata_field->name() == VariantDefs::kMetadataFieldName && + metadata_field->type()->id() == arrow::Type::BINARY && !metadata_field->nullable(); +} + +bool VariantTypeUtils::IsVariantMetadata( + const std::shared_ptr& metadata) { + if (!metadata) { + return false; + } + auto extension_name = metadata->Get(VariantDefs::kExtensionTypeKey); + return extension_name.ok() && *extension_name == VariantDefs::kExtensionTypeValue; +} + +std::shared_ptr VariantTypeUtils::UnshreddedStructType() { + auto value_field = + arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/false, + arrow::KeyValueMetadata::Make({DataField::FIELD_ID}, + {std::to_string(VariantDefs::kValueFieldId)})); + auto metadata_field = + arrow::field(VariantDefs::kMetadataFieldName, arrow::binary(), /*nullable=*/false, + arrow::KeyValueMetadata::Make( + {DataField::FIELD_ID}, {std::to_string(VariantDefs::kMetadataFieldId)})); + return arrow::struct_({value_field, metadata_field}); +} + +std::shared_ptr VariantTypeUtils::ToArrowField( + const std::string& field_name, bool nullable, + std::unordered_map metadata) { + metadata[VariantDefs::kExtensionTypeKey] = VariantDefs::kExtensionTypeValue; + return arrow::field(field_name, UnshreddedStructType(), nullable, + std::make_shared(metadata)); +} + +Status VariantTypeUtils::ValidateVariantShape(const std::shared_ptr& field) { + const auto& type = field->type(); + if (type->id() != arrow::Type::STRUCT) { + return Status::Invalid(fmt::format("Variant field '{}' must be a struct, but got {}", + field->name(), type->ToString())); + } + const auto& struct_type = std::static_pointer_cast(type); + if (struct_type->num_fields() != 2) { + return Status::Invalid( + fmt::format("Variant field '{}' must be a struct, " + "but got {}", + field->name(), type->ToString())); + } + const auto& value_field = struct_type->field(0); + const auto& metadata_field = struct_type->field(1); + if (value_field->name() != VariantDefs::kValueFieldName || + value_field->type()->id() != arrow::Type::BINARY || value_field->nullable() || + metadata_field->name() != VariantDefs::kMetadataFieldName || + metadata_field->type()->id() != arrow::Type::BINARY || metadata_field->nullable()) { + return Status::Invalid( + fmt::format("Variant field '{}' must be a struct, but got {}", + field->name(), type->ToString())); + } + return Status::OK(); +} + +bool VariantTypeUtils::ContainsVariantField(const std::shared_ptr& field) { + if (IsVariantField(field)) { + return true; + } + for (const auto& child : field->type()->fields()) { + if (ContainsVariantField(child)) { + return true; + } + } + return false; +} + +bool VariantTypeUtils::ContainsVariantField(const std::shared_ptr& schema) { + for (const auto& field : schema->fields()) { + if (ContainsVariantField(field)) { + return true; + } + } + return false; +} + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_type_utils.h b/src/paimon/common/data/variant/variant_type_utils.h new file mode 100644 index 000000000..0ffc3b1d4 --- /dev/null +++ b/src/paimon/common/data/variant/variant_type_utils.h @@ -0,0 +1,75 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/result.h" +#include "paimon/visibility.h" + +namespace arrow { +class DataType; +class Field; +class KeyValueMetadata; +class Schema; +} // namespace arrow + +namespace paimon { + +/// Utils for the Paimon Variant type, whose underlying Arrow representation is +/// `struct` marked with the +/// `paimon.extension.type = paimon.type.variant` field metadata (see `VariantDefs`). +class PAIMON_EXPORT VariantTypeUtils { + public: + VariantTypeUtils() = delete; + ~VariantTypeUtils() = delete; + + /// Whether the field is a Paimon Variant field (a struct with the variant metadata marker). + static bool IsVariantField(const std::shared_ptr& field); + + /// Whether `type` is the unshredded variant physical type + /// `struct` (field metadata ignored). + static bool IsUnshreddedVariantType(const std::shared_ptr& type); + + /// Whether the metadata carries the Paimon Variant extension type marker. + static bool IsVariantMetadata(const std::shared_ptr& metadata); + + /// The unshredded physical Arrow type of a variant field: + /// `struct` with paimon field ids 0/1 on + /// the children. + static std::shared_ptr UnshreddedStructType(); + + /// Creates a Variant Arrow field with the variant metadata marker. + static std::shared_ptr ToArrowField( + const std::string& field_name, bool nullable = true, + std::unordered_map metadata = {}); + + /// Validates that a variant-marked field has the expected physical shape: + /// `struct`. + static Status ValidateVariantShape(const std::shared_ptr& field); + + /// Whether the schema contains a variant field, at the top level or nested inside structs, + /// arrays or maps. + static bool ContainsVariantField(const std::shared_ptr& schema); + + /// Whether the field contains a variant field, itself included, at any nesting level. + static bool ContainsVariantField(const std::shared_ptr& field); +}; + +} // namespace paimon diff --git a/src/paimon/common/data/variant/variant_type_utils_test.cpp b/src/paimon/common/data/variant/variant_type_utils_test.cpp new file mode 100644 index 000000000..d01fe02f8 --- /dev/null +++ b/src/paimon/common/data/variant/variant_type_utils_test.cpp @@ -0,0 +1,102 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/common/data/variant/variant_type_utils.h" + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/utils/field_type_utils.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(VariantTypeUtilsTest, ToArrowFieldAndDetection) { + auto field = VariantTypeUtils::ToArrowField("v"); + ASSERT_EQ(field->name(), "v"); + ASSERT_TRUE(field->nullable()); + ASSERT_TRUE(VariantTypeUtils::IsVariantField(field)); + ASSERT_TRUE(VariantTypeUtils::IsVariantMetadata(field->metadata())); + ASSERT_OK(VariantTypeUtils::ValidateVariantShape(field)); + + auto struct_type = std::static_pointer_cast(field->type()); + ASSERT_EQ(struct_type->num_fields(), 2); + ASSERT_EQ(struct_type->field(0)->name(), VariantDefs::kValueFieldName); + ASSERT_EQ(struct_type->field(0)->type()->id(), arrow::Type::BINARY); + ASSERT_FALSE(struct_type->field(0)->nullable()); + ASSERT_EQ(struct_type->field(1)->name(), VariantDefs::kMetadataFieldName); + ASSERT_EQ(struct_type->field(1)->type()->id(), arrow::Type::BINARY); + ASSERT_FALSE(struct_type->field(1)->nullable()); + // The children carry paimon field ids 0/1 (mapped to parquet field ids on write). + ASSERT_EQ(struct_type->field(0)->metadata()->Get("paimon.id").ValueOr(""), "0"); + ASSERT_EQ(struct_type->field(1)->metadata()->Get("paimon.id").ValueOr(""), "1"); + + auto plain_struct = arrow::field("s", VariantTypeUtils::UnshreddedStructType()); + ASSERT_FALSE(VariantTypeUtils::IsVariantField(plain_struct)); + std::unordered_map metadata = { + {VariantDefs::kExtensionTypeKey, VariantDefs::kExtensionTypeValue}}; + auto marked_binary = arrow::field("b", arrow::binary(), true, + std::make_shared(metadata)); + ASSERT_FALSE(VariantTypeUtils::IsVariantField(marked_binary)); +} + +TEST(VariantTypeUtilsTest, ValidateVariantShapeRejectsWrongShape) { + std::unordered_map metadata = { + {VariantDefs::kExtensionTypeKey, VariantDefs::kExtensionTypeValue}}; + auto arrow_metadata = std::make_shared(metadata); + auto one_child = arrow::field( + "v", arrow::struct_({arrow::field("value", arrow::binary(), false)}), true, arrow_metadata); + ASSERT_NOK(VariantTypeUtils::ValidateVariantShape(one_child)); + auto nullable_children = + arrow::field("v", + arrow::struct_({arrow::field("value", arrow::binary(), true), + arrow::field("metadata", arrow::binary(), true)}), + true, arrow_metadata); + ASSERT_NOK(VariantTypeUtils::ValidateVariantShape(nullable_children)); + auto wrong_type = + arrow::field("v", + arrow::struct_({arrow::field("value", arrow::utf8(), false), + arrow::field("metadata", arrow::binary(), false)}), + true, arrow_metadata); + ASSERT_NOK(VariantTypeUtils::ValidateVariantShape(wrong_type)); +} + +TEST(VariantTypeUtilsTest, ContainsVariantField) { + auto variant_field = VariantTypeUtils::ToArrowField("v"); + auto plain = arrow::schema({arrow::field("a", arrow::int32())}); + ASSERT_FALSE(VariantTypeUtils::ContainsVariantField(plain)); + auto top_level = arrow::schema({arrow::field("a", arrow::int32()), variant_field}); + ASSERT_TRUE(VariantTypeUtils::ContainsVariantField(top_level)); + auto nested = + arrow::schema({arrow::field("row", arrow::struct_({arrow::field("inner", arrow::int32()), + VariantTypeUtils::ToArrowField("v")}))}); + ASSERT_TRUE(VariantTypeUtils::ContainsVariantField(nested)); + auto in_list = + arrow::schema({arrow::field("l", arrow::list(VariantTypeUtils::ToArrowField("item")))}); + ASSERT_TRUE(VariantTypeUtils::ContainsVariantField(in_list)); +} + +TEST(VariantTypeUtilsTest, FieldTypeConversion) { + auto variant_field = VariantTypeUtils::ToArrowField("v"); + ASSERT_OK_AND_ASSIGN(FieldType field_type, FieldTypeUtils::ConvertToFieldType(variant_field)); + ASSERT_EQ(field_type, FieldType::VARIANT); + ASSERT_EQ(FieldTypeUtils::FieldTypeToString(FieldType::VARIANT), "VARIANT"); + auto plain_struct = arrow::field("s", VariantTypeUtils::UnshreddedStructType()); + ASSERT_OK_AND_ASSIGN(FieldType plain_type, FieldTypeUtils::ConvertToFieldType(plain_struct)); + ASSERT_EQ(plain_type, FieldType::STRUCT); +} + +} // namespace paimon::test diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index 0186cd4ce..7d572f366 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -97,6 +97,15 @@ const char Options::MAP_STORAGE_LAYOUT[] = "map.storage-layout"; const char Options::MAP_SHARED_SHREDDING_MAX_COLUMNS[] = "map.shared-shredding.max-columns"; const char Options::MAP_SHARED_SHREDDING_COLUMN_PLACEMENT_POLICY[] = "map.shared-shredding.column-placement-policy"; +const char Options::VARIANT_SHREDDING_SCHEMA[] = "variant.shreddingSchema"; +const char Options::PARQUET_VARIANT_SHREDDING_SCHEMA[] = "parquet.variant.shreddingSchema"; +const char Options::VARIANT_INFER_SHREDDING_SCHEMA[] = "variant.inferShreddingSchema"; +const char Options::VARIANT_SHREDDING_MAX_SCHEMA_WIDTH[] = "variant.shredding.maxSchemaWidth"; +const char Options::VARIANT_SHREDDING_MAX_SCHEMA_DEPTH[] = "variant.shredding.maxSchemaDepth"; +const char Options::VARIANT_SHREDDING_MIN_FIELD_CARDINALITY_RATIO[] = + "variant.shredding.minFieldCardinalityRatio"; +const char Options::VARIANT_SHREDDING_MAX_INFER_BUFFER_ROW[] = + "variant.shredding.maxInferBufferRow"; const char Options::BLOB_AS_DESCRIPTOR[] = "blob-as-descriptor"; const char Options::BLOB_FIELD[] = "blob-field"; const char Options::BLOB_DESCRIPTOR_FIELD[] = "blob-descriptor-field"; diff --git a/src/paimon/common/types/data_type.cpp b/src/paimon/common/types/data_type.cpp index e3004a258..72137b89f 100644 --- a/src/paimon/common/types/data_type.cpp +++ b/src/paimon/common/types/data_type.cpp @@ -23,6 +23,7 @@ #include "arrow/util/checked_cast.h" #include "fmt/format.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/array_type.h" #include "paimon/common/types/map_type.h" #include "paimon/common/types/row_type.h" @@ -49,6 +50,11 @@ std::unique_ptr DataType::Create( case arrow::Type::type::LIST: return std::make_unique(type, nullable, metadata); case arrow::Type::type::STRUCT: + if (VariantTypeUtils::IsVariantMetadata(metadata)) { + // A variant field is physically a struct but is a scalar + // VARIANT type in the paimon type system, not a ROW type. + return std::unique_ptr(new DataType(type, nullable, metadata)); + } return std::make_unique(type, nullable, metadata); default: return std::unique_ptr(new DataType(type, nullable, metadata)); @@ -116,9 +122,16 @@ std::string DataType::DataTypeToString(const std::shared_ptr& t arrow::internal::checked_pointer_cast(type); return TimestampToString(timestamp_type); } + case arrow::Type::type::STRUCT: { + if (VariantTypeUtils::IsVariantMetadata(metadata_)) { + return "VARIANT"; + } + [[fallthrough]]; + } case arrow::Type::type::LARGE_BINARY: { // TODO(xinyu): change binary to large binary? - if (BlobUtils::IsBlobMetadata(metadata_)) { + if (type->id() == arrow::Type::type::LARGE_BINARY && + BlobUtils::IsBlobMetadata(metadata_)) { return "BLOB"; } [[fallthrough]]; diff --git a/src/paimon/common/types/data_type_json_parser.cpp b/src/paimon/common/types/data_type_json_parser.cpp index 58f9405b4..21fa21c0f 100644 --- a/src/paimon/common/types/data_type_json_parser.cpp +++ b/src/paimon/common/types/data_type_json_parser.cpp @@ -28,6 +28,7 @@ #include "fmt/format.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/rapidjson_util.h" @@ -80,9 +81,16 @@ struct Token { std::string value; }; +// Extension type attributes of a parsed atomic type. BLOB and VARIANT parse to plain arrow +// types (large_binary / struct) and need field-level metadata markers applied by the caller. +struct AtomicTypeAttributes { + bool is_blob = false; + bool is_variant = false; +}; + // nullptr is returned in the case of parsing failed Result> ParseAtomicType(const std::string& str, bool* nullable, - bool* is_blob); + AtomicTypeAttributes* attributes); std::vector Tokenize(const std::string& chars); bool IsWhitespace(char character); bool IsDelimiter(char character); @@ -136,6 +144,7 @@ enum class Keyword : int32_t { MAP, ROW, BLOB, + VARIANT, // NULL is keyword in c++ NULL_, RAW, @@ -184,6 +193,7 @@ const std::map& Keywords() { {"MAP", Keyword::MAP}, {"ROW", Keyword::ROW}, {"BLOB", Keyword::BLOB}, + {"VARIANT", Keyword::VARIANT}, {"NULL", Keyword::NULL_}, {"RAW", Keyword::RAW}, {"LEGACY", Keyword::LEGACY}, @@ -196,7 +206,8 @@ class TokenParser { TokenParser(const std::string& input_string, const std::vector& tokens) : input_string_(input_string), tokens_(tokens) {} - Result> ParseTokens(bool* nullable, bool* is_blob); + Result> ParseTokens(bool* nullable, + AtomicTypeAttributes* attributes); private: inline const Token& GetToken() const { @@ -225,9 +236,9 @@ class TokenParser { bool HasNextToken(const std::vector& types) const; bool HasNextToken(const std::vector& keywords) const; Result ParseNullability(); - Result> ParseTypeWithNullability(bool* nullable, - bool* is_blob); - Result> ParseTypeByKeyword(bool* is_blob); + Result> ParseTypeWithNullability( + bool* nullable, AtomicTypeAttributes* attributes); + Result> ParseTypeByKeyword(AtomicTypeAttributes* attributes); Result ParseStringLength(); template Result> ParseStringType(); @@ -245,11 +256,11 @@ class TokenParser { }; Result> ParseAtomicType(const std::string& str, bool* nullable, - bool* is_blob) { + AtomicTypeAttributes* attributes) { try { std::vector tokens = Tokenize(str); TokenParser converter(str, tokens); - return converter.ParseTokens(nullable, is_blob); + return converter.ParseTokens(nullable, attributes); } catch (...) { return Status::Invalid("parse atomic type failed."); } @@ -371,9 +382,10 @@ int32_t ConsumeIdentifier(const std::string& chars, int32_t cursor, std::ostring return cursor - 1; } -Result> TokenParser::ParseTokens(bool* nullable, bool* is_blob) { +Result> TokenParser::ParseTokens( + bool* nullable, AtomicTypeAttributes* attributes) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr type, - ParseTypeWithNullability(nullable, is_blob)); + ParseTypeWithNullability(nullable, attributes)); if (HasRemainingTokens()) { PAIMON_RETURN_NOT_OK(NextToken()); return Status::Invalid(fmt::format("Unexpected token: {}", GetToken().value)); @@ -452,9 +464,10 @@ Result TokenParser::ParseNullability() { return true; } -Result> TokenParser::ParseTypeWithNullability(bool* nullable, - bool* is_blob) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_type, ParseTypeByKeyword(is_blob)); +Result> TokenParser::ParseTypeWithNullability( + bool* nullable, AtomicTypeAttributes* attributes) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_type, + ParseTypeByKeyword(attributes)); PAIMON_ASSIGN_OR_RAISE(*nullable, ParseNullability()); // special case: suffix notation for ARRAY types if (HasNextToken({Keyword::ARRAY}) || HasNextToken({Keyword::MULTISET})) { @@ -463,7 +476,8 @@ Result> TokenParser::ParseTypeWithNullability(b return data_type; } -Result> TokenParser::ParseTypeByKeyword(bool* is_blob) { +Result> TokenParser::ParseTypeByKeyword( + AtomicTypeAttributes* attributes) { PAIMON_RETURN_NOT_OK(NextToken(TokenType::KEYWORD)); switch (TokenAsKeyword()) { case Keyword::CHAR: @@ -475,9 +489,13 @@ Result> TokenParser::ParseTypeByKeyword(bool* i case Keyword::BYTES: return arrow::binary(); case Keyword::BLOB: { - *is_blob = true; + attributes->is_blob = true; return arrow::large_binary(); } + case Keyword::VARIANT: { + attributes->is_variant = true; + return VariantTypeUtils::UnshreddedStructType(); + } case Keyword::STRING: return arrow::utf8(); case Keyword::BOOLEAN: @@ -612,11 +630,13 @@ Result> DataTypeJsonParser::ParseType( Result> DataTypeJsonParser::ParseAtomicTypeField( const std::string& name, const rapidjson::Value& type_json_value) { bool nullable = true; - bool is_blob = false; + AtomicTypeAttributes attributes; PAIMON_ASSIGN_OR_RAISE(std::shared_ptr type, - ParseAtomicType(type_json_value.GetString(), &nullable, &is_blob)); - if (is_blob) { + ParseAtomicType(type_json_value.GetString(), &nullable, &attributes)); + if (attributes.is_blob) { return BlobUtils::ToArrowField(name, nullable); + } else if (attributes.is_variant) { + return VariantTypeUtils::ToArrowField(name, nullable); } else { return arrow::field(name, type, nullable); } diff --git a/src/paimon/common/types/data_type_json_parser_test.cpp b/src/paimon/common/types/data_type_json_parser_test.cpp index 024335665..b8905d5c4 100644 --- a/src/paimon/common/types/data_type_json_parser_test.cpp +++ b/src/paimon/common/types/data_type_json_parser_test.cpp @@ -20,6 +20,7 @@ #include #include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -135,6 +136,17 @@ TEST(DataTypeJsonParserTest, ParseTypeAtomicTypeSuccess) { ASSERT_TRUE(field->type()->Equals(test_case.second)); } + // VARIANT parses to a variant-marked struct field. + for (const char* variant_str : {"VARIANT", "VARIANT NOT NULL"}) { + rapidjson::Document doc; + rapidjson::Value value(variant_str, doc.GetAllocator()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr field, + DataTypeJsonParser::ParseType("variant_field", value)); + ASSERT_TRUE(VariantTypeUtils::IsVariantField(field)); + ASSERT_EQ(field->nullable(), std::string(variant_str) == "VARIANT"); + ASSERT_TRUE(field->type()->Equals(VariantTypeUtils::UnshreddedStructType())); + } + // Invalid case { rapidjson::Document invalid_doc; diff --git a/src/paimon/common/types/data_type_test.cpp b/src/paimon/common/types/data_type_test.cpp index 2f0d99101..c565c0067 100644 --- a/src/paimon/common/types/data_type_test.cpp +++ b/src/paimon/common/types/data_type_test.cpp @@ -21,6 +21,7 @@ #include "arrow/api.h" #include "gtest/gtest.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/date_time_utils.h" namespace paimon::test { @@ -62,6 +63,19 @@ TEST(DataTypeTest, DataTypeToString) { DataType blob_type(blob_field->type(), blob_field->nullable(), blob_field->metadata()); ASSERT_EQ(blob_type.DataTypeToString(blob_field->type()), "BLOB"); } + { + std::shared_ptr variant_field = VariantTypeUtils::ToArrowField("f3_variant"); + DataType variant_type(variant_field->type(), variant_field->nullable(), + variant_field->metadata()); + ASSERT_EQ(variant_type.DataTypeToString(variant_field->type()), "VARIANT"); + // A variant field is a scalar VARIANT type, not a ROW type. + auto created = DataType::Create(variant_field->type(), variant_field->nullable(), + variant_field->metadata()); + rapidjson::Document doc; + auto json_value = created->ToJson(&doc.GetAllocator()); + ASSERT_TRUE(json_value.IsString()); + ASSERT_EQ(std::string(json_value.GetString()), "VARIANT"); + } ASSERT_EQ(dummy_data_type.DataTypeToString(arrow::date32()), "DATE"); auto decimal_type1 = arrow::decimal128(10, 2); diff --git a/src/paimon/common/utils/field_type_utils.h b/src/paimon/common/utils/field_type_utils.h index dcb0d5023..a812aaf2c 100644 --- a/src/paimon/common/utils/field_type_utils.h +++ b/src/paimon/common/utils/field_type_utils.h @@ -23,6 +23,7 @@ #include "arrow/api.h" #include "arrow/type_fwd.h" #include "fmt/format.h" +#include "paimon/common/data/variant/variant_defs.h" #include "paimon/defs.h" #include "paimon/result.h" #include "paimon/status.h" @@ -46,6 +47,18 @@ class FieldTypeUtils { (type == FieldType::BIGINT); } + /// Converts an arrow field to a `FieldType`, disambiguating metadata-marked extension types + /// (a VARIANT field is physically a STRUCT with the variant metadata marker). + static Result ConvertToFieldType(const std::shared_ptr& field) { + if (field->type()->id() == arrow::Type::type::STRUCT && field->HasMetadata()) { + auto extension_name = field->metadata()->Get(VariantDefs::kExtensionTypeKey); + if (extension_name.ok() && *extension_name == VariantDefs::kExtensionTypeValue) { + return FieldType::VARIANT; + } + } + return ConvertToFieldType(field->type()->id()); + } + static Result ConvertToFieldType(const arrow::Type::type& arrow_type) { switch (arrow_type) { case arrow::Type::type::BOOL: @@ -120,6 +133,8 @@ class FieldTypeUtils { return "MAP"; case FieldType::STRUCT: return "STRUCT"; + case FieldType::VARIANT: + return "VARIANT"; default: return "UNKNOWN, type id:" + std::to_string(static_cast(type)); } diff --git a/src/paimon/core/append/append_only_writer.cpp b/src/paimon/core/append/append_only_writer.cpp index 21e543008..43d1e2bd4 100644 --- a/src/paimon/core/append/append_only_writer.cpp +++ b/src/paimon/core/append/append_only_writer.cpp @@ -26,6 +26,7 @@ #include "arrow/c/helpers.h" #include "arrow/type.h" #include "arrow/util/key_value_metadata.h" +#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/types/row_kind.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -194,10 +195,11 @@ AppendOnlyWriter::RollingFileWriterResult AppendOnlyWriter::CreateRollingRowWrit AppendOnlyWriter::WriterFactory AppendOnlyWriter::GetDataFileWriterFactory( const std::shared_ptr& schema, const std::optional>& write_cols) const { - if (shredding_context_) { + if (auto plan_factory = ShreddingWritePlanFactories::SelectActive( + options_, schema, shredding_context_, memory_pool_)) { return std::make_shared( options_, schema_id_, schema, write_cols, seq_num_counter_, FileSource::Append(), - path_factory_, shredding_context_, memory_pool_); + path_factory_, plan_factory, memory_pool_); } return std::make_shared(options_, schema_id_, schema, write_cols, seq_num_counter_, FileSource::Append(), diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp index b98313256..bd040ad7f 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp @@ -1191,6 +1191,42 @@ Result CoreOptions::FieldCollectAggDistinct(const std::string& field_name) return distinct; } +std::optional CoreOptions::GetVariantShreddingSchema() const { + auto it = impl_->raw_options.find(Options::VARIANT_SHREDDING_SCHEMA); + if (it == impl_->raw_options.end()) { + it = impl_->raw_options.find(Options::PARQUET_VARIANT_SHREDDING_SCHEMA); + } + if (it == impl_->raw_options.end() || it->second.empty()) { + return std::nullopt; + } + return it->second; +} + +Result CoreOptions::VariantInferShreddingSchemaEnabled() const { + return OptionsUtils::GetValueFromMap(impl_->raw_options, + Options::VARIANT_INFER_SHREDDING_SCHEMA, false); +} + +Result CoreOptions::GetVariantShreddingMaxSchemaWidth() const { + return OptionsUtils::GetValueFromMap(impl_->raw_options, + Options::VARIANT_SHREDDING_MAX_SCHEMA_WIDTH, 300); +} + +Result CoreOptions::GetVariantShreddingMaxSchemaDepth() const { + return OptionsUtils::GetValueFromMap(impl_->raw_options, + Options::VARIANT_SHREDDING_MAX_SCHEMA_DEPTH, 50); +} + +Result CoreOptions::GetVariantShreddingMinFieldCardinalityRatio() const { + return OptionsUtils::GetValueFromMap( + impl_->raw_options, Options::VARIANT_SHREDDING_MIN_FIELD_CARDINALITY_RATIO, 0.1); +} + +Result CoreOptions::GetVariantShreddingMaxInferBufferRow() const { + return OptionsUtils::GetValueFromMap( + impl_->raw_options, Options::VARIANT_SHREDDING_MAX_INFER_BUFFER_ROW, 4096); +} + Result CoreOptions::GetMapStorageLayout(const std::string& field_name) const { std::string key = std::string(Options::FIELDS_PREFIX) + "." + field_name + "." + std::string(Options::MAP_STORAGE_LAYOUT); diff --git a/src/paimon/core/core_options.h b/src/paimon/core/core_options.h index b243c3c3d..4d50b6c42 100644 --- a/src/paimon/core/core_options.h +++ b/src/paimon/core/core_options.h @@ -128,6 +128,15 @@ class PAIMON_EXPORT CoreOptions { Result GetMapSharedShreddingColumnPlacementPolicy( const std::string& field_name) const; + /// The configured variant shredding schema JSON, if any (falls back to + /// "parquet.variant.shreddingSchema"). + std::optional GetVariantShreddingSchema() const; + Result VariantInferShreddingSchemaEnabled() const; + Result GetVariantShreddingMaxSchemaWidth() const; + Result GetVariantShreddingMaxSchemaDepth() const; + Result GetVariantShreddingMinFieldCardinalityRatio() const; + Result GetVariantShreddingMaxInferBufferRow() const; + bool DeletionVectorsEnabled() const; bool DeletionVectorsBitmap64() const; int64_t DeletionVectorTargetFileSize() const; diff --git a/src/paimon/core/io/infer_shredding_file_writer.h b/src/paimon/core/io/infer_shredding_file_writer.h new file mode 100644 index 000000000..3a278af4b --- /dev/null +++ b/src/paimon/core/io/infer_shredding_file_writer.h @@ -0,0 +1,175 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/data/shredding/shredding_write_plan_factory.h" +#include "paimon/core/io/single_file_writer.h" +#include "paimon/core/key_value.h" + +namespace paimon { + +/// A file writer that infers the shredding write plan from the first rows of the file. Incoming +/// batches are buffered until `ShreddingWritePlanFactory::InferBufferRowCount` rows have been +/// collected (or the writer is closed); the buffered batches are then sampled to create the +/// batch converter, the actual file writer is created with the resulting physical schema, and +/// the buffered batches are replayed into it. The file never rolls while buffering. +template +class InferShreddingFileWriter : public SingleFileWriter { + public: + using CreateInnerFn = std::function>>( + const std::shared_ptr&)>; + + InferShreddingFileWriter(const std::shared_ptr& logical_schema, + const std::shared_ptr& plan_factory, + const std::string& file_format_identifier, CreateInnerFn create_inner) + : SingleFileWriter(/*compression=*/"", std::function()), + logical_type_(arrow::struct_(logical_schema->fields())), + plan_factory_(plan_factory), + file_format_identifier_(file_format_identifier), + create_inner_(std::move(create_inner)) {} + + Status Write(T record) override { + if (plan_finalized_) { + return inner_->Write(std::move(record)); + } + PAIMON_RETURN_NOT_OK(Buffer(std::move(record))); + if (buffered_rows_ >= plan_factory_->InferBufferRowCount()) { + return FinalizePlanAndFlush(); + } + return Status::OK(); + } + + Status Close() override { + if (!plan_finalized_) { + PAIMON_RETURN_NOT_OK(FinalizePlanAndFlush()); + } + return inner_->Close(); + } + + Result GetResult() override { + if (!inner_) { + return Status::Invalid("Cannot access the result unless the writer is closed."); + } + return inner_->GetResult(); + } + + Result ReachTargetSize(bool suggested_check, int64_t target_size) override { + if (!plan_finalized_) { + // Never roll the file while rows are being buffered for inference. + return false; + } + return inner_->ReachTargetSize(suggested_check, target_size); + } + + Result::AbortExecutor> GetAbortExecutor() const override { + if (!inner_) { + return Status::Invalid("Writer should be closed!"); + } + return inner_->GetAbortExecutor(); + } + + std::string GetPath() const override { + return inner_ ? inner_->GetPath() : ""; + } + + void Abort() override { + if (inner_) { + inner_->Abort(); + } + } + + int64_t RecordCount() const override { + return plan_finalized_ ? inner_->RecordCount() : buffered_rows_; + } + + std::shared_ptr GetMetrics() const override { + return inner_ ? inner_->GetMetrics() : nullptr; + } + + private: + struct BufferedBatch { + std::shared_ptr batch; + // Holds the non-batch part of a KeyValueBatch record; unused for plain batches. + KeyValueBatch record_template; + }; + + Status Buffer(T record) { + BufferedBatch buffered; + ::ArrowArray* c_batch; + if constexpr (std::is_same_v) { + c_batch = record; + } else { + static_assert(std::is_same_v, + "InferShreddingFileWriter supports ::ArrowArray* and KeyValueBatch"); + buffered.record_template = std::move(record); + c_batch = buffered.record_template.batch.get(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(buffered.batch, + arrow::ImportArray(c_batch, logical_type_)); + buffered_rows_ += buffered.batch->length(); + buffered_batches_.push_back(std::move(buffered)); + return Status::OK(); + } + + Status FinalizePlanAndFlush() { + std::vector> samples; + samples.reserve(buffered_batches_.size()); + for (const auto& buffered : buffered_batches_) { + samples.push_back(buffered.batch); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr converter, + plan_factory_->CreateConverter(file_format_identifier_, samples)); + PAIMON_ASSIGN_OR_RAISE(inner_, create_inner_(converter)); + plan_finalized_ = true; + for (auto& buffered : buffered_batches_) { + if constexpr (std::is_same_v) { + ::ArrowArray c_batch; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*buffered.batch, &c_batch)); + PAIMON_RETURN_NOT_OK(inner_->Write(&c_batch)); + } else { + KeyValueBatch record = std::move(buffered.record_template); + record.batch = std::make_unique<::ArrowArray>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*buffered.batch, record.batch.get())); + PAIMON_RETURN_NOT_OK(inner_->Write(std::move(record))); + } + } + buffered_batches_.clear(); + return Status::OK(); + } + + std::shared_ptr logical_type_; + std::shared_ptr plan_factory_; + std::string file_format_identifier_; + CreateInnerFn create_inner_; + + std::vector buffered_batches_; + int64_t buffered_rows_ = 0; + bool plan_finalized_ = false; + std::unique_ptr> inner_; +}; + +} // namespace paimon diff --git a/src/paimon/core/io/infer_shredding_file_writer_test.cpp b/src/paimon/core/io/infer_shredding_file_writer_test.cpp new file mode 100644 index 000000000..4c758acd0 --- /dev/null +++ b/src/paimon/core/io/infer_shredding_file_writer_test.cpp @@ -0,0 +1,212 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "paimon/core/io/infer_shredding_file_writer.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_shredding_write_plan_factory.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/core_options.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/testing/utils/variant_test_data.h" + +namespace paimon::test { + +namespace { + +/// A file-less writer standing in for the actual data file writer: it applies the shredding +/// conversion like the injected converter lambda would and collects the written batches. +class CollectingFileWriter : public SingleFileWriter<::ArrowArray*, std::shared_ptr> { + public: + CollectingFileWriter(const std::shared_ptr& converter, + const std::shared_ptr& logical_type, + std::vector>* sink) + : SingleFileWriter("", std::function()), + converter_(converter), + logical_type_(logical_type), + sink_(sink) {} + + Status Write(::ArrowArray* record) override { + record_count_ += record->length; + std::shared_ptr array; + if (converter_) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowArray> physical, + converter_->Convert(record)); + auto physical_type = arrow::struct_(converter_->GetPhysicalSchema()->fields()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(array, + arrow::ImportArray(physical.get(), physical_type)); + } else { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(array, arrow::ImportArray(record, logical_type_)); + } + sink_->push_back(std::move(array)); + return Status::OK(); + } + + Status Close() override { + closed = true; + return Status::OK(); + } + + Result> GetResult() override { + return std::shared_ptr(nullptr); + } + + Result ReachTargetSize(bool suggested_check, int64_t target_size) override { + return false; + } + + Result GetAbortExecutor() const override { + return AbortExecutor(nullptr, ""); + } + + void Abort() override {} + + int64_t RecordCount() const override { + return record_count_; + } + + bool closed = false; + + private: + std::shared_ptr converter_; + std::shared_ptr logical_type_; + std::vector>* sink_; + int64_t record_count_ = 0; +}; + +} // namespace + +class InferShreddingFileWriterTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + std::vector fields = {DataField(1, arrow::field("id", arrow::int32())), + DataField(2, VariantTypeUtils::ToArrowField("v"))}; + schema_ = DataField::ConvertDataFieldsToArrowSchema(fields); + logical_type_ = arrow::struct_(schema_->fields()); + } + + std::shared_ptr BuildBatch(const std::vector& jsons) { + auto result = BuildVariantBatch(schema_->field(0), schema_->field(1), jsons, pool_); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return std::move(result).value(); + } + + Status WriteBatch(InferShreddingFileWriter<::ArrowArray*, std::shared_ptr>* w, + const std::vector& jsons) { + std::shared_ptr array = BuildBatch(jsons); + ::ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + return w->Write(&c_array); + } + + std::unique_ptr>> + MakeWriter(int32_t buffer_rows) { + std::map option_map = { + {"variant.inferShreddingSchema", "true"}, + {"variant.shredding.maxInferBufferRow", std::to_string(buffer_rows)}, + // Keep the manifest format resolvable in test binaries without the avro plugin. + {"manifest.format", "parquet"}}; + auto options = CoreOptions::FromMap(option_map); + EXPECT_TRUE(options.ok()) << options.status().ToString(); + auto plan_factory = + std::make_shared(options.value(), schema_, pool_); + auto create_inner = [this](const std::shared_ptr& converter) + -> Result< + std::unique_ptr>>> { + captured_converters_.push_back(converter); + auto writer = std::make_unique(converter, logical_type_, &sink_); + inner_ = writer.get(); + return std::unique_ptr>>( + std::move(writer)); + }; + return std::make_unique< + InferShreddingFileWriter<::ArrowArray*, std::shared_ptr>>( + schema_, plan_factory, "parquet", create_inner); + } + + protected: + std::shared_ptr pool_; + std::shared_ptr schema_; + std::shared_ptr logical_type_; + // The converters own the arrow pool backing the collected arrays; keep them declared first + // so the arrays are destroyed before the pool. + std::vector> captured_converters_; + std::vector> sink_; + CollectingFileWriter* inner_ = nullptr; +}; + +TEST_F(InferShreddingFileWriterTest, BuffersUntilThresholdThenReplays) { + auto writer = MakeWriter(/*buffer_rows=*/4); + // Never rolls while buffering. + ASSERT_OK_AND_ASSIGN(bool reach, writer->ReachTargetSize(true, 1)); + ASSERT_FALSE(reach); + + ASSERT_OK(WriteBatch( + writer.get(), {R"({"age": 35, "city": "Chicago"})", R"({"age": 25, "city": "Hangzhou"})"})); + ASSERT_TRUE(sink_.empty()); + ASSERT_EQ(writer->RecordCount(), 2); + + // Crossing the row threshold finalizes the plan and replays the buffered batches. + ASSERT_OK(WriteBatch( + writer.get(), {R"({"age": 18, "city": "Beijing"})", R"({"age": 60, "city": "Shanghai"})"})); + ASSERT_EQ(sink_.size(), 2); + ASSERT_EQ(captured_converters_.size(), 1); + ASSERT_NE(captured_converters_[0], nullptr); + const auto& physical_type = static_cast(*sink_[0]->type()); + const auto& variant_physical = + static_cast(*physical_type.GetFieldByName("v")->type()); + ASSERT_NE(variant_physical.GetFieldByName("typed_value"), nullptr); + + // Subsequent writes stream through the finalized writer directly. + ASSERT_OK(WriteBatch(writer.get(), {R"({"age": 1, "city": "Suzhou"})"})); + ASSERT_EQ(sink_.size(), 3); + ASSERT_EQ(writer->RecordCount(), 5); + + ASSERT_OK(writer->Close()); + ASSERT_TRUE(inner_->closed); +} + +TEST_F(InferShreddingFileWriterTest, CloseFlushesPartialBuffer) { + auto writer = MakeWriter(/*buffer_rows=*/100); + ASSERT_OK(WriteBatch(writer.get(), {R"({"age": 35, "city": "Chicago"})"})); + ASSERT_TRUE(sink_.empty()); + ASSERT_OK(writer->Close()); + ASSERT_EQ(sink_.size(), 1); + ASSERT_EQ(captured_converters_.size(), 1); + ASSERT_NE(captured_converters_[0], nullptr); + ASSERT_TRUE(inner_->closed); +} + +TEST_F(InferShreddingFileWriterTest, EmptyFileFallsBackToLogicalSchema) { + auto writer = MakeWriter(/*buffer_rows=*/4); + ASSERT_OK(writer->Close()); + // With no samples there is no useful shredding schema; the writer is created without a + // converter. + ASSERT_EQ(captured_converters_.size(), 1); + ASSERT_EQ(captured_converters_[0], nullptr); + ASSERT_TRUE(sink_.empty()); + ASSERT_TRUE(inner_->closed); +} + +} // namespace paimon::test diff --git a/src/paimon/core/io/rolling_file_writer.h b/src/paimon/core/io/rolling_file_writer.h index 7043fd4e0..ad5217f29 100644 --- a/src/paimon/core/io/rolling_file_writer.h +++ b/src/paimon/core/io/rolling_file_writer.h @@ -154,8 +154,10 @@ Status RollingFileWriter::CloseCurrentWriter() { if (current_writer_ == nullptr) { return Status::OK(); } - std::shared_ptr current_metrics = current_writer_->GetMetrics(); PAIMON_RETURN_NOT_OK(current_writer_->Close()); + // Read the metrics after Close(): writers that create their inner writer lazily (e.g. + // inferred shredding) only expose metrics once closed. + std::shared_ptr current_metrics = current_writer_->GetMetrics(); PAIMON_ASSIGN_OR_RAISE(auto abort_executor, current_writer_->GetAbortExecutor()); closed_writers_.push_back(abort_executor); PAIMON_ASSIGN_OR_RAISE(R result, current_writer_->GetResult()); diff --git a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp index 4bc98f6f5..2735442a2 100644 --- a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp @@ -19,13 +19,11 @@ #include #include "arrow/c/helpers.h" -#include "paimon/common/data/shredding/map_shared_shredding_batch_converter.h" -#include "paimon/common/data/shredding/map_shared_shredding_context.h" -#include "paimon/common/data/shredding/map_shared_shredding_utils.h" -#include "paimon/common/data/shredding/map_shredding_defs.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_file_writer.h" +#include "paimon/core/io/infer_shredding_file_writer.h" +#include "paimon/format/file_format.h" #include "paimon/fs/file_system.h" namespace paimon { @@ -36,20 +34,39 @@ ShreddingAppendDataFileWriterFactory::ShreddingAppendDataFileWriterFactory( const std::optional>& write_cols, const std::shared_ptr& seq_num_counter, FileSource file_source, const std::shared_ptr& path_factory, - const std::shared_ptr& shredding_context, + const std::shared_ptr& plan_factory, const std::shared_ptr& pool) : AppendDataFileWriterFactory(options, schema_id, write_schema, write_cols, seq_num_counter, file_source, path_factory, pool), - shredding_context_(shredding_context) {} + plan_factory_(plan_factory) {} Result>>> ShreddingAppendDataFileWriterFactory::CreateWriter() const { - if (!shredding_context_) { - return Status::Invalid("Shared-shredding append writer requires a shredding context."); + if (!plan_factory_) { + return Status::Invalid("Shredding append writer requires a write-plan factory."); + } + const std::string format_identifier = options_.GetFileFormat()->Identifier(); + if (plan_factory_->ShouldInferWritePlan()) { + auto create_inner = [this](const std::shared_ptr& converter) { + return CreateShreddedWriter(converter); + }; + return std::make_unique< + InferShreddingFileWriter<::ArrowArray*, std::shared_ptr>>( + write_schema_, plan_factory_, format_identifier, std::move(create_inner)); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr converter, + plan_factory_->CreateConverter(format_identifier, /*sample_batches=*/{})); + return CreateShreddedWriter(converter); +} + +Result>>> +ShreddingAppendDataFileWriterFactory::CreateShreddedWriter( + const std::shared_ptr& converter) const { + if (converter == nullptr) { + // No conversion is useful for this file; fall back to the plain writer. + return AppendDataFileWriterFactory::CreateWriter(); } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr converter, - MapSharedShreddingBatchConverter::Create( - write_schema_, shredding_context_, options_, pool_)); std::shared_ptr file_schema = converter->GetPhysicalSchema(); std::function batch_converter = [converter](::ArrowArray* input, ::ArrowArray* output) -> Status { @@ -66,9 +83,11 @@ ShreddingAppendDataFileWriterFactory::CreateWriter() const { pool_); PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); - writer->SetMetadataFinalizer(MapSharedShreddingUtils::BuildMetadataFinalizer( - converter, MapSharedShreddingDefine::kDefaultDictCompression, shredding_context_, - file_schema)); + ShreddingWritePlanFactory::MetadataFinalizer finalizer = + plan_factory_->CreateMetadataFinalizer(converter); + if (finalizer) { + writer->SetMetadataFinalizer(std::move(finalizer)); + } return std::unique_ptr>>( std::move(writer)); } diff --git a/src/paimon/core/io/shredding_append_data_file_writer_factory.h b/src/paimon/core/io/shredding_append_data_file_writer_factory.h index 299e1b5ec..589ce6a95 100644 --- a/src/paimon/core/io/shredding_append_data_file_writer_factory.h +++ b/src/paimon/core/io/shredding_append_data_file_writer_factory.h @@ -21,15 +21,18 @@ #include #include +#include "paimon/common/data/shredding/shredding_write_plan_factory.h" #include "paimon/core/io/append_data_file_writer_factory.h" namespace paimon { class DataFilePathFactory; class LongCounter; -class MapSharedShreddingContext; class MemoryPool; +/// Creates append data file writers that rewrite logical batches into a physical (shredded) +/// layout as planned by a `ShreddingWritePlanFactory` (MAP shared-shredding or VARIANT +/// shredding, configured or inferred). class ShreddingAppendDataFileWriterFactory : public AppendDataFileWriterFactory { public: ShreddingAppendDataFileWriterFactory( @@ -38,14 +41,17 @@ class ShreddingAppendDataFileWriterFactory : public AppendDataFileWriterFactory const std::optional>& write_cols, const std::shared_ptr& seq_num_counter, FileSource file_source, const std::shared_ptr& path_factory, - const std::shared_ptr& shredding_context, + const std::shared_ptr& plan_factory, const std::shared_ptr& pool); Result>>> CreateWriter() const override; private: - std::shared_ptr shredding_context_; + Result>>> + CreateShreddedWriter(const std::shared_ptr& converter) const; + + std::shared_ptr plan_factory_; }; } // namespace paimon diff --git a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp index 24f888326..129900a94 100644 --- a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.cpp @@ -16,16 +16,12 @@ #include "paimon/core/io/shredding_key_value_data_file_writer_factory.h" -#include #include #include "arrow/c/helpers.h" -#include "paimon/common/data/shredding/map_shared_shredding_batch_converter.h" -#include "paimon/common/data/shredding/map_shared_shredding_context.h" -#include "paimon/common/data/shredding/map_shared_shredding_utils.h" -#include "paimon/common/data/shredding/map_shredding_defs.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/io/infer_shredding_file_writer.h" #include "paimon/core/io/key_value_data_file_writer.h" #include "paimon/format/file_format.h" #include "paimon/fs/file_system.h" @@ -37,20 +33,40 @@ ShreddingKeyValueDataFileWriterFactory::ShreddingKeyValueDataFileWriterFactory( const std::shared_ptr& write_schema, int32_t level, FileSource file_source, const std::vector& primary_keys, const std::shared_ptr& path_factory, bool create_stats_extractor, - const std::shared_ptr& shredding_context, + const std::shared_ptr& plan_factory, const std::shared_ptr& pool) : KeyValueDataFileWriterFactory(options, schema_id, write_schema, level, file_source, primary_keys, path_factory, create_stats_extractor, pool), - shredding_context_(shredding_context) {} + plan_factory_(plan_factory) {} Result>>> ShreddingKeyValueDataFileWriterFactory::CreateWriter() const { - if (!shredding_context_) { - return Status::Invalid("Shared-shredding key-value writer requires a shredding context."); + if (!plan_factory_) { + return Status::Invalid("Shredding key-value writer requires a write-plan factory."); } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr converter, - MapSharedShreddingBatchConverter::Create( - write_schema_, shredding_context_, options_, pool_)); + const std::string format_identifier = options_.GetWriteFileFormat(level_)->Identifier(); + if (plan_factory_->ShouldInferWritePlan()) { + auto create_inner = [this](const std::shared_ptr& converter) { + return CreateShreddedWriter(converter); + }; + return std::make_unique< + InferShreddingFileWriter>>( + write_schema_, plan_factory_, format_identifier, std::move(create_inner)); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr converter, + plan_factory_->CreateConverter(format_identifier, /*sample_batches=*/{})); + return CreateShreddedWriter(converter); +} + +Result>>> +ShreddingKeyValueDataFileWriterFactory::CreateShreddedWriter( + const std::shared_ptr& converter) const { + if (converter == nullptr) { + // No conversion is useful for this file; fall back to the plain writer. + return KeyValueDataFileWriterFactory::CreateWriter(); + } + auto format = options_.GetWriteFileFormat(level_); std::shared_ptr file_schema = converter->GetPhysicalSchema(); std::function batch_converter = [converter](KeyValueBatch key_value_batch, ::ArrowArray* array) -> Status { @@ -59,8 +75,6 @@ ShreddingKeyValueDataFileWriterFactory::CreateWriter() const { ArrowArrayMove(physical.get(), array); return Status::OK(); }; - - auto format = options_.GetWriteFileFormat(level_); PAIMON_ASSIGN_OR_RAISE(WriterResources resources, CreateWriterResources(*format, file_schema, create_stats_extractor_)); auto writer = std::make_unique( @@ -69,9 +83,11 @@ ShreddingKeyValueDataFileWriterFactory::CreateWriter() const { path_factory_->IsExternalPath(), pool_); PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); - writer->SetMetadataFinalizer(MapSharedShreddingUtils::BuildMetadataFinalizer( - converter, MapSharedShreddingDefine::kDefaultDictCompression, shredding_context_, - file_schema)); + ShreddingWritePlanFactory::MetadataFinalizer finalizer = + plan_factory_->CreateMetadataFinalizer(converter); + if (finalizer) { + writer->SetMetadataFinalizer(std::move(finalizer)); + } return std::unique_ptr>>( std::move(writer)); } diff --git a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h index 4759fb28d..876287eff 100644 --- a/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h +++ b/src/paimon/core/io/shredding_key_value_data_file_writer_factory.h @@ -20,14 +20,17 @@ #include #include +#include "paimon/common/data/shredding/shredding_write_plan_factory.h" #include "paimon/core/io/key_value_data_file_writer_factory.h" namespace paimon { class DataFilePathFactory; -class MapSharedShreddingContext; class MemoryPool; +/// Creates key-value data file writers that rewrite logical batches into a physical (shredded) +/// layout as planned by a `ShreddingWritePlanFactory` (MAP shared-shredding or VARIANT +/// shredding, configured or inferred). class ShreddingKeyValueDataFileWriterFactory : public KeyValueDataFileWriterFactory { public: ShreddingKeyValueDataFileWriterFactory( @@ -35,14 +38,17 @@ class ShreddingKeyValueDataFileWriterFactory : public KeyValueDataFileWriterFact const std::shared_ptr& write_schema, int32_t level, FileSource file_source, const std::vector& primary_keys, const std::shared_ptr& path_factory, bool create_stats_extractor, - const std::shared_ptr& shredding_context, + const std::shared_ptr& plan_factory, const std::shared_ptr& pool); Result>>> CreateWriter() const override; private: - std::shared_ptr shredding_context_; + Result>>> + CreateShreddedWriter(const std::shared_ptr& converter) const; + + std::shared_ptr plan_factory_; }; } // namespace paimon diff --git a/src/paimon/core/io/single_file_writer.h b/src/paimon/core/io/single_file_writer.h index 0d66a571e..24ba08bb5 100644 --- a/src/paimon/core/io/single_file_writer.h +++ b/src/paimon/core/io/single_file_writer.h @@ -104,16 +104,16 @@ class SingleFileWriter : public FileWriter { return nullptr; } - Result ReachTargetSize(bool suggested_check, int64_t target_size); + virtual Result ReachTargetSize(bool suggested_check, int64_t target_size); - Result GetAbortExecutor() const { + virtual Result GetAbortExecutor() const { if (closed_ == false) { return Status::Invalid("Writer should be closed!"); } return AbortExecutor(fs_, path_); } - std::string GetPath() const { + virtual std::string GetPath() const { return path_; } diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp index e0c16af55..fc914a5b0 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_rewriter.cpp @@ -21,6 +21,7 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/common/data/shredding/map_shared_shredding_utils.h" +#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_writer_factory.h" @@ -152,11 +153,12 @@ MergeTreeCompactRewriter::CreateRollingRowWriter(int32_t level) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, CreateDataFilePathFactory(format->Identifier())); std::shared_ptr>> factory; - if (shredding_context_) { + if (auto plan_factory = ShreddingWritePlanFactories::SelectActive(options_, write_schema_, + shredding_context_, pool_)) { factory = std::make_shared( options_, schema_id_, write_schema_, level, FileSource::Compact(), trimmed_primary_keys_, data_file_path_factory, /*create_stats_extractor=*/true, - shredding_context_, pool_); + plan_factory, pool_); } else { factory = std::make_shared( options_, schema_id_, write_schema_, level, FileSource::Compact(), diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index 00130938c..ef82c86d9 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -25,6 +25,7 @@ #include "arrow/api.h" #include "arrow/c/abi.h" #include "arrow/c/helpers.h" +#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -321,11 +322,12 @@ Result MergeTreeWriter::DrainIncrement() { std::unique_ptr>> MergeTreeWriter::CreateRollingRowWriter() const { std::shared_ptr>> factory; - if (shredding_context_) { + if (auto plan_factory = ShreddingWritePlanFactories::SelectActive(options_, write_schema_, + shredding_context_, pool_)) { factory = std::make_shared( options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), - trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/true, - shredding_context_, pool_); + trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/true, plan_factory, + pool_); } else { factory = std::make_shared( options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 1e70c0554..9a8700d81 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -27,6 +27,9 @@ #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/shredding/map_shared_shredding_file_reader.h" #include "paimon/common/data/shredding/map_shared_shredding_utils.h" +#include "paimon/common/data/shredding/shredding_file_reader.h" +#include "paimon/common/data/variant/variant_shredding_read_plan_factory.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/reader/delegating_prefetch_reader.h" #include "paimon/common/reader/predicate_batch_reader.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" @@ -203,6 +206,8 @@ Result> AbstractSplitRead::CreateFieldMappingRe std::move(file_reader), read_schema)); file_reader = std::move(shared_shredding_result.first); skip_map_selected_keys_filter_field_ids = std::move(shared_shredding_result.second); + PAIMON_ASSIGN_OR_RAISE( + file_reader, ApplyVariantShreddingReaderIfNeeded(std::move(file_reader), read_schema)); } if (NeedCompleteRowTrackingFields(options_.RowTrackingEnabled(), read_schema)) { file_reader = std::make_unique( @@ -281,6 +286,32 @@ AbstractSplitRead::ApplySharedShreddingReaderIfNeeded( return std::make_pair(std::move(file_reader), std::move(handled_shared_shredding_field_ids)); } +Result> AbstractSplitRead::ApplyVariantShreddingReaderIfNeeded( + std::unique_ptr&& file_reader, + const std::shared_ptr& read_schema) const { + bool has_variant_field = false; + for (const auto& read_field : read_schema->fields()) { + if (VariantTypeUtils::IsVariantField(read_field)) { + has_variant_field = true; + break; + } + } + if (!has_variant_field) { + return std::move(file_reader); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, + file_reader->GetFileSchema()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr file_arrow_schema, + arrow::ImportSchema(file_schema.get())); + PAIMON_ASSIGN_OR_RAISE(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( + read_schema, file_arrow_schema, pool_)); + if (!plans.empty()) { + file_reader = + std::make_unique(std::move(file_reader), std::move(plans), pool_); + } + return std::move(file_reader); +} + Result> AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( const std::shared_ptr& data_schema, const std::optional>& write_cols) { diff --git a/src/paimon/core/operation/abstract_split_read.h b/src/paimon/core/operation/abstract_split_read.h index 4d49d78bc..89c94fbf4 100644 --- a/src/paimon/core/operation/abstract_split_read.h +++ b/src/paimon/core/operation/abstract_split_read.h @@ -113,6 +113,13 @@ class AbstractSplitRead : public SplitRead { ApplySharedShreddingReaderIfNeeded(std::unique_ptr&& file_reader, const std::shared_ptr& read_schema) const; + /// Wraps the reader with a `ShreddingFileReader` when any read variant column needs + /// reassembly or path extraction; a plain read of an unshredded variant column is passed + /// through untouched. + Result> ApplyVariantShreddingReaderIfNeeded( + std::unique_ptr&& file_reader, + const std::shared_ptr& read_schema) const; + static bool NeedCompleteRowTrackingFields(bool row_tracking_enabled, const std::shared_ptr& read_schema); diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index d59c3b980..41db77b20 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -22,6 +22,7 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/core/append/append_only_writer.h" @@ -227,10 +228,11 @@ AppendOnlyFileStoreWrite::WriterFactory AppendOnlyFileStoreWrite::GetDataFileWri const std::vector>& to_compact, const std::shared_ptr& shredding_context) const { auto seq_num_counter = std::make_shared(to_compact[0]->min_sequence_number); - if (shredding_context) { + if (auto plan_factory = + ShreddingWritePlanFactories::SelectActive(options_, schema, shredding_context, pool_)) { return std::make_shared( options_, table_schema_->Id(), schema, write_cols, seq_num_counter, - FileSource::Compact(), data_file_path_factory, shredding_context, pool_); + FileSource::Compact(), data_file_path_factory, plan_factory, pool_); } return std::make_shared( options_, table_schema_->Id(), schema, write_cols, seq_num_counter, FileSource::Compact(), diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index 504c598e5..c935e4140 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -23,6 +23,8 @@ #include "arrow/c/abi.h" #include "arrow/c/bridge.h" #include "fmt/format.h" +#include "paimon/common/data/variant/variant_access_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/predicate/predicate_validator.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" @@ -38,6 +40,14 @@ Result> InternalReadContext::AlignReadFieldWithTab const std::shared_ptr& table_field) { static const std::vector kReadMetadataWhitelist = {DataField::MAP_SELECTED_KEYS}; + if (VariantTypeUtils::IsVariantField(table_field) && + VariantAccessUtils::IsVariantAccessType(read_field->type())) { + // A variant column may be read as a variant-access projection: a struct whose children + // each carry a `__VARIANT_METADATA` description. Keep the projection type (including + // the children's descriptions) on the aligned field. + return table_field->WithType(read_field->type()); + } + if (read_field->type()->id() != table_field->type()->id()) { return Status::Invalid(fmt::format( "Read schema field '{}' type {} does not match table field type {}", read_field->name(), diff --git a/src/paimon/core/postpone/postpone_bucket_writer.cpp b/src/paimon/core/postpone/postpone_bucket_writer.cpp index 8cf2619a2..2934a8d00 100644 --- a/src/paimon/core/postpone/postpone_bucket_writer.cpp +++ b/src/paimon/core/postpone/postpone_bucket_writer.cpp @@ -30,6 +30,7 @@ #include "arrow/scalar.h" #include "arrow/util/checked_cast.h" #include "fmt/format.h" +#include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" @@ -261,11 +262,12 @@ PostponeBucketWriter::PrepareMinMaxKey( std::unique_ptr>> PostponeBucketWriter::CreateRollingRowWriter() const { std::shared_ptr>> factory; - if (shredding_context_) { + if (auto plan_factory = ShreddingWritePlanFactories::SelectActive(options_, write_schema_, + shredding_context_, pool_)) { factory = std::make_shared( options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), - trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/false, - shredding_context_, pool_); + trimmed_primary_keys_, path_factory_, /*create_stats_extractor=*/false, plan_factory, + pool_); } else { factory = std::make_shared( options_, schema_id_, write_schema_, /*level=*/0, FileSource::Append(), diff --git a/src/paimon/core/schema/arrow_schema_validator.cpp b/src/paimon/core/schema/arrow_schema_validator.cpp index 117d468b2..c9053475c 100644 --- a/src/paimon/core/schema/arrow_schema_validator.cpp +++ b/src/paimon/core/schema/arrow_schema_validator.cpp @@ -23,6 +23,8 @@ #include "arrow/util/checked_cast.h" #include "fmt/format.h" #include "paimon/common/data/blob_utils.h" +#include "paimon/common/data/variant/variant_access_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/decimal_utils.h" #include "paimon/common/utils/string_utils.h" @@ -116,6 +118,11 @@ Status ArrowSchemaValidator::ValidateDataTypeWithFieldId( break; } case arrow::Type::type::STRUCT: { + if (VariantTypeUtils::IsVariantMetadata(key_value_metadata)) { + // A variant struct is a leaf type: its value/metadata children carry fixed + // paimon field ids 0/1 which must not join the global field id uniqueness check. + break; + } arrow::FieldVector sub_fields = arrow::internal::checked_cast(type.get())->fields(); for (const auto& sub_field : sub_fields) { @@ -182,6 +189,16 @@ Status ArrowSchemaValidator::ValidateField(const std::shared_ptr& break; } case arrow::Type::type::STRUCT: { + if (VariantTypeUtils::IsVariantField(field)) { + if (VariantAccessUtils::IsVariantAccessType(field->type())) { + // A variant column read as a variant-access projection keeps the variant + // marker but replaces the type with the projection struct; its children are + // cast targets validated by the variant read plans. + break; + } + PAIMON_RETURN_NOT_OK(VariantTypeUtils::ValidateVariantShape(field)); + break; + } arrow::FieldVector arrow_fields = arrow::internal::checked_cast(*field->type()).fields(); for (const auto& sub_field : arrow_fields) { diff --git a/src/paimon/core/schema/arrow_schema_validator_test.cpp b/src/paimon/core/schema/arrow_schema_validator_test.cpp index 1d4057450..310955622 100644 --- a/src/paimon/core/schema/arrow_schema_validator_test.cpp +++ b/src/paimon/core/schema/arrow_schema_validator_test.cpp @@ -21,6 +21,9 @@ #include "arrow/type.h" #include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_access_utils.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/testing/utils/testharness.h" @@ -270,6 +273,47 @@ TEST(ArrowSchemaValidatorTest, ValidateDataTypeWithFieldId) { } } +TEST(ArrowSchemaValidatorTest, ValidateVariantField) { + // A variant field validates as a leaf: its fixed child ids 0/1 do not join the global field + // id uniqueness check, even when top-level fields use ids 0/1. + { + std::vector fields = {DataField(0, arrow::field("f0", arrow::utf8())), + DataField(1, VariantTypeUtils::ToArrowField("v1")), + DataField(2, VariantTypeUtils::ToArrowField("v2"))}; + auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(fields); + ASSERT_OK(ArrowSchemaValidator::ValidateSchemaWithFieldId(*arrow_schema)) + << ArrowSchemaValidator::ValidateSchemaWithFieldId(*arrow_schema).ToString(); + } + // A variant-access projection (variant marker + description-carrying children) validates + // as a leaf instead of being shape-checked. + { + auto age_child = + arrow::field("0", arrow::int64(), true, + arrow::KeyValueMetadata::Make( + {DataField::DESCRIPTION}, + {VariantAccessUtils::BuildVariantMetadata("$.age", true, "UTC")})); + std::unordered_map metadata = { + {VariantDefs::kExtensionTypeKey, VariantDefs::kExtensionTypeValue}}; + auto access_field = arrow::field("v", arrow::struct_({age_child}), true, + std::make_shared(metadata)); + auto arrow_schema = arrow::schema({access_field}); + ASSERT_OK(ArrowSchemaValidator::ValidateSchema(*arrow_schema)) + << ArrowSchemaValidator::ValidateSchema(*arrow_schema).ToString(); + } + // A variant-marked struct with the wrong physical shape is rejected. + { + std::unordered_map metadata = { + {VariantDefs::kExtensionTypeKey, VariantDefs::kExtensionTypeValue}}; + auto bad_variant = + arrow::field("v", + arrow::struct_({arrow::field("value", arrow::binary(), true), + arrow::field("metadata", arrow::binary(), true)}), + true, std::make_shared(metadata)); + auto arrow_schema = arrow::schema({bad_variant}); + ASSERT_NOK(ArrowSchemaValidator::ValidateSchema(*arrow_schema)); + } +} + TEST(ArrowSchemaValidatorTest, ContainTimestampWithTimezone) { auto timezone = DateTimeUtils::GetLocalTimezoneName(); { diff --git a/src/paimon/core/schema/schema_validation.cpp b/src/paimon/core/schema/schema_validation.cpp index d6684752d..437c691ec 100644 --- a/src/paimon/core/schema/schema_validation.cpp +++ b/src/paimon/core/schema/schema_validation.cpp @@ -31,6 +31,7 @@ #include "fmt/ranges.h" #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/shredding/map_shared_shredding_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/object_utils.h" @@ -536,6 +537,12 @@ Status SchemaValidation::ValidateMapStorageLayout(const TableSchema& schema, if (layout != MapStorageLayout::SHARED_SHREDDING) { continue; } + for (const auto& field : schema.Fields()) { + if (VariantTypeUtils::ContainsVariantField(field.ArrowField())) { + return Status::Invalid( + "MAP shared-shredding currently cannot be used with Variant fields."); + } + } // Column configured with shared-shredding must be MAP if (!MapSharedShreddingUtils::IsShreddingKeyMap(field_type)) { return Status::Invalid( diff --git a/src/paimon/core/schema/table_schema.cpp b/src/paimon/core/schema/table_schema.cpp index 34b8af51e..b991d59f1 100644 --- a/src/paimon/core/schema/table_schema.cpp +++ b/src/paimon/core/schema/table_schema.cpp @@ -25,6 +25,7 @@ #include "arrow/c/bridge.h" #include "arrow/util/checked_cast.h" #include "fmt/format.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/field_type_utils.h" @@ -93,6 +94,11 @@ Result> TableSchema::AssignFieldIdsRecursively( } auto type = field->type(); if (type->id() == arrow::Type::STRUCT) { + if (VariantTypeUtils::IsVariantField(field)) { + // A variant struct is a leaf type: its value/metadata children keep their fixed + // paimon field ids 0/1 (mapped to parquet field ids on write). + return metadata ? field->WithMergedMetadata(metadata) : field; + } auto struct_type = arrow::internal::checked_pointer_cast(field->type()); arrow::FieldVector new_childs; for (const auto& child : struct_type->fields()) { @@ -187,7 +193,7 @@ std::vector TableSchema::FieldNames() const { Result TableSchema::GetFieldType(const std::string& field_name) const { PAIMON_ASSIGN_OR_RAISE(DataField field, GetField(field_name)); - return FieldTypeUtils::ConvertToFieldType(field.Type()->id()); + return FieldTypeUtils::ConvertToFieldType(field.ArrowField()); } Result TableSchema::GetField(const std::string& field_name) const { diff --git a/src/paimon/core/schema/table_schema_test.cpp b/src/paimon/core/schema/table_schema_test.cpp index f8bd9eeed..4f486fa19 100644 --- a/src/paimon/core/schema/table_schema_test.cpp +++ b/src/paimon/core/schema/table_schema_test.cpp @@ -21,6 +21,7 @@ #include "arrow/api.h" #include "arrow/util/checked_cast.h" #include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/fs/local/local_file_system.h" @@ -133,6 +134,20 @@ TEST_F(TableSchemaTest, TestCreateWithAllFieldsHaveFieldId) { ASSERT_EQ(table_schema->Options(), options); } +TEST_F(TableSchemaTest, TestGetFieldTypeForVariant) { + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + VariantTypeUtils::ToArrowField("v")}; + ASSERT_OK_AND_ASSIGN(auto table_schema, + TableSchema::Create(/*schema_id=*/0, arrow::schema(fields), + /*partition_keys=*/{}, /*primary_keys=*/{}, + /*options=*/{})); + ASSERT_OK_AND_ASSIGN(FieldType id_type, table_schema->GetFieldType("id")); + ASSERT_EQ(id_type, FieldType::INT); + // The variant marker must disambiguate the physical struct into FieldType::VARIANT. + ASSERT_OK_AND_ASSIGN(FieldType variant_type, table_schema->GetFieldType("v")); + ASSERT_EQ(variant_type, FieldType::VARIANT); +} + TEST_F(TableSchemaTest, TestInvalidCreate) { // partial fields have field id arrow::FieldVector fields = MakeArrowField(3); diff --git a/src/paimon/core/utils/nested_projection_utils.cpp b/src/paimon/core/utils/nested_projection_utils.cpp index d456b6239..8cfd73036 100644 --- a/src/paimon/core/utils/nested_projection_utils.cpp +++ b/src/paimon/core/utils/nested_projection_utils.cpp @@ -28,6 +28,8 @@ #include "arrow/array/concatenate.h" #include "arrow/type.h" #include "fmt/format.h" +#include "paimon/common/data/variant/variant_access_utils.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/status.h" @@ -81,6 +83,11 @@ Result NestedProjectionUtils::HasNestedSubfieldProjectionType( const std::shared_ptr& read_type) { switch (file_type->id()) { case arrow::Type::STRUCT: { + if (VariantAccessUtils::IsVariantAccessType(read_type)) { + // A variant-access projection is resolved by the variant read plans, not by + // nested subfield projection. + return false; + } if (read_type->id() != arrow::Type::STRUCT) { return Status::Invalid(fmt::format( "HasNestedSubfieldProjectionType requires same nested type kind, but file " @@ -151,6 +158,12 @@ Result>> NestedProjectionUtils::P switch (read_type->id()) { case arrow::Type::STRUCT: { + if (VariantAccessUtils::IsVariantAccessType(read_type) && + VariantTypeUtils::IsUnshreddedVariantType(data_type)) { + // A variant-access projection replaces the variant column type; pass it through + // so the read path extracts the described paths. + return std::optional>(read_type); + } arrow::FieldVector pruned_fields; for (const auto& read_child : read_type->fields()) { PAIMON_ASSIGN_OR_RAISE(int32_t read_child_id, GetPaimonFieldId(read_child)); diff --git a/src/paimon/format/avro/avro_schema_converter.cpp b/src/paimon/format/avro/avro_schema_converter.cpp index 54c22d927..7a1d41e70 100644 --- a/src/paimon/format/avro/avro_schema_converter.cpp +++ b/src/paimon/format/avro/avro_schema_converter.cpp @@ -29,6 +29,7 @@ #include "avro/Types.hh" #include "avro/ValidSchema.hh" #include "fmt/format.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/format/avro/avro_file_format_factory.h" #include "paimon/format/avro/avro_utils.h" @@ -378,6 +379,9 @@ Result<::avro::Schema> AvroSchemaConverter::ArrowTypeToAvroSchema( Result<::avro::ValidSchema> AvroSchemaConverter::ArrowSchemaToAvroSchema( const std::shared_ptr& arrow_schema) { + if (VariantTypeUtils::ContainsVariantField(arrow_schema)) { + return Status::NotImplemented("Avro format does not support the VARIANT type"); + } // top level row name of avro record, the same as java paimon static const std::string kTopLevelRowName = "org.apache.paimon.avro.generated.record"; ::avro::RecordSchema record_schema(kTopLevelRowName); diff --git a/src/paimon/format/avro/avro_schema_converter_test.cpp b/src/paimon/format/avro/avro_schema_converter_test.cpp index c18e4bcb8..de7b96feb 100644 --- a/src/paimon/format/avro/avro_schema_converter_test.cpp +++ b/src/paimon/format/avro/avro_schema_converter_test.cpp @@ -20,6 +20,7 @@ #include "avro/Compiler.hh" #include "avro/ValidSchema.hh" #include "gtest/gtest.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/core/manifest/manifest_file_meta.h" #include "paimon/core/utils/versioned_object_serializer.h" @@ -244,4 +245,11 @@ TEST(AvroSchemaConverterTest, TestAvroSchemaToArrowDataTypeWithTimestampType) { ASSERT_TRUE(arrow_type->Equals(arrow::struct_(expected_fields))) << arrow_type->ToString(); } +TEST(AvroSchemaConverterTest, TestVariantNotSupported) { + auto arrow_schema = + arrow::schema({arrow::field("id", arrow::int32()), VariantTypeUtils::ToArrowField("v")}); + auto result = AvroSchemaConverter::ArrowSchemaToAvroSchema(arrow_schema); + ASSERT_TRUE(result.status().IsNotImplemented()) << result.status().ToString(); +} + } // namespace paimon::avro::test diff --git a/src/paimon/format/orc/orc_format_writer.cpp b/src/paimon/format/orc/orc_format_writer.cpp index 805647f67..1cb4643d1 100644 --- a/src/paimon/format/orc/orc_format_writer.cpp +++ b/src/paimon/format/orc/orc_format_writer.cpp @@ -36,6 +36,7 @@ #include "orc/Type.hh" #include "orc/Vector.hh" #include "orc/Writer.hh" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/options/memory_size.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -78,6 +79,9 @@ Result> OrcFormatWriter::Create( const std::map& options, const std::string& compression, int32_t batch_size, const std::shared_ptr& pool) { assert(output_stream); + if (VariantTypeUtils::ContainsVariantField(arrow::schema(schema.fields()))) { + return Status::NotImplemented("ORC format does not support the VARIANT type"); + } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::orc::Type> orc_type, OrcAdapter::GetOrcType(schema)); auto data_type = arrow::struct_(schema.fields()); try { diff --git a/src/paimon/format/orc/orc_format_writer_test.cpp b/src/paimon/format/orc/orc_format_writer_test.cpp index 83fac679e..5d1f513fe 100644 --- a/src/paimon/format/orc/orc_format_writer_test.cpp +++ b/src/paimon/format/orc/orc_format_writer_test.cpp @@ -37,6 +37,7 @@ #include "orc/Type.hh" #include "orc/Vector.hh" #include "orc/Writer.hh" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/format/orc/orc_format_defs.h" #include "paimon/format/orc/orc_input_stream_impl.h" #include "paimon/format/orc/orc_metrics.h" @@ -295,6 +296,23 @@ TEST_F(OrcFormatWriterTest, TestPrepareWriterOptions) { "invalid config, do not support writing timestamp with timezone in legacy format"); } } +TEST_F(OrcFormatWriterTest, TestVariantNotSupported) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + std::string file_name = test_root + "/variant.orc"; + auto arrow_schema = + arrow::schema({arrow::field("id", arrow::int32()), VariantTypeUtils::ToArrowField("v")}); + std::map options = {}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + file_system_->Create(file_name, /*overwrite=*/true)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr output_stream, + OrcOutputStreamImpl::Create(out)); + auto result = OrcFormatWriter::Create(std::move(output_stream), *arrow_schema, options, + /*compression=*/"lz4", /*batch_size=*/16, pool_); + ASSERT_TRUE(result.status().IsNotImplemented()) << result.status().ToString(); +} + // TODO(liancheng.lsz): add tests for GetEstimateLength } // namespace paimon::orc::test diff --git a/src/paimon/format/parquet/CMakeLists.txt b/src/paimon/format/parquet/CMakeLists.txt index 6e72e0c69..46ba42374 100644 --- a/src/paimon/format/parquet/CMakeLists.txt +++ b/src/paimon/format/parquet/CMakeLists.txt @@ -62,6 +62,7 @@ if(PAIMON_BUILD_TESTS) predicate_converter_test.cpp predicate_pushdown_test.cpp column_index_filter_test.cpp + variant_parquet_test.cpp STATIC_LINK_LIBS paimon_shared test_utils_static diff --git a/src/paimon/format/parquet/variant_parquet_test.cpp b/src/paimon/format/parquet/variant_parquet_test.cpp new file mode 100644 index 000000000..fdead0289 --- /dev/null +++ b/src/paimon/format/parquet/variant_parquet_test.cpp @@ -0,0 +1,568 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/io/file.h" +#include "gtest/gtest.h" +#include "paimon/common/data/shredding/shredding_file_reader.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_defs.h" +#include "paimon/common/data/variant/variant_schema.h" +#include "paimon/common/data/variant/variant_shredding_batch_converter.h" +#include "paimon/common/data/variant/variant_shredding_read_plan_factory.h" +#include "paimon/common/data/variant/variant_shredding_utils.h" +#include "paimon/common/data/variant/variant_shredding_write_plan.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/data/variant.h" +#include "paimon/format/parquet/parquet_field_id_converter.h" +#include "paimon/format/parquet/parquet_file_batch_reader.h" +#include "paimon/format/parquet/parquet_format_defs.h" +#include "paimon/format/parquet/parquet_format_writer.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/testing/utils/variant_test_data.h" +#include "parquet/arrow/reader.h" +#include "parquet/file_reader.h" +#include "parquet/metadata.h" +#include "parquet/properties.h" +#include "parquet/schema.h" + +namespace paimon::parquet::test { + +class VariantParquetTest : public ::testing::Test { + public: + void SetUp() override { + dir_ = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir_); + fs_ = std::make_shared(); + pool_ = GetDefaultPool(); + arrow_pool_ = GetArrowPool(pool_); + file_path_ = PathUtil::JoinPath(dir_->Str(), "variant.parquet"); + + std::vector fields = {DataField(1, arrow::field("id", arrow::int32())), + DataField(2, VariantTypeUtils::ToArrowField("v"))}; + paimon_schema_ = DataField::ConvertDataFieldsToArrowSchema(fields); + } + + std::shared_ptr BuildArray(const std::vector& jsons) { + auto result = paimon::test::BuildVariantBatch(paimon_schema_->field(0), + paimon_schema_->field(1), jsons, pool_); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return std::move(result).value(); + } + + void WriteFile(const std::shared_ptr& array) { + // The production write path maps paimon field ids to parquet field ids. + ASSERT_OK_AND_ASSIGN(std::shared_ptr write_schema, + ParquetFieldIdConverter::AddParquetIdsFromPaimonIds(paimon_schema_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path_, /*overwrite=*/true)); + ::parquet::WriterProperties::Builder builder; + auto writer_properties = builder.build(); + ASSERT_OK_AND_ASSIGN( + auto format_writer, + ParquetFormatWriter::Create(out, write_schema, writer_properties, + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_)); + auto arrow_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok()); + ASSERT_OK(format_writer->AddBatch(arrow_array.get())); + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + } + + // Writes `jsons` shredded according to the configured ROW-type shredding schema JSON. + void WriteShreddedFile(const std::vector& jsons, + const char* shredding_schema_json) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr plan, + VariantShreddingWritePlan::FromConfiguredSchema(paimon_schema_, shredding_schema_json)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr converter, + VariantShreddingBatchConverter::Create(plan, pool_)); + auto logical = BuildArray(jsons); + auto c_logical = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*logical, c_logical.get()).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr c_physical, + converter->Convert(c_logical.get())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr write_schema, + ParquetFieldIdConverter::AddParquetIdsFromPaimonIds(converter->GetPhysicalSchema())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path_, /*overwrite=*/true)); + ::parquet::WriterProperties::Builder builder; + auto writer_properties = builder.build(); + ASSERT_OK_AND_ASSIGN( + auto format_writer, + ParquetFormatWriter::Create(out, write_schema, writer_properties, + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_)); + ASSERT_OK(format_writer->AddBatch(c_physical.get())); + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + } + + // Opens the written file and returns the reader plus its imported file schema. + void OpenFile(std::unique_ptr* file_reader, + std::shared_ptr* file_schema) { + ASSERT_OK_AND_ASSIGN(auto input_stream, fs_->Open(file_path_)); + auto length = fs_->GetFileStatus(file_path_).value()->GetLen(); + auto in_stream = + std::make_unique(std::move(input_stream), arrow_pool_, length); + std::map options = {}; + ASSERT_OK_AND_ASSIGN(auto parquet_reader, ParquetFileBatchReader::Create( + std::move(in_stream), options, + /*batch_size=*/1024, + /*file_metadata=*/nullptr, arrow_pool_)); + *file_reader = std::move(parquet_reader); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> c_file_schema, + (*file_reader)->GetFileSchema()); + auto imported = arrow::ImportSchema(c_file_schema.get()); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + *file_schema = imported.ValueOrDie(); + } + + // Builds a variant-access projection field for column "v" via the public builder. + std::shared_ptr BuildAccessField( + const std::vector, std::string>>& accesses) { + VariantAccessBuilder builder; + for (const auto& [type, path] : accesses) { + auto c_target = std::make_unique(); + EXPECT_TRUE(arrow::ExportField(arrow::Field("t", type), c_target.get()).ok()); + EXPECT_OK(builder.AddField(c_target.get(), path, /*fail_on_error=*/false)); + } + auto c_field = builder.Build("v"); + EXPECT_TRUE(c_field.ok()) << c_field.status().ToString(); + auto imported = arrow::ImportField(c_field.value().get()); + EXPECT_TRUE(imported.ok()) << imported.status().ToString(); + return imported.ValueOrDie(); + } + + // Reads the whole file through the shredding reader with the given read schema and returns + // the `v` column. + void ReadVariantColumn(const std::shared_ptr& read_schema, + std::shared_ptr* v_column) { + std::unique_ptr file_reader; + std::shared_ptr file_schema; + OpenFile(&file_reader, &file_schema); + ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( + read_schema, file_schema, pool_)); + ASSERT_EQ(plans.size(), 1); + auto shredding_reader = + std::make_unique(std::move(file_reader), std::move(plans), pool_); + auto c_read_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_read_schema.get()).ok()); + ASSERT_OK(shredding_reader->SetReadSchema(c_read_schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(auto batch_with_bitmap, shredding_reader->NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch_with_bitmap)); + auto& [read_batch, bitmap] = batch_with_bitmap; + auto imported = arrow::ImportArray(read_batch.first.get(), read_batch.second.get()); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + auto result_struct = std::static_pointer_cast(imported.ValueOrDie()); + *v_column = std::static_pointer_cast(result_struct->field(1)); + shredding_reader->Close(); + // The assembled arrays borrow the reader's memory pool; keep the reader alive until the + // fixture is torn down (fixture members outlive test-body locals). + live_readers_.push_back(std::move(shredding_reader)); + } + + protected: + std::unique_ptr dir_; + std::shared_ptr fs_; + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; + std::string file_path_; + std::shared_ptr paimon_schema_; + std::vector> live_readers_; +}; + +namespace { + +constexpr const char* kAgeCityShreddingSchema = R"({ + "type": "ROW", + "fields": [ { + "id": 0, + "name": "v", + "type": { + "type": "ROW", + "fields": [ + {"id": 1, "name": "age", "type": "INT"}, + {"id": 2, "name": "city", "type": "STRING"} + ] + } + } ] +})"; + +} // namespace + +TEST_F(VariantParquetTest, PhysicalLayoutMatchesJava) { + auto array = BuildArray({R"({"a": 1, "b": "hello"})", nullptr, "[1,2,3]"}); + WriteFile(array); + + // The on-disk layout must match the Java ParquetSchemaConverter: an (unannotated) group + // with two REQUIRED BINARY fields `value` (id 0) and `metadata` (id 1). + auto file = arrow::io::ReadableFile::Open(file_path_, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> reader; + auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); + ASSERT_TRUE(status.ok()) << status.ToString(); + const ::parquet::SchemaDescriptor* schema = reader->parquet_reader()->metadata()->schema(); + ASSERT_EQ(schema->num_columns(), 3); + const auto* root = schema->group_node(); + ASSERT_EQ(root->field_count(), 2); + const auto& variant_group_node = root->field(1); + ASSERT_TRUE(variant_group_node->is_group()); + ASSERT_EQ(variant_group_node->name(), "v"); + ASSERT_EQ(variant_group_node->field_id(), 2); + ASSERT_EQ(variant_group_node->logical_type()->type(), ::parquet::LogicalType::Type::NONE); + const auto* variant_group = + static_cast(variant_group_node.get()); + ASSERT_EQ(variant_group->field_count(), 2); + const auto& value_node = variant_group->field(0); + ASSERT_EQ(value_node->name(), "value"); + ASSERT_TRUE(value_node->is_primitive()); + ASSERT_TRUE(value_node->is_required()); + ASSERT_EQ(value_node->field_id(), 0); + ASSERT_EQ( + static_cast(value_node.get())->physical_type(), + ::parquet::Type::BYTE_ARRAY); + const auto& metadata_node = variant_group->field(1); + ASSERT_EQ(metadata_node->name(), "metadata"); + ASSERT_TRUE(metadata_node->is_primitive()); + ASSERT_TRUE(metadata_node->is_required()); + ASSERT_EQ(metadata_node->field_id(), 1); +} + +TEST_F(VariantParquetTest, WriteAndReadRoundTrip) { + std::vector jsons = { + R"({"a": 1, "b": "hello"})", + nullptr, + "[1,2,3]", + "{\"nested\": {\"x\": [true, null, 1.5]}, \"s\": \"中文\"}", + "12345678901234", + "100.99", + }; + auto array = BuildArray(jsons); + WriteFile(array); + + { + // Sanity-check the raw file through the plain parquet-arrow reader: the struct child + // arrays must align with the logical rows. + auto file = arrow::io::ReadableFile::Open(file_path_, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> raw_reader; + ASSERT_TRUE( + ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &raw_reader).ok()); + std::shared_ptr table; + ASSERT_TRUE(raw_reader->ReadTable(&table).ok()); + auto raw_variant = std::static_pointer_cast(table->column(1)->chunk(0)); + auto raw_value = std::static_pointer_cast(raw_variant->field(0)); + for (size_t i = 0; i < jsons.size(); ++i) { + SCOPED_TRACE("raw row " + std::to_string(i)); + if (jsons[i] != nullptr) { + ASSERT_FALSE(raw_variant->IsNull(i)); + ASSERT_GT(raw_value->GetView(i).size(), 0); + } else { + ASSERT_TRUE(raw_variant->IsNull(i)); + } + } + } + + ASSERT_OK_AND_ASSIGN(auto input_stream, fs_->Open(file_path_)); + auto length = fs_->GetFileStatus(file_path_).value()->GetLen(); + auto in_stream = + std::make_unique(std::move(input_stream), arrow_pool_, length); + std::map options = {}; + ASSERT_OK_AND_ASSIGN(auto batch_reader, + ParquetFileBatchReader::Create(std::move(in_stream), options, + /*batch_size=*/1024, + /*file_metadata=*/nullptr, arrow_pool_)); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*paimon_schema_, c_schema.get()).ok()); + ASSERT_OK(batch_reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(auto result_chunked, + paimon::test::ReadResultCollector::CollectResult(batch_reader.get())); + batch_reader->Close(); + ASSERT_EQ(result_chunked->length(), static_cast(jsons.size())); + ASSERT_EQ(result_chunked->num_chunks(), 1); + auto result_struct = std::static_pointer_cast(result_chunked->chunk(0)); + + auto variant_column = std::static_pointer_cast(result_struct->field(1)); + ASSERT_EQ(variant_column->length(), static_cast(jsons.size())); + ASSERT_EQ(variant_column->field(0)->length(), variant_column->length()); + auto value_column = std::static_pointer_cast(variant_column->field(0)); + auto metadata_column = std::static_pointer_cast(variant_column->field(1)); + for (size_t i = 0; i < jsons.size(); ++i) { + SCOPED_TRACE("row " + std::to_string(i)); + if (jsons[i] == nullptr) { + ASSERT_TRUE(variant_column->IsNull(i)); + continue; + } + ASSERT_FALSE(variant_column->IsNull(i)); + auto value_view = value_column->GetView(i); + auto metadata_view = metadata_column->GetView(i); + SCOPED_TRACE("value size " + std::to_string(value_view.size()) + ", metadata size " + + std::to_string(metadata_view.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, + GenericVariant::Create(value_view, metadata_view, pool_)); + ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected, + GenericVariant::FromJson(jsons[i], pool_)); + ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); + ASSERT_EQ(actual_json, expected_json); + } +} + +TEST_F(VariantParquetTest, ShreddedWriteAndReadRoundTrip) { + std::vector jsons = { + R"({"age": 35, "city": "Hangzhou"})", + nullptr, + R"({"age": "not a number", "extra": [1, 2]})", + "[\"top level array\"]", + }; + WriteShreddedFile(jsons, kAgeCityShreddingSchema); + + { + std::unique_ptr file_reader; + std::shared_ptr file_schema; + OpenFile(&file_reader, &file_schema); + auto file_variant_field = file_schema->GetFieldByName("v"); + ASSERT_NE(file_variant_field, nullptr); + ASSERT_TRUE(VariantShreddingUtils::IsShreddedFileType(file_variant_field->type())) + << file_variant_field->type()->ToString(); + file_reader->Close(); + } + + // Reading the column as a plain VARIANT reassembles every physical shape back to the + // original logical value. + std::shared_ptr variant_column; + ReadVariantColumn(paimon_schema_, &variant_column); + ASSERT_EQ(variant_column->length(), static_cast(jsons.size())); + auto value_column = std::static_pointer_cast(variant_column->field(0)); + auto metadata_column = std::static_pointer_cast(variant_column->field(1)); + for (size_t i = 0; i < jsons.size(); ++i) { + SCOPED_TRACE("row " + std::to_string(i)); + if (jsons[i] == nullptr) { + ASSERT_TRUE(variant_column->IsNull(i)); + continue; + } + ASSERT_FALSE(variant_column->IsNull(i)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr variant, + GenericVariant::Create(value_column->GetView(i), metadata_column->GetView(i), pool_)); + ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected, + GenericVariant::FromJson(jsons[i], pool_)); + ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); + ASSERT_EQ(actual_json, expected_json); + } +} + +TEST_F(VariantParquetTest, VariantAccessReadMixedTypedAndBinary) { + std::vector jsons = {R"({"age": 35, "city": "Chicago"})", + R"({"age": 25, "other": "Hello"})", nullptr}; + WriteShreddedFile(jsons, kAgeCityShreddingSchema); + + auto access_field = BuildAccessField( + {{arrow::int64(), "$.age"}, {arrow::utf8(), "$.other"}, {arrow::utf8(), "$.missing"}}); + auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); + + // The plan prunes `typed_value` to the requested keys and keeps `value` because `$.other` + // and `$.missing` are not shredded. + { + std::unique_ptr file_reader; + std::shared_ptr file_schema; + OpenFile(&file_reader, &file_schema); + ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( + read_schema, file_schema, pool_)); + ASSERT_EQ(plans.size(), 1); + const auto& physical_type = + static_cast(*plans.at("v")->PhysicalField()->type()); + ASSERT_NE(physical_type.GetFieldByName(VariantDefs::kMetadataFieldName), nullptr); + ASSERT_NE(physical_type.GetFieldByName(VariantDefs::kValueFieldName), nullptr); + auto typed_value = physical_type.GetFieldByName(VariantDefs::kTypedValueFieldName); + ASSERT_NE(typed_value, nullptr); + const auto& typed_struct = static_cast(*typed_value->type()); + ASSERT_EQ(typed_struct.num_fields(), 1); + ASSERT_NE(typed_struct.GetFieldByName("age"), nullptr); + file_reader->Close(); + } + + std::shared_ptr v_column; + ReadVariantColumn(read_schema, &v_column); + ASSERT_EQ(v_column->length(), 3); + const auto& age = static_cast(*v_column->field(0)); + const auto& other = static_cast(*v_column->field(1)); + const auto& missing = static_cast(*v_column->field(2)); + ASSERT_EQ(age.Value(0), 35); + ASSERT_EQ(age.Value(1), 25); + ASSERT_TRUE(v_column->IsNull(2)); + ASSERT_TRUE(other.IsNull(0)); + ASSERT_EQ(other.GetString(1), "Hello"); + ASSERT_TRUE(missing.IsNull(0)); + ASSERT_TRUE(missing.IsNull(1)); +} + +TEST_F(VariantParquetTest, VariantAccessReadTypedOnlyPrunesValue) { + std::vector jsons = {R"({"age": 35, "city": "Chicago"})", + R"({"age": 25, "other": "Hello"})"}; + WriteShreddedFile(jsons, kAgeCityShreddingSchema); + + auto access_field = BuildAccessField({{arrow::int64(), "$.age"}, {arrow::utf8(), "$.city"}}); + auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); + + // All requested keys are shredded: neither `value` nor the unrequested typed keys are read. + { + std::unique_ptr file_reader; + std::shared_ptr file_schema; + OpenFile(&file_reader, &file_schema); + ASSERT_OK_AND_ASSIGN(auto plans, VariantShreddingReadPlanFactory::CreateReadPlans( + read_schema, file_schema, pool_)); + const auto& physical_type = + static_cast(*plans.at("v")->PhysicalField()->type()); + ASSERT_EQ(physical_type.GetFieldByName(VariantDefs::kValueFieldName), nullptr); + auto typed_value = physical_type.GetFieldByName(VariantDefs::kTypedValueFieldName); + ASSERT_NE(typed_value, nullptr); + ASSERT_EQ(typed_value->type()->num_fields(), 2); + file_reader->Close(); + } + + std::shared_ptr v_column; + ReadVariantColumn(read_schema, &v_column); + const auto& age = static_cast(*v_column->field(0)); + const auto& city = static_cast(*v_column->field(1)); + ASSERT_EQ(age.Value(0), 35); + ASSERT_EQ(age.Value(1), 25); + ASSERT_EQ(city.GetString(0), "Chicago"); + // Row 1 has no "city" key: the shredded field is missing, which reads as null. + ASSERT_TRUE(city.IsNull(1)); +} + +TEST_F(VariantParquetTest, VariantAccessReadUnshreddedFile) { + std::vector jsons = {R"({"age": 35, "city": "Chicago"})", + R"({"age": 25, "other": "Hello"})", nullptr}; + WriteFile(BuildArray(jsons)); + + auto access_field = BuildAccessField({{arrow::int64(), "$.age"}, {arrow::utf8(), "$.other"}}); + auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); + + std::shared_ptr v_column; + ReadVariantColumn(read_schema, &v_column); + ASSERT_EQ(v_column->length(), 3); + const auto& age = static_cast(*v_column->field(0)); + const auto& other = static_cast(*v_column->field(1)); + ASSERT_EQ(age.Value(0), 35); + ASSERT_EQ(age.Value(1), 25); + ASSERT_TRUE(v_column->IsNull(2)); + ASSERT_TRUE(other.IsNull(0)); + ASSERT_EQ(other.GetString(1), "Hello"); +} + +TEST_F(VariantParquetTest, VariantAccessReadSemicolonKey) { + // Object keys may contain the description delimiter; the description parser anchors on the + // trailing failOnError/timeZoneId tokens instead of splitting on every delimiter. + std::vector jsons = {R"({"a;b": 7})"}; + WriteFile(BuildArray(jsons)); + auto access_field = BuildAccessField({{arrow::int64(), "$['a;b']"}}); + auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); + std::shared_ptr v_column; + ReadVariantColumn(read_schema, &v_column); + ASSERT_EQ(static_cast(*v_column->field(0)).Value(0), 7); +} + +TEST_F(VariantParquetTest, VariantAccessReadVariantTarget) { + std::vector jsons = {R"({"user": {"name": "Paimon", "age": 1}})", + R"({"user": "flat"})"}; + WriteFile(BuildArray(jsons)); + + // A variant-marked target re-encodes the extracted sub-variant instead of casting it to a + // plain struct; the marker on the target field must survive AddField. + VariantAccessBuilder builder; + ASSERT_OK_AND_ASSIGN(auto variant_target, Variant::ArrowField("t")); + ASSERT_OK(builder.AddField(variant_target.get(), "$.user")); + ASSERT_OK_AND_ASSIGN(auto c_access_field, builder.Build("v")); + auto imported = arrow::ImportField(c_access_field.get()); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + auto read_schema = arrow::schema({paimon_schema_->field(0), imported.ValueOrDie()}); + + std::shared_ptr v_column; + ReadVariantColumn(read_schema, &v_column); + const auto& user = static_cast(*v_column->field(0)); + const auto& value_column = static_cast(*user.field(0)); + const auto& metadata_column = static_cast(*user.field(1)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr row0, + GenericVariant::Create(value_column.GetView(0), metadata_column.GetView(0), pool_)); + ASSERT_OK_AND_ASSIGN(std::string row0_json, row0->ToJson()); + ASSERT_EQ(row0_json, R"({"age":1,"name":"Paimon"})"); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr row1, + GenericVariant::Create(value_column.GetView(1), metadata_column.GetView(1), pool_)); + ASSERT_OK_AND_ASSIGN(std::string row1_json, row1->ToJson()); + ASSERT_EQ(row1_json, R"("flat")"); +} + +TEST_F(VariantParquetTest, VariantAccessReadNestedPath) { + const char* nested_shredding_schema = R"({ + "type": "ROW", + "fields": [ { + "id": 0, + "name": "v", + "type": { + "type": "ROW", + "fields": [ { + "id": 1, + "name": "address", + "type": { + "type": "ROW", + "fields": [ {"id": 2, "name": "city", "type": "STRING"} ] + } + } ] + } + } ] + })"; + std::vector jsons = {R"({"address": {"city": "Hangzhou"}})", + R"({"address": "oops"})", R"({"address": {"zip": 12345}})"}; + WriteShreddedFile(jsons, nested_shredding_schema); + + auto access_field = BuildAccessField({{arrow::utf8(), "$.address.city"}}); + auto read_schema = arrow::schema({paimon_schema_->field(0), access_field}); + + std::shared_ptr v_column; + ReadVariantColumn(read_schema, &v_column); + const auto& city = static_cast(*v_column->field(0)); + ASSERT_EQ(city.GetString(0), "Hangzhou"); + // Row 1's address is not an object; row 2's address has no "city" key. + ASSERT_TRUE(city.IsNull(1)); + ASSERT_TRUE(city.IsNull(2)); +} + +} // namespace paimon::parquet::test diff --git a/src/paimon/testing/utils/test_helper.h b/src/paimon/testing/utils/test_helper.h index 832055a51..2e51152fe 100644 --- a/src/paimon/testing/utils/test_helper.h +++ b/src/paimon/testing/utils/test_helper.h @@ -22,6 +22,7 @@ #include #include +#include "arrow/array/concatenate.h" #include "arrow/c/bridge.h" #include "arrow/ipc/api.h" #include "paimon/api.h" @@ -239,6 +240,40 @@ class TestHelper { return result_blobs; } + /// Reads all rows of the given splits and returns the raw result (including the leading + /// `_VALUE_KIND` column). Useful when the expected data cannot be expressed as JSON, e.g. + /// binary-encoded VARIANT columns. + Result> ReadResult( + const std::vector>& splits) { + return ReadResult(splits, /*read_schema=*/nullptr); + } + + /// Reads all rows of the given splits with an optional projected read schema. + Result> ReadResult( + const std::vector>& splits, + std::unique_ptr<::ArrowSchema> read_schema) { + ReadContextBuilder read_context_builder(table_path_); + read_context_builder.SetOptions(options_); + if (read_schema != nullptr) { + read_context_builder.SetReadSchema(std::move(read_schema)); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, + read_context_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); + PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr collected, + ReadResultCollector::CollectResult(batch_reader.get())); + if (collected->num_chunks() == 0) { + return collected; + } + // The collected batches borrow reader-owned buffers; copy them into the process pool + // while the reader is still alive so the returned result may outlive it. + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr copied, + arrow::Concatenate(collected->chunks(), arrow::default_memory_pool())); + return std::make_shared(copied); + } + Result ReadAndCheckResult(const std::shared_ptr& data_type, const std::vector>& splits, const std::string& expected_result) { diff --git a/src/paimon/testing/utils/variant_test_data.h b/src/paimon/testing/utils/variant_test_data.h new file mode 100644 index 000000000..35ed69467 --- /dev/null +++ b/src/paimon/testing/utils/variant_test_data.h @@ -0,0 +1,64 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "arrow/api.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace paimon::test { + +/// Builds a `[id, v]` struct array where `v` holds the variant encodings of `jsons` (nullptr +/// means a null variant) and `id` is a running int32 starting at `id_offset`. `variant_field` +/// must be a variant-marked field (see `VariantTypeUtils::ToArrowField`). +inline Result> BuildVariantBatch( + const std::shared_ptr& id_field, + const std::shared_ptr& variant_field, const std::vector& jsons, + const std::shared_ptr& pool, int32_t id_offset = 0) { + arrow::Int32Builder id_builder; + auto value_builder = std::make_shared(); + auto metadata_builder = std::make_shared(); + arrow::StructBuilder variant_builder(variant_field->type(), arrow::default_memory_pool(), + {value_builder, metadata_builder}); + for (size_t i = 0; i < jsons.size(); ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(id_builder.Append(id_offset + static_cast(i))); + if (jsons[i] == nullptr) { + // StructBuilder::AppendNull appends empty values to the child builders itself. + PAIMON_RETURN_NOT_OK_FROM_ARROW(variant_builder.AppendNull()); + } else { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr variant, + GenericVariant::FromJson(jsons[i], pool)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(variant_builder.Append()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Append(variant->RawValue())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata_builder->Append(variant->Metadata())); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr id_array, id_builder.Finish()); + std::shared_ptr variant_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(variant_builder.Finish(&variant_array)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr result, + arrow::StructArray::Make({id_array, variant_array}, {id_field, variant_field})); + return result; +} + +} // namespace paimon::test diff --git a/test/inte/CMakeLists.txt b/test/inte/CMakeLists.txt index 535394809..cb1d824c8 100644 --- a/test/inte/CMakeLists.txt +++ b/test/inte/CMakeLists.txt @@ -20,6 +20,13 @@ if(PAIMON_BUILD_TESTS) test_utils_static ${GTEST_LINK_TOOLCHAIN}) + add_paimon_test(variant_table_inte_test + STATIC_LINK_LIBS + paimon_shared + ${TEST_STATIC_LINK_LIBS} + test_utils_static + ${GTEST_LINK_TOOLCHAIN}) + add_paimon_test(data_evolution_table_test STATIC_LINK_LIBS paimon_shared diff --git a/test/inte/variant_table_inte_test.cpp b/test/inte/variant_table_inte_test.cpp new file mode 100644 index 000000000..1fa2d87bc --- /dev/null +++ b/test/inte/variant_table_inte_test.cpp @@ -0,0 +1,233 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed 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 "arrow/api.h" +#include "arrow/c/bridge.h" +#include "gtest/gtest.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/data/variant.h" +#include "paimon/defs.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/record_batch.h" +#include "paimon/table/source/startup_mode.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/test_helper.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/testing/utils/variant_test_data.h" + +namespace paimon::test { + +// End-to-end tests for tables with a VARIANT column: create, write, commit, scan and read. +class VariantTableInteTest : public ::testing::Test { + public: + void SetUp() override { + dir_ = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir_); + test_dir_ = dir_->Str(); + pool_ = GetDefaultPool(); + fields_ = {arrow::field("id", arrow::int32()), VariantTypeUtils::ToArrowField("v")}; + schema_ = arrow::schema(fields_); + } + + void TearDown() override { + dir_.reset(); + } + + std::shared_ptr BuildArray(const std::vector& jsons, + int32_t id_offset = 0) { + auto result = BuildVariantBatch(fields_[0], fields_[1], jsons, pool_, id_offset); + EXPECT_TRUE(result.ok()) << result.status().ToString(); + return std::move(result).value(); + } + + Result> MakeBatch( + const std::shared_ptr& array) { + ::ArrowArray arrow_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &arrow_array)); + RecordBatchBuilder batch_builder(&arrow_array); + return batch_builder.SetPartition({}).SetBucket(0).SetRowKinds({}).Finish(); + } + + // Reads all rows back and checks the variant column renders to `expected_jsons` (nullptr + // means a null variant). The read result carries a leading `_VALUE_KIND` column. + void ReadAndCheck(TestHelper* helper, const std::vector>& splits, + const std::vector& expected_ids, + const std::vector& expected_jsons) { + ASSERT_OK_AND_ASSIGN(auto result, helper->ReadResult(splits)); + ASSERT_EQ(result->num_chunks(), 1); + auto result_struct = std::static_pointer_cast(result->chunk(0)); + ASSERT_EQ(result_struct->length(), static_cast(expected_jsons.size())); + auto struct_type = std::static_pointer_cast(result_struct->type()); + int32_t id_index = struct_type->GetFieldIndex("id"); + int32_t variant_index = struct_type->GetFieldIndex("v"); + ASSERT_GE(id_index, 0); + ASSERT_GE(variant_index, 0); + auto id_column = + std::static_pointer_cast(result_struct->field(id_index)); + auto variant_column = + std::static_pointer_cast(result_struct->field(variant_index)); + auto value_column = std::static_pointer_cast(variant_column->field(0)); + auto metadata_column = + std::static_pointer_cast(variant_column->field(1)); + for (size_t i = 0; i < expected_jsons.size(); ++i) { + SCOPED_TRACE("row " + std::to_string(i)); + ASSERT_EQ(id_column->Value(i), expected_ids[i]); + if (expected_jsons[i] == nullptr) { + ASSERT_TRUE(variant_column->IsNull(i)); + continue; + } + ASSERT_FALSE(variant_column->IsNull(i)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr variant, + GenericVariant::Create(value_column->GetView(i), + metadata_column->GetView(i), pool_)); + ASSERT_OK_AND_ASSIGN(std::string actual_json, variant->ToJson()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected, + GenericVariant::FromJson(expected_jsons[i], pool_)); + ASSERT_OK_AND_ASSIGN(std::string expected_json, expected->ToJson()); + ASSERT_EQ(actual_json, expected_json); + } + } + + protected: + std::string test_dir_; + std::unique_ptr dir_; + std::shared_ptr pool_; + arrow::FieldVector fields_; + std::shared_ptr schema_; +}; + +TEST_F(VariantTableInteTest, TestAppendTable) { + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "-1"}, + }; + ASSERT_OK_AND_ASSIGN( + auto helper, TestHelper::Create(test_dir_, schema_, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); + std::vector jsons = { + R"({"age": 35, "city": "Hangzhou"})", + nullptr, + "[1, \"two\", 3.5, null, true]", + "{\"nested\": {\"x\": [1, 2]}, \"s\": \"中文\"}", + }; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeBatch(BuildArray(jsons))); + ASSERT_OK_AND_ASSIGN(auto commit_msgs, + helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt, + /*is_streaming=*/false)); + ReadAndCheck(helper.get(), splits, {0, 1, 2, 3}, jsons); +} + +TEST_F(VariantTableInteTest, TestPrimaryKeyTable) { + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "1"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(test_dir_, schema_, /*partition_keys=*/{}, + /*primary_keys=*/{"id"}, options, + /*is_streaming_mode=*/true)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch_1, + MakeBatch(BuildArray({"{\"a\": 1}", "{\"b\": 2}", nullptr}, /*id_offset=*/0))); + ASSERT_OK_AND_ASSIGN(auto commit_msgs_1, + helper->WriteAndCommit(std::move(batch_1), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_2, + MakeBatch(BuildArray({"{\"b\": \"updated\"}", "[42]"}, /*id_offset=*/1))); + ASSERT_OK_AND_ASSIGN(auto commit_msgs_2, + helper->WriteAndCommit(std::move(batch_2), /*commit_identifier=*/1, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + // The second batch overwrites ids 1 and 2, so the merged view holds three rows. + ReadAndCheck(helper.get(), splits, {0, 1, 2}, {"{\"a\": 1}", R"({"b": "updated"})", "[42]"}); +} + +TEST_F(VariantTableInteTest, TestVariantAccessRead) { + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "-1"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(test_dir_, schema_, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/false)); + std::vector jsons = {R"({"age": 35, "city": "Chicago"})", + R"({"age": 25, "other": "Hello"})", nullptr}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeBatch(BuildArray(jsons))); + ASSERT_OK_AND_ASSIGN(auto commit_msgs, + helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt, + /*is_streaming=*/false)); + + VariantAccessBuilder access_builder; + auto age_target = std::make_unique(); + ASSERT_TRUE(arrow::ExportField(arrow::Field("t", arrow::int64()), age_target.get()).ok()); + ASSERT_OK(access_builder.AddField(age_target.get(), "$.age", /*fail_on_error=*/false)); + auto other_target = std::make_unique(); + ASSERT_TRUE(arrow::ExportField(arrow::Field("t", arrow::utf8()), other_target.get()).ok()); + ASSERT_OK(access_builder.AddField(other_target.get(), "$.other", /*fail_on_error=*/false)); + ASSERT_OK_AND_ASSIGN(auto c_access_field, access_builder.Build("v")); + auto imported_access = arrow::ImportField(c_access_field.get()); + ASSERT_TRUE(imported_access.ok()) << imported_access.status().ToString(); + auto read_schema = arrow::schema({fields_[0], imported_access.ValueOrDie()}); + auto c_read_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*read_schema, c_read_schema.get()).ok()); + + ASSERT_OK_AND_ASSIGN(auto result, helper->ReadResult(splits, std::move(c_read_schema))); + ASSERT_EQ(result->num_chunks(), 1); + auto result_struct = std::static_pointer_cast(result->chunk(0)); + ASSERT_EQ(result_struct->length(), 3); + auto struct_type = std::static_pointer_cast(result_struct->type()); + auto v_column = std::static_pointer_cast( + result_struct->field(struct_type->GetFieldIndex("v"))); + const auto& age = static_cast(*v_column->field(0)); + const auto& other = static_cast(*v_column->field(1)); + ASSERT_EQ(age.Value(0), 35); + ASSERT_EQ(age.Value(1), 25); + ASSERT_TRUE(v_column->IsNull(2)); + ASSERT_TRUE(other.IsNull(0)); + ASSERT_EQ(other.GetString(1), "Hello"); +} + +TEST_F(VariantTableInteTest, TestOrcFormatRejected) { + std::map options = { + {Options::MANIFEST_FORMAT, "avro"}, + {Options::FILE_FORMAT, "orc"}, + {Options::BUCKET, "-1"}, + }; + ASSERT_OK_AND_ASSIGN( + auto helper, TestHelper::Create(test_dir_, schema_, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeBatch(BuildArray({"{\"a\": 1}"}))); + auto result = helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt); + ASSERT_FALSE(result.ok()); +} + +} // namespace paimon::test