Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/Storages/MergeTree/ExportPartitionUtils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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.
Expand Down
19 changes: 19 additions & 0 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/ChunkPartitioner.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand Down
24 changes: 24 additions & 0 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/Compaction.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
1 change: 1 addition & 0 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/Constant.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
10 changes: 8 additions & 2 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
10 changes: 7 additions & 3 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFile.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

#include <Common/logger_useful.h>
#include <fmt/format.h>
#include <fmt/ranges.h>


namespace DB::ErrorCodes
Expand Down Expand Up @@ -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>
Expand Down Expand Up @@ -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 += ", ";
}
Expand Down
7 changes: 6 additions & 1 deletion src/Storages/ObjectStorage/DataLakes/Iceberg/ManifestFile.h
Original file line number Diff line number Diff line change
Expand Up @@ -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>;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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_));
Expand All @@ -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_)
Expand All @@ -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>>())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 */);
Expand Down Expand Up @@ -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];
Expand All @@ -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)
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading
Loading