Move back to with_distributed_channel_resolver and let users inject either DistributedPhysicalOptimizerRule or their own rules

Author: gabotechs

benchmarks/src/tpch/run.rs

@@ -46,7 +46,8 @@ use datafusion_distributed::test_utils::localhost::{
     LocalHostChannelResolver, spawn_flight_service,
 };
 use datafusion_distributed::{
-    DistributedExt, DistributedSessionBuilder, DistributedSessionBuilderContext, NetworkBoundaryExt,
+    DistributedExt, DistributedPhysicalOptimizerRule, DistributedSessionBuilder,
+    DistributedSessionBuilderContext, NetworkBoundaryExt,
 };
 use log::info;
 use std::fs;
@@ -138,7 +139,8 @@ impl DistributedSessionBuilder for RunOpt {
             .with_default_features()
             .with_config(config)
             .with_distributed_user_codec(InMemoryCacheExecCodec)
-            .with_distributed_execution(LocalHostChannelResolver::new(self.workers.clone()))
+            .with_distributed_channel_resolver(LocalHostChannelResolver::new(self.workers.clone()))
+            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
             .with_distributed_option_extension_from_headers::<WarmingUpMarker>(&ctx.headers)?
             .with_distributed_files_per_task(
                 self.files_per_task.unwrap_or(get_available_parallelism()),

examples/in_memory_cluster.rs

@@ -7,11 +7,12 @@ use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use datafusion_distributed::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilderContext, create_flight_client,
+    DistributedPhysicalOptimizerRule, DistributedSessionBuilderContext, create_flight_client,
 };
 use futures::TryStreamExt;
 use hyper_util::rt::TokioIo;
 use std::error::Error;
+use std::sync::Arc;
 use structopt::StructOpt;
 use tonic::transport::{Endpoint, Server};
 
@@ -34,7 +35,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     let state = SessionStateBuilder::new()
         .with_default_features()
-        .with_distributed_execution(InMemoryChannelResolver::new())
+        .with_distributed_channel_resolver(InMemoryChannelResolver::new())
+        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
         .build();
 
     let ctx = SessionContext::from(state);
@@ -97,7 +99,7 @@ impl InMemoryChannelResolver {
                 async move {
                     let builder = SessionStateBuilder::new()
                         .with_default_features()
-                        .with_distributed_execution(this)
+                        .with_distributed_channel_resolver(this)
                         .with_runtime_env(ctx.runtime_env.clone());
                     Ok(builder.build())
                 }

examples/localhost_run.rs

@@ -6,9 +6,12 @@ use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt};
+use datafusion_distributed::{
+    BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedPhysicalOptimizerRule,
+};
 use futures::TryStreamExt;
 use std::error::Error;
+use std::sync::Arc;
 use structopt::StructOpt;
 use tonic::transport::Channel;
 use url::Url;
@@ -38,7 +41,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     let state = SessionStateBuilder::new()
         .with_default_features()
-        .with_distributed_execution(localhost_resolver)
+        .with_distributed_channel_resolver(localhost_resolver)
+        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
         .build();
 
     let ctx = SessionContext::from(state);

examples/localhost_worker.rs

@@ -38,7 +38,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
         async move {
             Ok(SessionStateBuilder::new()
                 .with_runtime_env(ctx.runtime_env)
-                .with_distributed_execution(local_host_resolver)
+                .with_distributed_channel_resolver(local_host_resolver)
                 .with_default_features()
                 .build())
         }

src/distributed_ext.rs

@@ -4,7 +4,7 @@ use crate::config_extension_ext::{
 };
 use crate::distributed_planner::set_distributed_task_estimator;
 use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc};
