Adapt planner so that plan distribution can be made async

Author: gabotechs

benchmarks/cdk/bin/worker.rs

@@ -10,8 +10,8 @@ use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::physical_plan::execute_stream;
 use datafusion::prelude::SessionContext;
 use datafusion_distributed::{
-    ChannelResolver, DistributedExt, DistributedMetricsFormat, DistributedPhysicalOptimizerRule,
-    Worker, WorkerResolver, display_plan_ascii, get_distributed_channel_resolver,
+    ChannelResolver, DistributedExt, DistributedMetricsFormat, SessionStateBuilderExt, Worker,
+    WorkerResolver, display_plan_ascii, get_distributed_channel_resolver,
     get_distributed_worker_resolver, rewrite_distributed_plan_with_metrics,
 };
 use futures::{StreamExt, TryFutureExt};
@@ -94,7 +94,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .with_default_features()
         .with_runtime_env(Arc::clone(&runtime_env))
         .with_distributed_worker_resolver(Ec2WorkerResolver::new())
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         .with_distributed_broadcast_joins(cmd.broadcast_joins)?
         .build();
     let ctx = SessionContext::from(state);

benchmarks/src/run.rs

@@ -28,9 +28,7 @@ use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::{collect, displayable};
 use datafusion::prelude::*;
 use datafusion_distributed::test_utils::localhost::LocalHostWorkerResolver;
-use datafusion_distributed::{
-    DistributedExt, DistributedPhysicalOptimizerRule, NetworkBoundaryExt, Worker,
-};
+use datafusion_distributed::{DistributedExt, NetworkBoundaryExt, SessionStateBuilderExt, Worker};
 use datafusion_distributed_benchmarks::datasets::{clickbench, register_tables, tpcds, tpch};
 use std::error::Error;
 use std::fs;
