Merge pull request #53754 from ClickHouse/fix-interpolate-subselect

Fix: require columns participating in interpolate expression
ClickHouse · Sep 4, 2023 · 217bfa0 · 217bfa0
2 parents ec628ee + 407b2e3
commit 217bfa0
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 2 deletions.
diff --git a/src/Interpreters/ExpressionAnalyzer.cpp b/src/Interpreters/ExpressionAnalyzer.cpp
@@ -1513,14 +1513,16 @@ ActionsDAGPtr SelectQueryExpressionAnalyzer::appendOrderBy(ExpressionActionsChai
         for (const auto & child : select_query->select()->children)
             select.insert(child->getAliasOrColumnName());
 
+        NameSet required_by_interpolate;
         /// collect columns required for interpolate expressions -
         /// interpolate expression can use any available column
-        auto find_columns = [&step, &select](IAST * function)
+        auto find_columns = [&step, &select, &required_by_interpolate](IAST * function)
         {
-            auto f_impl = [&step, &select](IAST * fn, auto fi)
+            auto f_impl = [&step, &select, &required_by_interpolate](IAST * fn, auto fi)
             {
                 if (auto * ident = fn->as<ASTIdentifier>())
                 {
+                    required_by_interpolate.insert(ident->getColumnName());
                     /// exclude columns from select expression - they are already available
                     if (!select.contains(ident->getColumnName()))
                         step.addRequiredOutput(ident->getColumnName());
@@ -1536,6 +1538,14 @@ ActionsDAGPtr SelectQueryExpressionAnalyzer::appendOrderBy(ExpressionActionsChai
 
         for (const auto & interpolate : interpolate_list->children)
             find_columns(interpolate->as<ASTInterpolateElement>()->expr.get());
+
+        if (!required_result_columns.empty())
+        {
+            NameSet required_result_columns_set(required_result_columns.begin(), required_result_columns.end());
+            for (const auto & name : required_by_interpolate)
+                if (!required_result_columns_set.contains(name))
+                    required_result_columns.push_back(name);
+        }
     }
 
     if (optimize_read_in_order)

diff --git a/src/Interpreters/TreeRewriter.cpp b/src/Interpreters/TreeRewriter.cpp
@@ -388,6 +388,44 @@ void removeUnneededColumnsFromSelectClause(ASTSelectQuery * select_query, const
     else
         return;
 
+    NameSet required_by_interpolate;
+
+    if (select_query->interpolate())
+    {
+        auto & children = select_query->interpolate()->children;
+        if (!children.empty())
+        {
+            NameToNameSetMap expressions;
+
+            auto interpolate_visitor = [](const ASTPtr ast, NameSet & columns) -> void
+            {
+                auto interpolate_visitor_impl = [](const ASTPtr node, NameSet & cols, auto self) -> void
+                {
+                    if (const auto * ident = node->as<ASTIdentifier>())
+                        cols.insert(ident->name());
+                    else if (const auto * func = node->as<ASTFunction>())
+                        for (const auto & elem : func->arguments->children)
+                            self(elem, cols, self);
+                };
+                interpolate_visitor_impl(ast, columns, interpolate_visitor_impl);
+            };
+
+            for (const auto & elem : children)
+            {
+                if (auto * interpolate = elem->as<ASTInterpolateElement>())
+                {
+                    NameSet needed_columns;
+                    interpolate_visitor(interpolate->expr, needed_columns);
+                    expressions.emplace(interpolate->column, std::move(needed_columns));
+                }
+            }
+
+            for (const auto & name : required_result_columns)
+                if (const auto it = expressions.find(name); it != expressions.end())
+                    required_by_interpolate.insert(it->second.begin(), it->second.end());
+        }
+    }
+
     ASTs new_elements;
     new_elements.reserve(elements.size());
 
@@ -403,6 +441,11 @@ void removeUnneededColumnsFromSelectClause(ASTSelectQuery * select_query, const
             new_elements.push_back(elem);
             --it->second;
         }
+        else if (required_by_interpolate.contains(name))
+        {
+            /// Columns required by interpolate expression are not always in the required_result_columns
+            new_elements.push_back(elem);
+        }
         else if (select_query->distinct || hasArrayJoin(elem))
         {
             /// ARRAY JOIN cannot be optimized out since it may change number of rows,

diff --git a/tests/queries/0_stateless/02863_interpolate_subquery.reference b/tests/queries/0_stateless/02863_interpolate_subquery.reference
diff --git a/tests/queries/0_stateless/02863_interpolate_subquery.sql b/tests/queries/0_stateless/02863_interpolate_subquery.sql
@@ -0,0 +1,7 @@
+-- https://github.com/ClickHouse/ClickHouse/issues/53640
+DROP TABLE IF EXISTS tab;
+CREATE TABLE tab (i UInt32, a UInt32) ENGINE=Memory;
+SELECT i, col1 FROM (
+    SELECT i, a AS col1, a AS col2 FROM tab ORDER BY i WITH FILL INTERPOLATE (col1 AS col1+col2, col2)
+);
+DROP TABLE tab;