-use crate::{ChannelResolver, DistributedConfig, DistributedPhysicalOptimizerRule, TaskEstimator};
+use crate::{ChannelResolver, DistributedConfig, TaskEstimator};
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::SessionStateBuilder;
@@ -178,14 +178,8 @@ pub trait DistributedExt: Sized {
     /// Same as [DistributedExt::set_distributed_user_codec] but with a dynamic argument.
     fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
 
-    /// Enables distributed execution. For this, several things happen:
-    ///
-    /// - Injects a [ChannelResolver] implementation for Distributed DataFusion to resolve worker
-    ///   nodes. When running in distributed mode, setting a [ChannelResolver] is required.
-    /// - Injects a [DistributedPhysicalOptimizerRule] rule that will inject network boundaries
-    ///   in the plan and will break it down into stages.
-    /// - Injects a [DistributedConfig] object with configuration about the amount of tasks that
-    ///   should be spawned while distributing the queries.
+    /// Injects a [ChannelResolver] implementation for Distributed DataFusion to resolve worker
+    /// nodes. When running in distributed mode, setting a [ChannelResolver] is required.
     ///
     /// Example:
     ///
@@ -196,7 +190,8 @@ pub trait DistributedExt: Sized {
     /// # use datafusion::execution::{SessionState, SessionStateBuilder};
     /// # use datafusion::prelude::SessionConfig;
     /// # use url::Url;
-    /// # use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedSessionBuilderContext};
+    /// # use std::sync::Arc;
+    /// # use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedPhysicalOptimizerRule, DistributedSessionBuilderContext};
     ///
     /// struct CustomChannelResolver;
     ///
@@ -211,25 +206,29 @@ pub trait DistributedExt: Sized {
     ///     }
     /// }
     ///
+    /// // This tweaks the SessionState so that it can plan for distributed queries and execute them.
     /// let state = SessionStateBuilder::new()
-    ///     .with_distributed_execution(CustomChannelResolver)
+    ///     .with_distributed_channel_resolver(CustomChannelResolver)
+    ///     // the DistributedPhysicalOptimizerRule also needs to be passed so that query plans
+    ///     // get distributed.
+    ///     .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
     ///     .build();
     ///
+    /// // This function can be provided to an ArrowFlightEndpoint so that, upon receiving a distributed
+    /// // part of a plan, it knows how to resolve gRPC channels from URLs for making network calls to other nodes.
     /// async fn build_state(ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
-    ///     // This function can be provided to an ArrowFlightEndpoint so that it knows how to
-    ///     // resolve tonic channels from URLs upon making network calls to other nodes.
     ///     Ok(SessionStateBuilder::new()
-    ///         .with_distributed_execution(CustomChannelResolver)
+    ///         .with_distributed_channel_resolver(CustomChannelResolver)
     ///         .build())
     /// }
     /// ```
-    fn with_distributed_execution<T: ChannelResolver + Send + Sync + 'static>(
+    fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
         self,
         resolver: T,
     ) -> Self;
 
-    /// Same as [DistributedExt::with_distributed_execution] but with an in-place mutation.
-    fn set_distributed_execution<T: ChannelResolver + Send + Sync + 'static>(
+    /// Same as [DistributedExt::with_distributed_channel_resolver] but with an in-place mutation.
+    fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
         &mut self,
         resolver: T,
     );
@@ -317,14 +316,12 @@ impl DistributedExt for SessionStateBuilder {
         set_distributed_user_codec_arc(self.config().get_or_insert_default(), codec)
     }
 
-    fn set_distributed_execution<T: ChannelResolver + Send + Sync + 'static>(
+    fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
         &mut self,
         resolver: T,
     ) {
         let cfg = self.config().get_or_insert_default();
         set_distributed_channel_resolver(cfg, resolver);
-        let rules = self.physical_optimizer_rules().get_or_insert_default();
-        rules.push(Arc::new(DistributedPhysicalOptimizerRule));
     }
 
     fn set_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(
@@ -372,9 +369,9 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($;self)]
             fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
 
-            #[call(set_distributed_execution)]
+            #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
-            fn with_distributed_execution<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
+            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
             #[call(set_distributed_task_estimator)]
             #[expr($;self)]

src/distributed_planner/distributed_physical_optimizer_rule.rs

