Skip to main content

datafusion_distributed/distributed_planner/
session_state_builder_ext.rs

1use crate::distributed_planner::distributed_query_planner::DistributedQueryPlanner;
2use datafusion::execution::SessionStateBuilder;
3use std::sync::Arc;
4
5/// Extension trait for [SessionStateBuilder].
6pub trait SessionStateBuilderExt {
7    /// Injects a [QueryPlanner] implementation that attempts to distribute the plan after the
8    /// normal planning passes are performed.
9    ///
10    /// It will wrap the existing query planner if one, so while setting up DataFusion's
11    /// [SessionStateBuilder], it's important to inject the custom user query planner implementation
12    /// with [SessionStateBuilderExt::with_distributed_planner] strictly *before* calling
13    /// [SessionStateBuilder::with_query_planner].
14    fn with_distributed_planner(self) -> Self;
15}
16
17impl SessionStateBuilderExt for SessionStateBuilder {
18    fn with_distributed_planner(mut self) -> Self {
19        self.config()
20            .get_or_insert_default()
21            .options_mut()
22            .optimizer
23            .enable_physical_uncorrelated_scalar_subquery = false;
24
25        let prev = std::mem::take(self.query_planner());
26        self.with_query_planner(Arc::new(DistributedQueryPlanner { prev }))
27    }
28}