From 868e86e9edefd29ee148abbc88a229716d486814 Mon Sep 17 00:00:00 2001
From: flynn <fenglv15@mails.ucas.ac.cn>
Date: Fri, 17 Mar 2023 16:27:24 +0000
Subject: [PATCH] Fix limit offset

---
 .../InterpreterSelectWithUnionQuery.cpp       | 28 +++++++++++--------
 .../01596_setting_limit_offset.reference      | 24 ++++++++++++++++
 .../01596_setting_limit_offset.sql            |  3 ++
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/src/Interpreters/InterpreterSelectWithUnionQuery.cpp b/src/Interpreters/InterpreterSelectWithUnionQuery.cpp
index bfa3d16bf29..025bb916720 100644
--- a/src/Interpreters/InterpreterSelectWithUnionQuery.cpp
+++ b/src/Interpreters/InterpreterSelectWithUnionQuery.cpp
@@ -2,22 +2,23 @@
 
 #include <Columns/getLeastSuperColumn.h>
 #include <Interpreters/Context.h>
+#include <Interpreters/InterpreterSelectIntersectExceptQuery.h>
 #include <Interpreters/InterpreterSelectQuery.h>
 #include <Interpreters/InterpreterSelectWithUnionQuery.h>
-#include <Interpreters/InterpreterSelectIntersectExceptQuery.h>
 #include <Interpreters/QueryLog.h>
+#include <Interpreters/evaluateConstantExpression.h>
+#include <Parsers/ASTSelectIntersectExceptQuery.h>
 #include <Parsers/ASTSelectQuery.h>
 #include <Parsers/ASTSelectWithUnionQuery.h>
-#include <Parsers/ASTSelectIntersectExceptQuery.h>
 #include <Parsers/queryToString.h>
 #include <Processors/QueryPlan/DistinctStep.h>
 #include <Processors/QueryPlan/ExpressionStep.h>
 #include <Processors/QueryPlan/IQueryPlanStep.h>
-#include <Processors/QueryPlan/QueryPlan.h>
-#include <Processors/QueryPlan/UnionStep.h>
 #include <Processors/QueryPlan/LimitStep.h>
 #include <Processors/QueryPlan/OffsetStep.h>
 #include <Processors/QueryPlan/Optimizations/QueryPlanOptimizationSettings.h>
+#include <Processors/QueryPlan/QueryPlan.h>
+#include <Processors/QueryPlan/UnionStep.h>
 #include <QueryPipeline/QueryPipelineBuilder.h>
 #include <Common/typeid_cast.h>
 
