Merge pull request ClickHouse#86880 from ClickHouse/backport/25.8/86357

Backport ClickHouse#86357 to 25.8: Add rows/bytes limit for inserted data files in delta lake

Author: clickhouse-gh[bot]

src/Core/Settings.cpp

@@ -6471,12 +6471,12 @@ Query Iceberg table using the snapshot that was current at a specific timestamp.
 )", 0) \
     DECLARE(Int64, iceberg_snapshot_id, 0, R"(
 Query Iceberg table using the specific snapshot id.
-)", 0) \
-    DECLARE(Bool, delta_lake_enable_expression_visitor_logging, false, R"(
-Enables Test level logs of DeltaLake expression visitor. These logs can be too verbose even for test logging.
 )", 0) \
     DECLARE(Bool, show_data_lake_catalogs_in_system_tables, true, R"(
 Enables showing data lake catalogs in system tables.
+)", 0) \
+    DECLARE(Bool, delta_lake_enable_expression_visitor_logging, false, R"(
+Enables Test level logs of DeltaLake expression visitor. These logs can be too verbose even for test logging.
 )", 0) \
     DECLARE(Int64, delta_lake_snapshot_version, -1, R"(
 Version of delta lake snapshot to read. Value -1 means to read latest version (value 0 is a valid snapshot version).
@@ -6486,6 +6486,12 @@ Enables throwing an exception if there was an error when analyzing scan predicat
 )", 0) \
     DECLARE(Bool, delta_lake_enable_engine_predicate, true, R"(
 Enables delta-kernel internal data pruning.
+)", 0) \
+    DECLARE(NonZeroUInt64, delta_lake_insert_max_rows_in_data_file, 100000, R"(
+Defines a rows limit for a single inserted data file in delta lake.
+)", 0) \
+    DECLARE(NonZeroUInt64, delta_lake_insert_max_bytes_in_data_file, 1_GiB, R"(
+Defines a bytes limit for a single inserted data file in delta lake.
 )", 0) \
     DECLARE(Bool, allow_experimental_delta_lake_writes, false, R"(
 Enables delta-kernel writes feature.

src/Core/SettingsChangesHistory.cpp

@@ -96,6 +96,8 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"per_part_index_stats", false, false, "New setting."},
             {"allow_experimental_iceberg_compaction", 0, 0, "New setting "},
             {"delta_lake_snapshot_version", -1, -1, "New setting"},
+            {"delta_lake_insert_max_bytes_in_data_file", 1_GiB, 1_GiB, "New setting."},
+            {"delta_lake_insert_max_rows_in_data_file", 100000, 100000, "New setting."},
             {"use_roaring_bitmap_iceberg_positional_deletes", false, false, "New setting"},
             {"iceberg_metadata_compression_method", "", "", "New setting"},
             {"allow_experimental_correlated_subqueries", false, true, "Mark correlated subqueries support as Beta."},

src/Storages/ObjectStorage/DataLakes/DeltaLake/DeltaLakePartitionedSink.cpp

@@ -6,6 +6,7 @@
 #include <Common/Arena.h>
 #include <Common/PODArray.h>
 #include <Core/UUID.h>
+#include <Core/Settings.h>
 
 #include <Formats/FormatFactory.h>
 #include <Processors/Formats/IOutputFormat.h>
@@ -30,6 +31,12 @@ namespace ErrorCodes
     extern const int INCORRECT_DATA;
 }
 
