move evaluation to binder and add test cases

risingwavelabs · Dec 22, 2024 · 9aab24b · 9aab24b
1 parent 2919bad
commit 9aab24b
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 62 deletions.
diff --git 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
@@ -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
@@ -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<ColumnOrder>,
-    pub limit: Option<ExprImpl>,
+    pub limit: Option<u64>,
     pub offset: Option<u64>,
     pub with_ties: bool,
     pub extra_order_exprs: Vec<ExprImpl>,
@@ -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
@@ -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() {