Merge pull request #1320 from Altinity/accept_table_function_as_destination_for_part_export

Accept table function as destination for part export

Author: MyroTk

docs/en/engines/table-engines/mergetree-family/part_export.md

@@ -21,6 +21,16 @@ SETTINGS allow_experimental_export_merge_tree_part = 1
          [, setting_name = value, ...]
 ```
 
+## Syntax with table function
+
+```sql
+ALTER TABLE [database.]table_name
+EXPORT PART 'part_name'
+TO TABLE FUNCTION s3(s3_conn, filename='table_function', partition_strategy...)
+SETTINGS allow_experimental_export_merge_tree_part = 1
+         [, setting_name = value, ...]
+```
+
 ### Parameters
 
 - **`table_name`**: The source MergeTree table containing the part to export
@@ -34,6 +44,8 @@ Source and destination tables must be 100% compatible:
 1. **Identical schemas** - same columns, types, and order
 2. **Matching partition keys** - partition expressions must be identical
 
+In case a table function is used as the destination, the schema can be omitted and it will be inferred from the source table.
+
 ## Settings
 
 ### `allow_experimental_export_merge_tree_part` (Required)
@@ -95,6 +107,20 @@ ALTER TABLE mt_table EXPORT PART '2021_2_2_0' TO TABLE s3_table
 SETTINGS allow_experimental_export_merge_tree_part = 1;
 ```
 
+### Table function export
+
+```sql
+-- Create source and destination tables
+CREATE TABLE mt_table (id UInt64, year UInt16)
+ENGINE = MergeTree() PARTITION BY year ORDER BY tuple();
+
+-- Insert and export
+INSERT INTO mt_table VALUES (1, 2020), (2, 2020), (3, 2021);
+
+ALTER TABLE mt_table EXPORT PART '2020_1_1_0' TO TABLE FUNCTION s3(s3_conn, filename='table_function', format=Parquet, partition_strategy='hive') PARTITION BY year
+SETTINGS allow_experimental_export_merge_tree_part = 1;
+```
+
 ## Monitoring
 
 ### Active Exports

src/Interpreters/InterpreterAlterQuery.cpp

@@ -576,7 +576,9 @@ AccessRightsElements InterpreterAlterQuery::getRequiredAccessForCommand(const AS
         case ASTAlterCommand::EXPORT_PART:
         {
             required_access.emplace_back(AccessType::ALTER_EXPORT_PART, database, table);
-            required_access.emplace_back(AccessType::INSERT, command.to_database, command.to_table);
+            /// For table functions, access control is handled by the table function itself
+            if (!command.to_table_function)
+                required_access.emplace_back(AccessType::INSERT, command.to_database, command.to_table);
             break;
         }
         case ASTAlterCommand::EXPORT_PARTITION:

src/Parsers/ASTAlterQuery.cpp

@@ -65,6 +65,10 @@ ASTPtr ASTAlterCommand::clone() const
         res->sql_security = res->children.emplace_back(sql_security->clone()).get();
     if (rename_to)
         res->rename_to = res->children.emplace_back(rename_to->clone()).get();
+    if (to_table_function)
+        res->to_table_function = res->children.emplace_back(to_table_function->clone()).get();
+    if (partition_by_expr)
+        res->partition_by_expr = res->children.emplace_back(partition_by_expr->clone()).get();
 
     return res;
 }
