From b08ffd41850bd6fa533377f3c19d50040216b7d3 Mon Sep 17 00:00:00 2001 From: sakhart Date: Wed, 24 Sep 2025 10:47:46 +0300 Subject: [PATCH 1/3] feat: added delete limit support --- datafusion/sql/src/query.rs | 40 ++++++++++++++++++++------------- datafusion/sql/src/statement.rs | 17 +++++++++----- 2 files changed, 35 insertions(+), 22 deletions(-) diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index a2b348bbed701..fafd968394bc2 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -86,15 +86,7 @@ impl SqlToRel<'_, S> { } let skip = match skip { - 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") - } + Some(skip_expr) => self.get_constant_usize_result(skip_expr.value, input.schema(), "OFFSET"), _ => Ok(0), }?; @@ -102,13 +94,7 @@ impl SqlToRel<'_, S> { Some(limit_expr) if limit_expr != sqlparser::ast::Expr::Value(Value::Null) => { - let expr = self.sql_to_expr( - limit_expr, - input.schema(), - &mut PlannerContext::new(), - )?; - let n = get_constant_result(&expr, "LIMIT")?; - Some(convert_usize_with_check(n, "LIMIT")?) + Some(self.get_constant_usize_result(limit_expr, input.schema(), "LIMIT")?) } _ => None, }; @@ -156,6 +142,28 @@ impl SqlToRel<'_, S> { _ => Ok(plan), } } + + /// Retrieves the constant usize result of an SQL expression, evaluating it if possible. + /// + /// This function takes an SQL expression, a related scheme and an argument name as input and returns + /// a `Result` indicating either the constant result of the expression or an + /// error if the expression cannot be evaluated. + /// + /// # Arguments + /// + /// * `expr` - An `SQLExpr` representing the expression to evaluate. + /// * `schema` - A related DataFusion schema to apply while converting an `expr` into a logical expression. + /// * `arg_name` - The name of the argument for error messages. + /// + /// # Returns + /// + /// * `Result` - An `Ok` variant containing the constant result if evaluation is successful, + /// or an `Err` variant containing an error message if evaluation fails. + pub(super) fn get_constant_usize_result(&self, expr: SQLExpr, schema: &datafusion_common::DFSchema, arg_name: &str) -> Result { + let expr = self.sql_to_expr(expr, schema, &mut PlannerContext::new())?; + let value = get_constant_result(&expr, arg_name)?; + convert_usize_with_check(value, arg_name) + } } /// Retrieves the constant result of an expression, evaluating it if possible. diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 840e31c88232e..1a789aa8f77a9 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -568,12 +568,8 @@ impl SqlToRel<'_, S> { plan_err!("Delete-order-by clause not yet supported")?; } - if limit.is_some() { - plan_err!("Delete-limit clause not yet supported")?; - } - let table_name = self.get_delete_target(from)?; - self.delete_to_plan(table_name, selection) + self.delete_to_plan(table_name, selection, limit) } Statement::StartTransaction { @@ -1223,6 +1219,7 @@ impl SqlToRel<'_, S> { &self, table_name: ObjectName, predicate_expr: Option, + limit: Option ) -> Result { // Do a table lookup to verify the table exists let table_ref = self.object_name_to_table_reference(table_name.clone())?; @@ -1237,7 +1234,7 @@ impl SqlToRel<'_, S> { .build()?; let mut planner_context = PlannerContext::new(); - let source = match predicate_expr { + let mut source = match predicate_expr { None => scan, Some(predicate_expr) => { let filter_expr = @@ -1254,6 +1251,14 @@ impl SqlToRel<'_, S> { } }; + if let Some(limit_expr) = limit { + let limit = (limit_expr != SQLExpr::Value(Value::Null)) + .then(|| self.get_constant_usize_result(limit_expr, source.schema(), "LIMIT")) + .transpose()?; + + source = LogicalPlanBuilder::from(source).limit(0, limit)?.build()? + } + let plan = LogicalPlan::Dml(DmlStatement::new( table_ref, table_source, From 02a2ce3d12e45820c86f853d66dffb91036d14ac Mon Sep 17 00:00:00 2001 From: sakhart Date: Wed, 24 Sep 2025 10:48:31 +0300 Subject: [PATCH 2/3] fix: used workspace versions of arrow and parquet in datafusion-cli --- datafusion-cli/Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 3b164b3ea1df0..e06a64e056691 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -29,7 +29,7 @@ rust-version = "1.82.0" readme = "README.md" [dependencies] -arrow = { version = "54.0.0" } +arrow = { workspace = true } async-trait = "0.1.73" aws-config = "1.5.5" # begin pin aws-sdk crates otherwise CI MSRV check fails @@ -57,7 +57,7 @@ futures = "0.3" mimalloc = { version = "0.1", default-features = false } object_store = { version = "0.11.0", features = ["aws", "gcp", "http"] } parking_lot = { version = "0.12" } -parquet = { version = "54.0.0", default-features = false } +parquet = { workspace = true, default-features = false } regex = "1.8" rustyline = "14.0" tokio = { version = "1.24", features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot", "signal"] } From 36d83e3217cbfe4275ac078bb50b9f5de0f20758 Mon Sep 17 00:00:00 2001 From: sakhart Date: Wed, 24 Sep 2025 11:20:01 +0300 Subject: [PATCH 3/3] test: delete limit usage --- datafusion/sql/tests/sql_integration.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 4bcaba810fa5e..ab3235cb67cef 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -621,6 +621,18 @@ Dml: op=[Delete] table=[person] quick_test(sql, plan); } +#[test] +fn plan_delete_limited() { + let sql = "delete from person limit 2"; + let plan = r#" +Dml: op=[Delete] table=[person] + Limit: skip=0, fetch=2 + TableScan: person + "# + .trim(); + quick_test(sql, plan); +} + #[test] fn select_column_does_not_exist() { let sql = "SELECT doesnotexist FROM person";