diff --git a/src/common/maybe_encoded.rs b/src/common/maybe_encoded.rs new file mode 100644 index 00000000..7f7cd228 --- /dev/null +++ b/src/common/maybe_encoded.rs @@ -0,0 +1,121 @@ +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. +/// Users are free to pass [MaybeEncoded::Encoded] or [MaybeEncoded::Decoded] at any +/// moment and Distributed DataFusion's code will internally know how to handle it. +#[derive(Clone)] +pub enum MaybeEncoded { + Encoded(Vec), + Decoded(T), +} + +impl MaybeEncoded { + /// Returns the decoded variant: + /// - If in `Decoded` state, it just passes through the content. + /// - If in `Encoded` state, it decodes using the provided callback. + pub(crate) fn decode_with(self, decode: impl FnOnce(Vec) -> Result) -> Result { + match self { + Self::Encoded(encoded) => decode(encoded), + Self::Decoded(decoded) => Ok(decoded), + } + } + + /// Returns the decoded variant: + /// - If in `Decoded` state, it just passes through the content. + /// - If in `Encoded` state, it throws an error. + pub(crate) fn try_decoded(self) -> Result { + match self { + Self::Encoded(_) => { + internal_err!( + "Expected MaybeDecoded::Decoded({}), but got MaybeEncoded::Decoded", + std::any::type_name::() + ) + } + Self::Decoded(decoded) => Ok(decoded), + } + } +} + +impl MaybeEncoded> { + /// Returns the encoded [ExecutionPlan] as protobuf bytes: + /// - If in `Decoded` state, it encodes it using the codecs registered in the [TaskContext]. + /// - If in `Encoded` state, it just passes through the content. + pub fn encode(self, ctx: &Arc) -> Result> { + match self { + Self::Encoded(encoded) => Ok(encoded), + Self::Decoded(plan) => { + let codec = DistributedCodec::new_combined_with_user(ctx.session_config()); + protobuf::PhysicalPlanNode::try_from_physical_plan(plan, &codec) + .map(|v| v.encode_to_vec()) + } + } + } + + /// Returns the decoded [ExecutionPlan]. + /// - If in `Decoded` state, it just passes through the content. + /// - If in `Encoded` state, it decodes it using the codecs registered in the [TaskContext]. + pub(crate) fn decode(self, task_ctx: &TaskContext) -> Result> { + 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 { + /// Returns the encoded [Partitioning] as protobuf bytes: + /// - If in `Decoded` state, it encodes it using the codecs registered in the [TaskContext]. + /// - If in `Encoded` state, it just passes through the content. + pub fn encode(self, ctx: &Arc) -> Result> { + 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()) + } + } + } + + /// Returns the decoded [Partitioning]. + /// - If in `Decoded` state, it just passes through the content. + /// - If in `Encoded` state, it decodes it using the codecs registered in the [TaskContext]. + pub fn decode(self, schema: SchemaRef, task_ctx: &TaskContext) -> Result { + 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")) + }) + } +} diff --git a/src/common/mod.rs b/src/common/mod.rs index 07463ca0..7dfd3a96 100644 --- a/src/common/mod.rs +++ b/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; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 8c862862..5df1f709 100644 --- a/src/coordinator/query_coordinator.rs +++ b/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, TaskKey, - WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, get_distributed_channel_resolver, - get_distributed_worker_resolver, + CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedTaskContext, + DistributedWorkUnitFeedContext, LoadInfo, MaybeEncoded, NetworkBoundaryExt, SetPlanRequest, + Stage, TaskKey, 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, + metrics: ExecutionPlanMetricsSet, coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, metrics_store: Option>, end_stream_notifier: Arc, @@ -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, + metrics_set: &'a ExecutionPlanMetricsSet, metrics: &'a CoordinatorToWorkerMetrics, metrics_store: &'a Option>, end_stream_notifier: &'a Arc, @@ -143,28 +143,23 @@ impl<'a> StageCoordinator<'a> { UnboundedReceiver, )> { 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, + metrics_set, + &task_ctx, + ) .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 @@ -453,7 +454,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, pub(super) instantiation_time: usize, } @@ -468,10 +468,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", @@ -501,8 +497,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. @@ -529,9 +524,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?}" ); diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index 57975991..c738069b 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -14,7 +14,6 @@ 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; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index 3b7bd3b1..235e4eee 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -1,7 +1,6 @@ use crate::execution_plans::SamplerExec; -use crate::protocol::ProducerHeadSpec; use crate::{ - BroadcastExec, DistributedCodec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, + BroadcastExec, MaybeEncoded, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage, }; use datafusion::arrow::datatypes::SchemaRef; @@ -10,13 +9,6 @@ use datafusion::execution::TaskContext; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; -use datafusion::prelude::SessionConfig; -use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; -use datafusion_proto::physical_plan::to_proto::serialize_partitioning; -use datafusion_proto::physical_plan::{DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext}; -use datafusion_proto::protobuf; -use datafusion_proto::protobuf::proto_error; -use prost::Message; use std::sync::Arc; /// This trait represents a node that introduces the necessity of a network boundary in the plan. @@ -40,13 +32,16 @@ pub trait NetworkBoundary: ExecutionPlan { /// Defines what shape should the head node of a stage have upon getting executed. Depending /// on the [NetworkBoundary] implementation, the stage below should have different head nodes. +#[derive(Clone)] pub enum ProducerHead { /// No specific head node is necessary. None, /// The head node should be a [BroadcastExec]. BroadcastExec { output_partitions: usize }, /// The head node should be a [RepartitionExec]. - RepartitionExec { partitioning: Partitioning }, + RepartitionExec { + partitioning: MaybeEncoded, + }, } /// Extension trait for downcasting dynamic types to [NetworkBoundary]. @@ -74,49 +69,13 @@ impl NetworkBoundaryExt for dyn ExecutionPlan { } impl ProducerHead { - pub(crate) fn to_spec(&self, cfg: &SessionConfig) -> Result { - match self { - Self::None => Ok(ProducerHeadSpec::None), - Self::BroadcastExec { output_partitions } => Ok(ProducerHeadSpec::BroadcastExec { - output_partitions: *output_partitions, - }), - Self::RepartitionExec { partitioning } => { - let partitioning = serialize_partitioning( - partitioning, - &DistributedCodec::new_combined_with_user(cfg), - &DefaultPhysicalProtoConverter {}, - ) - .map(|v| v.encode_to_vec())?; - Ok(ProducerHeadSpec::RepartitionExec { partitioning }) - } - } - } - - pub(crate) fn from_spec( - spec: &ProducerHeadSpec, - schema: SchemaRef, - ctx: &TaskContext, - ) -> Result { - match spec { - ProducerHeadSpec::None => Ok(Self::None), - ProducerHeadSpec::BroadcastExec { output_partitions } => Ok(Self::BroadcastExec { - output_partitions: *output_partitions, - }), - ProducerHeadSpec::RepartitionExec { partitioning } => { - let proto_partitioning = protobuf::Partitioning::decode(partitioning.as_slice()) - .map_err(|e| proto_error(e.to_string()))?; - let codec = DistributedCodec::new_combined_with_user(ctx.session_config()); - let decode_ctx = PhysicalPlanDecodeContext::new(ctx, &codec); - let partitioning = parse_protobuf_partitioning( - Some(&proto_partitioning), - &decode_ctx, - &schema, - &DefaultPhysicalProtoConverter {}, - )? - .ok_or_else(|| proto_error("Could not parse partitioning"))?; - Ok(Self::RepartitionExec { partitioning }) - } - } + pub(crate) fn resolve(self, schema: SchemaRef, ctx: &TaskContext) -> Result { + Ok(match self { + Self::RepartitionExec { partitioning } => Self::RepartitionExec { + partitioning: MaybeEncoded::Decoded(partitioning.decode(schema, ctx)?), + }, + v => v, + }) } /// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This @@ -135,9 +94,10 @@ impl ProducerHead { let partitions = input.output_partitioning().partition_count(); Arc::new(BroadcastExec::new(input, output_partitions / partitions)) } - ProducerHead::RepartitionExec { partitioning } => { - Arc::new(RepartitionExec::try_new(input, partitioning)?) - } + ProducerHead::RepartitionExec { partitioning } => Arc::new(RepartitionExec::try_new( + input, + partitioning.try_decoded()?, + )?), }; Ok(plan) } diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index c7a645ae..c11d3e6b 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -3,7 +3,7 @@ use crate::distributed_planner::ProducerHead; use crate::execution_plans::common::scale_partitioning; use crate::stage::{LocalStage, Stage}; use crate::worker::WorkerConnectionPool; -use crate::{DistributedTaskContext, NetworkBoundary}; +use crate::{DistributedTaskContext, MaybeEncoded, NetworkBoundary}; use datafusion::common::{Result, not_impl_err, plan_err}; use datafusion::error::DataFusionError; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -156,9 +156,10 @@ impl NetworkBoundary for NetworkShuffleExec { fn producer_head(&self, consumer_task_count: usize) -> ProducerHead { ProducerHead::RepartitionExec { - partitioning: scale_partitioning(&self.properties.partitioning, |prev| { - prev * consumer_task_count - }), + partitioning: MaybeEncoded::Decoded(scale_partitioning( + &self.properties.partitioning, + |prev| prev * consumer_task_count, + )), } } } diff --git a/src/lib.rs b/src/lib.rs index b132caf5..e4c4a7fc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,7 +20,7 @@ pub use arrow_ipc::CompressionType; pub use coordinator::DistributedExec; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ - DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, + DistributedConfig, NetworkBoundary, NetworkBoundaryExt, ProducerHead, SessionStateBuilderExt, }; pub use events::{ DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTasksEvent, @@ -47,12 +47,13 @@ pub mod test_utils; pub use protocol::grpc; pub use codec::DistributedCodec; +pub use common::MaybeEncoded; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, - GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, - WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + GetWorkerInfoResponse, LoadInfo, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, get_distributed_channel_resolver, }; pub use stage::{ diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 56612671..f8f5ff1d 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -8,9 +8,9 @@ use crate::grpc::on_drop_stream::on_drop_stream; use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, - MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHeadSpec, - SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, - WorkerChannel, WorkerToCoordinatorMsg, + MaxLatencyMetric, MaybeEncoded, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, + ProducerHead, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, + WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; use arrow_flight::decode::FlightRecordBatchStream; @@ -47,9 +47,21 @@ impl WorkerChannel for pb::worker_service_client::WorkerServiceClient, + metrics: ExecutionPlanMetricsSet, + ctx: &Arc, ) -> Result>> { - let input_stream = c2w_stream.map(encode_coordinator_to_worker_msg); + let set_plan_request = encode_set_plan_request(set_plan_request, ctx)?; + let plan_bytes_sent = set_plan_request.plan_proto.len(); + let input_stream = futures::stream::once(async move { + pb::CoordinatorToWorkerMsg { + inner: Some(pb::coordinator_to_worker_msg::Inner::SetPlanRequest( + set_plan_request, + )), + } + }) + .chain(c2w_stream.map(encode_coordinator_to_worker_msg)); let output_stream = self .coordinator_channel(Request::from_parts( @@ -65,6 +77,10 @@ impl WorkerChannel for pb::worker_service_client::WorkerServiceClient pb::execute_task_request::ProducerHead { - match head { - ProducerHeadSpec::None => pb::execute_task_request::ProducerHead::None(pb::NoneHead {}), - ProducerHeadSpec::BroadcastExec { output_partitions } => { +pub(super) fn encode_producer_head( + head: ProducerHead, + ctx: &Arc, +) -> Result { + Ok(match head { + ProducerHead::None => pb::execute_task_request::ProducerHead::None(pb::NoneHead {}), + ProducerHead::BroadcastExec { output_partitions } => { pb::execute_task_request::ProducerHead::Broadcast(pb::BroadcastExecHead { output_partitions: output_partitions as u64, }) } - ProducerHeadSpec::RepartitionExec { partitioning } => { + ProducerHead::RepartitionExec { partitioning } => { pb::execute_task_request::ProducerHead::Repartition(pb::RepartitionExecHead { - partitioning, + partitioning: partitioning.encode(ctx)?, }) } - } + }) } fn encode_coordinator_to_worker_msg(msg: CoordinatorToWorkerMsg) -> pb::CoordinatorToWorkerMsg { pb::CoordinatorToWorkerMsg { inner: Some(match msg { - CoordinatorToWorkerMsg::SetPlanRequest(request) => { - pb::coordinator_to_worker_msg::Inner::SetPlanRequest(encode_set_plan_request( - request, - )) - } CoordinatorToWorkerMsg::WorkUnitBatch(batch) => { pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(encode_work_unit_batch(batch)) } @@ -420,11 +432,15 @@ fn encode_coordinator_to_worker_msg(msg: CoordinatorToWorkerMsg) -> pb::Coordina } } -fn encode_set_plan_request(request: SetPlanRequest) -> pb::SetPlanRequest { - pb::SetPlanRequest { +fn encode_set_plan_request( + request: SetPlanRequest, + ctx: &Arc, +) -> Result { + let plan_proto = request.plan.encode(ctx)?; + Ok(pb::SetPlanRequest { task_key: Some(encode_task_key(request.task_key)), task_count: request.task_count as u64, - plan_proto: request.plan_proto, + plan_proto, work_unit_feed_declarations: request .work_unit_feed_declarations .into_iter() @@ -432,7 +448,7 @@ fn encode_set_plan_request(request: SetPlanRequest) -> pb::SetPlanRequest { .collect(), target_worker_url: request.target_worker_url.to_string(), query_start_time_ns: request.query_start_time_ns as u64, - } + }) } fn encode_work_unit_batch(batch: WorkUnitBatch) -> pb::WorkUnitBatch { @@ -445,7 +461,10 @@ fn encode_work_unit(work_unit: WorkUnitMsg) -> pb::WorkUnit { pb::WorkUnit { id: serialize_uuid(&work_unit.id), partition: work_unit.partition as u64, - body: work_unit.body, + body: match work_unit.body { + MaybeEncoded::Encoded(body) => body, + MaybeEncoded::Decoded(body) => body.encode_to_bytes(), + }, created_timestamp_unix_nanos: work_unit.created_timestamp_unix_nanos as u64, sent_timestamp_unix_nanos: work_unit.sent_timestamp_unix_nanos as u64, received_timestamp_unix_nanos: work_unit.received_timestamp_unix_nanos as u64, diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index 6fba14ee..1f0fdf0a 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -4,12 +4,11 @@ use super::metrics_proto::df_metrics_set_to_proto; use super::spawn_select_all::spawn_select_all; use crate::common::{deserialize_uuid, now_ns}; -use crate::protocol::ProducerHeadSpec; use crate::protocol::grpc::{ObservabilityServiceImpl, ObservabilityServiceServer}; use crate::{ - CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, SetPlanRequest, - TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, Worker, - WorkerResolver, WorkerToCoordinatorMsg, + CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, LoadInfo, MaybeEncoded, + ProducerHead, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, + WorkUnitMsg, Worker, WorkerResolver, WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; @@ -93,7 +92,21 @@ impl pb::worker_service_server::WorkerService for Worker { &self, request: Request>, ) -> Result, Status> { - let (metadata, _ext, body) = request.into_parts(); + let (metadata, _ext, mut body) = request.into_parts(); + + let msg = body + .message() + .await? + .ok_or_else(empty("Coordinator stream"))? + .inner + .ok_or_else(missing("CoordinatorToWorkerMsg.inner"))?; + let pb::coordinator_to_worker_msg::Inner::SetPlanRequest(set_plan_request) = msg else { + return Err(Status::invalid_argument( + "First Coordinator to Worker message must be SetPlanRequest", + )); + }; + + let set_plan_request = decode_set_plan_request(set_plan_request)?; let input_stream = body .map_err(map_status_to_datafusion_error) @@ -103,7 +116,7 @@ impl pb::worker_service_server::WorkerService for Worker { .boxed(); let output_stream = self - .coordinator_channel(metadata.into_headers(), input_stream) + .coordinator_channel(metadata.into_headers(), set_plan_request, input_stream) .await .map_err(datafusion_error_to_tonic_status)? .map(|msg| match msg { @@ -188,8 +201,10 @@ fn decode_coordinator_to_worker_msg( .inner .ok_or_else(missing("CoordinatorToWorkerMsg.inner"))? { - pb::coordinator_to_worker_msg::Inner::SetPlanRequest(request) => { - CoordinatorToWorkerMsg::SetPlanRequest(decode_set_plan_request(request)?) + pb::coordinator_to_worker_msg::Inner::SetPlanRequest(_) => { + return Err(Status::invalid_argument( + "SetPlanRequest must be the first coordinator message", + )); } pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(batch) => { CoordinatorToWorkerMsg::WorkUnitBatch(decode_work_unit_batch(batch)?) @@ -205,7 +220,7 @@ fn decode_set_plan_request(request: pb::SetPlanRequest) -> Result ProducerHeadSpec { +pub(super) fn decode_producer_head(proto: pb::execute_task_request::ProducerHead) -> ProducerHead { match proto { - pb::execute_task_request::ProducerHead::None(_) => ProducerHeadSpec::None, - pb::execute_task_request::ProducerHead::Broadcast(v) => ProducerHeadSpec::BroadcastExec { + pb::execute_task_request::ProducerHead::None(_) => ProducerHead::None, + pb::execute_task_request::ProducerHead::Broadcast(v) => ProducerHead::BroadcastExec { output_partitions: v.output_partitions as usize, }, - pb::execute_task_request::ProducerHead::Repartition(v) => { - ProducerHeadSpec::RepartitionExec { - partitioning: v.partitioning, - } - } + pb::execute_task_request::ProducerHead::Repartition(v) => ProducerHead::RepartitionExec { + partitioning: MaybeEncoded::Encoded(v.partitioning), + }, } } @@ -311,7 +322,7 @@ fn decode_work_unit(work_unit: pb::WorkUnit) -> Result { Ok(WorkUnitMsg { id: deserialize_uuid(&work_unit.id).map_err(datafusion_error_to_tonic_status)?, partition: work_unit.partition as usize, - body: work_unit.body, + body: MaybeEncoded::Encoded(work_unit.body), created_timestamp_unix_nanos: work_unit.created_timestamp_unix_nanos as usize, sent_timestamp_unix_nanos: work_unit.sent_timestamp_unix_nanos as usize, received_timestamp_unix_nanos: work_unit.received_timestamp_unix_nanos as usize, @@ -341,6 +352,10 @@ fn parse_url(value: &str, field: &'static str) -> Result { .map_err(|err| Status::invalid_argument(format!("Invalid field '{field}': {err}"))) } +fn empty(stream_name: &'static str) -> impl FnOnce() -> Status { + move || Status::invalid_argument(format!("Empty {stream_name}")) +} + fn missing(field: &'static str) -> impl FnOnce() -> Status { move || Status::invalid_argument(format!("Missing field '{field}'")) } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 6df0e11b..d026e7e2 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -9,6 +9,6 @@ pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, - LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, - WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + LoadInfo, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, + WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, }; diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 3ad3816c..0bea0bbe 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -1,7 +1,9 @@ +use crate::{MaybeEncoded, ProducerHead, WorkUnit}; use async_trait::async_trait; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::Result; use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use futures::stream::BoxStream; use http::HeaderMap; @@ -21,7 +23,10 @@ pub trait WorkerChannel: Send + Sync { async fn coordinator_channel( &mut self, headers: HeaderMap, + set_plan_request: SetPlanRequest, c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + metrics: ExecutionPlanMetricsSet, + task_ctx: &Arc, ) -> Result>>; /// Executes the requested partition range of a subplan previously sent by the coordinator channel. @@ -41,9 +46,6 @@ pub trait WorkerChannel: Send + Sync { } pub enum CoordinatorToWorkerMsg { - /// Sends a subplan to a worker so that a future ExecuteTask call can actually execute it. - /// The plan is identified by a TaskKey. - SetPlanRequest(SetPlanRequest), /// A batch of messages from a work unit feed belonging to different partitions from one node from the plan set in /// set_plan_request. A work unit feed is a per-partition stream of information that tells the node what should /// be executed within a partition, for example, a stream of file addresses that should be read. @@ -75,10 +77,7 @@ pub struct SetPlanRequest { /// The amount of tasks that share the same subplan. Necessary for building the DistributedTaskContext during execution. pub task_count: usize, /// The subplan the worker is expected to execute. - // TODO: this still forces implementations to pass a serialized plan. In-memory implementations - // might want to omit the serde step, so there should be a way to pass here a normal plan, and - // pass the serializer/deserialized separately instead of being coupled to protobuf serialization - pub plan_proto: Vec, + pub plan: MaybeEncoded>, /// Information about all the work unit feeds that will be streamed from coordinator to worker. /// This information is needed here because at the moment of setting the plan, all the appropriate /// channels for the incoming work unit feeds need to be constructed. @@ -104,7 +103,7 @@ pub struct WorkUnitMsg { /// The partition index within the node to which the work unit feed belongs to. pub partition: usize, /// Arbitrary user-defined data (e.g., a file address) necessary during execution. - pub body: Vec, + pub body: MaybeEncoded>, /// Unix timestamp in nanoseconds at which this message was created in the coordinator. pub created_timestamp_unix_nanos: usize, /// Unix timestamp in nanoseconds at which this message was sent by the coordinator. @@ -173,17 +172,7 @@ pub struct ExecuteTaskRequest { /// - A RepartitionExecHead implies a RepartitionExec at the head of the task. /// - A BroadcastExecHead implies a BroadcastExec at the head of the task. /// - A NoneHead does not need any specific head. - pub producer_head_spec: ProducerHeadSpec, -} - -#[derive(Clone)] -pub enum ProducerHeadSpec { - /// No specific head node is necessary. - None, - /// The head node should be a [BroadcastExec]. - BroadcastExec { output_partitions: usize }, - /// The head node should be a [RepartitionExec]. - RepartitionExec { partitioning: Vec }, + pub producer_head: ProducerHead, } pub struct GetWorkerInfoRequest {} diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index 18c51ffe..5b73ef73 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -1,11 +1,13 @@ use crate::common::now_ns; use crate::{ - BytesMetricExt, CoordinatorToWorkerMsg, LatencyMetricExt, WorkUnit, WorkUnitBatch, WorkUnitMsg, + BytesMetricExt, CoordinatorToWorkerMsg, LatencyMetricExt, MaybeEncoded, WorkUnit, + WorkUnitBatch, WorkUnitMsg, }; use datafusion::common::{HashMap, Result, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::MetricBuilder; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_proto::protobuf::proto_error; use futures::StreamExt; use futures::stream::BoxStream; use std::sync::{Arc, Mutex}; @@ -68,7 +70,7 @@ pub(crate) fn build_work_unit_batch_msg( Ok(WorkUnitMsg { id: *id, partition, - body: work_unit?.encode_to_bytes(), + body: MaybeEncoded::Decoded(work_unit?), created_timestamp_unix_nanos: now_ns(), sent_timestamp_unix_nanos: 0, received_timestamp_unix_nanos: 0, @@ -122,6 +124,7 @@ impl RemoteFeedProvider { let bdr = || MetricBuilder::new(&self.metrics); let bytes_transferred = bdr().bytes_counter("work_unit_bytes"); + let in_memory_transferred = bdr().global_counter("work_unit_in_memory_count"); let msg_count = bdr().global_counter("work_unit_count"); // Track end-to-end network latency distribution for all work units. let send_latency_max = bdr().max_latency("work_unit_send_latency_max"); @@ -159,11 +162,26 @@ impl RemoteFeedProvider { .map(move |work_unit_msg_or_err| { let mut work_unit_msg = work_unit_msg_or_err?; let timer = elapsed_compute.timer(); - let work_unit = T::decode(work_unit_msg.body.as_slice()) - .map_err(|err| datafusion_proto::protobuf::proto_error(format!("{err}"))); + let work_unit = match work_unit_msg.body { + MaybeEncoded::Encoded(bytes) => { + bytes_transferred.add_bytes(bytes.len()); + T::decode(bytes.as_slice()).map_err(|err| proto_error(format!("{err}")))? + } + MaybeEncoded::Decoded(work_unit) => { + let work_unit = work_unit.into_any(); + let Ok(work_unit) = work_unit.downcast::() else { + return exec_err!( + "Expected WorkUnit of type {}", + std::any::type_name::() + ); + }; + in_memory_transferred.add(1); + *work_unit + } + }; + timer.done(); work_unit_msg.processed_timestamp_unix_nanos = now_ns(); - let body_len = work_unit_msg.body.len(); let WorkUnitMsg { created_timestamp_unix_nanos: base, @@ -173,7 +191,6 @@ impl RemoteFeedProvider { .. } = work_unit_msg; - bytes_transferred.add_bytes(body_len); msg_count.add(1); send_latency_max.add_nanos(sent_timestamp_unix_nanos - base); @@ -185,7 +202,7 @@ impl RemoteFeedProvider { processed_latency_max.add_nanos(processed_timestamp_unix_nanos - base); processed_latency_p50.add_nanos(processed_timestamp_unix_nanos - base); - work_unit + Ok(work_unit) }) .boxed()) } diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index cf16e890..c9e773ca 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -4,17 +4,14 @@ use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_t use crate::worker::LocalWorkerContext; use crate::worker::task_data::TaskDataMetrics; use crate::{ - CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, - DistributedTaskContext, TaskData, TaskMetrics, Worker, WorkerQueryContext, - WorkerToCoordinatorMsg, + CoordinatorToWorkerMsg, DistributedConfig, DistributedExt, DistributedTaskContext, + SetPlanRequest, TaskData, TaskMetrics, Worker, WorkerQueryContext, WorkerToCoordinatorMsg, }; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err}; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; -use datafusion_proto::physical_plan::AsExecutionPlan; -use datafusion_proto::protobuf::PhysicalPlanNode; use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; use http::HeaderMap; @@ -26,17 +23,9 @@ impl Worker { pub async fn coordinator_channel( &self, headers: HeaderMap, - mut stream: BoxStream<'static, Result>, + request: SetPlanRequest, + stream: BoxStream<'static, Result>, ) -> Result>> { - // The first message must be a SetPlanRequest. - let Some(msg) = stream.try_next().await? else { - return internal_err!("Empty Coordinator stream"); - }; - - let CoordinatorToWorkerMsg::SetPlanRequest(request) = msg else { - return internal_err!("First Coordinator message must be SetPlanRequest"); - }; - let key = request.task_key; let entry = self @@ -83,10 +72,8 @@ impl Worker { }) .await?; - let codec = DistributedCodec::new_combined_with_user(session_state.config()); let task_ctx = session_state.task_ctx(); - let proto_node = PhysicalPlanNode::try_decode(request.plan_proto.as_ref())?; - let mut plan = proto_node.try_into_physical_plan(&task_ctx, &codec)?; + let mut plan = request.plan.decode(&task_ctx)?; for hook in self.hooks.on_plan.iter() { plan = hook(plan, session_state.config())?; @@ -135,11 +122,6 @@ impl Worker { let mut stream = stream.map_ok(set_work_unit_received_time); while let Some(Ok(msg)) = stream.next().await { match msg { - CoordinatorToWorkerMsg::SetPlanRequest(_) => { - // SetPlanRequest should be the first already polled message in the stream, - // if some reached here it means that something is wrong. - continue; - } CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) => { let Some(work_unit_senders) = work_unit_senders.as_mut() else { continue; diff --git a/src/worker/impl_execute_task.rs b/src/worker/impl_execute_task.rs index eff0bd98..75d7e6d5 100644 --- a/src/worker/impl_execute_task.rs +++ b/src/worker/impl_execute_task.rs @@ -40,7 +40,7 @@ impl Worker { .map_err(DataFusionError::Shared)?; task_data.task_data_metrics.mark_execution_started_once(); - let plan = task_data.plan(&request.producer_head_spec)?; + let plan = task_data.plan(request.producer_head)?; let task_ctx = task_data.task_ctx; let partition_count = plan.properties().partitioning.partition_count(); let plan_name = plan.name(); diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index c9085263..81320e91 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -1,8 +1,6 @@ use crate::common::OnceLockResult; use crate::common::now_ns; -use crate::distributed_planner::ProducerHead; -use crate::protocol::ProducerHeadSpec; -use crate::{MaxLatencyMetric, TaskMetrics}; +use crate::{MaxLatencyMetric, ProducerHead, TaskMetrics}; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; @@ -103,16 +101,9 @@ fn max_latency_metric(name: &'static str, value: &MaxLatencyMetric) -> Arc Result> { + pub(crate) fn plan(&self, producer_head: ProducerHead) -> Result> { let result = self.final_plan.get_or_init(|| { - let producer_head = ProducerHead::from_spec( - producer_head_spec, - self.base_plan.schema(), - &self.task_ctx, - )?; + let producer_head = producer_head.resolve(self.base_plan.schema(), &self.task_ctx)?; Ok(producer_head.insert(Arc::clone(&self.base_plan))?) }); diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 7fdb216a..354dc198 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -89,7 +89,7 @@ impl WorkerConnectionPool { task_key, target_partition, &target_url, - &producer_head, + producer_head.clone(), ctx, )? { return Ok(result @@ -125,7 +125,7 @@ impl WorkerConnectionPool { task_key, target_partition_start: target_partitions.start, target_partition_end: target_partitions.end, - producer_head_spec: producer_head.to_spec(ctx.session_config())?, + producer_head, }; let mut client = ch_resolver.get_worker_client_for_url(&target_url).await?; let headers = get_passthrough_headers(ctx.session_config()); @@ -175,7 +175,7 @@ impl WorkerConnectionPool { task_key: TaskKey, target_partition: usize, target_url: &Url, - producer_head: &ProducerHead, + producer_head: ProducerHead, ctx: &Arc, ) -> Result>>> { let Some(task_data_entries) = ctx @@ -196,7 +196,7 @@ impl WorkerConnectionPool { task_key, target_partition_start: target_partition, target_partition_end: target_partition + 1, - producer_head_spec: producer_head.to_spec(ctx.session_config())?, + producer_head, }; // The relevant entry from `task_data_entries` needs to be eagerly retrieved, it cannot be // left for until someone decides to start polling the returned `BoxStream`, otherwise,