Antalya-26.6: iceberg v3 multi-arg transforms ,added logic to read source_ids and not throw an error

src/Storages/MergeTree/ExportPartitionUtils.cpp

@@ -43,6 +43,7 @@
 #include <Storages/ColumnsDescription.h>
 
 #if USE_AVRO
+#include <Storages/ObjectStorage/DataLakes/Iceberg/ChunkPartitioner.h>
 #include <Storages/ObjectStorage/DataLakes/Iceberg/Constant.h>
 #include <Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.h>
 #include <Storages/ObjectStorage/DataLakes/Iceberg/ManifestFilesPruning.h>
@@ -942,6 +943,8 @@ namespace
         if (spec_size == 0)
             return;
 
+        ChunkPartitioner::checkPartitionSpecSupportedForWrites(spec_fields);
+
         /// Rebuild the destination spec as a ClickHouse partition key, the way the Iceberg read path does in
         /// ManifestFileIterator, so the same compatibility rule applies as for a plain object storage
         /// destination and the transform arguments keep the order the writer will use.

src/Storages/ObjectStorage/DataLakes/Iceberg/ChunkPartitioner.cpp

@@ -27,11 +27,30 @@ namespace ErrorCodes
     extern const int BAD_ARGUMENTS;
 }
 
+void ChunkPartitioner::checkPartitionSpecSupportedForWrites(const Poco::JSON::Array::Ptr & partition_specification)
+{
+    if (!partition_specification)
+        return;
+
+    for (size_t i = 0; i != partition_specification->size(); ++i)
+    {
+        /// Iceberg V3: multi-argument transforms use `source-ids` instead of `source-id`.
+        /// Writing with multi-arg transforms is not supported yet (hash semantics not finalized upstream).
+        if (partition_specification->getObject(static_cast<UInt32>(i))->has(Iceberg::f_source_ids))
+            throw Exception(
+                ErrorCodes::BAD_ARGUMENTS,
+                "Multi-argument partition transforms (source-ids) are not supported for writes. "
+                "Multi-argument transform evaluation is not yet implemented");
+    }
+}
+
 ChunkPartitioner::ChunkPartitioner(
     Poco::JSON::Array::Ptr partition_specification, Poco::JSON::Array::Ptr schema_fields, ContextPtr context, SharedHeader sample_block_)
     : sample_block(sample_block_)
     , max_partitions_count(context->getSettingsRef()[Setting::iceberg_insert_max_partitions])
 {
+    checkPartitionSpecSupportedForWrites(partition_specification);
+
     std::unordered_map<Int32, String> id_to_column;
     {
         for (size_t i = 0; i < schema_fields->size(); ++i)

src/Storages/ObjectStorage/DataLakes/Iceberg/ChunkPartitioner.h

@@ -24,6 +24,10 @@ class ChunkPartitioner
         ContextPtr context,
         SharedHeader sample_block_);
 
+    /// Throws if the partition spec cannot be evaluated on write (e.g. multi-argument `source-ids` transforms).
+    /// Callers that write objects before constructing a `ChunkPartitioner` must call this first.
+    static void checkPartitionSpecSupportedForWrites(const Poco::JSON::Array::Ptr & partition_specification);
+
     using PartitionKey = Row;
     struct PartitionKeyHasher
     {

src/Storages/ObjectStorage/DataLakes/Iceberg/Compaction.cpp

@@ -206,6 +206,29 @@ static Plan getPlan(
     if (initial_metadata_object->getValue<Int32>(Iceberg::f_format_version) < 2)
         throw Exception(ErrorCodes::BAD_ARGUMENTS, "Compaction is supported only for format_version 2.");
 
+    /// `writeMetadataFiles` partitions by the default spec only after the rewritten data files are written,
+    /// so an unsupported spec must be rejected here, before any object is created.
+    {
+        const auto default_spec_id = initial_metadata_object->getValue<Int64>(Iceberg::f_default_spec_id);
+        const auto partition_specs = initial_metadata_object->getArray(Iceberg::f_partition_specs);
+        Poco::JSON::Object::Ptr default_spec;
+        for (UInt32 i = 0; i < partition_specs->size(); ++i)
+        {
+            auto candidate = partition_specs->getObject(i);
+            if (candidate->getValue<Int64>(Iceberg::f_spec_id) == default_spec_id)
+            {
+                default_spec = candidate;
+                break;
+            }
+        }
+        if (!default_spec)
+            throw Exception(
+                ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION,
+                "Iceberg metadata does not contain a partition spec entry matching default-spec-id {}",
+                default_spec_id);
+        ChunkPartitioner::checkPartitionSpecSupportedForWrites(default_spec->getArray(Iceberg::f_fields));
+    }
+
     auto current_schema_id = initial_metadata_object->getValue<Int64>(Iceberg::f_current_schema_id);
     auto schemas = initial_metadata_object->getArray(Iceberg::f_schemas);
     Poco::JSON::Array::Ptr current_schema;
@@ -511,6 +534,7 @@ static bool writeConsolidatedManifestFile(
         ResolvedPartitionSpec resolved;
         resolved.spec = spec;
         auto spec_fields = spec->getArray(f_fields);
+        ChunkPartitioner::checkPartitionSpecSupportedForWrites(spec_fields);
 
         /// Partition field names and the schema source-ids they transform.
         std::vector<Int32> source_ids;

src/Storages/ObjectStorage/DataLakes/Iceberg/Constant.h

@@ -128,6 +128,7 @@ DEFINE_ICEBERG_FIELD_ALIAS(order_id, order-id);
 DEFINE_ICEBERG_FIELD_ALIAS(default_sort_order_id, default-sort-order-id);
 DEFINE_ICEBERG_FIELD_ALIAS(sort_orders, sort-orders);
 DEFINE_ICEBERG_FIELD_ALIAS(source_id, source-id);
+DEFINE_ICEBERG_FIELD_ALIAS(source_ids, source-ids);
 DEFINE_ICEBERG_FIELD_ALIAS(partition_transform, transform);
 DEFINE_ICEBERG_FIELD_ALIAS(partition_name, name);
 DEFINE_ICEBERG_FIELD_ALIAS(default_spec_id, default-spec-id);

src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.cpp

@@ -1846,8 +1846,14 @@ KeyDescription IcebergMetadata::getSortingKey(ContextPtr local_context, TableSta
 
     auto [schema, current_schema_id] = parseTableSchemaV2Method(metadata_object);
     auto result = getSortingKeyDescriptionFromMetadata(metadata_object, *persistent_components.schema_processor->getClickhouseTableSchemaById(current_schema_id), local_context);
-    auto sort_order_id = metadata_object->getValue<Int64>(f_default_sort_order_id);
-    result.sort_order_id = sort_order_id;
+    /// Only attach the sort-order id when the parsed key is non-empty.  When multi-arg or
+    /// unsupported transforms truncate the key to empty, pairing a stale sort_order_id with
+    /// an empty key would let `isDataSortedBySortingKey` incorrectly claim data is sorted.
+    if (result.definition_ast != nullptr)
+    {
+        auto sort_order_id = metadata_object->getValue<Int64>(f_default_sort_order_id);
+        result.sort_order_id = sort_order_id;
+    }
     return result;
 }
 

src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFile.cpp

@@ -9,6 +9,7 @@
 
 #include <Common/logger_useful.h>
 #include <fmt/format.h>
+#include <fmt/ranges.h>
 
 
 namespace DB::ErrorCodes
@@ -99,8 +100,10 @@ void requireDirectReferencedDataFileForPuffinDeletionVector(
 
 static std::strong_ordering operator<=>(const PartitionSpecsEntry & lhs, const PartitionSpecsEntry & rhs)
 {
-    return std::tie(lhs.source_id, lhs.transform_name, lhs.partition_name)
-        <=> std::tie(rhs.source_id, rhs.transform_name, rhs.partition_name);
+    if (auto cmp = lhs.source_ids <=> rhs.source_ids; cmp != std::strong_ordering::equal)
+        return cmp;
+    return std::tie(lhs.transform_name, lhs.partition_name)
+        <=> std::tie(rhs.transform_name, rhs.partition_name);
 }
 
 template <typename A>
@@ -138,7 +141,8 @@ static String dumpPartitionSpecification(const PartitionSpecification & partitio
         {
             const auto & entry = partition_specification[i];
             answer += fmt::format(
-                "(source id: {}, transform name: {}, partition name: {})", entry.source_id, entry.transform_name, entry.partition_name);
+                "(source ids: [{}], transform name: {}, partition name: {})",
+                fmt::join(entry.source_ids, ", "), entry.transform_name, entry.partition_name);
             if (i != partition_specification.size() - 1)
                 answer += ", ";
         }

src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFile.h

@@ -61,9 +61,14 @@ String FileContentTypeToString(FileContentType type);
 
 struct PartitionSpecsEntry
 {
-    Int32 source_id;
+    /// For single-argument transforms (V1/V2 and single-arg V3) this holds one element.
+    /// For multi-argument V3 transforms (e.g. bucket over multiple columns) it holds multiple.
+    std::vector<Int32> source_ids;
     String transform_name;
     String partition_name;
+
+    /// Convenience: true when the transform references more than one source column (V3 multi-arg).
+    bool isMultiArg() const { return source_ids.size() > 1; }
 };
 using PartitionSpecification = std::vector<PartitionSpecsEntry>;
 

src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFileIterator.cpp

@@ -254,28 +254,78 @@ std::shared_ptr<ManifestFileIterator> ManifestFileIterator::create(
     schema_processor.addIcebergTableSchema(schema_object, context_);
 
     PartitionSpecification partition_spec_vec;
+    /// Which spec-field positions made it into partition_key_ast (needed by the pruner to project
+    /// partition values, since the Avro tuple has one entry per spec field, not per kept field).
+    std::vector<size_t> partition_key_spec_positions;
     for (size_t i = 0; i != partition_specification->size(); ++i)
     {
         auto partition_specification_field = partition_specification->getObject(static_cast<UInt32>(i));
 
-        auto source_id = partition_specification_field->getValue<Int32>(f_source_id);
+        /// Iceberg V3 spec: partition fields use either singular `source-id` (single-arg transforms)
+        /// or `source-ids` (multi-arg transforms, e.g. bucket over multiple columns). They are mutually exclusive.
+        bool has_source_id = partition_specification_field->has(f_source_id);
+        bool has_source_ids = partition_specification_field->has(f_source_ids);
+
+        std::vector<Int32> source_ids;
+        if (has_source_id && has_source_ids)
+        {
+            throw Exception(
+                ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION,
+                "Partition field in manifest '{}' has both 'source-id' and 'source-ids' — they are mutually exclusive per the Iceberg spec",
+                path_to_manifest_file_);
+        }
+        else if (has_source_ids)
+        {
+            auto source_ids_array = partition_specification_field->getArray(f_source_ids);
+            for (UInt32 idx = 0; idx < source_ids_array->size(); ++idx)
+                source_ids.push_back(source_ids_array->getElement<Int32>(idx));
+        }
+        else if (has_source_id)
+        {
+            source_ids.push_back(partition_specification_field->getValue<Int32>(f_source_id));
+        }
+        else
+        {
+            throw Exception(
+                ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION,
+                "Partition field in manifest '{}' has neither 'source-id' nor 'source-ids'",
+                path_to_manifest_file_);
+        }
+
+        if (source_ids.empty())
+        {
+            throw Exception(
+                ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION,
+                "Partition field in manifest '{}' has an empty 'source-ids' array",
+                path_to_manifest_file_);
+        }
+
+        auto transform_name = partition_specification_field->getValue<String>(f_partition_transform);
+        auto partition_name = partition_specification_field->getValue<String>(f_partition_name);
+        partition_spec_vec.emplace_back(PartitionSpecsEntry{std::move(source_ids), transform_name, partition_name});
+
+        /// Multi-argument transforms (V3): we cannot evaluate the transform, so skip pruning for this field.
+        /// Per the Iceberg V3 spec: "all v3 readers are required to read tables with unknown transforms,
+        /// ignoring the unsupported partition fields when filtering."
+        if (partition_spec_vec.back().isMultiArg())
+            continue;
+
+        auto source_id = partition_spec_vec.back().source_ids[0];
         /// NOTE: tricky part to support RENAME column in partition key. Instead of some name
         /// we use column internal number as it's name.
         auto numeric_column_name = DB::backQuote(DB::toString(source_id));
         std::optional<DB::NameAndTypePair> manifest_file_column_characteristics
             = schema_processor.tryGetFieldCharacteristics(manifest_schema_id, source_id);
         if (!manifest_file_column_characteristics.has_value())
             continue;
-        auto transform_name = partition_specification_field->getValue<String>(f_partition_transform);
-        auto partition_name = partition_specification_field->getValue<String>(f_partition_name);
-        partition_spec_vec.emplace_back(source_id, transform_name, partition_name);
         auto partition_ast = getASTFromTransform(transform_name, numeric_column_name, context_->getSettingsRef()[Setting::iceberg_partition_timezone]);
         /// Unsupported partition key expression
         if (partition_ast == nullptr)
             continue;
 
         partition_key_ast->as<ASTFunction>()->arguments->children.emplace_back(std::move(partition_ast));
         partition_columns_description.emplace_back(numeric_column_name, removeNullable(manifest_file_column_characteristics->type));
+        partition_key_spec_positions.push_back(i);
     }
 
     std::optional<DB::KeyDescription> partition_key_description;
@@ -298,6 +348,7 @@ std::shared_ptr<ManifestFileIterator> ManifestFileIterator::create(
         manifest_schema_id,
         std::make_shared<const PartitionSpecification>(std::move(partition_spec_vec)),
         std::move(partition_key_description),
+        std::move(partition_key_spec_positions),
         total_rows,
         std::move(filter_dag_),
         table_snapshot_schema_id_));
@@ -316,6 +367,7 @@ ManifestFileIterator::ManifestFileIterator(
     Int32 manifest_schema_id_,
     std::shared_ptr<const PartitionSpecification> common_partition_specification_,
     std::optional<DB::KeyDescription> partition_key_description_,
+    std::vector<size_t> partition_key_spec_positions_,
     size_t total_rows_,
     std::shared_ptr<const ActionsDAG> filter_dag_,
     Int32 table_snapshot_schema_id_)
@@ -329,6 +381,7 @@ ManifestFileIterator::ManifestFileIterator(
     , manifest_schema_id(manifest_schema_id_)
     , common_partition_specification(std::move(common_partition_specification_))
     , partition_key_description(std::move(partition_key_description_))
+    , partition_key_spec_positions(std::move(partition_key_spec_positions_))
     , table_snapshot_schema_id(table_snapshot_schema_id_)
     , total_rows(total_rows_)
     , data_files_without_deleted(std::make_shared<std::vector<ProcessedManifestFileEntryPtr>>())

src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFileIterator.h

@@ -91,6 +91,9 @@ class ManifestFileIterator : public boost::noncopyable
 
     bool hasPartitionKey() const;
     const DB::KeyDescription & getPartitionKeyDescription() const;
+    /// Returns the spec-field positions (indices into the Avro partition tuple) that correspond
+    /// to the kept partition key columns.  The pruner uses this to project partition values.
+    const std::vector<size_t> & getPartitionKeySpecPositions() const { return partition_key_spec_positions; }
     /// Sums required per-file `record_count` for live entries of the given content type.
     /// Returns nullopt on negative `record_count` or Int64 overflow (fail closed).
     std::optional<Int64> getRowsCountInAllFilesExcludingDeleted(FileContentType content) const;
@@ -123,6 +126,7 @@ class ManifestFileIterator : public boost::noncopyable
         Int32 manifest_schema_id,
         std::shared_ptr<const PartitionSpecification> common_partition_specification,
         std::optional<DB::KeyDescription> partition_key_description,
+        std::vector<size_t> partition_key_spec_positions,
         size_t total_rows,
         std::shared_ptr<const ActionsDAG> filter_dag,
         Int32 table_snapshot_schema_id);
@@ -142,6 +146,8 @@ class ManifestFileIterator : public boost::noncopyable
     const Int32 manifest_schema_id;
     const std::shared_ptr<const PartitionSpecification> common_partition_specification;
     const std::optional<DB::KeyDescription> partition_key_description;
+    /// Indices (into the Avro partition tuple) of the spec fields that were kept for pruning.
+    const std::vector<size_t> partition_key_spec_positions;
     const Int32 table_snapshot_schema_id;
 
     /// Iteration state

src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFilesPruning.cpp

@@ -155,6 +155,7 @@ ManifestFilesPruner::ManifestFilesPruner(
     if (manifest_file.hasPartitionKey())
     {
         partition_key = &manifest_file.getPartitionKeyDescription();
+        partition_key_spec_positions = manifest_file.getPartitionKeySpecPositions();
         ActionsDAGWithInversionPushDown inverted_dag(transformed_dag->getOutputs().front(), context, /* boolean_context */ true);
         partition_key_condition.emplace(
             inverted_dag, context, partition_key->column_names, partition_key->expression, true /* single_point */);
@@ -247,7 +248,18 @@ PruningReturnStatus ManifestFilesPruner::canBePruned(
     if (partition_key_condition.has_value())
     {
         const auto & partition_value = entry->parsed_entry->partition_key_value;
-        std::vector<FieldRef> index_value(partition_value.begin(), partition_value.end());
+
+        /// Project: the Avro tuple has one entry per spec field, but the key only contains the
+        /// fields at the positions recorded in `partition_key_spec_positions`.  Pick exactly those.
+        std::vector<FieldRef> index_value;
+        index_value.reserve(partition_key_spec_positions.size());
+        for (size_t pos : partition_key_spec_positions)
+        {
+            if (pos >= partition_value.size())
+                return PruningReturnStatus::NOT_PRUNED; /// fail closed — cannot prune without the value
+            index_value.push_back(FieldRef(partition_value[pos]));
+        }
+
         for (size_t i = 0; i < index_value.size(); ++i)
         {
             auto & field = index_value[i];
@@ -263,7 +275,7 @@ PruningReturnStatus ManifestFilesPruner::canBePruned(
         }
 
         bool can_be_true = partition_key_condition->mayBeTrueInRange(
-            partition_value.size(), index_value.data(), index_value.data(), partition_key->data_types);
+            index_value.size(), index_value.data(), index_value.data(), partition_key->data_types);
 
         if (!can_be_true)
         {

src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFilesPruning.h

@@ -43,6 +43,8 @@ class ManifestFilesPruner
     Int32 current_schema_id;
     Int32 initial_schema_id;
     const DB::KeyDescription * partition_key;
+    /// Which positions in the Avro partition tuple correspond to the kept partition-key columns.
+    std::vector<size_t> partition_key_spec_positions;
     std::optional<DB::KeyCondition> partition_key_condition;
 
     std::unordered_map<Int32, DB::KeyCondition> min_max_key_conditions;