From 44381a6da7a4bac479ed04c683827e14661846a1 Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 00:17:15 +0800 Subject: [PATCH 01/10] feat: support expression in limit clause --- src/frontend/src/binder/query.rs | 10 +++--- src/frontend/src/planner/query.rs | 43 +++++++++++++++++++++++-- src/sqlparser/src/ast/query.rs | 4 +-- src/sqlparser/src/parser.rs | 12 +++---- src/tests/sqlsmith/src/sql_gen/query.rs | 8 +++-- 5 files changed, 58 insertions(+), 19 deletions(-) diff --git a/src/frontend/src/binder/query.rs b/src/frontend/src/binder/query.rs index f227b77c9df9..c70bd3a49c69 100644 --- a/src/frontend/src/binder/query.rs +++ b/src/frontend/src/binder/query.rs @@ -38,7 +38,7 @@ use crate::expr::{CorrelatedId, Depth, ExprImpl, ExprRewriter}; pub struct BoundQuery { pub body: BoundSetExpr, pub order: Vec, - pub limit: Option, + pub limit: Option, pub offset: Option, pub with_ties: bool, pub extra_order_exprs: Vec, @@ -174,13 +174,15 @@ impl Binder { ) => { with_ties = fetch_with_ties; match quantity { - Some(v) => Some(parse_non_negative_i64("LIMIT", &v)? as u64), - None => Some(1), + Some(v) => Some(v), + None => Some(Expr::Value(Value::Number("1".to_string()))), } } - (Some(limit), None) => Some(parse_non_negative_i64("LIMIT", &limit)? as u64), + (Some(limit), None) => Some(limit), (Some(_), Some(_)) => unreachable!(), // parse error }; + let limit = limit.map(|expr| self.bind_expr(expr)).transpose()?; + let offset = offset .map(|s| parse_non_negative_i64("OFFSET", &s)) .transpose()? diff --git a/src/frontend/src/planner/query.rs b/src/frontend/src/planner/query.rs index bbf152c36fdf..db116a0b0711 100644 --- a/src/frontend/src/planner/query.rs +++ b/src/frontend/src/planner/query.rs @@ -13,9 +13,11 @@ // limitations under the License. use fixedbitset::FixedBitSet; +use risingwave_common::types::DataType; use crate::binder::BoundQuery; -use crate::error::Result; +use crate::error::{ErrorCode, Result, RwError}; +use crate::expr::ExprImpl; use crate::optimizer::plan_node::{LogicalLimit, LogicalTopN}; use crate::optimizer::property::{Order, RequiredDist}; use crate::optimizer::PlanRoot; @@ -50,7 +52,44 @@ impl Planner { let func_dep = plan.functional_dependency(); order = func_dep.minimize_order_key(order, &[]); - let limit = limit.unwrap_or(LIMIT_ALL_COUNT); + let limit = limit.unwrap_or(ExprImpl::literal_bigint(LIMIT_ALL_COUNT as i64)); + if !limit.is_const() { + return Err(ErrorCode::ExprError( + format!( + "expects an integer or expression after LIMIT, but found:{:?}", + limit + ) + .into(), + ) + .into()); + } + let limit_original = limit.clone(); + let limit_err = ErrorCode::ExprError( + format!( + "expects an integer or expression after LIMIT, but found:{:?}", + limit_original + ) + .into(), + ); + let limit_cast_to_bigint = limit.cast_explicit(DataType::Int64).map_err(|_| { + RwError::from(ErrorCode::ExprError( + format!( + "expects an integer or expression after LIMIT, but found:{:?}", + limit_original + ) + .into(), + )) + })?; + let limit = match limit_cast_to_bigint.fold_const() { + Ok(datum) => match datum { + Some(datum) => datum.as_integral() as u64, + None => { + return Err(limit_err.into()); + } + }, + _ => return Err(limit_err.into()), + }; + let offset = offset.unwrap_or_default(); plan = if order.column_orders.is_empty() { // Should be rejected by parser. diff --git a/src/sqlparser/src/ast/query.rs b/src/sqlparser/src/ast/query.rs index 428fe4e4c541..310b9ceb882d 100644 --- a/src/sqlparser/src/ast/query.rs +++ b/src/sqlparser/src/ast/query.rs @@ -30,7 +30,7 @@ pub struct Query { /// ORDER BY pub order_by: Vec, /// `LIMIT { | ALL }` - pub limit: Option, + pub limit: Option, /// `OFFSET [ { ROW | ROWS } ]` /// /// `ROW` and `ROWS` are noise words that don't influence the effect of the clause. @@ -655,7 +655,7 @@ impl fmt::Display for OrderByExpr { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct Fetch { pub with_ties: bool, - pub quantity: Option, + pub quantity: Option, } impl fmt::Display for Fetch { diff --git a/src/sqlparser/src/parser.rs b/src/sqlparser/src/parser.rs index f9ee9f727543..6fde45db9148 100644 --- a/src/sqlparser/src/parser.rs +++ b/src/sqlparser/src/parser.rs @@ -5333,16 +5333,12 @@ impl Parser<'_> { } /// Parse a LIMIT clause - pub fn parse_limit(&mut self) -> PResult> { + pub fn parse_limit(&mut self) -> PResult> { if self.parse_keyword(Keyword::ALL) { Ok(None) } else { - let number = self.parse_number_value()?; - // TODO(Kexiang): support LIMIT expr - if self.consume_token(&Token::DoubleColon) { - self.expect_keyword(Keyword::BIGINT)? - } - Ok(Some(number)) + let expr = self.parse_expr()?; + Ok(Some(expr)) } } @@ -5366,7 +5362,7 @@ impl Parser<'_> { { None } else { - let quantity = self.parse_number_value()?; + let quantity = self.parse_expr()?; self.expect_one_of_keywords(&[Keyword::ROW, Keyword::ROWS])?; Some(quantity) }; diff --git a/src/tests/sqlsmith/src/sql_gen/query.rs b/src/tests/sqlsmith/src/sql_gen/query.rs index cfbfdb70cbf8..4e43204348d9 100644 --- a/src/tests/sqlsmith/src/sql_gen/query.rs +++ b/src/tests/sqlsmith/src/sql_gen/query.rs @@ -23,7 +23,7 @@ use rand::prelude::SliceRandom; use rand::Rng; use risingwave_common::types::DataType; use risingwave_sqlparser::ast::{ - Cte, Distinct, Expr, Ident, Query, Select, SelectItem, SetExpr, TableWithJoins, With, + Cte, Distinct, Expr, Ident, Query, Select, SelectItem, SetExpr, TableWithJoins, Value, With, }; use crate::sql_gen::utils::create_table_with_joins_from_table; @@ -160,10 +160,12 @@ impl SqlGenerator<'_, R> { } } - fn gen_limit(&mut self, has_order_by: bool) -> Option { + fn gen_limit(&mut self, has_order_by: bool) -> Option { if (!self.is_mview || has_order_by) && self.flip_coin() { let start = if self.is_mview { 1 } else { 0 }; - Some(self.rng.gen_range(start..=100).to_string()) + Some(Expr::Value(Value::Number( + self.rng.gen_range(start..=100).to_string(), + ))) } else { None } From 60266e501dac6ada53d74c8b88a6c71268143f27 Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 00:52:29 +0800 Subject: [PATCH 02/10] fix test --- src/sqlparser/tests/sqlparser_common.rs | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/src/sqlparser/tests/sqlparser_common.rs b/src/sqlparser/tests/sqlparser_common.rs index 496a34b34757..8f0b308e4032 100644 --- a/src/sqlparser/tests/sqlparser_common.rs +++ b/src/sqlparser/tests/sqlparser_common.rs @@ -206,17 +206,20 @@ fn parse_simple_select() { assert_eq!(select.distinct, Distinct::All); assert_eq!(3, select.projection.len()); let select = verified_query(sql); - assert_eq!(Some("5".to_owned()), select.limit); + assert_eq!( + Some(Expr::Value(Value::Number("5".to_string()))), + select.limit + ); } #[test] fn parse_limit_is_not_an_alias() { // In dialects supporting LIMIT it shouldn't be parsed as a table alias let ast = verified_query("SELECT id FROM customer LIMIT 1"); - assert_eq!(Some("1".to_owned()), ast.limit); + assert_eq!(Some(Expr::Value(Value::Number("1".to_string()))), ast.limit); let ast = verified_query("SELECT 1 LIMIT 5"); - assert_eq!(Some("5".to_owned()), ast.limit); + assert_eq!(Some(Expr::Value(Value::Number("5".to_string()))), ast.limit); } #[test] @@ -1068,7 +1071,10 @@ fn parse_select_order_by_limit() { ], select.order_by ); - assert_eq!(Some("2".to_owned()), select.limit); + assert_eq!( + Some(Expr::Value(Value::Number("2".to_string()))), + select.limit + ); } #[test] @@ -1091,7 +1097,10 @@ fn parse_select_order_by_nulls_order() { ], select.order_by ); - assert_eq!(Some("2".to_owned()), select.limit); + assert_eq!( + Some(Expr::Value(Value::Number("2".to_string()))), + select.limit + ); } #[test] @@ -3542,7 +3551,7 @@ fn parse_fetch() { let fetch_first_two_rows_only = Some(Fetch { with_ties: false, - quantity: Some("2".to_owned()), + quantity: Some(Expr::Value(Value::Number("2".to_string()))), }); let ast = verified_query("SELECT foo FROM bar FETCH FIRST 2 ROWS ONLY"); assert_eq!(ast.fetch, fetch_first_two_rows_only); @@ -3567,7 +3576,7 @@ fn parse_fetch() { ast.fetch, Some(Fetch { with_ties: true, - quantity: Some("2".to_owned()), + quantity: Some(Expr::Value(Value::Number("2".to_string()))), }) ); let ast = verified_query( From 5b2d41ffb3f78e5f0d2b6d5ad8d00899c6560e52 Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 02:21:25 +0800 Subject: [PATCH 03/10] fix style --- src/frontend/src/binder/query.rs | 2 +- src/frontend/src/planner/query.rs | 7 +------ src/sqlparser/tests/sqlparser_common.rs | 14 +++++++------- 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/src/frontend/src/binder/query.rs b/src/frontend/src/binder/query.rs index c70bd3a49c69..c7a1118662f6 100644 --- a/src/frontend/src/binder/query.rs +++ b/src/frontend/src/binder/query.rs @@ -175,7 +175,7 @@ impl Binder { with_ties = fetch_with_ties; match quantity { Some(v) => Some(v), - None => Some(Expr::Value(Value::Number("1".to_string()))), + None => Some(Expr::Value(Value::Number("1".to_owned()))), } } (Some(limit), None) => Some(limit), diff --git a/src/frontend/src/planner/query.rs b/src/frontend/src/planner/query.rs index db116a0b0711..fc125574bc68 100644 --- a/src/frontend/src/planner/query.rs +++ b/src/frontend/src/planner/query.rs @@ -81,12 +81,7 @@ impl Planner { )) })?; let limit = match limit_cast_to_bigint.fold_const() { - Ok(datum) => match datum { - Some(datum) => datum.as_integral() as u64, - None => { - return Err(limit_err.into()); - } - }, + Ok(Some(datum)) => datum.as_integral() as u64, _ => return Err(limit_err.into()), }; diff --git a/src/sqlparser/tests/sqlparser_common.rs b/src/sqlparser/tests/sqlparser_common.rs index 8f0b308e4032..9ce0b620bf96 100644 --- a/src/sqlparser/tests/sqlparser_common.rs +++ b/src/sqlparser/tests/sqlparser_common.rs @@ -207,7 +207,7 @@ fn parse_simple_select() { assert_eq!(3, select.projection.len()); let select = verified_query(sql); assert_eq!( - Some(Expr::Value(Value::Number("5".to_string()))), + Some(Expr::Value(Value::Number("5".to_owned()))), select.limit ); } @@ -216,10 +216,10 @@ fn parse_simple_select() { fn parse_limit_is_not_an_alias() { // In dialects supporting LIMIT it shouldn't be parsed as a table alias let ast = verified_query("SELECT id FROM customer LIMIT 1"); - assert_eq!(Some(Expr::Value(Value::Number("1".to_string()))), ast.limit); + assert_eq!(Some(Expr::Value(Value::Number("1".to_owned()))), ast.limit); let ast = verified_query("SELECT 1 LIMIT 5"); - assert_eq!(Some(Expr::Value(Value::Number("5".to_string()))), ast.limit); + assert_eq!(Some(Expr::Value(Value::Number("5".to_owned()))), ast.limit); } #[test] @@ -1072,7 +1072,7 @@ fn parse_select_order_by_limit() { select.order_by ); assert_eq!( - Some(Expr::Value(Value::Number("2".to_string()))), + Some(Expr::Value(Value::Number("2".to_owned()))), select.limit ); } @@ -1098,7 +1098,7 @@ fn parse_select_order_by_nulls_order() { select.order_by ); assert_eq!( - Some(Expr::Value(Value::Number("2".to_string()))), + Some(Expr::Value(Value::Number("2".to_owned()))), select.limit ); } @@ -3551,7 +3551,7 @@ fn parse_fetch() { let fetch_first_two_rows_only = Some(Fetch { with_ties: false, - quantity: Some(Expr::Value(Value::Number("2".to_string()))), + quantity: Some(Expr::Value(Value::Number("2".to_owned()))), }); let ast = verified_query("SELECT foo FROM bar FETCH FIRST 2 ROWS ONLY"); assert_eq!(ast.fetch, fetch_first_two_rows_only); @@ -3576,7 +3576,7 @@ fn parse_fetch() { ast.fetch, Some(Fetch { with_ties: true, - quantity: Some(Expr::Value(Value::Number("2".to_string()))), + quantity: Some(Expr::Value(Value::Number("2".to_owned()))), }) ); let ast = verified_query( From a32c07892356179b3023ef5724f6b77fedd27583 Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 12:21:08 +0800 Subject: [PATCH 04/10] fix error message and test cases --- src/frontend/src/planner/query.rs | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/src/frontend/src/planner/query.rs b/src/frontend/src/planner/query.rs index fc125574bc68..b6073370044a 100644 --- a/src/frontend/src/planner/query.rs +++ b/src/frontend/src/planner/query.rs @@ -63,26 +63,27 @@ impl Planner { ) .into()); } - let limit_original = limit.clone(); - let limit_err = ErrorCode::ExprError( - format!( - "expects an integer or expression after LIMIT, but found:{:?}", - limit_original - ) - .into(), - ); let limit_cast_to_bigint = limit.cast_explicit(DataType::Int64).map_err(|_| { RwError::from(ErrorCode::ExprError( - format!( - "expects an integer or expression after LIMIT, but found:{:?}", - limit_original - ) + "expects an integer or expression that can be evaluated to an integer after LIMIT" .into(), )) })?; let limit = match limit_cast_to_bigint.fold_const() { - Ok(Some(datum)) => datum.as_integral() as u64, - _ => return Err(limit_err.into()), + Ok(Some(datum)) => { + let value = datum.as_integral(); + if value < 0 { + return Err(ErrorCode::ExprError( + format!("LIMIT must not be negative, but found: {}", value).into(), + ) + .into()); + } + value as u64 + } + _ => return Err(ErrorCode::ExprError( + "expects an integer or expression that can be evaluated to an integer after LIMIT" + .into(), + ).into()), }; let offset = offset.unwrap_or_default(); From c5f4a0b9bab2d45c501922afa660a039310196be Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 15:48:23 +0800 Subject: [PATCH 05/10] revert changes for fetch and add test cases for limit --- e2e_test/batch/order/negative_offset.slt.part | 28 +++++++++++++++++++ src/frontend/src/binder/query.rs | 2 +- src/sqlparser/src/ast/query.rs | 2 +- src/sqlparser/src/parser.rs | 2 +- src/sqlparser/tests/sqlparser_common.rs | 4 +-- 5 files changed, 33 insertions(+), 5 deletions(-) diff --git a/e2e_test/batch/order/negative_offset.slt.part b/e2e_test/batch/order/negative_offset.slt.part index 88c5c47ed15a..38dadc4e1fa1 100644 --- a/e2e_test/batch/order/negative_offset.slt.part +++ b/e2e_test/batch/order/negative_offset.slt.part @@ -15,6 +15,34 @@ SELECT * FROM generate_series(0,10,1) LIMIT -3; statement error SELECT * FROM generate_series(0,10,1) LIMIT -1%; +statement ok +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1+2; +---- +0 +1 +2 + + +statement ok +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2; +--- +0 +1 +2 + + +statement ok +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit -3+4; +--- +0 + +statement ok +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2 OFFSET 2; +--- +2 +3 +4 + statement ok CREATE TABLE integers(k int); diff --git a/src/frontend/src/binder/query.rs b/src/frontend/src/binder/query.rs index c7a1118662f6..ab1de8e625ba 100644 --- a/src/frontend/src/binder/query.rs +++ b/src/frontend/src/binder/query.rs @@ -174,7 +174,7 @@ impl Binder { ) => { with_ties = fetch_with_ties; match quantity { - Some(v) => Some(v), + Some(v) => Some(Expr::Value(Value::Number(v))), None => Some(Expr::Value(Value::Number("1".to_owned()))), } } diff --git a/src/sqlparser/src/ast/query.rs b/src/sqlparser/src/ast/query.rs index 310b9ceb882d..925023dfa046 100644 --- a/src/sqlparser/src/ast/query.rs +++ b/src/sqlparser/src/ast/query.rs @@ -655,7 +655,7 @@ impl fmt::Display for OrderByExpr { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct Fetch { pub with_ties: bool, - pub quantity: Option, + pub quantity: Option, } impl fmt::Display for Fetch { diff --git a/src/sqlparser/src/parser.rs b/src/sqlparser/src/parser.rs index 6fde45db9148..7f3bf98f9d24 100644 --- a/src/sqlparser/src/parser.rs +++ b/src/sqlparser/src/parser.rs @@ -5362,7 +5362,7 @@ impl Parser<'_> { { None } else { - let quantity = self.parse_expr()?; + let quantity = self.parse_number_value()?; self.expect_one_of_keywords(&[Keyword::ROW, Keyword::ROWS])?; Some(quantity) }; diff --git a/src/sqlparser/tests/sqlparser_common.rs b/src/sqlparser/tests/sqlparser_common.rs index 9ce0b620bf96..9018a4f00510 100644 --- a/src/sqlparser/tests/sqlparser_common.rs +++ b/src/sqlparser/tests/sqlparser_common.rs @@ -3551,7 +3551,7 @@ fn parse_fetch() { let fetch_first_two_rows_only = Some(Fetch { with_ties: false, - quantity: Some(Expr::Value(Value::Number("2".to_owned()))), + quantity: Some("2".to_owned()), }); let ast = verified_query("SELECT foo FROM bar FETCH FIRST 2 ROWS ONLY"); assert_eq!(ast.fetch, fetch_first_two_rows_only); @@ -3576,7 +3576,7 @@ fn parse_fetch() { ast.fetch, Some(Fetch { with_ties: true, - quantity: Some(Expr::Value(Value::Number("2".to_owned()))), + quantity: Some("2".to_owned()), }) ); let ast = verified_query( From 3f29980bf6e539223262b86d20d1a0b58de7c4c6 Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 15:50:24 +0800 Subject: [PATCH 06/10] format --- src/frontend/src/planner/query.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/frontend/src/planner/query.rs b/src/frontend/src/planner/query.rs index b6073370044a..50b61f6403fc 100644 --- a/src/frontend/src/planner/query.rs +++ b/src/frontend/src/planner/query.rs @@ -55,10 +55,7 @@ impl Planner { let limit = limit.unwrap_or(ExprImpl::literal_bigint(LIMIT_ALL_COUNT as i64)); if !limit.is_const() { return Err(ErrorCode::ExprError( - format!( - "expects an integer or expression after LIMIT, but found:{:?}", - limit - ) + "expects an integer or expression that can be evaluated to an integer after LIMIT" .into(), ) .into()); From 36548a554efbf9970cab565adb3a3b4fd886ccea Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 16:22:00 +0800 Subject: [PATCH 07/10] fix test cases --- e2e_test/batch/order/negative_offset.slt.part | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/e2e_test/batch/order/negative_offset.slt.part b/e2e_test/batch/order/negative_offset.slt.part index 38dadc4e1fa1..8991ef1bb023 100644 --- a/e2e_test/batch/order/negative_offset.slt.part +++ b/e2e_test/batch/order/negative_offset.slt.part @@ -15,28 +15,26 @@ SELECT * FROM generate_series(0,10,1) LIMIT -3; statement error SELECT * FROM generate_series(0,10,1) LIMIT -1%; -statement ok +query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1+2; ---- 0 1 2 - -statement ok +query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2; --- 0 1 2 - -statement ok +query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit -3+4; --- 0 -statement ok +query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2 OFFSET 2; --- 2 From 2919bad73e5549b797b9bfef6b31629270108943 Mon Sep 17 00:00:00 2001 From: lmatz Date: Wed, 18 Dec 2024 17:19:08 +0800 Subject: [PATCH 08/10] fix format --- e2e_test/batch/order/negative_offset.slt.part | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/e2e_test/batch/order/negative_offset.slt.part b/e2e_test/batch/order/negative_offset.slt.part index 8991ef1bb023..8acfd257e25c 100644 --- a/e2e_test/batch/order/negative_offset.slt.part +++ b/e2e_test/batch/order/negative_offset.slt.part @@ -24,19 +24,19 @@ SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1+2; query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2; ---- +---- 0 1 2 query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit -3+4; ---- +---- 0 query I SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2 OFFSET 2; ---- +---- 2 3 4 From 9aab24b1a4d89e7dd7e8646263d59c2ffb9dafc9 Mon Sep 17 00:00:00 2001 From: lmatz Date: Sun, 22 Dec 2024 16:45:42 +0800 Subject: [PATCH 09/10] move evaluation to binder and add test cases --- e2e_test/batch/order/negative_offset.slt.part | 26 ------------- e2e_test/batch/order/test_limit.slt.part | 39 +++++++++++++++++++ src/frontend/src/binder/query.rs | 38 ++++++++++++++++-- src/frontend/src/planner/query.rs | 35 +---------------- 4 files changed, 76 insertions(+), 62 deletions(-) diff --git a/e2e_test/batch/order/negative_offset.slt.part b/e2e_test/batch/order/negative_offset.slt.part index 8acfd257e25c..88c5c47ed15a 100644 --- a/e2e_test/batch/order/negative_offset.slt.part +++ b/e2e_test/batch/order/negative_offset.slt.part @@ -15,32 +15,6 @@ SELECT * FROM generate_series(0,10,1) LIMIT -3; statement error SELECT * FROM generate_series(0,10,1) LIMIT -1%; -query I -SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1+2; ----- -0 -1 -2 - -query I -SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2; ----- -0 -1 -2 - -query I -SELECT * FROM generate_series(0,10,1) as t(v) order by v limit -3+4; ----- -0 - -query I -SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2 OFFSET 2; ----- -2 -3 -4 - statement ok CREATE TABLE integers(k int); diff --git a/e2e_test/batch/order/test_limit.slt.part b/e2e_test/batch/order/test_limit.slt.part index bb9db52868a6..6a7a4bc978fd 100644 --- a/e2e_test/batch/order/test_limit.slt.part +++ b/e2e_test/batch/order/test_limit.slt.part @@ -182,6 +182,45 @@ INSERT INTO integers VALUES (1), (2), (3), (4), (5); # 4 # 5 +# expression in the limit clause +query I +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1+2; +---- +0 +1 +2 + +query I +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2; +---- +0 +1 +2 + +query I +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit -3+4; +---- +0 + +query I +SELECT * FROM generate_series(0,10,1) as t(v) order by v limit 1.5*2 OFFSET 2; +---- +2 +3 +4 + +query I +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit NULL + 1; +---- +0 +1 +2 + +query I +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit 1.1::Decimal; +---- +0 + # Subqueries that return negative values statement error diff --git a/src/frontend/src/binder/query.rs b/src/frontend/src/binder/query.rs index ab1de8e625ba..c36d31cbfe1b 100644 --- a/src/frontend/src/binder/query.rs +++ b/src/frontend/src/binder/query.rs @@ -29,7 +29,7 @@ use super::statement::RewriteExprsRecursive; use super::BoundValues; use crate::binder::bind_context::{BindingCte, RecursiveUnion}; use crate::binder::{Binder, BoundSetExpr}; -use crate::error::{ErrorCode, Result}; +use crate::error::{ErrorCode, Result, RwError}; use crate::expr::{CorrelatedId, Depth, ExprImpl, ExprRewriter}; /// A validated sql query, including order and union. @@ -38,7 +38,7 @@ use crate::expr::{CorrelatedId, Depth, ExprImpl, ExprRewriter}; pub struct BoundQuery { pub body: BoundSetExpr, pub order: Vec, - pub limit: Option, + pub limit: Option, pub offset: Option, pub with_ties: bool, pub extra_order_exprs: Vec, @@ -181,7 +181,39 @@ impl Binder { (Some(limit), None) => Some(limit), (Some(_), Some(_)) => unreachable!(), // parse error }; - let limit = limit.map(|expr| self.bind_expr(expr)).transpose()?; + let limit_expr = limit.map(|expr| self.bind_expr(expr)).transpose()?; + let limit = if let Some(limit_expr) = limit_expr { + let limit_cast_to_bigint = limit_expr.cast_assign(DataType::Int64).map_err(|_| { + RwError::from(ErrorCode::ExprError( + "expects an integer or expression that can be evaluated to an integer after LIMIT" + .into(), + )) + })?; + let limit = match limit_cast_to_bigint.try_fold_const() { + Some(Ok(Some(datum))) => { + let value = datum.as_integral(); + if value < 0 { + return Err(ErrorCode::ExprError( + format!("LIMIT must not be negative, but found: {}", value).into(), + ) + .into()); + } + value as u64 + } + // If evaluated to NULL, we follow PG to treat NULL as no limit + Some(Ok(None)) => { + u64::MAX + } + // wrong type, not const, eval error all belongs to this branch + _ => return Err(ErrorCode::ExprError( + "expects an integer or expression that can be evaluated to an integer after LIMIT" + .into(), + ).into()), + }; + Some(limit) + } else { + None + }; let offset = offset .map(|s| parse_non_negative_i64("OFFSET", &s)) diff --git a/src/frontend/src/planner/query.rs b/src/frontend/src/planner/query.rs index 50b61f6403fc..099e31863256 100644 --- a/src/frontend/src/planner/query.rs +++ b/src/frontend/src/planner/query.rs @@ -13,11 +13,9 @@ // limitations under the License. use fixedbitset::FixedBitSet; -use risingwave_common::types::DataType; use crate::binder::BoundQuery; -use crate::error::{ErrorCode, Result, RwError}; -use crate::expr::ExprImpl; +use crate::error::Result; use crate::optimizer::plan_node::{LogicalLimit, LogicalTopN}; use crate::optimizer::property::{Order, RequiredDist}; use crate::optimizer::PlanRoot; @@ -52,36 +50,7 @@ impl Planner { let func_dep = plan.functional_dependency(); order = func_dep.minimize_order_key(order, &[]); - let limit = limit.unwrap_or(ExprImpl::literal_bigint(LIMIT_ALL_COUNT as i64)); - if !limit.is_const() { - return Err(ErrorCode::ExprError( - "expects an integer or expression that can be evaluated to an integer after LIMIT" - .into(), - ) - .into()); - } - let limit_cast_to_bigint = limit.cast_explicit(DataType::Int64).map_err(|_| { - RwError::from(ErrorCode::ExprError( - "expects an integer or expression that can be evaluated to an integer after LIMIT" - .into(), - )) - })?; - let limit = match limit_cast_to_bigint.fold_const() { - Ok(Some(datum)) => { - let value = datum.as_integral(); - if value < 0 { - return Err(ErrorCode::ExprError( - format!("LIMIT must not be negative, but found: {}", value).into(), - ) - .into()); - } - value as u64 - } - _ => return Err(ErrorCode::ExprError( - "expects an integer or expression that can be evaluated to an integer after LIMIT" - .into(), - ).into()), - }; + let limit = limit.unwrap_or(LIMIT_ALL_COUNT); let offset = offset.unwrap_or_default(); plan = if order.column_orders.is_empty() { From a4b261b82c6d13b1f206650757979e3cb26a7ede Mon Sep 17 00:00:00 2001 From: lmatz Date: Sun, 22 Dec 2024 17:05:33 +0800 Subject: [PATCH 10/10] add more test cases --- e2e_test/batch/order/test_limit.slt.part | 33 +++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/e2e_test/batch/order/test_limit.slt.part b/e2e_test/batch/order/test_limit.slt.part index 6a7a4bc978fd..9bcf832215ff 100644 --- a/e2e_test/batch/order/test_limit.slt.part +++ b/e2e_test/batch/order/test_limit.slt.part @@ -215,12 +215,43 @@ SELECT * FROM generate_series(0,3,1) as t(v) order by v limit NULL + 1; 0 1 2 +3 + +query I +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit NULL + '1'::bigint; +---- +0 +1 +2 +3 + +statement error +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit NULL + '1'; + +statement error +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit NULL + '1'::jsonb; query I -SELECT * FROM generate_series(0,3,1) as t(v) order by v limit 1.1::Decimal; +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit 2.1::Decimal; ---- 0 +1 + +statement error +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit '1'::jsonb; + + +query I +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit '2'; +---- +0 +1 +statement error +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit '-2'; + +statement error +SELECT * FROM generate_series(0,3,1) as t(v) order by v limit '2.2'; # Subqueries that return negative values statement error