Merge pull request ClickHouse#82131 from ClickHouse/fix-partiton-pruning-in-data-lake-cluster-functions

Fix partition pruning with data lake cluster functions

Author: kssenii

src/Planner/Planner.cpp

@@ -52,6 +52,7 @@
 #include <Storages/StorageDistributed.h>
 #include <Storages/StorageDummy.h>
 #include <Storages/StorageMerge.h>
+#include <Storages/ObjectStorage/StorageObjectStorageCluster.h>
 
 #include <AggregateFunctions/IAggregateFunction.h>
 
@@ -231,6 +232,11 @@ FiltersForTableExpressionMap collectFiltersForAnalysis(const QueryTreeNodePtr &
             collect_filters = true;
             break;
         }
+        if (typeid_cast<const StorageObjectStorageCluster *>(storage.get()))
+        {
+            collect_filters = true;
+            break;
+        }
     }
 
     if (!collect_filters)

src/Storages/IStorageCluster.cpp

@@ -117,7 +117,11 @@ void ReadFromCluster::createExtension(const ActionsDAG::Node * predicate, size_t
     if (extension)
         return;
 
-    extension = storage->getTaskIteratorExtension(predicate, context, number_of_replicas);
+    extension = storage->getTaskIteratorExtension(
+        predicate,
+        filter_actions_dag.has_value() ? &filter_actions_dag.value() : query_info.filter_actions_dag.get(),
+        context,
+        number_of_replicas);
 }
 
 /// The code executes on initiator

src/Storages/IStorageCluster.h

@@ -34,8 +34,13 @@ class IStorageCluster : public IStorage
         size_t /*num_streams*/) override;
 
     ClusterPtr getCluster(ContextPtr context) const;
+
     /// Query is needed for pruning by virtual columns (_file, _path)
-    virtual RemoteQueryExecutor::Extension getTaskIteratorExtension(const ActionsDAG::Node * predicate, const ContextPtr & context, size_t number_of_replicas) const = 0;
+    virtual RemoteQueryExecutor::Extension getTaskIteratorExtension(
+        const ActionsDAG::Node * predicate,
+        const ActionsDAG * filter_actions_dag,
+        const ContextPtr & context,
+        size_t number_of_replicas) const = 0;
 
     QueryProcessingStage::Enum getQueryProcessingStage(ContextPtr, QueryProcessingStage::Enum, const StorageSnapshotPtr &, SelectQueryInfo &) const override;
 

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -451,7 +451,7 @@ class ReadFromObjectStorageStep : public SourceStepWithFilter
         auto context = getContext();
         iterator_wrapper = StorageObjectStorageSource::createFileIterator(
             configuration, configuration->getQuerySettings(context), object_storage, distributed_processing,
-            context, predicate, filter_actions_dag, virtual_columns, info.hive_partition_columns_to_read_from_file_path, nullptr, context->getFileProgressCallback());
+            context, predicate, filter_actions_dag.has_value() ? &filter_actions_dag.value() : nullptr, virtual_columns, info.hive_partition_columns_to_read_from_file_path, nullptr, context->getFileProgressCallback());
     }
 };
 }

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -213,12 +213,27 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
     }
 }
 
+
 RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExtension(
-    const ActionsDAG::Node * predicate, const ContextPtr & local_context, const size_t number_of_replicas) const
+    const ActionsDAG::Node * predicate,
+    const ActionsDAG * filter,
+    const ContextPtr & local_context,
+    const size_t number_of_replicas) const
 {
     auto iterator = StorageObjectStorageSource::createFileIterator(
-        configuration, configuration->getQuerySettings(local_context), object_storage, /* distributed_processing */false,
-        local_context, predicate, {}, virtual_columns, hive_partition_columns_to_read_from_file_path, nullptr, local_context->getFileProgressCallback(), /*ignore_archive_globs=*/true, /*skip_object_metadata=*/true);
+        configuration,
+        configuration->getQuerySettings(local_context),
+        object_storage,
+        /* distributed_processing */false,
+        local_context,
+        predicate,
+        filter,
+        virtual_columns,
+		hive_partition_columns_to_read_from_file_path,
+        nullptr,
+        local_context->getFileProgressCallback(),
+        /*ignore_archive_globs=*/true,
+        /*skip_object_metadata=*/true);
 
     auto task_distributor = std::make_shared<StorageObjectStorageStableTaskDistributor>(iterator, number_of_replicas);
 

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -25,7 +25,10 @@ class StorageObjectStorageCluster : public IStorageCluster
     std::string getName() const override;
 
     RemoteQueryExecutor::Extension getTaskIteratorExtension(
-        const ActionsDAG::Node * predicate, const ContextPtr & context, size_t number_of_replicas) const override;
+        const ActionsDAG::Node * predicate,
+        const ActionsDAG * filter,
+        const ContextPtr & context,
+        size_t number_of_replicas) const override;
 
     String getPathSample(ContextPtr context);
 

