Distribute UNION operations (#262)

* Split channel resolver in two

* Simplify WorkerResolverExtension and ChannelResolverExtension

* Add default builder to ArrowFlightEndpoint

* Add some docs

* Listen to clippy

* Split get_flight_client_for_url in two

* Fix conflicts

* Remove unnecessary channel resolver

* Improve WorkerResolver docs

* Use one ChannelResolver per runtime

* Improve error reporting on client connection failure

* Add a from_session_builder method for constructing an InMemoryChannelResolver

* Add ChannelResolver and WorkerResolver default implementations for Arcs

* Make TPC-DS tests use DataFusion test dataset

* Remove non-working in-memory option from benchmarks

* Remove unnecessary utils folder

* Refactor benchmark folder

* Rename to prepare_tpch.rs

* Adapt benchmarks for TPC-DS

* Update benchmarks README.md

* Fix conflicts

* Use default session state builder

* Add ChildrenIsolatorUnionExec

* Add proto serde for ChildrenIsolatorUnionExec

* Wire up ChildrenIsolatorUnionExec to planner

* Add integration tests for distributed UNIONs

* Skip query 72 in TPC-DS benchmarks

* Allow setting children isolator unions

* Allow passing multiple queries in benchmarks

Author: gabotechs

benchmarks/get-tpcds.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+
+set -e
+
+SCALE_FACTOR=${SCALE_FACTOR:-1}
+PARTITIONS=${PARTITIONS:-16}
+
+echo "Generating TPC-DS dataset with SCALE_FACTOR=${SCALE_FACTOR} and PARTITIONS=${PARTITIONS}"
+
+# https://stackoverflow.com/questions/59895/how-do-i-get-the-directory-where-a-bash-script-is-located-from-within-the-script
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+DATA_DIR=${DATA_DIR:-$SCRIPT_DIR/data}
+CARGO_COMMAND=${CARGO_COMMAND:-"cargo run --release"}
+TPCDS_DIR="${DATA_DIR}/tpcds_sf${SCALE_FACTOR}"
+
+echo "Creating tpcds dataset at Scale Factor ${SCALE_FACTOR} in ${TPCDS_DIR}..."
+
+# Ensure the target data directory exists
+mkdir -p "${TPCDS_DIR}"
+
+$CARGO_COMMAND -- prepare-tpcds --output "${TPCDS_DIR}" --partitions "$PARTITIONS"

benchmarks/src/run.rs

@@ -58,8 +58,8 @@ use tonic::transport::Server;
 #[structopt(verbatim_doc_comment)]
 pub struct RunOpt {
     /// Query number. If not specified, runs all queries
-    #[structopt(short, long)]
-    pub query: Option<usize>,
+    #[structopt(short, long, use_delimiter = true)]
+    pub query: Vec<usize>,
 
     /// Path to data files
     #[structopt(parse(from_os_str), short = "p", long = "path")]
@@ -89,6 +89,10 @@ pub struct RunOpt {
     #[structopt(long)]
     cardinality_task_sf: Option<f64>,
 
+    /// Use children isolator UNIONs for distributing UNION operations.
+    #[structopt(long)]
+    children_isolator_unions: bool,
+
     /// Collects metrics across network boundaries
     #[structopt(long)]
     collect_metrics: bool,
@@ -143,7 +147,9 @@ impl Dataset {
             Dataset::Tpch => (1..22 + 1)
                 .map(|i| Ok((i as usize, tpch::get_test_tpch_query(i)?)))
                 .collect(),
-            Dataset::Tpcds => (1..99 + 1)
+            Dataset::Tpcds => (1..72)
+                // skip query 72, it's ridiculously slow
+                .chain(73..99 + 1)
                 .map(|i| Ok((i, tpcds::get_test_tpcds_query(i)?)))
                 .collect(),
             Dataset::Clickbench => (0..42 + 1)
@@ -203,6 +209,7 @@ impl RunOpt {
             .with_distributed_cardinality_effect_task_scale_factor(
                 self.cardinality_task_sf.unwrap_or(1.0),
             )?
+            .with_distributed_children_isolator_unions(self.children_isolator_unions)?
             .with_distributed_metrics_collection(self.collect_metrics)?
             .build();
         let ctx = SessionContext::new_with_state(state);
@@ -219,7 +226,7 @@ impl RunOpt {
         let dataset = Dataset::infer_from_data_path(path.clone())?;
 
         for (id, sql) in dataset.queries()? {
-            if self.query.is_some_and(|v| v != id) {
+            if !self.query.is_empty() && !self.query.contains(&id) {
                 continue;
             }
             let query_id = format!("{dataset:?} {id}");

src/distributed_ext.rs

@@ -418,6 +418,35 @@ pub trait DistributedExt: Sized {
 
     /// Same as [DistributedExt::with_distributed_metrics_collection] but with an in-place mutation.
     fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+
+    /// Enables children isolator unions for distributing UNION operations across as many tasks as
+    /// the sum of all the tasks required for each child.
+    ///
+    /// For example, if there is a UNION with 3 children, requiring one task each, it will result
+    /// in a plan with 3 tasks where each task runs one child:
+    ///
+    /// ```text
+    /// ┌─────────────────────────────┐┌─────────────────────────────┐┌─────────────────────────────┐
+    /// │           Task 1            ││           Task 2            ││           Task 3            │
+    /// │┌───────────────────────────┐││┌───────────────────────────┐││┌───────────────────────────┐│
+    /// ││ ChildrenIsolatorUnionExec ││││ ChildrenIsolatorUnionExec ││││ ChildrenIsolatorUnionExec ││
+    /// │└───▲─────────▲─────────▲───┘││└───▲─────────▲─────────▲───┘││└───▲─────────▲─────────▲───┘│
+    /// │    │                        ││              │              ││                        │    │
+    /// │┌───┴───┐ ┌  ─│ ─   ┌  ─│ ─  ││┌  ─│ ─   ┌───┴───┐ ┌  ─│ ─  ││┌  ─│ ─   ┌  ─│ ─   ┌───┴───┐│
+    /// ││Child 1│  Child 2│  Child 3│││ Child 1│ │Child 2│  Child 3│││ Child 1│  Child 2│ │Child 3││
+    /// │└───────┘ └  ─  ─   └  ─  ─  ││└  ─  ─   └───────┘ └  ─  ─  ││└  ─  ─   └  ─  ─   └───────┘│
+    /// └─────────────────────────────┘└─────────────────────────────┘└─────────────────────────────┘
+    /// ```
+    fn with_distributed_children_isolator_unions(
+        self,
+        enabled: bool,
+    ) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_children_isolator_unions] but with an in-place mutation.
+    fn set_distributed_children_isolator_unions(
+        &mut self,
+        enabled: bool,
+    ) -> Result<(), DataFusionError>;
 }
 
 impl DistributedExt for SessionConfig {
@@ -489,6 +518,15 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_children_isolator_unions(
+        &mut self,
+        enabled: bool,
+    ) -> Result<(), DataFusionError> {
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        d_cfg.children_isolator_unions = enabled;
+        Ok(())
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -530,6 +568,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_metrics_collection)]
             #[expr($?;Ok(self))]
             fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            #[call(set_distributed_children_isolator_unions)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -586,6 +628,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_metrics_collection)]
             #[expr($?;Ok(self))]
             fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_children_isolator_unions)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -642,6 +689,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_metrics_collection)]
             #[expr($?;Ok(self))]
             fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_children_isolator_unions)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -698,6 +750,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_metrics_collection)]
             #[expr($?;Ok(self))]
             fn with_distributed_metrics_collection(self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_children_isolator_unions(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_children_isolator_unions)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_children_isolator_unions(self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }

src/distributed_planner/distributed_config.rs

@@ -28,6 +28,14 @@ extensions_options! {
         /// batches over the wire.
         /// If set to 0, batch coalescing is disabled on network shuffle operations.
         pub shuffle_batch_size: usize, default = 8192
+        /// When encountering a UNION operation, isolate its children depending on the task context.
+        /// For example, on a UNION operation with 3 children running in 3 distributed tasks,
+        /// instead of executing the 3 children in each 3 tasks with a DistributedTaskContext of
+        /// 1/3, 2/3, and 3/3 respectively, Execute:
+        /// - The first child in the first task with a DistributedTaskContext of 1/1
+        /// - The second child in the second task with a DistributedTaskContext of 1/1
+        /// - The third child in the third task with a DistributedTaskContext of 1/1
+        pub children_isolator_unions: bool, default = true
         /// Propagate collected metrics from all nodes in the plan across network boundaries
         /// so that they can be reconstructed on the head node of the plan.
         pub collect_metrics: bool, default = true