Plain merge tree partition export

src/Interpreters/InterpreterKillQueryQuery.cpp

@@ -265,65 +265,96 @@ BlockIO InterpreterKillQueryQuery::execute()
                 "Exporting merge tree partition is experimental. Set the server setting `allow_experimental_export_merge_tree_partition` to enable it");
         }
 
-        Block exports_block = getSelectResult(
-            "source_database, source_table, transaction_id, destination_database, destination_table, partition_id",
-            "system.replicated_partition_exports");
-        if (exports_block.empty())
-            return res_io;
+        const String export_columns = "source_database, source_table, transaction_id, destination_database, destination_table, partition_id";
+
+        /// Partition exports live in two system tables: `replicated_partition_exports` for
+        /// ReplicatedMergeTree and `partition_exports` for plain MergeTree. Query both so a single
+        /// KILL EXPORT PARTITION targets either engine. Each query applies the user WHERE to its own
+        /// table (which exposes all of its columns); a WHERE that references columns specific to the
+        /// other engine simply yields no matches / is tolerated.
+        Block replicated_exports_block = getSelectResult(export_columns, "system.replicated_partition_exports");
+        Block plain_exports_block;
+        try
+        {
+            plain_exports_block = getSelectResult(export_columns, "system.partition_exports");
+        }
+        catch (...)
+        {
+            tryLogCurrentException(getLogger("InterpreterKillQueryQuery"),
+                "KILL EXPORT PARTITION: could not read system.partition_exports (the WHERE may reference "
+                "columns that only exist for ReplicatedMergeTree); ignoring plain MergeTree tables");
+        }
 
-        const ColumnString & src_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_database").column);
-        const ColumnString & src_table_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_table").column);
-        const ColumnString & dst_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_database").column);
-        const ColumnString & dst_table_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_table").column);
-        const ColumnString & tx_col = typeid_cast<const ColumnString &>(*exports_block.getByName("transaction_id").column);
+        if (replicated_exports_block.empty() && plain_exports_block.empty())
+            return res_io;
 
-        auto header = exports_block.cloneEmpty();
+        /// Build the result header explicitly from the fixed projection so it does not depend on
+        /// whether either source block ended up with rows.
+        Block header;
+        for (const auto & column_name : {"source_database", "source_table", "transaction_id",
+                                         "destination_database", "destination_table", "partition_id"})
+            header.insert({ColumnString::create(), std::make_shared<DataTypeString>(), column_name});
         header.insert(0, {ColumnString::create(), std::make_shared<DataTypeString>(), "kill_status"});
 
         MutableColumns res_columns = header.cloneEmptyColumns();
         AccessRightsElements required_access_rights;
         auto access = getContext()->getAccess();
         bool access_denied = false;
 
