Backport #62567 to 24.3: Fix scalar subquery in LIMIT

ClickHouse · Apr 12, 2024 · ef62639 · ef62639
1 parent a4a66c4
commit ef62639
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/src/Analyzer/Passes/QueryAnalysisPass.cpp b/src/Analyzer/Passes/QueryAnalysisPass.cpp
@@ -5620,17 +5620,35 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
                 function_name,
                 scope.scope_node->formatASTForErrorMessage());
 
+        bool argument_is_constant = false;
         const auto * constant_node = function_argument->as<ConstantNode>();
         if (constant_node)
         {
             argument_column.column = constant_node->getResultType()->createColumnConst(1, constant_node->getValue());
             argument_column.type = constant_node->getResultType();
+            argument_is_constant = true;
         }
-        else
+        else if (const auto * get_scalar_function_node = function_argument->as<FunctionNode>();
+                get_scalar_function_node && get_scalar_function_node->getFunctionName() == "__getScalar")
         {
-            all_arguments_constants = false;
+            /// Allow constant folding through getScalar
+            const auto * get_scalar_const_arg = get_scalar_function_node->getArguments().getNodes().at(0)->as<ConstantNode>();
+            if (get_scalar_const_arg && scope.context->hasQueryContext())
+            {
+                auto query_context = scope.context->getQueryContext();
+                auto scalar_string = toString(get_scalar_const_arg->getValue());
+                if (query_context->hasScalar(scalar_string))
+                {
+                    auto scalar = query_context->getScalar(scalar_string);
+                    argument_column.column = ColumnConst::create(scalar.getByPosition(0).column, 1);
+                    argument_column.type = get_scalar_function_node->getResultType();
+                    argument_is_constant = true;
+                }
+            }
         }
 
+        all_arguments_constants &= argument_is_constant;
+
         argument_types.push_back(argument_column.type);
         argument_columns.emplace_back(std::move(argument_column));
     }

diff --git a/tests/queries/0_stateless/00834_limit_with_constant_expressions.reference b/tests/queries/0_stateless/00834_limit_with_constant_expressions.reference
@@ -19,3 +19,6 @@
 2
 3
 4
+0
+1
+2
diff --git a/tests/queries/0_stateless/00834_limit_with_constant_expressions.sql b/tests/queries/0_stateless/00834_limit_with_constant_expressions.sql
@@ -24,3 +24,5 @@ SELECT * FROM numbers(10) LIMIT LENGTH('NNN') + COS(0), toDate('0000-00-02'); --
 SELECT * FROM numbers(10) LIMIT a + 5 - a; -- { serverError 47 }
 SELECT * FROM numbers(10) LIMIT a + b; -- { serverError 47 }
 SELECT * FROM numbers(10) LIMIT 'Hello'; -- { serverError 440 }
+
+SELECT number from numbers(10) order by number limit (select sum(number), count() from numbers(3)).1;