fix buld squashed

Author: zvonand

src/Core/Settings.h

@@ -110,8 +110,7 @@ class WriteBuffer;
     M(CLASS_NAME, URI) \
     M(CLASS_NAME, VectorSearchFilterStrategy) \
     M(CLASS_NAME, GeoToH3ArgumentOrder) \
-    M(CLASS_NAME, ObjectStorageGranularityLevel) \
-    M(CLASS_NAME, DecorrelationJoinKind)
+    M(CLASS_NAME, ObjectStorageGranularityLevel)
 
 
 COMMON_SETTINGS_SUPPORTED_TYPES(Settings, DECLARE_SETTING_TRAIT)

src/Disks/ObjectStorages/IObjectStorage.h

@@ -150,6 +150,8 @@ struct PathWithMetadata
     std::optional<String> absolute_path;
     ObjectStoragePtr object_storage_to_use = nullptr;
 
+    FileBucketInfoPtr file_bucket_info;
+
     PathWithMetadata() = default;
 
     explicit PathWithMetadata(
@@ -189,6 +191,14 @@ struct PathWithMetadata
     void loadMetadata(ObjectStoragePtr object_storage, bool ignore_non_existent_file = true);
 
     ObjectStoragePtr getObjectStorage() const { return object_storage_to_use; }
+
+    String getIdentifier() const
+    {
+        String result = absolute_path.value_or(relative_path);
+        if (file_bucket_info)
+            result += file_bucket_info->getIdentifier();
+        return result;
+    }
 };
 
 struct ObjectKeyWithMetadata

src/Interpreters/ClusterFunctionReadTask.h

@@ -1,11 +1,8 @@
 #pragma once
 #include <Core/Types.h>
-<<<<<<< HEAD
-=======
 #include <Storages/ObjectStorage/DataLakes/DataLakeObjectMetadata.h>
 #include <Processors/Formats/IInputFormat.h>
 #include <Storages/ObjectStorage/DataLakes/Iceberg/IcebergDataObjectInfo.h>
->>>>>>> 4bed2ad0c69 (Merge pull request #87508 from scanhex12/distributed_execution_better_spread)
 #include <Storages/ObjectStorage/IObjectIterator.h>
 #include <Storages/ObjectStorage/DataLakes/DataLakeObjectMetadata.h>
 

src/Storages/ObjectStorage/DataLakes/IDataLakeMetadata.h

@@ -6,7 +6,7 @@
 #include <Core/Range.h>
 #include <Interpreters/ActionsDAG.h>
 #include <Processors/ISimpleTransform.h>
-#include <Storages/ObjectStorage/IObjectIterator.h>
+#include <Disks/ObjectStorages/IObjectStorage.h>
 #include <QueryPipeline/QueryPipelineBuilder.h>
 #include <Storages/AlterCommands.h>
 #include <Storages/MutationCommands.h>
@@ -94,10 +94,9 @@ struct StorageID;
 struct IObjectIterator;
 struct RelativePathWithMetadata;
 class IObjectStorage;
-struct ObjectInfo;
-using ObjectInfoPtr = std::shared_ptr<ObjectInfo>;
 using ObjectIterator = std::shared_ptr<IObjectIterator>;
 using ObjectStoragePtr = std::shared_ptr<IObjectStorage>;
+using ObjectInfoPtr = std::shared_ptr<PathWithMetadata>;
 
 class IDataLakeMetadata : boost::noncopyable
 {

src/Storages/ObjectStorage/IObjectIterator.cpp

@@ -76,4 +76,49 @@ ObjectInfoPtr ObjectIteratorWithPathAndFileFilter::next(size_t id)
     return {};
 }
 
+ObjectIteratorSplitByBuckets::ObjectIteratorSplitByBuckets(
+    ObjectIterator iterator_,
+    const String & format_,
+    ObjectStoragePtr object_storage_,
+    const ContextPtr & context_)
+    : WithContext(context_)
+    , iterator(iterator_)
+    , format(format_)
+    , object_storage(object_storage_)
+    , format_settings(getFormatSettings(context_))
+{
+}
+
+ObjectInfoPtr ObjectIteratorSplitByBuckets::next(size_t id)
+{
+    if (!pending_objects_info.empty())
+    {
+        auto result = pending_objects_info.front();
+        pending_objects_info.pop();
+        return result;
+    }
+    auto last_object_info = iterator->next(id);
+    if (!last_object_info)
+        return {};
+
+    auto buffer = createReadBuffer(*last_object_info, object_storage, getContext(), log);
+
+    auto splitter = FormatFactory::instance().getSplitter(format);
+    if (splitter)
+    {
+        size_t bucket_size = getContext()->getSettingsRef()[Setting::cluster_table_function_buckets_batch_size];
+        auto file_bucket_info = splitter->splitToBuckets(bucket_size, *buffer, format_settings);
+        for (const auto & file_bucket : file_bucket_info)
+        {
+            auto copy_object_info = *last_object_info;
+            copy_object_info.file_bucket_info = file_bucket;
+            pending_objects_info.push(std::make_shared<ObjectInfo>(copy_object_info));
+        }
+    }
+
+    auto result = pending_objects_info.front();
+    pending_objects_info.pop();
+    return result;
+}
+
 }

