S3 hive style writes

src/Core/Settings.cpp

@@ -5856,7 +5856,7 @@ The maximum number of rows in the right table to determine whether to rerange th
     DECLARE(Bool, allow_experimental_join_right_table_sorting, false, R"(
 If it is set to true, and the conditions of `join_to_sort_minimum_perkey_rows` and `join_to_sort_maximum_table_rows` are met, rerange the right table by key to improve the performance in left or inner hash join.
 )", EXPERIMENTAL) \
-    DECLARE(Bool, use_hive_partitioning, false, R"(
+    DECLARE(Bool, use_hive_partitioning, true, R"(
 When enabled, ClickHouse will detect Hive-style partitioning in path (`/name=value/`) in file-like table engines [File](../../engines/table-engines/special/file.md#hive-style-partitioning)/[S3](../../engines/table-engines/integrations/s3.md#hive-style-partitioning)/[URL](../../engines/table-engines/special/url.md#hive-style-partitioning)/[HDFS](../../engines/table-engines/integrations/hdfs.md#hive-style-partitioning)/[AzureBlobStorage](../../engines/table-engines/integrations/azureBlobStorage.md#hive-style-partitioning) and will allow to use partition columns as virtual columns in the query. These virtual columns will have the same names as in the partitioned path, but starting with `_`.
 )", EXPERIMENTAL)\
     \

src/Functions/generateSnowflakeID.cpp

@@ -1,3 +1,4 @@
+#include <Functions/generateSnowflakeID.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <Functions/FunctionFactory.h>
 #include <Functions/FunctionsRandom.h>
@@ -147,6 +148,13 @@ struct Data
 
 }
 
+uint64_t generateSnowflakeID()
+{
+    Data data;
+    SnowflakeId snowflake_id = data.reserveRange(getMachineId(), 1);
+    return fromSnowflakeId(snowflake_id);
+}
+
 class FunctionGenerateSnowflakeID : public IFunction
 {
 public:

src/Functions/generateSnowflakeID.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <cstdint>
+
+namespace DB
+{
+
+uint64_t generateSnowflakeID();
+
+}

src/Interpreters/InterpreterInsertQuery.cpp

@@ -136,7 +136,7 @@ StoragePtr InterpreterInsertQuery::getTable(ASTInsertQuery & query)
         }
 
         return table_function_ptr->execute(query.table_function, current_context, table_function_ptr->getName(),
-                                           /* cached_columns */ {}, /* use_global_context */ false, /* is_insert_query */true);
+                                           /* cached_columns */ {}, /* use_global_context */ false, &query);
     }
 
     if (query.table_id)

src/Storages/ObjectStorage/S3/Configuration.cpp

@@ -53,6 +53,61 @@ namespace ErrorCodes
     extern const int LOGICAL_ERROR;
 }
 
+namespace
+{
+
+    ASTs::const_iterator getValueForNamedArgument(const ASTs & arguments, const std::string & argument_name, Field & value)
+    {
+        for (const auto * arg_it = arguments.begin(); arg_it != arguments.end(); ++arg_it)
+        {
+            auto argument = *arg_it;
+            const auto * type_ast_function = argument->as<ASTFunction>();
+
+            if (!type_ast_function || type_ast_function->name != "equals" || !type_ast_function->arguments || type_ast_function->arguments->children.size() != 2)
+            {
+                continue;
+            }
+
+            const auto * name = type_ast_function->arguments->children[0]->as<ASTIdentifier>();
+
+            if (name && name->name() == argument_name)
+            {
+                const auto * ast_literal = type_ast_function->arguments->children[1]->as<ASTLiteral>();
+
+                if (!ast_literal)
+                {
+                    throw Exception(
+                        ErrorCodes::BAD_ARGUMENTS,
+                        "Wrong parameter type for '{}'",
+                        name->name());
+                }
+
+                value = ast_literal->value;
+
+                return arg_it;
+            }
+        }
+
+        return arguments.end();
+    }
+
+    template <typename T>
+    std::optional<T> extractNamedArgumentAndRemoveFromList(ASTs & arguments, const std::string & argument_name)
+    {
+        Field value;
+        const auto * p = getValueForNamedArgument(arguments, argument_name, value);
+
+        if (p == arguments.end())
+        {
+            return std::nullopt;
+        }
+
+        arguments.erase(p);
+
+        return value.safeGet<T>();
+    }
+}
+
 static const std::unordered_set<std::string_view> required_configuration_keys = {
     "url",
 };
