This is an automated email from the ASF dual-hosted git repository.
akurmustafa pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1b6ae8fcda Add support for constant expression evaluation in limit
(#9790)
1b6ae8fcda is described below
commit 1b6ae8fcdad8324c30e654eaf47eb1ae9ddcd964
Author: Mustafa Akur <[email protected]>
AuthorDate: Wed Mar 27 17:37:02 2024 +0300
Add support for constant expression evaluation in limit (#9790)
* Add support for constant expression evaluation in limit
* Use existing const evaluator
* Revert "Use existing const evaluator"
This reverts commit 99b7a552d6ebf03997afd7d06f48da1e2adc4d94.
* Update datafusion/sql/src/query.rs
Co-authored-by: Andrew Lamb <[email protected]>
* Add negative tests
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/sql/src/query.rs | 83 +++++++++++++++++++--------
datafusion/sqllogictest/test_files/select.slt | 24 +++++++-
2 files changed, 83 insertions(+), 24 deletions(-)
diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs
index ea8edd0771..eda8398c43 100644
--- a/datafusion/sql/src/query.rs
+++ b/datafusion/sql/src/query.rs
@@ -25,6 +25,7 @@ use datafusion_common::{
};
use datafusion_expr::{
CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan,
LogicalPlanBuilder,
+ Operator,
};
use sqlparser::ast::{
Expr as SQLExpr, Offset as SQLOffset, OrderByExpr, Query, SetExpr,
SetOperator,
@@ -221,37 +222,29 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
let skip = match skip {
- Some(skip_expr) => match self.sql_to_expr(
- skip_expr.value,
- input.schema(),
- &mut PlannerContext::new(),
- )? {
- Expr::Literal(ScalarValue::Int64(Some(s))) => {
- if s < 0 {
- return plan_err!("Offset must be >= 0, '{s}' was
provided.");
- }
- Ok(s as usize)
- }
- _ => plan_err!("Unexpected expression in OFFSET clause"),
- }?,
- _ => 0,
- };
+ Some(skip_expr) => {
+ let expr = self.sql_to_expr(
+ skip_expr.value,
+ input.schema(),
+ &mut PlannerContext::new(),
+ )?;
+ let n = get_constant_result(&expr, "OFFSET")?;
+ convert_usize_with_check(n, "OFFSET")
+ }
+ _ => Ok(0),
+ }?;
let fetch = match fetch {
Some(limit_expr)
if limit_expr != sqlparser::ast::Expr::Value(Value::Null) =>
{
- let n = match self.sql_to_expr(
+ let expr = self.sql_to_expr(
limit_expr,
input.schema(),
&mut PlannerContext::new(),
- )? {
- Expr::Literal(ScalarValue::Int64(Some(n))) if n >= 0 => {
- Ok(n as usize)
- }
- _ => plan_err!("LIMIT must not be negative"),
- }?;
- Some(n)
+ )?;
+ let n = get_constant_result(&expr, "LIMIT")?;
+ Some(convert_usize_with_check(n, "LIMIT")?)
}
_ => None,
};
@@ -283,3 +276,47 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
}
}
+
+/// Retrieves the constant result of an expression, evaluating it if possible.
+///
+/// This function takes an expression and an argument name as input and returns
+/// a `Result<i64>` indicating either the constant result of the expression or
an
+/// error if the expression cannot be evaluated.
+///
+/// # Arguments
+///
+/// * `expr` - An `Expr` representing the expression to evaluate.
+/// * `arg_name` - The name of the argument for error messages.
+///
+/// # Returns
+///
+/// * `Result<i64>` - An `Ok` variant containing the constant result if
evaluation is successful,
+/// or an `Err` variant containing an error message if evaluation fails.
+///
+/// <https://github.com/apache/arrow-datafusion/issues/9821> tracks a more
general solution
+fn get_constant_result(expr: &Expr, arg_name: &str) -> Result<i64> {
+ match expr {
+ Expr::Literal(ScalarValue::Int64(Some(s))) => Ok(*s),
+ Expr::BinaryExpr(binary_expr) => {
+ let lhs = get_constant_result(&binary_expr.left, arg_name)?;
+ let rhs = get_constant_result(&binary_expr.right, arg_name)?;
+ let res = match binary_expr.op {
+ Operator::Plus => lhs + rhs,
+ Operator::Minus => lhs - rhs,
+ Operator::Multiply => lhs * rhs,
+ _ => return plan_err!("Unsupported operator for {arg_name}
clause"),
+ };
+ Ok(res)
+ }
+ _ => plan_err!("Unexpected expression in {arg_name} clause"),
+ }
+}
+
+/// Converts an `i64` to `usize`, performing a boundary check.
+fn convert_usize_with_check(n: i64, arg_name: &str) -> Result<usize> {
+ if n < 0 {
+ plan_err!("{arg_name} must be >= 0, '{n}' was provided.")
+ } else {
+ Ok(n as usize)
+ }
+}
diff --git a/datafusion/sqllogictest/test_files/select.slt
b/datafusion/sqllogictest/test_files/select.slt
index 3d3e73e816..3a5c6497eb 100644
--- a/datafusion/sqllogictest/test_files/select.slt
+++ b/datafusion/sqllogictest/test_files/select.slt
@@ -550,9 +550,31 @@ select * from (select 1 a union all select 2) b order by a
limit 1;
1
# select limit clause invalid
-statement error DataFusion error: Error during planning: LIMIT must not be
negative
+statement error DataFusion error: Error during planning: LIMIT must be >= 0,
'\-1' was provided\.
select * from (select 1 a union all select 2) b order by a limit -1;
+# select limit with basic arithmetic
+query I
+select * from (select 1 a union all select 2) b order by a limit 1+1;
+----
+1
+2
+
+# select limit with basic arithmetic
+query I
+select * from (values (1)) LIMIT 10*100;
+----
+1
+
+# More complex expressions in the limit is not supported yet.
+# See issue: https://github.com/apache/arrow-datafusion/issues/9821
+statement error DataFusion error: Error during planning: Unsupported operator
for LIMIT clause
+select * from (values (1)) LIMIT 100/10;
+
+# More complex expressions in the limit is not supported yet.
+statement error DataFusion error: Error during planning: Unexpected expression
in LIMIT clause
+select * from (values (1)) LIMIT cast(column1 as tinyint);
+
# select limit clause
query I
select * from (select 1 a union all select 2) b order by a limit null;