@@ -377,11 +381,23 @@ void ASTAlterCommand::formatImpl(WriteBuffer & ostr, const FormatSettings & sett
         {
             case DataDestinationType::TABLE:
                 ostr << "TABLE ";
-                if (!to_database.empty())
+                if (to_table_function)
+                {
+                    ostr << "FUNCTION ";
+                    to_table_function->format(ostr, settings, state, frame);
+                    if (partition_by_expr)
+                    {
+                        ostr << " PARTITION BY ";
+                        partition_by_expr->format(ostr, settings, state, frame);
+                    }
+                }
+                else
                 {
-                    ostr << backQuoteIfNeed(to_database) << ".";
+                    if (!to_database.empty())
+                        ostr << backQuoteIfNeed(to_database) << ".";
+
+                    ostr << backQuoteIfNeed(to_table);
                 }
-                ostr << backQuoteIfNeed(to_table);
                 return;
             default:
                 break;
@@ -603,6 +619,8 @@ void ASTAlterCommand::forEachPointerToChild(std::function<void(IAST **, boost::i
     f(&select, nullptr);
     f(&sql_security, nullptr);
     f(&rename_to, nullptr);
+    f(&to_table_function, nullptr);
+    f(&partition_by_expr, nullptr);
 }
 
 

src/Parsers/ASTAlterQuery.h

@@ -221,6 +221,9 @@ class ASTAlterCommand : public IAST
     /// MOVE PARTITION partition TO TABLE db.table
     String to_database;
     String to_table;
+    /// EXPORT PART/PARTITION to TABLE FUNCTION (e.g., s3())
+    IAST * to_table_function = nullptr;
+    IAST * partition_by_expr = nullptr;
 
     String snapshot_name;
     IAST * snapshot_desc;

src/Parsers/ParserAlterQuery.cpp

@@ -95,6 +95,7 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
     ParserKeyword s_unfreeze(Keyword::UNFREEZE);
     ParserKeyword s_unlock_snapshot(Keyword::UNLOCK_SNAPSHOT);
     ParserKeyword s_partition(Keyword::PARTITION);
+    ParserKeyword s_partition_by(Keyword::PARTITION_BY);
 
     ParserKeyword s_first(Keyword::FIRST);
     ParserKeyword s_after(Keyword::AFTER);
@@ -109,6 +110,7 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
     ParserKeyword s_to_volume(Keyword::TO_VOLUME);
     ParserKeyword s_to_table(Keyword::TO_TABLE);
     ParserKeyword s_to_shard(Keyword::TO_SHARD);
+    ParserKeyword s_function(Keyword::FUNCTION);
 
     ParserKeyword s_delete(Keyword::DELETE);
     ParserKeyword s_update(Keyword::UPDATE);
@@ -179,6 +181,8 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
     ASTPtr command_rename_to;
     ASTPtr command_sql_security;
     ASTPtr command_snapshot_desc;
+    ASTPtr export_table_function;
+    ASTPtr export_table_function_partition_by_expr;
 
     if (with_round_bracket)
     {
@@ -556,9 +560,27 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
                     return false;
                 }
 
-                if (!parseDatabaseAndTableName(pos, expected, command->to_database, command->to_table))
-                    return false;
-                command->move_destination_type = DataDestinationType::TABLE;
+                if (s_function.ignore(pos, expected))
+                {
+                    ParserFunction table_function_parser(/*allow_function_parameters=*/true, /*is_table_function=*/true);
+
+                    if (!table_function_parser.parse(pos, export_table_function, expected))
+                        return false;
+
+                    if (s_partition_by.ignore(pos, expected))
+                        if (!parser_exp_elem.parse(pos, export_table_function_partition_by_expr, expected))
+                            return false;
+
+                    command->to_table_function = export_table_function.get();
+                    command->partition_by_expr = export_table_function_partition_by_expr.get();
+                    command->move_destination_type = DataDestinationType::TABLE;
+                }
+                else
+                {
+                    if (!parseDatabaseAndTableName(pos, expected, command->to_database, command->to_table))
+                        return false;
+                    command->move_destination_type = DataDestinationType::TABLE;
+                }
             }
             else if (s_export_partition.ignore(pos, expected))
             {
@@ -1121,6 +1143,10 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
         command->rename_to = command->children.emplace_back(std::move(command_rename_to)).get();
     if (command_snapshot_desc)
         command->snapshot_desc = command->children.emplace_back(std::move(command_snapshot_desc)).get();
+    if (export_table_function)
+        command->to_table_function = command->children.emplace_back(std::move(export_table_function)).get();
+    if (export_table_function_partition_by_expr)
+        command->partition_by_expr = command->children.emplace_back(std::move(export_table_function_partition_by_expr)).get();
 
     return true;
 }

src/Storages/MergeTree/ExportPartTask.cpp

@@ -111,19 +111,12 @@ bool ExportPartTask::executeStep()
         block_with_partition_values = manifest.data_part->minmax_idx->getBlock(storage);
     }
 
-    auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest.destination_storage_id, local_context);
-    if (!destination_storage)
-    {
-        std::lock_guard inner_lock(storage.export_manifests_mutex);
-
-        const auto destination_storage_id_name = manifest.destination_storage_id.getNameForLogs();
-        storage.export_manifests.erase(manifest);
-        throw Exception(ErrorCodes::UNKNOWN_TABLE, "Failed to reconstruct destination storage: {}", destination_storage_id_name);
-    }
+    const auto & destination_storage = manifest.destination_storage_ptr;
+    const auto destination_storage_id = destination_storage->getStorageID();
 
     auto exports_list_entry = storage.getContext()->getExportsList().insert(
         getStorageID(),
-        manifest.destination_storage_id,
+        destination_storage_id,
         manifest.data_part->getBytesOnDisk(),
         manifest.data_part->name,
         std::vector<std::string>{},

src/Storages/MergeTree/MergeTreeData.cpp

@@ -4,6 +4,7 @@
 
 #include <Storages/MergeTree/ExportList.h>
 #include <Access/AccessControl.h>
+#include <TableFunctions/TableFunctionFactory.h>
 #include <AggregateFunctions/AggregateFunctionCount.h>
 #include <Analyzer/QueryTreeBuilder.h>
 #include <Analyzer/Utils.h>
@@ -6444,9 +6445,45 @@ void MergeTreeData::exportPartToTable(const PartitionCommand & command, ContextP
 
     const auto part_name = command.partition->as<ASTLiteral &>().value.safeGet<String>();
 
-    const auto database_name = query_context->resolveDatabase(command.to_database);
+    if (!command.to_table_function)
+    {
+        const auto database_name = query_context->resolveDatabase(command.to_database);
+        exportPartToTable(part_name, StorageID{database_name, command.to_table}, generateSnowflakeIDString(), query_context);
+
+        return;
+    }
+
+    auto table_function_ast = command.to_table_function;
+    auto table_function_ptr = TableFunctionFactory::instance().get(command.to_table_function, query_context);
+
+    if (table_function_ptr->needStructureHint())
+    {
+        const auto source_metadata_ptr = getInMemoryMetadataPtr();
+
+        /// Grab only the readable columns from the source metadata to skip ephemeral columns
+        const auto readable_columns = ColumnsDescription(source_metadata_ptr->getColumns().getReadable());
+        table_function_ptr->setStructureHint(readable_columns);
+    }
+
+    if (command.partition_by_expr)
+    {
+        table_function_ptr->setPartitionBy(command.partition_by_expr);
+    }
+
+    auto dest_storage = table_function_ptr->execute(
+        table_function_ast,
+        query_context,
+        table_function_ptr->getName(),
+        /* cached_columns */ {},
+        /* use_global_context */ false,
+        /* is_insert_query */ true);
+
+    if (!dest_storage)
+    {
+        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Failed to reconstruct destination storage");
+    }
 
-    exportPartToTable(part_name, StorageID{database_name, command.to_table}, generateSnowflakeIDString(), query_context);
+    exportPartToTable(part_name, dest_storage, generateSnowflakeIDString(), query_context);
 }
 
 void MergeTreeData::exportPartToTable(
@@ -6464,6 +6501,17 @@ void MergeTreeData::exportPartToTable(
         throw Exception(ErrorCodes::BAD_ARGUMENTS, "Exporting to the same table is not allowed");
     }
 
+    exportPartToTable(part_name, dest_storage, transaction_id, query_context, allow_outdated_parts, completion_callback);
+}
+
+void MergeTreeData::exportPartToTable(
+    const std::string & part_name,
+    const StoragePtr & dest_storage,
+    const String & transaction_id,
+    ContextPtr query_context,
+    bool allow_outdated_parts,
+    std::function<void(MergeTreePartExportManifest::CompletionCallbackResult)> completion_callback)
+{
     if (!dest_storage->supportsImport())
         throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Destination storage {} does not support MergeTree parts or uses unsupported partitioning", dest_storage->getName());
 
@@ -6533,7 +6581,7 @@ void MergeTreeData::exportPartToTable(
     {
         const auto format_settings = getFormatSettings(query_context);
         MergeTreePartExportManifest manifest(
-            dest_storage->getStorageID(),
+            dest_storage,
             part,
             transaction_id,
             query_context->getCurrentQueryId(),
@@ -6546,8 +6594,7 @@ void MergeTreeData::exportPartToTable(
 
         if (!export_manifests.emplace(std::move(manifest)).second)
         {
-            throw Exception(ErrorCodes::ABORTED, "Data part '{}' is already being exported to table '{}'",
-                            part_name, dest_storage->getStorageID().getFullTableName());
+            throw Exception(ErrorCodes::ABORTED, "Data part '{}' is already being exported", part_name);
         }
     }
 
@@ -9038,8 +9085,9 @@ std::vector<MergeTreeExportStatus> MergeTreeData::getExportsStatus() const
 
         status.source_database = source_database;
         status.source_table = source_table;
-        status.destination_database = manifest.destination_storage_id.database_name;
-        status.destination_table = manifest.destination_storage_id.table_name;
+        const auto destination_storage_id = manifest.destination_storage_ptr->getStorageID();
+        status.destination_database = destination_storage_id.database_name;
+        status.destination_table = destination_storage_id.table_name;
         status.create_time = manifest.create_time;
         status.part_name = manifest.data_part->name;
 

src/Storages/MergeTree/MergeTreeData.h

@@ -1025,6 +1025,14 @@ class MergeTreeData : public IStorage, public WithMutableContext
 
     void exportPartToTable(const PartitionCommand & command, ContextPtr query_context);
 
+    void exportPartToTable(
+        const std::string & part_name,
+        const StoragePtr & destination_storage,
+        const String & transaction_id,
+        ContextPtr query_context,
+        bool allow_outdated_parts = false,
+        std::function<void(MergeTreePartExportManifest::CompletionCallbackResult)> completion_callback = {});
+
     void exportPartToTable(
         const std::string & part_name,
         const StorageID & destination_storage_id,

src/Storages/MergeTree/MergeTreePartExportManifest.h

@@ -6,6 +6,7 @@
 #include <QueryPipeline/QueryPipeline.h>
 #include <optional>
 #include <Core/Settings.h>
+#include <Storages/IStorage.h>
 
 namespace DB
 {
@@ -43,15 +44,15 @@ struct MergeTreePartExportManifest
     };
 
     MergeTreePartExportManifest(
-        const StorageID & destination_storage_id_,
+        const StoragePtr destination_storage_ptr_,
         const DataPartPtr & data_part_,
         const String & transaction_id_,
         const String & query_id_,
         FileAlreadyExistsPolicy file_already_exists_policy_,
         const Settings & settings_,
         const StorageMetadataPtr & metadata_snapshot_,
         std::function<void(CompletionCallbackResult)> completion_callback_ = {})
-        : destination_storage_id(destination_storage_id_),
+        : destination_storage_ptr(destination_storage_ptr_),
           data_part(data_part_),
           transaction_id(transaction_id_),
           query_id(query_id_),
@@ -61,7 +62,7 @@ struct MergeTreePartExportManifest
           completion_callback(completion_callback_),
           create_time(time(nullptr)) {}
 
-    StorageID destination_storage_id;
+    StoragePtr destination_storage_ptr;
     DataPartPtr data_part;
     /// Used for killing the export.
     String transaction_id;
@@ -81,20 +82,12 @@ struct MergeTreePartExportManifest
 
     bool operator<(const MergeTreePartExportManifest & rhs) const 
     {
-        // Lexicographic comparison: first compare destination storage, then part name
-        auto lhs_storage = destination_storage_id.getQualifiedName();
-        auto rhs_storage = rhs.destination_storage_id.getQualifiedName();
-        
-        if (lhs_storage != rhs_storage)
-            return lhs_storage < rhs_storage;
-            
         return data_part->name < rhs.data_part->name;
     }
 
     bool operator==(const MergeTreePartExportManifest & rhs) const 
     {
-        return destination_storage_id.getQualifiedName() == rhs.destination_storage_id.getQualifiedName()
-            && data_part->name == rhs.data_part->name;
+        return data_part->name == rhs.data_part->name;
     }
 };