Optional serialization/deserialization for WorkerChannel

Author: gabotechs

src/common/maybe_encoded.rs

@@ -0,0 +1,100 @@
+use crate::DistributedCodec;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::{Result, internal_err};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::Partitioning;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning;
+use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
+use datafusion_proto::physical_plan::{
+    AsExecutionPlan, DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext,
+};
+use datafusion_proto::protobuf;
+use datafusion_proto::protobuf::proto_error;
+use prost::Message;
+use std::sync::Arc;
+
+/// A value that a transport may either leave encoded or materialize in memory.
+#[derive(Clone)]
+pub enum MaybeEncoded<T> {
+    Encoded(Vec<u8>),
+    Decoded(T),
+}
+
+impl<T> MaybeEncoded<T> {
+    pub fn decode_with(self, decode: impl FnOnce(Vec<u8>) -> Result<T>) -> Result<T> {
+        match self {
+            Self::Encoded(encoded) => decode(encoded),
+            Self::Decoded(decoded) => Ok(decoded),
+        }
+    }
+
+    pub fn try_decoded(self) -> Result<T> {
+        match self {
+            Self::Encoded(_) => {
+                internal_err!("Expected {} to be decoded", std::any::type_name::<T>())
+            }
+            Self::Decoded(decoded) => Ok(decoded),
+        }
+    }
+}
+
+impl MaybeEncoded<Arc<dyn ExecutionPlan>> {
+    pub fn encode(self, ctx: &Arc<TaskContext>) -> Result<Vec<u8>> {
+        match self {
+            Self::Encoded(encoded) => Ok(encoded),
+            Self::Decoded(plan) => {
+                let codec = DistributedCodec::new_combined_with_user(ctx.session_config());
+                Ok(
+                    protobuf::PhysicalPlanNode::try_from_physical_plan(plan, &codec)?
+                        .encode_to_vec(),
+                )
+            }
+        }
+    }
+
+    pub fn decode(self, task_ctx: &TaskContext) -> Result<Arc<dyn ExecutionPlan>> {
+        self.decode_with(|encoded| {
+            let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config());
+            let proto_node = protobuf::PhysicalPlanNode::try_decode(encoded.as_ref())?;
+            proto_node.try_into_physical_plan(task_ctx, &codec)
+        })
+    }
+}
+
+impl MaybeEncoded<Partitioning> {
+    pub fn encode(self, ctx: &Arc<TaskContext>) -> Result<Vec<u8>> {
+        match self {
+            Self::Encoded(encoded) => Ok(encoded),
+            Self::Decoded(partitioning) => {
+                let codec = DistributedCodec::new_combined_with_user(ctx.session_config());
+                Ok(serialize_partitioning(
+                    &partitioning,
+                    &codec,
+                    // I think nobody cares about this being the default PhysicalProtoConverter.
+                    // If someone does, please open an issue.
+                    &DefaultPhysicalProtoConverter {},
+                )?
+                .encode_to_vec())
+            }
+        }
+    }
+
+    pub fn decode(self, schema: SchemaRef, task_ctx: &TaskContext) -> Result<Partitioning> {
+        self.decode_with(|encoded| {
+            let proto_partitioning = protobuf::Partitioning::decode(encoded.as_slice())
+                .map_err(|err| proto_error(err.to_string()))?;
+            let codec = DistributedCodec::new_combined_with_user(task_ctx.session_config());
+            let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx, &codec);
+            parse_protobuf_partitioning(
+                Some(&proto_partitioning),
+                &decode_ctx,
+                &schema,
+                // I think nobody cares about this being the default PhysicalProtoConverter.
+                // If someone does, please open an issue.
+                &DefaultPhysicalProtoConverter {},
+            )?
+            .ok_or_else(|| proto_error("Could not parse partitioning"))
+        })
+    }
+}

src/common/mod.rs

@@ -1,4 +1,5 @@
 mod children_helpers;
+mod maybe_encoded;
 mod once_lock;
 mod recursion;
 mod task_context_helpers;
@@ -7,6 +8,7 @@ mod uuid;
 mod vec;
 
 pub(crate) use children_helpers::require_one_child;
+pub use maybe_encoded::MaybeEncoded;
 pub(crate) use once_lock::OnceLockResult;
 pub(crate) use recursion::TreeNodeExt;
 pub(crate) use task_context_helpers::task_ctx_with_extension;

