Rework SessionBuilder

Author: gabotechs

benchmarks/src/tpch/run.rs

@@ -36,14 +36,16 @@ use datafusion::datasource::listing::{
 };
 use datafusion::datasource::{MemTable, TableProvider};
 use datafusion::error::{DataFusionError, Result};
-use datafusion::execution::SessionStateBuilder;
+use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::{collect, displayable};
 use datafusion::prelude::*;
 
 use crate::util::{print_memory_stats, BenchmarkRun, CommonOpt, QueryResult};
 use datafusion_distributed::test_utils::localhost::start_localhost_context;
-use datafusion_distributed::{DistributedPhysicalOptimizerRule, SessionBuilder};
+use datafusion_distributed::{
+    DistributedPhysicalOptimizerRule, DistributedSessionBuilder, DistributedSessionBuilderContext,
+};
 use log::info;
 use structopt::StructOpt;
 
@@ -110,11 +112,13 @@ pub struct RunOpt {
 }
 
 #[async_trait]
-impl SessionBuilder for RunOpt {
-    fn session_state_builder(
+impl DistributedSessionBuilder for RunOpt {
+    async fn build_session_state(
         &self,
-        mut builder: SessionStateBuilder,
-    ) -> Result<SessionStateBuilder, DataFusionError> {
+        _ctx: DistributedSessionBuilderContext,
+    ) -> Result<SessionState, DataFusionError> {
+        let mut builder = SessionStateBuilder::new().with_default_features();
+
         let mut config = self
             .common
             .config()?
@@ -145,17 +149,14 @@ impl SessionBuilder for RunOpt {
             builder = builder.with_physical_optimizer_rule(Arc::new(rule));
         }
 
-        Ok(builder
+        let state = builder
             .with_config(config)
-            .with_runtime_env(rt_builder.build_arc()?))
-    }
+            .with_runtime_env(rt_builder.build_arc()?)
+            .build();
 
-    async fn session_context(
-        &self,
-        ctx: SessionContext,
-    ) -> std::result::Result<SessionContext, DataFusionError> {
+        let ctx = SessionContext::from(state);
         self.register_tables(&ctx).await?;
-        Ok(ctx)
+        Ok(ctx.state())
     }
 }
 

src/common/mod.rs

@@ -1,2 +1,3 @@
+#[allow(unused)]
 pub mod ttl_map;
 pub mod util;

src/config_extension_ext.rs

@@ -27,9 +27,9 @@ pub trait ConfigExtensionExt {
     /// # use async_trait::async_trait;
     /// # use datafusion::common::{extensions_options, DataFusionError};
     /// # use datafusion::config::ConfigExtension;
-    /// # use datafusion::execution::SessionState;
+    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
     /// # use datafusion::prelude::SessionConfig;
-    /// # use datafusion_distributed::{ConfigExtensionExt, SessionBuilder};
+    /// # use datafusion_distributed::{ConfigExtensionExt, DistributedSessionBuilder, DistributedSessionBuilderContext};
     ///
     /// extensions_options! {
     ///     pub struct CustomExtension {
@@ -52,11 +52,13 @@ pub trait ConfigExtensionExt {
     /// struct MyCustomSessionBuilder;
     ///
     /// #[async_trait]
-    /// impl SessionBuilder for MyCustomSessionBuilder {
-    ///     async fn session_state(&self, mut state: SessionState) -> Result<SessionState, DataFusionError> {
+    /// impl DistributedSessionBuilder for MyCustomSessionBuilder {
+    ///     async fn build_session_state(&self, ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
+    ///         let mut state = SessionStateBuilder::new().build();
+    ///
     ///         // while providing this MyCustomSessionBuilder to an Arrow Flight endpoint, it will
     ///         // know how to deserialize the CustomExtension from the gRPC metadata.
-    ///         state.retrieve_distributed_option_extension::<CustomExtension>()?;
+    ///         state.retrieve_distributed_option_extension::<CustomExtension>(&ctx.headers)?;
     ///         Ok(state)
     ///     }
     /// }
@@ -76,9 +78,9 @@ pub trait ConfigExtensionExt {
     /// # use async_trait::async_trait;
     /// # use datafusion::common::{extensions_options, DataFusionError};
     /// # use datafusion::config::ConfigExtension;
-    /// # use datafusion::execution::SessionState;
+    /// # use datafusion::execution::{SessionState, SessionStateBuilder};
     /// # use datafusion::prelude::SessionConfig;
-    /// # use datafusion_distributed::{ConfigExtensionExt, SessionBuilder};
+    /// # use datafusion_distributed::{ConfigExtensionExt, DistributedSessionBuilder, DistributedSessionBuilderContext};
     ///
     /// extensions_options! {
     ///     pub struct CustomExtension {
@@ -101,17 +103,19 @@ pub trait ConfigExtensionExt {
     /// struct MyCustomSessionBuilder;
     ///
     /// #[async_trait]
-    /// impl SessionBuilder for MyCustomSessionBuilder {
-    ///     async fn session_state(&self, mut state: SessionState) -> Result<SessionState, DataFusionError> {
+    /// impl DistributedSessionBuilder for MyCustomSessionBuilder {
+    ///     async fn build_session_state(&self, ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
+    ///         let mut state = SessionStateBuilder::new().build();
     ///         // while providing this MyCustomSessionBuilder to an Arrow Flight endpoint, it will
     ///         // know how to deserialize the CustomExtension from the gRPC metadata.
-    ///         state.retrieve_distributed_option_extension::<CustomExtension>()?;
+    ///         state.retrieve_distributed_option_extension::<CustomExtension>(&ctx.headers)?;
     ///         Ok(state)
     ///     }
     /// }
     /// ```
     fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(
         &mut self,
+        headers: &HeaderMap,
     ) -> Result<(), DataFusionError>;
 }
 
@@ -153,14 +157,11 @@ impl ConfigExtensionExt for SessionConfig {
 
     fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(
         &mut self,
+        headers: &HeaderMap,
     ) -> Result<(), DataFusionError> {
-        let Some(flight_metadata) = self.get_extension::<ContextGrpcMetadata>() else {
-            return Ok(());
-        };
-
         let mut result = T::default();
         let mut found_some = false;
-        for (k, v) in flight_metadata.0.iter() {
+        for (k, v) in headers.iter() {
             let key = k.as_str().trim_start_matches(FLIGHT_METADATA_CONFIG_PREFIX);
             let prefix = format!("{}.", T::PREFIX);
             if key.starts_with(&prefix) {
@@ -185,7 +186,7 @@ impl ConfigExtensionExt for SessionStateBuilder {
     delegate! {
         to self.config().get_or_insert_default() {
             fn add_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T) -> Result<(), DataFusionError>;
-            fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(&mut self) -> Result<(), DataFusionError>;
+            fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
         }
     }
 }
@@ -194,7 +195,7 @@ impl ConfigExtensionExt for SessionState {
     delegate! {
         to self.config_mut() {
             fn add_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T) -> Result<(), DataFusionError>;
-            fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(&mut self) -> Result<(), DataFusionError>;
+            fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
         }
     }
 }
@@ -203,7 +204,7 @@ impl ConfigExtensionExt for SessionContext {
     delegate! {
         to self.state_ref().write().config_mut() {
             fn add_distributed_option_extension<T: ConfigExtension + Default>(&mut self, t: T) -> Result<(), DataFusionError>;
-            fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(&mut self) -> Result<(), DataFusionError>;
+            fn retrieve_distributed_option_extension<T: ConfigExtension + Default>(&mut self, h: &HeaderMap) -> Result<(), DataFusionError>;
         }
     }
 }
@@ -212,17 +213,6 @@ impl ConfigExtensionExt for SessionContext {
 pub(crate) struct ContextGrpcMetadata(pub HeaderMap);
 
 impl ContextGrpcMetadata {
-    pub(crate) fn from_headers(metadata: HeaderMap) -> Self {
-        let mut new = HeaderMap::new();
-        for (k, v) in metadata.into_iter() {
-            let Some(k) = k else { continue };
-            if k.as_str().starts_with(FLIGHT_METADATA_CONFIG_PREFIX) {
-                new.insert(k, v);
-            }
-        }
-        Self(new)
-    }
-
     fn merge(mut self, other: Self) -> Self {
         for (k, v) in other.0.into_iter() {
             let Some(k) = k else { continue };
@@ -252,10 +242,9 @@ mod tests {
         opt.baz = true;
 
         config.add_distributed_option_extension(opt)?;
-
+        let metadata = config.get_extension::<ContextGrpcMetadata>().unwrap();
         let mut new_config = SessionConfig::new();
-        new_config.set_extension(config.get_extension::<ContextGrpcMetadata>().unwrap());
-        new_config.retrieve_distributed_option_extension::<CustomExtension>()?;
+        new_config.retrieve_distributed_option_extension::<CustomExtension>(&metadata.0)?;
 
         let opt = get_ext::<CustomExtension>(&config);
         let new_opt = get_ext::<CustomExtension>(&new_config);
@@ -317,7 +306,7 @@ mod tests {
     fn test_propagate_no_metadata() -> Result<(), Box<dyn std::error::Error>> {
         let mut config = SessionConfig::new();
 
-        config.retrieve_distributed_option_extension::<CustomExtension>()?;
+        config.retrieve_distributed_option_extension::<CustomExtension>(&Default::default())?;
 
         let extension = config.options().extensions.get::<CustomExtension>();
         assert!(extension.is_none());
@@ -330,13 +319,11 @@ mod tests {
         let mut config = SessionConfig::new();
         let mut header_map = HeaderMap::new();
         header_map.insert(
-            HeaderName::from_str("x-datafusion-distributed-other.setting").unwrap(),
+            HeaderName::from_str("x-datafusion-distributed-config-other.setting").unwrap(),
             HeaderValue::from_str("value").unwrap(),
         );
 
-        let flight_metadata = ContextGrpcMetadata::from_headers(header_map);
-        config.set_extension(std::sync::Arc::new(flight_metadata));
-        config.retrieve_distributed_option_extension::<CustomExtension>()?;
+        config.retrieve_distributed_option_extension::<CustomExtension>(&header_map)?;
 
         let extension = config.options().extensions.get::<CustomExtension>();
         assert!(extension.is_none());
@@ -384,9 +371,8 @@ mod tests {
         );
 
         let mut new_config = SessionConfig::new();
-        new_config.set_extension(flight_metadata);
-        new_config.retrieve_distributed_option_extension::<CustomExtension>()?;
-        new_config.retrieve_distributed_option_extension::<AnotherExtension>()?;
+        new_config.retrieve_distributed_option_extension::<CustomExtension>(&metadata)?;
+        new_config.retrieve_distributed_option_extension::<AnotherExtension>(&metadata)?;
 
         let propagated_custom = get_ext::<CustomExtension>(&new_config);
         let propagated_another = get_ext::<AnotherExtension>(&new_config);

src/flight_service/do_get.rs

@@ -3,16 +3,15 @@ use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::datafusion_error_to_tonic_status;
 use crate::flight_service::service::ArrowFlightEndpoint;
+use crate::flight_service::session_builder::DistributedSessionBuilderContext;
 use crate::plan::{DistributedCodec, PartitionGroup};
 use crate::stage::{stage_from_proto, ExecutionStage, ExecutionStageProto};
 use crate::user_provided_codec::get_user_codec;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::Ticket;
-use datafusion::execution::{SessionState, SessionStateBuilder};
-use datafusion::optimizer::OptimizerConfig;
-use datafusion::prelude::SessionConfig;
+use datafusion::execution::SessionState;
 use futures::TryStreamExt;
 use prost::Message;
 use std::sync::Arc;
@@ -90,7 +89,7 @@ impl ArrowFlightEndpoint {
     async fn get_state_and_stage(
         &self,
         doget: DoGet,
-        metadata: MetadataMap,
+        metadata_map: MetadataMap,
     ) -> Result<(SessionState, Arc<ExecutionStage>), Status> {
         let key = doget
             .stage_key
@@ -106,54 +105,34 @@ impl ArrowFlightEndpoint {
                     .stage_proto
                     .ok_or(Status::invalid_argument("DoGet is missing the stage proto"))?;
 
-                let mut config = SessionConfig::default();
-                config.set_extension(Arc::new(ContextGrpcMetadata::from_headers(
-                    metadata.into_headers(),
-                )));
-
-                let state_builder = SessionStateBuilder::new()
-                    .with_runtime_env(Arc::clone(&self.runtime))
-                    .with_config(config)
-                    .with_default_features();
-
-                let state_builder = self
-                    .session_builder
-                    .session_state_builder(state_builder)
-                    .map_err(|err| datafusion_error_to_tonic_status(&err))?;
-
-                let state = state_builder.build();
+                let headers = metadata_map.into_headers();
                 let mut state = self
                     .session_builder
-                    .session_state(state)
+                    .build_session_state(DistributedSessionBuilderContext {
+                        runtime_env: Arc::clone(&self.runtime),
+                        headers: headers.clone(),
+                    })
                     .await
                     .map_err(|err| datafusion_error_to_tonic_status(&err))?;
 
-                let function_registry =
-                    state.function_registry().ok_or(Status::invalid_argument(
-                        "FunctionRegistry not present in newly built SessionState",
-                    ))?;
-
                 let mut combined_codec = ComposedPhysicalExtensionCodec::default();
                 combined_codec.push(DistributedCodec);
                 if let Some(ref user_codec) = get_user_codec(state.config()) {
                     combined_codec.push_arc(Arc::clone(user_codec));
                 }
 
-                let stage = stage_from_proto(
-                    stage_proto,
-                    function_registry,
-                    self.runtime.as_ref(),
-                    &combined_codec,
-                )
-                .map(Arc::new)
-                .map_err(|err| {
-                    Status::invalid_argument(format!("Cannot decode stage proto: {err}"))
-                })?;
+                let stage =
+                    stage_from_proto(stage_proto, &state, self.runtime.as_ref(), &combined_codec)
+                        .map(Arc::new)
+                        .map_err(|err| {
+                            Status::invalid_argument(format!("Cannot decode stage proto: {err}"))
+                        })?;
 
                 // Add the extensions that might be required for ExecutionPlan nodes in the plan
                 let config = state.config_mut();
                 config.set_extension(Arc::clone(&self.channel_manager));
                 config.set_extension(stage.clone());
+                config.set_extension(Arc::new(ContextGrpcMetadata(headers)));
 
                 Ok::<_, Status>((state, stage))
             })

src/flight_service/mod.rs

@@ -5,4 +5,6 @@ mod session_builder;
 pub(crate) use do_get::DoGet;
 
 pub use service::{ArrowFlightEndpoint, StageKey};
-pub use session_builder::{NoopSessionBuilder, SessionBuilder};
+pub use session_builder::{
+    DefaultSessionBuilder, DistributedSessionBuilder, DistributedSessionBuilderContext,
+};

src/flight_service/service.rs

@@ -1,6 +1,6 @@
 use crate::channel_manager::ChannelManager;
-use crate::flight_service::session_builder::NoopSessionBuilder;
-use crate::flight_service::SessionBuilder;
+use crate::flight_service::session_builder::DefaultSessionBuilder;
+use crate::flight_service::DistributedSessionBuilder;
 use crate::stage::ExecutionStage;
 use crate::ChannelResolver;
 use arrow_flight::flight_service_server::FlightService;
@@ -36,7 +36,7 @@ pub struct ArrowFlightEndpoint {
     pub(super) runtime: Arc<RuntimeEnv>,
     #[allow(clippy::type_complexity)]
     pub(super) stages: DashMap<StageKey, Arc<OnceCell<(SessionState, Arc<ExecutionStage>)>>>,
-    pub(super) session_builder: Arc<dyn SessionBuilder + Send + Sync>,
+    pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
 }
 
 impl ArrowFlightEndpoint {
@@ -45,13 +45,13 @@ impl ArrowFlightEndpoint {
             channel_manager: Arc::new(ChannelManager::new(channel_resolver)),
             runtime: Arc::new(RuntimeEnv::default()),
             stages: DashMap::new(),
-            session_builder: Arc::new(NoopSessionBuilder),
+            session_builder: Arc::new(DefaultSessionBuilder),
         }
     }
 
     pub fn with_session_builder(
         &mut self,
-        session_builder: impl SessionBuilder + Send + Sync + 'static,
+        session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
     ) {
         self.session_builder = Arc::new(session_builder);
     }