From 18a0cffd8e4f6596c7b385100a24fced8dd763e7 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 23 Dec 2024 09:48:05 -0500 Subject: [PATCH] Fix recursive-protection feature flag --- datafusion/common/Cargo.toml | 1 - datafusion/core/Cargo.toml | 7 +++ datafusion/expr/Cargo.toml | 1 - datafusion/optimizer/Cargo.toml | 1 - datafusion/physical-optimizer/Cargo.toml | 1 - datafusion/sql/Cargo.toml | 2 +- datafusion/sql/src/lib.rs | 1 + datafusion/sql/src/query.rs | 10 ++-- datafusion/sql/src/stack.rs | 63 ++++++++++++++++++++++++ 9 files changed, 78 insertions(+), 9 deletions(-) create mode 100644 datafusion/sql/src/stack.rs diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 918f0cd583f7..9253647da12c 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -36,7 +36,6 @@ name = "datafusion_common" path = "src/lib.rs" [features] -default = ["recursive-protection"] avro = ["apache-avro"] backtrace = [] pyarrow = ["pyo3", "arrow/pyarrow", "parquet"] diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index dca40ab3d67a..83b4ab68c89b 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -69,6 +69,13 @@ pyarrow = ["datafusion-common/pyarrow", "parquet"] regex_expressions = [ "datafusion-functions/regex_expressions", ] +recusive-protection = [ + "datafusion-common/recursive-protection", + "datafusion-expr/recursive-protection", + "datafusion-optimizer/recursive-protection", + "datafusion-physical-optimizer/recursive-protection", + "datafusion-sql/recursive-protection" +] serde = ["arrow-schema/serde"] string_expressions = ["datafusion-functions/string_expressions"] unicode_expressions = [ diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 403a80972c3b..191ab1cf6915 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -36,7 +36,6 @@ name = "datafusion_expr" path = "src/lib.rs" [features] -default = ["recursive-protection"] recursive-protection = ["dep:recursive"] [dependencies] diff --git a/datafusion/optimizer/Cargo.toml b/datafusion/optimizer/Cargo.toml index 3032c67682b1..853be74c23e6 100644 --- a/datafusion/optimizer/Cargo.toml +++ b/datafusion/optimizer/Cargo.toml @@ -36,7 +36,6 @@ name = "datafusion_optimizer" path = "src/lib.rs" [features] -default = ["recursive-protection"] recursive-protection = ["dep:recursive"] [dependencies] diff --git a/datafusion/physical-optimizer/Cargo.toml b/datafusion/physical-optimizer/Cargo.toml index c964ca47e6a0..8b55e8f05b26 100644 --- a/datafusion/physical-optimizer/Cargo.toml +++ b/datafusion/physical-optimizer/Cargo.toml @@ -32,7 +32,6 @@ rust-version = { workspace = true } workspace = true [features] -default = ["recursive-protection"] recursive-protection = ["dep:recursive"] [dependencies] diff --git a/datafusion/sql/Cargo.toml b/datafusion/sql/Cargo.toml index c6500e974206..f75d6b6ddfcc 100644 --- a/datafusion/sql/Cargo.toml +++ b/datafusion/sql/Cargo.toml @@ -36,7 +36,7 @@ name = "datafusion_sql" path = "src/lib.rs" [features] -default = ["unicode_expressions", "unparser", "recursive-protection"] +default = ["unicode_expressions", "unparser"] unicode_expressions = [] unparser = [] recursive-protection = ["dep:recursive"] diff --git a/datafusion/sql/src/lib.rs b/datafusion/sql/src/lib.rs index a5d538989453..f20560fe7c90 100644 --- a/datafusion/sql/src/lib.rs +++ b/datafusion/sql/src/lib.rs @@ -43,6 +43,7 @@ mod query; mod relation; mod select; mod set_expr; +mod stack; mod statement; #[cfg(feature = "unparser")] pub mod unparser; diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 2e115d140ea8..9d5a54d90b2c 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; +use crate::stack::StackGuard; use datafusion_common::{not_impl_err, Constraints, DFSchema, Result}; use datafusion_expr::expr::Sort; use datafusion_expr::{ @@ -62,10 +63,11 @@ 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 min_stack_size = recursive::get_minimum_stack_size(); - recursive::set_minimum_stack_size(256 * 1024); - let plan = self.set_expr_to_plan(other, planner_context)?; - recursive::set_minimum_stack_size(min_stack_size); + let plan = { + // scope for dropping _guard + let _guard = StackGuard::new(256 * 1024); + self.set_expr_to_plan(other, planner_context) + }?; let oby_exprs = to_order_by_exprs(query.order_by)?; let order_by_rex = self.order_by_to_sort_expr( oby_exprs, diff --git a/datafusion/sql/src/stack.rs b/datafusion/sql/src/stack.rs new file mode 100644 index 000000000000..57b4aa60ed2d --- /dev/null +++ b/datafusion/sql/src/stack.rs @@ -0,0 +1,63 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// 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. +#[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, + } + + 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); + } + } +} + +/// A stub implementation of the stack guard when the recursive protection +/// feature is not enabled +#[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; + + 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 + } + } +}