src/coordinator/query_coordinator.rs

@@ -10,11 +10,10 @@ use crate::work_unit_feed::WorkUnitFeedRegistry;
 use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time};
 use crate::worker::LocalWorkerContext;
 use crate::{
-    BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg,
-    DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext,
-    DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage,
-    TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg,
-    get_distributed_channel_resolver, get_distributed_worker_resolver,
+    CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedTaskContext,
+    DistributedWorkUnitFeedContext, LoadInfo, MaybeEncoded, NetworkBoundaryExt, SetPlanRequest,
+    Stage, TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration,
+    WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_worker_resolver,
 };
 use datafusion::common::DataFusionError;
 use datafusion::common::instant::Instant;
@@ -25,10 +24,7 @@ use datafusion::execution::TaskContext;
 use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionConfig;
-use datafusion_proto::physical_plan::AsExecutionPlan;
-use datafusion_proto::protobuf::PhysicalPlanNode;
 use futures::{Stream, StreamExt, TryStreamExt};
-use prost::Message;
 use rand::Rng;
 use std::ops::DerefMut;
 use std::sync::{Arc, Mutex};
@@ -49,6 +45,7 @@ const WORK_UNIT_FEED_CHUNK_SIZE: usize = 256;
 /// [StageCoordinator] scoped to each individual stage.
 pub(super) struct QueryCoordinator {
     task_ctx: Arc<TaskContext>,
+    metrics: ExecutionPlanMetricsSet,
     coordinator_to_worker_metrics: CoordinatorToWorkerMetrics,
     metrics_store: Option<Arc<MetricsStore>>,
     end_stream_notifier: Arc<Notify>,
@@ -64,6 +61,7 @@ impl QueryCoordinator {
     ) -> Self {
         Self {
             task_ctx,
+            metrics: metrics_set.clone(),
             metrics_store,
             coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set),
             end_stream_notifier: Arc::new(Notify::new()),
@@ -80,6 +78,7 @@ impl QueryCoordinator {
             stage_id: stage.num,
             task_count: stage.tasks,
             task_ctx: &self.task_ctx,
+            metrics_set: &self.metrics,
             metrics: &self.coordinator_to_worker_metrics,
             metrics_store: &self.metrics_store,
             end_stream_notifier: &self.end_stream_notifier,
@@ -124,6 +123,7 @@ pub(super) struct StageCoordinator<'a> {
     stage_id: usize,
     task_count: usize,
     task_ctx: &'a Arc<TaskContext>,
+    metrics_set: &'a ExecutionPlanMetricsSet,
     metrics: &'a CoordinatorToWorkerMetrics,
     metrics_store: &'a Option<Arc<MetricsStore>>,
     end_stream_notifier: &'a Arc<Notify>,
@@ -143,28 +143,23 @@ impl<'a> StageCoordinator<'a> {
         UnboundedReceiver<WorkerToCoordinatorMsg>,
     )> {
         let session_config = self.task_ctx.session_config();
-        let codec = DistributedCodec::new_combined_with_user(session_config);
 
         let (specialized, work_unit_feed_declarations) = self.task_specialized_plan(task_i)?;
 
-        let plan_proto =
-            PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec();
-        let plan_size = plan_proto.len();
-
         let task_key = TaskKey {
             query_id: self.query_id,
             stage_id: self.stage_id,
             task_number: task_i,
         };
 
-        let msg = CoordinatorToWorkerMsg::SetPlanRequest(SetPlanRequest {
+        let set_plan_request = SetPlanRequest {
             task_key,
             task_count: self.task_count,
-            plan_proto,
+            plan: MaybeEncoded::Decoded(specialized),
             work_unit_feed_declarations,
             target_worker_url: url.clone(),
             query_start_time_ns: self.metrics.instantiation_time,
-        });
+        };
 
         let (coordinator_to_worker_tx, coordinator_to_worker_rx) =
             tokio::sync::mpsc::unbounded_channel();
@@ -176,8 +171,7 @@ impl<'a> StageCoordinator<'a> {
         let mut headers = get_config_extension_propagation_headers(session_config)?;
         headers.extend(get_passthrough_headers(session_config));
 
-        let coordinator_to_worker_stream = futures::stream::once(async { msg })
-            .chain(UnboundedReceiverStream::new(coordinator_to_worker_rx))
+        let coordinator_to_worker_stream = UnboundedReceiverStream::new(coordinator_to_worker_rx)
             .map(set_work_unit_send_time)
             // Keep the request side of the channel open until the query ends: this tail emits
             // no messages and only completes, once the `Notify` fires. Workers interpret this
@@ -193,15 +187,22 @@ impl<'a> StageCoordinator<'a> {
             .boxed();
 
         let metrics = self.metrics.clone();
+        let metrics_set = self.metrics_set.clone();
+        let task_ctx = Arc::clone(self.task_ctx);
 
         self.join_set.lock().unwrap().spawn(async move {
             let start = Instant::now();
             let mut client = channel_resolver.get_worker_client_for_url(&url).await?;
             let mut worker_to_coordinator_stream = client
-                .coordinator_channel(headers, coordinator_to_worker_stream)
+                .coordinator_channel(
+                    headers,
+                    set_plan_request,
+                    coordinator_to_worker_stream,
+                    &task_ctx,
+                    metrics_set,
+                )
                 .await?;
             metrics.plan_send_latency.record(&start);
-            metrics.plan_bytes_sent.add_bytes(plan_size);
             while let Some(msg) = worker_to_coordinator_stream.try_next().await? {
                 if worker_to_coordinator_tx.send(msg).is_err() {
                     break; // receiver dropped
@@ -454,7 +455,6 @@ impl Drop for NotifyGuard {
 /// Metrics that measure network details about communications between [DistributedExec] and a worker.
 #[derive(Clone)]
 pub(super) struct CoordinatorToWorkerMetrics {
-    pub(super) plan_bytes_sent: BytesCounterMetric,
     pub(super) plan_send_latency: Arc<LatencyMetric>,
     pub(super) instantiation_time: usize,
 }
@@ -469,10 +469,6 @@ fn with_task_id_label(builder: MetricBuilder) -> MetricBuilder {
 impl CoordinatorToWorkerMetrics {
     pub(super) fn new(metrics: &ExecutionPlanMetricsSet) -> Self {
         Self {
-            // Metric that measures to total sum of bytes worth of subplans sent.
-            plan_bytes_sent: MetricBuilder::new(metrics)
-                .with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0"))
-                .bytes_counter("plan_bytes_sent"),
             // Latency statistics about the network calls issued to the workers for feeding subplans.
             plan_send_latency: Arc::new(LatencyMetric::new(
                 "plan_send_latency",
@@ -502,8 +498,7 @@ mod tests {
     /// use-after-free that only reproduced in optimized abort builds. Passing
     /// the builder as a named `fn` ([`with_task_id_label`]) sidesteps it.
     ///
-    /// `CoordinatorToWorkerMetrics::new` registers three labeled metrics —
-    /// `plan_bytes_sent` and the latency `_max`/`_avg` pair — each of which
+    /// `CoordinatorToWorkerMetrics::new` registers the latency `_max`/`_avg` pair, each of which
     /// must own a distinct heap buffer holding exactly its single `task_id`
     /// label. The miscompile is observable as the `_avg` buffer aliasing the
     /// `_max` buffer with length 2.
@@ -530,9 +525,9 @@ mod tests {
 
             assert_eq!(
                 labeled.len(),
-                3,
-                "iteration {iteration}: expected 3 labeled metrics \
-                 (plan_bytes_sent, plan_send_latency_max, plan_send_latency_avg), \
+                2,
+                "iteration {iteration}: expected 2 labeled metrics \
+                 (plan_send_latency_max, plan_send_latency_avg), \
                  got {labeled:x?}"
             );
 

src/distributed_planner/mod.rs

@@ -14,8 +14,7 @@ pub use distributed_config::DistributedConfig;
 pub(crate) use inject_network_boundaries::{
     InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, inject_network_boundaries,
 };
-pub(crate) use network_boundary::ProducerHead;
-pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt};
+pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt, ProducerHead};
 pub use session_state_builder_ext::SessionStateBuilderExt;
 pub(crate) use statistics::calculate_cost;
 pub(crate) use task_estimator::{CombinedTaskEstimator, set_distributed_task_estimator};