ISSUES-12293 allow push predicate when subquery contains with clause (#12663)

src/Core/Settings.h

@@ -323,7 +323,8 @@ struct Settings : public SettingsCollection<Settings>
     M(SettingLogsLevel, send_logs_level, LogsLevel::fatal, "Send server text logs with specified minimum level to client. Valid values: 'trace', 'debug', 'information', 'warning', 'error', 'fatal', 'none'", 0) \
     M(SettingBool, enable_optimize_predicate_expression, 1, "If it is set to true, optimize predicates to subqueries.", 0) \
     M(SettingBool, enable_optimize_predicate_expression_to_final_subquery, 1, "Allow push predicate to final subquery.", 0) \
-    \
+    M(SettingBool, allow_push_predicate_when_subquery_contains_with, 1, "Allows push predicate when subquery contains WITH clause", 0) \
+\
     M(SettingUInt64, low_cardinality_max_dictionary_size, 8192, "Maximum size (in rows) of shared global dictionary for LowCardinality type.", 0) \
     M(SettingBool, low_cardinality_use_single_dictionary_for_part, false, "LowCardinality type serialization setting. If is true, than will use additional keys when global dictionary overflows. Otherwise, will create several shared dictionaries.", 0) \
     M(SettingBool, decimal_check_overflow, true, "Check overflow of decimal arithmetic/comparison operations", 0) \

src/Interpreters/PredicateExpressionsOptimizer.cpp

@@ -21,6 +21,7 @@ PredicateExpressionsOptimizer::PredicateExpressionsOptimizer(
     const Context & context_, const TablesWithColumns & tables_with_columns_, const Settings & settings)
     : enable_optimize_predicate_expression(settings.enable_optimize_predicate_expression)
     , enable_optimize_predicate_expression_to_final_subquery(settings.enable_optimize_predicate_expression_to_final_subquery)
+    , allow_push_predicate_when_subquery_contains_with(settings.allow_push_predicate_when_subquery_contains_with)
     , context(context_)
     , tables_with_columns(tables_with_columns_)
 {
@@ -151,7 +152,8 @@ bool PredicateExpressionsOptimizer::tryRewritePredicatesToTable(ASTPtr & table_e
     if (!table_predicates.empty())
     {
         auto optimize_final = enable_optimize_predicate_expression_to_final_subquery;
-        PredicateRewriteVisitor::Data data(context, table_predicates, std::move(table_columns), optimize_final);
+        auto optimize_with = allow_push_predicate_when_subquery_contains_with;
+        PredicateRewriteVisitor::Data data(context, table_predicates, std::move(table_columns), optimize_final, optimize_with);
 
         PredicateRewriteVisitor(data).visit(table_element);
         return data.is_rewrite;

src/Interpreters/PredicateExpressionsOptimizer.h

@@ -25,6 +25,7 @@ public:
 private:
     const bool enable_optimize_predicate_expression;
     const bool enable_optimize_predicate_expression_to_final_subquery;
+    const bool allow_push_predicate_when_subquery_contains_with;
     const Context & context;
     const TablesWithColumns & tables_with_columns;
 

src/Interpreters/PredicateRewriteVisitor.cpp

@@ -17,8 +17,8 @@ namespace DB
 {
 
 PredicateRewriteVisitorData::PredicateRewriteVisitorData(
-    const Context & context_, const ASTs & predicates_, Names && column_names_, bool optimize_final_)
-    : context(context_), predicates(predicates_), column_names(column_names_), optimize_final(optimize_final_)
+    const Context & context_, const ASTs & predicates_, Names && column_names_, bool optimize_final_, bool optimize_with_)
+    : context(context_), predicates(predicates_), column_names(column_names_), optimize_final(optimize_final_), optimize_with(optimize_with_)
 {
 }
 
@@ -85,7 +85,8 @@ static void cleanAliasAndCollectIdentifiers(ASTPtr & predicate, std::vector<ASTI
 bool PredicateRewriteVisitorData::rewriteSubquery(ASTSelectQuery & subquery, const Names & outer_columns, const Names & inner_columns)
 {
     if ((!optimize_final && subquery.final())
-        || subquery.with() || subquery.withFill()
+        || (!optimize_with && subquery.with())
+        || subquery.withFill()
         || subquery.limitBy() || subquery.limitLength()
         || hasStatefulFunction(subquery.select(), context))
         return false;

src/Interpreters/PredicateRewriteVisitor.h

@@ -24,13 +24,14 @@ public:
         return true;
     }
 
-    PredicateRewriteVisitorData(const Context & context_, const ASTs & predicates_, Names && column_names_, bool optimize_final_);
+    PredicateRewriteVisitorData(const Context & context_, const ASTs & predicates_, Names && column_names_, bool optimize_final_, bool optimize_with_);
 
 private:
     const Context & context;
     const ASTs & predicates;
     const Names column_names;
     bool optimize_final;
+    bool optimize_with;
 
     void visitFirstInternalSelect(ASTSelectQuery & select_query, ASTPtr &);
 

tests/queries/0_stateless/01414_push_predicate_when_contains_with_clause.reference

@@ -0,0 +1,15 @@
+999	1998
+999	1998
+SELECT 
+    number,
+    square_number
+FROM 
+(
+    WITH number * 2 AS square_number
+    SELECT 
+        number,
+        square_number
+    FROM numbers_indexed
+    WHERE number = 999
+) AS squares
+WHERE number = 999

tests/queries/0_stateless/01414_push_predicate_when_contains_with_clause.sql

@@ -0,0 +1,17 @@
+DROP TABLE IF EXISTS numbers_indexed;
+DROP TABLE IF EXISTS squares;
+
+CREATE TABLE numbers_indexed Engine=MergeTree ORDER BY number PARTITION BY bitShiftRight(number,8) SETTINGS index_granularity=8 AS SELECT * FROM numbers(16384);
+
+CREATE VIEW squares AS WITH number*2 AS square_number SELECT number, square_number FROM numbers_indexed;
+
+SET max_rows_to_read=8, read_overflow_mode='throw';
+
+WITH number * 2 AS square_number SELECT number, square_number FROM numbers_indexed WHERE number = 999;
+
+SELECT * FROM squares WHERE number = 999;
+
+EXPLAIN SYNTAX SELECT number, square_number FROM ( WITH number * 2 AS square_number SELECT number, square_number FROM numbers_indexed) AS squares WHERE number = 999;
+
+DROP TABLE IF EXISTS squares;
+DROP TABLE IF EXISTS numbers_indexed;