src/Storages/ObjectStorage/StorageObjectStorageSource.cpp

@@ -130,7 +130,7 @@ std::shared_ptr<IObjectIterator> StorageObjectStorageSource::createFileIterator(
     bool distributed_processing,
     const ContextPtr & local_context,
     const ActionsDAG::Node * predicate,
-    const std::optional<ActionsDAG> & filter_actions_dag,
+    const ActionsDAG * filter_actions_dag,
     const NamesAndTypesList & virtual_columns,
     const NamesAndTypesList & hive_columns,
     ObjectInfos * read_keys,
@@ -171,7 +171,7 @@ std::shared_ptr<IObjectIterator> StorageObjectStorageSource::createFileIterator(
     else if (configuration->supportsFileIterator())
     {
         return configuration->iterate(
-            filter_actions_dag.has_value() ? &filter_actions_dag.value() : nullptr,
+            filter_actions_dag,
             file_progress_callback,
             query_settings.list_object_keys_size,
             local_context);

src/Storages/ObjectStorage/StorageObjectStorageSource.h

@@ -57,7 +57,7 @@ class StorageObjectStorageSource : public SourceWithKeyCondition
         bool distributed_processing,
         const ContextPtr & local_context,
         const ActionsDAG::Node * predicate,
-        const std::optional<ActionsDAG> & filter_actions_dag,
+        const ActionsDAG * filter_actions_dag,
         const NamesAndTypesList & virtual_columns,
         const NamesAndTypesList & hive_columns,
         ObjectInfos * read_keys,

src/Storages/StorageDistributed.cpp

@@ -1318,7 +1318,8 @@ std::optional<QueryPipeline> StorageDistributed::distributedWriteFromClusterStor
 
     /// Select query is needed for pruining on virtual columns
     auto number_of_replicas = static_cast<UInt64>(cluster->getShardsInfo().size());
-    auto extension = src_storage_cluster.getTaskIteratorExtension(predicate, local_context, number_of_replicas);
+    auto extension = src_storage_cluster.getTaskIteratorExtension(
+        predicate, filter.has_value() ? &filter.value() : nullptr, local_context, number_of_replicas);
 
     /// Here we take addresses from destination cluster and assume source table exists on these nodes
     size_t replica_index = 0;

src/Storages/StorageFileCluster.cpp

@@ -100,7 +100,11 @@ void StorageFileCluster::updateQueryToSendIfNeeded(DB::ASTPtr & query, const Sto
     );
 }
 
-RemoteQueryExecutor::Extension StorageFileCluster::getTaskIteratorExtension(const ActionsDAG::Node * predicate, const ContextPtr & context, const size_t) const
+RemoteQueryExecutor::Extension StorageFileCluster::getTaskIteratorExtension(
+    const ActionsDAG::Node * predicate,
+    const ActionsDAG * /* filter */,
+    const ContextPtr & context,
+    const size_t) const
 {
     auto iterator = std::make_shared<StorageFileSource::FilesIterator>(paths, std::nullopt, predicate, getVirtualsList(), hive_partition_columns_to_read_from_file_path, context);
     auto callback = std::make_shared<TaskIterator>([iter = std::move(iterator)](size_t) mutable -> String { return iter->next(); });