@@ -74,7 +129,9 @@ static const std::unordered_set<std::string_view> optional_configuration_keys =
     "max_single_part_upload_size",
     "max_connections",
     "expiration_window_seconds",
-    "no_sign_request"
+    "no_sign_request",
+    "partition_strategy",
+    "hive_partition_strategy_write_partition_columns_into_files"
 };
 
 String StorageS3Configuration::getDataSourceDescription() const
@@ -172,6 +229,9 @@ void StorageS3Configuration::fromNamedCollection(const NamedCollection & collect
     else
         url = S3::URI(collection.get<String>("url"), settings[Setting::allow_archive_path_syntax]);
 
+    partition_strategy = collection.getOrDefault<String>("partition_strategy", "auto");
+    hive_partition_strategy_write_partition_columns_into_files = collection.getOrDefault<bool>("hive_partition_strategy_write_partition_columns_into_files", false);
+
     auth_settings[S3AuthSetting::access_key_id] = collection.getOrDefault<String>("access_key_id", "");
     auth_settings[S3AuthSetting::secret_access_key] = collection.getOrDefault<String>("secret_access_key", "");
     auth_settings[S3AuthSetting::use_environment_credentials] = collection.getOrDefault<UInt64>("use_environment_credentials", 1);
@@ -191,8 +251,16 @@ void StorageS3Configuration::fromNamedCollection(const NamedCollection & collect
 
 void StorageS3Configuration::fromAST(ASTs & args, ContextPtr context, bool with_structure)
 {
-    size_t count = StorageURL::evalArgsAndCollectHeaders(args, headers_from_ast, context);
+    /*
+     * Calls to `extractNamedArgumentAndRemoveFromList` need to happen before count is determined:
+     * `size_t count = StorageURL::evalArgsAndCollectHeaders`
+     * This is because extractNamedArgumentAndRemoveFromList alters the list of arguments
+     * */
+    partition_strategy = extractNamedArgumentAndRemoveFromList<std::string>(args, "partition_strategy").value_or("auto");
+    hive_partition_strategy_write_partition_columns_into_files = extractNamedArgumentAndRemoveFromList<bool>(args, "hive_partition_strategy_write_partition_columns_into_files").value_or(false);
 
+    size_t count = StorageURL::evalArgsAndCollectHeaders(args, headers_from_ast, context);
+
     if (count == 0 || count > getMaxNumberOfArguments(with_structure))
         throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
             "Storage S3 requires 1 to {} arguments. All supported signatures:\n{}",

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -42,6 +42,7 @@ namespace ErrorCodes
     extern const int DATABASE_ACCESS_DENIED;
     extern const int NOT_IMPLEMENTED;
     extern const int LOGICAL_ERROR;
+    extern const int BAD_ARGUMENTS;
 }
 
 namespace StorageObjectStorageSetting
@@ -139,6 +140,21 @@ StorageObjectStorage::StorageObjectStorage(
 
     setVirtuals(VirtualColumnUtils::getVirtualsForFileLikeStorage(metadata.columns, context, sample_path, format_settings));
     setInMemoryMetadata(metadata);
+
+    // perhaps it is worth adding some extra safeguards for cases like
+    // create table s3_table engine=s3('{_partition_id}'); -- partition id wildcard set, but no partition expression
+    // create table s3_table engine=s3(partition_strategy='hive'); -- partition strategy set, but no partition expression
+    if (partition_by)
+    {
+        partition_strategy = PartitionStrategyFactory::get(
+                partition_by,
+                metadata.getSampleBlock(),
+                context,
+                configuration->format,
+                configuration->withPartitionWildcard(),
+                configuration->partition_strategy,
+                configuration->hive_partition_strategy_write_partition_columns_into_files);
+    }
 }
 
 String StorageObjectStorage::getName() const
@@ -357,7 +373,7 @@ void StorageObjectStorage::read(
 }
 
 SinkToStoragePtr StorageObjectStorage::write(
-    const ASTPtr & query,
+    const ASTPtr &,
     const StorageMetadataPtr & metadata_snapshot,
     ContextPtr local_context,
     bool /* async_insert */)
@@ -380,22 +396,10 @@ SinkToStoragePtr StorageObjectStorage::write(
                         configuration->getPath());
     }
 
-    if (configuration->withPartitionWildcard())
+    if (partition_strategy)
     {
-        ASTPtr partition_by_ast = nullptr;
-        if (auto insert_query = std::dynamic_pointer_cast<ASTInsertQuery>(query))
-        {
-            if (insert_query->partition_by)
-                partition_by_ast = insert_query->partition_by;
-            else
-                partition_by_ast = partition_by;
-        }
-
-        if (partition_by_ast)
-        {
-            return std::make_shared<PartitionedStorageObjectStorageSink>(
-                object_storage, configuration, format_settings, sample_block, local_context, partition_by_ast);
-        }
+        return std::make_shared<PartitionedStorageObjectStorageSink>(
+            partition_strategy, object_storage, configuration, format_settings, sample_block, local_context);
     }
 
     auto paths = configuration->getPaths();
@@ -567,6 +571,10 @@ void StorageObjectStorage::Configuration::initialize(
     else
         configuration.fromAST(engine_args, local_context, with_table_structure);
 
+    if (configuration.isNamespaceWithGlobs())
+        throw Exception(ErrorCodes::BAD_ARGUMENTS,
+                        "Expression can not have wildcards inside {} name", configuration.getNamespaceType());
+
     if (configuration.format == "auto")
     {
         if (configuration.isDataLakeConfiguration())

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -18,6 +18,7 @@ namespace DB
 class ReadBufferIterator;
 class SchemaCache;
 class NamedCollection;
+struct PartitionStrategy;
 
 namespace ErrorCodes
 {
@@ -153,6 +154,7 @@ class StorageObjectStorage : public IStorage
     const std::optional<FormatSettings> format_settings;
     const ASTPtr partition_by;
     const bool distributed_processing;
+    std::shared_ptr<PartitionStrategy> partition_strategy;
 
     LoggerPtr log;
 };
@@ -244,6 +246,8 @@ class StorageObjectStorage::Configuration
     String format = "auto";
     String compression_method = "auto";
     String structure = "auto";
+    std::string partition_strategy = "auto";
+    bool hive_partition_strategy_write_partition_columns_into_files = false;
 
     virtual void update(ObjectStoragePtr object_storage, ContextPtr local_context);
     void updateIfRequired(ObjectStoragePtr object_storage, ContextPtr local_context);

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -92,8 +92,8 @@ StorageObjectStorageCluster::StorageObjectStorageCluster(
     if (sample_path.empty() && context_->getSettingsRef()[Setting::use_hive_partitioning])
         sample_path = getPathSample(metadata, context_);
 
-    setInMemoryMetadata(metadata);
     setVirtuals(VirtualColumnUtils::getVirtualsForFileLikeStorage(metadata.columns, context_, sample_path));
+    setInMemoryMetadata(metadata);
 
     pure_storage = std::make_shared<StorageObjectStorage>(
         configuration,

src/Storages/ObjectStorage/StorageObjectStorageSink.cpp

@@ -20,6 +20,34 @@ namespace ErrorCodes
     extern const int BAD_ARGUMENTS;
 }
 
+namespace
+{
+    void validateKey(const String & str)
+    {
+        /// See:
+        /// - https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
+        /// - https://cloud.ibm.com/apidocs/cos/cos-compatibility#putobject
+
+        if (str.empty() || str.size() > 1024)
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Incorrect key length (not empty, max 1023 characters), got: {}", str.size());
+
+        if (!UTF8::isValidUTF8(reinterpret_cast<const UInt8 *>(str.data()), str.size()))
+            throw Exception(ErrorCodes::CANNOT_PARSE_TEXT, "Incorrect non-UTF8 sequence in key");
+
+        PartitionedSink::validatePartitionKey(str, true);
+    }
+
+    void validateNamespace(const String & str, PartitionedStorageObjectStorageSink::ConfigurationPtr configuration)
+    {
+        configuration->validateNamespace(str);
+
+        if (!UTF8::isValidUTF8(reinterpret_cast<const UInt8 *>(str.data()), str.size()))
+            throw Exception(ErrorCodes::CANNOT_PARSE_TEXT, "Incorrect non-UTF8 sequence in bucket name");
+
+        PartitionedSink::validatePartitionKey(str, false);
+    }
+}
+
 StorageObjectStorageSink::StorageObjectStorageSink(
     ObjectStoragePtr object_storage,
     ConfigurationPtr configuration,
@@ -97,13 +125,13 @@ void StorageObjectStorageSink::cancelBuffers()
 }
 
 PartitionedStorageObjectStorageSink::PartitionedStorageObjectStorageSink(
+    std::shared_ptr<PartitionStrategy> partition_strategy_,
     ObjectStoragePtr object_storage_,
     ConfigurationPtr configuration_,
     std::optional<FormatSettings> format_settings_,
     const Block & sample_block_,
-    ContextPtr context_,
-    const ASTPtr & partition_by)
-    : PartitionedSink(partition_by, context_, sample_block_)
+    ContextPtr context_)
+    : PartitionedSink(partition_strategy_, context_, sample_block_)
     , object_storage(object_storage_)
     , configuration(configuration_)
     , query_settings(configuration_->getQuerySettings(context_))
@@ -121,51 +149,25 @@ StorageObjectStorageSink::~StorageObjectStorageSink()
 
 SinkPtr PartitionedStorageObjectStorageSink::createSinkForPartition(const String & partition_id)
 {
-    auto partition_bucket = replaceWildcards(configuration->getNamespace(), partition_id);
-    validateNamespace(partition_bucket);
+    auto file_path = getPartitionStrategy()->getPath(configuration->getPath(), partition_id);
 
-    auto partition_key = replaceWildcards(configuration->getPath(), partition_id);
-    validateKey(partition_key);
+    validateNamespace(configuration->getNamespace(), configuration);
+    validateKey(file_path);
 
     if (auto new_key = checkAndGetNewFileOnInsertIfNeeded(
-            *object_storage, *configuration, query_settings, partition_key, /* sequence_number */1))
+            *object_storage, *configuration, query_settings, file_path, /* sequence_number */1))
     {
-        partition_key = *new_key;
+        file_path = *new_key;
     }
 
     return std::make_shared<StorageObjectStorageSink>(
         object_storage,
         configuration,
         format_settings,
-        sample_block,
+        getPartitionStrategy()->getBlockWithoutPartitionColumnsIfNeeded(),
         context,
-        partition_key
+        file_path
     );
 }
 
-void PartitionedStorageObjectStorageSink::validateKey(const String & str)
-{
-    /// See:
-    /// - https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
-    /// - https://cloud.ibm.com/apidocs/cos/cos-compatibility#putobject
-
-    if (str.empty() || str.size() > 1024)
-        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Incorrect key length (not empty, max 1023 characters), got: {}", str.size());
-
-    if (!UTF8::isValidUTF8(reinterpret_cast<const UInt8 *>(str.data()), str.size()))
-        throw Exception(ErrorCodes::CANNOT_PARSE_TEXT, "Incorrect non-UTF8 sequence in key");
-
-    validatePartitionKey(str, true);
-}
-
-void PartitionedStorageObjectStorageSink::validateNamespace(const String & str)
-{
-    configuration->validateNamespace(str);
-
-    if (!UTF8::isValidUTF8(reinterpret_cast<const UInt8 *>(str.data()), str.size()))
-        throw Exception(ErrorCodes::CANNOT_PARSE_TEXT, "Incorrect non-UTF8 sequence in bucket name");
-
-    validatePartitionKey(str, false);
-}
-
 }

src/Storages/ObjectStorage/StorageObjectStorageSink.h

@@ -42,19 +42,16 @@ class PartitionedStorageObjectStorageSink : public PartitionedSink
     using ConfigurationPtr = StorageObjectStorage::ConfigurationPtr;
 
     PartitionedStorageObjectStorageSink(
+        std::shared_ptr<PartitionStrategy> partition_strategy_,
         ObjectStoragePtr object_storage_,
         ConfigurationPtr configuration_,
         std::optional<FormatSettings> format_settings_,
         const Block & sample_block_,
-        ContextPtr context_,
-        const ASTPtr & partition_by);
+        ContextPtr context_);
 
     SinkPtr createSinkForPartition(const String & partition_id) override;
 
 private:
-    void validateKey(const String & str);
-    void validateNamespace(const String & str);
-
     ObjectStoragePtr object_storage;
     ConfigurationPtr configuration;