fix(pipeline): async phases read max_passes from config instead of hardcoding 8

Author: notfilippo

datafusion/core/src/execution/context/mod.rs

@@ -749,10 +749,9 @@ impl SessionContext {
                     Arc::clone(self.state().config().options()),
                 )
                 .without_query_execution_start_time();
-                let plan = self.state().optimizer().optimize(
+                let plan = self.state().optimize_with_config(
                     Arc::unwrap_or_clone(input),
                     &optimizer_context,
-                    |_1, _2| {},
                 )?;
                 self.state
                     .write()

datafusion/core/src/execution/session_state.rs

@@ -61,12 +61,13 @@ datafusion-physical-expr = { workspace = true }
 indexmap = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
+async-trait = { workspace = true }
 recursive = { workspace = true, optional = true }
 regex = { workspace = true }
 regex-syntax = "0.8.9"
 
 [dev-dependencies]
-async-trait = { workspace = true }
+
 criterion = { workspace = true }
 ctor = { workspace = true }
 datafusion-functions-aggregate = { workspace = true }

datafusion/optimizer/Cargo.toml

@@ -20,6 +20,7 @@
 use std::fmt::Debug;
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use log::debug;
 
 use datafusion_common::Result;
@@ -61,6 +62,27 @@ pub trait AnalyzerRule: Debug {
     fn name(&self) -> &str;
 }
 
+/// Like [`AnalyzerRule`], but may perform async operations (e.g. remote catalog
+/// lookups, schema resolution) during analysis.
+///
+/// Async analyzer rules run as a pre-analysis phase in
+/// [`SessionState::create_physical_plan`], before the synchronous [`Analyzer`].
+/// By default no rules are registered, so the phase is a no-op.
+///
+/// [`SessionState::create_physical_plan`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.create_physical_plan
+#[async_trait]
+pub trait AsyncAnalyzerRule: Debug + Send + Sync {
+    /// Rewrite `plan`, possibly performing async I/O.
+    async fn analyze(
+        &self,
+        plan: LogicalPlan,
+        config: &ConfigOptions,
+    ) -> Result<LogicalPlan>;
+
+    /// A human readable name for this rule.
+    fn name(&self) -> &str;
+}
+
 /// Rule-based Analyzer.
 ///
 /// Applies [`FunctionRewrite`]s and [`AnalyzerRule`]s to transform a

datafusion/optimizer/src/analyzer/mod.rs

@@ -75,9 +75,17 @@ pub mod utils;
 #[cfg(test)]
 pub mod test;
 
-pub use analyzer::{Analyzer, AnalyzerRule};
+pub mod logical_pipeline;
+
+pub use analyzer::{Analyzer, AnalyzerRule, AsyncAnalyzerRule};
+pub use logical_pipeline::{
+    AsyncAnalysisPhase, AsyncOptimizationPhase, AsyncPhase, DEFAULT_ANALYSIS_PHASE,
+    DEFAULT_OPTIMIZATION_PHASE, LogicalPlanningPipeline, Phase, Strategy,
+    SyncAnalysisPhase, SyncOptimizationPhase, SyncPhase,
+};
 pub use optimizer::{
-    ApplyOrder, Optimizer, OptimizerConfig, OptimizerContext, OptimizerRule,
+    ApplyOrder, AsyncOptimizerRule, Optimizer, OptimizerConfig, OptimizerContext,
+    OptimizerRule,
 };
 
 pub(crate) mod join_key_set;

datafusion/optimizer/src/lib.rs

@@ -0,0 +1,129 @@
+// 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.
+
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::LogicalPlan;
+
+use crate::analyzer::AsyncAnalyzerRule;
+use crate::optimizer::AsyncOptimizerRule;
+
+use super::Strategy;
+
+/// A named, activatable phase whose rules are applied asynchronously.
+///
+/// Use [`AsyncAnalysisPhase`] for [`AsyncAnalyzerRule`]s and
+/// [`AsyncOptimizationPhase`] for [`AsyncOptimizerRule`]s.
+pub struct AsyncPhase<T: ?Sized> {
+    pub(super) name: String,
+    pub enabled: bool,
+    pub strategy: Strategy,
+    pub rules: Vec<Arc<T>>,
+}
+
+impl<T: ?Sized> Clone for AsyncPhase<T> {
+    fn clone(&self) -> Self {
+        Self {
+            name: self.name.clone(),
+            enabled: self.enabled,
+            strategy: self.strategy.clone(),
+            rules: self.rules.clone(),
+        }
+    }
+}
+
+impl<T: ?Sized + Debug> Debug for AsyncPhase<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AsyncPhase")
+            .field("name", &self.name)
+            .field("enabled", &self.enabled)
+            .field("strategy", &self.strategy)
+            .field("rules", &self.rules)
+            .finish()
+    }
+}
+
+impl<T: ?Sized> AsyncPhase<T> {
+    pub fn new(name: impl Into<String>, strategy: Strategy) -> Self {
+        Self {
+            name: name.into(),
+            enabled: true,
+            strategy,
+            rules: Vec::new(),
+        }
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+}
+
+impl AsyncPhase<dyn AsyncAnalyzerRule + Send + Sync> {
+    pub async fn apply(
+        &self,
+        plan: LogicalPlan,
+        config: &ConfigOptions,
+    ) -> Result<LogicalPlan> {
+        if !self.enabled || self.rules.is_empty() {
+            return Ok(plan);
+        }
+        let passes = match &self.strategy {
+            Strategy::Once => 1,
+            Strategy::FixedPoint { max_passes } => {
+                max_passes.unwrap_or_else(|| config.optimizer.max_passes)
+            }
+        };
+        let mut plan = plan;
+        for _ in 0..passes {
+            for rule in &self.rules {
+                plan = rule.analyze(plan, config).await?;
+            }
+        }
+        Ok(plan)
+    }
+}
+
+impl AsyncPhase<dyn AsyncOptimizerRule + Send + Sync> {
+    pub async fn apply(
+        &self,
+        plan: LogicalPlan,
+        config: &ConfigOptions,
+    ) -> Result<LogicalPlan> {
+        if !self.enabled || self.rules.is_empty() {
+            return Ok(plan);
+        }
+        let passes = match &self.strategy {
+            Strategy::Once => 1,
+            Strategy::FixedPoint { max_passes } => {
+                max_passes.unwrap_or_else(|| config.optimizer.max_passes)
+            }
+        };
+        let mut plan = plan;
+        for _ in 0..passes {
+            for rule in &self.rules {
+                plan = rule.rewrite(plan, config).await?;
+            }
+        }
+        Ok(plan)
+    }
+}
+
+pub type AsyncAnalysisPhase = AsyncPhase<dyn AsyncAnalyzerRule + Send + Sync>;
+pub type AsyncOptimizationPhase = AsyncPhase<dyn AsyncOptimizerRule + Send + Sync>;

datafusion/optimizer/src/logical_pipeline/async_phase.rs

@@ -0,0 +1,98 @@
+// 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.
+
+//! [`LogicalPlanningPipeline`]: an ordered sequence of named [`Phase`]s that
+//! transform a [`LogicalPlan`] before physical planning.
+//!
+//! Inspired by Spark SQL's `Batch(name, strategy, rules)` model.
+
+mod async_phase;
+mod pipeline;
+mod sync_phase;
+
+pub use async_phase::{AsyncAnalysisPhase, AsyncOptimizationPhase, AsyncPhase};
+pub use pipeline::LogicalPlanningPipeline;
+pub use sync_phase::{SyncAnalysisPhase, SyncOptimizationPhase, SyncPhase};
+
+/// Name of the default synchronous analysis phase.
+pub const DEFAULT_ANALYSIS_PHASE: &str = "analysis";
+
+/// Name of the default synchronous optimization phase.
+pub const DEFAULT_OPTIMIZATION_PHASE: &str = "optimization";
+
+/// How many times a phase's rules are applied.
+///
+/// # Choosing a strategy
+///
+/// Use [`Strategy::Once`] when rules are idempotent and a single pass always
+/// suffices (e.g. a post-processing cleanup that must run exactly one time).
+///
+/// Use [`Strategy::FixedPoint`] when rules interact: one rule's output may
+/// unlock another rule's pattern. The phase repeats until the plan signature
+/// stops changing or the pass limit is hit. This is the correct choice for
+/// both the default analysis and optimization phases.
+///
+/// Note: `FixedPoint { max_passes: Some(1) }` is mechanically equivalent to
+/// `Once` (both execute rules exactly one time), but carries different intent.
+/// `Once` means "one pass by design"; `FixedPoint { max_passes: Some(1) }`
+/// means "convergence loop, hard-capped at one iteration."
+#[derive(Clone, Debug)]
+pub enum Strategy {
+    /// Apply each rule exactly once. No convergence check is performed.
+    Once,
+    /// Repeat until the plan stops changing or `max_passes` iterations are
+    /// reached. Convergence is detected via [`LogicalPlanSignature`] so the
+    /// loop exits early as soon as a full pass produces no change.
+    ///
+    /// `max_passes: None` defers to `config.optimizer.max_passes` at runtime.
+    ///
+    /// [`LogicalPlanSignature`]: crate::plan_signature::LogicalPlanSignature
+    FixedPoint { max_passes: Option<usize> },
+}
+
+/// One step in a [`LogicalPlanningPipeline`].
+#[derive(Clone, Debug)]
+pub enum Phase {
+    SyncAnalysis(SyncAnalysisPhase),
+    SyncOptimization(SyncOptimizationPhase),
+    AsyncAnalysis(AsyncAnalysisPhase),
+    AsyncOptimization(AsyncOptimizationPhase),
+}
+
+impl Phase {
+    pub fn name(&self) -> &str {
+        match self {
+            Phase::SyncAnalysis(p) => p.name(),
+            Phase::SyncOptimization(p) => p.name(),
+            Phase::AsyncAnalysis(p) => p.name(),
+            Phase::AsyncOptimization(p) => p.name(),
+        }
+    }
+
+    pub fn is_async(&self) -> bool {
+        matches!(self, Phase::AsyncAnalysis(_) | Phase::AsyncOptimization(_))
+    }
+
+    pub fn is_enabled(&self) -> bool {
+        match self {
+            Phase::SyncAnalysis(p) => p.enabled,
+            Phase::SyncOptimization(p) => p.enabled,
+            Phase::AsyncAnalysis(p) => p.enabled,
+            Phase::AsyncOptimization(p) => p.enabled,
+        }
+    }
+}