Merge pull request #6 from robtandy/rob.tandy/arrow-flight

randomly assign stages to exchanges

Author: robtandy

datafusion_ray/core.py

@@ -22,6 +22,7 @@
 import uuid
 import os
 import time
+import random
 
 from datafusion_ray._datafusion_ray_internal import (
     RayContext as RayContextInternal,
@@ -35,6 +36,7 @@ class RayDataFrame:
     def __init__(
         self,
         ray_internal_df: RayDataFrameInternal,
+        num_partitions,
         batch_size=8192,
         isolate_parititions=False,
         bucket: str | None = None,
@@ -48,6 +50,7 @@ def __init__(
         self.isolate_partitions = isolate_parititions
         self.bucket = bucket
         self.num_exchangers = num_exchangers
+        self.num_partitions = num_partitions
 
     def stages(self):
         # create our coordinator now, which we need to create stages
@@ -56,7 +59,12 @@ def stages(self):
 
             self.coord = RayStageCoordinator.options(
                 name="RayQueryCoordinator:" + self.coordinator_id,
-            ).remote(self.coordinator_id, len(self._stages), self.num_exchangers)
+            ).remote(
+                self.coordinator_id,
+                len(self._stages),
+                self.num_exchangers,
+                self.num_partitions,
+            )
 
             ray.get(self.coord.start_up.remote())
             print("ray coord started up")
@@ -74,7 +82,7 @@ def collect(self) -> list[pa.RecordBatch]:
 
             last_stage = max([stage.stage_id for stage in self._stages])
 
-            ref = self.coord.get_exchanger_addr.remote(last_stage)
+            ref = self.coord.get_exchanger_addr.remote(last_stage, partition=0)
             self.create_ray_stages()
             t3 = time.time()
             print(f"creating ray stage actors took {t3 -t2}s")
@@ -86,7 +94,7 @@ def collect(self) -> list[pa.RecordBatch]:
             )
 
             print("calling df execute")
-            reader = self.df.execute({last_stage: addr})
+            reader = self.df.execute({(last_stage, 0): addr})
             print("called df execute, got reader")
             self._batches = list(reader)
             self.coord.all_done.remote()
@@ -187,6 +195,7 @@ def sql(self, query: str) -> RayDataFrame:
         df = self.ctx.sql(query, coordinator_id)
         return RayDataFrame(
             df,
+            self.ctx.get_target_partitions(),
             self.batch_size,
             self.isolate_partitions,
             self.bucket,
@@ -205,12 +214,17 @@ def set(self, option: str, value: str) -> None:
 @ray.remote(num_cpus=0)
 class RayStageCoordinator:
     def __init__(
-        self, coordinator_id: str, num_stages: int, num_exchangers: int
+        self,
+        coordinator_id: str,
+        num_stages: int,
+        num_exchangers: int,
+        num_partitions: int,
     ) -> None:
         self.my_id = coordinator_id
         self.stages = {}
         self.num_stages = num_stages
         self.num_exchangers = num_exchangers
+        self.num_partitions = num_partitions
         self.runtime_env = {}
 
     def start_up(self):
@@ -221,28 +235,27 @@ def start_up(self):
             RayExchanger.remote(f"Exchanger #{i}") for i in range(self.num_exchangers)
         ]
 
-        stages_per_exchanger = max(1, self.num_stages // self.num_exchangers)
-        print("Stages per exchanger: ", stages_per_exchanger)
-
         refs = [exchange.start_up.remote() for exchange in self.xs]
 
         # ensure we've done the necessary initialization before continuing
         ray.wait(refs, num_returns=len(refs))
         print("all exchanges started up")
 
-        self.exchanges = {}
+        # for each possible stage, and partition, assign it to an exchanger
         self.exchange_addrs = {}
-        for i in range(self.num_stages):
-            exchanger_i = min(len(self.xs) - 1, i // stages_per_exchanger)
-            print("exchanger_i = ", exchanger_i)
-            self.exchanges[i] = self.xs[exchanger_i]
-            self.exchange_addrs[i] = ray.get(self.xs[exchanger_i].addr.remote())
+        for stage_num in range(self.num_stages):
+            for partition_num in range(self.num_partitions):
+                exchanger_idx = random.choice(range(self.num_exchangers))
+                self.exchange_addrs[(stage_num, partition_num)] = ray.get(
+                    self.xs[exchanger_idx].addr.remote()
+                )
+        print(self.exchange_addrs)
 
         # don't wait for these
         [exchange.serve.remote() for exchange in self.xs]
 
-    def get_exchanger_addr(self, stage_num: int):
-        return self.exchange_addrs[stage_num]
+    def get_exchanger_addr(self, stage_num: int, partition: int):
+        return self.exchange_addrs[(stage_num, partition)]
 
     def all_done(self):
         print("calling exchangers all done")
@@ -269,16 +282,6 @@ def new_stage(
     ):
         stage_key = f"{stage_id}-{shadow_partition}"
         try:
-            if stage_key in self.stages:
-                print(f"already started stage {stage_key}")
-                return self.stages[stage_key]
-
-            exchange_addr = self.exchange_addrs[stage_id]
-
-            input_exchange_addrs = {
-                input_stage_id: self.exchange_addrs[input_stage_id]
-                for input_stage_id in input_stage_ids
-            }
 
             print(f"creating new stage {stage_key} from bytes {len(plan_bytes)}")
             stage = RayStage.options(
@@ -287,8 +290,7 @@ def new_stage(
             ).remote(
                 stage_id,
                 plan_bytes,
-                exchange_addr,
-                input_exchange_addrs,
+                self.exchange_addrs,
                 fraction,
                 shadow_partition,
                 bucket,
@@ -337,36 +339,31 @@ def __init__(
         self,
         stage_id: str,
         plan_bytes: bytes,
-        exchanger_addr: str,
-        input_exchange_addrs: dict[int, str],
+        exchanger_addrs: dict[tuple[int, int], str],
         fraction: float,
         shadow_partition=None,
         bucket: str | None = None,
     ):
 
         from datafusion_ray._datafusion_ray_internal import PyStage
 
+        self.shadow_partition = shadow_partition
+        shadow = (
+            f", shadowing:{self.shadow_partition}"
+            if self.shadow_partition is not None
+            else ""
+        )
+
         try:
             self.stage_id = stage_id
             self.pystage = PyStage(
                 stage_id,
                 plan_bytes,
-                exchanger_addr,
-                input_exchange_addrs,
+                exchanger_addrs,
                 shadow_partition,
                 bucket,
                 fraction,
             )
-            self.shadow_partition = shadow_partition
-            shadow = (
-                f", shadowing:{self.shadow_partition}"
-                if self.shadow_partition is not None
-                else ""
-            )
-
-            print(
-                f"RayStage[{self.stage_id}{shadow}] Sending to {exchanger_addr}, consuming from {input_exchange_addrs}"
-            )
         except Exception as e:
             print(
                 f"RayStage[{self.stage_id}{shadow}] Unhandled Exception in init: {e}!"

src/context.rs

@@ -124,6 +124,14 @@ impl RayContext {
         Ok(())
     }
 
+    pub fn get_target_partitions(&self) -> usize {
+        let state = self.ctx.state_ref();
+        let guard = state.read();
+        let config = guard.config();
+        let options = config.options();
+        options.execution.target_partitions
+    }
+
     pub fn set_coordinator_id(&self, id: String) -> PyResult<()> {
         let state = self.ctx.state_ref();
         let mut guard = state.write();

src/dataframe.rs

@@ -42,7 +42,7 @@ use tonic::transport::Channel;
 
 use crate::isolator::PartitionIsolatorExec;
 use crate::max_rows::MaxRowsExec;
-use crate::pystage::ExchangeFlightClient;
+use crate::pystage::ExchangeAddrs;
 use crate::ray_stage::RayStageExec;
 use crate::ray_stage_reader::RayStageReaderExec;
 use crate::util::make_client;
@@ -198,21 +198,9 @@ impl RayDataFrame {
 
     pub fn execute(
         &self,
-        py: Python,
-        in_exchange_addrs: HashMap<usize, String>,
+        exchange_addrs: HashMap<(usize, usize), String>,
     ) -> PyResult<PyRecordBatchStream> {
-        // TODO: consolidate this code
-        //
-        let in_client_map: HashMap<usize, FlightClient> = in_exchange_addrs
-            .iter()
-            .map(|(stage_num, addr)| {
-                let client = make_client(py, addr).to_py_err()?;
-                Ok::<_, PyErr>((stage_num.clone(), client))
-            })
-            .collect::<Result<HashMap<_, _>, PyErr>>()?;
-
-        let config =
-            SessionConfig::new().with_extension(Arc::new(ExchangeFlightClient(in_client_map)));
+        let config = SessionConfig::new().with_extension(Arc::new(ExchangeAddrs(exchange_addrs)));
 
         let state = SessionStateBuilder::new()
             .with_default_features()

src/pystage.rs

@@ -14,7 +14,7 @@ use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::common::{internal_datafusion_err, internal_err};
 use datafusion::physical_plan::{collect, displayable};
 use datafusion_python::physical_plan::PyExecutionPlan;
-use futures::future::try_join_all;
+use futures::future::{join_all, try_join_all};
 use futures::stream::FuturesUnordered;
 use object_store::aws::AmazonS3Builder;
 use prost::Message;
@@ -34,12 +34,12 @@ use pyo3::prelude::*;
 
 use anyhow::Result;
 
-pub struct ExchangeFlightClient(pub HashMap<usize, FlightClient>);
+/// a map of (stage_id, partition_id) to FlightClient used to speak to that Exchanger
+pub(crate) struct ExchangeAddrs(pub HashMap<(usize, usize), String>);
 
 #[pyclass]
 pub struct PyStage {
     name: String,
-    out_client: FlightClient,
     #[pyo3(get)]
     stage_id: usize,
     pub(crate) plan: Arc<dyn ExecutionPlan>,
@@ -51,13 +51,12 @@ pub struct PyStage {
 #[pymethods]
 impl PyStage {
     #[new]
-    #[pyo3(signature = (stage_id, plan_bytes, out_exchange_addr, in_exchange_addrs, shadow_partition_number=None, bucket=None, fraction=1.0))]
+    #[pyo3(signature = (stage_id, plan_bytes, exchange_addrs, shadow_partition_number=None, bucket=None, fraction=1.0))]
     pub fn from_bytes(
         py: Python,
         stage_id: usize,
         plan_bytes: Vec<u8>,
-        out_exchange_addr: String,
-        in_exchange_addrs: HashMap<usize, String>,
+        exchange_addrs: HashMap<(usize, usize), String>,
         shadow_partition_number: Option<usize>,
         bucket: Option<String>,
         fraction: f64,
@@ -69,19 +68,9 @@ impl PyStage {
                 .map(|s| s.to_string())
                 .unwrap_or("n/a".into())
         );
-        let out_client = make_client(py, &out_exchange_addr)?;
-
-        // make our own clone as FlightClient is not Clone, but inner is
-        let out_client_map: HashMap<usize, FlightClient> = in_exchange_addrs
-            .iter()
-            .map(|(stage_num, addr)| {
-                let client = make_client(py, addr).to_py_err()?;
-                Ok::<_, PyErr>((stage_num.clone(), client))
-            })
-            .collect::<Result<HashMap<_, _>, PyErr>>()?;
 
         let mut config =
-            SessionConfig::new().with_extension(Arc::new(ExchangeFlightClient(out_client_map)));
+            SessionConfig::new().with_extension(Arc::new(ExchangeAddrs(exchange_addrs)));
 
         // this only matters if the plan includes an PartitionIsolatorExec
         // and will be ignored otherwise
@@ -116,7 +105,6 @@ impl PyStage {
 
         Ok(Self {
             name,
-            out_client,
             stage_id,
             plan,
             ctx,
@@ -128,39 +116,47 @@ impl PyStage {
     pub fn execute(&mut self, py: Python) -> PyResult<()> {
         println!("{} executing", self.name);
 
+        let addrs = &self
+            .ctx
+            .state()
+            .config()
+            .get_extension::<ExchangeAddrs>()
+            .ok_or(internal_datafusion_err!("Flight Client not in context"))?
+            .clone()
+            .0;
+
         let futs = (0..self.num_output_partitions()).map(|partition| {
             let ctx = self.ctx.task_ctx();
-            // make our own clone as FlightClient is not Clone, but inner is
-            let inner = self.out_client.inner().clone();
-            let client_clone = FlightClient::new_from_inner(inner);
             let plan = self.plan.clone();
-            let stage_id = self.stage_id;
-            let fraction = self.fraction;
-            let shadow_partition_number = self.shadow_partition_number;
-
-            tokio::spawn(consume_stage(
-                stage_id,
-                shadow_partition_number,
-                fraction,
-                ctx,
-                partition,
-                plan,
-                client_clone,
-            ))
+            let stage_id = self.stage_id.clone();
+            let fraction = self.fraction.clone();
+            let shadow_partition_number = self.shadow_partition_number.clone();
+
+            async move {
+                // TODO propagate these errors appropriately
+                let client = addrs
+                    .get(&(stage_id, partition))
+                    .map(|addr| make_client(addr))
+                    .expect("cannot find addr")
+                    .await
+                    .expect("cannot make client");
+
+                tokio::spawn(consume_stage(
+                    stage_id,
+                    shadow_partition_number,
+                    fraction,
+                    ctx,
+                    partition,
+                    plan,
+                    client,
+                ));
+            }
         });
 
-        let name = self.name.clone();
-        let fut = async {
-            match try_join_all(futs).await {
-                Ok(_) => Ok(()),
-                Err(e) => {
-                    println!("{name}:ERROR executing {e}");
-                    Err(e)
-                }
-            }
-        };
+        let fut = join_all(futs);
 
-        wait_for_future(py, fut).to_py_err()
+        wait_for_future(py, fut);
+        Ok(())
     }
 
     pub fn num_output_partitions(&self) -> usize {