From 4af4c5ad8e67d6569f406ec78a742b713b26fd15 Mon Sep 17 00:00:00 2001 From: Andrew Xie Date: Tue, 27 Oct 2026 00:09:57 -0400 Subject: [PATCH] Backport improve iceberg writes for unity catalog Signed-off-by: Andrew Xie --- src/Databases/DataLake/RestCatalog.cpp | 6 +- .../DataLakes/Iceberg/AvroSchema.h | 7 + .../DataLakes/Iceberg/FileNamesGenerator.cpp | 6 + .../DataLakes/Iceberg/FileNamesGenerator.h | 8 + .../DataLakes/Iceberg/IcebergWrites.cpp | 171 ++++++++++-------- .../DataLakes/Iceberg/MetadataGenerator.cpp | 11 +- .../DataLakes/Iceberg/MetadataGenerator.h | 2 +- .../ObjectStorage/DataLakes/Iceberg/Utils.cpp | 2 +- 8 files changed, 127 insertions(+), 86 deletions(-) diff --git a/src/Databases/DataLake/RestCatalog.cpp b/src/Databases/DataLake/RestCatalog.cpp index 29994d7916899..2879bdbd249e2 100644 --- a/src/Databases/DataLake/RestCatalog.cpp +++ b/src/Databases/DataLake/RestCatalog.cpp @@ -1144,7 +1144,7 @@ void RestCatalog::sendRequest(const String & endpoint, Poco::JSON::Object::Ptr r void RestCatalog::createNamespaceIfNotExists(const String & namespace_name, const String & location) const { - const std::string endpoint = base_url / config.prefix / "namespaces"; + const std::string endpoint = (base_url / config.prefix / NAMESPACES_ENDPOINT).generic_string(); Poco::JSON::Object::Ptr request_body = new Poco::JSON::Object; { @@ -1178,7 +1178,7 @@ void RestCatalog::createTable(const String & namespace_name, const String & tabl createNamespaceIfNotExists(namespace_name, metadata_content->getValue("location")); - const std::string endpoint = base_url / config.prefix / "namespaces" / namespace_name / "tables"; + const std::string endpoint = (base_url / config.prefix / NAMESPACES_ENDPOINT / encodeNamespaceForURI(namespace_name) / "tables").generic_string(); Poco::JSON::Object::Ptr request_body = new Poco::JSON::Object; request_body->set("name", table_name); @@ -1217,7 +1217,7 @@ void RestCatalog::createTable(const String & namespace_name, const String & tabl bool RestCatalog::updateMetadata(const String & namespace_name, const String & table_name, const String & /*new_metadata_path*/, Poco::JSON::Object::Ptr new_snapshot) const { - const std::string endpoint = base_url / config.prefix / "namespaces" / namespace_name / "tables" / table_name; + const std::string endpoint = (base_url / config.prefix / NAMESPACES_ENDPOINT / encodeNamespaceForURI(namespace_name) / "tables" / table_name).generic_string(); Poco::JSON::Object::Ptr request_body = new Poco::JSON::Object; { diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/AvroSchema.h b/src/Storages/ObjectStorage/DataLakes/Iceberg/AvroSchema.h index 0c09df4d9e776..97c832760a11d 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/AvroSchema.h +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/AvroSchema.h @@ -187,6 +187,13 @@ static constexpr const char * manifest_list_v2_schema = R"( ], "doc": "Summary for each partition", "field-id": 507 + }, + { + "name": "key_metadata", + "type": ["null", "bytes"], + "doc": "Encryption key metadata blob", + "default": null, + "field-id": 519 } ] } diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.cpp b/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.cpp index b0b1c3e5ef682..d549153b23935 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.cpp +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.cpp @@ -27,6 +27,7 @@ FileNamesGenerator::FileNamesGenerator(const FileNamesGenerator & other) { initial_version = other.initial_version; table_location = other.table_location; + data_location = other.data_location; use_uuid_in_metadata = other.use_uuid_in_metadata; compression_method = other.compression_method; format_name = other.format_name; @@ -39,6 +40,7 @@ FileNamesGenerator & FileNamesGenerator::operator=(const FileNamesGenerator & ot initial_version = other.initial_version; table_location = other.table_location; + data_location = other.data_location; use_uuid_in_metadata = other.use_uuid_in_metadata; compression_method = other.compression_method; format_name = other.format_name; @@ -49,6 +51,8 @@ FileNamesGenerator & FileNamesGenerator::operator=(const FileNamesGenerator & ot Iceberg::IcebergPathFromMetadata FileNamesGenerator::generateDataFileName() { auto uuid_str = uuid_generator.createRandom().toString(); + if (!data_location.empty()) + return Iceberg::IcebergPathFromMetadata(fmt::format("{}/data-{}.{}", data_location, uuid_str, format_name)); return Iceberg::IcebergPathFromMetadata(fmt::format("{}data/data-{}.{}", table_location, uuid_str, format_name)); } @@ -97,6 +101,8 @@ Iceberg::IcebergPathFromMetadata FileNamesGenerator::generateVersionHint() Iceberg::IcebergPathFromMetadata FileNamesGenerator::generatePositionDeleteFile() { auto uuid_str = uuid_generator.createRandom().toString(); + if (!data_location.empty()) + return Iceberg::IcebergPathFromMetadata(fmt::format("{}/{}-deletes.{}", data_location, uuid_str, format_name)); return Iceberg::IcebergPathFromMetadata(fmt::format("{}data/{}-deletes.{}", table_location, uuid_str, format_name)); } diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.h b/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.h index b26a5c86ffd42..7a4533e67d54c 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.h +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.h @@ -51,9 +51,17 @@ class FileNamesGenerator void setVersion(Int32 initial_version_) { initial_version = initial_version_; } void setCompressionMethod(CompressionMethod compression_method_) { compression_method = compression_method_; } + void setDataLocation(String data_location_) + { + data_location = std::move(data_location_); + if (!data_location.empty() && data_location.back() == '/') + data_location.pop_back(); + } + private: Poco::UUIDGenerator uuid_generator; String table_location; + String data_location; /// Optional override from `write.data.path` table property bool use_uuid_in_metadata = false; CompressionMethod compression_method = CompressionMethod::None; String format_name; diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergWrites.cpp b/src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergWrites.cpp index acd374a71477d..ff0d5a19933e2 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergWrites.cpp +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergWrites.cpp @@ -56,7 +56,6 @@ #include #include #include - #include #include #include @@ -102,6 +101,7 @@ namespace ErrorCodes { extern const int LOGICAL_ERROR; extern const int BAD_ARGUMENTS; + extern const int DATALAKE_DATABASE_ERROR; extern const int NOT_IMPLEMENTED; extern const int ICEBERG_SPECIFICATION_VIOLATION; } @@ -513,7 +513,6 @@ void generateManifestFile( int current_schema_id = metadata->getValue(Iceberg::f_current_schema_id); Poco::JSON::Object::Ptr schema_object_to_write = metadata->getArray(Iceberg::f_schemas)->getObject(current_schema_id); Poco::JSON::Stringifier::stringify(schema_object_to_write, oss, 4); - std::string json_representation = removeEscapedSlashes(oss.str()); auto adapter = std::make_unique(buf); @@ -523,7 +522,7 @@ void generateManifestFile( writer.setMetadata(Iceberg::f_format_version, std::to_string(version)); std::ostringstream oss_partition_spec; // STYLE_CHECK_ALLOW_STD_STRING_STREAM - Poco::JSON::Stringifier::stringify(partition_spec->getArray(Iceberg::f_fields), oss_partition_spec, 4); + Poco::JSON::Stringifier::stringify(partition_spec->getArray(Iceberg::f_fields), oss_partition_spec); writer.setMetadata(Iceberg::f_partition_spec, oss_partition_spec.str()); writer.setMetadata(Iceberg::f_manifest_partition_spec_id, std::to_string(partition_spec_id)); for (size_t file_idx = 0; file_idx < data_file_names.size(); ++file_idx) @@ -637,7 +636,7 @@ void generateManifestFile( { avro::GenericDatum record_datum(schema_element); auto & record = record_datum.value(); - record.field(Iceberg::f_key) = static_cast(field_id); + record.field(Iceberg::f_key) = static_cast(field_id); record.field(Iceberg::f_value) = dump_function(field_id, value); record_values.value().push_back(record_datum); } @@ -909,8 +908,8 @@ void generateManifestList( else { entry.field(Iceberg::f_added_files_count) = 1; - entry.field(Iceberg::f_existing_files_count) - = summary->getValue(Iceberg::f_total_data_files); + /// This manifest only contains newly added files; no pre-existing entries. + entry.field(Iceberg::f_existing_files_count) = 0; entry.field(Iceberg::f_deleted_files_count) = 0; if (summary->has(Iceberg::f_added_position_deletes)) @@ -966,90 +965,91 @@ void generateManifestList( while (reader.read(datum)) { - if (version == 1) + if (datum.type() != avro::AVRO_RECORD) + throw Exception( + ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, + "Manifest list {} contains an entry with Avro type {}, but a record is required", + relative_path_with_metadata.getPath(), + static_cast(datum.type())); + + const avro::GenericRecord & old_entry = datum.value(); + + auto validate_field_type = [&](const String & field_name, avro::Type expected_type) -> const avro::GenericDatum & { - if (datum.type() != avro::AVRO_RECORD) + if (!old_entry.hasField(field_name)) throw Exception( ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, - "Manifest list {} contains an entry with Avro type {}, but a record is required", + "Manifest list {} entry is missing required field '{}'", relative_path_with_metadata.getPath(), - static_cast(datum.type())); + field_name); - const avro::GenericRecord & old_entry = datum.value(); + const avro::GenericDatum & field = old_entry.field(field_name); + if (field.type() != expected_type) + throw Exception( + ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, + "Manifest list {} field '{}' has Avro type {}, but type {} is required", + relative_path_with_metadata.getPath(), + field_name, + static_cast(field.type()), + static_cast(expected_type)); - auto validate_field_type = [&](const String & field_name, avro::Type expected_type) -> const avro::GenericDatum & - { - if (!old_entry.hasField(field_name)) - throw Exception( - ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, - "Manifest list {} entry is missing required field '{}'", - relative_path_with_metadata.getPath(), - field_name); - - const avro::GenericDatum & field = old_entry.field(field_name); - if (field.type() != expected_type) - throw Exception( - ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, - "Manifest list {} field '{}' has Avro type {}, but type {} is required", - relative_path_with_metadata.getPath(), - field_name, - static_cast(field.type()), - static_cast(expected_type)); - - return field; - }; - - const avro::GenericDatum & old_manifest_path = validate_field_type(Iceberg::f_manifest_path, avro::AVRO_STRING); - - avro::GenericDatum new_datum(schema.root()); - avro::GenericRecord & new_entry = new_datum.value(); - - auto copy_required_field = [&](const String & field_name, avro::Type expected_type) - { - new_entry.field(field_name) = validate_field_type(field_name, expected_type); - }; - - new_entry.field(f_manifest_path) = old_manifest_path; - copy_required_field(Iceberg::f_manifest_length, avro::AVRO_LONG); - copy_required_field(Iceberg::f_partition_spec_id, avro::AVRO_INT); - /// iceberg-spark changed `f_added_snapshot_id` from 'null, long' to 'long' (apache/iceberg#11626); rewrite with the new schema in case we read the old type. - if (old_entry.hasField(Iceberg::f_added_snapshot_id)) - { - const avro::GenericDatum & old_added_snapshot_id_entry = old_entry.field(Iceberg::f_added_snapshot_id); - if (old_added_snapshot_id_entry.type() != avro::AVRO_LONG) - throw Exception( - ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, - "Manifest list {} field '{}' has Avro type {}, but a non-null long is required", - relative_path_with_metadata.getPath(), - Iceberg::f_added_snapshot_id, - static_cast(old_added_snapshot_id_entry.type())); - - new_entry.field(f_added_snapshot_id) = old_added_snapshot_id_entry.value(); - } - else - /// This only happens when we read data written by a old version of iceberg, which violent the spec of iceberg. + return field; + }; + + const avro::GenericDatum & old_manifest_path = validate_field_type(Iceberg::f_manifest_path, avro::AVRO_STRING); + + avro::GenericDatum new_datum(schema.root()); + avro::GenericRecord & new_entry = new_datum.value(); + + auto copy_required_field = [&](const String & field_name, avro::Type expected_type) + { + new_entry.field(field_name) = validate_field_type(field_name, expected_type); + }; + + new_entry.field(f_manifest_path) = old_manifest_path; + copy_required_field(Iceberg::f_manifest_length, avro::AVRO_LONG); + copy_required_field(Iceberg::f_partition_spec_id, avro::AVRO_INT); + /// iceberg-spark changed `f_added_snapshot_id` from 'null, long' to 'long' (apache/iceberg#11626); rewrite with the new schema in case we read the old type. + if (old_entry.hasField(Iceberg::f_added_snapshot_id)) + { + const avro::GenericDatum & old_added_snapshot_id_entry = old_entry.field(Iceberg::f_added_snapshot_id); + if (old_added_snapshot_id_entry.type() != avro::AVRO_LONG) throw Exception( ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, - "Manifest list {} has null value for field '{}', but it is required", + "Manifest list {} field '{}' has Avro type {}, but a non-null long is required", relative_path_with_metadata.getPath(), - Iceberg::f_added_snapshot_id); - auto add_field_to_datum = [&](const String & field) - { - if (old_entry.hasField(field)) - new_entry.field(field) = old_entry.field(field); - }; - add_field_to_datum(Iceberg::f_added_files_count); - add_field_to_datum(Iceberg::f_existing_files_count); - add_field_to_datum(Iceberg::f_deleted_files_count); - add_field_to_datum(Iceberg::f_partitions); - add_field_to_datum(Iceberg::f_added_rows_count); - add_field_to_datum(Iceberg::f_existing_rows_count); - add_field_to_datum(Iceberg::f_deleted_rows_count); - add_field_to_datum(Iceberg::f_key_metadata); - writer.write(new_datum); + Iceberg::f_added_snapshot_id, + static_cast(old_added_snapshot_id_entry.type())); + + new_entry.field(f_added_snapshot_id) = old_added_snapshot_id_entry.value(); } else - writer.write(datum); + /// This only happens when we read data written by a old version of iceberg, which violates the spec of iceberg. + throw Exception( + ErrorCodes::ICEBERG_SPECIFICATION_VIOLATION, + "Manifest list {} has null value for field '{}', but it is required", + relative_path_with_metadata.getPath(), + Iceberg::f_added_snapshot_id); + auto add_field_to_datum = [&](const String & field) + { + if (old_entry.hasField(field)) + new_entry.field(field) = old_entry.field(field); + }; + add_field_to_datum(Iceberg::f_added_files_count); + add_field_to_datum(Iceberg::f_existing_files_count); + add_field_to_datum(Iceberg::f_deleted_files_count); + add_field_to_datum(Iceberg::f_partitions); + add_field_to_datum(Iceberg::f_added_rows_count); + add_field_to_datum(Iceberg::f_existing_rows_count); + add_field_to_datum(Iceberg::f_deleted_rows_count); + add_field_to_datum(Iceberg::f_key_metadata); + if (version == 2) + { + add_field_to_datum(Iceberg::f_content); + add_field_to_datum(Iceberg::f_sequence_number); + add_field_to_datum(Iceberg::f_min_sequence_number); + } + writer.write(new_datum); } break; } @@ -1105,6 +1105,13 @@ IcebergStorageSink::IcebergStorageSink( filename_generator.setVersion(last_version + 1); + if (metadata->has(Iceberg::f_properties)) + { + auto properties = metadata->getObject(Iceberg::f_properties); + if (properties && properties->has("write.data.path")) + filename_generator.setDataLocation(properties->getValue("write.data.path")); + } + partition_spec_id = metadata->getValue(Iceberg::f_default_spec_id); auto partitions_specs = metadata->getArray(Iceberg::f_partition_specs); @@ -1277,12 +1284,18 @@ void IcebergStorageSink::finalizeBuffers() /// TODO: there's a chance that initializeMetadata() doesn't succeed within MAX_TRANSACTION_RETRIES without throwing, perhaps we should fail in this case size_t i = 0; + bool successed_write = false; while (i < MAX_TRANSACTION_RETRIES) { if (initializeMetadata()) + { + successed_write = true; break; + } ++i; } + if (!successed_write) + throw Exception(ErrorCodes::DATALAKE_DATABASE_ERROR, "Write into iceberg was not successful"); } void IcebergStorageSink::releaseBuffers() diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.cpp b/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.cpp index 7d2b2ac6c5507..e702abca8e10d 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.cpp +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -10,6 +11,7 @@ #include #include +#include #if USE_AVRO @@ -72,12 +74,17 @@ bool checkValidSchemaEvolution(Poco::Dynamic::Var old_type, Poco::Dynamic::Var n MetadataGenerator::MetadataGenerator(Poco::JSON::Object::Ptr metadata_object_) : metadata_object(metadata_object_) , gen(randomSeed()) - , dis(0, INT32_MAX) + , dis(1, std::numeric_limits::max()) { } Int64 MetadataGenerator::getMaxSequenceNumber() { + /// Use the authoritative top-level field per Iceberg V2 spec. + /// Iterating snapshots is unreliable when catalogs prune snapshot history. + if (metadata_object->has(Iceberg::f_last_sequence_number)) + return metadata_object->getValue(Iceberg::f_last_sequence_number); + auto snapshots = metadata_object->get(Iceberg::f_snapshots).extract(); Int64 max_seq_number = 0; @@ -122,7 +129,7 @@ MetadataGenerator::NextMetadataResult MetadataGenerator::generateNextMetadata( if (format_version > 1) { auto sequence_number = getMaxSequenceNumber() + 1; - new_snapshot->set(Iceberg::f_metadata_sequence_number, getMaxSequenceNumber() + 1); + new_snapshot->set(Iceberg::f_metadata_sequence_number, sequence_number); metadata_object->set(Iceberg::f_last_sequence_number, sequence_number); } Int64 snapshot_id = user_defined_snapshot_id.value_or(static_cast(dis(gen))); diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.h b/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.h index fa9c340b86f0a..3576837f6f70a 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.h +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/MetadataGenerator.h @@ -51,7 +51,7 @@ class MetadataGenerator Poco::JSON::Object::Ptr metadata_object; pcg64_fast gen; - std::uniform_int_distribution dis; + std::uniform_int_distribution dis; Int64 getMaxSequenceNumber(); Poco::JSON::Object::Ptr getParentSnapshot(Int64 parent_snapshot_id); diff --git a/src/Storages/ObjectStorage/DataLakes/Iceberg/Utils.cpp b/src/Storages/ObjectStorage/DataLakes/Iceberg/Utils.cpp index eb0157c80645f..98752765c1554 100644 --- a/src/Storages/ObjectStorage/DataLakes/Iceberg/Utils.cpp +++ b/src/Storages/ObjectStorage/DataLakes/Iceberg/Utils.cpp @@ -103,7 +103,7 @@ namespace DB::Setting /// Hard to imagine a hint file larger than 10 MB static constexpr size_t MAX_HINT_FILE_SIZE = 10 * 1024 * 1024; -static constexpr auto MAX_TRANSACTION_RETRIES = 100; +static constexpr auto MAX_TRANSACTION_RETRIES = 1000; namespace DB::Iceberg {