feat: support expression in limit clause (#19834)

ci/workflows/pull-request.yml

@@ -607,7 +607,7 @@ steps:
           config: ci/docker-compose.yml
           mount-buildkite-agent: true
       - ./ci/plugins/upload-failure-logs
-    timeout_in_minutes: 20
+    timeout_in_minutes: 22
     retry: *auto-retry
 
   - label: "end-to-end test (deterministic simulation)"

e2e_test/batch/order/test_limit.slt.part

@@ -182,6 +182,127 @@ 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
+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';
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: Failed to bind expression: NULL + '1'
+  2: Bind error: function add(unknown, unknown) is not unique
+HINT:  Could not choose a best candidate function. You might need to add explicit type casts.
+
+
+statement error
+SELECT * FROM generate_series(0,3,1) as t(v) order by v limit NULL + '1'::jsonb;
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: Failed to bind expression: NULL + CAST('1' AS JSONB)
+  2: function add(unknown, jsonb) does not exist
+
+
+
+query I
+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;
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: Expr error
+  2: expects an integer or expression that can be evaluated to an integer after LIMIT
+
+
+
+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';
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: Expr error
+  2: LIMIT must not be negative, but found: -2
+
+
+statement error
+SELECT * FROM generate_series(0,3,1) as t(v) order by v limit '2.2';
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: Expr error
+  2: expects an integer or expression that can be evaluated to an integer after LIMIT,
+but the evaluation of the expression returns error:Expr error: error while evaluating expression `str_parse('2.2')`: Parse error: bigint invalid digit found in string
+
+
+statement error
+select 1 limit generate_series(1, 2);
+----
+db error: ERROR: Failed to run the query
+
+Caused by these errors (recent errors listed first):
+  1: Expr error
+  2: expects an integer or expression that can be evaluated to an integer after LIMIT, but found non-const expression
+
 
 # Subqueries that return negative values
 statement error

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.
@@ -174,13 +174,55 @@ 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(Expr::Value(Value::Number(v))),
+                    None => Some(Expr::Value(Value::Number("1".to_owned()))),
                 }
             }
-            (Some(limit), None) => Some(parse_non_negative_i64("LIMIT", &limit)? as u64),
+            (Some(limit), None) => Some(limit),
             (Some(_), Some(_)) => unreachable!(), // parse error
         };
+        let limit_expr = limit.map(|expr| self.bind_expr(expr)).transpose()?;
+        let limit = if let Some(limit_expr) = limit_expr {
+            // wrong type error is handled here
+            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_int64();
+                    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
+                }
+                // not const error
+                None => return Err(ErrorCode::ExprError(
+                    "expects an integer or expression that can be evaluated to an integer after LIMIT, but found non-const expression"
+                        .into(),
+                ).into()),
+                // eval error
+                Some(Err(e)) => {
+                    return Err(ErrorCode::ExprError(
+                        format!("expects an integer or expression that can be evaluated to an integer after LIMIT,\nbut the evaluation of the expression returns error:{}", e.as_report()
+                        ).into(),
+                    ).into())
+                }
+            };
+            Some(limit)
+        } else {
+            None
+        };
+
         let offset = offset
             .map(|s| parse_non_negative_i64("OFFSET", &s))
             .transpose()?

src/frontend/src/planner/query.rs

@@ -51,6 +51,7 @@ impl Planner {
             order = func_dep.minimize_order_key(order, &[]);
 
             let limit = limit.unwrap_or(LIMIT_ALL_COUNT);
+
             let offset = offset.unwrap_or_default();
             plan = if order.column_orders.is_empty() {
                 // Should be rejected by parser.

src/sqlparser/src/ast/query.rs

@@ -30,7 +30,7 @@ pub struct Query {
     /// ORDER BY
     pub order_by: Vec<OrderByExpr>,
     /// `LIMIT { <N> | ALL }`
-    pub limit: Option<String>,
+    pub limit: Option<Expr>,
     /// `OFFSET <N> [ { ROW | ROWS } ]`
     ///
     /// `ROW` and `ROWS` are noise words that don't influence the effect of the clause.

src/sqlparser/src/parser.rs

@@ -5357,16 +5357,12 @@ impl Parser<'_> {
     }
 
     /// Parse a LIMIT clause
-    pub fn parse_limit(&mut self) -> PResult<Option<String>> {
+    pub fn parse_limit(&mut self) -> PResult<Option<Expr>> {
         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))
         }
     }
 

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_owned()))),
+        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_owned()))), 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_owned()))), 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_owned()))),
+        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_owned()))),
+        select.limit
+    );
 }
 
 #[test]

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<R: Rng> SqlGenerator<'_, R> {
         }
     }
 
-    fn gen_limit(&mut self, has_order_by: bool) -> Option<String> {
+    fn gen_limit(&mut self, has_order_by: bool) -> Option<Expr> {
         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
         }