-        for (size_t i = 0; i < exports_block.rows(); ++i)
+        auto process_block = [&](const Block & exports_block)
         {
-            const auto src_database = src_db_col.getDataAt(i);
-            const auto src_table = src_table_col.getDataAt(i);
-            const auto dst_database = dst_db_col.getDataAt(i);
-            const auto dst_table = dst_table_col.getDataAt(i);
+            if (exports_block.empty())
+                return;
 
-            const auto table_id = StorageID{std::string{src_database}, std::string{src_table}};
-            const auto transaction_id = tx_col.getDataAt(i);
+            const ColumnString & src_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_database").column);
+            const ColumnString & src_table_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_table").column);
+            const ColumnString & dst_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_database").column);
+            const ColumnString & dst_table_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_table").column);
+            const ColumnString & tx_col = typeid_cast<const ColumnString &>(*exports_block.getByName("transaction_id").column);
 
-            CancellationCode code = CancellationCode::Unknown;
-            if (!query.test)
+            for (size_t i = 0; i < exports_block.rows(); ++i)
             {
-                auto storage = DatabaseCatalog::instance().tryGetTable(table_id, getContext());
-                if (!storage)
-                    code = CancellationCode::NotFound;
-                else
-                {
-                    ASTAlterCommand alter_command{};
-                    alter_command.type = ASTAlterCommand::EXPORT_PARTITION;
-                    alter_command.move_destination_type = DataDestinationType::TABLE;
-                    alter_command.from_database = src_database;
-                    alter_command.from_table = src_table;
-                    alter_command.to_database = dst_database;
-                    alter_command.to_table = dst_table;
+                const auto src_database = src_db_col.getDataAt(i);
+                const auto src_table = src_table_col.getDataAt(i);
+                const auto dst_database = dst_db_col.getDataAt(i);
+                const auto dst_table = dst_table_col.getDataAt(i);
 
-                    required_access_rights = InterpreterAlterQuery::getRequiredAccessForCommand(
-                        alter_command, table_id.database_name, table_id.table_name);
-                    if (!access->isGranted(required_access_rights))
+                const auto table_id = StorageID{std::string{src_database}, std::string{src_table}};
+                const auto transaction_id = tx_col.getDataAt(i);
+
+                CancellationCode code = CancellationCode::Unknown;
+                if (!query.test)
+                {
+                    auto storage = DatabaseCatalog::instance().tryGetTable(table_id, getContext());
+                    if (!storage)
+                        code = CancellationCode::NotFound;
+                    else
                     {
-                        access_denied = true;
-                        continue;
+                        ASTAlterCommand alter_command{};
+                        alter_command.type = ASTAlterCommand::EXPORT_PARTITION;
+                        alter_command.move_destination_type = DataDestinationType::TABLE;
+                        alter_command.from_database = src_database;
+                        alter_command.from_table = src_table;
+                        alter_command.to_database = dst_database;
+                        alter_command.to_table = dst_table;
+
+                        required_access_rights = InterpreterAlterQuery::getRequiredAccessForCommand(
+                            alter_command, table_id.database_name, table_id.table_name);
+                        if (!access->isGranted(required_access_rights))
+                        {
+                            access_denied = true;
+                            continue;
+                        }
+                        code = storage->killExportPartition(std::string{transaction_id});
                     }
-                    code = storage->killExportPartition(std::string{transaction_id});
                 }
+
+                insertResultRow(i, code, exports_block, header, res_columns);
             }
+        };
 
-            insertResultRow(i, code, exports_block, header, res_columns);
-        }
+        process_block(replicated_exports_block);
+        process_block(plain_exports_block);
 
         if (res_columns[0]->empty() && access_denied)
             throw Exception(ErrorCodes::ACCESS_DENIED, "Not allowed to kill export partition. "

src/Storages/MergeTree/ExportPartitionUtils.cpp

@@ -7,6 +7,10 @@
 #include "Storages/ExportReplicatedMergeTreePartitionManifest.h"
 #include "Storages/ExportReplicatedMergeTreePartitionTaskEntry.h"
 #include <Storages/MergeTree/MergeTreeData.h>
+#include <Storages/MergeTree/MergeTreePartitionExportTask.h>
+#include <Storages/ObjectStorage/StorageObjectStorage.h>
+#include <Storages/ObjectStorage/StorageObjectStorageCluster.h>
+#include <Parsers/IAST.h>
 #include <filesystem>
 #include <thread>
 #include <unordered_set>
@@ -20,6 +24,7 @@
 #if USE_AVRO
 #include <Storages/ObjectStorage/DataLakes/Iceberg/Constant.h>
 #include <Storages/ObjectStorage/DataLakes/Iceberg/Utils.h>
+#include <Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.h>
 #endif
 
 namespace ProfileEvents
@@ -75,6 +80,7 @@ namespace ErrorCodes
 namespace Setting
 {
     extern const SettingsBool export_merge_tree_part_allow_lossy_cast;
+    extern const SettingsBool allow_insert_into_iceberg;
 }
 
 namespace FailPoints
@@ -147,7 +153,8 @@ namespace ExportPartitionUtils
         return parts.front()->minmax_idx->getBlock(storage);
     }
 
-    ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ExportReplicatedMergeTreePartitionManifest & manifest)
+    template <typename ManifestT>
+    ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ManifestT & manifest)
     {
         auto context_copy = Context::createCopy(context);
         context_copy->makeQueryContextForExportPart();
@@ -187,6 +194,102 @@ namespace ExportPartitionUtils
 	    return context_copy;
     }
 