+namespace Setting
+{
+    extern const SettingsNonZeroUInt64 delta_lake_insert_max_rows_in_data_file;
+    extern const SettingsNonZeroUInt64 delta_lake_insert_max_bytes_in_data_file;
+}
+
 namespace
 {
     /// Given partition columns list,
@@ -74,6 +81,8 @@ DeltaLakePartitionedSink::DeltaLakePartitionedSink(
     , object_storage(object_storage_)
     , format_settings(format_settings_)
     , configuration(configuration_)
+    , data_file_max_rows(context_->getSettingsRef()[Setting::delta_lake_insert_max_rows_in_data_file])
+    , data_file_max_bytes(context_->getSettingsRef()[Setting::delta_lake_insert_max_bytes_in_data_file])
     , partition_strategy(createPartitionStrategy(partition_columns, getHeader(), context_))
     , delta_transaction(delta_transaction_)
 {
@@ -131,7 +140,7 @@ void DeltaLakePartitionedSink::consume(Chunk & chunk)
                 partition_index_to_chunk.emplace_back(Columns(), partition_column->size());
         }
 
-        for (size_t partition_index = 0; partition_index  < partitions_size; ++partition_index)
+        for (size_t partition_index = 0; partition_index < partitions_size; ++partition_index)
         {
             auto & partition_chunk = partition_index_to_chunk[partition_index];
             partition_chunk.addColumn(std::move(partition_index_to_column_split[partition_index]));
@@ -140,68 +149,86 @@ void DeltaLakePartitionedSink::consume(Chunk & chunk)
 
     for (const auto & [partition_key, partition_index] : partition_id_to_chunk_index)
     {
-        auto partition_data = getPartitionDataForPartitionKey(partition_key);
+        auto & data_files = getPartitionDataForPartitionKey(partition_key)->data_files;
         auto & partition_chunk = partition_index_to_chunk[partition_index];
-        partition_data->sink->consume(partition_chunk);
-        partition_data->size += partition_chunk.bytes();
+
+        if (data_files.empty()
+            || data_files.back().written_rows >= data_file_max_rows
+            || data_files.back().written_bytes >= data_file_max_bytes)
+        {
+            data_files.emplace_back(createSinkForPartition(partition_key));
+            total_data_files_count += 1;
+        }
+        auto & data_file = data_files.back();
+        data_file.written_bytes += partition_chunk.bytes();
+        data_file.written_rows += partition_chunk.getNumRows();
+        data_file.sink->consume(partition_chunk);
     }
 }
 
-DeltaLakePartitionedSink::PartitionDataPtr
+DeltaLakePartitionedSink::PartitionInfoPtr
 DeltaLakePartitionedSink::getPartitionDataForPartitionKey(StringRef partition_key)
 {
-    auto it = partition_id_to_sink.find(partition_key);
-    if (it == partition_id_to_sink.end())
-    {
-        auto data = std::make_shared<PartitionData>();
-        auto data_prefix = std::filesystem::path(delta_transaction->getDataPath()) / partition_key.toString();
-        data->path = DeltaLake::generateWritePath(std::move(data_prefix), configuration->format);
-
-        data->sink = std::make_shared<StorageObjectStorageSink>(
-            data->path,
-            object_storage,
-            configuration,
-            format_settings,
-            std::make_shared<Block>(partition_strategy->getFormatHeader()),
-            getContext()
-        );
-        std::tie(it, std::ignore) = partition_id_to_sink.emplace(partition_key, std::move(data));
-    }
+    auto it = partitions_data.find(partition_key);
+    if (it == partitions_data.end())
+        std::tie(it, std::ignore) = partitions_data.emplace(partition_key, std::make_shared<PartitionInfo>(partition_key));
     return it->second;
 }
 
+DeltaLakePartitionedSink::StorageSinkPtr
+DeltaLakePartitionedSink::createSinkForPartition(StringRef partition_key)
+{
+    auto data_prefix = std::filesystem::path(delta_transaction->getDataPath()) / partition_key.toString();
+    return std::make_unique<StorageObjectStorageSink>(
+        DeltaLake::generateWritePath(std::move(data_prefix), configuration->format),
+        object_storage,
+        configuration,
+        format_settings,
+        std::make_shared<Block>(partition_strategy->getFormatHeader()),
+        getContext());
+}
+
 void DeltaLakePartitionedSink::onFinish()
 {
-    if (isCancelled() || partition_id_to_sink.empty())
+    if (isCancelled() || partitions_data.empty())
         return;
 
-    for (auto & [_, data] : partition_id_to_sink)
-        data->sink->onFinish();
-
-    LOG_TEST(log, "Written to {} sinks", partition_id_to_sink.size());
+    std::vector<DeltaLake::WriteTransaction::CommitFile> files;
+    files.reserve(total_data_files_count);
+    const auto data_prefix = delta_transaction->getDataPath();
 
-    try
+    for (auto & [_, partition_info] : partitions_data)
     {
-        std::vector<DeltaLake::WriteTransaction::CommitFile> files;
-        files.reserve(partition_id_to_sink.size());
-        const auto data_prefix = delta_transaction->getDataPath();
-        for (auto & [_, data] : partition_id_to_sink)
+        auto & [partition_key, data_files] = *partition_info;
+        auto partition_key_str = partition_key.toString();
+        auto keys_and_values = HivePartitioningUtils::parseHivePartitioningKeysAndValues(partition_key_str);
+        Map partition_values;
+        partition_values.reserve(keys_and_values.size());
+        for (const auto & [key, value] : keys_and_values)
+            partition_values.emplace_back(DB::Tuple({key, value}));
+
+        for (const auto & [sink, written_bytes, written_rows] : data_files)
         {
-            auto keys_and_values = HivePartitioningUtils::parseHivePartitioningKeysAndValues(data->path);
-            Map partition_values;
-            partition_values.reserve(keys_and_values.size());
-            for (const auto & [key, value] : keys_and_values)
-                partition_values.emplace_back(DB::Tuple({key, value}));
-
-            files.emplace_back(data->path.substr(data_prefix.size()), data->size, partition_values);
+            sink->onFinish();
+            files.emplace_back(
+                sink->getPath().substr(data_prefix.size()),
+                sink->getFileSize(),
+                partition_values);
         }
+    }
+
+    LOG_TEST(log, "Written {} data files", total_data_files_count);
+
+    try
+    {
         delta_transaction->commit(files);
     }
     catch (...)
     {
-        for (auto & [_, data] : partition_id_to_sink)
+        for (auto & [_, partition_info] : partitions_data)
         {
-            object_storage->removeObjectIfExists(StoredObject(data->path));
+            for (const auto & [sink, written_bytes, written_rows] : partition_info->data_files)
+                object_storage->removeObjectIfExists(StoredObject(sink->getPath()));
         }
         throw;
     }

src/Storages/ObjectStorage/DataLakes/DeltaLake/DeltaLakePartitionedSink.h

@@ -9,6 +9,7 @@
 #include <absl/container/flat_hash_map.h>
 #include <Processors/Sinks/SinkToStorage.h>
 #include <Storages/ObjectStorage/IObjectIterator.h>
+#include <Storages/ObjectStorage/StorageObjectStorageSink.h>
 #include <Storages/IPartitionStrategy.h>
 
 
@@ -49,26 +50,42 @@ class DeltaLakePartitionedSink : public SinkToStorage, private WithContext
     void onFinish() override;
 
 private:
-    struct PartitionData
+    using StorageSinkPtr = std::unique_ptr<StorageObjectStorageSink>;
+
+    struct DataFileInfo
+    {
+        explicit DataFileInfo(StorageSinkPtr sink_) : sink(std::move(sink_)) {}
+
+        StorageSinkPtr sink;
+        size_t written_bytes = 0;
+        size_t written_rows = 0;
+    };
+    struct PartitionInfo
     {
-        SinkPtr sink;
-        std::string path;
-        size_t size = 0;
+        explicit PartitionInfo(StringRef partition_key_) : partition_key(partition_key_) {}
+
+        const StringRef partition_key;
+        std::vector<DataFileInfo> data_files;
     };
-    using PartitionDataPtr = std::shared_ptr<PartitionData>;
-    PartitionDataPtr getPartitionDataForPartitionKey(StringRef partition_key);
+    using PartitionInfoPtr = std::shared_ptr<PartitionInfo>;
 
     const LoggerPtr log;
     const Names partition_columns;
     const ObjectStoragePtr object_storage;
     const std::optional<FormatSettings> format_settings;
     const StorageObjectStorageConfigurationPtr configuration;
+    const size_t data_file_max_rows;
+    const size_t data_file_max_bytes;
     const std::unique_ptr<IPartitionStrategy> partition_strategy;
     const DeltaLake::WriteTransactionPtr delta_transaction;
 
-    absl::flat_hash_map<StringRef, PartitionDataPtr> partition_id_to_sink;
+    absl::flat_hash_map<StringRef, PartitionInfoPtr> partitions_data;
+    size_t total_data_files_count = 0;
     IColumn::Selector chunk_row_index_to_partition_index;
     Arena partition_keys_arena;
+
+    StorageSinkPtr createSinkForPartition(StringRef partition_key);
+    PartitionInfoPtr getPartitionDataForPartitionKey(StringRef partition_key);
 };
 
 }

src/Storages/ObjectStorage/DataLakes/DeltaLake/DeltaLakeSink.cpp

@@ -1,14 +1,21 @@
+#include <Core/Settings.h>
 #include <Storages/ObjectStorage/DataLakes/DeltaLake/DeltaLakeSink.h>
 
 #if USE_DELTA_KERNEL_RS
 #include <Common/logger_useful.h>
+#include <Interpreters/Context.h>
 #include <Storages/ObjectStorage/DataLakes/DeltaLakeMetadataDeltaKernel.h>
 #include <Storages/ObjectStorage/DataLakes/DeltaLake/WriteTransaction.h>
 #include <Storages/ObjectStorage/DataLakes/DeltaLake/KernelUtils.h>
 
 
 namespace DB
 {
+namespace Setting
+{
+    extern const SettingsNonZeroUInt64 delta_lake_insert_max_rows_in_data_file;
+    extern const SettingsNonZeroUInt64 delta_lake_insert_max_bytes_in_data_file;
+}
 
 DeltaLakeSink::DeltaLakeSink(
     DeltaLake::WriteTransactionPtr delta_transaction_,
@@ -17,49 +24,78 @@ DeltaLakeSink::DeltaLakeSink(
     ContextPtr context_,
     SharedHeader sample_block_,
     const std::optional<FormatSettings> & format_settings_)
-    : StorageObjectStorageSink(
-        DeltaLake::generateWritePath(
-            delta_transaction_->getDataPath(),
-            configuration_->format),
-        object_storage_,
-        configuration_,
-        format_settings_,
-        sample_block_,
-        context_)
+    : SinkToStorage(sample_block_)
+    , WithContext(context_)
     , delta_transaction(delta_transaction_)
     , object_storage(object_storage_)
+    , configuration(configuration_)
+    , format_settings(format_settings_)
+    , sample_block(sample_block_)
+    , data_file_max_rows(context_->getSettingsRef()[Setting::delta_lake_insert_max_rows_in_data_file])
+    , data_file_max_bytes(context_->getSettingsRef()[Setting::delta_lake_insert_max_bytes_in_data_file])
 {
     delta_transaction->validateSchema(getHeader());
 }
 
+DeltaLakeSink::StorageSinkPtr DeltaLakeSink::createStorageSink() const
+{
+    return std::make_unique<StorageObjectStorageSink>(
+        DeltaLake::generateWritePath(
+            delta_transaction->getDataPath(),
+            configuration->format),
+        object_storage,
+        configuration,
+        format_settings,
+        sample_block,
+        getContext());
+}
+
 void DeltaLakeSink::consume(Chunk & chunk)
 {
     if (isCancelled())
         return;
 
-    written_bytes += chunk.bytes();
-    StorageObjectStorageSink::consume(chunk);
+    if (data_files.empty()
+        || data_files.back().written_bytes >= data_file_max_bytes
+        || data_files.back().written_rows >= data_file_max_rows)
+    {
+        data_files.emplace_back(createStorageSink());
+    }
+
+    auto & data_file = data_files.back();
+    data_file.written_bytes += chunk.bytes();
+    data_file.written_rows += chunk.getNumRows();
+    data_file.sink->consume(chunk);
 }
 
 void DeltaLakeSink::onFinish()
 {
     if (isCancelled())
         return;
 
-    StorageObjectStorageSink::onFinish();
+    std::vector<DeltaLake::WriteTransaction::CommitFile> files;
+    files.reserve(data_files.size());
+    for (const auto & [sink, written_bytes, written_rows] : data_files)
+    {
+        sink->onFinish();
+        auto file_location = sink->getPath().substr(delta_transaction->getDataPath().size());
+        auto file_size = sink->getFileSize();
+        files.emplace_back(std::move(file_location), file_size, Map{});
+    }
 
     try
     {
-        std::vector<DeltaLake::WriteTransaction::CommitFile> files;
-        auto file_location = getPath().substr(delta_transaction->getDataPath().size());
-        files.emplace_back(std::move(file_location), written_bytes, Map{});
         delta_transaction->commit(files);
     }
     catch (...)
     {
-        /// FIXME: this should be just removeObject,
-        /// but IObjectStorage does not have such method.
-        object_storage->removeObjectIfExists(StoredObject(getPath()));
+        for (const auto & [sink, written_bytes, written_rows] : data_files)
+        {
+            /// FIXME: this should be just removeObject,
+            /// but IObjectStorage does not have such method.
+           object_storage->removeObjectIfExists(StoredObject(sink->getPath()));
+
+        }
         throw;
     }
 }