Skip to content

Antalya 26.3 Backport of #98162 - Improve iceberg writes for unity catalog #2714

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xieandrew wants to merge 1 commit into
base: antalya-26.3
Choose a base branch
from
Open

Antalya 26.3 Backport of #98162 - Improve iceberg writes for unity catalog #2714

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 3 additions & 3 deletions src/Databases/DataLake/RestCatalog.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
{
Expand Down Expand Up @@ -1178,7 +1178,7 @@ void RestCatalog::createTable(const String & namespace_name, const String & tabl

createNamespaceIfNotExists(namespace_name, metadata_content->getValue<String>("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);
Expand Down Expand Up @@ -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;
{
Expand Down
7 changes: 7 additions & 0 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/AvroSchema.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
]
}
Expand Down
6 changes: 6 additions & 0 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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));
}

Expand Down Expand Up @@ -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));
}

Expand Down
8 changes: 8 additions & 0 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/FileNamesGenerator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
171 changes: 92 additions & 79 deletions src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergWrites.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,6 @@
#include <Common/isValidUTF8.h>
#include <Common/quoteString.h>
#include <Common/randomSeed.h>

#include <cmath>
#include <cstdint>
#include <memory>
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -513,7 +513,6 @@ void generateManifestFile(
int current_schema_id = metadata->getValue<Int32>(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<OutputStreamWriteBufferAdapter>(buf);
Expand All @@ -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)
Expand Down Expand Up @@ -637,7 +636,7 @@ void generateManifestFile(
{
avro::GenericDatum record_datum(schema_element);
auto & record = record_datum.value<avro::GenericRecord>();
record.field(Iceberg::f_key) = static_cast<Int64>(field_id);
record.field(Iceberg::f_key) = static_cast<Int32>(field_id);
record.field(Iceberg::f_value) = dump_function(field_id, value);
record_values.value().push_back(record_datum);
}
Expand Down Expand Up @@ -909,8 +908,8 @@ void generateManifestList(
else
{
entry.field(Iceberg::f_added_files_count) = 1;
entry.field(Iceberg::f_existing_files_count)
= summary->getValue<Int32>(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))
Expand Down Expand Up @@ -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<int>(datum.type()));

const avro::GenericRecord & old_entry = datum.value<avro::GenericRecord>();

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<int>(datum.type()));
field_name);

const avro::GenericRecord & old_entry = datum.value<avro::GenericRecord>();
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<int>(field.type()),
static_cast<int>(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<int>(field.type()),
static_cast<int>(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<avro::GenericRecord>();

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<int>(old_added_snapshot_id_entry.type()));

new_entry.field(f_added_snapshot_id) = old_added_snapshot_id_entry.value<Int64>();
}
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<avro::GenericRecord>();

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<int>(old_added_snapshot_id_entry.type()));

new_entry.field(f_added_snapshot_id) = old_added_snapshot_id_entry.value<Int64>();
}
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;
}
Expand Down Expand Up @@ -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<String>("write.data.path"));
}

partition_spec_id = metadata->getValue<Int64>(Iceberg::f_default_spec_id);
auto partitions_specs = metadata->getArray(Iceberg::f_partition_specs);

Expand Down Expand Up @@ -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()
Expand Down
Loading
Loading