@@ -106,32 +107,35 @@ InterpreterSelectWithUnionQuery::InterpreterSelectWithUnionQuery(
             const ASTPtr limit_offset_ast = select_query->limitOffset();
             if (limit_offset_ast)
             {
-                limit_offset = limit_offset_ast->as<ASTLiteral &>().value.safeGet<UInt64>();
+                limit_offset = evaluateConstantExpressionAsLiteral(limit_offset_ast, context)->as<ASTLiteral &>().value.safeGet<UInt64>();
                 UInt64 new_limit_offset = settings.offset + limit_offset;
-                limit_offset_ast->as<ASTLiteral &>().value = Field(new_limit_offset);
+                ASTPtr new_limit_offset_ast = std::make_shared<ASTLiteral>(new_limit_offset);
+                select_query->setExpression(ASTSelectQuery::Expression::LIMIT_OFFSET, std::move(new_limit_offset_ast));
             }
             else if (settings.offset)
             {
-                ASTPtr new_limit_offset_ast = std::make_shared<ASTLiteral>(Field(static_cast<UInt64>(settings.offset)));
+                ASTPtr new_limit_offset_ast = std::make_shared<ASTLiteral>(settings.offset.value);
                 select_query->setExpression(ASTSelectQuery::Expression::LIMIT_OFFSET, std::move(new_limit_offset_ast));
             }
 
             const ASTPtr limit_length_ast = select_query->limitLength();
             if (limit_length_ast)
             {
-                limit_length = limit_length_ast->as<ASTLiteral &>().value.safeGet<UInt64>();
+                limit_length = evaluateConstantExpressionAsLiteral(limit_length_ast, context)->as<ASTLiteral &>().value.safeGet<UInt64>();
 
                 UInt64 new_limit_length = 0;
                 if (settings.offset == 0)
-                    new_limit_length = std::min(limit_length, static_cast<UInt64>(settings.limit));
+                    new_limit_length = std::min(limit_length, settings.limit.value);
                 else if (settings.offset < limit_length)
-                    new_limit_length =  settings.limit ? std::min(static_cast<UInt64>(settings.limit), limit_length - settings.offset) : (limit_length - settings.offset);
+                    new_limit_length = settings.limit ? std::min(settings.limit.value, limit_length - settings.offset.value)
+                                                      : (limit_length - settings.offset.value);
 
-                limit_length_ast->as<ASTLiteral &>().value = Field(new_limit_length);
+                ASTPtr new_limit_length_ast = std::make_shared<ASTLiteral>(new_limit_length);
+                select_query->setExpression(ASTSelectQuery::Expression::LIMIT_LENGTH, std::move(new_limit_length_ast));
             }
             else if (settings.limit)
             {
-                ASTPtr new_limit_length_ast = std::make_shared<ASTLiteral>(Field(static_cast<UInt64>(settings.limit)));
+                ASTPtr new_limit_length_ast = std::make_shared<ASTLiteral>(settings.limit.value);
                 select_query->setExpression(ASTSelectQuery::Expression::LIMIT_LENGTH, std::move(new_limit_length_ast));
             }
 
diff --git a/tests/queries/0_stateless/01596_setting_limit_offset.reference b/tests/queries/0_stateless/01596_setting_limit_offset.reference
index 96483268d43..fe6390b172c 100644
--- a/tests/queries/0_stateless/01596_setting_limit_offset.reference
+++ b/tests/queries/0_stateless/01596_setting_limit_offset.reference
@@ -22,6 +22,10 @@
 107
 108
 109
+102
+103
+104
+105
 105
 106
 107
@@ -38,6 +42,26 @@
 64
 64
 60
+60
+60
+61
+61
+62
+62
+63
+63
+64
+64
+60
+60
+61
+61
+62
+62
+63
+63
+64
+64
 35
 35
 36
diff --git a/tests/queries/0_stateless/01596_setting_limit_offset.sql b/tests/queries/0_stateless/01596_setting_limit_offset.sql
index 3c91e3542bb..0c2ab5fb4dc 100644
--- a/tests/queries/0_stateless/01596_setting_limit_offset.sql
+++ b/tests/queries/0_stateless/01596_setting_limit_offset.sql
@@ -11,6 +11,7 @@ SELECT * FROM test OFFSET 20; -- 5 rows
 SELECT * FROM (SELECT i FROM test LIMIT 10 OFFSET 50) TMP; -- 5 rows
 SELECT * FROM test LIMIT 4 OFFSET 192; -- 4 rows
 SELECT * FROM test LIMIT 10 OFFSET 195; -- 5 rows
+SELECT * FROM test LIMIT 2*2 OFFSET 192;
 
 -- Only set offset
 SET limit = 0;
@@ -21,6 +22,8 @@ SELECT * FROM test LIMIT 100; -- no result
 SET offset = 10;
 SELECT * FROM test LIMIT 20 OFFSET 100; -- 10 rows
 SELECT * FROM test LIMIT 11 OFFSET 100; -- 1 rows
+SELECT * FROM test LIMIT 20 OFFSET 10*10;
+SELECT * FROM test LIMIT 4*5 OFFSET 10*10;
 
 -- offset and limit together
 SET limit = 10;