@@ -178,7 +176,7 @@ impl RunOpt {
             .with_default_features()
             .with_config(self.config()?)
             .with_distributed_worker_resolver(LocalHostWorkerResolver::new(self.workers.clone()))
-            .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+            .with_distributed_planner()
             .with_distributed_files_per_task(
                 self.files_per_task.unwrap_or(get_available_parallelism()),
             )?

cli/src/main.rs

@@ -36,7 +36,7 @@ use datafusion_cli::{
 use datafusion_distributed::test_utils::in_memory_channel_resolver::{
     InMemoryChannelResolver, InMemoryWorkerResolver,
 };
-use datafusion_distributed::{DistributedExt, DistributedPhysicalOptimizerRule};
+use datafusion_distributed::{DistributedExt, SessionStateBuilderExt};
 use std::env;
 use std::path::Path;
 use std::process::ExitCode;
@@ -148,7 +148,7 @@ async fn main_inner() -> Result<()> {
         .with_default_features()
         .with_config(session_config)
         .with_runtime_env(runtime_env)
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         .with_distributed_worker_resolver(InMemoryWorkerResolver::new(16))
         .with_distributed_channel_resolver(InMemoryChannelResolver::default())
         .build();

console/examples/console_run.rs

@@ -4,11 +4,10 @@ use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use datafusion_distributed::{
-    DistributedExt, DistributedPhysicalOptimizerRule, WorkerResolver, display_plan_ascii,
+    DistributedExt, SessionStateBuilderExt, WorkerResolver, display_plan_ascii,
 };
 use futures::TryStreamExt;
 use std::error::Error;
-use std::sync::Arc;
 use structopt::StructOpt;
 use url::Url;
 
@@ -38,7 +37,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let state = SessionStateBuilder::new()
         .with_default_features()
         .with_distributed_worker_resolver(localhost_resolver)
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         .with_distributed_files_per_task(1)?
         .build();
 

console/examples/tpcds_runner.rs

@@ -4,7 +4,7 @@ use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::execute_stream;
 use datafusion::prelude::SessionContext;
 use datafusion_distributed::{
-    DistributedExt, DistributedMetricsFormat, DistributedPhysicalOptimizerRule, WorkerResolver,
+    DistributedExt, DistributedMetricsFormat, SessionStateBuilderExt, WorkerResolver,
     display_plan_ascii,
 };
 use datafusion_distributed_benchmarks::datasets::{register_tables, tpcds};
@@ -97,7 +97,7 @@ async fn run_queries(
     let state = SessionStateBuilder::new()
         .with_default_features()
         .with_distributed_worker_resolver(localhost_resolver)
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         .build();
 
     let ctx = SessionContext::from(state);

docs/source/user-guide/channel-resolver.md

@@ -43,7 +43,7 @@ async fn main() {
     let state = SessionStateBuilder::new()
         // these two are mandatory.
         .with_distributed_worker_resolver(todo!())
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         // the CustomChannelResolver needs to be passed here once...
         .with_distributed_channel_resolver(channel_resolver.clone())
         .build();

docs/source/user-guide/concepts.md

@@ -24,10 +24,10 @@ You'll see these concepts mentioned extensively across the documentation and the
 
 Some other more tangible concepts are the structs and traits exposed publicly, the most important are:
 
-## [DistributedPhysicalOptimizerRule](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/distributed_planner/distributed_physical_optimizer_rule.rs)
+## [SessionStateBuilderExt](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/distributed_planner/session_state_builder_ext.rs)
 
-A physical optimizer rule that transforms single-node DataFusion query plans into distributed query plans. It reads
-a fully formed physical plan and injects the appropriate nodes to execute the query in a distributed fashion.
+An extension trait for `SessionStateBuilder` that provides `with_distributed_planner()`. This registers a custom
+query planner that transforms single-node DataFusion query plans into distributed query plans after physical planning.
 
 It builds the distributed plan from bottom to top, injecting network boundaries at appropriate locations based on
 the nodes present in the original plan.

docs/source/user-guide/getting-started.md

@@ -49,7 +49,7 @@ impl WorkerResolver for LocalhostWorkerResolver {
 }
 ```
 
-Register both the `WorkerResolver` implementation and the `DistributedPhysicalOptimizerRule` in DataFusion's
+Register the `WorkerResolver` implementation and the distributed planner in DataFusion's
 `SessionStateBuilder` to enable distributed query planning:
 
 ```rs
@@ -59,7 +59,7 @@ let localhost_worker_resolver = LocalhostWorkerResolver {
 
 let state = SessionStateBuilder::new()
     .with_distributed_worker_resolver(localhost_worker_resolver)
-    .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+    .with_distributed_planner()
     .build();
 
 let ctx = SessionContext::from(state);

docs/source/user-guide/worker-resolver.md

@@ -24,7 +24,7 @@ impl WorkerResolver for CustomWorkerResolver {
 async fn main() {
     let state = SessionStateBuilder::new()
         .with_distributed_worker_resolver(CustomWorkerResolver)
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         .build();
 }
 ```

examples/custom_execution_plan.rs

@@ -39,8 +39,8 @@ use datafusion_distributed::test_utils::in_memory_channel_resolver::{
     InMemoryChannelResolver, InMemoryWorkerResolver,
 };
 use datafusion_distributed::{
-    DistributedExt, DistributedPhysicalOptimizerRule, DistributedTaskContext, TaskEstimation,
-    TaskEstimator, WorkerQueryContext, display_plan_ascii,
+    DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
+    WorkerQueryContext, display_plan_ascii,
 };
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf;
@@ -382,7 +382,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with_config(config)
         .with_distributed_worker_resolver(worker_resolver)
         .with_distributed_channel_resolver(channel_resolver)
-        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
+        .with_distributed_planner()
         .with_distributed_user_codec(NumbersExecCodec)
         .with_distributed_task_estimator(NumbersTaskEstimator)
         .build();