@@ -403,12 +403,13 @@ impl<T: ?Sized> Referenced<'_, T> {
 
 #[cfg(test)]
 mod tests {
-    use crate::DistributedExt;
     use crate::test_utils::in_memory_channel_resolver::InMemoryChannelResolver;
     use crate::test_utils::parquet::register_parquet_tables;
+    use crate::{DistributedExt, DistributedPhysicalOptimizerRule};
     use crate::{assert_snapshot, display_plan_ascii};
     use datafusion::execution::SessionStateBuilder;
     use datafusion::prelude::{SessionConfig, SessionContext};
+    use std::sync::Arc;
     /* shema for the "weather" table
 
      MinTemp [type=DOUBLE] [repetitiontype=OPTIONAL]
@@ -441,7 +442,7 @@ mod tests {
         SELECT * FROM weather
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @"DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet");
@@ -453,7 +454,7 @@ mod tests {
         SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -485,7 +486,7 @@ mod tests {
         SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(2))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(2))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -517,7 +518,7 @@ mod tests {
         SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(0))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(0))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -540,7 +541,7 @@ mod tests {
         SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
                 .with_distributed_cardinality_effect_task_scale_factor(3.0)
                 .unwrap()
         })
@@ -571,7 +572,7 @@ mod tests {
         SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
                 .with_distributed_files_per_task(3)
                 .unwrap()
         })
@@ -596,7 +597,7 @@ mod tests {
         SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -628,7 +629,7 @@ mod tests {
         SELECT a."MinTemp", b."MaxTemp" FROM weather a LEFT JOIN weather b ON a."RainToday" = b."RainToday"
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -666,7 +667,7 @@ mod tests {
         ON a."RainTomorrow" = b."RainTomorrow"
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -714,7 +715,7 @@ mod tests {
         SELECT * FROM weather ORDER BY "MinTemp" DESC
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -736,7 +737,7 @@ mod tests {
         SELECT DISTINCT "RainToday", "WindGustDir" FROM weather
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -765,7 +766,7 @@ mod tests {
         SHOW COLUMNS from weather
         "#;
         let plan = sql_to_explain(query, |b| {
-            b.with_distributed_execution(InMemoryChannelResolver::new(3))
+            b.with_distributed_channel_resolver(InMemoryChannelResolver::new(3))
         })
         .await;
         assert_snapshot!(plan, @r"
@@ -788,6 +789,7 @@ mod tests {
 
         let builder = SessionStateBuilder::new()
             .with_default_features()
+            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
             .with_config(config);
 
         let state = f(builder).build();

src/metrics/task_metrics_collector.rs

@@ -127,12 +127,12 @@ mod tests {
     use datafusion::arrow::record_batch::RecordBatch;
     use futures::StreamExt;
 
-    use crate::DistributedExt;
     use crate::execution_plans::DistributedExec;
     use crate::metrics::proto::metrics_set_proto_to_df;
     use crate::test_utils::in_memory_channel_resolver::InMemoryChannelResolver;
     use crate::test_utils::plans::{count_plan_nodes, get_stages_and_stage_keys};
     use crate::test_utils::session_context::register_temp_parquet_table;
+    use crate::{DistributedExt, DistributedPhysicalOptimizerRule};
     use datafusion::execution::{SessionStateBuilder, context::SessionContext};
     use datafusion::prelude::SessionConfig;
     use datafusion::{
@@ -151,7 +151,8 @@ mod tests {
         let state = SessionStateBuilder::new()
             .with_default_features()
             .with_config(config)
-            .with_distributed_execution(InMemoryChannelResolver::new(10))
+            .with_distributed_channel_resolver(InMemoryChannelResolver::new(10))
+            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
             .with_distributed_task_estimator(2)
             .build();
 

src/metrics/task_metrics_rewriter.rs

@@ -193,7 +193,6 @@ pub fn stage_metrics_rewriter(
 
 #[cfg(test)]
 mod tests {
-    use crate::DistributedExec;
     use crate::PartitionIsolatorExec;
     use crate::metrics::proto::{
         MetricsSetProto, df_metrics_set_to_proto, metrics_set_proto_to_df,
@@ -205,6 +204,7 @@ mod tests {
     use crate::test_utils::metrics::make_test_metrics_set_proto_from_seed;
     use crate::test_utils::plans::count_plan_nodes;
     use crate::test_utils::session_context::register_temp_parquet_table;
+    use crate::{DistributedExec, DistributedPhysicalOptimizerRule};
     use crate::{NetworkBoundaryExt, Stage};
     use bytes::Bytes;
     use datafusion::arrow::array::{Int32Array, StringArray};
@@ -244,7 +244,8 @@ mod tests {
 
         if distributed {
             builder = builder
-                .with_distributed_execution(InMemoryChannelResolver::new(10))
+                .with_distributed_channel_resolver(InMemoryChannelResolver::new(10))
+                .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
                 .with_distributed_task_estimator(2)
         }
 

src/test_utils/in_memory_channel_resolver.rs

@@ -45,7 +45,7 @@ impl InMemoryChannelResolver {
                 async move {
                     let builder = SessionStateBuilder::new()
                         .with_default_features()
-                        .with_distributed_execution(this)
+                        .with_distributed_channel_resolver(this)
                         .with_runtime_env(ctx.runtime_env.clone());
                     Ok(builder.build())
                 }

src/test_utils/localhost.rs

@@ -1,6 +1,6 @@
 use crate::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilder, DistributedSessionBuilderContext,
+    DistributedPhysicalOptimizerRule, DistributedSessionBuilder, DistributedSessionBuilderContext,
     MappedDistributedSessionBuilderExt, create_flight_client,
 };
 use arrow_flight::flight_service_client::FlightServiceClient;
@@ -53,7 +53,10 @@ where
     let channel_resolver = LocalHostChannelResolver::new(ports.clone());
     let session_builder = session_builder.map(move |builder: SessionStateBuilder| {
         let channel_resolver = channel_resolver.clone();
-        Ok(builder.with_distributed_execution(channel_resolver).build())
+        Ok(builder
+            .with_distributed_channel_resolver(channel_resolver)
+            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+            .build())
     });
     let mut join_set = JoinSet::new();
     for listener in listeners {