fix: avoid global SQL stack guard mutation in unparser

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"

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 }

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<S: ContextProvider> 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,

datafusion/sql/src/set_expr.rs

@@ -25,70 +25,71 @@ use datafusion_expr::{LogicalPlan, LogicalPlanBuilder};
 use sqlparser::ast::{SetExpr, SetOperator, SetQuantifier, Spanned};
 
 impl<S: ContextProvider> 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<LogicalPlan> {
-        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<bool> {

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<R>(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<R>(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);
     }
 }