+    template ContextPtr getContextCopyWithTaskSettings<ExportReplicatedMergeTreePartitionManifest>(
+        const ContextPtr &, const ExportReplicatedMergeTreePartitionManifest &);
+    template ContextPtr getContextCopyWithTaskSettings<MergeTreePartitionExportTask>(
+        const ContextPtr &, const MergeTreePartitionExportTask &);
+
+    std::string extractDestinationIcebergMetadataJson(
+        const StorageMetadataPtr & source_metadata,
+        const StorageID & source_storage_id,
+        const StoragePtr & dest_storage,
+        const ContextPtr & context)
+    {
+        if (dest_storage->getStorageID() == source_storage_id)
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Exporting to the same table is not allowed");
+
+        if (!dest_storage->supportsImport(context))
+            throw Exception(ErrorCodes::NOT_IMPLEMENTED,
+                "Destination storage {} does not support MergeTree parts or uses unsupported partitioning",
+                dest_storage->getName());
+
+        const auto destination_metadata = dest_storage->getInMemoryMetadataPtr();
+
+        /// Positional CAST matching, like `INSERT INTO dest SELECT * FROM src`.
+        verifyExportSchemaCastable(source_metadata, destination_metadata, dest_storage->getStorageID(), context);
+
+        auto query_to_string = [](const ASTPtr & ast) { return ast ? ast->formatWithSecretsOneLine() : ""; };
+
+        /// Iceberg partition compatibility is checked below; for non-data-lake targets we only need
+        /// the partition-key ASTs to match (partition-column types follow the lossy-cast gate).
+        if (!dest_storage->isDataLake())
+        {
+            if (query_to_string(source_metadata->getPartitionKeyAST()) != query_to_string(destination_metadata->getPartitionKeyAST()))
+                throw Exception(ErrorCodes::BAD_ARGUMENTS, "Tables have different partition key");
+            return {};
+        }
+
+#if USE_AVRO
+        auto * object_storage = dynamic_cast<StorageObjectStorage *>(dest_storage.get());
+        auto * object_storage_cluster = dynamic_cast<StorageObjectStorageCluster *>(dest_storage.get());
+
+        /// in theory this should never happen, but just in case
+        if (!object_storage && !object_storage_cluster)
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Destination storage {} is not a StorageObjectStorage", dest_storage->getName());
+
+        IcebergMetadata * iceberg_metadata = nullptr;
+        if (object_storage)
+            iceberg_metadata = dynamic_cast<IcebergMetadata *>(object_storage->getExternalMetadata(context));
+        else if (object_storage_cluster)
+            iceberg_metadata = dynamic_cast<IcebergMetadata *>(object_storage_cluster->getExternalMetadata(context));
+        if (!iceberg_metadata)
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Destination storage {} is a data lake but not an iceberg table", dest_storage->getName());
+
+        if (!context->getSettingsRef()[Setting::allow_insert_into_iceberg])
+            throw Exception(ErrorCodes::SUPPORT_IS_DISABLED,
+                "Iceberg writes are experimental. "
+                "To allow its usage, enable the setting `allow_insert_into_iceberg`.");
+
+        const auto metadata_object = iceberg_metadata->getMetadataJSON(context);
+
+        verifyIcebergPartitionCompatibility(metadata_object, source_metadata->getPartitionKeyAST());
+
+        std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
+        oss.exceptions(std::ios::failbit);
+        metadata_object->stringify(oss);
+        return oss.str();
+#else
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Data lake export requires Avro support");
+#endif
+    }
+
+    void commitExportedPaths(
+        const String & transaction_id,
+        const String & partition_id,
+        const String & iceberg_metadata_json,
+        bool write_full_path_in_iceberg_metadata,
+        const std::vector<std::string> & exported_paths,
+        const StoragePtr & destination_storage,
+        MergeTreeData & source_storage,
+        const ContextPtr & context_in)
+    {
+        auto context = Context::createCopy(context_in);
+        context->setSetting("write_full_path_in_iceberg_metadata", write_full_path_in_iceberg_metadata);
+
+        IStorage::IcebergCommitExportPartitionArguments iceberg_args;
+
+        if (!iceberg_metadata_json.empty())
+        {
+            iceberg_args.metadata_json_string = iceberg_metadata_json;
+            if (source_storage.getInMemoryMetadataPtr()->hasPartitionKey())
+                iceberg_args.partition_source_block =
+                    getPartitionSourceBlockForIcebergCommit(source_storage, partition_id);
+        }
+
+        destination_storage->commitExportPartitionTransaction(
+            transaction_id, partition_id, exported_paths, iceberg_args, context);
+    }
+
     /// Collect all the exported paths from the processed parts
     /// If multiRead is supported by the keeper implementation, it is done in a single request
     /// Otherwise, multiple async requests are sent
