Export partition to apache iceberg

src/Common/FailPoint.cpp

@@ -119,6 +119,7 @@ static struct InitFiu
     ONCE(disk_object_storage_fail_precommit_metadata_transaction) \
     REGULAR(slowdown_parallel_replicas_local_plan_read) \
     ONCE(iceberg_writes_cleanup) \
+    ONCE(iceberg_export_after_commit_before_zk_completed) \
     ONCE(backup_add_empty_memory_table) \
     PAUSEABLE_ONCE(backup_pause_on_start) \
     PAUSEABLE_ONCE(restore_pause_on_start) \

src/Storages/ExportReplicatedMergeTreePartitionManifest.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <base/types.h>
+#include <Core/Field.h>
 #include <Interpreters/StorageID.h>
 #include <Poco/JSON/Object.h>
 #include <Poco/JSON/Array.h>
@@ -118,6 +119,14 @@ struct ExportReplicatedMergeTreePartitionManifest
     MergeTreePartExportManifest::FileAlreadyExistsPolicy file_already_exists_policy;
     String filename_pattern;
     bool lock_inside_the_task; /// todo temporary
+    bool write_full_path_in_iceberg_metadata = false;
+    String iceberg_metadata_json;
+    /// Iceberg-only: JSON array of partition column values (after transforms) for this partition.
+    /// Columns and types are derived at commit time from iceberg_metadata_json; only values are persisted.
+    String partition_values_json;
+    /// Transient: parsed form of partition_values_json, populated by fromJsonString.
+    /// Not serialized to ZooKeeper. Used at commit time to avoid re-parsing JSON on each commit attempt.
+    std::vector<Field> partition_values;
 
     std::string toJsonString() const
     {
@@ -129,7 +138,17 @@ struct ExportReplicatedMergeTreePartitionManifest
         json.set("destination_table", destination_table);
         json.set("source_replica", source_replica);
         json.set("number_of_parts", number_of_parts);
-
+
+        if (!iceberg_metadata_json.empty())
+        {
+            json.set("iceberg_metadata_json", iceberg_metadata_json);
+        }
+
+        if (!partition_values_json.empty())
+        {
+            json.set("partition_values_json", partition_values_json);
+        }
+
         Poco::JSON::Array::Ptr parts_array = new Poco::JSON::Array();
         for (const auto & part : parts)
             parts_array->add(part);
@@ -145,6 +164,7 @@ struct ExportReplicatedMergeTreePartitionManifest
         json.set("max_retries", max_retries);
         json.set("ttl_seconds", ttl_seconds);
         json.set("lock_inside_the_task", lock_inside_the_task);
+        json.set("write_full_path_in_iceberg_metadata", write_full_path_in_iceberg_metadata);
         std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
         oss.exceptions(std::ios::failbit);
         Poco::JSON::Stringifier::stringify(json, oss);
@@ -166,6 +186,28 @@ struct ExportReplicatedMergeTreePartitionManifest
         manifest.source_replica = json->getValue<String>("source_replica");
         manifest.number_of_parts = json->getValue<size_t>("number_of_parts");
         manifest.max_retries = json->getValue<size_t>("max_retries");
+
+        if (json->has("iceberg_metadata_json"))
+        {
+            manifest.iceberg_metadata_json = json->getValue<String>("iceberg_metadata_json");
+        }
+
+        if (json->has("partition_values_json"))
+        {
+            manifest.partition_values_json = json->getValue<String>("partition_values_json");
+
+            Poco::JSON::Parser val_parser;
+            auto arr = val_parser.parse(manifest.partition_values_json).extract<Poco::JSON::Array::Ptr>();
+            for (size_t i = 0; i < arr->size(); ++i)
+            {
+                Poco::Dynamic::Var var = arr->get(static_cast<unsigned int>(i));
+                if (var.isString())
+                    manifest.partition_values.push_back(Field(var.extract<String>()));
+                else
+                    manifest.partition_values.push_back(Field(var.convert<Int64>()));
+            }
+        }
+
         auto parts_array = json->getArray("parts");
         for (size_t i = 0; i < parts_array->size(); ++i)
             manifest.parts.push_back(parts_array->getElement<String>(static_cast<unsigned int>(i)));
@@ -192,6 +234,8 @@ struct ExportReplicatedMergeTreePartitionManifest
 
         manifest.lock_inside_the_task = json->getValue<bool>("lock_inside_the_task");
 
+        manifest.write_full_path_in_iceberg_metadata = json->getValue<bool>("write_full_path_in_iceberg_metadata");
+
         return manifest;
     }
 };

