s3Cluster hive partitioning

src/Planner/Planner.cpp

@@ -37,6 +37,7 @@
 #include <Processors/QueryPlan/WindowStep.h>
 #include <Processors/QueryPlan/ReadNothingStep.h>
 #include <Processors/QueryPlan/ReadFromRecursiveCTEStep.h>
+#include <Processors/QueryPlan/ObjectFilterStep.h>
 #include <QueryPipeline/QueryPipelineBuilder.h>
 
 #include <Interpreters/Context.h>
@@ -133,6 +134,7 @@ namespace Setting
     extern const SettingsUInt64 min_count_to_compile_aggregate_expression;
     extern const SettingsBool enable_software_prefetch_in_aggregation;
     extern const SettingsBool optimize_group_by_constant_keys;
+    extern const SettingsBool use_hive_partitioning;
 }
 
 namespace ServerSetting
@@ -413,6 +415,19 @@ void addFilterStep(QueryPlan & query_plan,
     query_plan.addStep(std::move(where_step));
 }
 
+void addObjectFilterStep(QueryPlan & query_plan,
+    FilterAnalysisResult & filter_analysis_result,
+    const std::string & step_description)
+{
+    auto actions = std::move(filter_analysis_result.filter_actions->dag);
+
+    auto where_step = std::make_unique<ObjectFilterStep>(query_plan.getCurrentHeader(),
+        std::move(actions),
+        filter_analysis_result.filter_column_name);
+    where_step->setStepDescription(step_description);
+    query_plan.addStep(std::move(where_step));
+}
+
 Aggregator::Params getAggregatorParams(const PlannerContextPtr & planner_context,
     const AggregationAnalysisResult & aggregation_analysis_result,
     const QueryAnalysisResult & query_analysis_result,
@@ -1670,6 +1685,16 @@ void Planner::buildPlanForQueryNode()
 
     if (query_processing_info.isSecondStage() || query_processing_info.isFromAggregationState())
     {
+        if (settings[Setting::use_hive_partitioning]
+            && !query_processing_info.isFirstStage()
+            && expression_analysis_result.hasWhere())
+        {
+            if (typeid_cast<ReadFromCluster *>(query_plan.getRootNode()->step.get()))
+            {
+                addObjectFilterStep(query_plan, expression_analysis_result.getWhere(), "WHERE");
+            }
+        }
+
         if (query_processing_info.isFromAggregationState())
         {
             /// Aggregation was performed on remote shards

src/Processors/QueryPlan/ObjectFilterStep.cpp

@@ -0,0 +1,63 @@
+#include <Processors/QueryPlan/ObjectFilterStep.h>
+#include <Processors/QueryPlan/QueryPlanStepRegistry.h>
+#include <Processors/QueryPlan/Serialization.h>
+#include <Processors/Transforms/FilterTransform.h>
+#include <IO/Operators.h>
+
+#include <memory>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int INCORRECT_DATA;
+}
+
+ObjectFilterStep::ObjectFilterStep(
+    const Header & input_header_,
+    ActionsDAG actions_dag_,
+    String filter_column_name_)
+    : actions_dag(std::move(actions_dag_))
+    , filter_column_name(std::move(filter_column_name_))
+{
+    input_headers.emplace_back(input_header_);
+    output_header = input_headers.front();
+}
+
+QueryPipelineBuilderPtr ObjectFilterStep::updatePipeline(QueryPipelineBuilders pipelines, const BuildQueryPipelineSettings & /* settings */)
+{
+    return std::move(pipelines.front());
+}
+
+void ObjectFilterStep::updateOutputHeader()
+{
+    output_header = input_headers.front();
+}
+
+void ObjectFilterStep::serialize(Serialization & ctx) const
+{
+    writeStringBinary(filter_column_name, ctx.out);
+
+    actions_dag.serialize(ctx.out, ctx.registry);
+}
+
+std::unique_ptr<IQueryPlanStep> ObjectFilterStep::deserialize(Deserialization & ctx)
+{
+    if (ctx.input_headers.size() != 1)
+        throw Exception(ErrorCodes::INCORRECT_DATA, "ObjectFilterStep must have one input stream");
+
+    String filter_column_name;
+    readStringBinary(filter_column_name, ctx.in);
+
+    ActionsDAG actions_dag = ActionsDAG::deserialize(ctx.in, ctx.registry, ctx.context);
+
+    return std::make_unique<ObjectFilterStep>(ctx.input_headers.front(), std::move(actions_dag), std::move(filter_column_name));
+}
+
+void registerObjectFilterStep(QueryPlanStepRegistry & registry)
+{
+    registry.registerStep("ObjectFilter", ObjectFilterStep::deserialize);
+}
+
+}

src/Processors/QueryPlan/ObjectFilterStep.h

@@ -0,0 +1,35 @@
+#pragma once
+#include <Processors/QueryPlan/IQueryPlanStep.h>
+#include <Interpreters/ActionsDAG.h>
+
+namespace DB
+{
+
+/// Implements WHERE operation.
+class ObjectFilterStep : public IQueryPlanStep
+{
+public:
+    ObjectFilterStep(
+        const Header & input_header_,
+        ActionsDAG actions_dag_,
+        String filter_column_name_);
+
+    String getName() const override { return "ObjectFilter"; }
+    QueryPipelineBuilderPtr updatePipeline(QueryPipelineBuilders pipelines, const BuildQueryPipelineSettings & settings) override;
+
+    const ActionsDAG & getExpression() const { return actions_dag; }
+    ActionsDAG & getExpression() { return actions_dag; }
+    const String & getFilterColumnName() const { return filter_column_name; }
+
+    void serialize(Serialization & ctx) const override;
+
+    static std::unique_ptr<IQueryPlanStep> deserialize(Deserialization & ctx);
+
+private:
+    void updateOutputHeader() override;
+
+    ActionsDAG actions_dag;
+    String filter_column_name;
+};
+
+}

src/Processors/QueryPlan/Optimizations/optimizePrimaryKeyConditionAndLimit.cpp

@@ -3,6 +3,7 @@
 #include <Processors/QueryPlan/FilterStep.h>
 #include <Processors/QueryPlan/LimitStep.h>
 #include <Processors/QueryPlan/SourceStepWithFilter.h>
+#include <Processors/QueryPlan/ObjectFilterStep.h>
 
 namespace DB::QueryPlanOptimizations
 {
@@ -41,6 +42,10 @@ void optimizePrimaryKeyConditionAndLimit(const Stack & stack)
             /// So this is likely not needed.
             continue;
         }
+        else if (auto * object_filter_step = typeid_cast<ObjectFilterStep *>(iter->node->step.get()))
+        {
+            source_step_with_filter->addFilter(object_filter_step->getExpression().clone(), object_filter_step->getFilterColumnName());
+        }
         else
         {
             break;

src/Processors/QueryPlan/QueryPlanStepRegistry.cpp

@@ -48,6 +48,7 @@ void registerOffsetStep(QueryPlanStepRegistry & registry);
 void registerFilterStep(QueryPlanStepRegistry & registry);
 void registerTotalsHavingStep(QueryPlanStepRegistry & registry);
 void registerExtremesStep(QueryPlanStepRegistry & registry);
+void registerObjectFilterStep(QueryPlanStepRegistry & registry);
 
 void QueryPlanStepRegistry::registerPlanSteps()
 {
@@ -65,6 +66,7 @@ void QueryPlanStepRegistry::registerPlanSteps()
     registerFilterStep(registry);
     registerTotalsHavingStep(registry);
     registerExtremesStep(registry);
+    registerObjectFilterStep(registry);
 }
 
 }

src/Storages/IStorageCluster.cpp

@@ -15,7 +15,6 @@
 #include <Parsers/queryToString.h>
 #include <Processors/Sources/NullSource.h>
 #include <Processors/Sources/RemoteSource.h>
-#include <Processors/QueryPlan/SourceStepWithFilter.h>
 #include <QueryPipeline/narrowPipe.h>
 #include <QueryPipeline/Pipe.h>
 #include <QueryPipeline/RemoteQueryExecutor.h>
@@ -47,51 +46,6 @@ IStorageCluster::IStorageCluster(
 {
 }
 
-class ReadFromCluster : public SourceStepWithFilter
-{
-public:
-    std::string getName() const override { return "ReadFromCluster"; }
-    void initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &) override;
-    void applyFilters(ActionDAGNodes added_filter_nodes) override;
-
-    ReadFromCluster(
-        const Names & column_names_,
-        const SelectQueryInfo & query_info_,
-        const StorageSnapshotPtr & storage_snapshot_,
-        const ContextPtr & context_,
-        Block sample_block,
-        std::shared_ptr<IStorageCluster> storage_,
-        ASTPtr query_to_send_,
-        QueryProcessingStage::Enum processed_stage_,
-        ClusterPtr cluster_,
-        LoggerPtr log_)
-        : SourceStepWithFilter(
-            std::move(sample_block),
-            column_names_,
-            query_info_,
-            storage_snapshot_,
-            context_)
-        , storage(std::move(storage_))
-        , query_to_send(std::move(query_to_send_))
-        , processed_stage(processed_stage_)
-        , cluster(std::move(cluster_))
-        , log(log_)
-    {
-    }
-
-private:
-    std::shared_ptr<IStorageCluster> storage;
-    ASTPtr query_to_send;
-    QueryProcessingStage::Enum processed_stage;
-    ClusterPtr cluster;
-    LoggerPtr log;
-
-    std::optional<RemoteQueryExecutor::Extension> extension;
-
-    void createExtension(const ActionsDAG::Node * predicate);
-    ContextPtr updateSettings(const Settings & settings);
-};
-
 void ReadFromCluster::applyFilters(ActionDAGNodes added_filter_nodes)
 {
     SourceStepWithFilter::applyFilters(std::move(added_filter_nodes));

src/Storages/IStorageCluster.h

@@ -4,6 +4,7 @@
 #include <Interpreters/Cluster.h>
 #include <QueryPipeline/RemoteQueryExecutor.h>
 #include <Parsers/ASTExpressionList.h>
+#include <Processors/QueryPlan/SourceStepWithFilter.h>
 
 namespace DB
 {
@@ -52,4 +53,50 @@ class IStorageCluster : public IStorage
 };
 
 
+class ReadFromCluster : public SourceStepWithFilter
+{
+public:
+    std::string getName() const override { return "ReadFromCluster"; }
+    void initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &) override;
+    void applyFilters(ActionDAGNodes added_filter_nodes) override;
+
+    ReadFromCluster(
+        const Names & column_names_,
+        const SelectQueryInfo & query_info_,
+        const StorageSnapshotPtr & storage_snapshot_,
+        const ContextPtr & context_,
+        Block sample_block,
+        std::shared_ptr<IStorageCluster> storage_,
+        ASTPtr query_to_send_,
+        QueryProcessingStage::Enum processed_stage_,
+        ClusterPtr cluster_,
+        LoggerPtr log_)
+        : SourceStepWithFilter(
+            std::move(sample_block),
+            column_names_,
+            query_info_,
+            storage_snapshot_,
+            context_)
+        , storage(std::move(storage_))
+        , query_to_send(std::move(query_to_send_))
+        , processed_stage(processed_stage_)
+        , cluster(std::move(cluster_))
+        , log(log_)
+    {
+    }
+
+private:
+    std::shared_ptr<IStorageCluster> storage;
+    ASTPtr query_to_send;
+    QueryProcessingStage::Enum processed_stage;
+    ClusterPtr cluster;
+    LoggerPtr log;
+
+    std::optional<RemoteQueryExecutor::Extension> extension;
+
+    void createExtension(const ActionsDAG::Node * predicate);
+    ContextPtr updateSettings(const Settings & settings);
+};
+
+
 }

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -124,7 +124,7 @@ RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExten
 {
     auto iterator = StorageObjectStorageSource::createFileIterator(
         configuration, configuration->getQuerySettings(local_context), object_storage, /* distributed_processing */false,
-        local_context, predicate, virtual_columns, nullptr, local_context->getFileProgressCallback());
+        local_context, predicate, getVirtualsList(), nullptr, local_context->getFileProgressCallback());
 
     auto callback = std::make_shared<std::function<String()>>([iterator]() mutable -> String
     {

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -40,7 +40,6 @@ class StorageObjectStorageCluster : public IStorageCluster
     const String engine_name;
     const StorageObjectStorage::ConfigurationPtr configuration;
     const ObjectStoragePtr object_storage;
-    NamesAndTypesList virtual_columns;
     bool cluster_name_in_settings;
 };