fix execution by storing stages in a OnceCell

Author: robtandy

Cargo.lock

@@ -22,6 +22,7 @@ dashmap = "6.1.0"
 prost = "0.13.5"
 rand = "0.8.5"
 object_store = "0.12.3"
+async-stream = "0.3.6"
 
 [dev-dependencies]
 insta = { version = "1.43.1", features = ["filters"] }

Cargo.toml

@@ -1,28 +1,39 @@
 use crate::composed_extension_codec::ComposedPhysicalExtensionCodec;
 use crate::errors::datafusion_error_to_tonic_status;
 use crate::flight_service::service::ArrowFlightEndpoint;
-use crate::plan::DistributedCodec;
-use crate::stage::{stage_from_proto, ExecutionStageProto};
+use crate::plan::{DistributedCodec, PartitionGroup};
+use crate::stage::{stage_from_proto, ExecutionStage, ExecutionStageProto, StageKey};
 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::SessionStateBuilder;
+use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::optimizer::OptimizerConfig;
+use datafusion::physical_plan::ExecutionPlan;
 use futures::TryStreamExt;
 use prost::Message;
 use std::sync::Arc;
+use tokio::sync::OnceCell;
 use tonic::{Request, Response, Status};
 
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct DoGet {
     /// The ExecutionStage that we are going to execute
     #[prost(message, optional, tag = "1")]
     pub stage_proto: Option<ExecutionStageProto>,
-    /// the partition of the stage to execute
+    /// The index to the task within the stage that we want to execute
     #[prost(uint64, tag = "2")]
+    pub task_number: u64,
+    /// the partition number we want to execute
+    #[prost(uint64, tag = "3")]
     pub partition: u64,
+    /// The stage key that identifies the stage.  This is useful to keep
+    /// outside of the stage proto as it is used to store the stage
+    /// and we may not need to deserialize the entire stage proto
+    /// if we already have stored it
+    #[prost(message, optional, tag = "4")]
+    pub stage_key: Option<StageKey>,
 }
 
 impl ArrowFlightEndpoint {
@@ -35,42 +46,28 @@ impl ArrowFlightEndpoint {
             Status::invalid_argument(format!("Cannot decode DoGet message: {err}"))
         })?;
 
-        let stage_msg = doget
-            .stage_proto
-            .ok_or(Status::invalid_argument("DoGet is missing the stage proto"))?;
+        let partition = doget.partition as usize;
+        let task_number = doget.task_number as usize;
+        let (mut state, stage) = self.get_state_and_stage(doget).await?;
 
-        let state_builder = SessionStateBuilder::new()
-            .with_runtime_env(Arc::clone(&self.runtime))
-            .with_default_features();
+        // find out which partition group we are executing
+        let task = stage
+            .tasks
+            .get(task_number)
+            .ok_or(Status::invalid_argument(format!(
+                "Task number {} not found in stage {}",
+                task_number,
+                stage.name()
+            )))?;
 
-        let mut state = self.session_builder.on_new_session(state_builder).build();
+        let partition_group =
+            PartitionGroup(task.partition_group.iter().map(|p| *p as usize).collect());
+        state.config_mut().set_extension(Arc::new(partition_group));
 
-        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 mut stage = stage_from_proto(
-            stage_msg,
-            function_registry,
-            &self.runtime.as_ref(),
-            &combined_codec,
-        )
-        .map_err(|err| Status::invalid_argument(format!("Cannot decode stage proto: {err}")))?;
-        let inner_plan = Arc::clone(&stage.plan);
-
-        // 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(Arc::new(stage));
+        let inner_plan = stage.plan.clone();
 
         let stream = inner_plan
-            .execute(doget.partition as usize, state.task_ctx())
+            .execute(partition, state.task_ctx())
             .map_err(|err| Status::internal(format!("Error executing stage plan: {err:#?}")))?;
 
         let flight_data_stream = FlightDataEncoderBuilder::new()
@@ -86,4 +83,59 @@ impl ArrowFlightEndpoint {
             },
         ))))
     }
+
+    async fn get_state_and_stage(
+        &self,
+        doget: DoGet,
+    ) -> Result<(SessionState, Arc<ExecutionStage>), Status> {
+        let key = doget
+            .stage_key
+            .ok_or(Status::invalid_argument("DoGet is missing the stage key"))?;
+        let once_stage = self.stages.entry(key).or_default();
+
+        let (state, stage) = once_stage
+            .get_or_try_init(|| async {
+                let stage_proto = doget
+                    .stage_proto
+                    .ok_or(Status::invalid_argument("DoGet is missing the stage proto"))?;
+
+                let state_builder = SessionStateBuilder::new()
+                    .with_runtime_env(Arc::clone(&self.runtime))
+                    .with_default_features();
+
+                let mut state = self.session_builder.on_new_session(state_builder).build();
+
+                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}"))
+                })?;
+
+                // 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());
+
+                Ok::<_, Status>((state, stage))
+            })
+            .await?;
+
+        Ok((state.clone(), stage.clone()))
+    }
 }

src/flight_service/do_get.rs

@@ -1,7 +1,6 @@
 mod do_get;
 mod service;
 mod session_builder;
-mod stream_partitioner_registry;
 
 pub(crate) use do_get::DoGet;
 

src/flight_service/mod.rs

@@ -1,23 +1,26 @@
 use crate::channel_manager::ChannelManager;
 use crate::flight_service::session_builder::NoopSessionBuilder;
-use crate::flight_service::stream_partitioner_registry::StreamPartitionerRegistry;
 use crate::flight_service::SessionBuilder;
+use crate::stage::{ExecutionStage, StageKey};
 use crate::ChannelResolver;
 use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::{
     Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
     HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
 };
 use async_trait::async_trait;
+use dashmap::DashMap;
 use datafusion::execution::runtime_env::RuntimeEnv;
+use datafusion::execution::SessionState;
 use futures::stream::BoxStream;
 use std::sync::Arc;
+use tokio::sync::OnceCell;
 use tonic::{Request, Response, Status, Streaming};
 
 pub struct ArrowFlightEndpoint {
     pub(super) channel_manager: Arc<ChannelManager>,
     pub(super) runtime: Arc<RuntimeEnv>,
-    pub(super) partitioner_registry: Arc<StreamPartitionerRegistry>,
+    pub(super) stages: DashMap<StageKey, OnceCell<(SessionState, Arc<ExecutionStage>)>>,
     pub(super) session_builder: Arc<dyn SessionBuilder + Send + Sync>,
 }
 
@@ -26,7 +29,7 @@ impl ArrowFlightEndpoint {
         Self {
             channel_manager: Arc::new(ChannelManager::new(channel_resolver)),
             runtime: Arc::new(RuntimeEnv::default()),
-            partitioner_registry: Arc::new(StreamPartitionerRegistry::default()),
+            stages: DashMap::new(),
             session_builder: Arc::new(NoopSessionBuilder),
         }
     }