src/Storages/IStorage.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <Core/Field.h>
 #include <Core/Names.h>
 #include <Core/QueryProcessingStage.h>
 #include <Databases/IDatabase.h>
@@ -20,6 +21,7 @@
 #include <Common/RWLock.h>
 #include <Common/TypePromotion.h>
 #include <DataTypes/Serializations/SerializationInfo.h>
+#include <Poco/JSON/Object.h>
 
 #include <expected>
 #include <optional>
@@ -476,7 +478,7 @@ class IStorage : public std::enable_shared_from_this<IStorage>, public TypePromo
         ContextPtr /*context*/,
         bool /*async_insert*/);
 
-    virtual bool supportsImport() const
+    virtual bool supportsImport(ContextPtr) const
     {
       return false;
     }
@@ -493,16 +495,28 @@ It is currently only implemented in StorageObjectStorage.
         bool /* overwrite_if_exists */,
         std::size_t /* max_bytes_per_file */,
         std::size_t /* max_rows_per_file */,
+        const std::optional<std::string> & /* iceberg_metadata_json_string */,
         const std::optional<FormatSettings> & /* format_settings */,
         ContextPtr /* context */)
     {
       throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Import is not implemented for storage {}", getName());
     }
 
+    struct IcebergCommitExportPartitionArguments
+    {
+      std::string metadata_json_string;
+      /// Partition column values (after transforms). Callers are responsible for
+      /// populating this: the partition-export path parses them from the persisted
+      /// JSON string, while the direct EXPORT PART path reads them from the part's
+      /// partition key.
+      std::vector<Field> partition_values;
+    };
+
     virtual void commitExportPartitionTransaction(
       const String & /* transaction_id */,
       const String & /* partition_id */,
       const Strings & /* exported_paths */,
+      const IcebergCommitExportPartitionArguments & /* iceberg_commit_export_partition_arguments */,
       ContextPtr /* local_context */)
   {
       throw Exception(ErrorCodes::NOT_IMPLEMENTED, "commitExportPartitionTransaction is not implemented for storage type {}", getName());

src/Storages/MergeTree/ExportPartTask.cpp

@@ -21,6 +21,7 @@
 #include <Common/ProfileEventsScope.h>
 #include <Databases/DatabaseReplicated.h>
 #include <Storages/MergeTree/ExportList.h>
+#include <Storages/IStorage.h>
 #include <Formats/FormatFactory.h>
 #include <Databases/enableAllExperimentalSettings.h>
 #include <Processors/Sinks/SinkToStorage.h>
@@ -218,6 +219,7 @@ bool ExportPartTask::executeStep()
             manifest.file_already_exists_policy == MergeTreePartExportManifest::FileAlreadyExistsPolicy::overwrite,
             manifest.settings[Setting::export_merge_tree_part_max_bytes_per_file],
             manifest.settings[Setting::export_merge_tree_part_max_rows_per_file],
+            manifest.iceberg_metadata_json,
             getFormatSettings(local_context),
             local_context);
 
@@ -295,6 +297,23 @@ bool ExportPartTask::executeStep()
             throw Exception(ErrorCodes::QUERY_WAS_CANCELLED, "Export part was cancelled");
         }
 
+        /// For the direct EXPORT PART → Iceberg path there is no deferred-commit callback
+        /// (the partition-export path provides one that writes to ZooKeeper).
+        /// Commit the Iceberg metadata inline here so the rows become visible immediately.
+        if (destination_storage->isDataLake() && !manifest.completion_callback)
+        {
+            IStorage::IcebergCommitExportPartitionArguments iceberg_args;
+            iceberg_args.metadata_json_string = manifest.iceberg_metadata_json;
+            iceberg_args.partition_values = manifest.partition_values;
+
+            destination_storage->commitExportPartitionTransaction(
+                manifest.transaction_id,
+                manifest.data_part->info.getPartitionId(),
+                (*exports_list_entry)->destination_file_paths,
+                iceberg_args,
+                local_context);
+        }
+
         std::lock_guard inner_lock(storage.export_manifests_mutex);
         storage.writePartLog(
             PartLogElement::Type::EXPORT_PART,

src/Storages/MergeTree/ExportPartitionManifestUpdatingTask.cpp

@@ -34,7 +34,7 @@ namespace
         const zkutil::ZooKeeperPtr & zk,
         const std::string & entry_path,
         const LoggerPtr & log,
-        const ContextPtr & context,
+        const ContextPtr & storage_context,
         const std::string & key,
         const ExportReplicatedMergeTreePartitionManifest & metadata,
         const time_t now,
@@ -58,6 +58,8 @@ namespace
         }
         else if (is_pending)
         {
+            auto context = ExportPartitionUtils::getContextCopyWithTaskSettings(storage_context, metadata);
+
             ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
             ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperGetChildren);
             std::vector<std::string> parts_in_processing_or_pending;

