Merge pull request #1083 from Altinity/antalya_25_8_always_support_import_on_object_storage

Enmk · web-flow · commit 410f2f383a3e · 2025-10-16T18:50:29.000+02:00
diff --git a/src/Storages/MergeTree/IMergeTreeDataPart.cpp b/src/Storages/MergeTree/IMergeTreeDataPart.cpp
@@ -279,6 +279,42 @@ String IMergeTreeDataPart::MinMaxIndex::getFileColumnName(const String & column_
     return stream_name;
 }
 
+Block IMergeTreeDataPart::MinMaxIndex::getBlock(const MergeTreeData & data) const
+{
+    if (!initialized)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Attempt to get block from uninitialized MinMax index.");
+
+    Block block;
+
+    const auto metadata_snapshot = data.getInMemoryMetadataPtr();
+    const auto & partition_key = metadata_snapshot->getPartitionKey();
+
+    const auto minmax_column_names = data.getMinMaxColumnsNames(partition_key);
+    const auto minmax_column_types = data.getMinMaxColumnsTypes(partition_key);
+    const auto minmax_idx_size = minmax_column_types.size();
+
+    for (size_t i = 0; i < minmax_idx_size; ++i)
+    {
+        const auto & data_type = minmax_column_types[i];
+        const auto & column_name = minmax_column_names[i];
+
+        const auto column = data_type->createColumn();
+
+        auto range = hyperrectangle.at(i);
+        range.shrinkToIncludedIfPossible();
+
+        const auto & min_val = range.left;
+        const auto & max_val = range.right;
+
+        column->insert(min_val);
+        column->insert(max_val);
+
+        block.insert(ColumnWithTypeAndName(column->getPtr(), data_type, column_name));
+    }
+
+    return block;
+}
+
 void IMergeTreeDataPart::incrementStateMetric(MergeTreeDataPartState state_) const
 {
     switch (state_)
diff --git a/src/Storages/MergeTree/IMergeTreeDataPart.h b/src/Storages/MergeTree/IMergeTreeDataPart.h
@@ -362,6 +362,8 @@ class IMergeTreeDataPart : public std::enable_shared_from_this<IMergeTreeDataPar
         static String getFileColumnName(const String & column_name, const MergeTreeSettingsPtr & storage_settings_);
         /// For Load
         static String getFileColumnName(const String & column_name, const Checksums & checksums_);
+
+        Block getBlock(const MergeTreeData & data) const;
     };
 
     using MinMaxIndexPtr = std::shared_ptr<MinMaxIndex>;
