diff --git a/Cargo.lock b/Cargo.lock index d3f67f0601371..57df1b848150e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2625,6 +2625,7 @@ dependencies = [ "regex", "rstest", "sqlparser", + "stacker", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index df9cdfed40e00..769e5bd710949 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -195,6 +195,7 @@ rstest = "0.26.1" serde_json = "1" sha2 = "^0.11.0" sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor"] } +stacker = "0.1.24" strum = "0.28.0" strum_macros = "0.28.0" tempfile = "3" diff --git a/datafusion/sql/Cargo.toml b/datafusion/sql/Cargo.toml index cc299ce507099..318f8934639aa 100644 --- a/datafusion/sql/Cargo.toml +++ b/datafusion/sql/Cargo.toml @@ -44,7 +44,7 @@ name = "datafusion_sql" default = ["unicode_expressions", "unparser"] unicode_expressions = [] unparser = [] -recursive_protection = ["dep:recursive"] +recursive_protection = ["dep:recursive", "dep:stacker"] # Note the sql planner should not depend directly on the datafusion-function packages # so that it can be used in a standalone manner with other function implementations. @@ -62,6 +62,7 @@ log = { workspace = true } recursive = { workspace = true, optional = true } regex = { workspace = true } sqlparser = { workspace = true } +stacker = { workspace = true, optional = true } [dev-dependencies] ctor = { workspace = true } diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 76124cbc7eb59..e2b9e4d2d5305 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -19,7 +19,6 @@ use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use crate::stack::StackGuard; use datafusion_common::{Constraints, DFSchema, Result, not_impl_err}; use datafusion_expr::expr::{Sort, WildcardOptions}; @@ -81,11 +80,9 @@ impl SqlToRel<'_, S> { // The functions called from `set_expr_to_plan()` need more than 128KB // stack in debug builds as investigated in: // https://github.com/apache/datafusion/pull/13310#discussion_r1836813902 - let plan = { - // scope for dropping _guard - let _guard = StackGuard::new(256 * 1024); + let plan = crate::stack::maybe_grow(|| { self.set_expr_to_plan(other, planner_context) - }?; + })?; let oby_exprs = to_order_by_exprs(order_by)?; let order_by_rex = self.order_by_to_sort_expr( oby_exprs, diff --git a/datafusion/sql/src/set_expr.rs b/datafusion/sql/src/set_expr.rs index dc8e4f14d1ee8..51b11f3087095 100644 --- a/datafusion/sql/src/set_expr.rs +++ b/datafusion/sql/src/set_expr.rs @@ -25,70 +25,71 @@ use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{SetExpr, SetOperator, SetQuantifier, Spanned}; impl SqlToRel<'_, S> { - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub(super) fn set_expr_to_plan( &self, set_expr: SetExpr, planner_context: &mut PlannerContext, ) -> Result { - let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); - match set_expr { - SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), - SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), - SetExpr::SetOperation { - op, - left, - right, - set_quantifier, - } => { - let left_span = Span::try_from_sqlparser_span(left.span()); - let right_span = Span::try_from_sqlparser_span(right.span()); - let left_plan = self.set_expr_to_plan(*left, planner_context); - // Store the left plan's schema so that the right side can - // alias duplicate expressions to match. Skip for BY NAME - // operations since those match columns by name, not position. - if let Ok(plan) = &left_plan - && plan.schema().fields().len() > 1 - && !matches!( - set_quantifier, - SetQuantifier::ByName - | SetQuantifier::AllByName - | SetQuantifier::DistinctByName - ) - { - planner_context - .set_set_expr_left_schema(Some(Arc::clone(plan.schema()))); - } - let right_plan = self.set_expr_to_plan(*right, planner_context); - planner_context.set_set_expr_left_schema(None); - let (left_plan, right_plan) = match (left_plan, right_plan) { - (Ok(left_plan), Ok(right_plan)) => (left_plan, right_plan), - (Err(left_err), Err(right_err)) => { - return Err(DataFusionError::Collection(vec![ - left_err, right_err, - ])); + crate::stack::maybe_grow(|| { + let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); + match set_expr { + SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), + SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), + SetExpr::SetOperation { + op, + left, + right, + set_quantifier, + } => { + let left_span = Span::try_from_sqlparser_span(left.span()); + let right_span = Span::try_from_sqlparser_span(right.span()); + let left_plan = self.set_expr_to_plan(*left, planner_context); + // Store the left plan's schema so that the right side can + // alias duplicate expressions to match. Skip for BY NAME + // operations since those match columns by name, not position. + if let Ok(plan) = &left_plan + && plan.schema().fields().len() > 1 + && !matches!( + set_quantifier, + SetQuantifier::ByName + | SetQuantifier::AllByName + | SetQuantifier::DistinctByName + ) + { + planner_context + .set_set_expr_left_schema(Some(Arc::clone(plan.schema()))); } - (Err(err), _) | (_, Err(err)) => { - return Err(err); + let right_plan = self.set_expr_to_plan(*right, planner_context); + planner_context.set_set_expr_left_schema(None); + let (left_plan, right_plan) = match (left_plan, right_plan) { + (Ok(left_plan), Ok(right_plan)) => (left_plan, right_plan), + (Err(left_err), Err(right_err)) => { + return Err(DataFusionError::Collection(vec![ + left_err, right_err, + ])); + } + (Err(err), _) | (_, Err(err)) => { + return Err(err); + } + }; + if !(set_quantifier == SetQuantifier::ByName + || set_quantifier == SetQuantifier::AllByName) + { + self.validate_set_expr_num_of_columns( + op, + left_span, + right_span, + &left_plan, + &right_plan, + set_expr_span, + )?; } - }; - if !(set_quantifier == SetQuantifier::ByName - || set_quantifier == SetQuantifier::AllByName) - { - self.validate_set_expr_num_of_columns( - op, - left_span, - right_span, - &left_plan, - &right_plan, - set_expr_span, - )?; + self.set_operation_to_plan(op, left_plan, right_plan, set_quantifier) } - self.set_operation_to_plan(op, left_plan, right_plan, set_quantifier) + SetExpr::Query(q) => self.query_to_plan(*q, planner_context), + _ => not_impl_err!("Query {set_expr} not implemented yet"), } - SetExpr::Query(q) => self.query_to_plan(*q, planner_context), - _ => not_impl_err!("Query {set_expr} not implemented yet"), - } + }) } pub(super) fn is_union_all(set_quantifier: SetQuantifier) -> Result { diff --git a/datafusion/sql/src/stack.rs b/datafusion/sql/src/stack.rs index b7d5eebdd7188..ed3bf1553ebfc 100644 --- a/datafusion/sql/src/stack.rs +++ b/datafusion/sql/src/stack.rs @@ -15,49 +15,43 @@ // specific language governing permissions and limitations // under the License. -pub use inner::StackGuard; - -/// A guard that sets the minimum stack size for the current thread to `min_stack_size` bytes. +/// The local red zone used by SQL recursive entry points. +/// +/// Some SQL planner and unparser recursion paths need more than `recursive`'s +/// default 128 KiB red zone in debug builds. Keep this value local to each +/// stack-growth checkpoint rather than mutating `recursive`'s process-global +/// minimum stack size. #[cfg(feature = "recursive_protection")] -mod inner { - /// Sets the stack size to `min_stack_size` bytes on call to `new()` and - /// resets to the previous value when this structure is dropped. - pub struct StackGuard { - previous_stack_size: usize, - } +pub(crate) const SQL_RECURSION_RED_ZONE: usize = 256 * 1024; - impl StackGuard { - /// Sets the stack size to `min_stack_size` bytes on call to `new()` and - /// resets to the previous value when this structure is dropped. - pub fn new(min_stack_size: usize) -> Self { - let previous_stack_size = recursive::get_minimum_stack_size(); - recursive::set_minimum_stack_size(min_stack_size); - Self { - previous_stack_size, - } - } - } - - impl Drop for StackGuard { - fn drop(&mut self) { - recursive::set_minimum_stack_size(self.previous_stack_size); - } - } +/// Runs `callback` on a stack with enough space for SQL recursive entry points. +#[cfg(feature = "recursive_protection")] +#[inline] +pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { + stacker::maybe_grow( + SQL_RECURSION_RED_ZONE, + recursive::get_stack_allocation_size(), + callback, + ) } -/// A stub implementation of the stack guard when the recursive protection -/// feature is not enabled +/// Runs `callback` without stack growth when recursive protection is disabled. #[cfg(not(feature = "recursive_protection"))] -mod inner { - /// A stub implementation of the stack guard when the recursive protection - /// feature is not enabled that does nothing - pub struct StackGuard; +#[inline] +pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { + callback() +} + +#[cfg(all(test, feature = "recursive_protection"))] +mod tests { + use super::*; + + #[test] + fn maybe_grow_does_not_mutate_recursive_minimum_stack_size() { + let before = recursive::get_minimum_stack_size(); + let observed = maybe_grow(recursive::get_minimum_stack_size); - impl StackGuard { - /// A stub implementation of the stack guard when the recursive protection - /// feature is not enabled - pub fn new(_min_stack_size: usize) -> Self { - Self - } + assert_eq!(observed, before); + assert_eq!(recursive::get_minimum_stack_size(), before); } } diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index f5690ee797598..33457a1515645 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -31,7 +31,6 @@ use std::vec; use super::Unparser; use super::dialect::{DistinctFromStyle, IntervalStyle}; -use crate::stack::StackGuard; use arrow::array::{ ArrayRef, Date32Array, Date64Array, PrimitiveArray, types::{ @@ -100,26 +99,25 @@ impl Unparser<'_> { // default `recursive` red zone, so without raising the minimum stack // size the stack-growing trampoline engages too late and the OS stack // overflows on deeply nested expressions (issue #23056). The size - // mirrors the planner's `StackGuard` usage in `query.rs`. - let _guard = StackGuard::new(256 * 1024); - self.expr_to_sql_with_nesting(expr) + // mirrors the planner's stack-growth usage in `query.rs`. + crate::stack::maybe_grow(|| self.expr_to_sql_with_nesting(expr)) } /// Recursive entry point shared by the public [`Self::expr_to_sql`] and the /// internal recursion sites (scalar-function arguments, arrays, maps, and /// dialect scalar-function overrides). /// - /// This carries the `recursive` annotation so every nesting level becomes a - /// stack-growth checkpoint. Internal recursion must call this rather than - /// the public [`Self::expr_to_sql`]: the public entry point is not - /// annotated and would re-install the [`StackGuard`] on every level. - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] + /// This is a stack-growth checkpoint. Internal recursion must call this + /// rather than the public [`Self::expr_to_sql`]: the public entry point + /// would re-enter the public stack-growth boundary on every level. pub(crate) fn expr_to_sql_with_nesting(&self, expr: &Expr) -> Result { - let mut root_expr = self.expr_to_sql_inner(expr)?; - if self.pretty { - root_expr = self.remove_unnecessary_nesting(root_expr, LOWEST, LOWEST); - } - Ok(root_expr) + crate::stack::maybe_grow(|| { + let mut root_expr = self.expr_to_sql_inner(expr)?; + if self.pretty { + root_expr = self.remove_unnecessary_nesting(root_expr, LOWEST, LOWEST); + } + Ok(root_expr) + }) } fn distinct_from_to_sql( @@ -155,9 +153,8 @@ impl Unparser<'_> { } } - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn expr_to_sql_inner(&self, expr: &Expr) -> Result { - match expr { + crate::stack::maybe_grow(|| match expr { Expr::InList(InList { expr, list, @@ -658,7 +655,7 @@ impl Unparser<'_> { Expr::LambdaVariable(l) => Ok(ast::Expr::Identifier( self.new_ident_quoted_if_needs(l.name.clone()), )), - } + }) } pub fn scalar_function_to_sql( @@ -1016,14 +1013,13 @@ impl Unparser<'_> { /// /// Also note that when fetching the precedence of a nested expression, we ignore other nested /// expressions, so precedence of expr `(a * (b + c))` equals `*` and not `+`. - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn remove_unnecessary_nesting( &self, expr: ast::Expr, left_op: &BinaryOperator, right_op: &BinaryOperator, ) -> ast::Expr { - match expr { + crate::stack::maybe_grow(|| match expr { ast::Expr::Nested(nested) => { let surrounding_precedence = self .sql_op_precedence(left_op) @@ -1076,7 +1072,7 @@ impl Unparser<'_> { self.remove_unnecessary_nesting(*expr, left_op, IS), )), _ => expr, - } + }) } fn inner_precedence(&self, expr: &ast::Expr) -> u8 { @@ -3401,6 +3397,44 @@ mod tests { handle.join().expect("unparsing thread should not panic"); } + #[cfg(feature = "recursive_protection")] + #[test] + fn test_expr_to_sql_does_not_mutate_recursive_minimum_stack_size() -> Result<()> { + const DEFAULT_RECURSIVE_RED_ZONE: usize = 128 * 1024; + + let previous_minimum = recursive::get_minimum_stack_size(); + recursive::set_minimum_stack_size(DEFAULT_RECURSIVE_RED_ZONE); + + let observed_minimum = Arc::new(std::sync::atomic::AtomicUsize::new(usize::MAX)); + let dialect = DuckDBDialect::new().with_custom_scalar_overrides(vec![( + "dummy_udf", + Box::new({ + let observed_minimum = Arc::clone(&observed_minimum); + move |unparser: &Unparser, args: &[Expr]| { + observed_minimum.store( + recursive::get_minimum_stack_size(), + std::sync::atomic::Ordering::Relaxed, + ); + unparser.scalar_function_to_sql("dummy_udf", args).map(Some) + } + }) as ScalarFnToSqlHandler, + )]); + let expr = ScalarUDF::new_from_impl(DummyUDF::new()).call(vec![col("a")]); + + let result = Unparser::new(&dialect).expr_to_sql(&expr); + let final_minimum = recursive::get_minimum_stack_size(); + recursive::set_minimum_stack_size(previous_minimum); + + result?; + assert_eq!( + observed_minimum.load(std::sync::atomic::Ordering::Relaxed), + DEFAULT_RECURSIVE_RED_ZONE + ); + assert_eq!(final_minimum, DEFAULT_RECURSIVE_RED_ZONE); + + Ok(()) + } + #[test] fn test_window_func_support_window_frame() -> Result<()> { let default_dialect: Arc =