src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp

@@ -40,25 +40,57 @@ namespace ErrorCodes
 
 namespace
 {
-    ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ExportReplicatedMergeTreePartitionManifest & manifest)
+
+
+    Coordination::Requests getErrorRequests(
+        const std::filesystem::path & export_path,
+        const std::string & replica_name,
+        const zkutil::ZooKeeperPtr & zk,
+        const LoggerPtr & log,
+        const std::string & part_name,
+        const std::optional<Exception> & exception
+    )
     {
-        auto context_copy = Context::createCopy(context);
-        context_copy->makeQueryContextForExportPart();
-        context_copy->setCurrentQueryId(manifest.query_id);
-        context_copy->setSetting("output_format_parallel_formatting", manifest.parallel_formatting);
-        context_copy->setSetting("output_format_parquet_parallel_encoding", manifest.parquet_parallel_encoding);
-        context_copy->setSetting("max_threads", manifest.max_threads);
-        context_copy->setSetting("export_merge_tree_part_file_already_exists_policy", String(magic_enum::enum_name(manifest.file_already_exists_policy)));
-        context_copy->setSetting("export_merge_tree_part_max_bytes_per_file", manifest.max_bytes_per_file);
-        context_copy->setSetting("export_merge_tree_part_max_rows_per_file", manifest.max_rows_per_file);
-
-        /// always skip pending mutations and patch parts because we already validated the parts during query processing
-        context_copy->setSetting("export_merge_tree_part_throw_on_pending_mutations", false);
-        context_copy->setSetting("export_merge_tree_part_throw_on_pending_patch_parts", false);
-
-        context_copy->setSetting("export_merge_tree_part_filename_pattern", manifest.filename_pattern);
-
-	return context_copy;
+        Coordination::Requests ops;
+
+        const auto exceptions_per_replica_path = export_path / "exceptions_per_replica" / replica_name;
+        const auto count_path = exceptions_per_replica_path / "count";
+        const auto last_exception_path = exceptions_per_replica_path / "last_exception";
+
+        ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
+        ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperExists);
+        if (zk->exists(exceptions_per_replica_path))
+        {
+            LOG_INFO(log, "ExportPartition scheduler task: Exceptions per replica path exists, no need to create it");
+            std::string num_exceptions_string;
+            if (zk->tryGet(count_path, num_exceptions_string))
+            {
+                const auto num_exceptions = parse<size_t>(num_exceptions_string) + 1;
+                ops.emplace_back(zkutil::makeSetRequest(count_path, std::to_string(num_exceptions), -1));
+            }
+            else
+            {
+                /// TODO maybe we should find a better way to handle this case, not urgent
+                LOG_INFO(log, "ExportPartition scheduler task: Failed to get number of exceptions, will not increment it");
+            }
+
+            ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
+            ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperGet);
+
+            ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "part", part_name, -1));
+            ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "exception", exception->message(), -1));
+        }
+        else
+        {
+            LOG_INFO(log, "ExportPartition scheduler task: Exceptions per replica path does not exist, will create it");
+            ops.emplace_back(zkutil::makeCreateRequest(exceptions_per_replica_path, "", zkutil::CreateMode::Persistent));
+            ops.emplace_back(zkutil::makeCreateRequest(count_path, "1", zkutil::CreateMode::Persistent));
+            ops.emplace_back(zkutil::makeCreateRequest(last_exception_path, "", zkutil::CreateMode::Persistent));
+            ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "part", part_name, zkutil::CreateMode::Persistent));
+            ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "exception", exception->message(), zkutil::CreateMode::Persistent));
+        }
+
+        return ops;
     }
 }
 
@@ -199,7 +231,7 @@ void ExportPartitionTaskScheduler::run()
 
             LOG_INFO(storage.log, "ExportPartition scheduler task: Scheduling part export: {}", zk_part_name);
 
-            auto context = getContextCopyWithTaskSettings(storage.getContext(), manifest);
+            auto context = ExportPartitionUtils::getContextCopyWithTaskSettings(storage.getContext(), manifest);
 
             /// todo arthur this code path does not perform all the validations a simple part export does because we are not calling exportPartToTable directly.
             /// the schema and everything else has been validated when the export partition task was created, but nothing prevents the destination table from being
@@ -217,11 +249,13 @@ void ExportPartitionTaskScheduler::run()
                     context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value,
                     context->getSettingsCopy(),
                     storage.getInMemoryMetadataPtr(),