diff --git a/src/Storages/MergeTree/MergeTreeData.cpp b/src/Storages/MergeTree/MergeTreeData.cpp
@@ -6268,16 +6268,13 @@ void MergeTreeData::exportPartToTableImpl(
 
     MergeTreeSequentialSourceType read_type = MergeTreeSequentialSourceType::Export;
 
-    NamesAndTypesList partition_columns;
+    Block block_with_partition_values;
     if (metadata_snapshot->hasPartitionKey())
     {
-        const auto & partition_key = metadata_snapshot->getPartitionKey();
-        if (!partition_key.column_names.empty())
-            partition_columns = partition_key.expression->getRequiredColumnsWithTypes();
+        /// todo arthur do I need to init minmax_idx?
+        block_with_partition_values = manifest.data_part->minmax_idx->getBlock(*this);
     }
 
-    auto block_with_partition_values = manifest.data_part->partition.getBlockWithPartitionValues(partition_columns);
-
     auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest.destination_storage_id, getContext());
     if (!destination_storage)
     {
diff --git a/src/Storages/ObjectStorage/ObjectStorageFilePathGenerator.h b/src/Storages/ObjectStorage/ObjectStorageFilePathGenerator.h
@@ -4,6 +4,7 @@
 #include <Storages/PartitionedSink.h>
 #include <Poco/String.h>
 #include <Functions/generateSnowflakeID.h>
+#include <boost/algorithm/string/replace.hpp>
 
 namespace DB
 {
@@ -19,12 +20,15 @@ namespace DB
 
     struct ObjectStorageWildcardFilePathGenerator : ObjectStorageFilePathGenerator
     {
+        static constexpr const char * FILE_WILDCARD = "{_file}";
         explicit ObjectStorageWildcardFilePathGenerator(const std::string & raw_path_) : raw_path(raw_path_) {}
 
         using ObjectStorageFilePathGenerator::getPathForWrite;  // Bring base class overloads into scope
-        std::string getPathForWrite(const std::string & partition_id, const std::string & /* file_name_override */) const override
+        std::string getPathForWrite(const std::string & partition_id, const std::string & file_name_override) const override
         {
-            return PartitionedSink::replaceWildcards(raw_path, partition_id);
+            const auto partition_replaced_path = PartitionedSink::replaceWildcards(raw_path, partition_id);
+            const auto final_path = boost::replace_all_copy(partition_replaced_path, FILE_WILDCARD, file_name_override);
+            return final_path;
         }
 
         std::string getPathForRead() const override
diff --git a/src/Storages/ObjectStorage/StorageObjectStorage.cpp b/src/Storages/ObjectStorage/StorageObjectStorage.cpp
@@ -468,7 +468,13 @@ bool StorageObjectStorage::optimize(
 
 bool StorageObjectStorage::supportsImport() const
 {
-    return configuration->partition_strategy != nullptr && configuration->partition_strategy_type == PartitionStrategyFactory::StrategyType::HIVE;
+    if (!configuration->partition_strategy)
+        return false;
+
+    if (configuration->partition_strategy_type == PartitionStrategyFactory::StrategyType::WILDCARD)
+        return configuration->getRawPath().hasExportFilenameWildcard();
+
+    return configuration->partition_strategy_type == PartitionStrategyFactory::StrategyType::HIVE;
 }
 
 SinkToStoragePtr StorageObjectStorage::import(
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageConfiguration.cpp b/src/Storages/ObjectStorage/StorageObjectStorageConfiguration.cpp
@@ -159,6 +159,12 @@ bool StorageObjectStorageConfiguration::Path::hasPartitionWildcard() const
     return path.find(PARTITION_ID_WILDCARD) != String::npos;
 }
 
+bool StorageObjectStorageConfiguration::Path::hasExportFilenameWildcard() const
+{
+    return path.find(ObjectStorageWildcardFilePathGenerator::FILE_WILDCARD) != String::npos;
+}
+
+
 bool StorageObjectStorageConfiguration::Path::hasGlobsIgnorePartitionWildcard() const
 {
     if (!hasPartitionWildcard())
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageConfiguration.h b/src/Storages/ObjectStorage/StorageObjectStorageConfiguration.h
@@ -62,6 +62,7 @@ class StorageObjectStorageConfiguration
         std::string path;
 
         bool hasPartitionWildcard() const;
+        bool hasExportFilenameWildcard() const;
         bool hasGlobsIgnorePartitionWildcard() const;
         bool hasGlobs() const;
         std::string cutGlobs(bool supports_partial_prefix) const;
diff --git a/tests/queries/0_stateless/03572_export_merge_tree_part_to_object_storage.reference b/tests/queries/0_stateless/03572_export_merge_tree_part_to_object_storage.reference
@@ -14,3 +14,20 @@
 2	2020
 3	2020
 4	2021
+---- Export 2020_1_1_0 and 2021_2_2_0 to wildcard table
+---- Both data parts should appear
+1	2020
+2	2020
+3	2020
+4	2021
+---- Export the same part again, it should be idempotent
+1	2020
+2	2020
+3	2020
+4	2021
+---- Export 2020_1_1_0 and 2021_2_2_0 to wildcard table with partition expression with function
+---- Both data parts should appear
+1	2020
+2	2020
+3	2020
+4	2021
diff --git a/tests/queries/0_stateless/03572_export_merge_tree_part_to_object_storage.sh b/tests/queries/0_stateless/03572_export_merge_tree_part_to_object_storage.sh
@@ -1,19 +1,21 @@
 #!/usr/bin/env bash
-# Tags: no-fasttest
 
 CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 # shellcheck source=../shell_config.sh
 . "$CURDIR"/../shell_config.sh
 
 mt_table="mt_table_${RANDOM}"
+mt_table_partition_expression_with_function="mt_table_partition_expression_with_function_${RANDOM}"
 s3_table="s3_table_${RANDOM}"
+s3_table_wildcard="s3_table_wildcard_${RANDOM}"
+s3_table_wildcard_partition_expression_with_function="s3_table_wildcard_partition_expression_with_function_${RANDOM}"
 mt_table_roundtrip="mt_table_roundtrip_${RANDOM}"
 
 query() {
     $CLICKHOUSE_CLIENT --query "$1"
 }
 
-query "DROP TABLE IF EXISTS $mt_table, $s3_table, $mt_table_roundtrip"
+query "DROP TABLE IF EXISTS $mt_table, $s3_table, $mt_table_roundtrip, $s3_table_wildcard, $s3_table_wildcard_partition_expression_with_function, $mt_table_partition_expression_with_function"
 
 query "CREATE TABLE $mt_table (id UInt64, year UInt16) ENGINE = MergeTree() PARTITION BY year ORDER BY tuple()"
 query "CREATE TABLE $s3_table (id UInt64, year UInt16) ENGINE = S3(s3_conn, filename='$s3_table', format=Parquet, partition_strategy='hive') PARTITION BY year"
@@ -36,4 +38,31 @@ query "CREATE TABLE $mt_table_roundtrip ENGINE = MergeTree() PARTITION BY year O
 echo "---- Data in roundtrip MergeTree table (should match s3_table)"
 query "SELECT * FROM $s3_table ORDER BY id"
 
-query "DROP TABLE IF EXISTS $mt_table, $s3_table, $mt_table_roundtrip"
+query "CREATE TABLE $s3_table_wildcard (id UInt64, year UInt16) ENGINE = S3(s3_conn, filename='$s3_table_wildcard/{_partition_id}/{_file}.parquet', format=Parquet, partition_strategy='wildcard') PARTITION BY year"
+
+echo "---- Export 2020_1_1_0 and 2021_2_2_0 to wildcard table"
+query "ALTER TABLE $mt_table EXPORT PART '2020_1_1_0' TO TABLE $s3_table_wildcard SETTINGS allow_experimental_export_merge_tree_part = 1"
+query "ALTER TABLE $mt_table EXPORT PART '2021_2_2_0' TO TABLE $s3_table_wildcard SETTINGS allow_experimental_export_merge_tree_part = 1"
+
+echo "---- Both data parts should appear"
+query "SELECT * FROM s3(s3_conn, filename='$s3_table_wildcard/**.parquet') ORDER BY id"
+
+echo "---- Export the same part again, it should be idempotent"
+query "ALTER TABLE $mt_table EXPORT PART '2020_1_1_0' TO TABLE $s3_table_wildcard SETTINGS allow_experimental_export_merge_tree_part = 1"
+
+query "SELECT * FROM s3(s3_conn, filename='$s3_table_wildcard/**.parquet') ORDER BY id"
+
+query "CREATE TABLE $mt_table_partition_expression_with_function (id UInt64, year UInt16) ENGINE = MergeTree() PARTITION BY toString(year) ORDER BY tuple()"
+query "CREATE TABLE $s3_table_wildcard_partition_expression_with_function (id UInt64, year UInt16) ENGINE = S3(s3_conn, filename='$s3_table_wildcard_partition_expression_with_function/{_partition_id}/{_file}.parquet', format=Parquet, partition_strategy='wildcard') PARTITION BY toString(year)"
+
+# insert
+query "INSERT INTO $mt_table_partition_expression_with_function VALUES (1, 2020), (2, 2020), (3, 2020), (4, 2021)"
+
+echo "---- Export 2020_1_1_0 and 2021_2_2_0 to wildcard table with partition expression with function"
+query "ALTER TABLE $mt_table_partition_expression_with_function EXPORT PART 'cb217c742dc7d143b61583011996a160_1_1_0' TO TABLE $s3_table_wildcard_partition_expression_with_function SETTINGS allow_experimental_export_merge_tree_part = 1"
+query "ALTER TABLE $mt_table_partition_expression_with_function EXPORT PART '3be6d49ecf9749a383964bc6fab22d10_2_2_0' TO TABLE $s3_table_wildcard_partition_expression_with_function SETTINGS allow_experimental_export_merge_tree_part = 1"
+
+echo "---- Both data parts should appear"
+query "SELECT * FROM s3(s3_conn, filename='$s3_table_wildcard_partition_expression_with_function/**.parquet') ORDER BY id"
+
+query "DROP TABLE IF EXISTS $mt_table, $s3_table, $mt_table_roundtrip, $s3_table_wildcard, $s3_table_wildcard_partition_expression_with_function, $mt_table_partition_expression_with_function"

Original file line number	Diff line number	Diff line change
`@@ -6268,16 +6268,13 @@ void MergeTreeData::exportPartToTableImpl(`
`6268`	`6268`
`6269`	`6269`	`MergeTreeSequentialSourceType read_type = MergeTreeSequentialSourceType::Export;`
`6270`	`6270`
`6271`		`- NamesAndTypesList partition_columns;`
	`6271`	`+ Block block_with_partition_values;`
`6272`	`6272`	`if (metadata_snapshot->hasPartitionKey())`
`6273`	`6273`	`{`
`6274`		`- const auto & partition_key = metadata_snapshot->getPartitionKey();`
`6275`		`- if (!partition_key.column_names.empty())`
`6276`		`- partition_columns = partition_key.expression->getRequiredColumnsWithTypes();`
	`6274`	`+ /// todo arthur do I need to init minmax_idx?`
	`6275`	`+ block_with_partition_values = manifest.data_part->minmax_idx->getBlock(*this);`
`6277`	`6276`	`}`
`6278`	`6277`
`6279`		`- auto block_with_partition_values = manifest.data_part->partition.getBlockWithPartitionValues(partition_columns);`
`6280`		`-`
`6281`	`6278`	`auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest.destination_storage_id, getContext());`
`6282`	`6279`	`if (!destination_storage)`
`6283`	`6280`	`{`
Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@`
`4`	`4`	`#include <Storages/PartitionedSink.h>`
`5`	`5`	`#include <Poco/String.h>`
`6`	`6`	`#include <Functions/generateSnowflakeID.h>`
	`7`	`+#include <boost/algorithm/string/replace.hpp>`
`7`	`8`
`8`	`9`	`namespace DB`
`9`	`10`	`{`
`@@ -19,12 +20,15 @@ namespace DB`
`19`	`20`
`20`	`21`	`struct ObjectStorageWildcardFilePathGenerator : ObjectStorageFilePathGenerator`
`21`	`22`	`{`
	`23`	`+ static constexpr const char * FILE_WILDCARD = "{_file}";`
`22`	`24`	`explicit ObjectStorageWildcardFilePathGenerator(const std::string & raw_path_) : raw_path(raw_path_) {}`
`23`	`25`
`24`	`26`	`using ObjectStorageFilePathGenerator::getPathForWrite; // Bring base class overloads into scope`
`25`		`- std::string getPathForWrite(const std::string & partition_id, const std::string & /* file_name_override */) const override`
	`27`	`+ std::string getPathForWrite(const std::string & partition_id, const std::string & file_name_override) const override`
`26`	`28`	`{`
`27`		`- return PartitionedSink::replaceWildcards(raw_path, partition_id);`
	`29`	`+ const auto partition_replaced_path = PartitionedSink::replaceWildcards(raw_path, partition_id);`
	`30`	`+ const auto final_path = boost::replace_all_copy(partition_replaced_path, FILE_WILDCARD, file_name_override);`
	`31`	`+ return final_path;`
`28`	`32`	`}`
`29`	`33`
`30`	`34`	`std::string getPathForRead() const override`
Original file line number	Diff line number	Diff line change
`@@ -468,7 +468,13 @@ bool StorageObjectStorage::optimize(`
`468`	`468`
`469`	`469`	`bool StorageObjectStorage::supportsImport() const`
`470`	`470`	`{`
`471`		`- return configuration->partition_strategy != nullptr && configuration->partition_strategy_type == PartitionStrategyFactory::StrategyType::HIVE;`
	`471`	`+ if (!configuration->partition_strategy)`
	`472`	`+ return false;`
	`473`	`+`
	`474`	`+ if (configuration->partition_strategy_type == PartitionStrategyFactory::StrategyType::WILDCARD)`
	`475`	`+ return configuration->getRawPath().hasExportFilenameWildcard();`
	`476`	`+`
	`477`	`+ return configuration->partition_strategy_type == PartitionStrategyFactory::StrategyType::HIVE;`
`472`	`478`	`}`
`473`	`479`
`474`	`480`	`SinkToStoragePtr StorageObjectStorage::import(`
Original file line number	Diff line number	Diff line change
`@@ -159,6 +159,12 @@ bool StorageObjectStorageConfiguration::Path::hasPartitionWildcard() const`
`159`	`159`	`return path.find(PARTITION_ID_WILDCARD) != String::npos;`
`160`	`160`	`}`
`161`	`161`
	`162`	`+bool StorageObjectStorageConfiguration::Path::hasExportFilenameWildcard() const`
	`163`	`+{`
	`164`	`+ return path.find(ObjectStorageWildcardFilePathGenerator::FILE_WILDCARD) != String::npos;`
	`165`	`+}`
	`166`	`+`
	`167`	`+`
`162`	`168`	`bool StorageObjectStorageConfiguration::Path::hasGlobsIgnorePartitionWildcard() const`
`163`	`169`	`{`
`164`	`170`	`if (!hasPartitionWildcard())`