@@ -253,9 +356,6 @@ namespace ExportPartitionUtils
         MergeTreeData & source_storage,
         const String & replica_name)
     {
-        auto context = Context::createCopy(context_in);
-        context->setSetting("write_full_path_in_iceberg_metadata", manifest.write_full_path_in_iceberg_metadata);
-
         /// Failpoint used by integration tests to force persistent commit failure and exercise
         /// the commit-attempts budget / FAILED state transition.
         fiu_do_on(FailPoints::export_partition_commit_always_throw,
@@ -302,17 +402,15 @@ namespace ExportPartitionUtils
             throw Exception(ErrorCodes::CORRUPTED_DATA, "ExportPartition: Reached the commit phase, but exported paths size is less than the number of parts, will not commit export. This might be a bug");
         }
 
-        IStorage::IcebergCommitExportPartitionArguments iceberg_args;
-
-        if (!manifest.iceberg_metadata_json.empty())
-        {
-            iceberg_args.metadata_json_string = manifest.iceberg_metadata_json;
-            if (source_storage.getInMemoryMetadataPtr()->hasPartitionKey())
-                iceberg_args.partition_source_block =
-                    getPartitionSourceBlockForIcebergCommit(source_storage, manifest.partition_id);
-        }
-
-        destination_storage->commitExportPartitionTransaction(manifest.transaction_id, manifest.partition_id, exported_paths, iceberg_args, context);
+        commitExportedPaths(
+            manifest.transaction_id,
+            manifest.partition_id,
+            manifest.iceberg_metadata_json,
+            manifest.write_full_path_in_iceberg_metadata,
+            exported_paths,
+            destination_storage,
+            source_storage,
+            context_in);
 
         /// Failpoint to simulate a crash after the Iceberg commit succeeds but before
         /// ZooKeeper is updated to COMPLETED. Used by idempotency integration tests.

src/Storages/MergeTree/ExportPartitionUtils.h

@@ -27,7 +27,36 @@ namespace ExportPartitionUtils
 
     std::vector<std::string> getExportedPaths(const LoggerPtr & log, const zkutil::ZooKeeperPtr & zk, const std::string & export_path);
 
-    ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ExportReplicatedMergeTreePartitionManifest & manifest);
+    /// Build a query context carrying the export task's persisted settings. Templated on the
+    /// descriptor type so it serves both the replicated manifest (backed by ZooKeeper) and the
+    /// plain `MergeTreePartitionExportTask` (backed by disk); both expose the same setting fields.
+    template <typename ManifestT>
+    ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ManifestT & manifest);
+
+    /// Validates that `dest_storage` is a legal export target for a source table with the given
+    /// metadata: it must support imports, be positionally castable, and (for non-data-lake targets)
+    /// share the same partition key. For Iceberg destinations it additionally verifies partition-spec
+    /// compatibility and returns the serialized destination `metadata.json`; returns an empty string
+    /// for non-data-lake destinations. Throws on any incompatibility.
+    std::string extractDestinationIcebergMetadataJson(
+        const StorageMetadataPtr & source_metadata,
+        const StorageID & source_storage_id,
+        const StoragePtr & dest_storage,
+        const ContextPtr & context);
+
+    /// ZooKeeper-free commit core shared by the replicated and plain partition-export paths:
+    /// assembles the Iceberg commit arguments (deriving the partition source block from the source
+    /// table when the destination is a data lake with a partition key) and invokes
+    /// `commitExportPartitionTransaction` on the destination storage.
+    void commitExportedPaths(
+        const String & transaction_id,
+        const String & partition_id,
+        const String & iceberg_metadata_json,
+        bool write_full_path_in_iceberg_metadata,
+        const std::vector<std::string> & exported_paths,
+        const StoragePtr & destination_storage,
+        MergeTreeData & source_storage,
+        const ContextPtr & context);
 
     /// Returns the representative source partition-key columns (the first active local part's
     /// minmax block) for the given partition_id. The destination recomputes the Iceberg partition