+                    manifest.iceberg_metadata_json,
                     [this, key, zk_part_name, manifest, destination_storage]
                     (MergeTreePartExportManifest::CompletionCallbackResult result)
                     {
                         handlePartExportCompletion(key, zk_part_name, manifest, destination_storage, result);
-                    });
+                    },
+                    manifest.partition_values);
 
                 part_export_manifest.task = std::make_shared<ExportPartFromPartitionExportTask>(storage, key, part_export_manifest);
 
@@ -263,13 +297,15 @@ void ExportPartitionTaskScheduler::run()
                         part->name,
                         destination_storage_id,
                         manifest.transaction_id,
-                        getContextCopyWithTaskSettings(storage.getContext(), manifest),
+                        context,
+                        manifest.iceberg_metadata_json,
                         /*allow_outdated_parts*/ true,
                         [this, key, zk_part_name, manifest, destination_storage]
                         (MergeTreePartExportManifest::CompletionCallbackResult result)
                         {
                             handlePartExportCompletion(key, zk_part_name, manifest, destination_storage, result);
-                        });
+                        },
+                        manifest.partition_values);
 
                     scheduled_exports_count++;
                 }
@@ -342,7 +378,25 @@ void ExportPartitionTaskScheduler::handlePartExportSuccess(
 
     LOG_INFO(storage.log, "ExportPartition scheduler task: All parts are processed, will try to commit export partition");
 
-    ExportPartitionUtils::commit(manifest, destination_storage, zk, storage.log.load(), export_path, storage.getContext());
+    try
+    {
+        auto context = ExportPartitionUtils::getContextCopyWithTaskSettings(storage.getContext(), manifest);
+        ExportPartitionUtils::commit(manifest, destination_storage, zk, storage.log.load(), export_path, context);
+    }
+    catch (const Exception & e)
+    {
+        const auto error_requests = getErrorRequests(export_path, storage.replica_name, zk, storage.log.load(), part_name, e);
+
+        LOG_INFO(storage.log, "ExportPartition scheduler task: Caught exception while committing export partition, {}", e.message());
+
+        ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
+        ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperMulti);
+        Coordination::Responses responses;
+        if (Coordination::Error::ZOK != zk->tryMulti(error_requests, responses))
+        {
+            LOG_INFO(storage.log, "ExportPartition scheduler task: Failed to update zookeeper with the commit exception");
+        }
+    }
 }
 
 void ExportPartitionTaskScheduler::handlePartExportFailure(
@@ -450,42 +504,8 @@ void ExportPartitionTaskScheduler::handlePartExportFailure(
         LOG_INFO(storage.log, "ExportPartition scheduler task: Retry count limit not exceeded for part {}, will increment retry count", part_name);
     }
 
-    const auto exceptions_per_replica_path = export_path / "exceptions_per_replica" / storage.replica_name;
-    const auto count_path = exceptions_per_replica_path / "count";
-    const auto last_exception_path = exceptions_per_replica_path / "last_exception";
-
-    ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
-    ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperExists);
-    if (zk->exists(exceptions_per_replica_path))
-    {
-        LOG_INFO(storage.log, "ExportPartition scheduler task: Exceptions per replica path exists, no need to create it");
-        std::string num_exceptions_string;
-        if (zk->tryGet(count_path, num_exceptions_string))
-        {
-            const auto num_exceptions = parse<size_t>(num_exceptions_string) + 1;
-            ops.emplace_back(zkutil::makeSetRequest(count_path, std::to_string(num_exceptions), -1));
-        }
-        else
-        {
-            /// TODO maybe we should find a better way to handle this case, not urgent
-            LOG_INFO(storage.log, "ExportPartition scheduler task: Failed to get number of exceptions, will not increment it");
-        }
-
-        ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
-        ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperGet);
-
-        ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "part", part_name, -1));
-        ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "exception", exception->message(), -1));
-    }
-    else
-    {
-        LOG_INFO(storage.log, "ExportPartition scheduler task: Exceptions per replica path does not exist, will create it");
-        ops.emplace_back(zkutil::makeCreateRequest(exceptions_per_replica_path, "", zkutil::CreateMode::Persistent));
-        ops.emplace_back(zkutil::makeCreateRequest(count_path, "1", zkutil::CreateMode::Persistent));
-        ops.emplace_back(zkutil::makeCreateRequest(last_exception_path, "", zkutil::CreateMode::Persistent));
-        ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "part", part_name, zkutil::CreateMode::Persistent));
-        ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "exception", exception->message(), zkutil::CreateMode::Persistent));
-    }
+    auto error_requests = getErrorRequests(export_path, storage.replica_name, zk, storage.log.load(), part_name, exception);
+    ops.insert(ops.end(), error_requests.begin(), error_requests.end());
 
     ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
     ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperMulti);