Stable 24.8 Backport of #123436 - Fix ALIAS row policies on legacy distributed reads

src/Interpreters/InterpreterSelectQuery.cpp

@@ -151,10 +151,13 @@ FilterDAGInfoPtr generateFilterActions(
     select_ast->setExpression(ASTSelectQuery::Expression::SELECT, std::make_shared<ASTExpressionList>());
     auto expr_list = select_ast->select();
 
-    /// The first column is our filter expression.
-    /// the row_policy_filter_expression should be cloned, because it may be changed by TreeRewriter.
-    /// which make it possible an invalid expression, although it may be valid in whole select.
-    expr_list->children.push_back(row_policy_filter_expression->clone());
+    /// The first column is our filter expression. Clone it because `TreeRewriter` can change the AST.
+    auto filter_expression = row_policy_filter_expression->clone();
+
+    /// `TreeRewriter` expands table aliases only on the initiator, but these filters are also
+    /// created on shards. An `ALIAS` used by a row policy must be evaluated before the read filter.
+    replaceAliasColumnsInQuery(filter_expression, metadata_snapshot->getColumns(), {}, context);
+    expr_list->children.push_back(std::move(filter_expression));
 
     /// Keep columns that are required after the filter actions.
     for (const auto & column_str : prerequisite_columns)
@@ -2127,7 +2130,7 @@ void InterpreterSelectQuery::addPrewhereAliasActions()
     }
 
     /// Set of all (including ALIAS) required columns for PREWHERE
-    auto get_prewhere_columns = [&]()
+    auto get_prewhere_columns = [&](bool include_row_level_filter)
     {
         NameSet columns;
 
@@ -2138,7 +2141,7 @@ void InterpreterSelectQuery::addPrewhereAliasActions()
             columns.insert(prewhere_required_columns.begin(), prewhere_required_columns.end());
         }
 
-        if (row_level_filter)
+        if (row_level_filter && include_row_level_filter)
         {
             auto row_level_required_columns = row_level_filter->actions.getRequiredColumns().getNames();
             columns.insert(row_level_required_columns.begin(), row_level_required_columns.end());
@@ -2156,14 +2159,14 @@ void InterpreterSelectQuery::addPrewhereAliasActions()
     /// before any other executions.
     if (alias_columns_required)
     {
-        NameSet required_columns_from_prewhere = get_prewhere_columns();
+        /// The row-level filter runs before `PREWHERE`, but its inputs are not produced by
+        /// `PREWHERE`. Keep them in the alias actions for queries that still need them.
+        NameSet required_columns_from_prewhere = get_prewhere_columns(/*include_row_level_filter=*/ false);
         NameSet required_aliases_from_prewhere; /// Set of ALIAS required columns for PREWHERE
 
         /// Expression, that contains all raw required columns
         ASTPtr required_columns_all_expr = std::make_shared<ASTExpressionList>();
 
-        /// Expression, that contains raw required columns for PREWHERE
-        ASTPtr required_columns_from_prewhere_expr = std::make_shared<ASTExpressionList>();
 
         /// Sort out already known required columns between expressions,
         /// also populate `required_aliases_from_prewhere`.
@@ -2189,8 +2192,6 @@ void InterpreterSelectQuery::addPrewhereAliasActions()
 
             if (required_columns_from_prewhere.contains(column))
             {
-                required_columns_from_prewhere_expr->children.emplace_back(std::move(column_expr));
-
                 if (is_alias)
                     required_aliases_from_prewhere.insert(column);
             }
@@ -2258,7 +2259,7 @@ void InterpreterSelectQuery::addPrewhereAliasActions()
     const auto & supported_prewhere_columns = storage->supportedPrewhereColumns();
     if (supported_prewhere_columns.has_value())
     {
-        NameSet required_columns_from_prewhere = get_prewhere_columns();
+        NameSet required_columns_from_prewhere = get_prewhere_columns(/*include_row_level_filter=*/ true);
 
         for (const auto & column_name : required_columns_from_prewhere)
         {

tests/queries/0_stateless/05316_row_policy_alias_column_distributed.reference

@@ -0,0 +1,38 @@
+baseline local
+1
+2
+3
+4
+baseline distributed
+1
+2
+3
+4
+policy local
+2
+policy distributed local replica
+2
+policy distributed remote replica
+2
+policy distributed with PREWHERE
+2
+policy local without alias optimization
+2
+policy local with aliases selected
+2	20	ok20
+policy local with aliases selected and PREWHERE
+2	20	ok20
+policy distributed with aliases selected
+2	20	ok20
+policy distributed with aliases selected and PREWHERE
+2	20	ok20
+policy distributed new analyzer
+2
+policy with only an ALIAS predicate
+1	2
+policy with inline expression
+2
+additional filter on ALIAS
+2
+3
+4

tests/queries/0_stateless/05316_row_policy_alias_column_distributed.sql

@@ -0,0 +1,92 @@
+-- Tags: distributed
+
+SET enable_analyzer = 0;
+
+DROP ROW POLICY IF EXISTS rp ON t;
+DROP TABLE IF EXISTS t_dist;
+DROP TABLE IF EXISTS t;
+
+CREATE TABLE t
+(
+    k UInt64,
+    team String,
+    a UInt64 ALIAS k * 10,
+    a2 String ALIAS concat(team, toString(a))
+) ENGINE = MergeTree ORDER BY k;
+
+INSERT INTO t VALUES (1, 'ok'), (2, 'ok'), (3, 'ok'), (4, 'no');
+
+CREATE TABLE t_dist AS t
+ENGINE = Distributed('test_shard_localhost', currentDatabase(), 't', rand());
+
+SELECT 'baseline local';
+SELECT k FROM t ORDER BY k;
+
+SELECT 'baseline distributed';
+SELECT k FROM t_dist ORDER BY k SETTINGS prefer_localhost_replica = 0;
+
+-- 1 fails the `ALIAS` condition; 3 fails the nested `ALIAS` condition;
+-- 4 passes both `ALIAS` conditions but fails the physical team condition.
+CREATE ROW POLICY rp ON t FOR SELECT
+USING team = 'ok' AND a >= 20 AND a2 != 'ok30' TO ALL;
+
+SELECT 'policy local';
+SELECT k FROM t ORDER BY k;
+
+SELECT 'policy distributed local replica';
+SELECT k FROM t_dist ORDER BY k SETTINGS prefer_localhost_replica = 1;
+
+SELECT 'policy distributed remote replica';
+SELECT k FROM t_dist ORDER BY k SETTINGS prefer_localhost_replica = 0;
+
+SELECT 'policy distributed with PREWHERE';
+SELECT k FROM t_dist PREWHERE k >= 1 WHERE k <= 4 ORDER BY k
+SETTINGS prefer_localhost_replica = 0;
+
+SELECT 'policy local without alias optimization';
+SELECT k FROM t ORDER BY k SETTINGS optimize_respect_aliases = 0;
+
+SELECT 'policy local with aliases selected';
+SELECT k, a, a2 FROM t ORDER BY k SETTINGS optimize_respect_aliases = 0;
+
+-- `PREWHERE` uses team, while both the policy and the selected aliases need k.
+-- The alias step must preserve k even when a separate `PREWHERE` step exists.
+SELECT 'policy local with aliases selected and PREWHERE';
+SELECT k, a, a2 FROM t PREWHERE team != 'no' ORDER BY k
+SETTINGS optimize_respect_aliases = 0;
+
+SELECT 'policy distributed with aliases selected';
+SELECT k, a, a2 FROM t_dist ORDER BY k SETTINGS prefer_localhost_replica = 0;
+
+SELECT 'policy distributed with aliases selected and PREWHERE';
+SELECT k, a, a2 FROM t_dist PREWHERE team != 'no' ORDER BY k
+SETTINGS prefer_localhost_replica = 0, optimize_respect_aliases = 0;
+
+SELECT 'policy distributed new analyzer';
+SELECT k FROM t_dist ORDER BY k
+SETTINGS enable_analyzer = 1, prefer_localhost_replica = 0;
+
+-- A policy whose only predicate references an `ALIAS` still needs its value
+-- before the row-level filter, even when the `SELECT` only needs `count`.
+CREATE ROW POLICY OR REPLACE rp ON t FOR SELECT
+USING a2 = 'ok20' TO ALL;
+
+SELECT 'policy with only an ALIAS predicate';
+SELECT count(), sum(k) FROM t_dist SETTINGS prefer_localhost_replica = 0;
+
+-- The equivalent expression is a control for the alias substitution.
+CREATE ROW POLICY OR REPLACE rp ON t FOR SELECT
+USING team = 'ok' AND k * 10 >= 20 AND concat(team, toString(k * 10)) != 'ok30' TO ALL;
+
+SELECT 'policy with inline expression';
+SELECT k FROM t_dist ORDER BY k SETTINGS prefer_localhost_replica = 0;
+
+DROP ROW POLICY rp ON t;
+
+-- Additional table filters use the same filter-action construction path.
+SELECT 'additional filter on ALIAS';
+SELECT k FROM t ORDER BY k
+SETTINGS optimize_respect_aliases = 0, additional_table_filters = {'t': 'a >= 20'};
+
+DROP TABLE t_dist;
+DROP TABLE t;