src/Storages/ObjectStorage/StorageObjectStorageConfiguration.h

@@ -27,6 +27,7 @@ class IDataLakeMetadata;
 struct IObjectIterator;
 using SinkToStoragePtr = std::shared_ptr<SinkToStorage>;
 using ObjectIterator = std::shared_ptr<IObjectIterator>;
+using ObjectInfoPtr = std::shared_ptr<PathWithMetadata>;
 
 namespace ErrorCodes
 {
@@ -281,17 +282,16 @@ class StorageObjectStorageConfiguration
         return false;
     }
 
+    String format = "auto";
+    String compression_method = "auto";
+    String structure = "auto";
+
     PartitionStrategyFactory::StrategyType partition_strategy_type = PartitionStrategyFactory::StrategyType::NONE;
     std::shared_ptr<IPartitionStrategy> partition_strategy;
     /// Whether partition column values are contained in the actual data.
     /// And alternative is with hive partitioning, when they are contained in file path.
     bool partition_columns_in_data_file = true;
 
-private:
-    String format = "auto";
-    String compression_method = "auto";
-    String structure = "auto";
-
 protected:
     bool initialized = false;
 

src/Storages/StorageURL.cpp

@@ -807,10 +807,10 @@ std::function<void(std::ostream &)> IStorageURLBase::getReadPOSTDataCallback(
 
 namespace
 {
-    class ReadBufferIterator : public IReadBufferIterator, WithContext
+    class URLReadBufferIterator : public IReadBufferIterator, WithContext
     {
     public:
-        ReadBufferIterator(
+        URLReadBufferIterator(
             const std::vector<String> & urls_to_check_,
             std::optional<String> format_,
             const CompressionMethod & compression_method_,
@@ -1054,7 +1054,7 @@ std::pair<ColumnsDescription, String> IStorageURLBase::getTableStructureAndForma
     else
         urls_to_check = {uri};
 
-    ReadBufferIterator read_buffer_iterator(urls_to_check, format, compression_method, headers, format_settings, context);
+    URLReadBufferIterator read_buffer_iterator(urls_to_check, format, compression_method, headers, format_settings, context);
     if (format)
         return {readSchemaFromFormat(*format, format_settings, read_buffer_iterator, context), *format};
     return detectFormatAndReadSchema(format_settings, read_buffer_iterator, context);

tests/integration/test_storage_iceberg/test.py

@@ -3925,7 +3925,7 @@ def check_validity_and_get_prunned_files(select_expression):
     )
 
 
-    
+
 def test_iceberg_write_minmax(started_cluster):
     instance = started_cluster.instances["node1"]
     TABLE_NAME = "test_iceberg_write_minmax_" + get_uuid_str()
@@ -3939,3 +3939,95 @@ def test_iceberg_write_minmax(started_cluster):
 
     res = instance.query(f"SELECT x,y FROM {TABLE_NAME} WHERE y=2 ORDER BY ALL").strip()
     assert res == "1\t2"
+
+
+@pytest.mark.parametrize("format_version", ["1", "2"])
+@pytest.mark.parametrize("storage_type", ["s3", "azure"])
+@pytest.mark.parametrize("cluster_table_function_buckets_batch_size", [0, 100, 1000])
+@pytest.mark.parametrize("input_format_parquet_use_native_reader_v3", [0, 1])
+def test_cluster_table_function_split_by_row_groups(started_cluster_iceberg_with_spark, format_version, storage_type, cluster_table_function_buckets_batch_size,input_format_parquet_use_native_reader_v3):
+    instance = started_cluster_iceberg_with_spark.instances["node1"]
+    spark = started_cluster_iceberg_with_spark.spark_session
+
+    TABLE_NAME = (
+        "test_iceberg_cluster_"
+        + format_version
+        + "_"
+        + storage_type
+        + "_"
+        + get_uuid_str()
+    )
+
+    def add_df(mode):
+        write_iceberg_from_df(
+            spark,
+            generate_data(spark, 0, 100000),
+            TABLE_NAME,
+            mode=mode,
+            format_version=format_version,
+        )
+
+        files = default_upload_directory(
+            started_cluster_iceberg_with_spark,
+            storage_type,
+            f"/iceberg_data/default/{TABLE_NAME}/",
+            f"/iceberg_data/default/{TABLE_NAME}/",
+        )
+
+        logging.info(f"Adding another dataframe. result files: {files}")
+
+        return files
+
+    files = add_df(mode="overwrite")
+    for i in range(1, 5 * len(started_cluster_iceberg_with_spark.instances)):
+        files = add_df(mode="append")
+
+    clusters = instance.query(f"SELECT * FROM system.clusters")
+    logging.info(f"Clusters setup: {clusters}")
+
+    # Regular Query only node1
+    table_function_expr = get_creation_expression(
+        storage_type, TABLE_NAME, started_cluster_iceberg_with_spark, table_function=True
+    )
+    select_regular = (
+        instance.query(f"SELECT * FROM {table_function_expr} ORDER BY ALL").strip().split()
+    )
+
+    # Cluster Query with node1 as coordinator
+    table_function_expr_cluster = get_creation_expression(
+        storage_type,
+        TABLE_NAME,
+        started_cluster_iceberg_with_spark,
+        table_function=True,
+        run_on_cluster=True,
+    )
+    instance.query("SYSTEM FLUSH LOGS")
+
+    def get_buffers_count(func):
+        buffers_count_before = int(
+            instance.query(
+                f"SELECT sum(ProfileEvents['EngineFileLikeReadFiles']) FROM system.query_log WHERE type = 'QueryFinish'"
+            )
+        )
+
+        func()
+        instance.query("SYSTEM FLUSH LOGS")
+        buffers_count = int(
+            instance.query(
+                f"SELECT sum(ProfileEvents['EngineFileLikeReadFiles']) FROM system.query_log WHERE type = 'QueryFinish'"
+            )
+        )
+        return buffers_count - buffers_count_before
+
+    select_cluster = (
+        instance.query(f"SELECT * FROM {table_function_expr_cluster} ORDER BY ALL SETTINGS input_format_parquet_use_native_reader_v3={input_format_parquet_use_native_reader_v3},cluster_table_function_split_granularity='bucket', cluster_table_function_buckets_batch_size={cluster_table_function_buckets_batch_size}").strip().split()
+    )
+
+    # Simple size check
+    assert len(select_cluster) == len(select_regular)
+    # Actual check
+    assert select_cluster == select_regular
+
+    buffers_count_with_splitted_tasks = get_buffers_count(lambda: instance.query(f"SELECT * FROM {table_function_expr_cluster} ORDER BY ALL SETTINGS input_format_parquet_use_native_reader_v3={input_format_parquet_use_native_reader_v3},cluster_table_function_split_granularity='bucket', cluster_table_function_buckets_batch_size={cluster_table_function_buckets_batch_size}").strip().split())
+    buffers_count_default = get_buffers_count(lambda: instance.query(f"SELECT * FROM {table_function_expr_cluster} ORDER BY ALL SETTINGS input_format_parquet_use_native_reader_v3={input_format_parquet_use_native_reader_v3}, cluster_table_function_buckets_batch_size={cluster_table_function_buckets_batch_size}").strip().split())
+    assert buffers_count_with_splitted_tasks > buffers_count_default