From becc5be6a4b4e6120610ee69df2aa229d3fc8189 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Sat, 30 May 2026 08:36:32 +0200 Subject: [PATCH 01/20] Add NetworkBoundaryBuilder argument to inject_network_boundaries.rs From 77b00f665ed28541f89aff21f61da054153000ff Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Mon, 8 Jun 2026 12:47:20 +0200 Subject: [PATCH 02/20] Add a cost model for the physical layer --- src/distributed_planner/distributed_config.rs | 4 + src/distributed_planner/mod.rs | 1 + .../statistics/compute_per_node.rs | 1029 +++++++++++++++++ src/distributed_planner/statistics/cost.rs | 29 + .../statistics/default_bytes_for_datatype.rs | 137 +++ src/distributed_planner/statistics/mod.rs | 7 + .../statistics/plan_statistics.rs | 160 +++ src/test_utils/plans.rs | 39 + 8 files changed, 1406 insertions(+) create mode 100644 src/distributed_planner/statistics/compute_per_node.rs create mode 100644 src/distributed_planner/statistics/cost.rs create mode 100644 src/distributed_planner/statistics/default_bytes_for_datatype.rs create mode 100644 src/distributed_planner/statistics/mod.rs create mode 100644 src/distributed_planner/statistics/plan_statistics.rs diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index d6faee39..6795fd95 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -66,6 +66,10 @@ extensions_options! { /// budget will still be admitted (otherwise we would livelock), so the actual peak per /// connection is `worker_connection_buffer_budget_bytes + max_message_size`. pub worker_connection_buffer_budget_bytes: usize, default = 64 * 1024 * 1024 + /// Distributed DataFusion relies on row count estimation in order to infer how many workers + /// should be used in serving the query. Some plans might not implement any kind of row count + /// estimation, and this parameter sets the default estimated row count for those plans. + pub default_estimated_row_count: Option, default = Some(0) /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to /// estimate how many tasks should be spawned for the [Stage] containing the leaf node. pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default() diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index a1a7c9fb..f12c684c 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -7,6 +7,7 @@ mod partial_reduce_below_network_shuffles; mod prepare_network_boundaries; mod push_fetch_into_network_coalesce; mod session_state_builder_ext; +mod statistics; mod task_estimator; pub use distributed_config::DistributedConfig; diff --git a/src/distributed_planner/statistics/compute_per_node.rs b/src/distributed_planner/statistics/compute_per_node.rs new file mode 100644 index 00000000..27b31176 --- /dev/null +++ b/src/distributed_planner/statistics/compute_per_node.rs @@ -0,0 +1,1029 @@ +use crate::BroadcastExec; +use crate::execution_plans::ChildrenIsolatorUnionExec; +use datafusion::catalog::memory::DataSourceExec; +use datafusion::common::{JoinSide, Statistics}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::aggregates::AggregateExec; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{Column, Literal}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion::physical_plan::joins::{ + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, SymmetricHashJoinExec, +}; +use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; +use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +#[derive(Clone, PartialEq, Eq)] +pub(super) enum Complexity { + /// Constant complexity + Constant(usize), + /// Linear with a specific column from a specific child. + Linear(LinearComplexity), + /// NLogM + Log(Box, Box), + /// N+M + Plus(Box, Box), + /// N*M + Multiply(Box, Box), +} + +#[derive(Clone, PartialEq, Eq)] +pub(super) enum LinearComplexity { + /// Depends on linearly with the input column with the provided index + Column(usize), + /// Depends on linearly with the all the input columns + AllColumns, + /// Depends on linearly with the input column with the provided index from the left child + ColumnFromLeft(usize), + /// Depends on linearly with the all the input columns from the left child + AllColumnsFromLeft, + /// Depends on linearly with the input column with the provided index from the right child + ColumnFromRight(usize), + /// Depends on linearly with the all the input columns from the right child + AllColumnsFromRight, + /// Depends on linearly with the all the output columns + AllOutputColumns, +} + +impl Complexity { + fn log(self, other: Self) -> Self { + match (self, other) { + (Self::Constant(n), Self::Constant(m)) => { + Self::Constant(n * (m as f64).log2() as usize) + } + (s, o) => Self::Log(Box::new(s), Box::new(o)), + } + } + + fn plus(self, other: Self) -> Self { + match (self, other) { + (Self::Constant(n), Self::Constant(m)) => Self::Constant(n + m), + // (A + k1) + k2 = A + (k1 + k2): bubble constants rightward so they can fold + (Self::Plus(a, b), Self::Constant(m)) if matches!(*b, Self::Constant(_)) => { + (*a).plus((*b).plus(Self::Constant(m))) + } + (s, o) if s == o => Self::Constant(2).multiply(s), + (s, o) => Self::Plus(Box::new(s), Box::new(o)), + } + } + + fn multiply(self, other: Self) -> Self { + match (self, other) { + (Self::Constant(n), Self::Constant(m)) => Self::Constant(n * m), + (s, o) => Self::Multiply(Box::new(s), Box::new(o)), + } + } + + /// Computes the total bytes processed given per-child row counts. + /// Returns None if statistics are unavailable for any required input. + pub(super) fn cost( + &self, + output_stat: &Arc, + input_stats: &[Arc], + ) -> Option { + Some(match self { + Self::Constant(v) => *v, + Self::Linear(linear) => match linear { + LinearComplexity::Column(i) => { + let col_stats = &input_stats.first()?.column_statistics; + *col_stats.get(*i)?.byte_size.get_value()? + } + LinearComplexity::AllColumns => { + *input_stats.first()?.total_byte_size.get_value()? + } + LinearComplexity::ColumnFromLeft(i) => { + let col_stats = &input_stats.first()?.column_statistics; + *col_stats.get(*i)?.byte_size.get_value()? + } + LinearComplexity::AllColumnsFromLeft => { + *input_stats.first()?.total_byte_size.get_value()? + } + LinearComplexity::ColumnFromRight(i) => { + let col_stats = &input_stats.last()?.column_statistics; + *col_stats.get(*i)?.byte_size.get_value()? + } + LinearComplexity::AllColumnsFromRight => { + *input_stats.last()?.total_byte_size.get_value()? + } + LinearComplexity::AllOutputColumns => *output_stat.total_byte_size.get_value()?, + }, + Self::Log(n, m) => { + let n = n.cost(output_stat, input_stats)?; + let m = m.cost(output_stat, input_stats)?; + // `ilog2` panics on 0, which happens whenever the logged input has zero estimated + // bytes/rows (e.g. an empty or fully-pruned relation). Flooring at 1 makes log2 + // contribute 0 there, i.e. sorting/merging nothing costs nothing. + n * m.checked_ilog2().unwrap_or(0) as usize + } + Self::Plus(n, m) => { + n.cost(output_stat, input_stats)? + m.cost(output_stat, input_stats)? + } + Self::Multiply(n, m) => { + n.cost(output_stat, input_stats)? * m.cost(output_stat, input_stats)? + } + }) + } +} + +impl Debug for Complexity { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + fn trim_parenthesis(dbg: &Complexity) -> String { + let s = format!("{dbg:?}"); + if s.starts_with('(') && s.ends_with(')') { + s[1..s.len() - 1].to_string() + } else { + s + } + } + match self { + Self::Constant(v) => write!(f, "{v}"), + Self::Linear(linear) => match linear { + LinearComplexity::Column(i) => write!(f, "Col{i}"), + LinearComplexity::AllColumns => write!(f, "Cols"), + LinearComplexity::ColumnFromLeft(i) => write!(f, "left_Col{i}"), + LinearComplexity::AllColumnsFromLeft => write!(f, "left_Cols"), + LinearComplexity::ColumnFromRight(i) => write!(f, "right_Col{i}"), + LinearComplexity::AllColumnsFromRight => write!(f, "right_Cols"), + LinearComplexity::AllOutputColumns => write!(f, "out_Cols"), + }, + Self::Log(n, m) => write!(f, "{n:?}*Log({m:?})"), + Self::Plus(n, m) => { + if matches!(n.as_ref(), &Self::Plus(_, _)) { + write!(f, "({}+{m:?})", trim_parenthesis(n)) + } else { + write!(f, "({n:?}+{m:?})") + } + } + Self::Multiply(n, m) => { + if matches!(n.as_ref(), &Self::Multiply(_, _)) { + write!(f, "({}*{m:?})", trim_parenthesis(n)) + } else { + write!(f, "({n:?}*{m:?})") + } + } + } + } +} + +/// Calculates what's the cost, expressed as a number, per input row for each input children. +/// +/// The Vec return has equal size to `node.children()`, and determines how many each input needs +/// to be processed +pub(super) fn calculate_compute_complexity(node: &Arc) -> Complexity { + // NestedLoopJoinExec: O(n*m) - evaluates join condition for each pair of rows + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/joins/nested_loop_join.rs + if let Some(node) = node.downcast_ref::() { + // Assume we need to do read all input rows one by one. + let n = Complexity::Linear(LinearComplexity::AllColumnsFromLeft); + let m = Complexity::Linear(LinearComplexity::AllColumnsFromRight); + let mut c = n.multiply(m); + // The join condition is evaluated on every (left, right) pair. We can't express the + // exact per-pair cost (it would be filter_cost * n * m), so we add the filter columns + // as a lower-bound refinement; the O(n*m) materialization term above already dominates. + if let Some(filter) = node.filter() { + c = c.plus(join_filter_complexity(filter)); + } + return c; + } + + // CrossJoinExec: O(n*m) - produces Cartesian product of all row pairs + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/joins/cross_join.rs + if let Some(_node) = node.downcast_ref::() { + // Assume we need to do read all input rows one by one. + let n = Complexity::Linear(LinearComplexity::AllColumnsFromLeft); + let m = Complexity::Linear(LinearComplexity::AllColumnsFromRight); + return n.multiply(m); + } + + // SortExec: O(n log n) - uses lexsort_to_indices, may spill to disk + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/sorts/sort.rs + if let Some(node) = node.downcast_ref::() { + // All the rows will need to be copied one by one. + let mut n = Complexity::Linear(LinearComplexity::AllColumns); + // The sort comparators read every sort key on every row, so even a plain column key costs + // its bytes (a wide UTF8 key is far costlier to compare than an int). + for expr in node.expr() { + n = n.plus(hashed_or_sorted_key_complexity(&expr.expr)) + } + return n.clone().log(n); + } + + // HashJoinExec: hash table build (O(n)) + probe (O(m)) + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/joins/hash_join/exec.rs + if let Some(join) = node.downcast_ref::() { + // Build side (left): concat_batches copies all data (2x read), plus hash table storage, + // plus hashing left join keys. + let mut c = Complexity::Linear(LinearComplexity::AllColumnsFromLeft) + .plus(Complexity::Linear(LinearComplexity::AllColumnsFromLeft)); + for (left_key, _) in join.on() { + c = c.plus(join_key_complexity(left_key, true)); + } + // Probe side (right): read all columns + hash right join keys + c = c.plus(Complexity::Linear(LinearComplexity::AllColumnsFromRight)); + for (_, right_key) in join.on() { + c = c.plus(join_key_complexity(right_key, false)); + } + // Optional join filter evaluated on candidate matches during the probe. + if let Some(filter) = join.filter() { + c = c.plus(join_filter_complexity(filter)); + } + return c; + } + + // SortMergeJoinExec: merge of sorted streams with comparisons + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs + // Unlike hash join, sort-merge doesn't buffer all data or build hash tables. It streams + // through both sorted inputs with O(max_group_size) memory, using partial_cmp comparisons + // (no hashing). Per-row cost is just key comparisons + optional filter evaluation. + if let Some(node) = node.downcast_ref::() { + let mut c: Option = None; + // Left side: compare join keys during merge + for (left_key, _) in node.on() { + let key = join_key_complexity(left_key, true); + c = Some(match c { + Some(existing) => existing.plus(key), + None => key, + }); + } + // Right side: compare join keys during merge + for (_, right_key) in node.on() { + let key = join_key_complexity(right_key, false); + c = Some(match c { + Some(existing) => existing.plus(key), + None => key, + }); + } + // Optional join filter evaluated on matched pairs during the merge. + if let Some(filter) = node.filter() { + let f = join_filter_complexity(filter); + c = Some(match c { + Some(existing) => existing.plus(f), + None => f, + }); + } + return c.unwrap_or(Complexity::Constant(1)); + } + + // SymmetricHashJoinExec: streaming join with hash tables on both sides + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/joins/symmetric_hash_join.rs + // More expensive than HashJoinExec: both sides maintain hash tables, concat_batches + // runs on every incoming batch (not once at end), plus pruning interval computation + // and HashSet tracking for visited rows. + if let Some(node) = node.downcast_ref::() { + // Both sides: concat_batches on every batch (2x read) + hash table + hash keys + let mut c = Complexity::Linear(LinearComplexity::AllColumnsFromLeft) + .plus(Complexity::Linear(LinearComplexity::AllColumnsFromLeft)); + for (left_key, _) in node.on() { + c = c.plus(join_key_complexity(left_key, true)); + } + c = c + .plus(Complexity::Linear(LinearComplexity::AllColumnsFromRight)) + .plus(Complexity::Linear(LinearComplexity::AllColumnsFromRight)); + for (_, right_key) in node.on() { + c = c.plus(join_key_complexity(right_key, false)); + } + // Optional join filter evaluated on matched pairs as batches stream in. + if let Some(filter) = node.filter() { + c = c.plus(join_filter_complexity(filter)); + } + return c; + } + + // Aggregation: hash group-by keys + accumulate aggregate inputs + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/aggregates/mod.rs + if let Some(agg) = node.downcast_ref::() { + // Base: read all input columns for accumulation + let mut c = Complexity::Linear(LinearComplexity::AllColumns); + // Additional: evaluate and hash group-by key expressions + for (expr, _) in agg.group_expr().expr() { + c = c.plus(hashed_or_sorted_key_complexity(expr)); + } + // Per-aggregate filter expressions (e.g. COUNT(*) FILTER (WHERE ...)) + for filter in agg.filter_expr().iter().flatten() { + c = c.plus(expression_complexity(filter)); + } + return c; + } + + // Window functions: buffer partitions, compute aggregates over windows + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/windows/window_agg_exec.rs + if let Some(node) = node.downcast_ref::() { + // Read all input data + evaluate/hash partition key expressions + let mut c = Complexity::Linear(LinearComplexity::AllColumns); + for expr in node.partition_keys() { + c = c.plus(hashed_or_sorted_key_complexity(&expr)); + } + return c; + } + + if let Some(node) = node.downcast_ref::() { + let mut c = Complexity::Linear(LinearComplexity::AllColumns); + for expr in node.partition_keys() { + c = c.plus(hashed_or_sorted_key_complexity(&expr)); + } + return c; + } + + // SortPreservingMergeExec: merges pre-sorted streams with comparisons + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs + // K-way merge: O(N log K) comparisons on sort key expressions + if let Some(node) = node.downcast_ref::() { + // need to copy all rows... + let mut n = Complexity::Linear(LinearComplexity::AllColumns); + // and compare the sort keys on all of them; a plain column key still costs its bytes. + for expr in node.expr() { + n = n.plus(hashed_or_sorted_key_complexity(&expr.expr)) + } + return n; + } + + // FilterExec: evaluates predicate expression per row + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/filter.rs + // Cost depends on predicate complexity - LIKE/Regex operations are expensive + if let Some(node) = node.downcast_ref::() { + // It needs to perform a copy operation just to the output rows... + let n = Complexity::Linear(LinearComplexity::AllOutputColumns); + // ...and predicate evaluation on all input rows. + return n.plus(expression_complexity(node.predicate())); + } + + // ProjectionExec: cost depends on whether it's simple columns or expressions + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/projection.rs + if let Some(node) = node.downcast_ref::() { + let mut n: Option = None; + for expr in node.expr() { + n = if let Some(n) = n { + Some(n.plus(expression_complexity(&expr.expr))) + } else { + Some(expression_complexity(&expr.expr)) + }; + } + return n.unwrap_or(Complexity::Constant(1)); + } + + // RepartitionExec with Hash: computes hash per row + take_arrays + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/repartition/mod.rs + if let Some(node) = node.downcast_ref::() { + // It needs to copy all the data for chunking it to the different output partitions... + let mut n = Complexity::Linear(LinearComplexity::AllColumns); + // And it might need to compute a hash per row based on the provided expressions; hashing a + // plain column key still costs its bytes. + match node.partitioning() { + Partitioning::Hash(expressions, _) => { + for expr in expressions { + n = n.plus(hashed_or_sorted_key_complexity(expr)) + } + } + Partitioning::RoundRobinBatch(_) => {} + Partitioning::UnknownPartitioning(_) => {} + }; + return n; + } + + // DataSourceExec: Produces data, so assume that it's an O(N) operation over all the columns. + if node.is::() { + return Complexity::Linear(LinearComplexity::AllOutputColumns); + } + + // Limit: just counts rows and stops early. + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/limit.rs + if node.is::() || node.is::() { + return Complexity::Constant(1); + } + + // CoalescePartitionsExec: receives batches from partitions, just passes through the record + // batches in a zero copy manner. + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/coalesce_partitions.rs + if node.is::() { + return Complexity::Constant(1); + } + + // BroadcastExec: This node does not do any computation, does not even read the data. + if node.is::() { + return Complexity::Constant(1); + } + + // UnionExec: combines multiple input streams, no processing + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/union.rs + if node.is::() || node.is::() { + return Complexity::Constant(1); + } + + // InterleaveExec: round-robin merging of inputs, no processing + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/union.rs + if node.is::() { + return Complexity::Constant(1); + } + + // EmptyExec: produces no data + // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/empty.rs + if node.is::() { + return Complexity::Constant(1); + } + + // For unknown node types, assume we have to do an O(N) operation over all the rows. + Complexity::Linear(LinearComplexity::AllOutputColumns) +} + +struct BytesPerRow { + processed: Option, + cols_read: Vec, +} + +fn expression_complexity(expression: &Arc) -> Complexity { + _expression_complexity(expression) + .processed + .unwrap_or(Complexity::Constant(1)) +} + +/// Computes the complexity of processing a join key expression, including the cost of +/// reading the leaf columns from the appropriate child (left or right). +/// Unlike `expression_complexity`, this accounts for the cost of hashing/comparing +/// simple column references (which have zero evaluation cost but real I/O cost). +fn join_key_complexity(expression: &Arc, from_left: bool) -> Complexity { + let bpr = _expression_complexity(expression); + let mut result: Option = None; + for col_idx in &bpr.cols_read { + let linear = if from_left { + LinearComplexity::ColumnFromLeft(*col_idx) + } else { + LinearComplexity::ColumnFromRight(*col_idx) + }; + result = Some(match result { + Some(r) => r.plus(Complexity::Linear(linear)), + None => Complexity::Linear(linear), + }); + } + result.unwrap_or(Complexity::Constant(1)) +} + +/// Computes the per-row processing cost of a join filter predicate. +/// +/// A `JoinFilter` is evaluated against an intermediate batch whose columns are described by +/// `column_indices`: intermediate column `i` originates from the left or right child at some +/// original index. `expression_complexity` returns a `Complexity` whose `LinearComplexity::Column` +/// terms reference those intermediate indices, so we remap each of them back onto the +/// corresponding child column before the cost can be evaluated against child statistics. +fn join_filter_complexity(filter: &JoinFilter) -> Complexity { + remap_filter_columns( + expression_complexity(filter.expression()), + filter.column_indices(), + ) +} + +/// Rewrites a `Complexity` built from a join filter's intermediate schema so that every +/// `LinearComplexity::Column` term refers to the left/right child column it actually reads. +/// Columns belonging to neither side (the mark-join sentinel) carry no child bytes, so they +/// collapse to a constant. +fn remap_filter_columns(c: Complexity, column_indices: &[ColumnIndex]) -> Complexity { + match c { + Complexity::Constant(v) => Complexity::Constant(v), + Complexity::Linear(LinearComplexity::Column(i)) => match column_indices.get(i) { + Some(ColumnIndex { + index, + side: JoinSide::Left, + }) => Complexity::Linear(LinearComplexity::ColumnFromLeft(*index)), + Some(ColumnIndex { + index, + side: JoinSide::Right, + }) => Complexity::Linear(LinearComplexity::ColumnFromRight(*index)), + _ => Complexity::Constant(1), + }, + // `expression_complexity` only ever emits `Column` linear terms, but keep the rest + // intact so the remapping stays total. + Complexity::Linear(other) => Complexity::Linear(other), + Complexity::Log(n, m) => { + remap_filter_columns(*n, column_indices).log(remap_filter_columns(*m, column_indices)) + } + Complexity::Plus(n, m) => { + remap_filter_columns(*n, column_indices).plus(remap_filter_columns(*m, column_indices)) + } + Complexity::Multiply(n, m) => remap_filter_columns(*n, column_indices) + .multiply(remap_filter_columns(*m, column_indices)), + } +} + +/// Cost of using an expression as a hashing or comparison key. +/// +/// Unlike `expression_complexity` (which only counts the CPU of *evaluating* a `PhysicalExpr`, +/// so a bare column passthrough is free), this charges the bytes of each underlying leaf column. +/// The hashing/comparison itself is performed by the operator — hash-table build, partition +/// hashing, sort comparators — not by any expression in the plan, and its cost scales with the +/// key's byte width. Use it for group-by keys, hash-partition keys and sort keys. +fn hashed_or_sorted_key_complexity(expression: &Arc) -> Complexity { + let bpr = _expression_complexity(expression); + let mut result: Option = None; + for col_idx in &bpr.cols_read { + result = Some(match result { + Some(r) => r.plus(Complexity::Linear(LinearComplexity::Column(*col_idx))), + None => Complexity::Linear(LinearComplexity::Column(*col_idx)), + }); + } + result.unwrap_or(Complexity::Constant(1)) +} + +fn _expression_complexity(expression: &Arc) -> BytesPerRow { + if let Some(col) = expression.downcast_ref::() { + BytesPerRow { + processed: None, + cols_read: vec![col.index()], + } + } else if expression.is::() { + BytesPerRow { + processed: None, + cols_read: vec![], + } + } else { + // Generic handler for all other expressions: CastExpr, TryCastExpr, CaseExpr, + // InListExpr, IsNullExpr, IsNotNullExpr, NotExpr, NegativeExpr, LikeExpr, + // ScalarFunctionExpr, AsyncFuncExpr, etc. + let mut bytes_per_row = BytesPerRow { + processed: None, + cols_read: vec![], + }; + // This operation processes the result of every child once. We model its per-row cost as + // the sum of (1) the processing already incurred inside each child sub-expression and + // (2) one linear pass over each leaf column feeding the child. A leaf column therefore + // contributes once per operation sitting above it, i.e. its bytes are weighted by its + // depth in the expression tree. Carrying (1) is what keeps nested operations + // (e.g. the `+` in `(a + b) * c`) from being silently dropped. + for child in expression.children() { + let c = _expression_complexity(child); + if let Some(child_processed) = c.processed { + bytes_per_row.processed = Some(match bytes_per_row.processed.take() { + Some(processed) => processed.plus(child_processed), + None => child_processed, + }); + } + for col_read in &c.cols_read { + bytes_per_row.processed = Some(match bytes_per_row.processed.take() { + Some(processed) => { + processed.plus(Complexity::Linear(LinearComplexity::Column(*col_read))) + } + None => Complexity::Linear(LinearComplexity::Column(*col_read)), + }); + } + bytes_per_row.cols_read.extend(&c.cols_read); + } + bytes_per_row + } +} + +#[cfg(test)] +mod tests { + use crate::assert_snapshot; + use crate::distributed_planner::statistics::compute_per_node::calculate_compute_complexity; + use crate::test_utils::plans::TestPlanBuilder; + use datafusion::common::tree_node::{Transformed, TreeNode}; + use datafusion::physical_plan::{ExecutionPlan, displayable}; + use std::cell::RefCell; + use std::sync::Arc; + /* schema for the "weather" table + + MinTemp [type=DOUBLE] [repetitiontype=OPTIONAL] + MaxTemp [type=DOUBLE] [repetitiontype=OPTIONAL] + Rainfall [type=DOUBLE] [repetitiontype=OPTIONAL] + Evaporation [type=DOUBLE] [repetitiontype=OPTIONAL] + Sunshine [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + WindGustDir [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + WindGustSpeed [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + WindDir9am [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + WindDir3pm [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + WindSpeed9am [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + WindSpeed3pm [type=INT64] [convertedtype=INT_64] [repetitiontype=OPTIONAL] + Humidity9am [type=INT64] [convertedtype=INT_64] [repetitiontype=OPTIONAL] + Humidity3pm [type=INT64] [convertedtype=INT_64] [repetitiontype=OPTIONAL] + Pressure9am [type=DOUBLE] [repetitiontype=OPTIONAL] + Pressure3pm [type=DOUBLE] [repetitiontype=OPTIONAL] + Cloud9am [type=INT64] [convertedtype=INT_64] [repetitiontype=OPTIONAL] + Cloud3pm [type=INT64] [convertedtype=INT_64] [repetitiontype=OPTIONAL] + Temp9am [type=DOUBLE] [repetitiontype=OPTIONAL] + Temp3pm [type=DOUBLE] [repetitiontype=OPTIONAL] + RainToday [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + RISK_MM [type=DOUBLE] [repetitiontype=OPTIONAL] + RainTomorrow [type=BYTE_ARRAY] [convertedtype=UTF8] [repetitiontype=OPTIONAL] + */ + + // DataSourceExec: produces data, modeled as O(N) over all output columns. + #[tokio::test] + async fn data_source_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT "MinTemp" FROM weather"#) + .await; + assert_snapshot!(plan_costs(plan), @"O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet"); + } + + // FilterExec: copies the output rows + evaluates the predicate over the input rows. + #[tokio::test] + async fn filter_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT * FROM weather WHERE "MinTemp" > 5"#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O((out_Cols+Col0)) | FilterExec: MinTemp@0 > 5 + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, predicate=MinTemp@0 > 5, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 > 5, required_guarantees=[] + "); + } + + // ProjectionExec: cost is the sum of its expressions; plain column passthroughs are free. + #[tokio::test] + async fn projection_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT "MinTemp" + "MaxTemp" AS s FROM weather"#) + .await; + assert_snapshot!(plan_costs(plan), @"O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp@0 + MaxTemp@1 as s], file_type=parquet"); + } + + // AggregateExec: reads all input columns + hashes the group-by keys. + #[tokio::test] + async fn aggregate_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT "RainToday", COUNT(*) FROM weather GROUP BY "RainToday""#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(2) | ProjectionExec: expr=[RainToday@0 as RainToday, count(Int64(1))@1 as count(*)] + O((Cols+Col0)) | AggregateExec: mode=Single, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet + "); + } + + // SortExec: O(n log n) copy + sort-key evaluation. + #[tokio::test] + async fn sort_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT * FROM weather ORDER BY "WindGustDir""#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O((Cols+Col5)*Log((Cols+Col5))) | SortExec: expr=[WindGustDir@5 ASC NULLS LAST], preserve_partitioning=[false] + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, sort_order_for_reorder=[WindGustDir@5 ASC NULLS LAST] + "); + } + + // SortPreservingMergeExec: appears when several pre-sorted partitions are merged. + #[tokio::test] + async fn sort_preserving_merge_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan(r#"SELECT * FROM weather ORDER BY "WindGustDir""#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O((Cols+Col5)) | SortPreservingMergeExec: [WindGustDir@5 ASC NULLS LAST] + O((Cols+Col5)*Log((Cols+Col5))) | SortExec: expr=[WindGustDir@5 ASC NULLS LAST], preserve_partitioning=[true] + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, sort_order_for_reorder=[WindGustDir@5 ASC NULLS LAST] + "); + } + + // RepartitionExec (Hash): copies all data + hashes the partition keys. + #[tokio::test] + async fn repartition_hash_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan(r#"SELECT "RainToday", COUNT(*) FROM weather GROUP BY "RainToday""#) + .await; + + assert_snapshot!(plan_costs(plan), @r" + O(2) | ProjectionExec: expr=[RainToday@0 as RainToday, count(Int64(1))@1 as count(*)] + O((Cols+Col0)) | AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + O((Cols+Col0)) | RepartitionExec: partitioning=Hash([RainToday@0], 4), input_partitions=3 + O((Cols+Col0)) | AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet + "); + } + + // HashJoinExec: build side (2x read + key hash) + probe side (read + key hash). + #[tokio::test] + async fn hash_join_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a JOIN weather b ON a."RainToday" = b."RainToday" + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(((2*left_Cols)+left_Col1+right_Cols+right_Col1)) | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(RainToday@1, RainToday@1)], projection=[MinTemp@0, MaxTemp@2] + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday], file_type=parquet + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet, predicate=DynamicFilter [ empty ] + "); + } + + // HashJoinExec with a residual filter: the equi-predicate becomes the hash join key while the + // inequality (`a.MinTemp > b.MaxTemp`) becomes a JoinFilter over an intermediate schema, so + // the cost must include the left/right columns the filter reads, not just the join keys. + #[tokio::test] + async fn hash_join_with_filter_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a + JOIN weather b + ON a."RainToday" = b."RainToday" + AND a."MinTemp" > b."MaxTemp" + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(((2*left_Cols)+left_Col1+right_Cols+right_Col1+(left_Col0+right_Col0))) | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(RainToday@1, RainToday@1)], filter=MinTemp@0 > MaxTemp@1, projection=[MinTemp@0, MaxTemp@2] + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday], file_type=parquet + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet, predicate=DynamicFilter [ empty ] + "); + } + + // CrossJoinExec: O(n*m) Cartesian product over all columns of both sides. + #[tokio::test] + async fn cross_join_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT a."MinTemp", b."MaxTemp" FROM weather a CROSS JOIN weather b"#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O((left_Cols*right_Cols)) | CrossJoinExec + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp], file_type=parquet + "); + } + + // NestedLoopJoinExec: produced when a join has no equi-key, only an inequality filter. + #[tokio::test] + async fn nested_loop_join_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a JOIN weather b ON a."MinTemp" > b."MaxTemp" + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(((left_Cols*right_Cols)+(left_Col0+right_Col0))) | NestedLoopJoinExec: join_type=Inner, filter=MinTemp@0 > MaxTemp@1 + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp], file_type=parquet + "); + } + + // SortMergeJoinExec: produced when hash joins are disabled; streams both sorted inputs. + // Requires target_partitions > 1 + repartition_joins + !prefer_hash_join (see physical_planner). + #[tokio::test] + async fn sort_merge_join_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .information_schema(true) + .prefer_hash_joins(false) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a JOIN weather b ON a."RainToday" = b."RainToday" + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(2) | ProjectionExec: expr=[MinTemp@0 as MinTemp, MaxTemp@2 as MaxTemp] + O((left_Col1+right_Col1)) | SortMergeJoinExec: join_type=Inner, on=[(RainToday@1, RainToday@1)] + O((Cols+Col1)*Log((Cols+Col1))) | SortExec: expr=[RainToday@1 ASC], preserve_partitioning=[true] + O((Cols+Col1)) | RepartitionExec: partitioning=Hash([RainToday@1], 4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday], file_type=parquet + O((Cols+Col1)*Log((Cols+Col1))) | SortExec: expr=[RainToday@1 ASC], preserve_partitioning=[true] + O((Cols+Col1)) | RepartitionExec: partitioning=Hash([RainToday@1], 4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet + "); + } + + // BoundedWindowAggExec: window function with an ORDER BY frame (RANK). + #[tokio::test] + async fn bounded_window_agg_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT RANK() OVER (PARTITION BY "RainToday" ORDER BY "MaxTemp") FROM weather + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r#" + O(1) | ProjectionExec: expr=[rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] + O((Cols+Col1)) | BoundedWindowAggExec: wdw=[rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + O((Cols+Col1+Col0)*Log((Cols+Col1+Col0))) | SortExec: expr=[RainToday@1 ASC NULLS LAST, MaxTemp@0 ASC NULLS LAST], preserve_partitioning=[false] + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000002.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000000.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet, sort_order_for_reorder=[RainToday@1 ASC NULLS LAST, MaxTemp@0 ASC NULLS LAST] + "#); + } + + // WindowAggExec: window aggregate without an ORDER BY (unbounded frame). + #[tokio::test] + async fn window_agg_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#"SELECT SUM("Rainfall") OVER (PARTITION BY "WindGustDir") FROM weather"#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r#" + O(1) | ProjectionExec: expr=[sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] + O((Cols+Col1)) | WindowAggExec: wdw=[sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + O((Cols+Col1)*Log((Cols+Col1))) | SortExec: expr=[WindGustDir@1 ASC NULLS LAST], preserve_partitioning=[false] + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[Rainfall, WindGustDir], file_type=parquet, sort_order_for_reorder=[WindGustDir@1 ASC NULLS LAST] + "#); + } + + // UnionExec: combines input streams with no per-row processing. + #[tokio::test] + async fn union_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT "MinTemp" AS t FROM weather + UNION ALL + SELECT "MaxTemp" AS t FROM weather + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(1) | UnionExec + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp@0 as t], file_type=parquet + O(out_Cols) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp@1 as t], file_type=parquet + "); + } + + // AggregateExec with no GROUP BY + CoalescePartitionsExec: the filter prevents the planner from + // answering COUNT(*) straight from parquet metadata, so a real partial aggregate runs per + // partition and is merged through a CoalescePartitionsExec before the single final aggregate. + #[tokio::test] + async fn aggregate_no_group_by_and_coalesce_partitions() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan(r#"SELECT COUNT(*) FROM weather WHERE "MinTemp" > 5"#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(1) | ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + O(Cols) | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + O(1) | CoalescePartitionsExec + O(Cols) | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + O((out_Cols+Col0)) | FilterExec: MinTemp@0 > 5, projection=[] + O(Cols) | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet, predicate=MinTemp@0 > 5, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 > 5, required_guarantees=[] + "); + } + + // GlobalLimitExec: an OFFSET can't be pushed down as a per-partition fetch, so a GlobalLimitExec + // is materialized. (LocalLimitExec shares this exact cost branch but the DF53 planner prefers to + // carry `fetch` on CoalescePartitionsExec rather than emit a separate LocalLimitExec node.) + #[tokio::test] + async fn global_limit_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan(r#"SELECT * FROM weather WHERE "MinTemp" > 5 LIMIT 10 OFFSET 5"#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(1) | GlobalLimitExec: skip=5, fetch=10 + O(1) | CoalescePartitionsExec: fetch=15 + O((out_Cols+Col0)) | FilterExec: MinTemp@0 > 5, fetch=15 + O(Cols) | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, predicate=MinTemp@0 > 5, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 > 5, required_guarantees=[] + "); + } + + // EmptyExec: an always-false predicate collapses to an empty relation that produces no data. + #[tokio::test] + async fn empty_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT "MinTemp" FROM weather WHERE 1 = 0"#) + .await; + assert_snapshot!(plan_costs(plan), @"O(1) | EmptyExec"); + } + + // RoundRobin RepartitionExec: has no hash keys, so it takes the bare all-columns copy path. + #[tokio::test] + async fn round_robin_repartition_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan(r#"SELECT * FROM weather WHERE "MinTemp" > 5"#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O((out_Cols+Col0)) | FilterExec: MinTemp@0 > 5 + O(Cols) | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, predicate=MinTemp@0 > 5, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 > 5, required_guarantees=[] + "); + } + + // HashJoinExec in Partitioned mode: a distinct planner path from CollectLeft. The cost formula + // is the same (build 2x read + key hash, probe read + key hash), now over hash-repartitioned + // inputs rather than a collected left side. + #[tokio::test] + async fn partitioned_hash_join_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .information_schema(true) + // Zero the single-partition thresholds so the planner uses Partitioned mode (hash + // repartition on both sides) instead of collecting the left side. + .hash_join_single_partition_threshold(0) + .hash_join_single_partition_threshold_rows(0) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a JOIN weather b ON a."RainToday" = b."RainToday" + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(((2*left_Cols)+left_Col1+right_Cols+right_Col1)) | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(RainToday@1, RainToday@1)], projection=[MinTemp@0, MaxTemp@2] + O((Cols+Col1)) | RepartitionExec: partitioning=Hash([RainToday@1], 4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday], file_type=parquet + O((Cols+Col1)) | RepartitionExec: partitioning=Hash([RainToday@1], 4), input_partitions=3 + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet, predicate=DynamicFilter [ empty ] + "); + } + + // InterleaveExec: unioning two identically hash-partitioned aggregates lets the planner + // interleave the partitions instead of concatenating streams. + #[tokio::test] + async fn interleave_exec() { + let plan = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan( + r#" + SELECT "RainToday" AS k, COUNT(*) AS c FROM weather GROUP BY "RainToday" + UNION ALL + SELECT "RainTomorrow" AS k, COUNT(*) AS c FROM weather GROUP BY "RainTomorrow" + "#, + ) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(1) | InterleaveExec + O(2) | ProjectionExec: expr=[RainToday@0 as k, count(Int64(1))@1 as c] + O((Cols+Col0)) | AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + O((Cols+Col0)) | RepartitionExec: partitioning=Hash([RainToday@0], 4), input_partitions=3 + O((Cols+Col0)) | AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet + O(2) | ProjectionExec: expr=[RainTomorrow@0 as k, count(Int64(1))@1 as c] + O((Cols+Col0)) | AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[count(Int64(1))] + O((Cols+Col0)) | RepartitionExec: partitioning=Hash([RainTomorrow@0], 4), input_partitions=3 + O((Cols+Col0)) | AggregateExec: mode=Partial, gby=[RainTomorrow@0 as RainTomorrow], aggr=[count(Int64(1))] + O(out_Cols) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[RainTomorrow], file_type=parquet + "); + } + + // Default fallback for unhandled nodes: `SELECT 1` plans a PlaceholderRowExec, which has no + // dedicated branch and therefore takes the catch-all O(N)-over-output-columns estimate. This is + // intentionally conservative; for a 1-row placeholder the output byte size is tiny anyway. + #[tokio::test] + async fn default_fallback_unhandled_node() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT 1"#) + .await; + assert_snapshot!(plan_costs(plan), @r" + O(1) | ProjectionExec: expr=[1 as Int64(1)] + O(out_Cols) | PlaceholderRowExec + "); + } + + // NOTE: BroadcastExec and ChildrenIsolatorUnionExec are inserted only by the distributed + // planner (not by plain DataFusion planning), and SymmetricHashJoinExec requires unbounded + // streaming inputs — none are reachable from these parquet-backed queries. + + fn plan_costs(plan: Arc) -> String { + let mut display = String::new(); + let depth = RefCell::new(0); + plan.transform_down_up( + |plan| { + let indent = " ".repeat(*depth.borrow()); + // `one_line()` renders just this node (with its full config), not its children. + let node = displayable(plan.as_ref()).one_line().to_string(); + display += &format!( + "{indent}O({:?}) | {}\n", + calculate_compute_complexity(&plan), + node.trim_end() + ); + *depth.borrow_mut() += 1; + Ok(Transformed::no(plan)) + }, + |plan| { + *depth.borrow_mut() -= 1; + Ok(Transformed::no(plan)) + }, + ) + .expect("Cannot fail"); + display + } +} diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs new file mode 100644 index 00000000..befb8de7 --- /dev/null +++ b/src/distributed_planner/statistics/cost.rs @@ -0,0 +1,29 @@ +use crate::DistributedConfig; +use crate::distributed_planner::statistics::compute_per_node::calculate_compute_complexity; +use crate::distributed_planner::statistics::plan_statistics::plan_statistics; +use datafusion::common::Result; +use datafusion::physical_plan::{ExecutionPlan, Statistics}; +use std::sync::Arc; + +pub(crate) fn calculate_cost( + plan: &Arc, + cfg: &DistributedConfig, +) -> Result { + f(plan, cfg).map(|(cost, _stats)| cost) +} + +fn f(plan: &Arc, d_cfg: &DistributedConfig) -> Result<(usize, Arc)> { + let children = plan.children(); + let mut child_stats = Vec::with_capacity(children.len()); + let mut acc_cost = 0; + for child in children { + let (cost, child_stat) = f(child, d_cfg)?; + acc_cost += cost; + child_stats.push(child_stat); + } + + let stats = plan_statistics(plan, &child_stats, d_cfg)?; + let complexity = calculate_compute_complexity(plan); + acc_cost += complexity.cost(&stats, &child_stats).unwrap_or(0); + Ok((acc_cost, stats)) +} diff --git a/src/distributed_planner/statistics/default_bytes_for_datatype.rs b/src/distributed_planner/statistics/default_bytes_for_datatype.rs new file mode 100644 index 00000000..994fd6ea --- /dev/null +++ b/src/distributed_planner/statistics/default_bytes_for_datatype.rs @@ -0,0 +1,137 @@ +use datafusion::arrow::datatypes::{DataType, IntervalUnit}; + +/// Default data size estimate for variable-width columns when no statistics are available. +/// +/// Reference: Trino's PlanNodeStatsEstimate.java:40 +/// https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L40 +const DEFAULT_DATA_SIZE_PER_COLUMN: usize = 50; + +/// This function returns the amount of bytes each row is estimated to occupy. +/// +/// The estimation follows Trino's approach for calculating output size per row: +/// - For fixed-width (primitive) types: uses the type's fixed byte width +/// - For variable-width types: uses a default estimate plus offset overhead +/// - Accounts for validity bitmap overhead (1 bit per value, rounded to 1 byte per row) +/// +/// DataFusion has `Statistics::calculate_total_byte_size()` which uses `DataType::primitive_width()`, +/// but it returns `Precision::Absent` (unknown) when encountering any non-primitive type: +/// https://github.com/apache/datafusion/blob/branch-52/datafusion/common/src/stats.rs#L326-L347 +/// +/// For distributed query planning, we need estimates even for variable-width types to make +/// cost-based decisions about data shuffling and task count assignation. This implementation +/// provides estimates for all types following Trino's cost model. +/// +/// Reference: Trino's PlanNodeStatsEstimate.getOutputSizeForSymbol() +/// https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L89-L114 +pub(super) fn default_bytes_for_datatype(data_type: &DataType) -> usize { + // 1 byte for validity bitmap per row (Arrow uses 1 bit, but we round up for estimation). + // Trino calls this the "is null" boolean array. + // Reference: PlanNodeStatsEstimate.java:98-99 + // https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L98-L99 + const VALIDITY_OVERHEAD: usize = 1; + + // Handle non-primitive types. + // NOTE: The cases below are Arrow-specific adaptations. Trino only distinguishes between + // FixedWidthType and variable-width types, using Integer.BYTES (4) for offsets. + // Reference: PlanNodeStatsEstimate.java:108-109 + // https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L108-L109 + match data_type { + // Primitive types from data_type.primitive_width() + DataType::Int8 => VALIDITY_OVERHEAD + 1, + DataType::Int16 => VALIDITY_OVERHEAD + 2, + DataType::Int32 => VALIDITY_OVERHEAD + 4, + DataType::Int64 => VALIDITY_OVERHEAD + 8, + DataType::UInt8 => VALIDITY_OVERHEAD + 1, + DataType::UInt16 => VALIDITY_OVERHEAD + 2, + DataType::UInt32 => VALIDITY_OVERHEAD + 4, + DataType::UInt64 => VALIDITY_OVERHEAD + 8, + DataType::Float16 => VALIDITY_OVERHEAD + 2, + DataType::Float32 => VALIDITY_OVERHEAD + 4, + DataType::Float64 => VALIDITY_OVERHEAD + 8, + DataType::Timestamp(_, _) => VALIDITY_OVERHEAD + 8, + DataType::Date32 => VALIDITY_OVERHEAD + 4, + DataType::Date64 => VALIDITY_OVERHEAD + 8, + DataType::Time32(_) => VALIDITY_OVERHEAD + 4, + DataType::Time64(_) => VALIDITY_OVERHEAD + 8, + DataType::Duration(_) => VALIDITY_OVERHEAD + 8, + DataType::Interval(IntervalUnit::YearMonth) => VALIDITY_OVERHEAD + 4, + DataType::Interval(IntervalUnit::DayTime) => VALIDITY_OVERHEAD + 8, + DataType::Interval(IntervalUnit::MonthDayNano) => VALIDITY_OVERHEAD + 16, + DataType::Decimal32(_, _) => VALIDITY_OVERHEAD + 4, + DataType::Decimal64(_, _) => VALIDITY_OVERHEAD + 8, + DataType::Decimal128(_, _) => VALIDITY_OVERHEAD + 16, + DataType::Decimal256(_, _) => VALIDITY_OVERHEAD + 32, + // Null type has no data (Arrow-specific) + DataType::Null => 0, + + // Boolean is stored as bits (1/8 byte per value), but we round up (Arrow-specific) + DataType::Boolean => VALIDITY_OVERHEAD + 1, + + // Fixed-size binary: just the fixed size + validity (Arrow-specific) + DataType::FixedSizeBinary(size) => VALIDITY_OVERHEAD + (*size as usize), + + // Fixed-size list: fixed count * element size (Arrow-specific) + DataType::FixedSizeList(field, size) => { + VALIDITY_OVERHEAD + (*size as usize) * default_bytes_for_datatype(field.data_type()) + } + + // Struct: sum of all child field sizes (Arrow-specific) + // Trino would treat ROW types as variable-width + DataType::Struct(fields) => fields + .iter() + .map(|f| default_bytes_for_datatype(f.data_type())) + .sum(), + + // Dictionary-encoded: just the key indices, values are shared across rows (Arrow-specific) + // Trino doesn't have dictionary encoding at the type level + DataType::Dictionary(key_type, _value_type) => default_bytes_for_datatype(key_type), + + // Union: type_id (1 byte) + max child size (Arrow-specific) + DataType::Union(fields, _) => { + let max_child_size = fields + .iter() + .map(|(_, f)| default_bytes_for_datatype(f.data_type())) + .max() + .unwrap_or(0); + 1 + max_child_size + } + + // Run-end encoded: estimate as if it were the value type (Arrow-specific) + // Actual compression depends on data distribution + DataType::RunEndEncoded(_, values) => default_bytes_for_datatype(values.data_type()), + + // Variable-width string/binary types. + // Offset size follows Trino's Integer.BYTES (4 bytes). + // Reference: PlanNodeStatsEstimate.java:109 + DataType::Utf8 | DataType::Binary => { + VALIDITY_OVERHEAD + size_of::() + DEFAULT_DATA_SIZE_PER_COLUMN + } + // Large variants use i64 offsets (Arrow-specific, Trino doesn't have large variants) + DataType::LargeUtf8 | DataType::LargeBinary => { + VALIDITY_OVERHEAD + size_of::() + DEFAULT_DATA_SIZE_PER_COLUMN + } + // View types use 16-byte inline representation (Arrow-specific) + // Reference: https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout + DataType::Utf8View | DataType::BinaryView => VALIDITY_OVERHEAD + 16, + + // List types (Arrow-specific adaptation) + // Spark assumes 1 element average for collections (SPARK-18853). Trino treats them + // as flat variable-width with 50-byte default. We follow Spark's 1-element assumption + // to avoid massive overestimation (e.g. Map was 605 bytes with 10 elements). + DataType::List(field) => { + VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) + } + DataType::LargeList(field) => { + VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) + } + DataType::ListView(field) | DataType::LargeListView(field) => { + VALIDITY_OVERHEAD + 8 + default_bytes_for_datatype(field.data_type()) + } + + // Map type: stored as List> (Arrow-specific) + // Uses same 1-element assumption as List types (following Spark). + DataType::Map(field, _) => { + VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) + } // Fallback for any other types - use Trino's default + } +} diff --git a/src/distributed_planner/statistics/mod.rs b/src/distributed_planner/statistics/mod.rs new file mode 100644 index 00000000..9a23b794 --- /dev/null +++ b/src/distributed_planner/statistics/mod.rs @@ -0,0 +1,7 @@ +mod compute_per_node; +mod cost; +mod default_bytes_for_datatype; +mod plan_statistics; + +#[allow(unused)] // will be used in a follow-up PR. +pub(crate) use cost::calculate_cost; diff --git a/src/distributed_planner/statistics/plan_statistics.rs b/src/distributed_planner/statistics/plan_statistics.rs new file mode 100644 index 00000000..b1888df1 --- /dev/null +++ b/src/distributed_planner/statistics/plan_statistics.rs @@ -0,0 +1,160 @@ +use crate::DistributedConfig; +use crate::distributed_planner::statistics::default_bytes_for_datatype::default_bytes_for_datatype; +use datafusion::common::stats::Precision; +use datafusion::common::{Statistics, not_impl_err, plan_err}; +use datafusion::config::ConfigOptions; +use datafusion::error::Result; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use delegate::delegate; +use itertools::Itertools; +use std::fmt::Formatter; +use std::sync::Arc; + +/// Uses upstream DataFusion stats system with some small overrides. +pub(super) fn plan_statistics( + node: &Arc, + children_stats: &[Arc], + opts: &DistributedConfig, +) -> Result> { + let mut stats = partition_statistics_with_children_override(node, None, children_stats)?; + + // If rows are absent, be conservative and assume that all the rows from all the children + // are going to be returned. + if matches!(stats.num_rows, Precision::Absent) { + let num_rows = children_stats + .iter() + .flat_map(|v| v.num_rows.get_value()) + .sum1::(); + let num_rows = if let Some(num_rows) = num_rows { + num_rows + } else if let Some(default) = opts.default_estimated_row_count { + default + } else { + return plan_err!( + "{} does not provide row stats, and none of its children [{}] provides a row count", + node.name(), + node.children() + .iter() + .map(|v| v.name()) + .collect::>() + .join(", ") + ); + }; + stats.num_rows = Precision::Inexact(num_rows) + } + + let schema = node.schema(); + + for (i, col_stats) in &mut stats.column_statistics.iter_mut().enumerate() { + let rows = stats.num_rows.get_value().unwrap_or(&0); + + // If some of the NDVs are not present in one of the column-level stats, assume the + // worst and use the same as the input number of rows. + if matches!(col_stats.distinct_count, Precision::Absent) { + col_stats.distinct_count = Precision::Inexact(*rows); + } + + // If the per-column byte size stats are not present, estimate the byte size based on the + // data type and the row count. + let Some(dt) = schema.fields.get(i).map(|v| v.data_type()) else { + return plan_err!("Field with index {i} not present in schema: {schema:?}"); + }; + + // If it turns out that we do not have `byte_size` stats, but we do have an estimated number + // of rows, do a best-effort in trying to infer the byte size for each column. + if matches!(col_stats.byte_size, Precision::Absent) { + col_stats.byte_size = Precision::Inexact(default_bytes_for_datatype(dt) * rows) + } + } + + // If bytes are absent, let's just infer them based on the schema and the + // number of rows. + if matches!(stats.total_byte_size, Precision::Absent) { + let mut total_byte_size = 0; + for col_stats in &stats.column_statistics { + total_byte_size += col_stats.byte_size.get_value().unwrap_or(&0); + } + stats.total_byte_size = Precision::Inexact(total_byte_size); + } + + Ok(Arc::new(stats)) +} + +// FIXME: because of limitations the the statistics API on DataFusion, we need to resource to +// this sketchy way of overriding child statistics, as we cannot just provide our own. +// If we don't do this: +// 1. we cannot tell nodes to compute statistics based on the ones we provide. +// 2. we recompute statistics unnecessarily across the plan +// This is tracked by https://github.com/apache/datafusion/issues/20184 upstream, and until +// that one is solved, we need to resource to this wrapper. +fn partition_statistics_with_children_override( + node: &Arc, + partition: Option, + child_stats: &[Arc], +) -> Result { + // DataFusion stats system is not very mature yet. This override layer brings in changes + // that might not have already been released or informed overrides. + let statistics_wrapped_children = child_stats + .iter() + .zip(node.children()) + .map(|(stats, child)| StatisticsWrapper { + inner: Arc::clone(child), + stats: Arc::clone(stats), + }) + .map(|v| Arc::new(v) as _) + .collect(); + + let stats = Arc::clone(node) + .with_new_children(statistics_wrapped_children)? + .partition_statistics(partition)?; + + Ok(stats.as_ref().clone()) +} + +#[derive(Debug)] +struct StatisticsWrapper { + stats: Arc, + inner: Arc, +} + +impl DisplayAs for StatisticsWrapper { + delegate! { + to self.inner { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result; + } + } +} + +impl ExecutionPlan for StatisticsWrapper { + fn partition_statistics(&self, partition: Option) -> Result> { + if partition.is_some() { + return plan_err!("StatisticsWrapper not prepared for partition-specific stats"); + } + Ok(Arc::clone(&self.stats)) + } + + delegate! { + to self.inner { + fn name(&self) -> &str; + fn properties(&self) -> &Arc; + fn maintains_input_order(&self) -> Vec; + fn benefits_from_input_partitioning(&self) -> Vec; + fn children(&self) -> Vec<&Arc>; + fn repartitioned(&self, _target_partitions: usize, _config: &ConfigOptions) -> Result>>; + fn execute(&self, partition: usize, context: Arc) -> Result; + fn supports_limit_pushdown(&self) -> bool; + fn with_fetch(&self, _limit: Option) -> Option>; + fn fetch(&self) -> Option; + fn cardinality_effect(&self) -> CardinalityEffect; + } + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + not_impl_err!("with_new_children not implemented") + } +} diff --git a/src/test_utils/plans.rs b/src/test_utils/plans.rs index c5245697..7d299489 100644 --- a/src/test_utils/plans.rs +++ b/src/test_utils/plans.rs @@ -139,6 +139,9 @@ pub(crate) struct TestPlanBuilder { distributed_partial_reduce: Option, distributed_children_isolator_unions: Option, distributed_max_tasks_per_stage: Option, + prefer_hash_join: Option, + hash_join_single_partition_threshold: Option, + hash_join_single_partition_threshold_rows: Option, } #[cfg(test)] @@ -156,6 +159,9 @@ impl TestPlanBuilder { distributed_partial_reduce: None, distributed_children_isolator_unions: None, distributed_max_tasks_per_stage: None, + prefer_hash_join: None, + hash_join_single_partition_threshold: None, + hash_join_single_partition_threshold_rows: None, } } @@ -192,6 +198,21 @@ impl TestPlanBuilder { self } + pub fn prefer_hash_joins(mut self, prefer_hash_joins: bool) -> Self { + self.prefer_hash_join = Some(prefer_hash_joins); + self + } + + pub fn hash_join_single_partition_threshold(mut self, v: usize) -> Self { + self.hash_join_single_partition_threshold = Some(v); + self + } + + pub fn hash_join_single_partition_threshold_rows(mut self, v: usize) -> Self { + self.hash_join_single_partition_threshold_rows = Some(v); + self + } + pub fn broadcast_joins(mut self, enabled: bool) -> Self { self.broadcast_joins = enabled; self @@ -250,6 +271,21 @@ impl TestPlanBuilder { if let Some(enabled) = self.information_schema { config = config.with_information_schema(enabled); } + if let Some(enabled) = self.prefer_hash_join { + config.options_mut().optimizer.prefer_hash_join = enabled + } + if let Some(value) = self.hash_join_single_partition_threshold { + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = value + } + if let Some(value) = self.hash_join_single_partition_threshold_rows { + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = value + } config } @@ -310,6 +346,9 @@ impl Default for TestPlanBuilder { distributed_partial_reduce: None, distributed_children_isolator_unions: None, distributed_max_tasks_per_stage: None, + prefer_hash_join: None, + hash_join_single_partition_threshold: None, + hash_join_single_partition_threshold_rows: None, } } } From 5f586ce7a7ae7b3920fb238e1707dd3287947c07 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Thu, 25 Jun 2026 10:04:01 +0200 Subject: [PATCH 03/20] Use saturating_add and saturating_mul --- .../statistics/compute_per_node.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/distributed_planner/statistics/compute_per_node.rs b/src/distributed_planner/statistics/compute_per_node.rs index 27b31176..36490aee 100644 --- a/src/distributed_planner/statistics/compute_per_node.rs +++ b/src/distributed_planner/statistics/compute_per_node.rs @@ -125,12 +125,12 @@ impl Complexity { // contribute 0 there, i.e. sorting/merging nothing costs nothing. n * m.checked_ilog2().unwrap_or(0) as usize } - Self::Plus(n, m) => { - n.cost(output_stat, input_stats)? + m.cost(output_stat, input_stats)? - } - Self::Multiply(n, m) => { - n.cost(output_stat, input_stats)? * m.cost(output_stat, input_stats)? - } + Self::Plus(n, m) => n + .cost(output_stat, input_stats)? + .saturating_add(m.cost(output_stat, input_stats)?), + Self::Multiply(n, m) => n + .cost(output_stat, input_stats)? + .saturating_mul(m.cost(output_stat, input_stats)?), }) } } From 1257d58758577b6a0789c9bf8c0bc9cf5c7f1e91 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Thu, 25 Jun 2026 10:04:16 +0200 Subject: [PATCH 04/20] Remove default workarounds --- src/distributed_planner/distributed_config.rs | 4 - src/distributed_planner/statistics/cost.rs | 20 ++--- .../statistics/plan_statistics.rs | 76 +------------------ 3 files changed, 11 insertions(+), 89 deletions(-) diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index 6795fd95..d6faee39 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -66,10 +66,6 @@ extensions_options! { /// budget will still be admitted (otherwise we would livelock), so the actual peak per /// connection is `worker_connection_buffer_budget_bytes + max_message_size`. pub worker_connection_buffer_budget_bytes: usize, default = 64 * 1024 * 1024 - /// Distributed DataFusion relies on row count estimation in order to infer how many workers - /// should be used in serving the query. Some plans might not implement any kind of row count - /// estimation, and this parameter sets the default estimated row count for those plans. - pub default_estimated_row_count: Option, default = Some(0) /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to /// estimate how many tasks should be spawned for the [Stage] containing the leaf node. pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default() diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs index befb8de7..96d995d1 100644 --- a/src/distributed_planner/statistics/cost.rs +++ b/src/distributed_planner/statistics/cost.rs @@ -1,28 +1,28 @@ -use crate::DistributedConfig; use crate::distributed_planner::statistics::compute_per_node::calculate_compute_complexity; -use crate::distributed_planner::statistics::plan_statistics::plan_statistics; +use crate::distributed_planner::statistics::plan_statistics::partition_statistics_with_children_override; use datafusion::common::Result; use datafusion::physical_plan::{ExecutionPlan, Statistics}; use std::sync::Arc; -pub(crate) fn calculate_cost( - plan: &Arc, - cfg: &DistributedConfig, -) -> Result { - f(plan, cfg).map(|(cost, _stats)| cost) +pub(crate) fn calculate_cost(plan: &Arc) -> Result { + f(plan).map(|(cost, _stats)| cost) } -fn f(plan: &Arc, d_cfg: &DistributedConfig) -> Result<(usize, Arc)> { +fn f(plan: &Arc) -> Result<(usize, Arc)> { let children = plan.children(); let mut child_stats = Vec::with_capacity(children.len()); let mut acc_cost = 0; for child in children { - let (cost, child_stat) = f(child, d_cfg)?; + let (cost, child_stat) = f(child)?; acc_cost += cost; child_stats.push(child_stat); } - let stats = plan_statistics(plan, &child_stats, d_cfg)?; + let stats = Arc::new(partition_statistics_with_children_override( + plan, + None, + &child_stats, + )?); let complexity = calculate_compute_complexity(plan); acc_cost += complexity.cost(&stats, &child_stats).unwrap_or(0); Ok((acc_cost, stats)) diff --git a/src/distributed_planner/statistics/plan_statistics.rs b/src/distributed_planner/statistics/plan_statistics.rs index b1888df1..82a8209d 100644 --- a/src/distributed_planner/statistics/plan_statistics.rs +++ b/src/distributed_planner/statistics/plan_statistics.rs @@ -1,6 +1,3 @@ -use crate::DistributedConfig; -use crate::distributed_planner::statistics::default_bytes_for_datatype::default_bytes_for_datatype; -use datafusion::common::stats::Precision; use datafusion::common::{Statistics, not_impl_err, plan_err}; use datafusion::config::ConfigOptions; use datafusion::error::Result; @@ -8,80 +5,9 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use delegate::delegate; -use itertools::Itertools; use std::fmt::Formatter; use std::sync::Arc; -/// Uses upstream DataFusion stats system with some small overrides. -pub(super) fn plan_statistics( - node: &Arc, - children_stats: &[Arc], - opts: &DistributedConfig, -) -> Result> { - let mut stats = partition_statistics_with_children_override(node, None, children_stats)?; - - // If rows are absent, be conservative and assume that all the rows from all the children - // are going to be returned. - if matches!(stats.num_rows, Precision::Absent) { - let num_rows = children_stats - .iter() - .flat_map(|v| v.num_rows.get_value()) - .sum1::(); - let num_rows = if let Some(num_rows) = num_rows { - num_rows - } else if let Some(default) = opts.default_estimated_row_count { - default - } else { - return plan_err!( - "{} does not provide row stats, and none of its children [{}] provides a row count", - node.name(), - node.children() - .iter() - .map(|v| v.name()) - .collect::>() - .join(", ") - ); - }; - stats.num_rows = Precision::Inexact(num_rows) - } - - let schema = node.schema(); - - for (i, col_stats) in &mut stats.column_statistics.iter_mut().enumerate() { - let rows = stats.num_rows.get_value().unwrap_or(&0); - - // If some of the NDVs are not present in one of the column-level stats, assume the - // worst and use the same as the input number of rows. - if matches!(col_stats.distinct_count, Precision::Absent) { - col_stats.distinct_count = Precision::Inexact(*rows); - } - - // If the per-column byte size stats are not present, estimate the byte size based on the - // data type and the row count. - let Some(dt) = schema.fields.get(i).map(|v| v.data_type()) else { - return plan_err!("Field with index {i} not present in schema: {schema:?}"); - }; - - // If it turns out that we do not have `byte_size` stats, but we do have an estimated number - // of rows, do a best-effort in trying to infer the byte size for each column. - if matches!(col_stats.byte_size, Precision::Absent) { - col_stats.byte_size = Precision::Inexact(default_bytes_for_datatype(dt) * rows) - } - } - - // If bytes are absent, let's just infer them based on the schema and the - // number of rows. - if matches!(stats.total_byte_size, Precision::Absent) { - let mut total_byte_size = 0; - for col_stats in &stats.column_statistics { - total_byte_size += col_stats.byte_size.get_value().unwrap_or(&0); - } - stats.total_byte_size = Precision::Inexact(total_byte_size); - } - - Ok(Arc::new(stats)) -} - // FIXME: because of limitations the the statistics API on DataFusion, we need to resource to // this sketchy way of overriding child statistics, as we cannot just provide our own. // If we don't do this: @@ -89,7 +15,7 @@ pub(super) fn plan_statistics( // 2. we recompute statistics unnecessarily across the plan // This is tracked by https://github.com/apache/datafusion/issues/20184 upstream, and until // that one is solved, we need to resource to this wrapper. -fn partition_statistics_with_children_override( +pub(crate) fn partition_statistics_with_children_override( node: &Arc, partition: Option, child_stats: &[Arc], From 8934aa2f864969530afa0ce28df3c0a48ae4f312 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Thu, 25 Jun 2026 10:15:44 +0200 Subject: [PATCH 05/20] Factor out complexity --- .../statistics/complexity.rs | 155 ++++++++++++++++ ...{compute_per_node.rs => complexity_cpu.rs} | 167 +----------------- src/distributed_planner/statistics/cost.rs | 4 +- src/distributed_planner/statistics/mod.rs | 3 +- 4 files changed, 165 insertions(+), 164 deletions(-) create mode 100644 src/distributed_planner/statistics/complexity.rs rename src/distributed_planner/statistics/{compute_per_node.rs => complexity_cpu.rs} (87%) diff --git a/src/distributed_planner/statistics/complexity.rs b/src/distributed_planner/statistics/complexity.rs new file mode 100644 index 00000000..46659e03 --- /dev/null +++ b/src/distributed_planner/statistics/complexity.rs @@ -0,0 +1,155 @@ +use datafusion::common::Statistics; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +#[derive(Clone, PartialEq, Eq)] +pub(super) enum Complexity { + /// Constant complexity + Constant(usize), + /// Linear with a specific column from a specific child. + Linear(LinearComplexity), + /// NLogM + Log(Box, Box), + /// N+M + Plus(Box, Box), + /// N*M + Multiply(Box, Box), +} + +#[derive(Clone, PartialEq, Eq)] +pub(super) enum LinearComplexity { + /// Depends on linearly with the input column with the provided index + Column(usize), + /// Depends on linearly with the all the input columns + AllColumns, + /// Depends on linearly with the input column with the provided index from the left child + ColumnFromLeft(usize), + /// Depends on linearly with the all the input columns from the left child + AllColumnsFromLeft, + /// Depends on linearly with the input column with the provided index from the right child + ColumnFromRight(usize), + /// Depends on linearly with the all the input columns from the right child + AllColumnsFromRight, + /// Depends on linearly with the all the output columns + AllOutputColumns, +} + +impl Complexity { + pub(super) fn log(self, other: Self) -> Self { + match (self, other) { + (Self::Constant(n), Self::Constant(m)) => { + Self::Constant(n * (m as f64).log2() as usize) + } + (s, o) => Self::Log(Box::new(s), Box::new(o)), + } + } + + pub(super) fn plus(self, other: Self) -> Self { + match (self, other) { + (Self::Constant(n), Self::Constant(m)) => Self::Constant(n + m), + // (A + k1) + k2 = A + (k1 + k2): bubble constants rightward so they can fold + (Self::Plus(a, b), Self::Constant(m)) if matches!(*b, Self::Constant(_)) => { + (*a).plus((*b).plus(Self::Constant(m))) + } + (s, o) if s == o => Self::Constant(2).multiply(s), + (s, o) => Self::Plus(Box::new(s), Box::new(o)), + } + } + + pub(super) fn multiply(self, other: Self) -> Self { + match (self, other) { + (Self::Constant(n), Self::Constant(m)) => Self::Constant(n * m), + (s, o) => Self::Multiply(Box::new(s), Box::new(o)), + } + } + + /// Computes the total bytes processed given per-child row counts. + /// Returns None if statistics are unavailable for any required input. + pub(super) fn cost( + &self, + output_stat: &Arc, + input_stats: &[Arc], + ) -> Option { + Some(match self { + Self::Constant(v) => *v, + Self::Linear(linear) => match linear { + LinearComplexity::Column(i) => { + let col_stats = &input_stats.first()?.column_statistics; + *col_stats.get(*i)?.byte_size.get_value()? + } + LinearComplexity::AllColumns => { + *input_stats.first()?.total_byte_size.get_value()? + } + LinearComplexity::ColumnFromLeft(i) => { + let col_stats = &input_stats.first()?.column_statistics; + *col_stats.get(*i)?.byte_size.get_value()? + } + LinearComplexity::AllColumnsFromLeft => { + *input_stats.first()?.total_byte_size.get_value()? + } + LinearComplexity::ColumnFromRight(i) => { + let col_stats = &input_stats.last()?.column_statistics; + *col_stats.get(*i)?.byte_size.get_value()? + } + LinearComplexity::AllColumnsFromRight => { + *input_stats.last()?.total_byte_size.get_value()? + } + LinearComplexity::AllOutputColumns => *output_stat.total_byte_size.get_value()?, + }, + Self::Log(n, m) => { + let n = n.cost(output_stat, input_stats)?; + let m = m.cost(output_stat, input_stats)?; + // `ilog2` panics on 0, which happens whenever the logged input has zero estimated + // bytes/rows (e.g. an empty or fully-pruned relation). Flooring at 1 makes log2 + // contribute 0 there, i.e. sorting/merging nothing costs nothing. + n * m.checked_ilog2().unwrap_or(0) as usize + } + Self::Plus(n, m) => n + .cost(output_stat, input_stats)? + .saturating_add(m.cost(output_stat, input_stats)?), + Self::Multiply(n, m) => n + .cost(output_stat, input_stats)? + .saturating_mul(m.cost(output_stat, input_stats)?), + }) + } +} + +impl Debug for Complexity { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + fn trim_parenthesis(dbg: &Complexity) -> String { + let s = format!("{dbg:?}"); + if s.starts_with('(') && s.ends_with(')') { + s[1..s.len() - 1].to_string() + } else { + s + } + } + match self { + Self::Constant(v) => write!(f, "{v}"), + Self::Linear(linear) => match linear { + LinearComplexity::Column(i) => write!(f, "Col{i}"), + LinearComplexity::AllColumns => write!(f, "Cols"), + LinearComplexity::ColumnFromLeft(i) => write!(f, "left_Col{i}"), + LinearComplexity::AllColumnsFromLeft => write!(f, "left_Cols"), + LinearComplexity::ColumnFromRight(i) => write!(f, "right_Col{i}"), + LinearComplexity::AllColumnsFromRight => write!(f, "right_Cols"), + LinearComplexity::AllOutputColumns => write!(f, "out_Cols"), + }, + Self::Log(n, m) => write!(f, "{n:?}*Log({m:?})"), + Self::Plus(n, m) => { + if matches!(n.as_ref(), &Self::Plus(_, _)) { + write!(f, "({}+{m:?})", trim_parenthesis(n)) + } else { + write!(f, "({n:?}+{m:?})") + } + } + Self::Multiply(n, m) => { + if matches!(n.as_ref(), &Self::Multiply(_, _)) { + write!(f, "({}*{m:?})", trim_parenthesis(n)) + } else { + write!(f, "({n:?}*{m:?})") + } + } + } + } +} diff --git a/src/distributed_planner/statistics/compute_per_node.rs b/src/distributed_planner/statistics/complexity_cpu.rs similarity index 87% rename from src/distributed_planner/statistics/compute_per_node.rs rename to src/distributed_planner/statistics/complexity_cpu.rs index 36490aee..07c96ffe 100644 --- a/src/distributed_planner/statistics/compute_per_node.rs +++ b/src/distributed_planner/statistics/complexity_cpu.rs @@ -1,7 +1,8 @@ use crate::BroadcastExec; +use crate::distributed_planner::statistics::complexity::{Complexity, LinearComplexity}; use crate::execution_plans::ChildrenIsolatorUnionExec; use datafusion::catalog::memory::DataSourceExec; -use datafusion::common::{JoinSide, Statistics}; +use datafusion::common::JoinSide; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::aggregates::AggregateExec; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; @@ -20,166 +21,10 @@ use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMerge use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion::physical_plan::{ExecutionPlan, Partitioning}; -use std::fmt::{Debug, Formatter}; use std::sync::Arc; -#[derive(Clone, PartialEq, Eq)] -pub(super) enum Complexity { - /// Constant complexity - Constant(usize), - /// Linear with a specific column from a specific child. - Linear(LinearComplexity), - /// NLogM - Log(Box, Box), - /// N+M - Plus(Box, Box), - /// N*M - Multiply(Box, Box), -} - -#[derive(Clone, PartialEq, Eq)] -pub(super) enum LinearComplexity { - /// Depends on linearly with the input column with the provided index - Column(usize), - /// Depends on linearly with the all the input columns - AllColumns, - /// Depends on linearly with the input column with the provided index from the left child - ColumnFromLeft(usize), - /// Depends on linearly with the all the input columns from the left child - AllColumnsFromLeft, - /// Depends on linearly with the input column with the provided index from the right child - ColumnFromRight(usize), - /// Depends on linearly with the all the input columns from the right child - AllColumnsFromRight, - /// Depends on linearly with the all the output columns - AllOutputColumns, -} - -impl Complexity { - fn log(self, other: Self) -> Self { - match (self, other) { - (Self::Constant(n), Self::Constant(m)) => { - Self::Constant(n * (m as f64).log2() as usize) - } - (s, o) => Self::Log(Box::new(s), Box::new(o)), - } - } - - fn plus(self, other: Self) -> Self { - match (self, other) { - (Self::Constant(n), Self::Constant(m)) => Self::Constant(n + m), - // (A + k1) + k2 = A + (k1 + k2): bubble constants rightward so they can fold - (Self::Plus(a, b), Self::Constant(m)) if matches!(*b, Self::Constant(_)) => { - (*a).plus((*b).plus(Self::Constant(m))) - } - (s, o) if s == o => Self::Constant(2).multiply(s), - (s, o) => Self::Plus(Box::new(s), Box::new(o)), - } - } - - fn multiply(self, other: Self) -> Self { - match (self, other) { - (Self::Constant(n), Self::Constant(m)) => Self::Constant(n * m), - (s, o) => Self::Multiply(Box::new(s), Box::new(o)), - } - } - - /// Computes the total bytes processed given per-child row counts. - /// Returns None if statistics are unavailable for any required input. - pub(super) fn cost( - &self, - output_stat: &Arc, - input_stats: &[Arc], - ) -> Option { - Some(match self { - Self::Constant(v) => *v, - Self::Linear(linear) => match linear { - LinearComplexity::Column(i) => { - let col_stats = &input_stats.first()?.column_statistics; - *col_stats.get(*i)?.byte_size.get_value()? - } - LinearComplexity::AllColumns => { - *input_stats.first()?.total_byte_size.get_value()? - } - LinearComplexity::ColumnFromLeft(i) => { - let col_stats = &input_stats.first()?.column_statistics; - *col_stats.get(*i)?.byte_size.get_value()? - } - LinearComplexity::AllColumnsFromLeft => { - *input_stats.first()?.total_byte_size.get_value()? - } - LinearComplexity::ColumnFromRight(i) => { - let col_stats = &input_stats.last()?.column_statistics; - *col_stats.get(*i)?.byte_size.get_value()? - } - LinearComplexity::AllColumnsFromRight => { - *input_stats.last()?.total_byte_size.get_value()? - } - LinearComplexity::AllOutputColumns => *output_stat.total_byte_size.get_value()?, - }, - Self::Log(n, m) => { - let n = n.cost(output_stat, input_stats)?; - let m = m.cost(output_stat, input_stats)?; - // `ilog2` panics on 0, which happens whenever the logged input has zero estimated - // bytes/rows (e.g. an empty or fully-pruned relation). Flooring at 1 makes log2 - // contribute 0 there, i.e. sorting/merging nothing costs nothing. - n * m.checked_ilog2().unwrap_or(0) as usize - } - Self::Plus(n, m) => n - .cost(output_stat, input_stats)? - .saturating_add(m.cost(output_stat, input_stats)?), - Self::Multiply(n, m) => n - .cost(output_stat, input_stats)? - .saturating_mul(m.cost(output_stat, input_stats)?), - }) - } -} - -impl Debug for Complexity { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - fn trim_parenthesis(dbg: &Complexity) -> String { - let s = format!("{dbg:?}"); - if s.starts_with('(') && s.ends_with(')') { - s[1..s.len() - 1].to_string() - } else { - s - } - } - match self { - Self::Constant(v) => write!(f, "{v}"), - Self::Linear(linear) => match linear { - LinearComplexity::Column(i) => write!(f, "Col{i}"), - LinearComplexity::AllColumns => write!(f, "Cols"), - LinearComplexity::ColumnFromLeft(i) => write!(f, "left_Col{i}"), - LinearComplexity::AllColumnsFromLeft => write!(f, "left_Cols"), - LinearComplexity::ColumnFromRight(i) => write!(f, "right_Col{i}"), - LinearComplexity::AllColumnsFromRight => write!(f, "right_Cols"), - LinearComplexity::AllOutputColumns => write!(f, "out_Cols"), - }, - Self::Log(n, m) => write!(f, "{n:?}*Log({m:?})"), - Self::Plus(n, m) => { - if matches!(n.as_ref(), &Self::Plus(_, _)) { - write!(f, "({}+{m:?})", trim_parenthesis(n)) - } else { - write!(f, "({n:?}+{m:?})") - } - } - Self::Multiply(n, m) => { - if matches!(n.as_ref(), &Self::Multiply(_, _)) { - write!(f, "({}*{m:?})", trim_parenthesis(n)) - } else { - write!(f, "({n:?}*{m:?})") - } - } - } - } -} - -/// Calculates what's the cost, expressed as a number, per input row for each input children. -/// -/// The Vec return has equal size to `node.children()`, and determines how many each input needs -/// to be processed -pub(super) fn calculate_compute_complexity(node: &Arc) -> Complexity { +/// Calculates the CPU cost for the provided node, without recursing into children. +pub(super) fn complexity_cpu(node: &Arc) -> Complexity { // NestedLoopJoinExec: O(n*m) - evaluates join condition for each pair of rows // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/joins/nested_loop_join.rs if let Some(node) = node.downcast_ref::() { @@ -582,7 +427,7 @@ fn _expression_complexity(expression: &Arc) -> BytesPerRow { #[cfg(test)] mod tests { use crate::assert_snapshot; - use crate::distributed_planner::statistics::compute_per_node::calculate_compute_complexity; + use crate::distributed_planner::statistics::complexity_cpu::complexity_cpu; use crate::test_utils::plans::TestPlanBuilder; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::physical_plan::{ExecutionPlan, displayable}; @@ -1012,7 +857,7 @@ mod tests { let node = displayable(plan.as_ref()).one_line().to_string(); display += &format!( "{indent}O({:?}) | {}\n", - calculate_compute_complexity(&plan), + complexity_cpu(&plan), node.trim_end() ); *depth.borrow_mut() += 1; diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs index 96d995d1..7ec3daec 100644 --- a/src/distributed_planner/statistics/cost.rs +++ b/src/distributed_planner/statistics/cost.rs @@ -1,4 +1,4 @@ -use crate::distributed_planner::statistics::compute_per_node::calculate_compute_complexity; +use crate::distributed_planner::statistics::complexity_cpu::complexity_cpu; use crate::distributed_planner::statistics::plan_statistics::partition_statistics_with_children_override; use datafusion::common::Result; use datafusion::physical_plan::{ExecutionPlan, Statistics}; @@ -23,7 +23,7 @@ fn f(plan: &Arc) -> Result<(usize, Arc)> { None, &child_stats, )?); - let complexity = calculate_compute_complexity(plan); + let complexity = complexity_cpu(plan); acc_cost += complexity.cost(&stats, &child_stats).unwrap_or(0); Ok((acc_cost, stats)) } diff --git a/src/distributed_planner/statistics/mod.rs b/src/distributed_planner/statistics/mod.rs index 9a23b794..3439d997 100644 --- a/src/distributed_planner/statistics/mod.rs +++ b/src/distributed_planner/statistics/mod.rs @@ -1,4 +1,5 @@ -mod compute_per_node; +mod complexity; +mod complexity_cpu; mod cost; mod default_bytes_for_datatype; mod plan_statistics; From 9517fbd49b668aabad356261198363967d37d992 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Thu, 25 Jun 2026 10:45:32 +0200 Subject: [PATCH 06/20] Introduce memory and network complexity --- .../statistics/complexity_memory.rs | 256 ++++++++++++++++++ .../statistics/complexity_network.rs | 13 + src/distributed_planner/statistics/cost.rs | 32 ++- src/distributed_planner/statistics/mod.rs | 2 + 4 files changed, 298 insertions(+), 5 deletions(-) create mode 100644 src/distributed_planner/statistics/complexity_memory.rs create mode 100644 src/distributed_planner/statistics/complexity_network.rs diff --git a/src/distributed_planner/statistics/complexity_memory.rs b/src/distributed_planner/statistics/complexity_memory.rs new file mode 100644 index 00000000..6458dfcc --- /dev/null +++ b/src/distributed_planner/statistics/complexity_memory.rs @@ -0,0 +1,256 @@ +use crate::BroadcastExec; +use crate::distributed_planner::statistics::complexity::{Complexity, LinearComplexity}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::aggregates::AggregateExec; +use datafusion::physical_plan::joins::{ + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, SymmetricHashJoinExec, +}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; +use std::sync::Arc; + +/// Calculates the memory cost for the provided node, without recursing into children. +pub(super) fn complexity_memory(node: &Arc) -> Complexity { + // NestedLoopJoinExec buffers the left/build side and streams the right/probe side. Its CPU + // cost is O(left * right), but its retained memory is bounded by the build side plus small + // output/bitmap buffers. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/joins/nested_loop_join.rs + if node.is::() { + return Complexity::Linear(LinearComplexity::AllColumnsFromLeft); + } + + // CrossJoinExec also loads the left/build side into memory once and combines it with the + // streamed right side. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/joins/cross_join.rs + if node.is::() { + return Complexity::Linear(LinearComplexity::AllColumnsFromLeft); + } + + // SortExec buffers input batches. For a full sort, DataFusion may need both input and sorted + // output working space before spilling. For TopK, the retained heap is capped by the fetch + // output, so use the output statistics instead of full input size. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/sorts/sort.rs + if let Some(node) = node.downcast_ref::() { + if node.fetch().is_some() { + return Complexity::Linear(LinearComplexity::AllOutputColumns); + } + + let input = Complexity::Linear(LinearComplexity::AllColumns); + return input.clone().plus(input); + } + + // HashJoinExec retains the left/build side in hash-table state. The right/probe side streams. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/joins/hash_join/exec.rs + if node.is::() { + return Complexity::Linear(LinearComplexity::AllColumnsFromLeft); + } + + // SortMergeJoinExec streams sorted inputs and only keeps bounded merge state. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs + if node.is::() { + return Complexity::Constant(0); + } + + // SymmetricHashJoinExec keeps hash tables for both inputs. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/joins/symmetric_hash_join.rs + if node.is::() { + return Complexity::Linear(LinearComplexity::AllColumnsFromLeft) + .plus(Complexity::Linear(LinearComplexity::AllColumnsFromRight)); + } + + // Hash/group aggregation retains one accumulator state per group. Without a GROUP BY the + // retained state is just the fixed set of aggregate accumulators. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/aggregates/mod.rs + if let Some(agg) = node.downcast_ref::() { + if agg.group_expr().is_true_no_grouping() { + return Complexity::Constant(1); + } + + return Complexity::Linear(LinearComplexity::AllOutputColumns); + } + + // WindowAggExec can retain full partitions for unbounded frames. BoundedWindowAggExec is only + // selected when every window expression reports bounded memory, so keep that fixed-size. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/windows/window_agg_exec.rs + if node.is::() { + return Complexity::Linear(LinearComplexity::AllColumns); + } + + if node.is::() { + return Complexity::Constant(1); + } + + // SortPreservingMergeExec performs a streaming K-way merge. + // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs + if node.is::() { + return Complexity::Constant(0); + } + + // BroadcastExec retains batches so several consumers can replay the same input partition. + if node.is::() { + return Complexity::Linear(LinearComplexity::AllOutputColumns); + } + + Complexity::Constant(0) +} + +#[cfg(test)] +mod tests { + use crate::assert_snapshot; + use crate::distributed_planner::statistics::complexity_memory::complexity_memory; + use crate::test_utils::plans::TestPlanBuilder; + use datafusion::common::tree_node::{Transformed, TreeNode}; + use datafusion::physical_plan::{ExecutionPlan, displayable}; + use std::cell::RefCell; + use std::sync::Arc; + + #[tokio::test] + async fn hash_join_buffers_build_side() { + let plan = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a JOIN weather b ON a."RainToday" = b."RainToday" + "#, + ) + .await; + assert_snapshot!(plan_memory(plan), @r" + M(left_Cols) | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(RainToday@1, RainToday@1)], projection=[MinTemp@0, MaxTemp@2] + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday], file_type=parquet + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet, predicate=DynamicFilter [ empty ] + "); + } + + #[tokio::test] + async fn cross_and_nested_loop_join_buffer_build_side() { + let cross = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT a."MinTemp", b."MaxTemp" FROM weather a CROSS JOIN weather b"#) + .await; + assert_snapshot!(plan_memory(cross), @r" + M(left_Cols) | CrossJoinExec + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp], file_type=parquet + "); + + let nested = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a JOIN weather b ON a."MinTemp" > b."MaxTemp" + "#, + ) + .await; + assert_snapshot!(plan_memory(nested), @r" + M(left_Cols) | NestedLoopJoinExec: join_type=Inner, filter=MinTemp@0 > MaxTemp@1 + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MaxTemp], file_type=parquet + "); + } + + #[tokio::test] + async fn sort_buffers_input_and_topk_buffers_output() { + let full_sort = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT * FROM weather ORDER BY "WindGustDir""#) + .await; + assert_snapshot!(plan_memory(full_sort), @r" + M((2*Cols)) | SortExec: expr=[WindGustDir@5 ASC NULLS LAST], preserve_partitioning=[false] + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, sort_order_for_reorder=[WindGustDir@5 ASC NULLS LAST] + "); + + let topk = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT * FROM weather ORDER BY "WindGustDir" LIMIT 10"#) + .await; + assert_snapshot!(plan_memory(topk), @r" + M(out_Cols) | SortExec: TopK(fetch=10), expr=[WindGustDir@5 ASC NULLS LAST], preserve_partitioning=[false] + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[MinTemp, MaxTemp, Rainfall, Evaporation, Sunshine, WindGustDir, WindGustSpeed, WindDir9am, WindDir3pm, WindSpeed9am, WindSpeed3pm, Humidity9am, Humidity3pm, Pressure9am, Pressure3pm, Cloud9am, Cloud3pm, Temp9am, Temp3pm, RainToday, RISK_MM, RainTomorrow], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[WindGustDir@5 ASC NULLS LAST] + "); + } + + #[tokio::test] + async fn aggregate_memory_tracks_group_state() { + let grouped = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan(r#"SELECT "RainToday", COUNT(*) FROM weather GROUP BY "RainToday""#) + .await; + assert_snapshot!(plan_memory(grouped), @r" + M(0) | ProjectionExec: expr=[RainToday@0 as RainToday, count(Int64(1))@1 as count(*)] + M(out_Cols) | AggregateExec: mode=Single, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))] + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet + "); + + let no_grouping = TestPlanBuilder::new() + .target_partitions(4) + .physical_plan(r#"SELECT COUNT(*) FROM weather WHERE "MinTemp" > 5"#) + .await; + assert_snapshot!(plan_memory(no_grouping), @r" + M(0) | ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + M(1) | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + M(0) | CoalescePartitionsExec + M(1) | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + M(0) | FilterExec: MinTemp@0 > 5, projection=[] + M(0) | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=3 + M(0) | DataSourceExec: file_groups={3 groups: [[/testdata/weather/result-000000.parquet], [/testdata/weather/result-000001.parquet], [/testdata/weather/result-000002.parquet]]}, projection=[MinTemp], file_type=parquet, predicate=MinTemp@0 > 5, pruning_predicate=MinTemp_null_count@1 != row_count@2 AND MinTemp_max@0 > 5, required_guarantees=[] + "); + } + + #[tokio::test] + async fn window_memory_distinguishes_unbounded_and_bounded() { + let unbounded = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#"SELECT SUM("Rainfall") OVER (PARTITION BY "WindGustDir") FROM weather"#, + ) + .await; + assert_snapshot!(plan_memory(unbounded), @r#" + M(0) | ProjectionExec: expr=[sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] + M(Cols) | WindowAggExec: wdw=[sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(weather.Rainfall) PARTITION BY [weather.WindGustDir] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + M((2*Cols)) | SortExec: expr=[WindGustDir@1 ASC NULLS LAST], preserve_partitioning=[false] + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000000.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000002.parquet]]}, projection=[Rainfall, WindGustDir], file_type=parquet, sort_order_for_reorder=[WindGustDir@1 ASC NULLS LAST] + "#); + + let bounded = TestPlanBuilder::new() + .target_partitions(1) + .physical_plan( + r#" + SELECT RANK() OVER (PARTITION BY "RainToday" ORDER BY "MaxTemp") FROM weather + "#, + ) + .await; + assert_snapshot!(plan_memory(bounded), @r#" + M(0) | ProjectionExec: expr=[rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] + M(1) | BoundedWindowAggExec: wdw=[rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [weather.RainToday] ORDER BY [weather.MaxTemp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + M((2*Cols)) | SortExec: expr=[RainToday@1 ASC NULLS LAST, MaxTemp@0 ASC NULLS LAST], preserve_partitioning=[false] + M(0) | DataSourceExec: file_groups={1 group: [[/testdata/weather/result-000002.parquet, /testdata/weather/result-000001.parquet, /testdata/weather/result-000000.parquet]]}, projection=[MaxTemp, RainToday], file_type=parquet, sort_order_for_reorder=[RainToday@1 ASC NULLS LAST, MaxTemp@0 ASC NULLS LAST] + "#); + } + + fn plan_memory(plan: Arc) -> String { + let mut display = String::new(); + let depth = RefCell::new(0); + plan.transform_down_up( + |plan| { + let indent = " ".repeat(*depth.borrow()); + let node = displayable(plan.as_ref()).one_line().to_string(); + display += &format!( + "{indent}M({:?}) | {}\n", + complexity_memory(&plan), + node.trim_end() + ); + *depth.borrow_mut() += 1; + Ok(Transformed::no(plan)) + }, + |plan| { + *depth.borrow_mut() -= 1; + Ok(Transformed::no(plan)) + }, + ) + .expect("Cannot fail"); + display + } +} diff --git a/src/distributed_planner/statistics/complexity_network.rs b/src/distributed_planner/statistics/complexity_network.rs new file mode 100644 index 00000000..9a246633 --- /dev/null +++ b/src/distributed_planner/statistics/complexity_network.rs @@ -0,0 +1,13 @@ +use crate::NetworkBoundaryExt; +use crate::distributed_planner::statistics::complexity::{Complexity, LinearComplexity}; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; + +/// Calculates the memory cost for the provided node, without recursing into children. +pub(super) fn complexity_network(node: &Arc) -> Complexity { + if node.is_network_boundary() { + return Complexity::Linear(LinearComplexity::AllOutputColumns); + } + + Complexity::Constant(0) +} diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs index 7ec3daec..c643e44f 100644 --- a/src/distributed_planner/statistics/cost.rs +++ b/src/distributed_planner/statistics/cost.rs @@ -1,17 +1,35 @@ use crate::distributed_planner::statistics::complexity_cpu::complexity_cpu; +use crate::distributed_planner::statistics::complexity_memory::complexity_memory; +use crate::distributed_planner::statistics::complexity_network::complexity_network; use crate::distributed_planner::statistics::plan_statistics::partition_statistics_with_children_override; use datafusion::common::Result; use datafusion::physical_plan::{ExecutionPlan, Statistics}; +use std::ops::AddAssign; use std::sync::Arc; -pub(crate) fn calculate_cost(plan: &Arc) -> Result { +#[derive(Default)] +pub(crate) struct Cost { + pub(crate) cpu: usize, + pub(crate) memory: usize, + pub(crate) network: usize, +} + +impl AddAssign for Cost { + fn add_assign(&mut self, rhs: Self) { + self.cpu += rhs.cpu; + self.memory += rhs.memory; + self.network += rhs.network; + } +} + +pub(crate) fn calculate_cost(plan: &Arc) -> Result { f(plan).map(|(cost, _stats)| cost) } -fn f(plan: &Arc) -> Result<(usize, Arc)> { +fn f(plan: &Arc) -> Result<(Cost, Arc)> { let children = plan.children(); let mut child_stats = Vec::with_capacity(children.len()); - let mut acc_cost = 0; + let mut acc_cost = Cost::default(); for child in children { let (cost, child_stat) = f(child)?; acc_cost += cost; @@ -23,7 +41,11 @@ fn f(plan: &Arc) -> Result<(usize, Arc)> { None, &child_stats, )?); - let complexity = complexity_cpu(plan); - acc_cost += complexity.cost(&stats, &child_stats).unwrap_or(0); + let cpu = complexity_cpu(plan); + acc_cost.cpu += cpu.cost(&stats, &child_stats).unwrap_or(0); + let memory = complexity_memory(plan); + acc_cost.memory += memory.cost(&stats, &child_stats).unwrap_or(0); + let network = complexity_network(plan); + acc_cost.network += network.cost(&stats, &child_stats).unwrap_or(0); Ok((acc_cost, stats)) } diff --git a/src/distributed_planner/statistics/mod.rs b/src/distributed_planner/statistics/mod.rs index 3439d997..59f1ec04 100644 --- a/src/distributed_planner/statistics/mod.rs +++ b/src/distributed_planner/statistics/mod.rs @@ -1,5 +1,7 @@ mod complexity; mod complexity_cpu; +mod complexity_memory; +mod complexity_network; mod cost; mod default_bytes_for_datatype; mod plan_statistics; From f8fb0f3a71fd9ef40a9ca60dd5d1e9d280b26228 Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Thu, 25 Jun 2026 10:57:24 +0200 Subject: [PATCH 07/20] Remove unused file --- .../statistics/default_bytes_for_datatype.rs | 137 ------------------ src/distributed_planner/statistics/mod.rs | 1 - 2 files changed, 138 deletions(-) delete mode 100644 src/distributed_planner/statistics/default_bytes_for_datatype.rs diff --git a/src/distributed_planner/statistics/default_bytes_for_datatype.rs b/src/distributed_planner/statistics/default_bytes_for_datatype.rs deleted file mode 100644 index 994fd6ea..00000000 --- a/src/distributed_planner/statistics/default_bytes_for_datatype.rs +++ /dev/null @@ -1,137 +0,0 @@ -use datafusion::arrow::datatypes::{DataType, IntervalUnit}; - -/// Default data size estimate for variable-width columns when no statistics are available. -/// -/// Reference: Trino's PlanNodeStatsEstimate.java:40 -/// https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L40 -const DEFAULT_DATA_SIZE_PER_COLUMN: usize = 50; - -/// This function returns the amount of bytes each row is estimated to occupy. -/// -/// The estimation follows Trino's approach for calculating output size per row: -/// - For fixed-width (primitive) types: uses the type's fixed byte width -/// - For variable-width types: uses a default estimate plus offset overhead -/// - Accounts for validity bitmap overhead (1 bit per value, rounded to 1 byte per row) -/// -/// DataFusion has `Statistics::calculate_total_byte_size()` which uses `DataType::primitive_width()`, -/// but it returns `Precision::Absent` (unknown) when encountering any non-primitive type: -/// https://github.com/apache/datafusion/blob/branch-52/datafusion/common/src/stats.rs#L326-L347 -/// -/// For distributed query planning, we need estimates even for variable-width types to make -/// cost-based decisions about data shuffling and task count assignation. This implementation -/// provides estimates for all types following Trino's cost model. -/// -/// Reference: Trino's PlanNodeStatsEstimate.getOutputSizeForSymbol() -/// https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L89-L114 -pub(super) fn default_bytes_for_datatype(data_type: &DataType) -> usize { - // 1 byte for validity bitmap per row (Arrow uses 1 bit, but we round up for estimation). - // Trino calls this the "is null" boolean array. - // Reference: PlanNodeStatsEstimate.java:98-99 - // https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L98-L99 - const VALIDITY_OVERHEAD: usize = 1; - - // Handle non-primitive types. - // NOTE: The cases below are Arrow-specific adaptations. Trino only distinguishes between - // FixedWidthType and variable-width types, using Integer.BYTES (4) for offsets. - // Reference: PlanNodeStatsEstimate.java:108-109 - // https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L108-L109 - match data_type { - // Primitive types from data_type.primitive_width() - DataType::Int8 => VALIDITY_OVERHEAD + 1, - DataType::Int16 => VALIDITY_OVERHEAD + 2, - DataType::Int32 => VALIDITY_OVERHEAD + 4, - DataType::Int64 => VALIDITY_OVERHEAD + 8, - DataType::UInt8 => VALIDITY_OVERHEAD + 1, - DataType::UInt16 => VALIDITY_OVERHEAD + 2, - DataType::UInt32 => VALIDITY_OVERHEAD + 4, - DataType::UInt64 => VALIDITY_OVERHEAD + 8, - DataType::Float16 => VALIDITY_OVERHEAD + 2, - DataType::Float32 => VALIDITY_OVERHEAD + 4, - DataType::Float64 => VALIDITY_OVERHEAD + 8, - DataType::Timestamp(_, _) => VALIDITY_OVERHEAD + 8, - DataType::Date32 => VALIDITY_OVERHEAD + 4, - DataType::Date64 => VALIDITY_OVERHEAD + 8, - DataType::Time32(_) => VALIDITY_OVERHEAD + 4, - DataType::Time64(_) => VALIDITY_OVERHEAD + 8, - DataType::Duration(_) => VALIDITY_OVERHEAD + 8, - DataType::Interval(IntervalUnit::YearMonth) => VALIDITY_OVERHEAD + 4, - DataType::Interval(IntervalUnit::DayTime) => VALIDITY_OVERHEAD + 8, - DataType::Interval(IntervalUnit::MonthDayNano) => VALIDITY_OVERHEAD + 16, - DataType::Decimal32(_, _) => VALIDITY_OVERHEAD + 4, - DataType::Decimal64(_, _) => VALIDITY_OVERHEAD + 8, - DataType::Decimal128(_, _) => VALIDITY_OVERHEAD + 16, - DataType::Decimal256(_, _) => VALIDITY_OVERHEAD + 32, - // Null type has no data (Arrow-specific) - DataType::Null => 0, - - // Boolean is stored as bits (1/8 byte per value), but we round up (Arrow-specific) - DataType::Boolean => VALIDITY_OVERHEAD + 1, - - // Fixed-size binary: just the fixed size + validity (Arrow-specific) - DataType::FixedSizeBinary(size) => VALIDITY_OVERHEAD + (*size as usize), - - // Fixed-size list: fixed count * element size (Arrow-specific) - DataType::FixedSizeList(field, size) => { - VALIDITY_OVERHEAD + (*size as usize) * default_bytes_for_datatype(field.data_type()) - } - - // Struct: sum of all child field sizes (Arrow-specific) - // Trino would treat ROW types as variable-width - DataType::Struct(fields) => fields - .iter() - .map(|f| default_bytes_for_datatype(f.data_type())) - .sum(), - - // Dictionary-encoded: just the key indices, values are shared across rows (Arrow-specific) - // Trino doesn't have dictionary encoding at the type level - DataType::Dictionary(key_type, _value_type) => default_bytes_for_datatype(key_type), - - // Union: type_id (1 byte) + max child size (Arrow-specific) - DataType::Union(fields, _) => { - let max_child_size = fields - .iter() - .map(|(_, f)| default_bytes_for_datatype(f.data_type())) - .max() - .unwrap_or(0); - 1 + max_child_size - } - - // Run-end encoded: estimate as if it were the value type (Arrow-specific) - // Actual compression depends on data distribution - DataType::RunEndEncoded(_, values) => default_bytes_for_datatype(values.data_type()), - - // Variable-width string/binary types. - // Offset size follows Trino's Integer.BYTES (4 bytes). - // Reference: PlanNodeStatsEstimate.java:109 - DataType::Utf8 | DataType::Binary => { - VALIDITY_OVERHEAD + size_of::() + DEFAULT_DATA_SIZE_PER_COLUMN - } - // Large variants use i64 offsets (Arrow-specific, Trino doesn't have large variants) - DataType::LargeUtf8 | DataType::LargeBinary => { - VALIDITY_OVERHEAD + size_of::() + DEFAULT_DATA_SIZE_PER_COLUMN - } - // View types use 16-byte inline representation (Arrow-specific) - // Reference: https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout - DataType::Utf8View | DataType::BinaryView => VALIDITY_OVERHEAD + 16, - - // List types (Arrow-specific adaptation) - // Spark assumes 1 element average for collections (SPARK-18853). Trino treats them - // as flat variable-width with 50-byte default. We follow Spark's 1-element assumption - // to avoid massive overestimation (e.g. Map was 605 bytes with 10 elements). - DataType::List(field) => { - VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) - } - DataType::LargeList(field) => { - VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) - } - DataType::ListView(field) | DataType::LargeListView(field) => { - VALIDITY_OVERHEAD + 8 + default_bytes_for_datatype(field.data_type()) - } - - // Map type: stored as List> (Arrow-specific) - // Uses same 1-element assumption as List types (following Spark). - DataType::Map(field, _) => { - VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) - } // Fallback for any other types - use Trino's default - } -} diff --git a/src/distributed_planner/statistics/mod.rs b/src/distributed_planner/statistics/mod.rs index 59f1ec04..77044eb5 100644 --- a/src/distributed_planner/statistics/mod.rs +++ b/src/distributed_planner/statistics/mod.rs @@ -3,7 +3,6 @@ mod complexity_cpu; mod complexity_memory; mod complexity_network; mod cost; -mod default_bytes_for_datatype; mod plan_statistics; #[allow(unused)] // will be used in a follow-up PR. From 74c74f2d627ffdd819101ce901bd714fde14c26b Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Thu, 25 Jun 2026 13:09:02 +0200 Subject: [PATCH 08/20] Bring back defensive defaults --- src/distributed_planner/statistics/cost.rs | 8 +- .../statistics/default_bytes_for_datatype.rs | 137 ++++++++++++++++++ src/distributed_planner/statistics/mod.rs | 1 + .../statistics/plan_statistics.rs | 69 ++++++++- 4 files changed, 208 insertions(+), 7 deletions(-) create mode 100644 src/distributed_planner/statistics/default_bytes_for_datatype.rs diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs index c643e44f..fb926d62 100644 --- a/src/distributed_planner/statistics/cost.rs +++ b/src/distributed_planner/statistics/cost.rs @@ -1,7 +1,7 @@ use crate::distributed_planner::statistics::complexity_cpu::complexity_cpu; use crate::distributed_planner::statistics::complexity_memory::complexity_memory; use crate::distributed_planner::statistics::complexity_network::complexity_network; -use crate::distributed_planner::statistics::plan_statistics::partition_statistics_with_children_override; +use crate::distributed_planner::statistics::plan_statistics::plan_statistics; use datafusion::common::Result; use datafusion::physical_plan::{ExecutionPlan, Statistics}; use std::ops::AddAssign; @@ -36,11 +36,7 @@ fn f(plan: &Arc) -> Result<(Cost, Arc)> { child_stats.push(child_stat); } - let stats = Arc::new(partition_statistics_with_children_override( - plan, - None, - &child_stats, - )?); + let stats = plan_statistics(plan, &child_stats)?; let cpu = complexity_cpu(plan); acc_cost.cpu += cpu.cost(&stats, &child_stats).unwrap_or(0); let memory = complexity_memory(plan); diff --git a/src/distributed_planner/statistics/default_bytes_for_datatype.rs b/src/distributed_planner/statistics/default_bytes_for_datatype.rs new file mode 100644 index 00000000..994fd6ea --- /dev/null +++ b/src/distributed_planner/statistics/default_bytes_for_datatype.rs @@ -0,0 +1,137 @@ +use datafusion::arrow::datatypes::{DataType, IntervalUnit}; + +/// Default data size estimate for variable-width columns when no statistics are available. +/// +/// Reference: Trino's PlanNodeStatsEstimate.java:40 +/// https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L40 +const DEFAULT_DATA_SIZE_PER_COLUMN: usize = 50; + +/// This function returns the amount of bytes each row is estimated to occupy. +/// +/// The estimation follows Trino's approach for calculating output size per row: +/// - For fixed-width (primitive) types: uses the type's fixed byte width +/// - For variable-width types: uses a default estimate plus offset overhead +/// - Accounts for validity bitmap overhead (1 bit per value, rounded to 1 byte per row) +/// +/// DataFusion has `Statistics::calculate_total_byte_size()` which uses `DataType::primitive_width()`, +/// but it returns `Precision::Absent` (unknown) when encountering any non-primitive type: +/// https://github.com/apache/datafusion/blob/branch-52/datafusion/common/src/stats.rs#L326-L347 +/// +/// For distributed query planning, we need estimates even for variable-width types to make +/// cost-based decisions about data shuffling and task count assignation. This implementation +/// provides estimates for all types following Trino's cost model. +/// +/// Reference: Trino's PlanNodeStatsEstimate.getOutputSizeForSymbol() +/// https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L89-L114 +pub(super) fn default_bytes_for_datatype(data_type: &DataType) -> usize { + // 1 byte for validity bitmap per row (Arrow uses 1 bit, but we round up for estimation). + // Trino calls this the "is null" boolean array. + // Reference: PlanNodeStatsEstimate.java:98-99 + // https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L98-L99 + const VALIDITY_OVERHEAD: usize = 1; + + // Handle non-primitive types. + // NOTE: The cases below are Arrow-specific adaptations. Trino only distinguishes between + // FixedWidthType and variable-width types, using Integer.BYTES (4) for offsets. + // Reference: PlanNodeStatsEstimate.java:108-109 + // https://github.com/trinodb/trino/blob/458/core/trino-main/src/main/java/io/trino/cost/PlanNodeStatsEstimate.java#L108-L109 + match data_type { + // Primitive types from data_type.primitive_width() + DataType::Int8 => VALIDITY_OVERHEAD + 1, + DataType::Int16 => VALIDITY_OVERHEAD + 2, + DataType::Int32 => VALIDITY_OVERHEAD + 4, + DataType::Int64 => VALIDITY_OVERHEAD + 8, + DataType::UInt8 => VALIDITY_OVERHEAD + 1, + DataType::UInt16 => VALIDITY_OVERHEAD + 2, + DataType::UInt32 => VALIDITY_OVERHEAD + 4, + DataType::UInt64 => VALIDITY_OVERHEAD + 8, + DataType::Float16 => VALIDITY_OVERHEAD + 2, + DataType::Float32 => VALIDITY_OVERHEAD + 4, + DataType::Float64 => VALIDITY_OVERHEAD + 8, + DataType::Timestamp(_, _) => VALIDITY_OVERHEAD + 8, + DataType::Date32 => VALIDITY_OVERHEAD + 4, + DataType::Date64 => VALIDITY_OVERHEAD + 8, + DataType::Time32(_) => VALIDITY_OVERHEAD + 4, + DataType::Time64(_) => VALIDITY_OVERHEAD + 8, + DataType::Duration(_) => VALIDITY_OVERHEAD + 8, + DataType::Interval(IntervalUnit::YearMonth) => VALIDITY_OVERHEAD + 4, + DataType::Interval(IntervalUnit::DayTime) => VALIDITY_OVERHEAD + 8, + DataType::Interval(IntervalUnit::MonthDayNano) => VALIDITY_OVERHEAD + 16, + DataType::Decimal32(_, _) => VALIDITY_OVERHEAD + 4, + DataType::Decimal64(_, _) => VALIDITY_OVERHEAD + 8, + DataType::Decimal128(_, _) => VALIDITY_OVERHEAD + 16, + DataType::Decimal256(_, _) => VALIDITY_OVERHEAD + 32, + // Null type has no data (Arrow-specific) + DataType::Null => 0, + + // Boolean is stored as bits (1/8 byte per value), but we round up (Arrow-specific) + DataType::Boolean => VALIDITY_OVERHEAD + 1, + + // Fixed-size binary: just the fixed size + validity (Arrow-specific) + DataType::FixedSizeBinary(size) => VALIDITY_OVERHEAD + (*size as usize), + + // Fixed-size list: fixed count * element size (Arrow-specific) + DataType::FixedSizeList(field, size) => { + VALIDITY_OVERHEAD + (*size as usize) * default_bytes_for_datatype(field.data_type()) + } + + // Struct: sum of all child field sizes (Arrow-specific) + // Trino would treat ROW types as variable-width + DataType::Struct(fields) => fields + .iter() + .map(|f| default_bytes_for_datatype(f.data_type())) + .sum(), + + // Dictionary-encoded: just the key indices, values are shared across rows (Arrow-specific) + // Trino doesn't have dictionary encoding at the type level + DataType::Dictionary(key_type, _value_type) => default_bytes_for_datatype(key_type), + + // Union: type_id (1 byte) + max child size (Arrow-specific) + DataType::Union(fields, _) => { + let max_child_size = fields + .iter() + .map(|(_, f)| default_bytes_for_datatype(f.data_type())) + .max() + .unwrap_or(0); + 1 + max_child_size + } + + // Run-end encoded: estimate as if it were the value type (Arrow-specific) + // Actual compression depends on data distribution + DataType::RunEndEncoded(_, values) => default_bytes_for_datatype(values.data_type()), + + // Variable-width string/binary types. + // Offset size follows Trino's Integer.BYTES (4 bytes). + // Reference: PlanNodeStatsEstimate.java:109 + DataType::Utf8 | DataType::Binary => { + VALIDITY_OVERHEAD + size_of::() + DEFAULT_DATA_SIZE_PER_COLUMN + } + // Large variants use i64 offsets (Arrow-specific, Trino doesn't have large variants) + DataType::LargeUtf8 | DataType::LargeBinary => { + VALIDITY_OVERHEAD + size_of::() + DEFAULT_DATA_SIZE_PER_COLUMN + } + // View types use 16-byte inline representation (Arrow-specific) + // Reference: https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout + DataType::Utf8View | DataType::BinaryView => VALIDITY_OVERHEAD + 16, + + // List types (Arrow-specific adaptation) + // Spark assumes 1 element average for collections (SPARK-18853). Trino treats them + // as flat variable-width with 50-byte default. We follow Spark's 1-element assumption + // to avoid massive overestimation (e.g. Map was 605 bytes with 10 elements). + DataType::List(field) => { + VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) + } + DataType::LargeList(field) => { + VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) + } + DataType::ListView(field) | DataType::LargeListView(field) => { + VALIDITY_OVERHEAD + 8 + default_bytes_for_datatype(field.data_type()) + } + + // Map type: stored as List> (Arrow-specific) + // Uses same 1-element assumption as List types (following Spark). + DataType::Map(field, _) => { + VALIDITY_OVERHEAD + size_of::() + default_bytes_for_datatype(field.data_type()) + } // Fallback for any other types - use Trino's default + } +} diff --git a/src/distributed_planner/statistics/mod.rs b/src/distributed_planner/statistics/mod.rs index 77044eb5..59f1ec04 100644 --- a/src/distributed_planner/statistics/mod.rs +++ b/src/distributed_planner/statistics/mod.rs @@ -3,6 +3,7 @@ mod complexity_cpu; mod complexity_memory; mod complexity_network; mod cost; +mod default_bytes_for_datatype; mod plan_statistics; #[allow(unused)] // will be used in a follow-up PR. diff --git a/src/distributed_planner/statistics/plan_statistics.rs b/src/distributed_planner/statistics/plan_statistics.rs index 82a8209d..2265fb4e 100644 --- a/src/distributed_planner/statistics/plan_statistics.rs +++ b/src/distributed_planner/statistics/plan_statistics.rs @@ -1,3 +1,5 @@ +use crate::distributed_planner::statistics::default_bytes_for_datatype::default_bytes_for_datatype; +use datafusion::common::stats::Precision; use datafusion::common::{Statistics, not_impl_err, plan_err}; use datafusion::config::ConfigOptions; use datafusion::error::Result; @@ -5,9 +7,74 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use delegate::delegate; +use itertools::Itertools; use std::fmt::Formatter; use std::sync::Arc; +/// Ratio applied to the total number of rows for calculating the fallback value for NDVs. +/// +/// If NDVs are absent for a specific column, this ratio kicks in and is applied to the estimated +/// number of rows for calculating the final NDV value. +const FALLBACK_NDV_RATIO: f64 = 0.5; + +/// Uses upstream DataFusion stats system with some small overrides. +pub(super) fn plan_statistics( + node: &Arc, + children_stats: &[Arc], +) -> Result> { + let mut stats = partition_statistics_with_children_override(node, None, children_stats)?; + + // If rows are absent, but the children declares rows, be conservative and assume that the node + // is not going to reduce cardinality and that the row count stays the same. + if matches!(stats.num_rows, Precision::Absent) + && let Some(child_rows) = children_stats + .iter() + .flat_map(|v| v.num_rows.get_value()) + .sum1::() + { + stats.num_rows = Precision::Inexact(child_rows) + } + + let schema = node.schema(); + + for (i, col_stats) in &mut stats.column_statistics.iter_mut().enumerate() { + let Some(rows) = stats.num_rows.get_value() else { + break; + }; + + // If some of the NDVs are not present in one of the column-level stats, assume the + // worst and use the same as the input number of rows. + if matches!(col_stats.distinct_count, Precision::Absent) { + let fallback_ndv = ((*rows as f64) * FALLBACK_NDV_RATIO) as usize; + col_stats.distinct_count = Precision::Inexact(fallback_ndv); + } + + // If the per-column byte size stats are not present, estimate the byte size based on the + // data type and the row count. + let Some(dt) = schema.fields.get(i).map(|v| v.data_type()) else { + return plan_err!("Field with index {i} not present in schema: {schema:?}"); + }; + + // If it turns out that we do not have `byte_size` stats, but we do have an estimated number + // of rows, do a best-effort in trying to infer the byte size for each column. + if matches!(col_stats.byte_size, Precision::Absent) { + col_stats.byte_size = Precision::Inexact(default_bytes_for_datatype(dt) * rows) + } + } + + // If bytes are absent, let's just infer them based on the schema and the + // number of rows. + if matches!(stats.total_byte_size, Precision::Absent) { + let mut total_byte_size = 0; + for col_stats in &stats.column_statistics { + total_byte_size += col_stats.byte_size.get_value().unwrap_or(&0); + } + stats.total_byte_size = Precision::Inexact(total_byte_size); + } + + Ok(Arc::new(stats)) +} + // FIXME: because of limitations the the statistics API on DataFusion, we need to resource to // this sketchy way of overriding child statistics, as we cannot just provide our own. // If we don't do this: @@ -15,7 +82,7 @@ use std::sync::Arc; // 2. we recompute statistics unnecessarily across the plan // This is tracked by https://github.com/apache/datafusion/issues/20184 upstream, and until // that one is solved, we need to resource to this wrapper. -pub(crate) fn partition_statistics_with_children_override( +fn partition_statistics_with_children_override( node: &Arc, partition: Option, child_stats: &[Arc], From a06c9563fceb2d02c99fd6abaa3fe0a80ea1e531 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 29 Jun 2026 12:56:46 +0200 Subject: [PATCH 09/20] Address comments --- .../statistics/complexity.rs | 12 ++++------ .../statistics/complexity_cpu.rs | 24 +++++++++---------- .../statistics/complexity_memory.rs | 10 ++++---- .../statistics/complexity_network.rs | 21 ++++++++++++---- src/distributed_planner/statistics/cost.rs | 18 +++++++++----- .../statistics/plan_statistics.rs | 11 +++++---- 6 files changed, 57 insertions(+), 39 deletions(-) diff --git a/src/distributed_planner/statistics/complexity.rs b/src/distributed_planner/statistics/complexity.rs index 46659e03..b03c6c7e 100644 --- a/src/distributed_planner/statistics/complexity.rs +++ b/src/distributed_planner/statistics/complexity.rs @@ -2,10 +2,10 @@ use datafusion::common::Statistics; use std::fmt::{Debug, Formatter}; use std::sync::Arc; -#[derive(Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq)] pub(super) enum Complexity { /// Constant complexity - Constant(usize), + Constant(f32), /// Linear with a specific column from a specific child. Linear(LinearComplexity), /// NLogM @@ -37,9 +37,7 @@ pub(super) enum LinearComplexity { impl Complexity { pub(super) fn log(self, other: Self) -> Self { match (self, other) { - (Self::Constant(n), Self::Constant(m)) => { - Self::Constant(n * (m as f64).log2() as usize) - } + (Self::Constant(n), Self::Constant(m)) => Self::Constant(n * m.log2()), (s, o) => Self::Log(Box::new(s), Box::new(o)), } } @@ -51,7 +49,7 @@ impl Complexity { (Self::Plus(a, b), Self::Constant(m)) if matches!(*b, Self::Constant(_)) => { (*a).plus((*b).plus(Self::Constant(m))) } - (s, o) if s == o => Self::Constant(2).multiply(s), + (s, o) if s == o => Self::Constant(2.).multiply(s), (s, o) => Self::Plus(Box::new(s), Box::new(o)), } } @@ -71,7 +69,7 @@ impl Complexity { input_stats: &[Arc], ) -> Option { Some(match self { - Self::Constant(v) => *v, + Self::Constant(v) => *v as usize, Self::Linear(linear) => match linear { LinearComplexity::Column(i) => { let col_stats = &input_stats.first()?.column_statistics; diff --git a/src/distributed_planner/statistics/complexity_cpu.rs b/src/distributed_planner/statistics/complexity_cpu.rs index 07c96ffe..49ed6e77 100644 --- a/src/distributed_planner/statistics/complexity_cpu.rs +++ b/src/distributed_planner/statistics/complexity_cpu.rs @@ -116,7 +116,7 @@ pub(super) fn complexity_cpu(node: &Arc) -> Complexity { None => f, }); } - return c.unwrap_or(Complexity::Constant(1)); + return c.unwrap_or(Complexity::Constant(1.)); } // SymmetricHashJoinExec: streaming join with hash tables on both sides @@ -213,7 +213,7 @@ pub(super) fn complexity_cpu(node: &Arc) -> Complexity { Some(expression_complexity(&expr.expr)) }; } - return n.unwrap_or(Complexity::Constant(1)); + return n.unwrap_or(Complexity::Constant(1.)); } // RepartitionExec with Hash: computes hash per row + take_arrays @@ -243,37 +243,37 @@ pub(super) fn complexity_cpu(node: &Arc) -> Complexity { // Limit: just counts rows and stops early. // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/limit.rs if node.is::() || node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } // CoalescePartitionsExec: receives batches from partitions, just passes through the record // batches in a zero copy manner. // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/coalesce_partitions.rs if node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.0); } // BroadcastExec: This node does not do any computation, does not even read the data. if node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } // UnionExec: combines multiple input streams, no processing // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/union.rs if node.is::() || node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } // InterleaveExec: round-robin merging of inputs, no processing // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/union.rs if node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } // EmptyExec: produces no data // https://github.com/apache/datafusion/blob/branch-52/datafusion/physical-plan/src/empty.rs if node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } // For unknown node types, assume we have to do an O(N) operation over all the rows. @@ -288,7 +288,7 @@ struct BytesPerRow { fn expression_complexity(expression: &Arc) -> Complexity { _expression_complexity(expression) .processed - .unwrap_or(Complexity::Constant(1)) + .unwrap_or(Complexity::Constant(1.)) } /// Computes the complexity of processing a join key expression, including the cost of @@ -309,7 +309,7 @@ fn join_key_complexity(expression: &Arc, from_left: bool) -> C None => Complexity::Linear(linear), }); } - result.unwrap_or(Complexity::Constant(1)) + result.unwrap_or(Complexity::Constant(1.)) } /// Computes the per-row processing cost of a join filter predicate. @@ -342,7 +342,7 @@ fn remap_filter_columns(c: Complexity, column_indices: &[ColumnIndex]) -> Comple index, side: JoinSide::Right, }) => Complexity::Linear(LinearComplexity::ColumnFromRight(*index)), - _ => Complexity::Constant(1), + _ => Complexity::Constant(1.), }, // `expression_complexity` only ever emits `Column` linear terms, but keep the rest // intact so the remapping stays total. @@ -374,7 +374,7 @@ fn hashed_or_sorted_key_complexity(expression: &Arc) -> Comple None => Complexity::Linear(LinearComplexity::Column(*col_idx)), }); } - result.unwrap_or(Complexity::Constant(1)) + result.unwrap_or(Complexity::Constant(1.)) } fn _expression_complexity(expression: &Arc) -> BytesPerRow { diff --git a/src/distributed_planner/statistics/complexity_memory.rs b/src/distributed_planner/statistics/complexity_memory.rs index 6458dfcc..f7e2a0c6 100644 --- a/src/distributed_planner/statistics/complexity_memory.rs +++ b/src/distributed_planner/statistics/complexity_memory.rs @@ -49,7 +49,7 @@ pub(super) fn complexity_memory(node: &Arc) -> Complexity { // SortMergeJoinExec streams sorted inputs and only keeps bounded merge state. // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs if node.is::() { - return Complexity::Constant(0); + return Complexity::Constant(0.); } // SymmetricHashJoinExec keeps hash tables for both inputs. @@ -64,7 +64,7 @@ pub(super) fn complexity_memory(node: &Arc) -> Complexity { // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/aggregates/mod.rs if let Some(agg) = node.downcast_ref::() { if agg.group_expr().is_true_no_grouping() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } return Complexity::Linear(LinearComplexity::AllOutputColumns); @@ -78,13 +78,13 @@ pub(super) fn complexity_memory(node: &Arc) -> Complexity { } if node.is::() { - return Complexity::Constant(1); + return Complexity::Constant(1.); } // SortPreservingMergeExec performs a streaming K-way merge. // https://github.com/apache/datafusion/blob/branch-54/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs if node.is::() { - return Complexity::Constant(0); + return Complexity::Constant(0.); } // BroadcastExec retains batches so several consumers can replay the same input partition. @@ -92,7 +92,7 @@ pub(super) fn complexity_memory(node: &Arc) -> Complexity { return Complexity::Linear(LinearComplexity::AllOutputColumns); } - Complexity::Constant(0) + Complexity::Constant(0.) } #[cfg(test)] diff --git a/src/distributed_planner/statistics/complexity_network.rs b/src/distributed_planner/statistics/complexity_network.rs index 9a246633..c52b3393 100644 --- a/src/distributed_planner/statistics/complexity_network.rs +++ b/src/distributed_planner/statistics/complexity_network.rs @@ -1,13 +1,26 @@ -use crate::NetworkBoundaryExt; use crate::distributed_planner::statistics::complexity::{Complexity, LinearComplexity}; +use crate::{BroadcastExec, NetworkBoundaryExt}; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; -/// Calculates the memory cost for the provided node, without recursing into children. +/// Calculates the network cost for the provided node, without recursing into children. pub(super) fn complexity_network(node: &Arc) -> Complexity { if node.is_network_boundary() { - return Complexity::Linear(LinearComplexity::AllOutputColumns); + let bytes = Complexity::Linear(LinearComplexity::AllOutputColumns); + // A broadcast boundary replicates its input to every consumer task, so the bytes that + // cross the network scale with the consumer task count, unlike shuffle/coalesce which move + // each row once. The broadcast boundary's only child is always a BroadcastExec (enforced + // by NetworkBroadcastExec::try_new), which carries the replication factor. + if let Some(bcast) = node + .children() + .first() + .copied() + .and_then(|c| c.downcast_ref::()) + { + return bytes.multiply(Complexity::Constant(bcast.consumer_task_count() as f32)); + } + return bytes; } - Complexity::Constant(0) + Complexity::Constant(0.) } diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs index fb926d62..c169f6fe 100644 --- a/src/distributed_planner/statistics/cost.rs +++ b/src/distributed_planner/statistics/cost.rs @@ -16,9 +16,9 @@ pub(crate) struct Cost { impl AddAssign for Cost { fn add_assign(&mut self, rhs: Self) { - self.cpu += rhs.cpu; - self.memory += rhs.memory; - self.network += rhs.network; + self.cpu = self.cpu.saturating_add(rhs.cpu); + self.memory = self.memory.saturating_add(rhs.memory); + self.network = self.network.saturating_add(rhs.network); } } @@ -38,10 +38,16 @@ fn f(plan: &Arc) -> Result<(Cost, Arc)> { let stats = plan_statistics(plan, &child_stats)?; let cpu = complexity_cpu(plan); - acc_cost.cpu += cpu.cost(&stats, &child_stats).unwrap_or(0); + acc_cost.cpu = acc_cost + .cpu + .saturating_add(cpu.cost(&stats, &child_stats).unwrap_or(0)); let memory = complexity_memory(plan); - acc_cost.memory += memory.cost(&stats, &child_stats).unwrap_or(0); + acc_cost.memory = acc_cost + .memory + .saturating_add(memory.cost(&stats, &child_stats).unwrap_or(0)); let network = complexity_network(plan); - acc_cost.network += network.cost(&stats, &child_stats).unwrap_or(0); + acc_cost.network = acc_cost + .network + .saturating_add(network.cost(&stats, &child_stats).unwrap_or(0)); Ok((acc_cost, stats)) } diff --git a/src/distributed_planner/statistics/plan_statistics.rs b/src/distributed_planner/statistics/plan_statistics.rs index 2265fb4e..d3d54b8e 100644 --- a/src/distributed_planner/statistics/plan_statistics.rs +++ b/src/distributed_planner/statistics/plan_statistics.rs @@ -42,8 +42,7 @@ pub(super) fn plan_statistics( break; }; - // If some of the NDVs are not present in one of the column-level stats, assume the - // worst and use the same as the input number of rows. + // If a column's NDV is absent, fall back to a fraction of the row count if matches!(col_stats.distinct_count, Precision::Absent) { let fallback_ndv = ((*rows as f64) * FALLBACK_NDV_RATIO) as usize; col_stats.distinct_count = Precision::Inexact(fallback_ndv); @@ -58,16 +57,18 @@ pub(super) fn plan_statistics( // If it turns out that we do not have `byte_size` stats, but we do have an estimated number // of rows, do a best-effort in trying to infer the byte size for each column. if matches!(col_stats.byte_size, Precision::Absent) { - col_stats.byte_size = Precision::Inexact(default_bytes_for_datatype(dt) * rows) + col_stats.byte_size = + Precision::Inexact(default_bytes_for_datatype(dt).saturating_mul(*rows)) } } // If bytes are absent, let's just infer them based on the schema and the // number of rows. if matches!(stats.total_byte_size, Precision::Absent) { - let mut total_byte_size = 0; + let mut total_byte_size: usize = 0; for col_stats in &stats.column_statistics { - total_byte_size += col_stats.byte_size.get_value().unwrap_or(&0); + total_byte_size = + total_byte_size.saturating_add(*col_stats.byte_size.get_value().unwrap_or(&0)); } stats.total_byte_size = Precision::Inexact(total_byte_size); } From 2203b3bfdd93b1ad258e60a6dc668c9911b7b9ff Mon Sep 17 00:00:00 2001 From: Gabriel Musat Mestre Date: Sat, 23 May 2026 14:26:35 +0200 Subject: [PATCH 10/20] Support dynamic task count assignation --- .github/workflows/ci.yml | 43 +- Cargo.lock | 1 + Cargo.toml | 1 + benchmarks/cdk/bin/datafusion-bench.ts | 10 + benchmarks/src/run.rs | 5 + src/common/mod.rs | 2 + src/common/recursion.rs | 1 + src/common/vec.rs | 80 +++ src/coordinator/distributed.rs | 8 +- src/coordinator/mod.rs | 1 + src/coordinator/prepare_dynamic_plan.rs | 350 +++++++++++ src/coordinator/prepare_static_plan.rs | 1 + src/coordinator/query_coordinator.rs | 52 +- src/distributed_ext.rs | 74 +++ src/distributed_planner/distributed_config.rs | 9 + .../distributed_query_planner.rs | 8 + .../inject_network_boundaries.rs | 12 +- src/distributed_planner/mod.rs | 6 +- src/distributed_planner/network_boundary.rs | 61 +- .../prepare_network_boundaries.rs | 7 +- .../benchmarks/shuffle_bench.rs | 1 + .../benchmarks/transport_bench.rs | 1 + src/execution_plans/mod.rs | 2 + src/execution_plans/network_broadcast.rs | 13 +- src/execution_plans/network_coalesce.rs | 20 +- src/execution_plans/network_shuffle.rs | 18 +- src/execution_plans/sampler.rs | 594 ++++++++++++++++++ src/metrics/bytes_metric.rs | 11 + src/metrics/task_metrics_rewriter.rs | 2 + src/protobuf/distributed_codec.rs | 26 +- src/stage.rs | 72 ++- src/work_unit_feed/remote_work_unit_feed.rs | 10 + src/worker/generated/worker.rs | 36 +- src/worker/impl_coordinator_channel.rs | 36 +- src/worker/task_data.rs | 7 +- src/worker/worker.proto | 27 + tests/clickbench_correctness_test.rs | 6 +- tests/metrics_collection.rs | 32 + tests/stateful_data_cleanup.rs | 23 +- tests/tpcds_correctness_test.rs | 6 +- tests/tpch_correctness_test.rs | 6 +- 41 files changed, 1600 insertions(+), 81 deletions(-) create mode 100644 src/common/vec.rs create mode 100644 src/coordinator/prepare_dynamic_plan.rs create mode 100644 src/execution_plans/sampler.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cc217c7b..3c5da3fd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,19 +41,33 @@ jobs: - uses: ./.github/actions/setup - run: cargo test --features integration - tpch-test: + tpch-correctness-test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + planning_mode: [ "adaptive", "static" ] + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/setup + - run: cargo test --features tpch --test tpch_correctness_test + env: + ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }} + + tpch-plans-test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - uses: ./.github/actions/setup - - run: cargo test --features tpch --test 'tpch_*' + - run: cargo test --features tpch --test tpch_plans_test tpcds-correctness-test: runs-on: ubuntu-latest strategy: fail-fast: false matrix: - shard: ["01", "02", "03", "04", "05", "06", "07", "08", "09", "10"] + shard: [ "01", "02", "03", "04", "05", "06", "07", "08", "09", "10" ] + planning_mode: [ "adaptive", "static" ] steps: - uses: actions/checkout@v4 - uses: ./.github/actions/setup @@ -62,6 +76,8 @@ jobs: path: testdata/tpcds/main.zip key: "main.zip" - run: cargo test --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }} + env: + ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }} tpcds-plans-test: runs-on: ubuntu-latest @@ -74,7 +90,24 @@ jobs: key: "main.zip" - run: cargo test --features tpcds --test tpcds_plans_test - clickbench-test: + clickbench-correctness-test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + planning_mode: [ "adaptive", "static" ] + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/setup + - uses: actions/cache@v4 + with: + path: testdata/clickbench/ + key: "data" + - run: cargo test --features clickbench --test clickbench_correctness_test + env: + ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }} + + clickbench-plans-test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -83,7 +116,7 @@ jobs: with: path: testdata/clickbench/ key: "data" - - run: cargo test --features clickbench --test 'clickbench_*' + - run: cargo test --features clickbench --test clickbench_plans_test format-check: runs-on: ubuntu-latest diff --git a/Cargo.lock b/Cargo.lock index ebe698d2..7044adaa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2208,6 +2208,7 @@ dependencies = [ "insta", "itertools 0.14.0", "moka", + "num-traits", "object_store", "parquet", "pin-project", diff --git a/Cargo.toml b/Cargo.toml index 4d6e3e7e..fd5ceab2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,7 @@ moka = { version = "0.12", features = ["sync", "future"] } crossbeam-queue = "0.3" sysinfo = { version = "0.30", optional = true } sketches-ddsketch = { version = "0.3", features = ["use_serde"] } +num-traits = "0.2" bincode = "1" tonic-prost = "0.14.2" diff --git a/benchmarks/cdk/bin/datafusion-bench.ts b/benchmarks/cdk/bin/datafusion-bench.ts index f5d15eea..f32af829 100644 --- a/benchmarks/cdk/bin/datafusion-bench.ts +++ b/benchmarks/cdk/bin/datafusion-bench.ts @@ -24,6 +24,8 @@ async function main() { .option('--max-tasks-per-stage ', 'Max tasks per stage', '0') .option('--repartition-file-min-size ', 'repartition_file_min_size DF option', '10485760' /* upstream default */) .option('--target-partitions ', 'target_partitions DF option', '8') + .option('--dynamic ', 'Use the dynamic task count assigner', 'false') + .option('--bytes-per-partition-per-second ', 'Target throughput in bytes per partition per second for the dynamic task count allocator', `${16 * 1024 * 1024}`) .option('--queries ', 'Specific queries to run', undefined) .option('--debug ', 'Print the generated plans to stdout') .option('--warmup ', 'Perform a warmup query before the benchmarks', 'true') @@ -46,6 +48,8 @@ async function main() { const childrenIsolatorUnions = options.childrenIsolatorUnions === 'true' || options.childrenIsolatorUnions === 1 const broadcastJoins = options.broadcastJoins === 'true' || options.broadcastJoins === 1 const partialReduce = options.partialReduce === 'true' || options.partialReduce === 1 + const dynamicTaskCount = options.dynamic === 'true' || options.dynamic === 1 + const bytesPerPartitionPerSecond = parseInt(options.bytesPerPartitionPerSecond) const debug = options.debug === true || options.debug === 'true' || options.debug === 1 const warmup = options.warmup === true || options.warmup === 'true' || options.warmup === 1 @@ -59,6 +63,8 @@ async function main() { compression, broadcastJoins, partialReduce, + dynamicTaskCount, + bytesPerPartitionPerSecond, maxTasksPerStage, repartitionFileMinSize, targetPartitions @@ -98,6 +104,8 @@ class DataFusionRunner implements BenchmarkRunner { childrenIsolatorUnions: boolean; broadcastJoins: boolean; partialReduce: boolean; + dynamicTaskCount: boolean; + bytesPerPartitionPerSecond: number; maxTasksPerStage: number; repartitionFileMinSize: number; targetPartitions: number; @@ -177,6 +185,8 @@ class DataFusionRunner implements BenchmarkRunner { SET distributed.children_isolator_unions=${this.options.childrenIsolatorUnions}; SET distributed.broadcast_joins=${this.options.broadcastJoins}; SET distributed.partial_reduce=${this.options.partialReduce}; + SET distributed.dynamic_task_count=${this.options.dynamicTaskCount}; + SET distributed.bytes_per_partition_per_second=${this.options.bytesPerPartitionPerSecond}; SET distributed.max_tasks_per_stage=${this.options.maxTasksPerStage}; SET datafusion.optimizer.repartition_file_min_size=${this.options.repartitionFileMinSize}; SET datafusion.execution.target_partitions=${this.options.targetPartitions}; diff --git a/benchmarks/src/run.rs b/benchmarks/src/run.rs index 94345f30..7ecaee98 100644 --- a/benchmarks/src/run.rs +++ b/benchmarks/src/run.rs @@ -106,6 +106,10 @@ pub struct RunOpt { #[structopt(long, default_value = "0")] max_tasks_per_stage: usize, + /// Activate dynamic task count + #[structopt(long)] + dynamic: bool, + /// Number of iterations of each test run #[structopt(short = "i", long = "iterations", default_value = "5")] iterations: usize, @@ -203,6 +207,7 @@ impl RunOpt { .with_distributed_cardinality_effect_task_scale_factor( self.cardinality_task_sf.unwrap_or(1.0), )? + .with_distributed_dynamic_task_count(self.dynamic)? .with_distributed_compression(match self.compression.as_str() { "zstd" => Some(CompressionType::ZSTD), "lz4" => Some(CompressionType::LZ4_FRAME), diff --git a/src/common/mod.rs b/src/common/mod.rs index bf9ed549..18cb28a1 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -5,6 +5,7 @@ mod recursion; mod task_context_helpers; mod time; mod uuid; +mod vec; pub(crate) use children_helpers::require_one_child; pub(crate) use on_drop_stream::on_drop_stream; @@ -13,3 +14,4 @@ pub(crate) use recursion::TreeNodeExt; pub(crate) use task_context_helpers::task_ctx_with_extension; pub(crate) use time::now_ns; pub(crate) use uuid::{deserialize_uuid, serialize_uuid}; +pub(crate) use vec::{element_wise_sum, vec_avg_reduce, vec_cast, vec_div, vec_mul}; diff --git a/src/common/recursion.rs b/src/common/recursion.rs index 2f2b9463..5c16116e 100644 --- a/src/common/recursion.rs +++ b/src/common/recursion.rs @@ -589,6 +589,7 @@ mod tests { query_id: uuid::Uuid::nil(), num: 0, workers: vec![], + runtime_stats: None, })) .unwrap() } diff --git a/src/common/vec.rs b/src/common/vec.rs new file mode 100644 index 00000000..ec62678b --- /dev/null +++ b/src/common/vec.rs @@ -0,0 +1,80 @@ +use datafusion::common::internal_err; +use datafusion::error::Result; +use num_traits::AsPrimitive; +use std::ops::{AddAssign, DivAssign, MulAssign}; + +/// Converts a slice of type `I` to a `Vec` using `as`-style primitive casting. +pub(crate) fn vec_cast(input: &[I]) -> Vec +where + I: AsPrimitive, + O: Copy + 'static, +{ + input.iter().map(|v| v.as_()).collect() +} + +/// Adds each element of `other` into the corresponding element of `one`, converting types via `AsPrimitive`. +pub(crate) fn element_wise_sum(mut one: Vec, other: &[O]) -> Result> +where + I: AddAssign + Copy + 'static, + O: AsPrimitive + 'static, +{ + if one.len() != other.len() { + return internal_err!("Cannot do an element wise sum of two vectors of different lengths"); + } + for i in 0..one.len() { + one[i] += other[i].as_(); + } + Ok(one) +} + +/// Multiplies every element of `one` by the scalar `other`, converting types via `AsPrimitive`. +pub(crate) fn vec_mul(mut one: Vec, other: O) -> Vec +where + I: MulAssign + Copy + 'static, + O: AsPrimitive + 'static, +{ + for el in one.iter_mut() { + *el *= other.as_(); + } + one +} + +/// Divides every element of `one` by the scalar `other`, converting types via `AsPrimitive`. +pub(crate) fn vec_div(mut one: Vec, other: O) -> Vec +where + I: DivAssign + Copy + 'static, + O: AsPrimitive + 'static, +{ + for el in one.iter_mut() { + *el /= other.as_(); + } + one +} + +/// Reduces a collection of same-length `f32` vectors into a single vector by averaging element-wise. +/// Empty inner vecs are skipped; returns an empty vec if all inputs are empty. +pub(crate) fn vec_avg_reduce(vecs: Vec>) -> Result> { + let sample_count = vecs.len(); + let mut iter = vecs.into_iter(); + let mut acc = loop { + let Some(v) = iter.next() else { + return Ok(vec![]); + }; + if !v.is_empty() { + break v; + } + }; + for v in iter { + if v.is_empty() { + continue; + } else if acc.len() != v.len() { + return internal_err!( + "vec_avg_reduce: length mismatch — first vec has {} elements, got {}", + acc.len(), + v.len() + ); + } + acc = element_wise_sum(acc, &v)?; + } + Ok(vec_div(acc, sample_count as f32)) +} diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index fe1bbff3..d7d62a08 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,5 +1,7 @@ +use crate::DistributedConfig; use crate::common::{require_one_child, serialize_uuid}; use crate::coordinator::metrics_store::MetricsStore; +use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; use crate::distributed_planner::NetworkBoundaryExt; @@ -198,7 +200,11 @@ impl ExecutionPlan for DistributedExec { builder.spawn(async move { let _guard = query_coordinator.end_query_guard(); - let result = prepare_static_plan(&query_coordinator, &base_plan)?; + let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?; + let result = match d_cfg.dynamic_task_count { + true => prepare_dynamic_plan(&query_coordinator, &base_plan).await?, + false => prepare_static_plan(&query_coordinator, &base_plan)?, + }; plan_for_viz .lock() diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index 8fe771d3..c1a8a8dd 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,6 +1,7 @@ mod distributed; mod latency_metric; mod metrics_store; +mod prepare_dynamic_plan; mod prepare_static_plan; mod query_coordinator; diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs new file mode 100644 index 00000000..d90fce5e --- /dev/null +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -0,0 +1,350 @@ +use crate::TaskCountAnnotation::{Desired, Maximum}; +use crate::common::{TreeNodeExt, element_wise_sum, vec_avg_reduce, vec_div, vec_mul}; +use crate::coordinator::distributed::PreparedPlan; +use crate::coordinator::query_coordinator::QueryCoordinator; +use crate::distributed_planner::{ + InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, ProducerHead, calculate_cost, + inject_network_boundaries, +}; +use crate::execution_plans::SamplerExec; +use crate::stage::{LocalStage, RemoteStage}; +use crate::worker::generated::worker as pb; +use crate::{BytesCounterMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage}; +use dashmap::DashMap; +use datafusion::common::stats::Precision; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::{Result, exec_err, plan_err}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Statistics, +}; +use futures::{Stream, StreamExt}; +use std::any::TypeId; +use std::sync::Arc; +use tokio_stream::wrappers::UnboundedReceiverStream; + +pub(super) async fn prepare_dynamic_plan( + query_coordinator: &QueryCoordinator, + base_plan: &Arc, +) -> Result { + let plans_for_viz = Arc::new(PlanReconstructor::default()); + + let head_stage = inject_network_boundaries( + Arc::clone(base_plan), + |mut input_stage: LocalStage, nb_type: TypeId, nb_ctx: &InjectNetworkBoundaryContext| { + let mut metrics = MetricsSet::new(); + + // At this point, input_stage.plan has two kind of leaf nodes: + // - The ones that naturally do not read from any children, like DataSourceExec + // - Network boundaries whose Stage was set to Stage::Remote by a previous iteration + // of this same function. + // Both types of leaf nodes contain very valuable and accurate statistics that are used + // here for computing an estimation of the compute cost (measured in bytes): + // - DataSourceExec (or natural leaf nodes) contain stats pulled directly from their + // data source, like parquet files. + // - Network boundaries contain statistics collected from runtime information, gathered + // by the SamplerExec injected by this same function. + let cost = calculate_cost(&input_stage.plan)?; + metrics.push(BytesCounterMetric::new_metric("cpu_cost", cost.cpu)); + metrics.push(BytesCounterMetric::new_metric("memory_cost", cost.memory)); + metrics.push(BytesCounterMetric::new_metric("network_cost", cost.network)); + let compute_based_task_count = cost + .cpu + .div_ceil(nb_ctx.d_cfg.bytes_per_partition_per_second.max(1)) + .div_ceil(input_stage.plan.output_partitioning().partition_count()) + .clamp(1, nb_ctx.max_tasks()?); + let task_count = nb_ctx + .task_count(&input_stage.plan)? + .merge(Desired(compute_based_task_count)); + + // Propagate the final task_count inferred based on runtime statistics and compute cost. + // Here is where leaf nodes are scaled up by TaskEstimator::scale_up_leaf_node, and the + // plan is finally left ready for distribution. + input_stage.plan = nb_ctx + .propagate_task_count_until_network_boundaries(&input_stage.plan, task_count)?; + input_stage.tasks = task_count.as_usize(); + // In order to infer the compute the cost of the stage above this one, here a sampler + // is injected to gather runtime statistics. + input_stage.plan = ProducerHead::insert_sampler(input_stage.plan)?; + + let mut stage_coordinator = query_coordinator.stage_coordinator(&input_stage); + + let mut workers = Vec::with_capacity(input_stage.tasks); + let mut load_info_rxs = Vec::with_capacity(input_stage.tasks); + + let routed_urls = if input_stage.tasks == 1 { + // If there's an input stage with a single worker, and the current stage is also + // going to run in a single worker, we want to co-locate them so that unnecessary + // network transfers are avoided. + match stage_coordinator.find_input_stage_with_single_url() { + Some(single_url) => vec![single_url], + None => stage_coordinator.routed_urls()?, + } + } else { + stage_coordinator.routed_urls()? + }; + + for (i, routed_url) in routed_urls.into_iter().enumerate() { + workers.push(routed_url.clone()); + // Spawns the task that feeds this subplan to this worker. There will be as + // many as this spawned tasks as workers. + let (worker_tx, worker_rx) = stage_coordinator.send_plan_task(i, routed_url)?; + load_info_rxs.push({ + let rx = stage_coordinator.worker_to_coordinator_task(i, worker_rx); + UnboundedReceiverStream::new(rx) + }); + stage_coordinator.coordinator_to_worker_task(i, worker_tx)?; + } + + let plans_for_viz = Arc::clone(&plans_for_viz); + Ok(async move { + let (stats, consumer_tc) = if nb_type == TypeId::of::() { + (None, Maximum(1)) + } else { + let stats = gather_runtime_statistics(load_info_rxs, &input_stage.plan).await?; + let sampled_bytes = *stats.total_byte_size.get_value().unwrap_or(&0); + metrics.push(BytesCounterMetric::new_metric( + "sampled_bytes", + sampled_bytes, + )); + // returning Desired(1) here is our way to tell the planner that we don't care + // about the task count assigned to the network boundary in the consumer stage, + // and we don't want it to affect other task count decisions. + (Some(Arc::new(stats)), Desired(1)) + }; + + // Capture the output partitioning of the (rescaled, sampler-wrapped) input plan + // before it's moved: the returned stage is remote and carries no plan to read it + // back from. + let input_properties = Arc::clone(input_stage.plan.properties()); + plans_for_viz.insert(input_stage.num, input_stage.plan, metrics); + Ok(NetworkBoundaryBuilderResult { + consumer_task_count: consumer_tc, + input_stage: Stage::Remote(RemoteStage { + query_id: input_stage.query_id, + num: input_stage.num, + workers, + runtime_stats: stats, + }), + input_properties, + }) + }) + }, + query_coordinator.session_config().options(), + ) + .await?; + + Ok(PreparedPlan { + plan_for_viz: plans_for_viz.reconstruct(&head_stage)?, + head_stage, + }) +} + +/// Reconstructs the plan dynamically as stages get transitioned to Remote and get sent to the +/// respective workers. +/// +/// As the [prepare_dynamic_plan] function recurses and progressively sends the plan to workers, the +/// original plan gets modified, and subplans belong to the different [Stage]s get lost as they get +/// transitioned to [Stage::Remote]. +/// +/// This struct is in charge of tracking the [prepare_dynamic_plan] process and storing the final +/// version of all the subplans so that it can be reconstructed into a fully blown plan for +/// visualization purposes. +#[derive(Default)] +struct PlanReconstructor { + stage_map: DashMap, MetricsSet)>, +} + +impl PlanReconstructor { + fn insert(&self, stage: usize, plan: Arc, metrics_set: MetricsSet) { + self.stage_map.insert(stage, (plan, metrics_set)); + } + + fn reconstruct(&self, head_stage: &Arc) -> Result> { + let head_stage = Arc::clone(head_stage); + let reconstructed = head_stage.transform_down_with_task_count(1, |plan, tc| { + let Some(nb) = plan.as_network_boundary() else { + return Ok(Transformed::no(plan)); + }; + let input_stage = nb.input_stage(); + let Some((_, entry)) = self.stage_map.remove(&input_stage.num()) else { + return exec_err!( + "Failed to retrieve plan for stage {} for visualization purposes", + input_stage.num() + ); + }; + let (plan_for_viz, metrics_set) = entry; + + let plan_for_viz = nb.producer_head(tc).insert(plan_for_viz)?; + + let nb = nb.with_input_stage(Stage::Local(LocalStage { + query_id: input_stage.query_id(), + num: input_stage.num(), + plan: plan_for_viz, + tasks: input_stage.task_count(), + metrics_set, + }))?; + + Ok(Transformed::yes(nb)) + })?; + Ok(reconstructed.data) + } +} + +/// Estimates the bytes per second flowing through a stage by reading sample information. +async fn gather_runtime_statistics( + per_task_load_info_stream: Vec + Unpin>, + plan: &Arc, +) -> Result { + const ESTIMATED_QUERY_TIME_S: usize = 10; + const BYTES_READY_SAMPLE_PERCENTAGE: f32 = 0.2; + const BYTES_PER_SECOND_SAMPLE_PERCENTAGE: f32 = 0.2; + + let Some(sampler) = find_sampler(plan) else { + return plan_err!("Mising SamplerExec while gathering load report"); + }; + let n_cols = sampler.schema().fields.len(); + + fn apply_pct(value: usize, pct: f32) -> usize { + (value as f32 * pct).round() as usize + } + + let partitions_per_task = sampler.partition_samplers.len(); + let task_count = per_task_load_info_stream.len(); + let total_partitions = partitions_per_task * task_count; + + let mut partitions_with_bytes_per_second_done = 0; + let mut partitions_with_bytes_ready_done = 0; + let mut partitions_done = 0; + let mut rows_ready = 0; + let mut rows_per_second = 0; + let mut per_col_bytes_ready = vec![0usize; n_cols]; + let mut per_col_bytes_per_second = vec![0usize; n_cols]; + + let mut ndv_pct = vec![]; + let mut null_pct = vec![]; + + let mut load_info_stream = futures::stream::select_all(per_task_load_info_stream); + while let Some(load_info) = load_info_stream.next().await { + rows_per_second += load_info.rows_per_second as usize; + rows_ready += load_info.rows_ready as usize; + per_col_bytes_per_second = element_wise_sum( + per_col_bytes_per_second, + &load_info.per_column_bytes_per_second, + )?; + per_col_bytes_ready = + element_wise_sum(per_col_bytes_ready, &load_info.per_column_bytes_ready)?; + ndv_pct.push(load_info.per_column_ndv_percentage); + null_pct.push(load_info.per_column_null_percentage); + + partitions_with_bytes_per_second_done += + load_info.per_column_bytes_per_second.iter().any(|v| *v > 0) as usize; + partitions_with_bytes_ready_done += + load_info.per_column_bytes_ready.iter().any(|v| *v > 0) as usize; + partitions_done += 1; + + // Short circuit if we collected enough bytes_ready measurements. + if partitions_with_bytes_ready_done + >= apply_pct(total_partitions, BYTES_READY_SAMPLE_PERCENTAGE).max(1) + { + break; + } + + // Short circuit if we collected enough bytes_per_second measurements. + if partitions_with_bytes_per_second_done + >= apply_pct(total_partitions, BYTES_PER_SECOND_SAMPLE_PERCENTAGE).max(1) + { + break; + } + + // Short circuit if there are no further partitions remaining to sample from. + if partitions_done == total_partitions { + break; + } + } + + if partitions_done == 0 { + return Ok(zero_stats(plan.schema().fields.len())); + } + + let per_col_bytes_ready = vec_div( + vec_mul(per_col_bytes_ready, total_partitions), + partitions_done, + ); + let per_col_bytes_per_second = vec_div( + vec_mul(per_col_bytes_per_second, total_partitions), + partitions_done, + ); + + let rows_ready = rows_ready * total_partitions / partitions_done; + let rows_per_second = rows_per_second * total_partitions / partitions_done; + + let total_num_rows = rows_ready + rows_per_second * ESTIMATED_QUERY_TIME_S; + + if total_num_rows == 0 { + return Ok(zero_stats(n_cols)); + } + + let per_col_byte_size = element_wise_sum( + per_col_bytes_ready, + &vec_mul(per_col_bytes_per_second, ESTIMATED_QUERY_TIME_S), + )?; + let total_byte_size: usize = per_col_byte_size.iter().sum(); + + let ndv_pct = vec_avg_reduce(ndv_pct)?; + if ndv_pct.len() != n_cols { + return plan_err!("Expected {n_cols} ndv values, but got {}", ndv_pct.len()); + } + let null_pct = vec_avg_reduce(null_pct)?; + if null_pct.len() != n_cols { + return plan_err!("Expected {n_cols} null values, but got {}", null_pct.len()); + } + + Ok(Statistics { + num_rows: Precision::Inexact(total_num_rows), + total_byte_size: Precision::Inexact(total_byte_size), + column_statistics: ndv_pct + .into_iter() + .zip(null_pct) + .zip(per_col_byte_size) + .map(|((ndv, null), col_bytes)| ColumnStatistics { + null_count: Precision::Inexact((null * total_num_rows as f32) as usize), + distinct_count: Precision::Inexact((ndv * total_num_rows as f32) as usize), + byte_size: Precision::Inexact(col_bytes), + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + }) + .collect(), + }) +} + +fn find_sampler(plan: &Arc) -> Option<&SamplerExec> { + let mut sampler = None; + plan.apply(|plan| { + if let Some(node) = plan.downcast_ref::() { + sampler = Some(node); + return Ok(TreeNodeRecursion::Stop); + }; + Ok(TreeNodeRecursion::Continue) + }) + .expect("Cannot fail"); + sampler +} + +fn zero_stats(n_cols: usize) -> Statistics { + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: (0..n_cols) + .map(|_| ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + distinct_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + }) + .collect(), + } +} diff --git a/src/coordinator/prepare_static_plan.rs b/src/coordinator/prepare_static_plan.rs index 65d74276..3bfc2d2f 100644 --- a/src/coordinator/prepare_static_plan.rs +++ b/src/coordinator/prepare_static_plan.rs @@ -50,6 +50,7 @@ pub(super) fn prepare_static_plan( query_id: stage.query_id, num: stage.num, workers, + runtime_stats: None, }, ))?)) })?; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 3d19c750..6f0010e6 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -12,17 +12,19 @@ use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; use crate::{ BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, - DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, TaskEstimator, - TaskKey, TaskRoutingContext, get_distributed_channel_resolver, get_distributed_worker_resolver, + DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, NetworkBoundaryExt, + Stage, TaskEstimator, TaskKey, TaskRoutingContext, get_distributed_channel_resolver, + get_distributed_worker_resolver, }; use datafusion::common::instant::Instant; use datafusion::common::runtime::JoinSet; -use datafusion::common::tree_node::{Transformed, TreeNodeRecursion}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{DataFusionError, exec_datafusion_err}; use datafusion::common::{Result, exec_err}; 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}; @@ -88,6 +90,11 @@ impl QueryCoordinator { } } + /// Returns the [SessionConfig] for the current query. + pub(super) fn session_config(&self) -> &SessionConfig { + self.task_ctx.session_config() + } + /// returns a guard that, when dropped, it signals all the coordinator->worker connections that /// the query is finished, ending them, and propagating the EOS to the workers so that they can /// clean up any remaining state. @@ -200,8 +207,8 @@ impl<'a> StageCoordinator<'a> { let mut worker_to_coordinator_stream = response.into_inner(); while let Some(msg_or_err) = worker_to_coordinator_stream.next().await { let msg = msg_or_err.map_err(|err| { - tonic_status_to_datafusion_error(err).unwrap_or_else(|| { - exec_datafusion_err!("Unknown error on worker to coordinator stream") + tonic_status_to_datafusion_error(&err).unwrap_or_else(|| { + exec_datafusion_err!("Unknown error on worker to coordinator stream: {err}") }) })?; if worker_to_coordinator_tx.send(msg).is_err() { @@ -221,13 +228,15 @@ impl<'a> StageCoordinator<'a> { &mut self, task_i: usize, mut worker_to_coordinator_rx: UnboundedReceiver, - ) { + ) -> UnboundedReceiver { let task_key = TaskKey { query_id: serialize_uuid(&self.query_id), stage_id: self.stage_id as u64, task_number: task_i as u64, }; let task_metrics = self.metrics_store.clone(); + let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut load_info_tx_opt = Some(load_info_tx); // Cannot use self.join_set because that's tied to the lifetime of the query, and the // metrics collection process might outlive the query's lifetime. @@ -242,9 +251,18 @@ impl<'a> StageCoordinator<'a> { task_metrics.insert(task_key.clone(), pre_order_metrics); } } + pb::worker_to_coordinator_msg::Inner::LoadInfo(load_info) => { + if let Some(tx) = &load_info_tx_opt { + let _ = tx.send(load_info); + } + } + pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { + let _ = load_info_tx_opt.take(); + } } } }); + load_info_rx } /// Spawns a background task in charge of sending messages to workers. Some things that are sent @@ -401,6 +419,28 @@ impl<'a> StageCoordinator<'a> { } Ok(routed_urls) } + + pub(super) fn find_input_stage_with_single_url(&self) -> Option { + let mut single_stage_url = None; + self.plan + .apply(|plan| { + let Some(nb) = plan.as_network_boundary() else { + return Ok(TreeNodeRecursion::Continue); + }; + + if let Stage::Remote(remote) = nb.input_stage() + && remote.workers.len() == 1 + { + single_stage_url = Some(remote.workers[0].clone()); + return Ok(TreeNodeRecursion::Stop); + } + + Ok(TreeNodeRecursion::Jump) + }) + .expect("Cannot fail"); + + single_stage_url + } } fn keep_stream_alive(notify: Arc) -> impl Stream + 'static { diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 17852a04..d33b18a6 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -577,6 +577,27 @@ pub trait DistributedExt: Sized { P: WorkUnitFeedProvider + 'static, P::WorkUnit: 'static, F: Fn(&T) -> Option<&WorkUnitFeed

> + Send + Sync + 'static; + + /// Dynamically allocates tasks to the different stages based on runtime statistics + /// collected during execution. + fn with_distributed_dynamic_task_count(self, enabled: bool) -> Result; + + /// Same as [DistributedExt::with_distributed_dynamic_task_count] but with an in-place mutation. + fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>; + + /// Target throughput in bytes per partition per second used by the dynamic task count + /// allocator to decide how many tasks to assign to each stage based on runtime statistics. + fn with_distributed_bytes_per_partition_per_second( + self, + bytes_per_partition_per_second: usize, + ) -> Result; + + /// Same as [DistributedExt::with_distributed_bytes_per_partition_per_second] but with an + /// in-place mutation. + fn set_distributed_bytes_per_partition_per_second( + &mut self, + bytes_per_partition_per_second: usize, + ) -> Result<(), DataFusionError>; } impl DistributedExt for SessionConfig { @@ -722,6 +743,21 @@ impl DistributedExt for SessionConfig { }) } + fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError> { + let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; + d_cfg.dynamic_task_count = enabled; + Ok(()) + } + + fn set_distributed_bytes_per_partition_per_second( + &mut self, + bytes_per_partition_per_second: usize, + ) -> Result<(), DataFusionError> { + let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; + d_cfg.bytes_per_partition_per_second = bytes_per_partition_per_second; + Ok(()) + } + delegate! { to self { #[call(set_distributed_option_extension)] @@ -804,6 +840,14 @@ impl DistributedExt for SessionConfig { P: WorkUnitFeedProvider + 'static, P::WorkUnit: 'static, F: Fn(&T) -> Option<&WorkUnitFeed

> + Send + Sync + 'static; + + #[call(set_distributed_dynamic_task_count)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result; + + #[call(set_distributed_bytes_per_partition_per_second)] + #[expr($?;Ok(self))] + fn with_distributed_bytes_per_partition_per_second(mut self, bytes_per_partition_per_second: usize) -> Result; } } } @@ -915,6 +959,16 @@ impl DistributedExt for SessionStateBuilder { P: WorkUnitFeedProvider + 'static, P::WorkUnit: 'static, F: Fn(&T) -> Option<&WorkUnitFeed

> + Send + Sync + 'static; + + fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_task_count)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result; + + fn set_distributed_bytes_per_partition_per_second(&mut self, bytes_per_partition_per_second: usize) -> Result<(), DataFusionError>; + #[call(set_distributed_bytes_per_partition_per_second)] + #[expr($?;Ok(self))] + fn with_distributed_bytes_per_partition_per_second(mut self, bytes_per_partition_per_second: usize) -> Result; } } } @@ -1026,6 +1080,16 @@ impl DistributedExt for SessionState { P: WorkUnitFeedProvider + 'static, P::WorkUnit: 'static, F: Fn(&T) -> Option<&WorkUnitFeed

> + Send + Sync + 'static; + + fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_task_count)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result; + + fn set_distributed_bytes_per_partition_per_second(&mut self, bytes_per_partition_per_second: usize) -> Result<(), DataFusionError>; + #[call(set_distributed_bytes_per_partition_per_second)] + #[expr($?;Ok(self))] + fn with_distributed_bytes_per_partition_per_second(mut self, bytes_per_partition_per_second: usize) -> Result; } } } @@ -1137,6 +1201,16 @@ impl DistributedExt for SessionContext { P: WorkUnitFeedProvider + 'static, P::WorkUnit: 'static, F: Fn(&T) -> Option<&WorkUnitFeed

> + Send + Sync + 'static; + + fn set_distributed_dynamic_task_count(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_task_count)] + #[expr($?;Ok(self))] + fn with_distributed_dynamic_task_count(self, enabled: bool) -> Result; + + fn set_distributed_bytes_per_partition_per_second(&mut self, bytes_per_partition_per_second: usize) -> Result<(), DataFusionError>; + #[call(set_distributed_bytes_per_partition_per_second)] + #[expr($?;Ok(self))] + fn with_distributed_bytes_per_partition_per_second(self, bytes_per_partition_per_second: usize) -> Result; } } } diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index d6faee39..e01b5ff4 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -66,6 +66,15 @@ extensions_options! { /// budget will still be admitted (otherwise we would livelock), so the actual peak per /// connection is `worker_connection_buffer_budget_bytes + max_message_size`. pub worker_connection_buffer_budget_bytes: usize, default = 64 * 1024 * 1024 + /// Calculates the task count of the different stages at execution time, based on runtime + /// information collected by sampling at the head of the stages. + /// + /// With this option enabled, the shape of the distributed plan is only known after fully + /// executing it, as it's dynamically created on the fly during execution. + pub dynamic_task_count: bool, default = false + /// If `dynamic_task_count` is enabled, this value is the amount of bytes/second each + /// partition is expected to handle. Lower values will result in greater parallelism. + pub bytes_per_partition_per_second: usize, default = 16 * 1024 * 1024 /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to /// estimate how many tasks should be spawned for the [Stage] containing the leaf node. pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default() diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index f5ac6168..8e2d8f2c 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/src/distributed_planner/distributed_query_planner.rs @@ -108,6 +108,14 @@ impl QueryPlanner for DistributedQueryPlanner { plan = insert_broadcast_execs(plan, cfg)?; + if d_cfg.dynamic_task_count { + // The task count will be decided dynamically at execution time. + return Ok(Arc::new( + DistributedExec::new(plan).with_metrics_collection(d_cfg.collect_metrics), + )); + } + + // Compute per-node task counts and inject `Network*Exec` nodes at the stage boundaries. plan = inject_network_boundaries(plan, CardinalityBasedNetworkBoundaryBuilder, cfg).await?; plan = prepare_network_boundaries(plan)?; diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 21898440..78fbe6b9 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -152,8 +152,9 @@ pub(crate) async fn inject_network_boundaries( #[derive(Clone)] pub(crate) struct InjectNetworkBoundaryContext<'a> { + pub(crate) d_cfg: &'a DistributedConfig, + cfg: &'a ConfigOptions, - d_cfg: &'a DistributedConfig, nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync), task_counts: &'a Mutex>, query_id: Uuid, @@ -161,7 +162,7 @@ pub(crate) struct InjectNetworkBoundaryContext<'a> { } impl<'a> InjectNetworkBoundaryContext<'a> { - fn max_tasks(&self) -> Result { + pub(crate) fn max_tasks(&self) -> Result { Ok(match self.d_cfg.max_tasks_per_stage { 0 => self .d_cfg @@ -190,7 +191,7 @@ impl<'a> InjectNetworkBoundaryContext<'a> { plan } - fn task_count(&self, plan: &Arc) -> Result { + pub(crate) fn task_count(&self, plan: &Arc) -> Result { let Some(task_count) = self .task_counts .lock() @@ -294,6 +295,7 @@ async fn _inject_network_boundaries( num: nb_ctx.fetch_add_stage_id(), plan: nb_ctx.plan_with_task_count(plan, task_count), tasks: task_count.as_usize(), + metrics_set: Default::default(), }; let result = nb_ctx .nb_builder @@ -323,6 +325,7 @@ async fn _inject_network_boundaries( num: nb_ctx.fetch_add_stage_id(), plan: nb_ctx.plan_with_task_count(plan, task_count), tasks: task_count.as_usize(), + metrics_set: Default::default(), }; let result = nb_ctx .nb_builder @@ -339,6 +342,7 @@ async fn _inject_network_boundaries( num: nb_ctx.fetch_add_stage_id(), plan: nb_ctx.plan_with_task_count(plan, task_count), tasks: task_count.as_usize(), + metrics_set: Default::default(), }; let result = nb_ctx .nb_builder @@ -409,7 +413,7 @@ async fn _inject_network_boundaries( /// - **Everything else**: recurse into children with the same `task_count`, then rebuild the /// node with the rebuilt children. impl InjectNetworkBoundaryContext<'_> { - fn propagate_task_count_until_network_boundaries( + pub(crate) fn propagate_task_count_until_network_boundaries( &self, plan: &Arc, task_count: TaskCountAnnotation, diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index f12c684c..18cfe6e0 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -11,8 +11,12 @@ mod statistics; mod task_estimator; 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(crate) use network_boundary::{ProducerHead, insert_producer_head}; pub use session_state_builder_ext::SessionStateBuilderExt; +pub(crate) use statistics::calculate_cost; pub(crate) use task_estimator::set_distributed_task_estimator; pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext}; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index 858a7e00..e6479e7e 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -1,3 +1,4 @@ +use crate::execution_plans::SamplerExec; use crate::{BroadcastExec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage}; use datafusion::common::Result; use datafusion::physical_expr::Partitioning; @@ -13,7 +14,7 @@ pub trait NetworkBoundary: ExecutionPlan { /// information to perform any internal transformations necessary for distributed execution. /// /// Typically, [NetworkBoundary]s will use this call for transitioning from "Pending" to "ready". - fn with_input_stage(&self, input_stage: Stage) -> Result>; + fn with_input_stage(&self, input_stage: Stage) -> Result>; /// Returns the assigned input [Stage], if any. fn input_stage(&self) -> &Stage; @@ -59,28 +60,40 @@ impl NetworkBoundaryExt for dyn ExecutionPlan { } } -/// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This -/// can be called both during planning and lazily at runtime. -pub(crate) fn insert_producer_head( - input: Arc, - head: ProducerHead, -) -> Result> { - let input = if let Some(r_exec) = input.downcast_ref::() { - Arc::clone(r_exec.input()) - } else if let Some(b_exec) = input.downcast_ref::() { - Arc::clone(b_exec.input()) - } else { - input - }; - let plan = match head { - ProducerHead::None => input, - ProducerHead::BroadcastExec { output_partitions } => { - 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)?) +impl ProducerHead { + /// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This + /// can be called both during planning and lazily at runtime. + pub(crate) fn insert(self, input: Arc) -> Result> { + let input = if let Some(r_exec) = input.downcast_ref::() { + Arc::clone(r_exec.input()) + } else if let Some(b_exec) = input.downcast_ref::() { + Arc::clone(b_exec.input()) + } else { + input + }; + let plan = match self { + ProducerHead::None => input, + ProducerHead::BroadcastExec { output_partitions } => { + 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)?) + } + }; + Ok(plan) + } + + /// Injects a [SamplerExec] right below a [RepartitionExec] or [BroadcastExec]. + pub(crate) fn insert_sampler(input: Arc) -> Result> { + if let Some(r_exec) = input.downcast_ref::() { + let child = Arc::clone(r_exec.input()); + input.with_new_children(vec![Arc::new(SamplerExec::new(child))]) + } else if let Some(b_exec) = input.downcast_ref::() { + let child = Arc::clone(b_exec.input()); + input.with_new_children(vec![Arc::new(SamplerExec::new(child))]) + } else { + Ok(input) } - }; - Ok(plan) + } } diff --git a/src/distributed_planner/prepare_network_boundaries.rs b/src/distributed_planner/prepare_network_boundaries.rs index 1b6f1ac7..dd070793 100644 --- a/src/distributed_planner/prepare_network_boundaries.rs +++ b/src/distributed_planner/prepare_network_boundaries.rs @@ -1,5 +1,4 @@ use crate::common::TreeNodeExt; -use crate::distributed_planner::network_boundary::insert_producer_head; use crate::stage::LocalStage; use crate::{NetworkBoundaryExt, Stage}; use datafusion::common::Result; @@ -35,8 +34,9 @@ pub(crate) fn prepare_network_boundaries( // 2) Scale up the head node of the input stage in order to account for the amount of partition // and consumer count above it. - let plan = - insert_producer_head(Arc::clone(&input_stage.plan), nb.producer_head(task_count))?; + let plan = nb + .producer_head(task_count) + .insert(Arc::clone(&input_stage.plan))?; // 3) Make sure the input stage can be uniquely identified with a stage index and query id. // If there were already some `query_id` and `num` that's fine. @@ -45,6 +45,7 @@ pub(crate) fn prepare_network_boundaries( num: stage_id, plan, tasks: input_stage.tasks, + metrics_set: Default::default(), }))?; stage_id += 1; Ok(Transformed::yes(nb)) diff --git a/src/execution_plans/benchmarks/shuffle_bench.rs b/src/execution_plans/benchmarks/shuffle_bench.rs index 8406a052..4f46f1be 100644 --- a/src/execution_plans/benchmarks/shuffle_bench.rs +++ b/src/execution_plans/benchmarks/shuffle_bench.rs @@ -216,6 +216,7 @@ impl ShuffleFixture { query_id, num: 0, workers: self.input_stage_workers.clone(), + runtime_stats: None, }); let mut join_set = JoinSet::default(); diff --git a/src/execution_plans/benchmarks/transport_bench.rs b/src/execution_plans/benchmarks/transport_bench.rs index 847c307b..a386c7bd 100644 --- a/src/execution_plans/benchmarks/transport_bench.rs +++ b/src/execution_plans/benchmarks/transport_bench.rs @@ -266,6 +266,7 @@ impl TransportFixture { query_id, num: 0, workers: self.input_stage_tasks.clone(), + runtime_stats: None, }); let mut join_set = JoinSet::default(); diff --git a/src/execution_plans/mod.rs b/src/execution_plans/mod.rs index a1ea6316..ecdcbc35 100644 --- a/src/execution_plans/mod.rs +++ b/src/execution_plans/mod.rs @@ -6,6 +6,7 @@ mod metrics; mod network_broadcast; mod network_coalesce; mod network_shuffle; +mod sampler; #[cfg(any(test, feature = "integration"))] pub mod benchmarks; @@ -18,3 +19,4 @@ pub(crate) use metrics::MetricsWrapperExec; pub use network_broadcast::NetworkBroadcastExec; pub use network_coalesce::NetworkCoalesceExec; pub use network_shuffle::NetworkShuffleExec; +pub use sampler::SamplerExec; diff --git a/src/execution_plans/network_broadcast.rs b/src/execution_plans/network_broadcast.rs index f9dee080..251ed29c 100644 --- a/src/execution_plans/network_broadcast.rs +++ b/src/execution_plans/network_broadcast.rs @@ -9,7 +9,7 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr_common::metrics::MetricsSet; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, }; use std::fmt::Formatter; use std::sync::Arc; @@ -153,6 +153,7 @@ impl NetworkBroadcastExec { num: 0, plan: input, tasks: producer_tasks, + metrics_set: Default::default(), }), input_properties, )) @@ -160,7 +161,7 @@ impl NetworkBroadcastExec { } impl NetworkBoundary for NetworkBroadcastExec { - fn with_input_stage(&self, input_stage: Stage) -> Result> { + fn with_input_stage(&self, input_stage: Stage) -> Result> { let mut self_clone = self.clone(); self_clone.worker_connections = WorkerConnectionPool::new(input_stage.task_count()); self_clone.input_stage = input_stage; @@ -268,4 +269,12 @@ impl ExecutionPlan for NetworkBroadcastExec { fn metrics(&self) -> Option { Some(self.worker_connections.metrics.clone_inner()) } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.input_stage.partition_statistics( + partition, + self.properties.output_partitioning().partition_count(), + self.schema(), + ) + } } diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 8fb06e74..1582d08c 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -12,7 +12,7 @@ use datafusion::physical_plan::limit::LocalLimitExec; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, PlanProperties, - internal_err, + Statistics, internal_err, }; use std::fmt::{Debug, Formatter}; use std::sync::Arc; @@ -131,6 +131,7 @@ impl NetworkCoalesceExec { num: 0, plan: input, tasks: producer_tasks, + metrics_set: Default::default(), }), input_properties, consumer_tasks, @@ -157,6 +158,7 @@ impl NetworkCoalesceExec { num: local.num, plan: input_with_fetch, tasks: local.tasks, + metrics_set: Default::default(), }); Ok(Arc::new(self_clone)) } @@ -167,7 +169,7 @@ impl NetworkBoundary for NetworkCoalesceExec { &self.input_stage } - fn with_input_stage(&self, input_stage: Stage) -> Result> { + fn with_input_stage(&self, input_stage: Stage) -> Result> { let mut self_clone = self.clone(); self_clone.properties = scale_partitioning_props(self_clone.properties(), |p| { p * input_stage.task_count() / self_clone.input_stage.task_count().max(1) @@ -247,10 +249,8 @@ impl ExecutionPlan for NetworkCoalesceExec { ); } - let partitions_per_task = self - .properties() - .partitioning - .partition_count() + let out_partitions = self.properties().partitioning.partition_count(); + let partitions_per_task = out_partitions .checked_div( self.input_stage .task_count() @@ -311,6 +311,14 @@ impl ExecutionPlan for NetworkCoalesceExec { fn metrics(&self) -> Option { Some(self.worker_connections.metrics.clone_inner()) } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.input_stage.partition_statistics( + partition, + self.properties.output_partitioning().partition_count(), + self.schema(), + ) + } } #[derive(Debug, Clone, Copy)] diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 157cfd99..2a575b42 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -11,7 +11,9 @@ use datafusion::physical_expr::Partitioning; use datafusion::physical_expr_common::metrics::MetricsSet; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics, +}; use std::fmt::Formatter; use std::sync::Arc; use uuid::Uuid; @@ -133,6 +135,7 @@ impl NetworkShuffleExec { num: 0, plan: input, tasks: producer_tasks, + metrics_set: Default::default(), }), input_properties, )) @@ -144,7 +147,7 @@ impl NetworkBoundary for NetworkShuffleExec { &self.input_stage } - fn with_input_stage(&self, input_stage: Stage) -> Result> { + fn with_input_stage(&self, input_stage: Stage) -> Result> { let mut self_clone = self.clone(); self_clone.worker_connections = WorkerConnectionPool::new(input_stage.task_count()); self_clone.input_stage = input_stage; @@ -217,7 +220,8 @@ impl ExecutionPlan for NetworkShuffleExec { }; let task_context = DistributedTaskContext::from_ctx(&context); - let off = self.properties.partitioning.partition_count() * task_context.task_index; + let out_partitions = self.properties.partitioning.partition_count(); + let off = out_partitions * task_context.task_index; let mut streams = Vec::with_capacity(remote_stage.workers.len()); for input_task_index in 0..remote_stage.workers.len() { @@ -242,4 +246,12 @@ impl ExecutionPlan for NetworkShuffleExec { fn metrics(&self) -> Option { Some(self.worker_connections.metrics.clone_inner()) } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.input_stage.partition_statistics( + partition, + self.properties.output_partitioning().partition_count(), + self.schema(), + ) + } } diff --git a/src/execution_plans/sampler.rs b/src/execution_plans/sampler.rs new file mode 100644 index 00000000..0e120219 --- /dev/null +++ b/src/execution_plans/sampler.rs @@ -0,0 +1,594 @@ +use crate::common::{require_one_child, vec_cast}; +use crate::worker::generated::worker as pb; +use crate::{ + BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, MaxGaugeMetric, + MaxLatencyMetric, P50LatencyMetric, +}; +use datafusion::arrow::array::Array; +use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::runtime::SpawnedTask; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::common::{DataFusionError, Result, exec_err}; +use datafusion::common::{HashSet, ScalarValue}; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr_common::metrics::{Gauge, MetricValue, MetricsSet}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; +use std::collections::VecDeque; +use std::fmt::{Debug, Formatter}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, LazyLock, Mutex, OnceLock}; +use std::task::{Context, Poll}; +use std::time::Instant; +use tokio::sync::oneshot; + +/// How many [RecordBatch]s to allow the input stream to yield synchronously (without yielding back +/// to tokio) before short-circuiting buffering. +const READY_CHUNK_LIMIT: usize = 256; +/// Maximum read of bytes per second allowed to be emitted. Reads greater than this will be +/// truncated to this max value, as it's assumed that [READY_CHUNK_LIMIT] was hit and no useful +/// measurement can actually be emitted. +const MAX_BYTES_PER_SECOND: usize = 512 * 1024 * 1024; +/// Maximum number of rows per second allowed to be emitted. Reads greater than this will be +/// truncated to this max value, as it's assumed that [READY_CHUNK_LIMIT] was hit and no useful +/// measurement can actually be emitted. +const MAX_ROWS_PER_SECOND: usize = 1024 * 1024; +/// Maximum number of rows sampled from the peek queue when estimating per-column NDV. +const NDV_MAX_ROWS_SAMPLE: usize = 1000; + +#[derive(Debug)] +pub struct SamplerExec { + pub(crate) input: Arc, + pub(crate) metric_set: ExecutionPlanMetricsSet, + pub(crate) partition_samplers: Vec, + pub(crate) execution_started: Arc, +} + +/// Metrics that quantify how long the sampler held data in memory before the consumer +/// (real execution) attached, plus the peak accumulated size reached. All metrics are shared +/// across the partition samplers; the latency metrics aggregate per-partition observations. +#[derive(Debug, Clone)] +pub(crate) struct SamplerExecMetrics { + /// Time since [SamplerExec::kick_off_first_sampler] was called until the first batch from + /// the input arrived + kick_off_to_fist_batch_p50: P50LatencyMetric, + kick_off_to_fist_batch_max: MaxLatencyMetric, + /// Time since [SamplerExec::kick_off_first_sampler] was called until the [pb::LoadInfo] message + /// was sent. + kick_off_to_load_info_sent_p50: P50LatencyMetric, + kick_off_to_load_info_sent_max: MaxLatencyMetric, + /// Time since [SamplerExec::kick_off_first_sampler] was called until the node was properly + /// executed with [SamplerExec::execute]. + kick_off_to_execution_p50: P50LatencyMetric, + kick_off_to_execution_max: MaxLatencyMetric, + /// Maximum number of record batches peeked by a sampler. + max_batches_peeked: MaxGaugeMetric, + /// Peak memory accumulated by any partition sampler during the sampling phase. + max_mem_used: Gauge, + /// Bytes per second flowing through the sampler node. + bytes_per_sec: BytesCounterMetric, + /// Bytes ready at the moment of reporting load info. + bytes_ready: BytesCounterMetric, + /// Elapsed compute while sampling. + elapsed_compute: Time, +} + +impl SamplerExecMetrics { + fn new(metric_set: &ExecutionPlanMetricsSet) -> Self { + let bdr = || MetricBuilder::new(metric_set); + Self { + kick_off_to_fist_batch_p50: bdr().p50_latency("kick_off_to_first_batch_p50"), + kick_off_to_fist_batch_max: bdr().max_latency("kick_off_to_first_batch_max"), + kick_off_to_load_info_sent_p50: bdr().p50_latency("kick_off_to_load_info_sent_p50"), + kick_off_to_load_info_sent_max: bdr().max_latency("kick_off_to_load_info_sent_max"), + kick_off_to_execution_p50: bdr().p50_latency("kick_off_to_execution_p50"), + kick_off_to_execution_max: bdr().max_latency("kick_off_to_execution_max"), + max_batches_peeked: bdr().max_gauge("max_batches_peeked"), + max_mem_used: bdr().global_gauge("max_mem_used"), + bytes_per_sec: bdr().bytes_counter("bytes_per_sec"), + bytes_ready: bdr().bytes_counter("bytes_ready"), + elapsed_compute: { + let time = Time::new(); + bdr().build(MetricValue::ElapsedCompute(time.clone())); + time + }, + } + } +} + +impl SamplerExec { + pub(crate) fn new(input: Arc) -> Self { + let metric_set = ExecutionPlanMetricsSet::new(); + let metric_set_clone = metric_set.clone(); + // Metrics need to be lazily initialized, otherwise the coordinator side will register + // them when they are never relevant there, they are just relevant in workers. + // + // If we don't do this, the [SamplerExec] constructed during planning will register its + // own zeroed SamplerExecMetrics in the ExecutionPlanMetricsSet, even if the metrics we care + // about are just the ones collected in workers during execution. + let metrics: Arc SamplerExecMetrics + Send>>> = + Arc::new(LazyLock::new(Box::new(move || { + SamplerExecMetrics::new(&metric_set_clone) + }))); + let partitions = input.properties().partitioning.partition_count(); + let execution_started = Arc::new(AtomicBool::new(false)); + let mut samplers = Vec::with_capacity(partitions); + for i in 0..partitions { + samplers.push(PartitionSampler { + partition_idx: i, + input: Arc::clone(&input), + stream: Mutex::new(None), + metrics: Arc::clone(&metrics), + kick_off_at: Arc::new(OnceLock::new()), + first_batch_at: Arc::new(OnceLock::new()), + load_info_sent_at: Arc::new(OnceLock::new()), + execution_started: Arc::clone(&execution_started), + }); + } + Self { + input, + metric_set, + partition_samplers: samplers, + execution_started, + } + } + + pub(crate) fn kick_off_first_sampler( + plan: Arc, + ctx: Arc, + ) -> Result>> { + let mut receivers = vec![]; + plan.apply(|plan| { + let Some(sampler) = plan.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + receivers.reserve(sampler.partition_samplers.len()); + for partition_sampler in &sampler.partition_samplers { + let rx = partition_sampler.kick_off(Arc::clone(&ctx))?; + receivers.push(rx); + } + Ok(TreeNodeRecursion::Stop) + })?; + Ok(receivers) + } +} + +pub(crate) struct PartitionSampler { + partition_idx: usize, + input: Arc, + stream: Mutex>, + execution_started: Arc, + + // Metrics state. + metrics: Arc SamplerExecMetrics + Send>>>, + /// Set when `kick_off` is invoked. Used at `execute()` time to record how long the + /// sampler sampled data before the consumer attached. + kick_off_at: Arc>, + /// Set the first time the producer task emits a `LoadInfo`. Used at `execute()` time + /// to record the gap between the first sample and the consumer starting. + first_batch_at: Arc>, + /// Set immediately after `sampling_tx.send()` succeeds. Used to measure the full + /// round-trip: LoadInfo sent → coordinator collects votes → downstream plan dispatched + /// → consumer calls execute(). + load_info_sent_at: Arc>, +} + +impl Debug for PartitionSampler { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PartitionSampler").finish() + } +} + +impl PartitionSampler { + fn start_stream(&self) -> Option { + let Some(kick_off_at) = self.kick_off_at.get() else { + return self.stream.lock().unwrap().take(); + }; + + // Time since this sampler was kicked off until the first batch arrived. + if let Some(t) = self.first_batch_at.get() { + let delay = t.saturating_duration_since(*kick_off_at); + self.metrics.kick_off_to_fist_batch_p50.add_duration(delay); + self.metrics.kick_off_to_fist_batch_max.add_duration(delay); + } + + // Time since the sampler was kicked off until the pb::LoadInfo message was sent. + if let Some(t) = self.load_info_sent_at.get() { + let delay = t.saturating_duration_since(*kick_off_at); + self.metrics + .kick_off_to_load_info_sent_p50 + .add_duration(delay); + self.metrics + .kick_off_to_load_info_sent_max + .add_duration(delay); + } + + // Time since the sampler was kicked off until it started executing. + let delay = kick_off_at.elapsed(); + self.metrics.kick_off_to_execution_p50.add_duration(delay); + self.metrics.kick_off_to_execution_max.add_duration(delay); + + self.stream.lock().unwrap().take() + } + + fn kick_off(&self, ctx: Arc) -> Result> { + let _ = self.kick_off_at.set(Instant::now()); + let (sampling_tx, sampling_rx) = oneshot::channel(); + + let input = Arc::clone(&self.input); + let partition_idx = self.partition_idx; + let schema = input.schema(); + let elapsed_compute = self.metrics.elapsed_compute.clone(); + let first_batch_at = Arc::clone(&self.first_batch_at); + let n_cols = self.input.schema().fields.len(); + + let reporter = LoadInfoDropHandler { + load_info: zero_load_info(partition_idx, n_cols), + sampling_tx: Some(sampling_tx), + bytes_per_second_metric: self.metrics.bytes_per_sec.clone(), + load_info_sent_at: Arc::clone(&self.load_info_sent_at), + bytes_ready_metric: self.metrics.bytes_ready.clone(), + omit: Arc::clone(&self.execution_started), + }; + + let mut peek = RecordBatchPeek { + peek: VecDeque::new(), + n_cols, + max_mem_used: self.metrics.max_mem_used.clone(), + max_batches_peeked: self.metrics.max_batches_peeked.clone(), + memory_reservation: Arc::new( + MemoryConsumer::new(format!("PartitionSampler[{partition_idx}]")) + .register(ctx.memory_pool()), + ), + first_batch_at: Arc::clone(&self.first_batch_at), + }; + + // Execute the input synchronously so any setup error surfaces before we + // spawn the producer task. + let mut input_stream = input.execute(partition_idx, ctx)?.fuse(); + + let task = SpawnedTask::spawn(async move { + // First, read at once all the RecordBatches that are ready to be yielded synchronously. + // Some downstream nodes will accumulate data in-memory, and will then yield several + // RecordBatches at once synchronously (without Poll::Pending gaps in between). + let mut chunked = (&mut input_stream).ready_chunks(READY_CHUNK_LIMIT); + let Some(batches) = chunked.next().await else { + // Not a single RecordBatch was produced, so let bytes_per_second=0 be sent as-is. + return Ok(peek.chain(input_stream).boxed()); + }; + let _elapsed_compute_timer = elapsed_compute.timer(); + for batch in batches { + let _ = first_batch_at.set(Instant::now()); + peek.push(batch?); + } + + // Peek whether there is more data to be produced. + if let Some(result) = input_stream.next().now_or_never() { + return if let Some(batch) = result { + // A batch was immediately available without hitting an async gap (the input is + // still yielding synchronously). store it so its rows are not lost. We cannot + // measure a meaningful arrival velocity in this case, so as before, assume the + // worst. + peek.push(batch?); + reporter.report(&peek, MAX_BYTES_PER_SECOND, MAX_ROWS_PER_SECOND); + Ok(peek.chain(input_stream).boxed()) + } else { + // No more batches to read, so no velocity measurement. + reporter.report(&peek, 0, 0); + Ok(peek.chain(input_stream).boxed()) + }; + } + + drop(_elapsed_compute_timer); + + // Wait for an async gap in order to measure data velocity. + let poll_start = Instant::now(); + let Some(batch) = input_stream.try_next().await? else { + let _elapsed_compute_timer = elapsed_compute.timer(); + // The last message was somehow the last message in the stream, but the stream did + // not end immediately. This is an unlikely scenario. + reporter.report(&peek, 0, 0); + return Ok(peek.chain(input_stream).boxed()); + }; + let _elapsed_compute_timer = elapsed_compute.timer(); + + let bytes_per_second = + (record_batch_size(&batch) as f32 / poll_start.elapsed().as_secs_f32()) as usize; + let rows_per_second = + (batch.num_rows() as f32 / poll_start.elapsed().as_secs_f32()) as usize; + + peek.push(batch); + + // Some RecordBatches where buffered, but there's more to be yielded, so both + // bytes_per_second and bytes_ready can be reported. + reporter.report(&peek, bytes_per_second, rows_per_second); + + Ok(peek.chain(input_stream).boxed()) + }); + + let stream = async move { + task.await + .map_err(|err| DataFusionError::Internal(err.to_string()))? + } + .try_flatten_stream(); + + self.stream + .lock() + .expect("poisoned lock") + .replace(Box::pin(RecordBatchStreamAdapter::new(schema, stream))); + + Ok(sampling_rx) + } +} + +/// Wraps a [pb::LoadInfo] and emits it on [Drop] through the provided [oneshot] channel. +/// +/// Emitting on drop ensures that it's always emitted. +struct LoadInfoDropHandler { + omit: Arc, + + load_info: pb::LoadInfo, + bytes_ready_metric: BytesCounterMetric, + bytes_per_second_metric: BytesCounterMetric, + sampling_tx: Option>, + load_info_sent_at: Arc>, +} + +impl LoadInfoDropHandler { + fn report(mut self, peek: &RecordBatchPeek, bps: usize, rps: usize) { + if self.omit.load(Ordering::Relaxed) { + return; + } + + self.set_per_col_bytes_ready(peek.per_col_bytes_ready()); + self.set_rows_ready(peek.rows_ready()); + self.set_per_col_ndv(peek.per_col_ndv()); + self.set_per_col_null_pct(peek.per_col_null_pct()); + self.set_per_col_bytes_per_second(bps); + self.set_rows_per_second(rps) + } + + fn set_per_col_bytes_ready(&mut self, bytes_ready: Vec) { + self.load_info.per_column_bytes_ready = vec_cast(&bytes_ready); + self.bytes_ready_metric.add_bytes(bytes_ready.iter().sum()); + } + + fn set_per_col_bytes_per_second(&mut self, total_bytes_per_second: usize) { + let per_col_ready: &[u64] = &self.load_info.per_column_bytes_ready; + let total_ready: u64 = per_col_ready.iter().sum(); + let per_col: Vec = if total_ready == 0 { + vec![total_bytes_per_second / per_col_ready.len().max(1); per_col_ready.len()] + } else { + per_col_ready + .iter() + .map(|&ready| { + (ready.saturating_mul(total_bytes_per_second as u64) / total_ready) as usize + }) + .collect() + }; + self.load_info.per_column_bytes_per_second = vec_cast(&per_col); + self.bytes_per_second_metric + .add_bytes(total_bytes_per_second); + } + + fn set_rows_ready(&mut self, rows_ready: usize) { + self.load_info.rows_ready = rows_ready as u64; + } + + fn set_rows_per_second(&mut self, rows_per_second: usize) { + self.load_info.rows_per_second = rows_per_second as u64; + } + + fn set_per_col_ndv(&mut self, per_column_ndv: Vec) { + self.load_info.per_column_ndv_percentage = per_column_ndv; + } + + fn set_per_col_null_pct(&mut self, per_column_null_pct: Vec) { + self.load_info.per_column_null_percentage = per_column_null_pct; + } +} + +impl Drop for LoadInfoDropHandler { + fn drop(&mut self) { + if self.omit.load(Ordering::Relaxed) { + return; + } + if let Some(sampling_tx) = self.sampling_tx.take() { + let _ = sampling_tx.send(std::mem::take(&mut self.load_info)); + let _ = self.load_info_sent_at.set(Instant::now()); + } + } +} + +fn zero_load_info(partition_idx: usize, n_cols: usize) -> pb::LoadInfo { + pb::LoadInfo { + partition: partition_idx as u64, + rows_per_second: 0, + rows_ready: 0, + per_column_bytes_per_second: vec![0; n_cols], + per_column_bytes_ready: vec![0; n_cols], + per_column_ndv_percentage: vec![0.0; n_cols], + per_column_null_percentage: vec![0.0; n_cols], + } +} + +struct RecordBatchPeek { + peek: VecDeque, + n_cols: usize, + max_batches_peeked: MaxGaugeMetric, + max_mem_used: Gauge, + memory_reservation: Arc, + first_batch_at: Arc>, +} + +impl RecordBatchPeek { + fn push(&mut self, batch: RecordBatch) { + let batch_size = record_batch_size(&batch); + if self.peek.is_empty() { + let _ = self.first_batch_at.set(Instant::now()); + } + self.max_mem_used.add(batch_size); + self.memory_reservation.grow(batch_size); + self.peek.push_back(batch); + self.max_batches_peeked.set_max(self.peek.len()); + } + + fn per_col_bytes_ready(&self) -> Vec { + let mut result = vec![0; self.n_cols]; + for batch in self.peek.iter() { + for (i, col) in batch.columns().iter().enumerate() { + result[i] += column_size(col) + } + } + result + } + + fn rows_ready(&self) -> usize { + self.peek.iter().map(|batch| batch.num_rows()).sum() + } + + fn per_col_ndv(&self) -> Vec { + let total_rows: usize = self.peek.iter().map(|b| b.num_rows()).sum(); + if total_rows == 0 { + return vec![0.0; self.n_cols]; + } + + // Build the list of flat row indices to sample, sorted for cache-friendly access. + let sampled_indices: Vec = if total_rows <= NDV_MAX_ROWS_SAMPLE { + (0..total_rows).collect() + } else { + let mut indices = + rand::seq::index::sample(&mut rand::rng(), total_rows, NDV_MAX_ROWS_SAMPLE) + .into_vec(); + indices.sort_unstable(); + indices + }; + let rows_sampled = sampled_indices.len(); + + let mut sets: Vec> = vec![HashSet::new(); self.n_cols]; + let mut flat_base = 0usize; + let mut sample_pos = 0usize; + + for batch in &self.peek { + let batch_end = flat_base + batch.num_rows(); + while sample_pos < sampled_indices.len() && sampled_indices[sample_pos] < batch_end { + let row = sampled_indices[sample_pos] - flat_base; + for (col_idx, set) in sets.iter_mut().enumerate() { + let col = batch.column(col_idx); + if !col.is_null(row) + && let Ok(v) = ScalarValue::try_from_array(col, row) + { + set.insert(v); + } + } + sample_pos += 1; + } + if sample_pos >= sampled_indices.len() { + break; + } + flat_base = batch_end; + } + + sets.into_iter() + .map(|s| s.len() as f32 / rows_sampled as f32) + .collect() + } + + fn per_col_null_pct(&self) -> Vec { + let total_rows: usize = self.peek.iter().map(|b| b.num_rows()).sum(); + if total_rows == 0 { + return vec![0.0; self.n_cols]; + } + let mut counts = vec![0usize; self.n_cols]; + for batch in &self.peek { + for (col_idx, count) in counts.iter_mut().enumerate() { + *count += batch.column(col_idx).null_count(); + } + } + counts + .into_iter() + .map(|c| c as f32 / total_rows as f32) + .collect() + } +} + +impl Stream for RecordBatchPeek { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + match self.as_mut().peek.pop_front() { + None => Poll::Ready(None), + Some(batch) => { + self.memory_reservation.shrink(record_batch_size(&batch)); + Poll::Ready(Some(Ok(batch))) + } + } + } + + fn size_hint(&self) -> (usize, Option) { + (self.peek.len(), Some(self.peek.len())) + } +} + +fn column_size(arr: &ArrayRef) -> usize { + arr.to_data().get_slice_memory_size().unwrap_or(0) +} + +fn record_batch_size(batch: &RecordBatch) -> usize { + batch.columns().iter().map(column_size).sum() +} + +impl DisplayAs for SamplerExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "SamplerExec: partitions={}", + self.partition_samplers.len() + ) + } +} + +impl ExecutionPlan for SamplerExec { + fn name(&self) -> &str { + "SamplerExec" + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new(require_one_child(children)?))) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + self.execution_started.store(true, Ordering::Relaxed); + let Some(partition_sampler) = self.partition_samplers.get(partition) else { + return exec_err!("Partition {partition} not available in SamplerExec"); + }; + let Some(stream) = partition_sampler.start_stream() else { + return exec_err!("SamplerExec[{partition}] was not kicked off"); + }; + Ok(stream) + } + + fn metrics(&self) -> Option { + Some(self.metric_set.clone_inner()) + } +} diff --git a/src/metrics/bytes_metric.rs b/src/metrics/bytes_metric.rs index 3bdd7f5a..5d8c7e6c 100644 --- a/src/metrics/bytes_metric.rs +++ b/src/metrics/bytes_metric.rs @@ -5,6 +5,7 @@ use std::{ sync::{Arc, atomic::AtomicUsize}, }; +use datafusion::physical_plan::Metric; use datafusion::{ common::human_readable_size, physical_plan::metrics::{CustomMetricValue, MetricBuilder, MetricValue}, @@ -50,6 +51,16 @@ impl Default for BytesCounterMetric { } impl BytesCounterMetric { + pub fn new_metric(name: impl Into>, bytes: usize) -> Arc { + Arc::new(Metric::new( + MetricValue::Custom { + name: name.into(), + value: Arc::new(BytesCounterMetric::from_value(bytes)), + }, + None, + )) + } + pub fn from_value(bytes: usize) -> Self { Self { bytes: Arc::new(AtomicUsize::new(bytes)), diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 3ba5c88c..87d93c96 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -81,6 +81,7 @@ pub async fn rewrite_distributed_plan_with_metrics( num: stage.num, plan: plan_with_metrics, tasks: stage.tasks, + metrics_set: stage.metrics_set.clone(), }))?; let network_boundary = MetricsWrapperExec::new(network_boundary, plan.metrics().unwrap_or_default()); @@ -415,6 +416,7 @@ mod tests { num: 2, plan, tasks: 4, + metrics_set: Default::default(), } } diff --git a/src/protobuf/distributed_codec.rs b/src/protobuf/distributed_codec.rs index 3cfec30a..1ea27541 100644 --- a/src/protobuf/distributed_codec.rs +++ b/src/protobuf/distributed_codec.rs @@ -1,9 +1,9 @@ use super::get_distributed_user_codecs; use crate::NetworkShuffleExec; -use crate::common::{deserialize_uuid, serialize_uuid}; +use crate::common::{deserialize_uuid, require_one_child, serialize_uuid}; use crate::execution_plans::{ BroadcastExec, ChildWeight, ChildrenIsolatorUnionExec, NetworkBroadcastExec, - NetworkCoalesceExec, + NetworkCoalesceExec, SamplerExec, }; use crate::stage::{LocalStage, RemoteStage, Stage}; use crate::worker::WorkerConnectionPool; @@ -74,6 +74,7 @@ impl PhysicalExtensionCodec for DistributedCodec { num: proto.num as usize, plan: input, tasks: proto.tasks.len(), + metrics_set: Default::default(), })) } else { let mut worker_urls = Vec::with_capacity(proto.tasks.len()); @@ -90,6 +91,7 @@ impl PhysicalExtensionCodec for DistributedCodec { query_id: deserialize_uuid(proto.query_id.as_ref())?, num: proto.num as usize, workers: worker_urls, + runtime_stats: None, })) } } @@ -233,6 +235,9 @@ impl PhysicalExtensionCodec for DistributedCodec { .collect(), })) } + DistributedExecNode::Sampler(SamplerExecProto {}) => { + Ok(Arc::new(SamplerExec::new(require_one_child(inputs)?))) + } } } @@ -349,6 +354,14 @@ impl PhysicalExtensionCodec for DistributedCodec { node: Some(DistributedExecNode::ChildrenIsolatorUnion(inner)), }; + wrapper.encode(buf).map_err(|e| proto_error(format!("{e}"))) + } else if let Some(_node) = node.downcast_ref::() { + let inner = SamplerExecProto {}; + + let wrapper = DistributedExecProto { + node: Some(DistributedExecNode::Sampler(inner)), + }; + wrapper.encode(buf).map_err(|e| proto_error(format!("{e}"))) } else { Err(proto_error(format!("Unexpected plan {}", node.name()))) @@ -380,7 +393,7 @@ pub struct ExecutionTaskProto { #[derive(Clone, PartialEq, ::prost::Message)] pub struct DistributedExecProto { - #[prost(oneof = "DistributedExecNode", tags = "1, 2, 3, 4, 5, 6")] + #[prost(oneof = "DistributedExecNode", tags = "1, 2, 3, 4, 5, 6, 7")] pub node: Option, } @@ -397,6 +410,8 @@ pub enum DistributedExecNode { NetworkBroadcast(NetworkBroadcastExecProto), #[prost(message, tag = "6")] Broadcast(BroadcastExecProto), + #[prost(message, tag = "7")] + Sampler(SamplerExecProto), } /// Protobuf representation of the [NetworkShuffleExec] physical node. It serves as @@ -509,6 +524,9 @@ pub struct BroadcastExecProto { pub consumer_task_count: u64, } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SamplerExecProto {} + fn new_network_broadcast_exec( partitioning: Partitioning, schema: SchemaRef, @@ -547,6 +565,7 @@ mod tests { query_id: Default::default(), num: 0, workers: vec![], + runtime_stats: None, }) } @@ -556,6 +575,7 @@ mod tests { num: 0, plan: empty_exec(), tasks: 1, + metrics_set: Default::default(), }) } diff --git a/src/stage.rs b/src/stage.rs index 0545d718..3ac39e7a 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -1,13 +1,15 @@ use crate::coordinator::{DistributedExec, MetricsStore}; use crate::execution_plans::{DistributedLeafExec, NetworkCoalesceExec}; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; -use datafusion::common::{HashMap, config_err}; +use datafusion::common::{HashMap, Statistics, config_err}; use datafusion::common::{exec_err, plan_err}; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::metrics::{Label, Metric, MetricsSet}; -use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; +use datafusion::physical_plan::{ + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, displayable, +}; use itertools::Either; use std::collections::VecDeque; use std::sync::Arc; @@ -84,6 +86,8 @@ pub struct LocalStage { pub plan: Arc, /// The number of tasks the stage has. pub tasks: usize, + /// Metrics collected by the coordinator + pub metrics_set: MetricsSet, } impl LocalStage { @@ -107,6 +111,8 @@ pub struct RemoteStage { pub num: usize, /// The worker URLs to which queries should be issued. pub workers: Vec, + /// Statistics collected at runtime, if any. + pub runtime_stats: Option>, } impl Stage { @@ -137,6 +143,63 @@ impl Stage { Self::Remote(_) => None, } } + + pub fn metrics(&self) -> MetricsSet { + match &self { + Self::Local(v) => v.metrics_set.clone(), + Self::Remote(_) => MetricsSet::new(), + } + } + + pub fn partition_statistics( + &self, + partition: Option, + partition_count: usize, + schema: SchemaRef, + ) -> Result> { + match self { + Stage::Local(local) => local.plan.partition_statistics(partition), + Stage::Remote(remote) => { + let Some(runtime_stats) = &remote.runtime_stats else { + return Ok(Arc::new(Statistics::new_unknown(&schema))); + }; + match partition { + None => Ok(Arc::clone(runtime_stats)), + Some(_) => Ok(Arc::new(multiply_stats( + runtime_stats, + 1.0 / partition_count as f32, + ))), + } + } + } + } +} + +fn multiply_stats(stats: &Statistics, f: f32) -> Statistics { + Statistics { + num_rows: multiply_precision(stats.num_rows, f), + total_byte_size: multiply_precision(stats.total_byte_size, f), + column_statistics: stats + .column_statistics + .iter() + .map(|col| ColumnStatistics { + null_count: multiply_precision(col.null_count, f), + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + distinct_count: multiply_precision(col.distinct_count, f), + byte_size: multiply_precision(col.byte_size, f), + }) + .collect(), + } +} + +fn multiply_precision(p: Precision, f: f32) -> Precision { + match p { + Precision::Exact(v) => Precision::Exact((v as f32 * f) as usize), + Precision::Inexact(v) => Precision::Inexact((v as f32 * f) as usize), + Precision::Absent => Precision::Absent, + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -161,7 +224,9 @@ use crate::metrics::proto::metric_proto_to_df; use crate::worker::generated::worker as pb; use crate::{DistributedMetricsFormat, NetworkShuffleExec, rewrite_distributed_plan_with_metrics}; use crate::{NetworkBoundary, NetworkBoundaryExt}; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::DataFusionError; +use datafusion::common::stats::Precision; use datafusion::physical_expr::Partitioning; /// Be able to display a nice tree for stages. /// @@ -373,7 +438,7 @@ fn gather_stage_header_metrics(stage: &Stage, metrics_store: &MetricsStore) -> M stage_id: stage.num() as u64, task_number: 0, }; - let mut all_metrics = MetricsSet::new(); + let mut all_metrics = stage.metrics(); while let Some(metrics_set) = metrics_store.get(&task_key).and_then(|v| v.task_metrics) { for mut metric in metrics_set.metrics { metric.labels.push(pb::Label { @@ -573,6 +638,7 @@ pub fn display_plan_graphviz(plan: Arc) -> Result { num: max_num + 1, plan: plan.clone(), tasks: 1, + metrics_set: MetricsSet::new(), }); all_stages.insert(0, &head_stage); diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index f914f228..1526508c 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -38,8 +38,18 @@ pub(crate) struct RemoteWorkUnitFeedRegistry { impl RemoteWorkUnitFeedRegistry { /// Creates all the receivers and senders for a specific [WorkUnit] Feed id. One feed per /// partition is created. + /// + /// Calling this twice with the same `id` is a coordinator bug — duplicate declarations + /// mean two plan nodes share a UUID, which would cause "already consumed" when both + /// nodes call `feed()`. We skip rather than overwrite so the coordinator-side duplicate + /// detection in `task_specialized_plan` surfaces the real error first. pub(crate) fn add(&mut self, id: Uuid, partitions: usize) { for partition in 0..partitions { + // Skip if already registered; overwriting would silently drop the existing + // receiver and cause a confusing "already consumed" error at execution time. + if self.receivers.contains_key(&(id, partition)) { + continue; + } let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); self.receivers.insert((id, partition), Mutex::new(Some(rx))); self.senders.insert((id, partition), tx); diff --git a/src/worker/generated/worker.rs b/src/worker/generated/worker.rs index fe7a0137..290261ed 100644 --- a/src/worker/generated/worker.rs +++ b/src/worker/generated/worker.rs @@ -24,7 +24,7 @@ pub mod coordinator_to_worker_msg { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct WorkerToCoordinatorMsg { - #[prost(oneof = "worker_to_coordinator_msg::Inner", tags = "1")] + #[prost(oneof = "worker_to_coordinator_msg::Inner", tags = "1, 2, 3")] pub inner: ::core::option::Option, } /// Nested message and enum types in `WorkerToCoordinatorMsg`. @@ -37,6 +37,12 @@ pub mod worker_to_coordinator_msg { /// metrics\[i\] is the set of metrics for plan node i in pre-order traversal order. #[prost(message, tag = "1")] TaskMetrics(super::TaskMetrics), + /// Load information reported by a task. This information is used for dynamically + /// sizing the number of workers involved in a query. + #[prost(message, tag = "2")] + LoadInfo(super::LoadInfo), + #[prost(bool, tag = "3")] + LoadInfoEos(bool), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -52,6 +58,34 @@ pub struct TaskMetrics { #[prost(message, optional, tag = "2")] pub task_metrics: ::core::option::Option, } +/// Load information reported for a specific partition with information about this +/// amount of data flowing through the plan. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct LoadInfo { + /// The partition index to which this message belongs to. + #[prost(uint64, tag = "1")] + pub partition: u64, + /// The amount of rows ready to be returned. + #[prost(uint64, tag = "2")] + pub rows_ready: u64, + /// The estimated velocity at which rows will flow through the node. If all the rows were + /// already accumulated, they will be reported by `rows_ready`, and this field will be 0. + #[prost(uint64, tag = "3")] + pub rows_per_second: u64, + /// The amount of bytes ready to be returned per column. + #[prost(uint64, repeated, tag = "4")] + pub per_column_bytes_ready: ::prost::alloc::vec::Vec, + /// The estimated velocity at which data will flow through each column. If all the bytes were + /// already accumulated, they will be reported by `bytes_ready`, and this field will be 0. + #[prost(uint64, repeated, tag = "5")] + pub per_column_bytes_per_second: ::prost::alloc::vec::Vec, + /// Approximate ratio of NDV for each column. + #[prost(float, repeated, tag = "6")] + pub per_column_ndv_percentage: ::prost::alloc::vec::Vec, + /// Approximate ratio of null count for each column. + #[prost(float, repeated, tag = "7")] + pub per_column_null_percentage: ::prost::alloc::vec::Vec, +} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct GetWorkerInfoRequest {} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index d1c2097f..fcb3944a 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,4 +1,5 @@ use crate::common::deserialize_uuid; +use crate::execution_plans::SamplerExec; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; @@ -17,6 +18,7 @@ use datafusion::execution::SessionStateBuilder; use datafusion::prelude::SessionConfig; use datafusion_proto::physical_plan::AsExecutionPlan; use datafusion_proto::protobuf::PhysicalPlanNode; +use futures::stream::FuturesUnordered; use futures::{FutureExt, StreamExt, TryStreamExt}; use std::sync::atomic::AtomicUsize; use std::sync::{Arc, OnceLock}; @@ -55,6 +57,7 @@ impl Worker { } let (metrics_tx, metrics_rx) = oneshot::channel(); + let mut load_info_rxs = vec![]; let task_data = || async { let headers = grpc_headers.into_headers(); @@ -98,6 +101,8 @@ impl Worker { for hook in self.hooks.on_plan.iter() { plan = hook(plan, session_state.config())?; } + load_info_rxs = + SamplerExec::kick_off_first_sampler(Arc::clone(&plan), Arc::clone(&task_ctx))?; // Initialize partition count to the number of partitions in the stage let total_partitions = plan.properties().partitioning.partition_count(); @@ -172,19 +177,34 @@ impl Worker { tokio::spawn(async move { task_data_entries.invalidate(&key).await }); }); + let load_info_stream = FuturesUnordered::from_iter(load_info_rxs) + .filter_map(async |load_info_or_channel_dropped| { + // This error can only happen if the pb::LoadInfo sender was dropped, which is fine. + let load_info = load_info_or_channel_dropped.ok()?; + Some(Ok(WorkerToCoordinatorMsg { + inner: Some(worker_to_coordinator_msg::Inner::LoadInfo(load_info)), + })) + }) + .chain(futures::stream::once(async move { + Ok(WorkerToCoordinatorMsg { + inner: Some(worker_to_coordinator_msg::Inner::LoadInfoEos(true)), + }) + })); + // Stream back the metrics once the task finishes executing. // The oneshot receiver resolves when impl_execute_task sends the collected // metrics after all partitions have finished or been dropped. let metrics_stream = metrics_rx.into_stream(); - let metrics_stream = metrics_stream.filter_map(|task_metrics| async move { - match task_metrics { - Ok(task_metrics) => Some(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics)), - }), - Err(_) => None, // channel dropped without sending any message - } + let metrics_stream = metrics_stream.filter_map(async |task_metrics_or_channel_dropped| { + let task_metrics = task_metrics_or_channel_dropped.ok()?; + Some(Ok(WorkerToCoordinatorMsg { + inner: Some(worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics)), + })) }); - Ok(Response::new(metrics_stream.map(Ok).boxed())) + + Ok(Response::new( + futures::stream::select(load_info_stream, metrics_stream).boxed(), + )) } } diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index 28f5ca5d..97b2e806 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -1,6 +1,7 @@ use crate::MaxLatencyMetric; -use crate::common::{OnceLockResult, now_ns}; -use crate::distributed_planner::{ProducerHead, insert_producer_head}; +use crate::common::OnceLockResult; +use crate::common::now_ns; +use crate::distributed_planner::ProducerHead; use crate::worker::generated::worker as pb; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; @@ -134,7 +135,7 @@ impl TaskData { let producer_head = ProducerHead::from_proto(producer_head, &self.base_plan.schema(), &self.task_ctx)?; - let plan = insert_producer_head(Arc::clone(&self.base_plan), producer_head)?; + let plan = producer_head.insert(Arc::clone(&self.base_plan))?; self.num_partitions_remaining.store( plan.output_partitioning().partition_count(), diff --git a/src/worker/worker.proto b/src/worker/worker.proto index bc1e3412..ef691b59 100644 --- a/src/worker/worker.proto +++ b/src/worker/worker.proto @@ -33,6 +33,12 @@ message WorkerToCoordinatorMsg { // ensuring metrics are never lost due to early stream termination. // metrics[i] is the set of metrics for plan node i in pre-order traversal order. TaskMetrics task_metrics = 1; + + // Load information reported by a task. This information is used for dynamically + // sizing the number of workers involved in a query. + LoadInfo load_info = 2; + + bool load_info_eos = 3; } } @@ -47,6 +53,27 @@ message TaskMetrics { MetricsSet task_metrics = 2; } +// Load information reported for a specific partition with information about this +// amount of data flowing through the plan. +message LoadInfo { + // The partition index to which this message belongs to. + uint64 partition = 1; + // The amount of rows ready to be returned. + uint64 rows_ready = 2; + // The estimated velocity at which rows will flow through the node. If all the rows were + // already accumulated, they will be reported by `rows_ready`, and this field will be 0. + uint64 rows_per_second = 3; + // The amount of bytes ready to be returned per column. + repeated uint64 per_column_bytes_ready = 4; + // The estimated velocity at which data will flow through each column. If all the bytes were + // already accumulated, they will be reported by `bytes_ready`, and this field will be 0. + repeated uint64 per_column_bytes_per_second = 5; + // Approximate ratio of NDV for each column. + repeated float per_column_ndv_percentage = 6; + // Approximate ratio of null count for each column. + repeated float per_column_null_percentage = 7; +} + message GetWorkerInfoRequest {} message GetWorkerInfoResponse { diff --git a/tests/clickbench_correctness_test.rs b/tests/clickbench_correctness_test.rs index 9df2a8d3..acf3daa8 100644 --- a/tests/clickbench_correctness_test.rs +++ b/tests/clickbench_correctness_test.rs @@ -18,6 +18,7 @@ mod tests { use std::sync::Arc; use tokio::sync::OnceCell; + const ADAPTIVE_ENV_VAR: &str = "ADAPTIVE"; const NUM_WORKERS: usize = 4; const PARTITIONS: usize = 3; const FILE_SCAN_CONFIG_BYTES_PER_PARTITION: usize = 1; @@ -289,12 +290,15 @@ mod tests { .options_mut() .execution .target_partitions = PARTITIONS; - let d_ctx = d_ctx + let mut d_ctx = d_ctx .with_distributed_file_scan_config_bytes_per_partition( FILE_SCAN_CONFIG_BYTES_PER_PARTITION, )? .with_distributed_cardinality_effect_task_scale_factor(CARDINALITY_TASK_COUNT_FACTOR)? .with_distributed_broadcast_joins(true)?; + if std::env::var(ADAPTIVE_ENV_VAR).unwrap_or_default() == "true" { + d_ctx.set_distributed_dynamic_task_count(true)?; + } register_tables(&s_ctx, &data_dir).await?; register_tables(&d_ctx, &data_dir).await?; diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index ccb43a57..40d05532 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -6,6 +6,7 @@ mod tests { use datafusion::common::{Result, assert_contains}; use datafusion::execution::SessionState; use datafusion::physical_plan::display::DisplayableExecutionPlan; + use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ExecutionPlan, execute_stream}; use datafusion::prelude::SessionContext; use datafusion_distributed::test_utils::localhost::start_localhost_context; @@ -341,6 +342,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_metrics_collection_dynamic() -> Result<(), Box> { + let (mut d_ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; + d_ctx.set_distributed_dynamic_task_count(true)?; + + let query = + r#"SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)"#; + + let s_ctx = SessionContext::default(); + let (s_physical, mut d_physical) = execute(&s_ctx, &d_ctx, query).await?; + d_physical = rewrite_with_metrics(d_physical, DistributedMetricsFormat::Aggregated).await; + println!("{}", display_plan_ascii(s_physical.as_ref(), true)); + println!("{}", display_plan_ascii(d_physical.as_ref(), true)); + + assert_metrics_equal::( + ["output_rows", "output_bytes"], + &s_physical, + &d_physical, + 0, + ); + + assert_metrics_equal::( + ["output_rows", "output_bytes"], + &s_physical, + &d_physical, + 0, + ); + + Ok(()) + } + /// Looks for an [ExecutionPlan] that matches the provided type parameter `T1` in /// the left node and `T2` in the right node and compares its metrics. /// There might be more than one, so `index` determines which one is compared. diff --git a/tests/stateful_data_cleanup.rs b/tests/stateful_data_cleanup.rs index a3fe7bca..fae892d2 100644 --- a/tests/stateful_data_cleanup.rs +++ b/tests/stateful_data_cleanup.rs @@ -20,13 +20,18 @@ mod tests { const TPCH_DATA_PARTS: usize = 16; const CARDINALITY_TASK_COUNT_FACTOR: f64 = 1.0; - #[test_case(false; "metrics_disabled")] - #[test_case(true; "metrics_enabled")] + #[test_case((false, false); "metrics_disabled_static_planner")] + #[test_case((true, false); "metrics_enabled_static_planner")] + #[test_case((false, true); "metrics_disabled_dynamic_planner")] + #[test_case((true, true); "metrics_enabled_dynamic_planner")] #[tokio::test(flavor = "multi_thread")] - async fn no_pending_tasks_if_dynamic_query_completes(collect_metrics: bool) -> Result<()> { + async fn no_pending_tasks_if_dynamic_query_completes( + (collect_metrics, adaptive): (bool, bool), + ) -> Result<()> { let (mut d_ctx, _guard, workers) = start_localhost_context(NUM_WORKERS, DefaultSessionBuilder).await; d_ctx.set_distributed_metrics_collection(collect_metrics)?; + d_ctx.set_distributed_dynamic_task_count(adaptive)?; run_tpch_query(d_ctx, "q1").await?; @@ -35,10 +40,18 @@ mod tests { Ok(()) } + #[test_case((false, false); "metrics_disabled_static_planner")] + #[test_case((true, false); "metrics_enabled_static_planner")] + #[test_case((false, true); "metrics_disabled_dynamic_planner")] + #[test_case((true, true); "metrics_enabled_dynamic_planner")] #[tokio::test(flavor = "multi_thread")] - async fn no_pending_tasks_if_query_aborts() -> Result<()> { - let (d_ctx, _guard, workers) = + async fn no_pending_tasks_if_query_aborts( + (collect_metrics, adaptive): (bool, bool), + ) -> Result<()> { + let (mut d_ctx, _guard, workers) = start_localhost_context(NUM_WORKERS, DefaultSessionBuilder).await; + d_ctx.set_distributed_metrics_collection(collect_metrics)?; + d_ctx.set_distributed_dynamic_task_count(adaptive)?; let _ = timeout(Duration::from_millis(100), run_tpch_query(d_ctx, "q1")).await; diff --git a/tests/tpcds_correctness_test.rs b/tests/tpcds_correctness_test.rs index bb6011a7..e4baeba9 100644 --- a/tests/tpcds_correctness_test.rs +++ b/tests/tpcds_correctness_test.rs @@ -18,6 +18,7 @@ mod tests { use std::sync::Arc; use tokio::sync::OnceCell; + const ADAPTIVE_ENV_VAR: &str = "ADAPTIVE"; const NUM_WORKERS: usize = 4; const PARTITIONS: usize = 3; const FILE_SCAN_CONFIG_BYTES_PER_PARTITION: usize = 1; @@ -575,12 +576,15 @@ mod tests { .options_mut() .execution .target_partitions = PARTITIONS; - let d_ctx = d_ctx + let mut d_ctx = d_ctx .with_distributed_file_scan_config_bytes_per_partition( FILE_SCAN_CONFIG_BYTES_PER_PARTITION, )? .with_distributed_cardinality_effect_task_scale_factor(CARDINALITY_TASK_COUNT_FACTOR)? .with_distributed_broadcast_joins(true)?; + if std::env::var(ADAPTIVE_ENV_VAR).unwrap_or_default() == "true" { + d_ctx.set_distributed_dynamic_task_count(true)?; + } register_tables(&s_ctx, &data_dir).await?; register_tables(&d_ctx, &data_dir).await?; diff --git a/tests/tpch_correctness_test.rs b/tests/tpch_correctness_test.rs index 3b6bc5a4..3ae089dc 100644 --- a/tests/tpch_correctness_test.rs +++ b/tests/tpch_correctness_test.rs @@ -12,6 +12,7 @@ mod tests { use std::path::Path; use tokio::sync::OnceCell; + const ADAPTIVE_ENV_VAR: &str = "ADAPTIVE"; const NUM_WORKERS: usize = 4; const PARTITIONS: usize = 6; const FILE_SCAN_CONFIG_BYTES_PER_PARTITION: usize = 1; @@ -139,7 +140,10 @@ mod tests { // in a non-distributed manner. For each query, it asserts that the results are identical. async fn test_tpch_query(sql: String) -> Result<(), Box> { let d_ctx = start_in_memory_context(NUM_WORKERS, DefaultSessionBuilder).await; - let d_ctx = d_ctx.with_distributed_broadcast_joins(true)?; + let mut d_ctx = d_ctx.with_distributed_broadcast_joins(true)?; + if std::env::var(ADAPTIVE_ENV_VAR).unwrap_or_default() == "true" { + d_ctx.set_distributed_dynamic_task_count(true)?; + } let d_ctx = d_ctx .with_distributed_file_scan_config_bytes_per_partition( From fa78bcaeed2036ac5a2cdd9e450fa12fb0c45836 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 29 Jun 2026 10:06:07 +0200 Subject: [PATCH 11/20] Abstract away the gRPC protocol --- .gitignore | 8 +- Cargo.toml | 23 +- console/src/worker.rs | 2 +- examples/in_memory_cluster.rs | 13 +- examples/localhost_versioned_run.rs | 12 +- src/{protobuf => codec}/distributed_codec.rs | 0 src/{protobuf => codec}/mod.rs | 6 - src/{protobuf => codec}/user_codec.rs | 0 src/common/mod.rs | 2 - src/common/time.rs | 9 +- src/coordinator/distributed.rs | 14 +- src/coordinator/metrics_store.rs | 13 +- src/coordinator/prepare_dynamic_plan.rs | 9 +- src/coordinator/query_coordinator.rs | 126 ++- src/distributed_ext.rs | 48 +- src/distributed_planner/distributed_config.rs | 3 +- src/distributed_planner/network_boundary.rs | 60 +- src/distributed_planner/task_estimator.rs | 2 +- src/execution_plans/benchmarks/fixture.rs | 14 +- .../benchmarks/transport_bench.rs | 6 +- src/execution_plans/network_broadcast.rs | 8 +- src/execution_plans/network_coalesce.rs | 5 +- src/execution_plans/network_shuffle.rs | 8 +- src/execution_plans/sampler.rs | 35 +- src/lib.rs | 45 +- src/metrics/latency_metric.rs | 12 +- src/metrics/mod.rs | 1 - src/metrics/task_metrics_collector.rs | 4 +- src/metrics/task_metrics_rewriter.rs | 69 +- src/networking/mod.rs | 11 - src/protobuf/producer_head.rs | 63 -- src/protocol/channel_resolver.rs | 98 +++ .../grpc}/channel_resolver.rs | 127 +-- .../grpc}/errors/arrow_error.rs | 2 +- .../grpc}/errors/datafusion_error.rs | 12 +- .../grpc}/errors/io_error.rs | 0 src/{protobuf => protocol/grpc}/errors/mod.rs | 2 +- .../grpc}/errors/objectstore_error.rs | 0 .../grpc}/errors/parquet_error.rs | 0 .../grpc}/errors/parser_error.rs | 0 .../grpc}/errors/schema_error.rs | 0 src/{worker => protocol/grpc}/gen/Cargo.toml | 0 .../grpc}/gen/regen.sh | 2 +- src/{worker => protocol/grpc}/gen/src/main.rs | 4 +- .../grpc}/generated/mod.rs | 0 .../grpc}/generated/worker.rs | 0 .../grpc/metrics_proto.rs} | 18 +- src/protocol/grpc/mod.rs | 21 + .../grpc}/observability/README.md | 2 +- .../grpc}/observability/gen/Cargo.toml | 0 src/protocol/grpc/observability/gen/regen.sh | 6 + .../grpc}/observability/gen/src/main.rs | 6 +- .../grpc}/observability/generated/mod.rs | 0 .../observability/generated/observability.rs | 2 +- src/{ => protocol/grpc}/observability/mod.rs | 0 .../observability/proto/observability.proto | 0 .../grpc}/observability/service.rs | 40 +- .../grpc}/on_drop_stream.rs | 0 .../grpc}/spawn_select_all.rs | 0 src/{worker => protocol/grpc}/worker.proto | 0 src/protocol/grpc/worker_client.rs | 766 ++++++++++++++++++ src/protocol/grpc/worker_service.rs | 422 ++++++++++ src/protocol/mod.rs | 14 + src/protocol/worker_channel.rs | 187 +++++ src/stage.rs | 34 +- src/test_utils/in_memory_channel_resolver.rs | 14 +- src/test_utils/metrics.rs | 65 +- src/test_utils/plans.rs | 10 +- src/work_unit_feed/remote_work_unit_feed.rs | 106 ++- src/work_unit_feed/work_unit_feed_provider.rs | 2 +- src/worker/gen/regen.sh | 6 - src/worker/impl_coordinator_channel.rs | 167 ++-- src/worker/impl_execute_task.rs | 297 +------ src/worker/mod.rs | 2 - src/worker/single_write_multi_read.rs | 1 + src/worker/task_data.rs | 94 +-- src/worker/test_utils/worker_handles.rs | 8 +- src/worker/worker_connection_pool.rs | 753 ++--------------- src/worker/worker_service.rs | 109 +-- src/{networking => }/worker_resolver.rs | 0 80 files changed, 2259 insertions(+), 1771 deletions(-) rename src/{protobuf => codec}/distributed_codec.rs (100%) rename src/{protobuf => codec}/mod.rs (55%) rename src/{protobuf => codec}/user_codec.rs (100%) delete mode 100644 src/networking/mod.rs delete mode 100644 src/protobuf/producer_head.rs create mode 100644 src/protocol/channel_resolver.rs rename src/{networking => protocol/grpc}/channel_resolver.rs (56%) rename src/{protobuf => protocol/grpc}/errors/arrow_error.rs (99%) rename src/{protobuf => protocol/grpc}/errors/datafusion_error.rs (97%) rename src/{protobuf => protocol/grpc}/errors/io_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/mod.rs (97%) rename src/{protobuf => protocol/grpc}/errors/objectstore_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/parquet_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/parser_error.rs (100%) rename src/{protobuf => protocol/grpc}/errors/schema_error.rs (100%) rename src/{worker => protocol/grpc}/gen/Cargo.toml (100%) rename src/{observability => protocol/grpc}/gen/regen.sh (57%) rename src/{worker => protocol/grpc}/gen/src/main.rs (87%) rename src/{worker => protocol/grpc}/generated/mod.rs (100%) rename src/{worker => protocol/grpc}/generated/worker.rs (100%) rename src/{metrics/proto.rs => protocol/grpc/metrics_proto.rs} (98%) create mode 100644 src/protocol/grpc/mod.rs rename src/{ => protocol/grpc}/observability/README.md (88%) rename src/{ => protocol/grpc}/observability/gen/Cargo.toml (100%) create mode 100755 src/protocol/grpc/observability/gen/regen.sh rename src/{ => protocol/grpc}/observability/gen/src/main.rs (76%) rename src/{ => protocol/grpc}/observability/generated/mod.rs (100%) rename src/{ => protocol/grpc}/observability/generated/observability.rs (99%) rename src/{ => protocol/grpc}/observability/mod.rs (100%) rename src/{ => protocol/grpc}/observability/proto/observability.proto (100%) rename src/{ => protocol/grpc}/observability/service.rs (80%) rename src/{common => protocol/grpc}/on_drop_stream.rs (100%) rename src/{worker => protocol/grpc}/spawn_select_all.rs (100%) rename src/{worker => protocol/grpc}/worker.proto (100%) create mode 100644 src/protocol/grpc/worker_client.rs create mode 100644 src/protocol/grpc/worker_service.rs create mode 100644 src/protocol/mod.rs create mode 100644 src/protocol/worker_channel.rs delete mode 100755 src/worker/gen/regen.sh rename src/{networking => }/worker_resolver.rs (100%) diff --git a/.gitignore b/.gitignore index 479b6815..70fa5beb 100644 --- a/.gitignore +++ b/.gitignore @@ -8,7 +8,7 @@ testdata/tpcds/* !testdata/tpcds/README.md testdata/clickbench/* !testdata/clickbench/queries -src/observability/gen/target/* -src/observability/gen/Cargo.lock -src/worker/gen/target/* -src/worker/gen/Cargo.lock +src/protocol/grpc/observability/gen/target/* +src/protocol/grpc/observability/gen/Cargo.lock +src/protocol/grpc/gen/target/* +src/protocol/grpc/gen/Cargo.lock diff --git a/Cargo.toml b/Cargo.toml index fd5ceab2..79366ed3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,13 +23,8 @@ datafusion = { workspace = true, features = [ "datetime_expressions", ] } datafusion-proto = { workspace = true } -arrow-flight = "58" -arrow-select = "58" -arrow-ipc = { version = "58", features = ["zstd"] } async-trait = "0.1.89" tokio = { version = "1.48", features = ["full"] } -tonic = { version = "0.14.1", features = ["transport"] } -tower = "0.5.2" http = "1.3.1" itertools = "0.14.0" futures = "0.3.31" @@ -52,6 +47,13 @@ num-traits = "0.2" bincode = "1" tonic-prost = "0.14.2" +# grpc-specific features +arrow-flight = { version = "58", optional = true } +arrow-select = { version = "58", optional = true } +arrow-ipc = { version = "58", features = ["zstd"], optional = true } +tonic = { version = "0.14.1", features = ["transport"], optional = true } +tower = { version = "0.5.2", optional = true } + # integration_tests deps insta = { version = "1.46.0", features = ["filters"], optional = true } parquet = { version = "58", optional = true } @@ -59,8 +61,17 @@ arrow = { version = "58", optional = true, features = ["test_utils"] } hyper-util = { version = "0.1.16", optional = true } [features] +default = ["grpc"] avro = ["datafusion/avro"] -integration = ["insta", "parquet", "arrow", "hyper-util"] +grpc = [ + "arrow-flight", + "arrow-select", + "arrow-ipc", + "tonic", + "tower" +] + +integration = ["insta", "parquet", "arrow", "hyper-util", "grpc"] system-metrics = ["sysinfo"] diff --git a/console/src/worker.rs b/console/src/worker.rs index d84fe602..4a23c5b9 100644 --- a/console/src/worker.rs +++ b/console/src/worker.rs @@ -1,5 +1,5 @@ use datafusion::common::{HashMap, HashSet}; -use datafusion_distributed::{ +use datafusion_distributed::grpc::{ GetClusterWorkersRequest, GetTaskProgressRequest, ObservabilityServiceClient, PingRequest, TaskProgress, TaskStatus, }; diff --git a/examples/in_memory_cluster.rs b/examples/in_memory_cluster.rs index 1d02737c..6a9d7ba1 100644 --- a/examples/in_memory_cluster.rs +++ b/examples/in_memory_cluster.rs @@ -4,9 +4,8 @@ use datafusion::common::DataFusionError; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_distributed::{ - BoxCloneSyncChannel, ChannelResolver, DistributedExt, SessionStateBuilderExt, Worker, - WorkerQueryContext, WorkerResolver, WorkerServiceClient, create_worker_client, - display_plan_ascii, + ChannelResolver, DistributedExt, SessionStateBuilderExt, Worker, WorkerChannel, + WorkerQueryContext, WorkerResolver, display_plan_ascii, grpc, }; use futures::TryStreamExt; use hyper_util::rt::TokioIo; @@ -66,7 +65,7 @@ const DUMMY_URL: &str = "http://localhost:50051"; /// tokio duplex rather than a TCP connection. #[derive(Clone)] struct InMemoryChannelResolver { - channel: WorkerServiceClient, + channel: grpc::BoxCloneSyncChannel, } impl InMemoryChannelResolver { @@ -84,7 +83,7 @@ impl InMemoryChannelResolver { })); let this = Self { - channel: create_worker_client(BoxCloneSyncChannel::new(channel)), + channel: grpc::BoxCloneSyncChannel::new(channel), }; let this_clone = this.clone(); @@ -110,8 +109,8 @@ impl ChannelResolver for InMemoryChannelResolver { async fn get_worker_client_for_url( &self, _: &url::Url, - ) -> Result, DataFusionError> { - Ok(self.channel.clone()) + ) -> Result, DataFusionError> { + Ok(grpc::create_worker_client(self.channel.clone())) } } diff --git a/examples/localhost_versioned_run.rs b/examples/localhost_versioned_run.rs index dc3f38c2..b10520cb 100644 --- a/examples/localhost_versioned_run.rs +++ b/examples/localhost_versioned_run.rs @@ -4,8 +4,8 @@ use datafusion::common::DataFusionError; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_distributed::{ - DefaultChannelResolver, DistributedExt, GetWorkerInfoRequest, SessionStateBuilderExt, - WorkerResolver, create_worker_client, display_plan_ascii, + DistributedExt, GetWorkerInfoRequest, SessionStateBuilderExt, WorkerResolver, + display_plan_ascii, grpc, }; use futures::TryStreamExt; use std::error::Error; @@ -40,7 +40,7 @@ struct Args { /// the `GetWorkerInfo` RPC. Returns `false` if the worker is unreachable, returns /// an error, or reports a different version. async fn worker_has_version( - channel_resolver: &DefaultChannelResolver, + channel_resolver: &grpc::DefaultChannelResolver, url: &Url, expected_version: &str, ) -> bool { @@ -48,12 +48,12 @@ async fn worker_has_version( return false; }; - let mut client = create_worker_client(channel); + let mut client = grpc::create_worker_client(channel); let Ok(response) = client.get_worker_info(GetWorkerInfoRequest {}).await else { return false; }; - response.into_inner().version == expected_version + response.version == expected_version } #[tokio::main] @@ -61,7 +61,7 @@ async fn main() -> Result<(), Box> { let args = Args::from_args(); let ports = if let Some(target_version) = &args.version { - let channel_resolver = DefaultChannelResolver::default(); + let channel_resolver = grpc::DefaultChannelResolver::default(); let mut compatible = Vec::new(); for &port in &args.cluster_ports { let url = Url::parse(&format!("http://localhost:{port}"))?; diff --git a/src/protobuf/distributed_codec.rs b/src/codec/distributed_codec.rs similarity index 100% rename from src/protobuf/distributed_codec.rs rename to src/codec/distributed_codec.rs diff --git a/src/protobuf/mod.rs b/src/codec/mod.rs similarity index 55% rename from src/protobuf/mod.rs rename to src/codec/mod.rs index 47deeb40..c4acaf1a 100644 --- a/src/protobuf/mod.rs +++ b/src/codec/mod.rs @@ -1,13 +1,7 @@ mod distributed_codec; -mod errors; -mod producer_head; mod user_codec; pub use distributed_codec::DistributedCodec; -pub(crate) use errors::{ - datafusion_error_to_tonic_status, map_flight_to_datafusion_error, - tonic_status_to_datafusion_error, -}; pub(crate) use user_codec::{ get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc, }; diff --git a/src/protobuf/user_codec.rs b/src/codec/user_codec.rs similarity index 100% rename from src/protobuf/user_codec.rs rename to src/codec/user_codec.rs diff --git a/src/common/mod.rs b/src/common/mod.rs index 18cb28a1..07463ca0 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,5 +1,4 @@ mod children_helpers; -mod on_drop_stream; mod once_lock; mod recursion; mod task_context_helpers; @@ -8,7 +7,6 @@ mod uuid; mod vec; pub(crate) use children_helpers::require_one_child; -pub(crate) use on_drop_stream::on_drop_stream; 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/common/time.rs b/src/common/time.rs index 07757f32..6efb572c 100644 --- a/src/common/time.rs +++ b/src/common/time.rs @@ -1,9 +1,14 @@ +use num_traits::AsPrimitive; use std::time::{SystemTime, UNIX_EPOCH}; /// Nanoseconds elapsed since UNIX epoch. -pub(crate) fn now_ns() -> u64 { +pub(crate) fn now_ns() -> T +where + T: 'static + Copy, + u128: AsPrimitive, +{ SystemTime::now() .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_nanos() as u64) + .map(|duration| duration.as_nanos().as_()) .expect("SystemTime before UNIX EPOCH!") } diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index d7d62a08..9c682328 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,11 +1,10 @@ -use crate::DistributedConfig; -use crate::common::{require_one_child, serialize_uuid}; +use crate::common::require_one_child; use crate::coordinator::metrics_store::MetricsStore; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; use crate::distributed_planner::NetworkBoundaryExt; -use crate::worker::generated::worker::TaskKey; +use crate::{DistributedConfig, TaskKey}; use datafusion::common::internal_datafusion_err; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{Result, exec_err}; @@ -94,9 +93,9 @@ impl DistributedExec { let stage = boundary.input_stage(); for i in 0..stage.task_count() { expected_keys.push(TaskKey { - query_id: serialize_uuid(&stage.query_id()), - stage_id: stage.num() as u64, - task_number: i as u64, + query_id: stage.query_id(), + stage_id: stage.num(), + task_number: i, }); } } @@ -198,7 +197,7 @@ impl ExecutionPlan for DistributedExec { let tx = builder.tx(); builder.spawn(async move { - let _guard = query_coordinator.end_query_guard(); + let guard = query_coordinator.end_query_guard(); let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?; let result = match d_cfg.dynamic_task_count { @@ -221,6 +220,7 @@ impl ExecutionPlan for DistributedExec { } } drop(tx); + drop(guard); query_coordinator.drain_pending_tasks().await?; Ok(()) }); diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index 4b9b0ffc..8ccc1f7c 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -1,9 +1,8 @@ -use crate::TaskKey; -use crate::worker::generated::worker as pb; +use crate::{TaskKey, TaskMetrics}; use datafusion::common::HashMap; use tokio::sync::watch; -type MetricsMap = HashMap; +type MetricsMap = HashMap; /// Stores the metrics collected from all worker tasks, and notifies waiters when new entries arrive. #[derive(Debug, Clone)] @@ -18,20 +17,18 @@ impl MetricsStore { Self { tx, rx } } - pub(crate) fn insert(&self, key: TaskKey, metrics: pb::TaskMetrics) { + pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { self.tx.send_modify(|map| { map.insert(key, metrics); }); } - pub(crate) fn get(&self, key: &TaskKey) -> Option { + pub(crate) fn get(&self, key: &TaskKey) -> Option { self.rx.borrow().get(key).cloned() } #[cfg(test)] - pub(crate) fn from_entries( - entries: impl IntoIterator, - ) -> Self { + pub(crate) fn from_entries(entries: impl IntoIterator) -> Self { let map: HashMap<_, _> = entries.into_iter().collect(); let (tx, rx) = watch::channel(map); Self { tx, rx } diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index d90fce5e..b22a91e5 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -8,8 +8,7 @@ use crate::distributed_planner::{ }; use crate::execution_plans::SamplerExec; use crate::stage::{LocalStage, RemoteStage}; -use crate::worker::generated::worker as pb; -use crate::{BytesCounterMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage}; +use crate::{BytesCounterMetric, LoadInfo, NetworkBoundaryExt, NetworkCoalesceExec, Stage}; use dashmap::DashMap; use datafusion::common::stats::Precision; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -193,7 +192,7 @@ impl PlanReconstructor { /// Estimates the bytes per second flowing through a stage by reading sample information. async fn gather_runtime_statistics( - per_task_load_info_stream: Vec + Unpin>, + per_task_load_info_stream: Vec + Unpin>, plan: &Arc, ) -> Result { const ESTIMATED_QUERY_TIME_S: usize = 10; @@ -226,8 +225,8 @@ async fn gather_runtime_statistics( let mut load_info_stream = futures::stream::select_all(per_task_load_info_stream); while let Some(load_info) = load_info_stream.next().await { - rows_per_second += load_info.rows_per_second as usize; - rows_ready += load_info.rows_ready as usize; + rows_per_second += load_info.rows_per_second; + rows_ready += load_info.rows_ready; per_col_bytes_per_second = element_wise_sum( per_col_bytes_per_second, &load_info.per_column_bytes_per_second, diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 6f0010e6..1e309410 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,25 +1,22 @@ -use crate::common::{TreeNodeExt, now_ns, serialize_uuid, task_ctx_with_extension}; +use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::MetricsStore; use crate::coordinator::latency_metric::LatencyMetric; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; -use crate::protobuf::tonic_status_to_datafusion_error; use crate::stage::LocalStage; use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; -use crate::worker::generated::worker as pb; -use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; -use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; use crate::{ - BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, - DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, NetworkBoundaryExt, - Stage, TaskEstimator, TaskKey, TaskRoutingContext, get_distributed_channel_resolver, - get_distributed_worker_resolver, + BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg, + DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedConfig, + DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, 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; use datafusion::common::runtime::JoinSet; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion::common::{DataFusionError, exec_datafusion_err}; use datafusion::common::{Result, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; @@ -27,8 +24,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion_proto::physical_plan::AsExecutionPlan; use datafusion_proto::protobuf::PhysicalPlanNode; -use futures::{Stream, StreamExt}; -use http::Extensions; +use futures::{Stream, StreamExt, TryStreamExt}; use prost::Message; use rand::Rng; use std::ops::DerefMut; @@ -36,13 +32,11 @@ use std::sync::{Arc, Mutex}; use tokio::sync::Notify; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; -use tonic::Request; -use tonic::metadata::MetadataMap; use url::Url; use uuid::Uuid; -/// How many [crate::WorkUnit] messages are allowed to be chunked synchronously together in order to -/// send fewer bigger [crate::WorkUnit] batches over the wire, reducing the overhead of sending many +/// How many [crate::WorkUnitMsg] messages are allowed to be chunked synchronously together in order to +/// send fewer bigger [crate::WorkUnitMsg] batches over the wire, reducing the overhead of sending many /// small batches. See [StreamExt::ready_chunks] docs for more details about how chunking works. const WORK_UNIT_FEED_CHUNK_SIZE: usize = 256; @@ -142,8 +136,8 @@ impl<'a> StageCoordinator<'a> { task_i: usize, url: Url, ) -> Result<( - UnboundedSender, - UnboundedReceiver, + UnboundedSender, + UnboundedReceiver, )> { let session_config = self.task_ctx.session_config(); let codec = DistributedCodec::new_combined_with_user(session_config); @@ -155,21 +149,20 @@ impl<'a> StageCoordinator<'a> { let plan_size = plan_proto.len(); let task_key = TaskKey { - query_id: serialize_uuid(&self.query_id), - stage_id: self.stage_id as u64, - task_number: task_i as u64, - }; - let msg = pb::CoordinatorToWorkerMsg { - inner: Some(Inner::SetPlanRequest(pb::SetPlanRequest { - plan_proto, - task_count: self.task_count as u64, - task_key: Some(task_key.clone()), - work_unit_feed_declarations, - target_worker_url: url.to_string(), - query_start_time_ns: self.metrics.instantiation_time, - })), + query_id: self.query_id, + stage_id: self.stage_id, + task_number: task_i, }; + let msg = CoordinatorToWorkerMsg::SetPlanRequest(SetPlanRequest { + task_key, + task_count: self.task_count, + plan_proto, + 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(); let (worker_to_coordinator_tx, worker_to_coordinator_rx) = @@ -180,37 +173,26 @@ impl<'a> StageCoordinator<'a> { let mut headers = get_config_extension_propagation_headers(session_config)?; headers.extend(get_passthrough_headers(session_config)); - let request = Request::from_parts( - MetadataMap::from_headers(headers), - Extensions::default(), - futures::stream::once(async { msg }) - .chain(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 - // EOS of this stream as a query finished/aborted signal. - .chain(keep_stream_alive(Arc::clone(self.end_stream_notifier))), - ); + let coordinator_to_worker_stream = futures::stream::once(async { msg }) + .chain(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 + // EOS of this stream as a query finished/aborted signal. + .chain(keep_stream_alive(Arc::clone(self.end_stream_notifier))) + .boxed(); let metrics = self.metrics.clone(); 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 response = client.coordinator_channel(request).await.map_err(|e| { - tonic_status_to_datafusion_error(&e).unwrap_or_else(|| { - exec_datafusion_err!("Error sending plan to worker {url}: {e}") - }) - })?; + let mut worker_to_coordinator_stream = client + .coordinator_channel(headers, coordinator_to_worker_stream) + .await?; metrics.plan_send_latency.record(&start); metrics.plan_bytes_sent.add_bytes(plan_size); - let mut worker_to_coordinator_stream = response.into_inner(); - while let Some(msg_or_err) = worker_to_coordinator_stream.next().await { - let msg = msg_or_err.map_err(|err| { - tonic_status_to_datafusion_error(&err).unwrap_or_else(|| { - exec_datafusion_err!("Unknown error on worker to coordinator stream: {err}") - }) - })?; + while let Some(msg) = worker_to_coordinator_stream.try_next().await? { if worker_to_coordinator_tx.send(msg).is_err() { break; // receiver dropped } @@ -227,12 +209,12 @@ impl<'a> StageCoordinator<'a> { pub(super) fn worker_to_coordinator_task( &mut self, task_i: usize, - mut worker_to_coordinator_rx: UnboundedReceiver, - ) -> UnboundedReceiver { + mut worker_to_coordinator_rx: UnboundedReceiver, + ) -> UnboundedReceiver { let task_key = TaskKey { - query_id: serialize_uuid(&self.query_id), - stage_id: self.stage_id as u64, - task_number: task_i as u64, + query_id: self.query_id, + stage_id: self.stage_id, + task_number: task_i, }; let task_metrics = self.metrics_store.clone(); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); @@ -243,20 +225,18 @@ impl<'a> StageCoordinator<'a> { #[allow(clippy::disallowed_methods)] tokio::spawn(async move { while let Some(msg) = worker_to_coordinator_rx.recv().await { - let Some(inner) = msg.inner else { continue }; - - match inner { - pb::worker_to_coordinator_msg::Inner::TaskMetrics(pre_order_metrics) => { + match msg { + WorkerToCoordinatorMsg::TaskMetrics(v) => { if let Some(task_metrics) = &task_metrics { - task_metrics.insert(task_key.clone(), pre_order_metrics); + task_metrics.insert(task_key, v); } } - pb::worker_to_coordinator_msg::Inner::LoadInfo(load_info) => { + WorkerToCoordinatorMsg::LoadInfo(load_info) => { if let Some(tx) = &load_info_tx_opt { let _ = tx.send(load_info); } } - pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { + WorkerToCoordinatorMsg::LoadInfoEos => { let _ = load_info_tx_opt.take(); } } @@ -271,7 +251,7 @@ impl<'a> StageCoordinator<'a> { pub(super) fn coordinator_to_worker_task( &mut self, task_i: usize, - tx: UnboundedSender, + tx: UnboundedSender, ) -> Result<()> { let session_config = self.task_ctx.session_config(); let d_cfg = DistributedConfig::from_config_options(session_config.options())?; @@ -322,12 +302,10 @@ impl<'a> StageCoordinator<'a> { Ok(TreeNodeRecursion::Continue) })?; - struct WorkUnitEosOnDrop(UnboundedSender); + struct WorkUnitEosOnDrop(UnboundedSender); impl Drop for WorkUnitEosOnDrop { fn drop(&mut self) { - let _ = self.0.send(pb::CoordinatorToWorkerMsg { - inner: Some(Inner::WorkUnitEos(true)), - }); + let _ = self.0.send(CoordinatorToWorkerMsg::WorkUnitEos); } } @@ -361,8 +339,8 @@ impl<'a> StageCoordinator<'a> { let transformed = plan.transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| { if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) { work_unit_feed_declarations.push(WorkUnitFeedDeclaration { - id: serialize_uuid(&wuf.id()), - partitions: plan.properties().partitioning.partition_count() as u64, + id: wuf.id(), + partitions: plan.properties().partitioning.partition_count(), }); }; @@ -460,7 +438,7 @@ impl Drop for NotifyGuard { pub(super) struct CoordinatorToWorkerMetrics { pub(super) plan_bytes_sent: BytesCounterMetric, pub(super) plan_send_latency: Arc, - pub(super) instantiation_time: u64, + pub(super) instantiation_time: usize, } impl CoordinatorToWorkerMetrics { diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index d33b18a6..9efdb420 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -1,16 +1,16 @@ +use crate::codec::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::config_extension_ext::{ set_distributed_option_extension, set_distributed_option_extension_from_headers, }; use crate::distributed_planner::set_distributed_task_estimator; -use crate::networking::{set_distributed_channel_resolver, set_distributed_worker_resolver}; use crate::passthrough_headers::set_passthrough_headers; -use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc}; +use crate::protocol::set_distributed_channel_resolver; use crate::work_unit_feed::set_distributed_work_unit_feed; +use crate::worker_resolver::set_distributed_worker_resolver; use crate::{ ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, WorkerResolver, }; -use arrow_ipc::CompressionType; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; use datafusion::execution::{SessionState, SessionStateBuilder}; @@ -202,7 +202,7 @@ pub trait DistributedExt: Sized { /// # use datafusion::prelude::SessionConfig; /// # use url::Url; /// # use std::sync::Arc; - /// # use datafusion_distributed::{BoxCloneSyncChannel, WorkerResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext}; + /// # use datafusion_distributed::{WorkerResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext}; /// /// struct CustomWorkerResolver; /// @@ -238,14 +238,14 @@ pub trait DistributedExt: Sized { /// # use datafusion::prelude::SessionConfig; /// # use url::Url; /// # use std::sync::Arc; - /// # use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt, SessionStateBuilderExt, WorkerQueryContext, WorkerServiceClient}; + /// # use datafusion_distributed::{ChannelResolver, DistributedExt, SessionStateBuilderExt, WorkerChannel, WorkerQueryContext, grpc}; /// /// struct CustomChannelResolver; /// /// #[async_trait] /// impl ChannelResolver for CustomChannelResolver { - /// async fn get_worker_client_for_url(&self, url: &Url) -> Result, DataFusionError> { - /// // Build a custom WorkerServiceClient wrapped with tower layers or something similar. + /// async fn get_worker_client_for_url(&self, url: &Url) -> Result, DataFusionError> { + /// // Build a custom worker client wrapped with tower layers or something similar. /// todo!() /// } /// } @@ -451,18 +451,20 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_broadcast_joins_enabled] but with an in-place mutation. fn set_distributed_broadcast_joins(&mut self, enabled: bool) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] /// The compression type to use for sending data over the wire. /// /// The default is [CompressionType::LZ4_FRAME]. fn with_distributed_compression( self, - compression: Option, + compression: Option, ) -> Result; + #[cfg(feature = "grpc")] /// Same as [DistributedExt::with_distributed_compression] but with an in-place mutation. fn set_distributed_compression( &mut self, - compression: Option, + compression: Option, ) -> Result<(), DataFusionError>; /// Overrides `datafusion.execution.batch_size` for worker-executed stages, letting users @@ -678,14 +680,15 @@ impl DistributedExt for SessionConfig { Ok(()) } + #[cfg(feature = "grpc")] fn set_distributed_compression( &mut self, - compression: Option, + compression: Option, ) -> Result<(), DataFusionError> { let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?; d_cfg.compression = match compression { - Some(CompressionType::ZSTD) => "zstd".to_string(), - Some(CompressionType::LZ4_FRAME) => "lz4".to_string(), + Some(arrow_ipc::CompressionType::ZSTD) => "zstd".to_string(), + Some(arrow_ipc::CompressionType::LZ4_FRAME) => "lz4".to_string(), _ => "none".to_string(), }; Ok(()) @@ -810,7 +813,8 @@ impl DistributedExt for SessionConfig { #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(mut self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(mut self, compression: Option) -> Result; #[call(set_distributed_shuffle_batch_size)] #[expr($?;Ok(self))] @@ -915,10 +919,12 @@ impl DistributedExt for SessionStateBuilder { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; - fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(mut self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(mut self, compression: Option) -> Result; fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>; #[call(set_distributed_shuffle_batch_size)] @@ -1036,10 +1042,12 @@ impl DistributedExt for SessionState { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(mut self, enabled: bool) -> Result; - fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(mut self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(mut self, compression: Option) -> Result; fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>; #[call(set_distributed_shuffle_batch_size)] @@ -1157,10 +1165,12 @@ impl DistributedExt for SessionContext { #[expr($?;Ok(self))] fn with_distributed_broadcast_joins(self, enabled: bool) -> Result; - fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; + #[cfg(feature = "grpc")] + fn set_distributed_compression(&mut self, compression: Option) -> Result<(), DataFusionError>; #[call(set_distributed_compression)] #[expr($?;Ok(self))] - fn with_distributed_compression(self, compression: Option) -> Result; + #[cfg(feature = "grpc")] + fn with_distributed_compression(self, compression: Option) -> Result; fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>; #[call(set_distributed_shuffle_batch_size)] diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index e01b5ff4..92cc9352 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -1,6 +1,7 @@ use crate::distributed_planner::task_estimator::CombinedTaskEstimator; -use crate::networking::{ChannelResolverExtension, WorkerResolverExtension}; +use crate::protocol::ChannelResolverExtension; use crate::work_unit_feed::WorkUnitFeedRegistry; +use crate::worker_resolver::WorkerResolverExtension; use crate::{TaskEstimator, WorkerResolver}; use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err}; use datafusion::config::{ConfigExtension, ConfigField, ConfigOptions, Visit}; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index e6479e7e..3b7bd3b1 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -1,9 +1,22 @@ use crate::execution_plans::SamplerExec; -use crate::{BroadcastExec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage}; +use crate::protocol::ProducerHeadSpec; +use crate::{ + BroadcastExec, DistributedCodec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, + Stage, +}; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::Result; +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. @@ -61,6 +74,51 @@ 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 }) + } + } + } + /// Ensures the head of the provided plan complies with the passed [ProducerHead] definition. This /// can be called both during planning and lazily at runtime. pub(crate) fn insert(self, input: Arc) -> Result> { diff --git a/src/distributed_planner/task_estimator.rs b/src/distributed_planner/task_estimator.rs index 1e792579..b915e6c1 100644 --- a/src/distributed_planner/task_estimator.rs +++ b/src/distributed_planner/task_estimator.rs @@ -388,9 +388,9 @@ impl TaskEstimator for CombinedTaskEstimator { #[cfg(test)] mod tests { use super::*; - use crate::networking::WorkerResolverExtension; use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; use crate::test_utils::parquet::register_parquet_tables; + use crate::worker_resolver::WorkerResolverExtension; use datafusion::error::DataFusionError; use datafusion::prelude::SessionContext; diff --git a/src/execution_plans/benchmarks/fixture.rs b/src/execution_plans/benchmarks/fixture.rs index 4423d7d0..d8b2cc88 100644 --- a/src/execution_plans/benchmarks/fixture.rs +++ b/src/execution_plans/benchmarks/fixture.rs @@ -1,5 +1,4 @@ -use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; -use crate::{BoxCloneSyncChannel, ChannelResolver, create_worker_client}; +use crate::{ChannelResolver, WorkerChannel, grpc}; use arrow::datatypes::DataType::{ Boolean, Dictionary, Float64, Int32, Int64, List, Timestamp, UInt8, Utf8, }; @@ -72,18 +71,17 @@ pub(super) fn rows_for_producer( /// tokio duplex rather than a TCP connection. #[derive(Clone)] pub(super) struct InMemoryChannelsResolver { - pub channels: Vec, + pub channels: Vec, } #[async_trait::async_trait] impl ChannelResolver for InMemoryChannelsResolver { - async fn get_worker_client_for_url( - &self, - url: &Url, - ) -> Result> { + async fn get_worker_client_for_url(&self, url: &Url) -> Result> { let Some(port) = url.port() else { return exec_err!("Missing port in url {url}"); }; - Ok(create_worker_client(self.channels[port as usize].clone())) + Ok(grpc::create_worker_client( + self.channels[port as usize].clone(), + )) } } diff --git a/src/execution_plans/benchmarks/transport_bench.rs b/src/execution_plans/benchmarks/transport_bench.rs index a386c7bd..ad159cd5 100644 --- a/src/execution_plans/benchmarks/transport_bench.rs +++ b/src/execution_plans/benchmarks/transport_bench.rs @@ -4,9 +4,7 @@ use super::fixture::{ use crate::common::task_ctx_with_extension; use crate::stage::RemoteStage; use crate::worker::test_utils::worker_handles::{MemoryWorkerHandle, TcpWorkerHandle}; -use crate::{ - DefaultChannelResolver, DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage, -}; +use crate::{DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage, grpc}; use arrow::datatypes::Schema; use arrow_ipc::CompressionType; use datafusion::common::Result; @@ -233,7 +231,7 @@ impl TransportBench { bench: self.clone(), schema, task_ctx: SessionStateBuilder::new() - .with_distributed_channel_resolver(DefaultChannelResolver::default()) + .with_distributed_channel_resolver(grpc::DefaultChannelResolver::default()) .with_distributed_compression(self.compression)? .build() .task_ctx(), diff --git a/src/execution_plans/network_broadcast.rs b/src/execution_plans/network_broadcast.rs index 251ed29c..aa84ea78 100644 --- a/src/execution_plans/network_broadcast.rs +++ b/src/execution_plans/network_broadcast.rs @@ -248,16 +248,14 @@ impl ExecutionPlan for NetworkBroadcastExec { let mut streams = Vec::with_capacity(self.input_stage.task_count()); for input_task_index in 0..self.input_stage.task_count() { - let worker_connection = self.worker_connections.get_or_init_worker_connection( + streams.push(self.worker_connections.execute( remote_stage, off..(off + self.properties.partitioning.partition_count()), input_task_index, + off + partition, self.producer_head(task_context.task_count), &context, - )?; - - let stream = worker_connection.execute(off + partition)?; - streams.push(stream); + )?); } Ok(Box::pin(RecordBatchStreamAdapter::new( diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 1582d08c..703843cc 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -292,16 +292,15 @@ impl ExecutionPlan for NetworkCoalesceExec { let target_task = group.start_task + input_task_offset; - let worker_connection = self.worker_connections.get_or_init_worker_connection( + let stream = self.worker_connections.execute( remote_stage, 0..partitions_per_task, target_task, + target_partition, self.producer_head(task_context.task_count), &context, )?; - let stream = worker_connection.execute(target_partition)?; - Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), stream, diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 2a575b42..c7a645ae 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -225,16 +225,14 @@ impl ExecutionPlan for NetworkShuffleExec { let mut streams = Vec::with_capacity(remote_stage.workers.len()); for input_task_index in 0..remote_stage.workers.len() { - let worker_connection = self.worker_connections.get_or_init_worker_connection( + streams.push(self.worker_connections.execute( remote_stage, off..(off + self.properties.partitioning.partition_count()), input_task_index, + off + partition, self.producer_head(task_context.task_count), &context, - )?; - - let stream = worker_connection.execute(off + partition)?; - streams.push(stream); + )?); } Ok(Box::pin(RecordBatchStreamAdapter::new( diff --git a/src/execution_plans/sampler.rs b/src/execution_plans/sampler.rs index 0e120219..836d9986 100644 --- a/src/execution_plans/sampler.rs +++ b/src/execution_plans/sampler.rs @@ -1,7 +1,6 @@ use crate::common::{require_one_child, vec_cast}; -use crate::worker::generated::worker as pb; use crate::{ - BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, MaxGaugeMetric, + BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, LoadInfo, MaxGaugeMetric, MaxLatencyMetric, P50LatencyMetric, }; use datafusion::arrow::array::Array; @@ -58,7 +57,7 @@ pub(crate) struct SamplerExecMetrics { /// the input arrived kick_off_to_fist_batch_p50: P50LatencyMetric, kick_off_to_fist_batch_max: MaxLatencyMetric, - /// Time since [SamplerExec::kick_off_first_sampler] was called until the [pb::LoadInfo] message + /// Time since [SamplerExec::kick_off_first_sampler] was called until the [LoadInfo] message /// was sent. kick_off_to_load_info_sent_p50: P50LatencyMetric, kick_off_to_load_info_sent_max: MaxLatencyMetric, @@ -141,7 +140,7 @@ impl SamplerExec { pub(crate) fn kick_off_first_sampler( plan: Arc, ctx: Arc, - ) -> Result>> { + ) -> Result>> { let mut receivers = vec![]; plan.apply(|plan| { let Some(sampler) = plan.downcast_ref::() else { @@ -197,7 +196,7 @@ impl PartitionSampler { self.metrics.kick_off_to_fist_batch_max.add_duration(delay); } - // Time since the sampler was kicked off until the pb::LoadInfo message was sent. + // Time since the sampler was kicked off until the LoadInfo message was sent. if let Some(t) = self.load_info_sent_at.get() { let delay = t.saturating_duration_since(*kick_off_at); self.metrics @@ -216,7 +215,7 @@ impl PartitionSampler { self.stream.lock().unwrap().take() } - fn kick_off(&self, ctx: Arc) -> Result> { + fn kick_off(&self, ctx: Arc) -> Result> { let _ = self.kick_off_at.set(Instant::now()); let (sampling_tx, sampling_rx) = oneshot::channel(); @@ -326,16 +325,16 @@ impl PartitionSampler { } } -/// Wraps a [pb::LoadInfo] and emits it on [Drop] through the provided [oneshot] channel. +/// Wraps a [LoadInfo] and emits it on [Drop] through the provided [oneshot] channel. /// /// Emitting on drop ensures that it's always emitted. struct LoadInfoDropHandler { omit: Arc, - load_info: pb::LoadInfo, + load_info: LoadInfo, bytes_ready_metric: BytesCounterMetric, bytes_per_second_metric: BytesCounterMetric, - sampling_tx: Option>, + sampling_tx: Option>, load_info_sent_at: Arc>, } @@ -359,16 +358,14 @@ impl LoadInfoDropHandler { } fn set_per_col_bytes_per_second(&mut self, total_bytes_per_second: usize) { - let per_col_ready: &[u64] = &self.load_info.per_column_bytes_ready; - let total_ready: u64 = per_col_ready.iter().sum(); + let per_col_ready = &self.load_info.per_column_bytes_ready; + let total_ready = per_col_ready.iter().sum::(); let per_col: Vec = if total_ready == 0 { vec![total_bytes_per_second / per_col_ready.len().max(1); per_col_ready.len()] } else { per_col_ready .iter() - .map(|&ready| { - (ready.saturating_mul(total_bytes_per_second as u64) / total_ready) as usize - }) + .map(|&ready| ready.saturating_mul(total_bytes_per_second) / total_ready) .collect() }; self.load_info.per_column_bytes_per_second = vec_cast(&per_col); @@ -377,11 +374,11 @@ impl LoadInfoDropHandler { } fn set_rows_ready(&mut self, rows_ready: usize) { - self.load_info.rows_ready = rows_ready as u64; + self.load_info.rows_ready = rows_ready; } fn set_rows_per_second(&mut self, rows_per_second: usize) { - self.load_info.rows_per_second = rows_per_second as u64; + self.load_info.rows_per_second = rows_per_second; } fn set_per_col_ndv(&mut self, per_column_ndv: Vec) { @@ -405,9 +402,9 @@ impl Drop for LoadInfoDropHandler { } } -fn zero_load_info(partition_idx: usize, n_cols: usize) -> pb::LoadInfo { - pb::LoadInfo { - partition: partition_idx as u64, +fn zero_load_info(partition_idx: usize, n_cols: usize) -> LoadInfo { + LoadInfo { + partition: partition_idx, rows_per_second: 0, rows_ready: 0, per_column_bytes_per_second: vec![0; n_cols], diff --git a/src/lib.rs b/src/lib.rs index d2a3a463..b714fe41 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,24 +1,21 @@ #![deny(clippy::all)] +mod codec; mod common; mod config_extension_ext; +mod coordinator; mod distributed_ext; +mod distributed_planner; mod execution_plans; mod metrics; mod passthrough_headers; +mod protocol; mod stage; -mod worker; - -mod distributed_planner; -mod networking; -mod observability; -mod protobuf; -pub use protobuf::DistributedCodec; -mod coordinator; -#[cfg(any(feature = "integration", test))] -pub mod test_utils; mod work_unit_feed; +mod worker; +mod worker_resolver; +#[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; pub use coordinator::DistributedExec; pub use distributed_ext::DistributedExt; @@ -36,9 +33,21 @@ pub use metrics::{ MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, rewrite_distributed_plan_with_metrics, }; -pub use networking::{ - BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, WorkerResolver, - create_worker_client, get_distributed_channel_resolver, get_distributed_worker_resolver, + +#[cfg(any(feature = "integration", test))] +pub mod test_utils; +#[cfg(feature = "grpc")] +pub use protocol::grpc; + +pub use codec::DistributedCodec; +pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; + +/// TODO: do not expose this yet. +pub use protocol::{ + ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, + GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, + WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, + get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, @@ -46,21 +55,11 @@ pub use stage::{ pub use work_unit_feed::{ DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, }; -pub use worker::generated::worker::worker_service_client::WorkerServiceClient; -pub use worker::generated::worker::worker_service_server::WorkerServiceServer; -pub use worker::generated::worker::{GetWorkerInfoRequest, GetWorkerInfoResponse, TaskKey}; pub use worker::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerQueryContext, WorkerSessionBuilder, }; -pub use observability::{ - GetClusterWorkersRequest, GetClusterWorkersResponse, GetTaskProgressRequest, - GetTaskProgressResponse, ObservabilityService, ObservabilityServiceClient, - ObservabilityServiceImpl, ObservabilityServiceServer, PingRequest, PingResponse, TaskProgress, - TaskStatus, WorkerMetrics, -}; - #[cfg(any(feature = "integration", test))] pub use execution_plans::benchmarks::{ LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, ShuffleBench, diff --git a/src/metrics/latency_metric.rs b/src/metrics/latency_metric.rs index b01f7adb..4d415f75 100644 --- a/src/metrics/latency_metric.rs +++ b/src/metrics/latency_metric.rs @@ -244,7 +244,7 @@ pub struct AvgLatencyMetric { } impl AvgLatencyMetric { - pub(crate) fn from_raw(nanos_sum: usize, count: usize) -> Self { + pub fn from_raw(nanos_sum: usize, count: usize) -> Self { Self { nanos_sum: Arc::new(AtomicUsize::new(nanos_sum)), count: Arc::new(AtomicUsize::new(count)), @@ -255,11 +255,11 @@ impl AvgLatencyMetric { self.nanos_sum.load(Relaxed) / self.count.load(Relaxed).max(1) } - pub(crate) fn nanos_sum(&self) -> usize { + pub fn nanos_sum(&self) -> usize { self.nanos_sum.load(Relaxed) } - pub(crate) fn count(&self) -> usize { + pub fn count(&self) -> usize { self.count.load(Relaxed) } @@ -410,7 +410,7 @@ macro_rules! percentile_latency_metric { } impl $name { - pub(crate) fn from_sketch(sketch: DDSketch) -> Self { + pub fn from_sketch(sketch: DDSketch) -> Self { Self { inner: Arc::new(Mutex::new(sketch)), } @@ -421,7 +421,7 @@ macro_rules! percentile_latency_metric { sketch.quantile($percentile).unwrap_or(None).unwrap_or(0.0) as usize } - pub(crate) fn serialize_sketch(&self) -> Result> { + pub fn serialize_sketch(&self) -> Result> { let sketch = self.inner.lock().unwrap(); bincode::serialize(&*sketch).map_err(|e| { datafusion::error::DataFusionError::Internal(format!( @@ -430,7 +430,7 @@ macro_rules! percentile_latency_metric { }) } - pub(crate) fn count(&self) -> usize { + pub fn count(&self) -> usize { let sketch = self.inner.lock().unwrap(); sketch.count() as usize } diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs index 6e88c231..fd113942 100644 --- a/src/metrics/mod.rs +++ b/src/metrics/mod.rs @@ -1,7 +1,6 @@ mod bytes_metric; mod latency_metric; mod max_gauge_metric; -pub(crate) mod proto; mod task_metrics_collector; mod task_metrics_rewriter; diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index 553db1cc..2d78998d 100644 --- a/src/metrics/task_metrics_collector.rs +++ b/src/metrics/task_metrics_collector.rs @@ -186,7 +186,7 @@ mod tests { // empty metrics (e.g., custom execution plans without metrics), which is fine - we // just verify that a metrics set exists for each node. The count assertion above // ensures all nodes are included in the metrics collection. - let stage = stages.get(&(expected_task_key.stage_id as usize)).unwrap(); + let stage = stages.get(&expected_task_key.stage_id).unwrap(); let stage_plan = stage.local_plan().unwrap(); assert_eq!( actual_metrics.pre_order_plan_metrics.len(), @@ -306,7 +306,7 @@ mod tests { sent metrics via the coordinator channel." ) }); - let stage = stages.get(&(expected_task_key.stage_id as usize)).unwrap(); + let stage = stages.get(&expected_task_key.stage_id).unwrap(); let stage_plan = stage.local_plan().unwrap(); assert_eq!( actual_metrics.pre_order_plan_metrics.len(), diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 87d93c96..e3c39530 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -1,13 +1,11 @@ -use crate::DistributedTaskContext; -use crate::common::{TreeNodeExt, serialize_uuid}; +use crate::common::TreeNodeExt; use crate::coordinator::{DistributedExec, MetricsStore}; use crate::distributed_planner::NetworkBoundaryExt; use crate::execution_plans::MetricsWrapperExec; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; use crate::metrics::collect_plan_metrics; -use crate::metrics::proto::metrics_set_proto_to_df; use crate::stage::{LocalStage, Stage}; -use crate::worker::generated::worker::TaskKey; +use crate::{DistributedTaskContext, TaskKey}; use datafusion::common::HashMap; use datafusion::common::plan_err; use datafusion::common::tree_node::Transformed; @@ -227,9 +225,9 @@ pub fn stage_metrics_rewriter( task_count: stage.tasks, }; let task_key = TaskKey { - query_id: serialize_uuid(&stage.query_id), - stage_id: stage.num as u64, - task_number: task_id as u64, + query_id: stage.query_id, + stage_id: stage.num, + task_number: task_id, }; let Some(task_metrics) = metrics_collection.get(&task_key) else { return internal_err!( @@ -248,8 +246,7 @@ pub fn stage_metrics_rewriter( ); } - let node_metrics_protos = task_metrics.pre_order_plan_metrics[per_task_counter].clone(); - let mut node_metrics = metrics_set_proto_to_df(&node_metrics_protos)?; + let mut node_metrics = task_metrics.pre_order_plan_metrics[per_task_counter].clone(); let rewrite_ctx = format.to_rewrite_ctx(task_id as u64); node_metrics = rewrite_ctx.maybe_rewrite_node_metics(node_metrics); @@ -284,41 +281,35 @@ pub fn stage_metrics_rewriter( #[cfg(test)] mod tests { + use crate::DistributedExt; use crate::coordinator::MetricsStore; use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL; - use crate::metrics::proto::{df_metrics_set_to_proto, metrics_set_proto_to_df}; + use crate::metrics::task_metrics_rewriter::MetricsWrapperExec; use crate::metrics::task_metrics_rewriter::{ annotate_metrics_set_with_task_id, stage_metrics_rewriter, }; use crate::metrics::{DistributedMetricsFormat, rewrite_distributed_plan_with_metrics}; + use crate::stage::LocalStage; use crate::test_utils::in_memory_channel_resolver::{ InMemoryChannelResolver, InMemoryWorkerResolver, }; - use crate::test_utils::metrics::make_test_metrics_set_proto_from_seed; + use crate::test_utils::metrics::make_test_metrics_set_from_seed; use crate::test_utils::plans::count_plan_nodes_up_to_network_boundary; use crate::test_utils::session_context::register_temp_parquet_table; - use crate::worker::generated::worker as pb; - use crate::{DistributedExec, SessionStateBuilderExt}; + use crate::{DistributedExec, SessionStateBuilderExt, TaskKey, TaskMetrics}; use datafusion::arrow::array::{Int32Array, StringArray}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::metrics::{Count, Label, Metric, MetricValue, MetricsSet}; - use test_case::test_case; - use datafusion::physical_plan::{ExecutionPlan, collect}; - use itertools::Itertools; - use uuid::Uuid; - - use crate::DistributedExt; - use crate::common::serialize_uuid; - use crate::metrics::task_metrics_rewriter::MetricsWrapperExec; - use crate::stage::LocalStage; - use crate::worker::generated::worker::TaskKey; - use datafusion::physical_plan::empty::EmptyExec; use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; + use itertools::Itertools; use std::sync::Arc; + use test_case::test_case; + use uuid::Uuid; async fn make_test_ctx() -> SessionContext { make_test_ctx_inner(false).await @@ -430,8 +421,10 @@ mod tests { fn metrics_set_eq(a: &MetricsSet, b: &MetricsSet) -> bool { println!("a: {a:?}"); println!("b: {b:?}"); - // Check equality by converting to proto representation. - df_metrics_set_to_proto(a).unwrap() == df_metrics_set_to_proto(b).unwrap() + a.iter().count() == b.iter().count() + && a.iter().zip(b.iter()).all(|(a, b)| { + a.value() == b.value() && a.partition() == b.partition() && a.labels() == b.labels() + }) } /// Asserts that we successfully re-write the metrics of a plan generated from the provided SQL query. @@ -457,20 +450,20 @@ mod tests { // Generate metrics for each task and store them in the map. let metrics_collection = MetricsStore::from_entries((0..stage.tasks).map(|task_id| { let task_key = TaskKey { - query_id: serialize_uuid(&stage.query_id), - stage_id: stage.num as u64, - task_number: task_id as u64, + query_id: stage.query_id, + stage_id: stage.num, + task_number: task_id, }; let metrics = (0..count_plan_nodes_up_to_network_boundary(&plan)) .map(|node_id| { - make_test_metrics_set_proto_from_seed( + make_test_metrics_set_from_seed( (node_id * task_id) as u64, num_metrics_per_task_per_node, ) }) - .collect::>(); - let task_metrics = pb::TaskMetrics { - task_metrics: None, + .collect::>(); + let task_metrics = TaskMetrics { + task_metrics: MetricsSet::new(), pre_order_plan_metrics: metrics, }; (task_key, task_metrics) @@ -501,9 +494,9 @@ mod tests { { let expected_task_node_metrics = metrics_collection .get(&TaskKey { - query_id: serialize_uuid(&stage.query_id), - stage_id: stage.num as u64, - task_number: task_id as u64, + query_id: stage.query_id, + stage_id: stage.num, + task_number: task_id, }) .unwrap() .pre_order_plan_metrics[node_id] @@ -513,9 +506,7 @@ mod tests { actual_task_node_metrics_set .for_each(|metric| actual_metrics_set.push(metric.clone())); - // Convert from proto to check for equality. - let mut expected_metrics_set = - metrics_set_proto_to_df(&expected_task_node_metrics).unwrap(); + let mut expected_metrics_set = expected_task_node_metrics; if format == DistributedMetricsFormat::PerTask { // Add task ids labels. We expect the actual metrics to be annotated by the diff --git a/src/networking/mod.rs b/src/networking/mod.rs deleted file mode 100644 index 6bab9ae0..00000000 --- a/src/networking/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -mod channel_resolver; -mod worker_resolver; - -pub use channel_resolver::{ - BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client, - get_distributed_channel_resolver, -}; -pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; - -pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; -pub(crate) use worker_resolver::{WorkerResolverExtension, set_distributed_worker_resolver}; diff --git a/src/protobuf/producer_head.rs b/src/protobuf/producer_head.rs deleted file mode 100644 index d4cbfc42..00000000 --- a/src/protobuf/producer_head.rs +++ /dev/null @@ -1,63 +0,0 @@ -use crate::DistributedCodec; -use crate::distributed_planner::ProducerHead; -use crate::worker::generated::worker::{ - BroadcastExecHead, NoneHead, RepartitionExecHead, execute_task_request as pb, -}; -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::Result; -use datafusion::execution::TaskContext; -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; - -impl ProducerHead { - pub(crate) fn from_proto( - proto: pb::ProducerHead, - input_schema: &SchemaRef, - ctx: &Arc, - ) -> Result { - Ok(match proto { - pb::ProducerHead::None(_) => ProducerHead::None, - pb::ProducerHead::Broadcast(v) => ProducerHead::BroadcastExec { - output_partitions: v.output_partitions as usize, - }, - pb::ProducerHead::Repartition(v) => { - let proto_partitioning = protobuf::Partitioning::decode(v.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, - input_schema, - &DefaultPhysicalProtoConverter {}, - )? - .ok_or_else(|| proto_error("Could not parse partitioning"))?; - - ProducerHead::RepartitionExec { partitioning } - } - }) - } - - pub(crate) fn to_proto(&self, ctx: &Arc) -> Result { - Ok(match self { - ProducerHead::None => pb::ProducerHead::None(NoneHead {}), - ProducerHead::BroadcastExec { output_partitions } => { - pb::ProducerHead::Broadcast(BroadcastExecHead { - output_partitions: *output_partitions as u64, - }) - } - ProducerHead::RepartitionExec { partitioning } => { - let codec = DistributedCodec::new_combined_with_user(ctx.session_config()); - let partitioning = - serialize_partitioning(partitioning, &codec, &DefaultPhysicalProtoConverter {}) - .map(|v| v.encode_to_vec())?; - pb::ProducerHead::Repartition(RepartitionExecHead { partitioning }) - } - }) - } -} diff --git a/src/protocol/channel_resolver.rs b/src/protocol/channel_resolver.rs new file mode 100644 index 00000000..ac16fefc --- /dev/null +++ b/src/protocol/channel_resolver.rs @@ -0,0 +1,98 @@ +use crate::config_extension_ext::set_distributed_option_extension; +#[cfg(feature = "grpc")] +use crate::protocol::grpc; +use crate::{DistributedConfig, WorkerChannel}; +use async_trait::async_trait; +use datafusion::common::DataFusionError; +use datafusion::execution::TaskContext; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; +use url::Url; + +/// Allows users to customize the way Worker clients are created. A common use case is to +/// wrap the client with tower layers or schedule it in an IO-specific tokio runtime. +/// +/// There is a default implementation of this trait that should be enough for the most common +/// use-cases. +/// +/// # Implementation Notes +/// - This is called per gRPC request, so implementors of this trait should make sure that +/// clients are reused across method calls instead of building a new Worker client every time. +/// +/// - When implementing `get_worker_client_for_url`, it is recommended to use the +/// [`create_worker_client`] helper function to ensure clients are configured with +/// appropriate message size limits for internal communication. This helps avoid message +/// size errors when transferring large datasets. +#[async_trait] +pub trait ChannelResolver { + /// For a given URL, get a Worker gRPC client for communicating to it. + /// + /// *WARNING*: This method is called for every gRPC request, so to not create + /// one client connection for each request, users are required to reuse generated clients. + /// It's recommended to rely on [DefaultChannelResolver] either by delegating method calls + /// to it or by copying the implementation. + /// + /// Consider using [`create_worker_client`] to create the client with appropriate + /// default message size limits. + async fn get_worker_client_for_url( + &self, + url: &Url, + ) -> Result, DataFusionError>; +} +pub(crate) fn set_distributed_channel_resolver( + cfg: &mut SessionConfig, + channel_resolver: impl ChannelResolver + Send + Sync + 'static, +) { + let opts = cfg.options_mut(); + let channel_resolver = ChannelResolverExtension(Some(Arc::new(channel_resolver))); + if let Some(distributed_cfg) = opts.extensions.get_mut::() { + distributed_cfg.__private_channel_resolver = channel_resolver; + } else { + set_distributed_option_extension( + cfg, + DistributedConfig { + __private_channel_resolver: channel_resolver, + ..Default::default() + }, + ) + } +} + +pub fn get_distributed_channel_resolver( + task_ctx: &TaskContext, +) -> Arc { + let opts = task_ctx.session_config().options(); + if let Some(distributed_cfg) = opts.extensions.get::() + && let Some(cr) = &distributed_cfg.__private_channel_resolver.0 + { + return Arc::clone(cr); + } + + #[cfg(feature = "grpc")] + { + let runtime_addr = Arc::as_ptr(&task_ctx.runtime_env()) as usize; + grpc::DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME.get_with(runtime_addr, || { + Arc::new(grpc::DefaultChannelResolver::default()) + }) + } + + #[cfg(not(feature = "grpc"))] + { + panic!( + "gRPC feature is not enabled, and no channel resolver was provided, so no default ChannelResolver can be provided" + ); + } +} + +#[async_trait] +impl ChannelResolver for Arc { + async fn get_worker_client_for_url( + &self, + url: &Url, + ) -> Result, DataFusionError> { + self.as_ref().get_worker_client_for_url(url).await + } +} + +#[derive(Clone, Default)] +pub(crate) struct ChannelResolverExtension(Option>); diff --git a/src/networking/channel_resolver.rs b/src/protocol/grpc/channel_resolver.rs similarity index 56% rename from src/networking/channel_resolver.rs rename to src/protocol/grpc/channel_resolver.rs index d15c8bef..b545e21b 100644 --- a/src/networking/channel_resolver.rs +++ b/src/protocol/grpc/channel_resolver.rs @@ -1,10 +1,7 @@ -use crate::DistributedConfig; -use crate::config_extension_ext::set_distributed_option_extension; -use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; +use crate::protocol::grpc::worker_client::create_worker_client; +use crate::{ChannelResolver, WorkerChannel}; use async_trait::async_trait; use datafusion::common::{DataFusionError, config_datafusion_err, exec_datafusion_err}; -use datafusion::execution::TaskContext; -use datafusion::prelude::SessionConfig; use futures::FutureExt; use futures::future::Shared; use std::sync::{Arc, LazyLock}; @@ -15,56 +12,6 @@ use tonic::transport::Channel; use tower::ServiceExt; use url::Url; -/// Allows users to customize the way Worker clients are created. A common use case is to -/// wrap the client with tower layers or schedule it in an IO-specific tokio runtime. -/// -/// There is a default implementation of this trait that should be enough for the most common -/// use-cases. -/// -/// # Implementation Notes -/// - This is called per gRPC request, so implementors of this trait should make sure that -/// clients are reused across method calls instead of building a new Worker client every time. -/// -/// - When implementing `get_worker_client_for_url`, it is recommended to use the -/// [`create_worker_client`] helper function to ensure clients are configured with -/// appropriate message size limits for internal communication. This helps avoid message -/// size errors when transferring large datasets. -#[async_trait] -pub trait ChannelResolver { - /// For a given URL, get a Worker gRPC client for communicating to it. - /// - /// *WARNING*: This method is called for every gRPC request, so to not create - /// one client connection for each request, users are required to reuse generated clients. - /// It's recommended to rely on [DefaultChannelResolver] either by delegating method calls - /// to it or by copying the implementation. - /// - /// Consider using [`create_worker_client`] to create the client with appropriate - /// default message size limits. - async fn get_worker_client_for_url( - &self, - url: &Url, - ) -> Result, DataFusionError>; -} - -pub(crate) fn set_distributed_channel_resolver( - cfg: &mut SessionConfig, - channel_resolver: impl ChannelResolver + Send + Sync + 'static, -) { - let opts = cfg.options_mut(); - let channel_resolver = ChannelResolverExtension(Some(Arc::new(channel_resolver))); - if let Some(distributed_cfg) = opts.extensions.get_mut::() { - distributed_cfg.__private_channel_resolver = channel_resolver; - } else { - set_distributed_option_extension( - cfg, - DistributedConfig { - __private_channel_resolver: channel_resolver, - ..Default::default() - }, - ) - } -} - // Unlike TaskContext, a DataFusion RuntimeEnv does not allow to introduce user-defined extensions. // For the default implementation of the ChannelResolvers, we cannot inject one DefaultChannelResolver // per TaskContext, as this holds reference to Tonic channels that must outlive a single TaskContext. @@ -72,27 +19,13 @@ pub(crate) fn set_distributed_channel_resolver( // The Tonic channels need to be established and reused under a whole RuntimeEnv scope, not a single // TaskContext, which forces us to put the default implementation in a static global variable that // stores and reuses tonic channels per RuntimeEnv's pointer address. -static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock< +pub(crate) static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock< moka::sync::Cache< /* Arc pointer address */ usize, /* ChannelResolver that reuses built channels */ Arc, >, > = LazyLock::new(|| moka::sync::Cache::builder().max_capacity(256).build()); -pub fn get_distributed_channel_resolver( - task_ctx: &TaskContext, -) -> Arc { - let opts = task_ctx.session_config().options(); - if let Some(distributed_cfg) = opts.extensions.get::() - && let Some(cr) = &distributed_cfg.__private_channel_resolver.0 - { - return Arc::clone(cr); - } - let runtime_addr = Arc::as_ptr(&task_ctx.runtime_env()) as usize; - DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME - .get_with(runtime_addr, || Arc::new(DefaultChannelResolver::default())) -} - pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService< http::Request, http::Response, @@ -101,9 +34,6 @@ pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService< type ChannelCacheValue = Shared>>; -#[derive(Clone, Default)] -pub(crate) struct ChannelResolverExtension(Option>); - /// Default implementation of a [ChannelResolver] that connects to the workers given the URL once /// and stores the connection instance in a TTI cache. /// @@ -159,8 +89,8 @@ impl DefaultChannelResolver { })?; Ok(BoxCloneSyncChannel::new(channel)) } - .boxed() - .shared() + .boxed() + .shared() }); channel.await.map_err(|err| { @@ -175,56 +105,11 @@ impl ChannelResolver for DefaultChannelResolver { async fn get_worker_client_for_url( &self, url: &Url, - ) -> Result, DataFusionError> { + ) -> Result, DataFusionError> { self.get_channel(url).await.map(create_worker_client) } } -#[async_trait] -impl ChannelResolver for Arc { - async fn get_worker_client_for_url( - &self, - url: &Url, - ) -> Result, DataFusionError> { - self.as_ref().get_worker_client_for_url(url).await - } -} - -/// Creates a [`WorkerServiceClient`] with high default message size limits. -/// -/// This is a convenience function that wraps [`WorkerServiceClient::new`] and configures -/// it with `max_decoding_message_size(usize::MAX)` and `max_encoding_message_size(usize::MAX)` -/// to avoid message size limitations for internal communication. -/// -/// Users implementing custom [`ChannelResolver`]s should use this function in their -/// `get_worker_client_for_url` implementations to ensure consistent behavior with built-in -/// implementations. -/// -/// # Example -/// -/// ```rust,ignore -/// use datafusion_distributed::{create_worker_client, BoxCloneSyncChannel, ChannelResolver}; -/// /// use tonic::transport::Channel; -/// -/// #[async_trait] -/// impl ChannelResolver for MyResolver { -/// async fn get_worker_client_for_url( -/// &self, -/// url: &Url, -/// ) -> Result, DataFusionError> { -/// let channel = Channel::from_shared(url.to_string())?.connect().await?; -/// Ok(create_worker_client(BoxCloneSyncChannel::new(channel))) -/// } -/// } -/// ``` -pub fn create_worker_client( - channel: BoxCloneSyncChannel, -) -> WorkerServiceClient { - WorkerServiceClient::new(channel) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX) -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/protobuf/errors/arrow_error.rs b/src/protocol/grpc/errors/arrow_error.rs similarity index 99% rename from src/protobuf/errors/arrow_error.rs rename to src/protocol/grpc/errors/arrow_error.rs index a8630bba..b78a950c 100644 --- a/src/protobuf/errors/arrow_error.rs +++ b/src/protocol/grpc/errors/arrow_error.rs @@ -1,6 +1,6 @@ use datafusion::arrow::error::ArrowError; -use crate::protobuf::errors::io_error::IoErrorProto; +use super::io_error::IoErrorProto; #[derive(Clone, PartialEq, ::prost::Message)] pub struct ArrowErrorProto { diff --git a/src/protobuf/errors/datafusion_error.rs b/src/protocol/grpc/errors/datafusion_error.rs similarity index 97% rename from src/protobuf/errors/datafusion_error.rs rename to src/protocol/grpc/errors/datafusion_error.rs index 1f5caede..71d41730 100644 --- a/src/protobuf/errors/datafusion_error.rs +++ b/src/protocol/grpc/errors/datafusion_error.rs @@ -1,9 +1,9 @@ -use crate::protobuf::errors::arrow_error::ArrowErrorProto; -use crate::protobuf::errors::io_error::IoErrorProto; -use crate::protobuf::errors::objectstore_error::ObjectStoreErrorProto; -use crate::protobuf::errors::parquet_error::ParquetErrorProto; -use crate::protobuf::errors::parser_error::ParserErrorProto; -use crate::protobuf::errors::schema_error::SchemaErrorProto; +use super::arrow_error::ArrowErrorProto; +use super::io_error::IoErrorProto; +use super::objectstore_error::ObjectStoreErrorProto; +use super::parquet_error::ParquetErrorProto; +use super::parser_error::ParserErrorProto; +use super::schema_error::SchemaErrorProto; use datafusion::common::{DataFusionError, Diagnostic}; use datafusion::logical_expr::sqlparser::parser::ParserError; use std::error::Error; diff --git a/src/protobuf/errors/io_error.rs b/src/protocol/grpc/errors/io_error.rs similarity index 100% rename from src/protobuf/errors/io_error.rs rename to src/protocol/grpc/errors/io_error.rs diff --git a/src/protobuf/errors/mod.rs b/src/protocol/grpc/errors/mod.rs similarity index 97% rename from src/protobuf/errors/mod.rs rename to src/protocol/grpc/errors/mod.rs index 0393964e..b940d7c5 100644 --- a/src/protobuf/errors/mod.rs +++ b/src/protocol/grpc/errors/mod.rs @@ -6,7 +6,7 @@ use datafusion::error::DataFusionError; use prost::Message; use std::borrow::Borrow; -use crate::protobuf::errors::datafusion_error::DataFusionErrorProto; +use super::errors::datafusion_error::DataFusionErrorProto; mod arrow_error; mod datafusion_error; diff --git a/src/protobuf/errors/objectstore_error.rs b/src/protocol/grpc/errors/objectstore_error.rs similarity index 100% rename from src/protobuf/errors/objectstore_error.rs rename to src/protocol/grpc/errors/objectstore_error.rs diff --git a/src/protobuf/errors/parquet_error.rs b/src/protocol/grpc/errors/parquet_error.rs similarity index 100% rename from src/protobuf/errors/parquet_error.rs rename to src/protocol/grpc/errors/parquet_error.rs diff --git a/src/protobuf/errors/parser_error.rs b/src/protocol/grpc/errors/parser_error.rs similarity index 100% rename from src/protobuf/errors/parser_error.rs rename to src/protocol/grpc/errors/parser_error.rs diff --git a/src/protobuf/errors/schema_error.rs b/src/protocol/grpc/errors/schema_error.rs similarity index 100% rename from src/protobuf/errors/schema_error.rs rename to src/protocol/grpc/errors/schema_error.rs diff --git a/src/worker/gen/Cargo.toml b/src/protocol/grpc/gen/Cargo.toml similarity index 100% rename from src/worker/gen/Cargo.toml rename to src/protocol/grpc/gen/Cargo.toml diff --git a/src/observability/gen/regen.sh b/src/protocol/grpc/gen/regen.sh similarity index 57% rename from src/observability/gen/regen.sh rename to src/protocol/grpc/gen/regen.sh index d5d48ef8..1a1d85ea 100755 --- a/src/observability/gen/regen.sh +++ b/src/protocol/grpc/gen/regen.sh @@ -3,4 +3,4 @@ set -e repo_root=$(git rev-parse --show-toplevel) -cd "$repo_root" && cargo run --manifest-path src/observability/gen/Cargo.toml +cd "$repo_root" && cargo run --manifest-path src/protocol/grpc/gen/Cargo.toml diff --git a/src/worker/gen/src/main.rs b/src/protocol/grpc/gen/src/main.rs similarity index 87% rename from src/worker/gen/src/main.rs rename to src/protocol/grpc/gen/src/main.rs index 8994d54d..7505aae5 100644 --- a/src/worker/gen/src/main.rs +++ b/src/protocol/grpc/gen/src/main.rs @@ -4,9 +4,9 @@ use std::fs; fn main() -> Result<(), Box> { let repo_root = env::current_dir()?; - let proto_dir = repo_root.join("src/worker"); + let proto_dir = repo_root.join("src/protocol/grpc"); let proto_file = proto_dir.join("worker.proto"); - let out_dir = repo_root.join("src/worker/generated"); + let out_dir = repo_root.join("src/protocol/grpc/generated"); fs::create_dir_all(&out_dir)?; diff --git a/src/worker/generated/mod.rs b/src/protocol/grpc/generated/mod.rs similarity index 100% rename from src/worker/generated/mod.rs rename to src/protocol/grpc/generated/mod.rs diff --git a/src/worker/generated/worker.rs b/src/protocol/grpc/generated/worker.rs similarity index 100% rename from src/worker/generated/worker.rs rename to src/protocol/grpc/generated/worker.rs diff --git a/src/metrics/proto.rs b/src/protocol/grpc/metrics_proto.rs similarity index 98% rename from src/metrics/proto.rs rename to src/protocol/grpc/metrics_proto.rs index c114a41c..6db8cda9 100644 --- a/src/metrics/proto.rs +++ b/src/protocol/grpc/metrics_proto.rs @@ -1,3 +1,4 @@ +use super::generated::worker as pb; use chrono::DateTime; use datafusion::common::internal_err; use datafusion::error::DataFusionError; @@ -8,15 +9,12 @@ use sketches_ddsketch::DDSketch; use std::borrow::Cow; use std::sync::Arc; -use super::bytes_metric::BytesCounterMetric; -use super::latency_metric::{ - AvgLatencyMetric, FirstLatencyMetric, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, - P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, +use crate::{ + AvgLatencyMetric, BytesCounterMetric, FirstLatencyMetric, MaxGaugeMetric, MaxLatencyMetric, + MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, }; -use crate::MaxGaugeMetric; -use crate::worker::generated::worker as pb; -/// df_metrics_set_to_proto converts a [datafusion::physical_plan::metrics::MetricsSet] to a [pb::MetricsSet]. +/// df_metrics_set_to_proto converts a [MetricsSet] to a [pb::MetricsSet]. /// Custom metrics are filtered out, but any other errors are returned. /// TODO(#140): Support custom metrics. pub fn df_metrics_set_to_proto( @@ -44,7 +42,7 @@ pub fn df_metrics_set_to_proto( Ok(pb::MetricsSet { metrics }) } -/// metrics_set_proto_to_df converts a [pb::MetricsSet] to a [datafusion::physical_plan::metrics::MetricsSet]. +/// metrics_set_proto_to_df converts a [pb::MetricsSet] to a [MetricsSet]. pub fn metrics_set_proto_to_df( metrics_set_proto: &pb::MetricsSet, ) -> Result { @@ -64,7 +62,7 @@ const CUSTOM_METRICS_NOT_SUPPORTED: &str = /// New DataFusion metrics that are not yet supported in proto conversion. const UNSUPPORTED_METRICS: &str = "metric type not supported in proto conversion"; -/// df_metric_to_proto converts a `datafusion::physical_plan::metrics::Metric` to a `pb::Metric`. It does not consume the Arc. +/// df_metric_to_proto converts a `Metric` to a `pb::Metric`. It does not consume the Arc. pub fn df_metric_to_proto(metric: Arc) -> Result { let partition = metric.partition().map(|p| p as u64); let labels = metric @@ -284,7 +282,7 @@ pub fn df_metric_to_proto(metric: Arc) -> Result Result, DataFusionError> { let partition = metric.partition.map(|p| p as usize); let labels = metric diff --git a/src/protocol/grpc/mod.rs b/src/protocol/grpc/mod.rs new file mode 100644 index 00000000..e432b607 --- /dev/null +++ b/src/protocol/grpc/mod.rs @@ -0,0 +1,21 @@ +mod channel_resolver; +mod errors; +mod generated; +mod metrics_proto; +mod observability; +mod on_drop_stream; +mod spawn_select_all; +mod worker_client; +mod worker_service; + +// TODO: this should not be exposed. +pub(crate) use channel_resolver::DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME; + +pub use channel_resolver::{BoxCloneSyncChannel, DefaultChannelResolver}; +pub use observability::{ + GetClusterWorkersRequest, GetClusterWorkersResponse, GetTaskProgressRequest, + GetTaskProgressResponse, ObservabilityService, ObservabilityServiceClient, + ObservabilityServiceImpl, ObservabilityServiceServer, PingRequest, PingResponse, TaskProgress, + TaskStatus, WorkerMetrics, +}; +pub use worker_client::create_worker_client; diff --git a/src/observability/README.md b/src/protocol/grpc/observability/README.md similarity index 88% rename from src/observability/README.md rename to src/protocol/grpc/observability/README.md index 73dd5389..94d9b840 100644 --- a/src/observability/README.md +++ b/src/protocol/grpc/observability/README.md @@ -8,5 +8,5 @@ for tonic/gRPC services. In the root of the datafusion-distribued repo, run: ```bash -./src/observability/gen/regen.sh +./src/protocol/grpc/observability/gen/regen.sh ``` diff --git a/src/observability/gen/Cargo.toml b/src/protocol/grpc/observability/gen/Cargo.toml similarity index 100% rename from src/observability/gen/Cargo.toml rename to src/protocol/grpc/observability/gen/Cargo.toml diff --git a/src/protocol/grpc/observability/gen/regen.sh b/src/protocol/grpc/observability/gen/regen.sh new file mode 100755 index 00000000..c95505eb --- /dev/null +++ b/src/protocol/grpc/observability/gen/regen.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash + +set -e + +repo_root=$(git rev-parse --show-toplevel) +cd "$repo_root" && cargo run --manifest-path src/protocol/grpc/observability/gen/Cargo.toml diff --git a/src/observability/gen/src/main.rs b/src/protocol/grpc/observability/gen/src/main.rs similarity index 76% rename from src/observability/gen/src/main.rs rename to src/protocol/grpc/observability/gen/src/main.rs index d414fc53..124f0f1b 100644 --- a/src/observability/gen/src/main.rs +++ b/src/protocol/grpc/observability/gen/src/main.rs @@ -4,9 +4,9 @@ use std::fs; fn main() -> Result<(), Box> { let repo_root = env::current_dir()?; - let proto_dir = repo_root.join("src/observability/proto"); + let proto_dir = repo_root.join("src/protocol/grpc/observability/proto"); let proto_file = proto_dir.join("observability.proto"); - let out_dir = repo_root.join("src/observability/generated"); + let out_dir = repo_root.join("src/protocol/grpc/observability/generated"); fs::create_dir_all(&out_dir)?; @@ -21,7 +21,7 @@ fn main() -> Result<(), Box> { .out_dir(&out_dir) .extern_path( ".observability.TaskKey", - "crate::worker::generated::worker::TaskKey", + "crate::protocol::grpc::generated::worker::TaskKey", ) .compile_protos(&[proto_file], &[proto_dir])?; diff --git a/src/observability/generated/mod.rs b/src/protocol/grpc/observability/generated/mod.rs similarity index 100% rename from src/observability/generated/mod.rs rename to src/protocol/grpc/observability/generated/mod.rs diff --git a/src/observability/generated/observability.rs b/src/protocol/grpc/observability/generated/observability.rs similarity index 99% rename from src/observability/generated/observability.rs rename to src/protocol/grpc/observability/generated/observability.rs index c8da3372..7706dd21 100644 --- a/src/observability/generated/observability.rs +++ b/src/protocol/grpc/observability/generated/observability.rs @@ -12,7 +12,7 @@ pub struct GetTaskProgressRequest {} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct TaskProgress { #[prost(message, optional, tag = "1")] - pub task_key: ::core::option::Option, + pub task_key: ::core::option::Option, #[prost(uint64, tag = "2")] pub total_partitions: u64, #[prost(uint64, tag = "3")] diff --git a/src/observability/mod.rs b/src/protocol/grpc/observability/mod.rs similarity index 100% rename from src/observability/mod.rs rename to src/protocol/grpc/observability/mod.rs diff --git a/src/observability/proto/observability.proto b/src/protocol/grpc/observability/proto/observability.proto similarity index 100% rename from src/observability/proto/observability.proto rename to src/protocol/grpc/observability/proto/observability.proto diff --git a/src/observability/service.rs b/src/protocol/grpc/observability/service.rs similarity index 80% rename from src/observability/service.rs rename to src/protocol/grpc/observability/service.rs index 7b7c5473..b7cc2120 100644 --- a/src/observability/service.rs +++ b/src/protocol/grpc/observability/service.rs @@ -2,13 +2,16 @@ use super::{ GetTaskProgressResponse, ObservabilityService, TaskProgress, TaskStatus, WorkerMetrics, generated::observability::{GetTaskProgressRequest, PingRequest, PingResponse}, }; -use crate::worker::generated::worker::TaskKey; +use crate::common::serialize_uuid; +use crate::grpc::{GetClusterWorkersRequest, GetClusterWorkersResponse}; +use crate::protocol::grpc::generated::worker as worker_pb; use crate::worker::{SingleWriteMultiRead, TaskData}; -use crate::{GetClusterWorkersRequest, GetClusterWorkersResponse, WorkerResolver}; +use crate::{TaskKey, WorkerResolver}; use datafusion::error::DataFusionError; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use moka::future::Cache; use std::sync::Arc; +use std::sync::atomic::Ordering; #[cfg(feature = "system-metrics")] use std::time::Duration; #[cfg(feature = "system-metrics")] @@ -99,7 +102,7 @@ impl ObservabilityService for ObservabilityServiceImpl { let output_rows = output_rows_from_plan(&task_data.base_plan); tasks.push(TaskProgress { - task_key: Some((*internal_key).clone()), + task_key: Some(task_key_to_proto(&internal_key)), total_partitions, completed_partitions, status: TaskStatus::Running as i32, @@ -139,7 +142,7 @@ impl ObservabilityServiceImpl { } #[cfg(feature = "system-metrics")] - return *self.system.borrow(); + *self.system.borrow() } } @@ -147,3 +150,30 @@ impl ObservabilityServiceImpl { fn output_rows_from_plan(plan: &Arc) -> u64 { plan.metrics().and_then(|m| m.output_rows()).unwrap_or(0) as u64 } + +fn task_key_to_proto(task_key: &TaskKey) -> worker_pb::TaskKey { + worker_pb::TaskKey { + query_id: serialize_uuid(&task_key.query_id), + stage_id: task_key.stage_id as u64, + task_number: task_key.task_number as u64, + } +} + +impl TaskData { + /// Returns the number of partitions remaining to be processed. + fn num_partitions_remaining(&self) -> usize { + self.num_partitions_remaining.load(Ordering::SeqCst) + } + + /// Returns the total number of partitions in this task. + fn total_partitions(&self) -> usize { + match self.final_plan.get() { + Some(Ok(plan)) => plan.output_partitioning().partition_count(), + _ => self + .base_plan + .properties() + .output_partitioning() + .partition_count(), + } + } +} diff --git a/src/common/on_drop_stream.rs b/src/protocol/grpc/on_drop_stream.rs similarity index 100% rename from src/common/on_drop_stream.rs rename to src/protocol/grpc/on_drop_stream.rs diff --git a/src/worker/spawn_select_all.rs b/src/protocol/grpc/spawn_select_all.rs similarity index 100% rename from src/worker/spawn_select_all.rs rename to src/protocol/grpc/spawn_select_all.rs diff --git a/src/worker/worker.proto b/src/protocol/grpc/worker.proto similarity index 100% rename from src/worker/worker.proto rename to src/protocol/grpc/worker.proto diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs new file mode 100644 index 00000000..2af90c59 --- /dev/null +++ b/src/protocol/grpc/worker_client.rs @@ -0,0 +1,766 @@ +use super::channel_resolver::BoxCloneSyncChannel; +use super::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; +use super::generated::worker as pb; +use super::metrics_proto::metrics_set_proto_to_df; +use crate::common::serialize_uuid; +use crate::grpc::generated::worker::FlightAppMetadata; +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, +}; +use arrow_flight::FlightData; +use arrow_flight::decode::FlightRecordBatchStream; +use arrow_flight::error::FlightError; +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::instant::Instant; +use datafusion::common::runtime::SpawnedTask; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::physical_expr_common::metrics::{Count, MetricBuilder, MetricValue, Time}; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::stream::BoxStream; +use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; +use http::{Extensions, HeaderMap}; +use pin_project::{pin_project, pinned_drop}; +use prost::Message; +use std::borrow::Cow; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::sync::Notify; +use tokio::sync::mpsc::UnboundedSender; +use tokio_stream::wrappers::UnboundedReceiverStream; +use tokio_util::sync::CancellationToken; +use tonic::metadata::MetadataMap; +use tonic::{Request, Status}; + +#[async_trait] +impl WorkerChannel for pb::worker_service_client::WorkerServiceClient { + async fn coordinator_channel( + &mut self, + headers: HeaderMap, + c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>> { + let input_stream = c2w_stream.map(encode_coordinator_to_worker_msg); + + let output_stream = self + .coordinator_channel(Request::from_parts( + MetadataMap::from_headers(headers), + Extensions::default(), + input_stream, + )) + .boxed() + .await + .map_err(map_status_to_datafusion_error)? + .into_inner() + .map_err(map_status_to_datafusion_error) + .map(|msg| decode_worker_to_coordinator_msg(msg?)) + .boxed(); + + Ok(output_stream) + } + + async fn execute_task( + &mut self, + headers: HeaderMap, + request: ExecuteTaskRequest, + metrics: ExecutionPlanMetricsSet, + ctx: &Arc, + ) -> Result>>> { + let d_cfg = DistributedConfig::from_session_config(ctx.session_config())?; + let buffer_budget_bytes = d_cfg.worker_connection_buffer_budget_bytes; + + // We are retaining record batches in memory until they are consumed, so we need to account + // for them in the memory pool. + let memory_reservation = + Arc::new(MemoryConsumer::new("WorkerConnection").register(ctx.memory_pool())); + let memory_reservation_clone = Arc::clone(&memory_reservation); + + // Track the maximum memory used to buffer recieved messages. + let mut curr_max_mem = 0; + let max_mem_used = MetricBuilder::new(&metrics).global_gauge("max_mem_used"); + // Track the total encoded size of all recieved messages. + let bytes_transferred = MetricBuilder::new(&metrics).bytes_counter("bytes_transferred"); + let msg_count = MetricBuilder::new(&metrics).global_counter("msg_count"); + // Track end-to-end network latency distribution for messages that actually arrive. + let mut latency_metrics = NetworkLatencyMetrics::new(&metrics); + // Track the total CPU time spent in polling messages over the network + decoding them. + let elapsed_compute = Time::new(); + let elapsed_compute_clone = elapsed_compute.clone(); + MetricBuilder::new(&metrics).build(MetricValue::ElapsedCompute(elapsed_compute.clone())); + + let target_partition_range = request.target_partition_start..request.target_partition_end; + let request = pb::ExecuteTaskRequest { + task_key: Some(encode_task_key(request.task_key)), + target_partition_start: request.target_partition_start as u64, + target_partition_end: request.target_partition_end as u64, + producer_head: Some(encode_producer_head_spec(request.producer_head_spec)), + }; + let metadata = MetadataMap::from_headers(headers); + + // The senders and receivers are unbounded queues used for multiplexing the record + // batches sent through the single gRPC stream into one stream per partition. They + // are unbounded to avoid head-of-line blocking: a single bounded queue could block + // the demux task and starve all sibling partitions even though they have capacity, + // which deadlocks queries with cross-partition dependencies. + // Total memory is bounded globally below via `mem_available_notify`. + let mut per_partition_tx = Vec::with_capacity(target_partition_range.len()); + let mut per_partition_rx = Vec::with_capacity(target_partition_range.len()); + for _partition in target_partition_range.clone() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::(); + per_partition_tx.push(tx); + per_partition_rx.push(rx); + } + + let mem_available_notify = Arc::new(Notify::new()); + let mem_available_notify_for_task = Arc::clone(&mem_available_notify); + + let first_poll_notify = Arc::new(Notify::new()); + let first_poll_notify_for_task = Arc::clone(&first_poll_notify); + + // Cancellation token allows us to stop the background task promptly when all partition + // streams are dropped (e.g., when the query is cancelled). + let cancel_token = CancellationToken::new(); + let cancel = cancel_token.clone(); + + let mut self_clone = self.clone(); + let request_for_task = request.clone(); + let metadata_for_task = metadata.clone(); + + // This task will pull data from all the partitions in `target_partition_range`, and will + // fan them out to the appropriate `per_partition_rx` based on the "partition" declared + // in each individual record batch flight metadata. + let task = SpawnedTask::spawn(async move { + tokio::select! { + biased; + _ = cancel.cancelled() => { + // If all SendableRecordBatchStreams canceled before any poll, we need to + // anyway trigger the task execution and cancel it immediately so that the + // cancellation is propagated also in the remote worker. Otherwise, it might + // hang forever waiting for someone to execute it. + let _ = self_clone.execute_task(Request::from_parts( + metadata_for_task.clone(), + Extensions::default(), + request_for_task.clone(), + )).await; + return + }, + _ = first_poll_notify_for_task.notified() => {} + } + + let request = Request::from_parts( + metadata_for_task, + Extensions::default(), + request_for_task, + ); + let mut interleaved_stream = match self_clone.execute_task(request).await { + Ok(v) => v.into_inner(), + Err(err) => return fanout(&per_partition_tx, err), + }; + + loop { + // Backpressure gate. Per-partition channels are unbounded, so we cap + // total in-flight buffered bytes here by pausing the gRPC pull when + // consumers haven't drained enough. This propagates flow control all + // the way back to the worker without coupling sibling partitions. + // We always allow a message through when reservation == 0 to avoid + // livelock if a single message is larger than the budget. + while memory_reservation.size() >= buffer_budget_bytes { + tokio::select! { + biased; + _ = cancel.cancelled() => return, + _ = mem_available_notify_for_task.notified() => {} + } + } + + // Check for cancellation while waiting for the next message. + let flight_data = tokio::select! { + biased; + _ = cancel.cancelled() => return, + msg = interleaved_stream.next() => { + match msg { + Some(Ok(v)) => v, + Some(Err(err)) => return fanout(&per_partition_tx, err), + None => return, // Stream exhausted + } + } + }; + + // Earliest time at which the msg was received. + let msg_received_time = SystemTime::now(); + + let flight_metadata = match FlightAppMetadata::decode(flight_data.app_metadata.as_ref()) { + Ok(v) => v, + Err(err) => { + return fanout(&per_partition_tx, Status::internal(err.to_string())); + } + }; + + // Update the running latency tracker. + let sent_time = UNIX_EPOCH + Duration::from_nanos(flight_metadata.created_timestamp_unix_nanos); + if flight_metadata.created_timestamp_unix_nanos > 0 + && let Ok(delta) = msg_received_time.duration_since(sent_time) { + latency_metrics.add_duration(delta); + } + + let partition = flight_metadata.partition as usize; + // the `per_partition_tx` variable is using a normal `Vec` for storing the + // channel transmitters, so we need to subtract the `target_partition_range.start` + // to the `partition` in order to offset it to the appropriate index. + let Some(sender_i) = partition.checked_sub(target_partition_range.start) else { + let msg = format!( + "Received partition {partition} in Flight metadata, but available partitions are {target_partition_range:?}" + ); + return fanout(&per_partition_tx, Status::internal(msg)); + }; + + let Some(o_tx) = per_partition_tx.get(sender_i) else { + let msg = format!( + "Received partition {partition} in Flight metadata, but available partitions are {target_partition_range:?}" + ); + return fanout(&per_partition_tx, Status::internal(msg)); + }; + + // We need to send the memory reservation in the same tuple as the actual message + // so that it gets dropped as soon as the message leaves the queue. Dropping the + // memory reservation means releasing the memory from the pool for that specific + // message + let size = flight_data.encoded_len(); + memory_reservation.grow(size); + + // Update memory related metrics. + msg_count.add(1); + bytes_transferred.add_bytes(size); + let curr_mem = memory_reservation.size(); + if curr_mem > curr_max_mem { + curr_max_mem = curr_mem; + max_mem_used.set(curr_max_mem); + } + + if o_tx.send(Ok((flight_data, flight_metadata))).is_err() { + // The receiver for this partition was dropped (e.g. a hash join partition + // completed early without consuming its probe side). Don't exit: other + // partitions multiplexed over the same gRPC stream still need their data. + // Undo the memory reservation that was grown for this dropped batch. + memory_reservation.shrink(size); + continue; + }; + } + }.with_elapsed_compute(elapsed_compute)); + + let task = Arc::new(task); + let not_consumed_streams = Arc::new(AtomicUsize::new(per_partition_rx.len())); + + let mut result = Vec::with_capacity(per_partition_rx.len()); + for partition_receiver in per_partition_rx { + let task = Arc::clone(&task); + let cancel_token = cancel_token.clone(); + + let first_poll_notify = Arc::clone(&first_poll_notify); + let stream = async move { + first_poll_notify.notify_one(); + UnboundedReceiverStream::new(partition_receiver) + } + .flatten_stream(); + + let stream = stream.map_err(|err| FlightError::Tonic(Box::new(err))); + let reservation = Arc::clone(&memory_reservation_clone); + let mem_available_notify = Arc::clone(&mem_available_notify); + let stream = stream.map_ok(move |(data, _meta)| { + reservation.shrink(data.encoded_len()); + // Wake the demux task in case it is blocked on the byte budget. + mem_available_notify.notify_one(); + let _ = &task; // <- keep the task that polls data from the network alive. + data + }); + let stream = FlightRecordBatchStream::new_from_flight_data(stream); + let stream = stream.map_err(map_flight_to_datafusion_error); + let stream = stream.with_elapsed_compute(elapsed_compute_clone.clone()); + + // When the stream is dropped, cancel the background task to ensure prompt cleanup. + let not_consumed_streams = Arc::clone(¬_consumed_streams); + result.push( + on_drop_stream(stream, move || { + let remaining_streams = not_consumed_streams.fetch_sub(1, Ordering::SeqCst) - 1; + if remaining_streams == 0 { + cancel_token.cancel(); + } + }) + .boxed(), + ); + } + + Ok(result) + } + + async fn get_worker_info( + &mut self, + _request: GetWorkerInfoRequest, + ) -> Result { + let response = self + .get_worker_info(pb::GetWorkerInfoRequest {}) + .await + .map_err(map_status_to_datafusion_error)? + .into_inner(); + Ok(GetWorkerInfoResponse { + version: response.version, + }) + } +} + +type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; + +struct NetworkLatencyMetrics { + metrics: ExecutionPlanMetricsSet, + values: Option, +} + +impl NetworkLatencyMetrics { + fn new(metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + metrics: metrics.clone(), + values: None, + } + } + + fn add_duration(&mut self, duration: Duration) { + self.values + .get_or_insert_with(|| NetworkLatencyMetricValues::new(&self.metrics)) + .add_duration(duration); + } +} + +struct NetworkLatencyMetricValues { + min_latency: MinLatencyMetric, + max_latency: MaxLatencyMetric, + p50_latency: P50LatencyMetric, + p95_latency: P95LatencyMetric, + first_latency: FirstLatencyMetric, + sum_latency: Time, + latency_count: Count, +} + +impl NetworkLatencyMetricValues { + fn new(metrics: &ExecutionPlanMetricsSet) -> Self { + let min_latency = MetricBuilder::new(metrics).min_latency("network_latency_min"); + let max_latency = MetricBuilder::new(metrics).max_latency("network_latency_max"); + let p50_latency = MetricBuilder::new(metrics).p50_latency("network_latency_p50"); + let p95_latency = MetricBuilder::new(metrics).p95_latency("network_latency_p95"); + let first_latency = MetricBuilder::new(metrics).first_latency("network_latency_first"); + let sum_latency = Time::new(); + MetricBuilder::new(metrics).build(MetricValue::Time { + name: Cow::Borrowed("network_latency_sum"), + time: sum_latency.clone(), + }); + let latency_count = MetricBuilder::new(metrics).counter("network_latency_count", 0); + + Self { + min_latency, + max_latency, + p50_latency, + p95_latency, + first_latency, + sum_latency, + latency_count, + } + } + + fn add_duration(&self, duration: Duration) { + self.min_latency.add_duration(duration); + self.max_latency.add_duration(duration); + self.p50_latency.add_duration(duration); + self.p95_latency.add_duration(duration); + self.first_latency.add_duration(duration); + self.sum_latency.add_duration(duration); + self.latency_count.add(1); + } +} + +pub(super) fn encode_producer_head_spec( + head: ProducerHeadSpec, +) -> pb::execute_task_request::ProducerHead { + match head { + ProducerHeadSpec::None => pb::execute_task_request::ProducerHead::None(pb::NoneHead {}), + ProducerHeadSpec::BroadcastExec { output_partitions } => { + pb::execute_task_request::ProducerHead::Broadcast(pb::BroadcastExecHead { + output_partitions: output_partitions as u64, + }) + } + ProducerHeadSpec::RepartitionExec { partitioning } => { + pb::execute_task_request::ProducerHead::Repartition(pb::RepartitionExecHead { + partitioning, + }) + } + } +} + +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)) + } + CoordinatorToWorkerMsg::WorkUnitEos => { + pb::coordinator_to_worker_msg::Inner::WorkUnitEos(true) + } + }), + } +} + +fn encode_set_plan_request(request: SetPlanRequest) -> pb::SetPlanRequest { + pb::SetPlanRequest { + task_key: Some(encode_task_key(request.task_key)), + task_count: request.task_count as u64, + plan_proto: request.plan_proto, + work_unit_feed_declarations: request + .work_unit_feed_declarations + .into_iter() + .map(encode_work_unit_feed_declaration) + .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 { + pb::WorkUnitBatch { + batch: batch.batch.into_iter().map(encode_work_unit).collect(), + } +} + +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, + 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, + processed_timestamp_unix_nanos: work_unit.processed_timestamp_unix_nanos as u64, + } +} + +fn encode_work_unit_feed_declaration( + declaration: WorkUnitFeedDeclaration, +) -> pb::set_plan_request::WorkUnitFeedDeclaration { + pb::set_plan_request::WorkUnitFeedDeclaration { + id: serialize_uuid(&declaration.id), + partitions: declaration.partitions as u64, + } +} + +fn encode_task_key(task_key: TaskKey) -> pb::TaskKey { + pb::TaskKey { + query_id: serialize_uuid(&task_key.query_id), + stage_id: task_key.stage_id as u64, + task_number: task_key.task_number as u64, + } +} + +fn decode_worker_to_coordinator_msg( + msg: pb::WorkerToCoordinatorMsg, +) -> Result { + Ok( + match msg + .inner + .ok_or_else(|| missing("WorkerToCoordinatorMsg.inner"))? + { + pb::worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics) => { + WorkerToCoordinatorMsg::TaskMetrics(decode_task_metrics(task_metrics)?) + } + pb::worker_to_coordinator_msg::Inner::LoadInfo(load_info) => { + WorkerToCoordinatorMsg::LoadInfo(decode_load_info(load_info)) + } + pb::worker_to_coordinator_msg::Inner::LoadInfoEos(_) => { + WorkerToCoordinatorMsg::LoadInfoEos + } + }, + ) +} + +fn decode_task_metrics(task_metrics: pb::TaskMetrics) -> Result { + Ok(TaskMetrics { + pre_order_plan_metrics: task_metrics + .pre_order_plan_metrics + .into_iter() + .map(|metrics_set| metrics_set_proto_to_df(&metrics_set)) + .collect::>()?, + task_metrics: metrics_set_proto_to_df( + &task_metrics + .task_metrics + .ok_or_else(|| missing("task_metrics"))?, + )?, + }) +} + +fn decode_load_info(load_info: pb::LoadInfo) -> LoadInfo { + LoadInfo { + partition: load_info.partition as usize, + rows_ready: load_info.rows_ready as usize, + rows_per_second: load_info.rows_per_second as usize, + per_column_bytes_ready: load_info + .per_column_bytes_ready + .into_iter() + .map(|bytes| bytes as usize) + .collect(), + per_column_bytes_per_second: load_info + .per_column_bytes_per_second + .into_iter() + .map(|bytes| bytes as usize) + .collect(), + per_column_ndv_percentage: load_info.per_column_ndv_percentage, + per_column_null_percentage: load_info.per_column_null_percentage, + } +} + +fn missing(field: &'static str) -> DataFusionError { + DataFusionError::Internal(format!("Missing field '{field}'")) +} + +fn fanout(o_txs: &[UnboundedSender], err: Status) { + for o_tx in o_txs { + let _ = o_tx.send(Err(err.clone())); + } +} + +/// Creates a [`WorkerServiceClient`] with high default message size limits. +/// +/// This is a convenience function that wraps [`WorkerServiceClient::new`] and configures +/// it with `max_decoding_message_size(usize::MAX)` and `max_encoding_message_size(usize::MAX)` +/// to avoid message size limitations for internal communication. +/// +/// Users implementing custom [`ChannelResolver`]s should use this function in their +/// `get_worker_client_for_url` implementations to ensure consistent behavior with built-in +/// implementations. +/// +/// # Example +/// +/// ```rust +/// # use datafusion::common::DataFusionError; +/// # use datafusion::error::Result; +/// # use tonic::transport::Channel; +/// # use url::Url; +/// # use datafusion_distributed::{ChannelResolver, WorkerChannel, grpc}; +/// +/// struct MyResolver; +/// +/// #[async_trait::async_trait] +/// impl ChannelResolver for MyResolver { +/// async fn get_worker_client_for_url(&self, url: &Url) -> Result> { +/// let channel = Channel::from_shared(url.to_string()) +/// .map_err(|err| DataFusionError::External(Box::new(err)))? +/// .connect() +/// .await +/// .map_err(|err| DataFusionError::External(Box::new(err)))?; +/// Ok(grpc::create_worker_client(grpc::BoxCloneSyncChannel::new(channel))) +/// } +/// } +/// ``` +pub fn create_worker_client(channel: BoxCloneSyncChannel) -> Box { + Box::new( + pb::worker_service_client::WorkerServiceClient::new(channel) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX), + ) +} + +trait ElapsedComputeFutureExt: Future + Sized { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture; +} + +trait ElapsedComputeStreamExt: Stream + Sized { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream; +} + +impl> ElapsedComputeFutureExt for F { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture { + ElapsedComputeFuture { + inner: self, + curr: Duration::default(), + elapsed_compute, + } + } +} + +impl> ElapsedComputeStreamExt for S { + fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream { + ElapsedComputeStream { + inner: self, + curr: Duration::default(), + elapsed_compute, + } + } +} + +#[pin_project(PinnedDrop)] +struct ElapsedComputeStream { + #[pin] + inner: T, + curr: Duration, + elapsed_compute: Time, +} + +/// Drop implementation that ensures that any accumulated time is properly dumped to the metric +/// in case the stream gets dropped before completion. +#[pinned_drop] +impl PinnedDrop for ElapsedComputeStream { + fn drop(self: Pin<&mut Self>) { + if self.curr > Duration::default() { + let self_projected = self.project(); + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + } + } +} + +impl> Stream for ElapsedComputeStream { + type Item = O; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let self_projected = self.project(); + let start = Instant::now(); + let result = self_projected.inner.poll_next(cx); + *self_projected.curr += start.elapsed(); + if result.is_ready() { + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + *self_projected.curr = Duration::default(); + } + result + } +} + +#[pin_project(PinnedDrop)] +struct ElapsedComputeFuture { + #[pin] + inner: T, + curr: Duration, + elapsed_compute: Time, +} + +/// Drop implementation that ensures that any accumulated time is properly dumped to the metric +/// in case the future gets dropped before completion. +#[pinned_drop] +impl PinnedDrop for ElapsedComputeFuture { + fn drop(self: Pin<&mut Self>) { + if self.curr > Duration::default() { + let self_projected = self.project(); + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + } + } +} + +impl> Future for ElapsedComputeFuture { + type Output = O; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let self_projected = self.project(); + let start = Instant::now(); + let result = self_projected.inner.poll(cx); + *self_projected.curr += start.elapsed(); + if result.is_ready() { + self_projected + .elapsed_compute + .add_duration(*self_projected.curr); + *self_projected.curr = Duration::default(); + } + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::StreamExt; + use futures::stream::unfold; + + #[tokio::test] + async fn elapsed_compute_future() { + async fn cheap() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + + async fn expensive() { + let mut _count = 0f64; + for i in 0..100000 { + tokio::task::yield_now().await; + _count /= i as f64 + } + } + + let cheap_time = Time::new(); + cheap().with_elapsed_compute(cheap_time.clone()).await; + println!("cheap future: {}", cheap_time.value()); + + let expensive_time = Time::new(); + expensive() + .with_elapsed_compute(expensive_time.clone()) + .await; + println!("expensive future: {}", expensive_time.value()); + + assert!(expensive_time.value() > cheap_time.value()); + } + + #[tokio::test] + async fn elapsed_compute_stream() { + fn cheap() -> impl Stream { + unfold(0i64, |state| async move { + if state < 10 { + tokio::time::sleep(Duration::from_micros(10)).await; + Some((state, state + 1)) + } else { + None + } + }) + } + + fn expensive() -> impl Stream { + unfold(0i64, |state| async move { + if state < 10 { + // Simulate expensive computation + let mut _count = 0f64; + for i in 1..100000 { + _count += (i as f64).sqrt(); + } + tokio::task::yield_now().await; + Some((state, state + 1)) + } else { + None + } + }) + } + + let cheap_time = Time::new(); + cheap() + .with_elapsed_compute(cheap_time.clone()) + .collect::>() + .await; + println!("cheap future: {}", cheap_time.value()); + + let expensive_time = Time::new(); + expensive() + .with_elapsed_compute(expensive_time.clone()) + .collect::>() + .await; + println!("expensive future: {}", expensive_time.value()); + + assert!(expensive_time.value() > cheap_time.value()); + } +} diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs new file mode 100644 index 00000000..fa677a10 --- /dev/null +++ b/src/protocol/grpc/worker_service.rs @@ -0,0 +1,422 @@ +use super::errors::{datafusion_error_to_tonic_status, map_status_to_datafusion_error}; +use super::generated::worker as pb; +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, +}; + +use arrow_flight::FlightData; +use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder}; +use arrow_flight::error::FlightError; +use arrow_select::dictionary::garbage_collect_any_dictionary; +use async_trait::async_trait; +use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; +use datafusion::arrow::ipc::CompressionType; +use datafusion::arrow::ipc::writer::IpcWriteOptions; +use datafusion::common::DataFusionError; +use datafusion::execution::SendableRecordBatchStream; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +use prost::Message; +use std::sync::Arc; +use tonic::{Request, Response, Status, Streaming}; +use url::Url; + +const RECORD_BATCH_BUFFER_SIZE: usize = 2; + +impl Worker { + /// Converts this [Worker] into a [`WorkerServiceServer`] with high default message size limits. + /// + /// This is a convenience method that wraps the endpoint in a [`WorkerServiceServer`] and + /// configures it with `max_decoding_message_size(usize::MAX)` and + /// `max_encoding_message_size(usize::MAX)` to avoid message size limitations for internal + /// communication. + /// + /// You can further customize the returned server by chaining additional tonic methods. + /// + /// # Example + /// + /// ``` + /// # use datafusion_distributed::Worker; + /// # use tonic::transport::Server; + /// # use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + /// # async fn f() { + /// + /// let worker = Worker::default(); + /// let server = worker.into_worker_server(); + /// + /// Server::builder() + /// .add_service(Worker::default().into_worker_server()) + /// .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)) + /// .await; + /// + /// # } + /// ``` + pub fn into_worker_server(self) -> pb::worker_service_server::WorkerServiceServer { + pb::worker_service_server::WorkerServiceServer::new(self) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX) + } + + /// Creates an [`ObservabilityServiceServer`] that exposes task progress and cluster + /// worker discovery via the provided [`WorkerResolver`]. + /// + /// The returned server is meant to be added to the same [`tonic::transport::Server`] as the + /// Flight service — gRPC multiplexes both services on a single port. + pub fn with_observability_service( + &self, + worker_resolver: Arc, + ) -> ObservabilityServiceServer { + ObservabilityServiceServer::new(ObservabilityServiceImpl::new( + self.task_data_entries.clone(), + worker_resolver, + )) + } +} + +/// Implementation of the `worker.proto` specification based on the generated Rust stubs. +/// +/// The methods are delegated to plan `impl Worker` implementations so that they can be implemented +/// in different files. +#[async_trait] +impl pb::worker_service_server::WorkerService for Worker { + type CoordinatorChannelStream = BoxStream<'static, Result>; + + async fn coordinator_channel( + &self, + request: Request>, + ) -> Result, Status> { + let (metadata, _ext, body) = request.into_parts(); + + let input_stream = body + .map_err(map_status_to_datafusion_error) + .map(move |msg| { + decode_coordinator_to_worker_msg(msg?).map_err(map_status_to_datafusion_error) + }) + .boxed(); + + let output_stream = self + .coordinator_channel(metadata.into_headers(), input_stream) + .await + .map_err(datafusion_error_to_tonic_status)? + .map(|msg| match msg { + Ok(msg) => encode_worker_to_coordinator_msg(msg), + Err(err) => Err(datafusion_error_to_tonic_status(err)), + }) + .boxed(); + + Ok(Response::new(output_stream)) + } + + type ExecuteTaskStream = BoxStream<'static, Result>; + + async fn execute_task( + &self, + request: Request, + ) -> Result, Status> { + let body = request.into_inner(); + let request = decode_execute_task_request(body).await?; + let partition_range = request.target_partition_start..request.target_partition_end; + + let (arrow_streams, task_ctx) = self + .execute_task(request) + .await + .map_err(datafusion_error_to_tonic_status)?; + + let d_cfg = DistributedConfig::from_config_options(task_ctx.session_config().options()) + .map_err(datafusion_error_to_tonic_status)?; + + let compression = match d_cfg.compression.as_str() { + "lz4" => Some(CompressionType::LZ4_FRAME), + "zstd" => Some(CompressionType::ZSTD), + "none" => None, + v => Err(Status::invalid_argument(format!( + "Unknown compression type {v}" + )))?, + }; + let mut flight_streams = Vec::with_capacity(arrow_streams.len()); + for (partition, arrow_stream) in partition_range.zip(arrow_streams) { + let flight_stream = + build_flight_data_stream(arrow_stream, compression)?.map(move |msg| { + // For each FlightData produced by this stream, mark it with the appropriate + // partition. This stream will be merged with several others from other partitions, + // so marking it with the original partition allows it to be deconstructed into + // the original per-partition streams in later steps. + let flight_data = pb::FlightAppMetadata { + partition: partition as u64, + created_timestamp_unix_nanos: now_ns::(), + }; + msg.map(|v| v.with_app_metadata(flight_data.encode_to_vec())) + }); + + flight_streams.push(flight_stream); + } + + // Merge all the per-partition streams into one. Each message in the stream is marked with + // the original partition, so they can be reconstructed at the other side of the boundary. + let memory_pool = Arc::clone(&task_ctx.runtime_env().memory_pool); + let stream = spawn_select_all(flight_streams, memory_pool, RECORD_BATCH_BUFFER_SIZE); + + Ok(Response::new(Box::pin(stream.map_err(|err| match err { + FlightError::Tonic(status) => *status, + _ => Status::internal(format!("Error during flight stream: {err}")), + })))) + } + + async fn get_worker_info( + &self, + _request: Request, + ) -> Result, Status> { + Ok(Response::new(pb::GetWorkerInfoResponse { + version: self.version().to_string(), + })) + } +} + +fn decode_coordinator_to_worker_msg( + msg: pb::CoordinatorToWorkerMsg, +) -> Result { + Ok( + match 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::WorkUnitBatch(batch) => { + CoordinatorToWorkerMsg::WorkUnitBatch(decode_work_unit_batch(batch)?) + } + pb::coordinator_to_worker_msg::Inner::WorkUnitEos(_) => { + CoordinatorToWorkerMsg::WorkUnitEos + } + }, + ) +} + +fn decode_set_plan_request(request: pb::SetPlanRequest) -> Result { + Ok(SetPlanRequest { + task_key: decode_task_key(request.task_key.ok_or_else(missing("task_key"))?)?, + task_count: request.task_count as usize, + plan_proto: request.plan_proto, + work_unit_feed_declarations: request + .work_unit_feed_declarations + .into_iter() + .map(decode_work_unit_feed_declaration) + .collect::>()?, + target_worker_url: parse_url(&request.target_worker_url, "target_worker_url")?, + query_start_time_ns: request.query_start_time_ns as usize, + }) +} + +async fn decode_execute_task_request( + request: pb::ExecuteTaskRequest, +) -> Result { + Ok(ExecuteTaskRequest { + task_key: decode_task_key(request.task_key.ok_or_else(missing("task_key"))?)?, + target_partition_start: request.target_partition_start as usize, + target_partition_end: request.target_partition_end as usize, + producer_head_spec: decode_producer_head_spec( + request.producer_head.ok_or_else(missing("producer_head"))?, + ), + }) +} + +pub(super) fn decode_producer_head_spec( + proto: pb::execute_task_request::ProducerHead, +) -> ProducerHeadSpec { + match proto { + pb::execute_task_request::ProducerHead::None(_) => ProducerHeadSpec::None, + pb::execute_task_request::ProducerHead::Broadcast(v) => ProducerHeadSpec::BroadcastExec { + output_partitions: v.output_partitions as usize, + }, + pb::execute_task_request::ProducerHead::Repartition(v) => { + ProducerHeadSpec::RepartitionExec { + partitioning: v.partitioning, + } + } + } +} + +fn encode_worker_to_coordinator_msg( + msg: WorkerToCoordinatorMsg, +) -> Result { + Ok(pb::WorkerToCoordinatorMsg { + inner: Some(match msg { + WorkerToCoordinatorMsg::TaskMetrics(task_metrics) => { + pb::worker_to_coordinator_msg::Inner::TaskMetrics(encode_task_metrics( + task_metrics, + )?) + } + WorkerToCoordinatorMsg::LoadInfo(load_info) => { + pb::worker_to_coordinator_msg::Inner::LoadInfo(encode_load_info(load_info)) + } + WorkerToCoordinatorMsg::LoadInfoEos => { + pb::worker_to_coordinator_msg::Inner::LoadInfoEos(true) + } + }), + }) +} + +fn encode_task_metrics(task_metrics: TaskMetrics) -> Result { + Ok(pb::TaskMetrics { + pre_order_plan_metrics: task_metrics + .pre_order_plan_metrics + .into_iter() + .map(|metrics_set| { + df_metrics_set_to_proto(&metrics_set).map_err(datafusion_error_to_tonic_status) + }) + .collect::>()?, + task_metrics: Some( + df_metrics_set_to_proto(&task_metrics.task_metrics) + .map_err(datafusion_error_to_tonic_status)?, + ), + }) +} + +fn encode_load_info(load_info: LoadInfo) -> pb::LoadInfo { + pb::LoadInfo { + partition: load_info.partition as u64, + rows_ready: load_info.rows_ready as u64, + rows_per_second: load_info.rows_per_second as u64, + per_column_bytes_ready: load_info + .per_column_bytes_ready + .into_iter() + .map(|bytes| bytes as u64) + .collect(), + per_column_bytes_per_second: load_info + .per_column_bytes_per_second + .into_iter() + .map(|bytes| bytes as u64) + .collect(), + per_column_ndv_percentage: load_info.per_column_ndv_percentage, + per_column_null_percentage: load_info.per_column_null_percentage, + } +} + +fn decode_work_unit_batch(batch: pb::WorkUnitBatch) -> Result { + Ok(WorkUnitBatch { + batch: batch + .batch + .into_iter() + .map(decode_work_unit) + .collect::>()?, + }) +} + +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, + 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, + processed_timestamp_unix_nanos: work_unit.processed_timestamp_unix_nanos as usize, + }) +} + +fn decode_work_unit_feed_declaration( + declaration: pb::set_plan_request::WorkUnitFeedDeclaration, +) -> Result { + Ok(WorkUnitFeedDeclaration { + id: deserialize_uuid(&declaration.id).map_err(datafusion_error_to_tonic_status)?, + partitions: declaration.partitions as usize, + }) +} + +fn decode_task_key(task_key: pb::TaskKey) -> Result { + Ok(TaskKey { + query_id: deserialize_uuid(&task_key.query_id).map_err(datafusion_error_to_tonic_status)?, + stage_id: task_key.stage_id as usize, + task_number: task_key.task_number as usize, + }) +} + +fn parse_url(value: &str, field: &'static str) -> Result { + Url::parse(value) + .map_err(|err| Status::invalid_argument(format!("Invalid field '{field}': {err}"))) +} + +fn missing(field: &'static str) -> impl FnOnce() -> Status { + move || Status::invalid_argument(format!("Missing field '{field}'")) +} + +fn build_flight_data_stream( + stream: SendableRecordBatchStream, + compression_type: Option, +) -> datafusion::common::Result { + let stream = FlightDataEncoderBuilder::new() + .with_options( + IpcWriteOptions::default() + .try_with_compression(compression_type) + .map_err(|err| Status::internal(err.to_string()))?, + ) + .with_schema(stream.schema()) + // This tells the encoder to send dictionaries across the wire as-is. + // The alternative (`DictionaryHandling::Hydrate`) would expand the dictionaries + // into their value types, which can potentially blow up the size of the data transfer. + // The main reason to use `DictionaryHandling::Hydrate` is for compatibility with clients + // that do not support dictionaries, but since we are using the same server/client on both + // sides, we can safely use `DictionaryHandling::Resend`. + // Note that we do garbage collection of unused dictionary values above, so we are not sending + // unused dictionary values over the wire. + .with_dictionary_handling(DictionaryHandling::Resend) + // Set max flight data size to unlimited. + // This requires servers and clients to also be configured to handle unlimited sizes. + // Using unlimited sizes avoids splitting RecordBatches into multiple FlightData messages, + // which could add significant overhead for large RecordBatches. + // The only reason to split them really is if the client/server are configured with a message size limit, + // which mainly makes sense in a public network scenario where you want to avoid DoS attacks. + // Since all of our Arrow Flight communication happens within trusted data plane networks, + // we can safely use unlimited sizes here. + .with_max_flight_data_size(usize::MAX) + .build( + stream + // Apply garbage collection of dictionary and view arrays before sending over the network + .and_then(|rb| std::future::ready(garbage_collect_arrays(rb))) + .map_err(|err| FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(err)))), + ); + Ok(stream) +} + +/// Garbage collects values sub-arrays. +/// +/// We apply this before sending RecordBatches over the network to avoid sending +/// values that are not referenced by any dictionary keys or buffers that are not used. +/// +/// Unused values can arise from operations such as filtering, where +/// some keys may no longer be referenced in the filtered result. +fn garbage_collect_arrays( + batch: RecordBatch, +) -> datafusion::common::Result { + let (schema, arrays, row_count) = batch.into_parts(); + + let arrays = arrays + .into_iter() + .map(|array| { + if let Some(array) = array.as_any_dictionary_opt() { + garbage_collect_any_dictionary(array) + } else if let Some(array) = array.as_string_view_opt() { + Ok(Arc::new(array.gc()) as Arc) + } else if let Some(array) = array.as_binary_view_opt() { + Ok(Arc::new(array.gc()) as Arc) + } else { + Ok(array) + } + }) + .collect::, _>>()?; + + Ok(RecordBatch::try_new_with_options( + schema, + arrays, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + )?) +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs new file mode 100644 index 00000000..5c6bdd1a --- /dev/null +++ b/src/protocol/mod.rs @@ -0,0 +1,14 @@ +#[cfg(feature = "grpc")] +pub mod grpc; + +mod channel_resolver; +mod worker_channel; + +pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; +pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; + +pub use worker_channel::{ + CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, + LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, + WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, +}; diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs new file mode 100644 index 00000000..41ee2a78 --- /dev/null +++ b/src/protocol/worker_channel.rs @@ -0,0 +1,187 @@ +use async_trait::async_trait; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use futures::stream::BoxStream; +use http::HeaderMap; +use std::sync::Arc; +use url::Url; +use uuid::Uuid; + +#[async_trait] +pub trait WorkerChannel: Send + Sync { + /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages + /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel + /// per task. + async fn coordinator_channel( + &mut self, + headers: HeaderMap, + c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>>; + + /// Executes the requested partition range of a subplan previously sent by the coordinator channel. + async fn execute_task( + &mut self, + headers: HeaderMap, + request: ExecuteTaskRequest, + metrics: ExecutionPlanMetricsSet, + task_ctx: &Arc, + ) -> Result>>>; + + /// Returns metadata about a worker. Currently only used for worker versioning. + async fn get_worker_info( + &mut self, + request: GetWorkerInfoRequest, + ) -> Result; +} + +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. + WorkUnitBatch(WorkUnitBatch), + /// Signals an EOS for WorkUnits. After this message is received, no more WorkUnits will be sent. + WorkUnitEos, +} + +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] +pub struct TaskKey { + /// Our query id. + pub query_id: Uuid, + /// Our stage id. + pub stage_id: usize, + /// The task number within the stage. + pub task_number: usize, +} + +pub struct WorkUnitFeedDeclaration { + /// Unique identifier of the node to which work unit feeds are expected to be streamed. + pub id: Uuid, + /// The amount of partitions expected to be streamed. + pub partitions: usize, +} + +pub struct SetPlanRequest { + /// The unique identifier of the task to which the subplan belongs to. + pub task_key: TaskKey, + /// 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, + /// 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. + /// + /// If no WorkUnitFeedExec nodes are present in the plan, this should be empty. + pub work_unit_feed_declarations: Vec, + /// The worker URL to which this message will go. The receiving worker will use this information + /// to identify itself, and avoid further calls in case it needs to call itself for executing tasks. + pub target_worker_url: Url, + /// Unix nanos when the query started as reported by the coordinator. Used for collecting temporal metrics + /// relative to when the query was fired in the coordinator. + pub query_start_time_ns: usize, +} + +pub struct WorkUnitBatch { + /// A batch of WorkUnits. + pub batch: Vec, +} + +pub struct WorkUnitMsg { + /// Identifier of the node to which this work unit feed belongs to. + pub id: Uuid, + /// 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, + /// 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. + pub sent_timestamp_unix_nanos: usize, + /// Unix timestamp in nanoseconds at which this message was received by a worker. + pub received_timestamp_unix_nanos: usize, + /// Unix timestamp in nanoseconds at which this message started being processed. + pub processed_timestamp_unix_nanos: usize, +} + +pub enum WorkerToCoordinatorMsg { + /// Sends the metrics collected during task execution back to the coordinator. + /// This is sent after all partitions of a task have finished (or been dropped), + /// ensuring metrics are never lost due to early stream termination. + /// metrics[i] is the set of metrics for plan node i in pre-order traversal order. + TaskMetrics(TaskMetrics), + /// Load information reported by a task. This information is used for dynamically + /// sizing the number of workers involved in a query. + LoadInfo(LoadInfo), + LoadInfoEos, +} + +#[derive(Clone, Debug)] +pub struct TaskMetrics { + /// Metrics for a single task's plan nodes in pre-order traversal order. + /// The TaskKey is implicit — it is determined by the SetPlanRequest that + /// opened this coordinator channel connection. + pub pre_order_plan_metrics: Vec, + /// Metrics related to the execution of a task within a stage. This metrics, instead of being + /// associated to a specific node, they are global to the task, like the time at which the plan + /// was fed by the coordinator to the worker. + pub task_metrics: MetricsSet, +} + +#[derive(Default)] +pub struct LoadInfo { + /// The partition index to which this message belongs to. + pub partition: usize, + /// The amount of rows ready to be returned. + pub rows_ready: usize, + /// The estimated velocity at which rows will flow through the node. If all the rows were + /// already accumulated, they will be reported by `rows_ready`, and this field will be 0. + pub rows_per_second: usize, + /// The amount of bytes ready to be returned per column. + pub per_column_bytes_ready: Vec, + /// The estimated velocity at which data will flow through each column. If all the bytes were + /// already accumulated, they will be reported by `bytes_ready`, and this field will be 0. + pub per_column_bytes_per_second: Vec, + /// Approximate ratio of NDV for each column. + pub per_column_ndv_percentage: Vec, + /// Approximate ratio of null count for each column. + pub per_column_null_percentage: Vec, +} + +pub struct ExecuteTaskRequest { + /// The unique identifier of the task that is going to get executed. + pub task_key: TaskKey, + /// The start of the partition range of the specified task that is going to be executed. + pub target_partition_start: usize, + /// The end of the partition range of the specified task that is going to be executed. + pub target_partition_end: usize, + /// The head node the requested task should have. Depending on the network boundary executing + /// the task, the head node should be prepared differently, for example: + /// - 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 struct GetWorkerInfoRequest {} + +pub struct GetWorkerInfoResponse { + pub version: String, +} diff --git a/src/stage.rs b/src/stage.rs index 3ac39e7a..2968ce8b 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -219,10 +219,9 @@ impl DistributedTaskContext { } } -use crate::common::serialize_uuid; -use crate::metrics::proto::metric_proto_to_df; -use crate::worker::generated::worker as pb; -use crate::{DistributedMetricsFormat, NetworkShuffleExec, rewrite_distributed_plan_with_metrics}; +use crate::{ + DistributedMetricsFormat, NetworkShuffleExec, TaskKey, rewrite_distributed_plan_with_metrics, +}; use crate::{NetworkBoundary, NetworkBoundaryExt}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::DataFusionError; @@ -433,21 +432,24 @@ fn display_inner_distributed_leaf( /// Gathers the metrics global to a stage. These metrics are not specific to any plan node, and /// are instead global to a whole stage. fn gather_stage_header_metrics(stage: &Stage, metrics_store: &MetricsStore) -> MetricsSet { - let mut task_key = pb::TaskKey { - query_id: serialize_uuid(&stage.query_id()), - stage_id: stage.num() as u64, + let mut task_key = TaskKey { + query_id: stage.query_id(), + stage_id: stage.num(), task_number: 0, }; let mut all_metrics = stage.metrics(); - while let Some(metrics_set) = metrics_store.get(&task_key).and_then(|v| v.task_metrics) { - for mut metric in metrics_set.metrics { - metric.labels.push(pb::Label { - name: DISTRIBUTED_DATAFUSION_TASK_ID_LABEL.to_string(), - value: task_key.task_number.to_string(), - }); - if let Ok(metric) = metric_proto_to_df(metric) { - all_metrics.push(metric) - }; + while let Some(metrics_set) = metrics_store.get(&task_key).map(|v| v.task_metrics) { + for metric in metrics_set.iter() { + let mut labels = metric.labels().to_vec(); + labels.push(Label::new( + DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, + task_key.task_number.to_string(), + )); + all_metrics.push(Arc::new(Metric::new_with_labels( + metric.value().clone(), + metric.partition(), + labels, + ))); } task_key.task_number += 1; } diff --git a/src/test_utils/in_memory_channel_resolver.rs b/src/test_utils/in_memory_channel_resolver.rs index 12ad84bb..92f710c7 100644 --- a/src/test_utils/in_memory_channel_resolver.rs +++ b/src/test_utils/in_memory_channel_resolver.rs @@ -1,8 +1,6 @@ -use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; use crate::{ - BoxCloneSyncChannel, ChannelResolver, DefaultSessionBuilder, DistributedExt, - MappedWorkerSessionBuilderExt, SessionStateBuilderExt, Worker, WorkerResolver, - WorkerSessionBuilder, create_worker_client, + ChannelResolver, DefaultSessionBuilder, DistributedExt, MappedWorkerSessionBuilderExt, + SessionStateBuilderExt, Worker, WorkerChannel, WorkerResolver, WorkerSessionBuilder, grpc, }; use async_trait::async_trait; use datafusion::common::DataFusionError; @@ -19,7 +17,7 @@ const DUMMY_URL_PREFIX: &str = "http://url-"; /// tokio duplex rather than a TCP connection. #[derive(Clone)] pub struct InMemoryChannelResolver { - channel: WorkerServiceClient, + channel: grpc::BoxCloneSyncChannel, } impl InMemoryChannelResolver { @@ -50,7 +48,7 @@ impl InMemoryChannelResolver { })); let this = Self { - channel: create_worker_client(BoxCloneSyncChannel::new(channel)), + channel: grpc::BoxCloneSyncChannel::new(channel), }; let this_clone = this.clone(); @@ -83,8 +81,8 @@ impl ChannelResolver for InMemoryChannelResolver { async fn get_worker_client_for_url( &self, _: &url::Url, - ) -> Result, DataFusionError> { - Ok(self.channel.clone()) + ) -> Result, DataFusionError> { + Ok(grpc::create_worker_client(self.channel.clone())) } } diff --git a/src/test_utils/metrics.rs b/src/test_utils/metrics.rs index 0677ffca..e8e5fe46 100644 --- a/src/test_utils/metrics.rs +++ b/src/test_utils/metrics.rs @@ -1,6 +1,7 @@ use crate::coordinator::DistributedExec; -use crate::worker::generated::worker as pb; +use chrono::{DateTime, Utc}; use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::metrics::{Count, Metric, MetricValue, MetricsSet, Time, Timestamp}; use std::sync::Arc; /// Waits until all worker tasks have reported their metrics back via the coordinator channel. @@ -11,43 +12,41 @@ pub async fn wait_for_all_metrics(plan: &Arc) { } /// creates a "distinct" set of metrics from the provided seed -pub fn make_test_metrics_set_proto_from_seed(seed: u64, num_metrics: usize) -> pb::MetricsSet { +pub fn make_test_metrics_set_from_seed(seed: u64, num_metrics: usize) -> MetricsSet { const TEST_TIMESTAMP: i64 = 1758200400000000000; // 2025-09-18 13:00:00 UTC - let mut result = pb::MetricsSet { metrics: vec![] }; + let mut result = MetricsSet::new(); for i in 0..num_metrics { - let value = seed + i as u64; - result.metrics.push(match i % 4 { - 0 => pb::Metric { - value: Some(pb::metric::Value::OutputRows(pb::OutputRows { value })), - labels: vec![], - partition: None, + let value = (seed + i as u64) as usize; + result.push(Arc::new(Metric::new( + match i % 4 { + 0 => { + let count = Count::new(); + count.add(value); + MetricValue::OutputRows(count) + } + 1 => { + let time = Time::new(); + time.add_duration(std::time::Duration::from_nanos(value as u64)); + MetricValue::ElapsedCompute(time) + } + 2 => MetricValue::StartTimestamp(timestamp_from_nanos( + TEST_TIMESTAMP + (value as i64 * 1_000_000_000), + )), + 3 => MetricValue::EndTimestamp(timestamp_from_nanos( + TEST_TIMESTAMP + (value as i64 * 1_000_000_000), + )), + _ => unreachable!(), }, - - 1 => pb::Metric { - value: Some(pb::metric::Value::ElapsedCompute(pb::ElapsedCompute { - value, - })), - labels: vec![], - partition: None, - }, - 2 => pb::Metric { - value: Some(pb::metric::Value::StartTimestamp(pb::StartTimestamp { - value: Some(TEST_TIMESTAMP + (value as i64 * 1_000_000_000)), - })), - labels: vec![], - partition: None, - }, - 3 => pb::Metric { - value: Some(pb::metric::Value::EndTimestamp(pb::EndTimestamp { - value: Some(TEST_TIMESTAMP + (value as i64 * 1_000_000_000)), - })), - labels: vec![], - partition: None, - }, - _ => unreachable!(), - }) + None, + ))) } result } + +fn timestamp_from_nanos(nanos: i64) -> Timestamp { + let timestamp = Timestamp::new(); + timestamp.set(DateTime::::from_timestamp_nanos(nanos)); + timestamp +} diff --git a/src/test_utils/plans.rs b/src/test_utils/plans.rs index 7d299489..aa56338e 100644 --- a/src/test_utils/plans.rs +++ b/src/test_utils/plans.rs @@ -1,15 +1,13 @@ #[cfg(test)] use super::parquet::register_parquet_tables; -use crate::NetworkBoundaryExt; -use crate::common::serialize_uuid; use crate::coordinator::DistributedExec; use crate::stage::Stage; -use crate::worker::generated::worker::TaskKey; #[cfg(test)] use crate::{ DistributedConfig, DistributedExt, SessionStateBuilderExt, TaskEstimation, TaskEstimator, display_plan_ascii, test_utils::in_memory_channel_resolver::InMemoryWorkerResolver, }; +use crate::{NetworkBoundaryExt, TaskKey}; #[cfg(test)] use datafusion::{ common::Result, @@ -67,9 +65,9 @@ pub fn get_stages_and_task_keys( // Add each task. for j in 0..stage.task_count() { task_keys.insert(TaskKey { - query_id: serialize_uuid(&stage.query_id()), - stage_id: stage.num() as u64, - task_number: j as u64, + query_id: stage.query_id(), + stage_id: stage.num(), + task_number: j, }); } diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index 1526508c..18c51ffe 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -1,11 +1,11 @@ -use crate::common::{now_ns, serialize_uuid}; -use crate::worker::generated::worker as pb; -use crate::{BytesMetricExt, LatencyMetricExt, WorkUnit}; +use crate::common::now_ns; +use crate::{ + BytesMetricExt, CoordinatorToWorkerMsg, LatencyMetricExt, 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}; @@ -13,8 +13,8 @@ use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; use uuid::Uuid; -pub(crate) type WorkUnitTx = UnboundedSender>; -pub(crate) type WorkUnitRx = UnboundedReceiver>; +pub(crate) type WorkUnitTx = UnboundedSender>; +pub(crate) type WorkUnitRx = UnboundedReceiver>; pub(crate) type RemoteWorkUnitFeedRxs = HashMap<(Uuid, usize), Mutex>>; pub(crate) type RemoteWorkUnitFeedTxs = HashMap<(Uuid, usize), WorkUnitTx>; @@ -36,7 +36,7 @@ pub(crate) struct RemoteWorkUnitFeedRegistry { } impl RemoteWorkUnitFeedRegistry { - /// Creates all the receivers and senders for a specific [WorkUnit] Feed id. One feed per + /// Creates all the receivers and senders for a specific [WorkUnitMsg] Feed id. One feed per /// partition is created. /// /// Calling this twice with the same `id` is a coordinator bug — duplicate declarations @@ -60,36 +60,27 @@ impl RemoteWorkUnitFeedRegistry { pub(crate) fn build_work_unit_batch_msg( id: &Uuid, work_unit_batch: Vec<(usize, Result>)>, -) -> Result { - Ok(pb::CoordinatorToWorkerMsg { - inner: Some(pb::coordinator_to_worker_msg::Inner::WorkUnitBatch( - pb::WorkUnitBatch { - batch: work_unit_batch - .into_iter() - .map(|(partition, work_unit)| { - Ok(pb::WorkUnit { - id: serialize_uuid(id), - partition: partition as u64, - body: work_unit?.encode_to_bytes(), - created_timestamp_unix_nanos: now_ns(), - sent_timestamp_unix_nanos: 0, - received_timestamp_unix_nanos: 0, - processed_timestamp_unix_nanos: 0, - }) - }) - .collect::>()?, - }, - )), - }) +) -> Result { + Ok(CoordinatorToWorkerMsg::WorkUnitBatch(WorkUnitBatch { + batch: work_unit_batch + .into_iter() + .map(|(partition, work_unit)| { + Ok(WorkUnitMsg { + id: *id, + partition, + body: work_unit?.encode_to_bytes(), + created_timestamp_unix_nanos: now_ns(), + sent_timestamp_unix_nanos: 0, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + }) + }) + .collect::>()?, + })) } -pub(crate) fn set_work_unit_send_time( - mut msg: pb::CoordinatorToWorkerMsg, -) -> pb::CoordinatorToWorkerMsg { - if let pb::CoordinatorToWorkerMsg { - inner: Some(pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(work_unit_batch)), - } = &mut msg - { +pub(crate) fn set_work_unit_send_time(mut msg: CoordinatorToWorkerMsg) -> CoordinatorToWorkerMsg { + if let CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) = &mut msg { for work_unit in &mut work_unit_batch.batch { work_unit.sent_timestamp_unix_nanos = now_ns(); } @@ -98,12 +89,9 @@ pub(crate) fn set_work_unit_send_time( } pub(crate) fn set_work_unit_received_time( - mut msg: pb::CoordinatorToWorkerMsg, -) -> pb::CoordinatorToWorkerMsg { - if let pb::CoordinatorToWorkerMsg { - inner: Some(pb::coordinator_to_worker_msg::Inner::WorkUnitBatch(work_unit_batch)), - } = &mut msg - { + mut msg: CoordinatorToWorkerMsg, +) -> CoordinatorToWorkerMsg { + if let CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) = &mut msg { for work_unit in &mut work_unit_batch.batch { work_unit.received_timestamp_unix_nanos = now_ns(); } @@ -111,7 +99,7 @@ pub(crate) fn set_work_unit_received_time( msg } -/// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnit]s coming over +/// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnitMsg]s coming over /// the wire from a [RemoteWorkUnitFeedRegistry]. /// /// Deserializing a [crate::WorkUnitFeed] with [crate::WorkUnitFeed::from_proto] always returns a @@ -126,7 +114,7 @@ pub(crate) struct RemoteFeedProvider { } impl RemoteFeedProvider { - pub(crate) fn feed( + pub(crate) fn feed( &self, partition: usize, ctx: Arc, @@ -168,36 +156,36 @@ impl RemoteFeedProvider { }; Ok(UnboundedReceiverStream::new(receiver) - .map(move |work_unit_or_err| { - let mut work_unit = work_unit_or_err?; + .map(move |work_unit_msg_or_err| { + let mut work_unit_msg = work_unit_msg_or_err?; let timer = elapsed_compute.timer(); - let result = T::decode(work_unit.body.as_slice()) - .map_err(|err| proto_error(format!("{err}"))); + let work_unit = T::decode(work_unit_msg.body.as_slice()) + .map_err(|err| datafusion_proto::protobuf::proto_error(format!("{err}"))); timer.done(); - work_unit.processed_timestamp_unix_nanos = now_ns(); + work_unit_msg.processed_timestamp_unix_nanos = now_ns(); + let body_len = work_unit_msg.body.len(); - let pb::WorkUnit { + let WorkUnitMsg { created_timestamp_unix_nanos: base, sent_timestamp_unix_nanos, received_timestamp_unix_nanos, processed_timestamp_unix_nanos, - body, .. - } = work_unit; + } = work_unit_msg; - bytes_transferred.add_bytes(body.len()); + bytes_transferred.add_bytes(body_len); msg_count.add(1); - send_latency_max.add_nanos((sent_timestamp_unix_nanos - base) as usize); - send_latency_p50.add_nanos((sent_timestamp_unix_nanos - base) as usize); + send_latency_max.add_nanos(sent_timestamp_unix_nanos - base); + send_latency_p50.add_nanos(sent_timestamp_unix_nanos - base); - received_latency_max.add_nanos((received_timestamp_unix_nanos - base) as usize); - received_latency_p50.add_nanos((received_timestamp_unix_nanos - base) as usize); + received_latency_max.add_nanos(received_timestamp_unix_nanos - base); + received_latency_p50.add_nanos(received_timestamp_unix_nanos - base); - processed_latency_max.add_nanos((processed_timestamp_unix_nanos - base) as usize); - processed_latency_p50.add_nanos((processed_timestamp_unix_nanos - base) as usize); + processed_latency_max.add_nanos(processed_timestamp_unix_nanos - base); + processed_latency_p50.add_nanos(processed_timestamp_unix_nanos - base); - result + work_unit }) .boxed()) } diff --git a/src/work_unit_feed/work_unit_feed_provider.rs b/src/work_unit_feed/work_unit_feed_provider.rs index aab04c88..7647b6a1 100644 --- a/src/work_unit_feed/work_unit_feed_provider.rs +++ b/src/work_unit_feed/work_unit_feed_provider.rs @@ -21,7 +21,7 @@ use std::sync::Arc; /// /// See [`WorkUnitFeedProvider::feed`] for the per-call contract. pub trait WorkUnitFeedProvider: Send + Sync + Debug { - type WorkUnit: WorkUnit + Default; + type WorkUnit: WorkUnit + Default + 'static; /// Builds a [`WorkUnit`] stream for the given `partition`. /// diff --git a/src/worker/gen/regen.sh b/src/worker/gen/regen.sh deleted file mode 100755 index 39e04fdc..00000000 --- a/src/worker/gen/regen.sh +++ /dev/null @@ -1,6 +0,0 @@ -#!/usr/bin/env bash - -set -e - -repo_root=$(git rev-parse --show-toplevel) -cd "$repo_root" && cargo run --manifest-path src/worker/gen/Cargo.toml diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index fcb3944a..b60edc37 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,77 +1,68 @@ -use crate::common::deserialize_uuid; +use crate::common::TreeNodeExt; use crate::execution_plans::SamplerExec; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; -use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; -use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; -use crate::worker::generated::worker::worker_service_server::WorkerService; -use crate::worker::generated::worker::{ - CoordinatorToWorkerMsg, WorkerToCoordinatorMsg, worker_to_coordinator_msg, -}; use crate::worker::task_data::TaskDataMetrics; use crate::{ - DistributedCodec, DistributedConfig, DistributedExt, DistributedTaskContext, TaskData, Worker, - WorkerQueryContext, + CoordinatorToWorkerMsg, DistributedCodec, DistributedConfig, DistributedExt, + DistributedTaskContext, TaskData, TaskMetrics, Worker, WorkerQueryContext, + WorkerToCoordinatorMsg, }; -use datafusion::common::DataFusionError; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_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::FuturesUnordered; +use futures::stream::{BoxStream, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; +use http::HeaderMap; use std::sync::atomic::AtomicUsize; -use std::sync::{Arc, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock}; use tokio::sync::oneshot; -use tonic::{Request, Response, Status, Streaming}; -use url::Url; +use tokio::sync::oneshot::Sender; impl Worker { - pub(super) async fn impl_coordinator_channel( + pub async fn coordinator_channel( &self, - request: Request>, - ) -> Result::CoordinatorChannelStream>, Status> { - let (grpc_headers, _ext, mut body) = request.into_parts(); - + headers: HeaderMap, + mut stream: BoxStream<'static, Result>, + ) -> Result>> { // The first message must be a SetPlanRequest. - let Some(msg) = body.next().await else { - return Err(Status::internal("Empty Coordinator stream")); + let Some(msg) = stream.try_next().await? else { + return internal_err!("Empty Coordinator stream"); }; - let Some(Inner::SetPlanRequest(request)) = msg?.inner else { - return Err(Status::internal( - "First Coordinator message must be SetPlanRequest", - )); + + let CoordinatorToWorkerMsg::SetPlanRequest(request) = msg else { + return internal_err!("First Coordinator message must be SetPlanRequest"); }; - let key = request.task_key.ok_or_else(missing("task_key"))?; + + let key = request.task_key; let entry = self .task_data_entries - .get_with(key.clone(), async { Default::default() }) + .get_with(key, async { Default::default() }) .await; let mut remote_work_unit_feed_registry = RemoteWorkUnitFeedRegistry::default(); - for WorkUnitFeedDeclaration { id, partitions } in &request.work_unit_feed_declarations { - if let Ok(id) = deserialize_uuid(id) { - remote_work_unit_feed_registry.add(id, *partitions as usize); - } + for decl in request.work_unit_feed_declarations { + remote_work_unit_feed_registry.add(decl.id, decl.partitions); } let (metrics_tx, metrics_rx) = oneshot::channel(); let mut load_info_rxs = vec![]; let task_data = || async { - let headers = grpc_headers.into_headers(); - let mut cfg = SessionConfig::default() .with_extension(Arc::new(remote_work_unit_feed_registry.receivers)) .with_extension(Arc::new(DistributedTaskContext { - task_index: key.task_number as usize, - task_count: request.task_count as usize, + task_index: request.task_key.task_number, + task_count: request.task_count, })) .with_extension(Arc::new(LocalWorkerContext { task_data_entries: Arc::clone(&self.task_data_entries), - self_url: Url::parse(&request.target_worker_url) - .map_err(|e| DataFusionError::External(Box::new(e)))?, + self_url: request.target_worker_url, })) .with_distributed_option_extension_from_headers::(&headers)?; @@ -119,38 +110,36 @@ impl Worker { }) }; - entry.write(task_data().await.map_err(Arc::new)).map_err(|_| { - Status::internal(format!( - "Logic error while setting plan for TaskKey {key:?}: the plan was set twice. This is a bug in datafusion-distributed, please report it." - )) - })?; + let task_data_result = task_data().await.map_err(Arc::new); + + entry + .write(task_data_result.clone()) + .map_err(|e| exec_datafusion_err!("{e}"))?; + + let task_data = task_data_result.map_err(DataFusionError::Shared)?; // Continue reading remaining messages (work unit feed data) in the background. let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders); let task_data_entries = Arc::clone(&self.task_data_entries); + let task_count = request.task_count; #[allow(clippy::disallowed_methods)] tokio::spawn(async move { - let mut body = body.map_ok(set_work_unit_received_time); - while let Some(Ok(msg)) = body.next().await { - let Some(msg) = msg.inner else { - continue; - }; + let mut stream = stream.map_ok(set_work_unit_received_time); + while let Some(Ok(msg)) = stream.next().await { match msg { - Inner::SetPlanRequest(_) => { + CoordinatorToWorkerMsg::SetPlanRequest(_) => { // SetPlanRequest should be the first already polled message in the stream, // if some reached here it means that something is wrong. continue; } - Inner::WorkUnitBatch(msg) => { + CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) => { let Some(work_unit_senders) = work_unit_senders.as_mut() else { continue; }; - for wu in msg.batch { - let Ok(id) = deserialize_uuid(&wu.id) else { - continue; - }; - let partition = wu.partition as usize; - let Some(tx) = work_unit_senders.get(&(id, partition)) else { + for wu in work_unit_batch.batch { + let id = wu.id; + let partition = wu.partition; + let Some(tx) = work_unit_senders.get(&(wu.id, partition)) else { continue; }; if tx.send(Ok(wu)).is_err() { @@ -161,7 +150,7 @@ impl Worker { } } } - Inner::WorkUnitEos(_) => { + CoordinatorToWorkerMsg::WorkUnitEos => { // No further work unit message will be received here, so drop all the // sender sides so that receiver sides see an EOS upon draining the // remaining messages. @@ -173,41 +162,69 @@ impl Worker { } } } - #[allow(clippy::disallowed_methods)] - tokio::spawn(async move { task_data_entries.invalidate(&key).await }); + + if let Some(Ok(plan)) = task_data.final_plan.get() { + let d_ctx = DistributedTaskContext { + task_index: key.task_number, + task_count, + }; + let task_data_metrics = &task_data.task_data_metrics; + task_data_metrics.mark_execution_finished(); + send_metrics_via_channel(&task_data.metrics_tx, plan, d_ctx, task_data_metrics); + } + task_data_entries.invalidate(&key).await }); let load_info_stream = FuturesUnordered::from_iter(load_info_rxs) .filter_map(async |load_info_or_channel_dropped| { // This error can only happen if the pb::LoadInfo sender was dropped, which is fine. let load_info = load_info_or_channel_dropped.ok()?; - Some(Ok(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::LoadInfo(load_info)), - })) + Some(WorkerToCoordinatorMsg::LoadInfo(load_info)) }) .chain(futures::stream::once(async move { - Ok(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::LoadInfoEos(true)), - }) + WorkerToCoordinatorMsg::LoadInfoEos })); - // Stream back the metrics once the task finishes executing. - // The oneshot receiver resolves when impl_execute_task sends the collected - // metrics after all partitions have finished or been dropped. + // Stream back metrics when the coordinator channel reaches EOS. At that point the + // coordinator has closed the query-scoped request stream, so any remaining task state can + // be finalized even if some partition streams were not dropped through the normal path. let metrics_stream = metrics_rx.into_stream(); let metrics_stream = metrics_stream.filter_map(async |task_metrics_or_channel_dropped| { let task_metrics = task_metrics_or_channel_dropped.ok()?; - Some(Ok(WorkerToCoordinatorMsg { - inner: Some(worker_to_coordinator_msg::Inner::TaskMetrics(task_metrics)), - })) + Some(WorkerToCoordinatorMsg::TaskMetrics(task_metrics)) }); - Ok(Response::new( - futures::stream::select(load_info_stream, metrics_stream).boxed(), - )) + Ok(futures::stream::select(load_info_stream, metrics_stream) + .map(Ok) + .boxed()) } } -fn missing(field: &'static str) -> impl FnOnce() -> Status { - move || Status::invalid_argument(format!("Missing field '{field}'")) +/// Collects metrics from the plan in pre-order traversal order and sends them via the +/// coordinator channel oneshot. +fn send_metrics_via_channel( + metrics_tx: &Arc>>>, + plan: &Arc, + dt_ctx: DistributedTaskContext, + task_data_metrics: &Arc, +) { + let mut pre_order_plan_metrics = vec![]; + let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { + pre_order_plan_metrics.push(node.metrics().unwrap_or_default()); + Ok(TreeNodeRecursion::Continue) + }); + + let tx = { + let mut guard = match metrics_tx.lock() { + Ok(g) => g, + Err(_) => return, + }; + guard.take() + }; + let Some(tx) = tx else { return }; + // Ignore send errors — the coordinator channel may have been dropped (e.g. query cancelled). + let _ = tx.send(TaskMetrics { + pre_order_plan_metrics, + task_metrics: task_data_metrics.to_metrics_set(), + }); } diff --git a/src/worker/impl_execute_task.rs b/src/worker/impl_execute_task.rs index 9be4e51b..10dc6c14 100644 --- a/src/worker/impl_execute_task.rs +++ b/src/worker/impl_execute_task.rs @@ -1,39 +1,14 @@ -use crate::common::{TreeNodeExt, now_ns, on_drop_stream}; -use crate::metrics::proto::df_metrics_set_to_proto; -use crate::protobuf::datafusion_error_to_tonic_status; -use crate::worker::generated::worker::{FlightAppMetadata, TaskMetrics}; +use crate::ExecuteTaskRequest; use crate::worker::worker_service::{TaskDataEntries, Worker}; -use crate::{DistributedConfig, DistributedTaskContext}; -use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder}; -use arrow_flight::error::FlightError; -use arrow_select::dictionary::garbage_collect_any_dictionary; -use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; -use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{Result, exec_err, internal_err}; - -use crate::worker::generated::worker::ExecuteTaskRequest; -use crate::worker::generated::worker::worker_service_server::WorkerService; -use crate::worker::spawn_select_all::spawn_select_all; -use crate::worker::task_data::TaskDataMetrics; -use datafusion::arrow::ipc::CompressionType; -use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::common::exec_datafusion_err; +use datafusion::common::{Result, exec_err}; use datafusion::error::DataFusionError; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use futures::TryStreamExt; -use prost::Message; use std::sync::Arc; -use std::sync::Mutex; -use std::sync::atomic::Ordering; use std::time::Duration; -use tokio::sync::oneshot::Sender; -use tokio_stream::StreamExt; -use tonic::{Request, Response, Status}; /// How many record batches to buffer from the plan execution. -const RECORD_BATCH_BUFFER_SIZE: usize = 2; const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; /// Builds several per-partition streams by retrieving the appropriate entry from [TaskDataEntries] @@ -41,234 +16,52 @@ const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; /// /// This method is async mainly for the key retrieval operation from [TaskDataEntries], but it does /// not start polling any stream, it just instantiates them. -pub(crate) async fn execute_local_task( - task_data_entries: &Arc, - body: ExecuteTaskRequest, -) -> Result<(Vec, Arc)> { - let Some(key) = body.task_key.as_ref().cloned() else { - return internal_err!("Missing task_key in LocalWorkerConnection"); - }; - let Some(producer_head) = body.producer_head.as_ref().cloned() else { - return internal_err!("Missing producer_head"); - }; - let entry = task_data_entries - .get_with(key.clone(), async { Default::default() }) - .await; - - // Other request is responsible for writing the plan that belongs to this TaskKey, so - // we'll resolve immediately if it was already there, or wait until it's ready. - let task_data = entry - .read(Duration::from_secs(WAIT_PLAN_TIMEOUT_SECS)) - .await - .map_err(|e| exec_datafusion_err!("Worker::execute_task timed-out while waiting for the plan to be set by the coordinator. ({e})"))? - .map_err(DataFusionError::Shared)?; - task_data.task_data_metrics.mark_execution_started_once(); - - let plan = task_data.plan(producer_head)?; - let task_ctx = task_data.task_ctx; - let d_cfg = DistributedConfig::from_config_options(task_ctx.session_config().options())?; - let d_ctx = *DistributedTaskContext::from_ctx(&task_ctx).as_ref(); - - let send_metrics = d_cfg.collect_metrics; - let partition_count = plan.properties().partitioning.partition_count(); - let plan_name = plan.name(); - - // Execute all the requested partitions at once, and collect all the streams so that they - // can be merged into a single one at the end of this function. - let n_streams = body.target_partition_end - body.target_partition_start; - let mut streams = Vec::with_capacity(n_streams as usize); - for partition in body.target_partition_start..body.target_partition_end { - if partition >= partition_count as u64 { - return exec_err!( - "partition {partition} not available. The head plan {plan_name} of the stage just has {partition_count} partitions" - ); - } - - let stream = plan.execute(partition as usize, Arc::clone(&task_ctx))?; - let stream_schema = plan.schema(); - - let plan = Arc::clone(&plan); - - let task_data_entries = Arc::clone(task_data_entries); - let num_partitions_remaining = Arc::clone(&task_data.num_partitions_remaining); - let metrics_tx = Arc::clone(&task_data.metrics_tx); - let task_data_metrics = Arc::clone(&task_data.task_data_metrics); - let key = key.clone(); - let stream = on_drop_stream(stream, move || { - // Stream was dropped before fully consumed -- see https://github.com/datafusion-contrib/datafusion-distributed/issues/412 - // Send metrics via the coordinator channel so they are not lost. - if num_partitions_remaining.fetch_sub(1, Ordering::SeqCst) == 1 { - // Fire-and-forget background tokio task to handle async - // invalidate() within synchronous on_drop_stream. - #[allow(clippy::disallowed_methods)] - tokio::spawn(async move { - task_data_entries.invalidate(&key).await; - }); - task_data_metrics.mark_execution_finished(); - if send_metrics { - send_metrics_via_channel(&metrics_tx, &plan, d_ctx, &task_data_metrics); - } - } - }); - streams.push(Box::pin(RecordBatchStreamAdapter::new(stream_schema, stream)) as _); +impl Worker { + pub async fn execute_task( + &self, + request: ExecuteTaskRequest, + ) -> Result<(Vec, Arc)> { + Self::execute_task_static(Arc::clone(&self.task_data_entries), request).await } - Ok((streams, task_ctx)) -} - -/// Builds several per-partition streams by retrieving the appropriate entry from [TaskDataEntries] -/// based on the task key extracted from the gRPC request. -/// -/// This method eagerly starts streaming data from the task, and communicates via channels the -/// produced [RecordBatch]s already encoded as Arrow Flight data. -pub(crate) async fn execute_remote_task( - task_data_entries: &Arc, - request: Request, -) -> Result::ExecuteTaskStream>, Status> { - let body = request.into_inner(); - let partition_range = body.target_partition_start..body.target_partition_end; - let (arrow_streams, task_ctx) = execute_local_task(task_data_entries, body) - .await - .map_err(datafusion_error_to_tonic_status)?; - - let d_cfg = DistributedConfig::from_config_options(task_ctx.session_config().options()) - .map_err(datafusion_error_to_tonic_status)?; + pub(crate) async fn execute_task_static( + task_data_entries: Arc, + request: ExecuteTaskRequest, + ) -> Result<(Vec, Arc)> { + let entry = task_data_entries + .get_with(request.task_key, async { Default::default() }) + .await; + + // Other request is responsible for writing the plan that belongs to this TaskKey, so + // we'll resolve immediately if it was already there, or wait until it's ready. + let task_data = entry + .read(Duration::from_secs(WAIT_PLAN_TIMEOUT_SECS)) + .await + .map_err(|e| exec_datafusion_err!("Worker::execute_task timed-out while waiting for the plan to be set by the coordinator. ({e})"))? + .map_err(DataFusionError::Shared)?; + task_data.task_data_metrics.mark_execution_started_once(); + + let plan = task_data.plan(&request.producer_head_spec)?; + let task_ctx = task_data.task_ctx; + let partition_count = plan.properties().partitioning.partition_count(); + let plan_name = plan.name(); + + // Execute all the requested partitions at once, and collect all the streams so that they + // can be merged into a single one at the end of this function. + let n_streams = request.target_partition_end - request.target_partition_start; + let mut streams = Vec::with_capacity(n_streams); + for partition in request.target_partition_start..request.target_partition_end { + if partition >= partition_count { + return exec_err!( + "partition {partition} not available. The head plan {plan_name} of the stage just has {partition_count} partitions" + ); + } - let compression = match d_cfg.compression.as_str() { - "lz4" => Some(CompressionType::LZ4_FRAME), - "zstd" => Some(CompressionType::ZSTD), - "none" => None, - v => Err(Status::invalid_argument(format!( - "Unknown compression type {v}" - )))?, - }; - let mut flight_streams = Vec::with_capacity(arrow_streams.len()); - for (partition, arrow_stream) in partition_range.zip(arrow_streams) { - let flight_stream = build_flight_data_stream(arrow_stream, compression)?.map(move |msg| { - // For each FlightData produced by this stream, mark it with the appropriate - // partition. This stream will be merged with several others from other partitions, - // so marking it with the original partition allows it to be deconstructed into - // the original per-partition streams in later steps. - let flight_data = FlightAppMetadata { - partition, - created_timestamp_unix_nanos: now_ns(), - }; - msg.map(|v| v.with_app_metadata(flight_data.encode_to_vec())) - }); + let stream = plan.execute(partition, Arc::clone(&task_ctx))?; + let stream_schema = plan.schema(); - flight_streams.push(flight_stream); + streams.push(Box::pin(RecordBatchStreamAdapter::new(stream_schema, stream)) as _); + } + Ok((streams, task_ctx)) } - - // Merge all the per-partition streams into one. Each message in the stream is marked with - // the original partition, so they can be reconstructed at the other side of the boundary. - let memory_pool = Arc::clone(&task_ctx.runtime_env().memory_pool); - let stream = spawn_select_all(flight_streams, memory_pool, RECORD_BATCH_BUFFER_SIZE); - - Ok(Response::new(Box::pin(stream.map_err(|err| match err { - FlightError::Tonic(status) => *status, - _ => Status::internal(format!("Error during flight stream: {err}")), - })))) -} - -fn build_flight_data_stream( - stream: SendableRecordBatchStream, - compression_type: Option, -) -> Result { - let stream = FlightDataEncoderBuilder::new() - .with_options( - IpcWriteOptions::default() - .try_with_compression(compression_type) - .map_err(|err| Status::internal(err.to_string()))?, - ) - .with_schema(stream.schema()) - // This tells the encoder to send dictionaries across the wire as-is. - // The alternative (`DictionaryHandling::Hydrate`) would expand the dictionaries - // into their value types, which can potentially blow up the size of the data transfer. - // The main reason to use `DictionaryHandling::Hydrate` is for compatibility with clients - // that do not support dictionaries, but since we are using the same server/client on both - // sides, we can safely use `DictionaryHandling::Resend`. - // Note that we do garbage collection of unused dictionary values above, so we are not sending - // unused dictionary values over the wire. - .with_dictionary_handling(DictionaryHandling::Resend) - // Set max flight data size to unlimited. - // This requires servers and clients to also be configured to handle unlimited sizes. - // Using unlimited sizes avoids splitting RecordBatches into multiple FlightData messages, - // which could add significant overhead for large RecordBatches. - // The only reason to split them really is if the client/server are configured with a message size limit, - // which mainly makes sense in a public network scenario where you want to avoid DoS attacks. - // Since all of our Arrow Flight communication happens within trusted data plane networks, - // we can safely use unlimited sizes here. - .with_max_flight_data_size(usize::MAX) - .build( - stream - // Apply garbage collection of dictionary and view arrays before sending over the network - .and_then(|rb| std::future::ready(garbage_collect_arrays(rb))) - .map_err(|err| FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(err)))), - ); - Ok(stream) -} - -/// Collects metrics from the plan in pre-order traversal order and sends them via the -/// coordinator channel oneshot. -fn send_metrics_via_channel( - metrics_tx: &Arc>>>, - plan: &Arc, - dt_ctx: DistributedTaskContext, - task_data_metrics: &Arc, -) { - let mut pre_order_plan_metrics = vec![]; - let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { - pre_order_plan_metrics.push( - node.metrics() - .and_then(|m| df_metrics_set_to_proto(&m).ok()) - .unwrap_or_default(), - ); - Ok(TreeNodeRecursion::Continue) - }); - - let tx = { - let mut guard = match metrics_tx.lock() { - Ok(g) => g, - Err(_) => return, - }; - guard.take() - }; - let Some(tx) = tx else { return }; - // Ignore send errors — the coordinator channel may have been dropped (e.g. query cancelled). - let _ = tx.send(TaskMetrics { - pre_order_plan_metrics, - task_metrics: Some(task_data_metrics.to_proto_metrics_set()), - }); -} - -/// Garbage collects values sub-arrays. -/// -/// We apply this before sending RecordBatches over the network to avoid sending -/// values that are not referenced by any dictionary keys or buffers that are not used. -/// -/// Unused values can arise from operations such as filtering, where -/// some keys may no longer be referenced in the filtered result. -fn garbage_collect_arrays(batch: RecordBatch) -> Result { - let (schema, arrays, row_count) = batch.into_parts(); - - let arrays = arrays - .into_iter() - .map(|array| { - if let Some(array) = array.as_any_dictionary_opt() { - garbage_collect_any_dictionary(array) - } else if let Some(array) = array.as_string_view_opt() { - Ok(Arc::new(array.gc()) as Arc) - } else if let Some(array) = array.as_binary_view_opt() { - Ok(Arc::new(array.gc()) as Arc) - } else { - Ok(array) - } - }) - .collect::, _>>()?; - - Ok(RecordBatch::try_new_with_options( - schema, - arrays, - &RecordBatchOptions::new().with_row_count(Some(row_count)), - )?) } diff --git a/src/worker/mod.rs b/src/worker/mod.rs index e89921fc..6e8ce7b9 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -1,9 +1,7 @@ -pub(crate) mod generated; mod impl_coordinator_channel; mod impl_execute_task; mod session_builder; mod single_write_multi_read; -mod spawn_select_all; mod task_data; #[cfg(any(test, feature = "integration"))] pub(crate) mod test_utils; diff --git a/src/worker/single_write_multi_read.rs b/src/worker/single_write_multi_read.rs index b9632352..30c4ce08 100644 --- a/src/worker/single_write_multi_read.rs +++ b/src/worker/single_write_multi_read.rs @@ -63,6 +63,7 @@ impl SingleWriteMultiRead { } /// Reads the current value, if any, not waiting for it to be set by a writer. + #[cfg(feature = "grpc")] pub(crate) fn read_now(&self) -> Option { self.rx.borrow().clone() } diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index 97b2e806..878c2f52 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -1,12 +1,13 @@ -use crate::MaxLatencyMetric; use crate::common::OnceLockResult; use crate::common::now_ns; use crate::distributed_planner::ProducerHead; -use crate::worker::generated::worker as pb; +use crate::protocol::ProducerHeadSpec; +use crate::{MaxLatencyMetric, TaskMetrics}; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; -use datafusion::physical_expr_common::metrics::CustomMetricValue; +use datafusion::physical_plan::metrics::{Metric, MetricValue, MetricsSet}; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use std::borrow::Cow; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; @@ -17,7 +18,7 @@ use tokio::sync::oneshot; /// by concurrent requests for the same task which execute separate partitions. pub struct TaskData { /// Task context suitable for execute different partitions from the same task. - pub(super) task_ctx: Arc, + pub(crate) task_ctx: Arc, pub(crate) base_plan: Arc, pub(crate) final_plan: Arc>>, /// `num_partitions_remaining` is initialized to the total number of partitions in the task (not @@ -25,11 +26,11 @@ pub struct TaskData { /// for this task. Once this count is zero, the task is likely complete. The task may not be /// complete because it's possible that the same partition was retried and this count was /// decremented more than once for the same partition. - pub(super) num_partitions_remaining: Arc, - /// Sender half of the metrics channel. `impl_execute_task` takes this (via `Option::take`) - /// once all partitions have finished or been dropped, sending the collected metrics back to - /// the coordinator through the `CoordinatorChannel` side channel. - pub(super) metrics_tx: Arc>>>, + pub(crate) num_partitions_remaining: Arc, + /// Sender half of the metrics channel. `impl_coordinator_channel` takes this (via + /// `Option::take`) when the coordinator channel reaches EOS, sending the collected metrics + /// back to the coordinator through the `CoordinatorChannel` side channel. + pub(super) metrics_tx: Arc>>>, /// Metrics related to the execution of a task within a stage. This metrics, instead of being /// associated to a specific node, they are global to the task, like the time at which the plan /// was fed by the coordinator to the worker. @@ -42,7 +43,7 @@ pub(crate) const PLAN_FINISHED_AT_METRIC: &str = "plan_finished_at"; #[derive(Debug)] pub(super) struct TaskDataMetrics { - pub(super) query_start_time_ns: u64, + pub(super) query_start_time_ns: usize, /// When the plan was set by the coordinator. pub(super) plan_added_at: MaxLatencyMetric, /// When the plan execution was triggered by the parent worker. @@ -52,10 +53,10 @@ pub(super) struct TaskDataMetrics { } impl TaskDataMetrics { - pub(super) fn new(query_start_time_ns: u64) -> Self { + pub(super) fn new(query_start_time_ns: usize) -> Self { let plan_added_at = MaxLatencyMetric::default(); plan_added_at.add_duration(Duration::from_nanos( - now_ns().saturating_sub(query_start_time_ns), + now_ns::().saturating_sub(query_start_time_ns as u64), )); Self { query_start_time_ns, @@ -68,72 +69,57 @@ impl TaskDataMetrics { pub(super) fn mark_execution_started_once(&self) { if self.plan_executed_at.value() == 0 { self.plan_executed_at.add_duration(Duration::from_nanos( - now_ns().saturating_sub(self.query_start_time_ns), + now_ns::().saturating_sub(self.query_start_time_ns as u64), )) } } pub(super) fn mark_execution_finished(&self) { self.plan_finished_at.add_duration(Duration::from_nanos( - now_ns().saturating_sub(self.query_start_time_ns), + now_ns::().saturating_sub(self.query_start_time_ns as u64), )) } - pub(super) fn to_proto_metrics_set(&self) -> pb::MetricsSet { - let mut task_metrics_set = pb::MetricsSet { metrics: vec![] }; - - fn new_metric(name: &str, value: usize) -> pb::Metric { - pb::Metric { - partition: None, - labels: vec![], - value: Some(pb::metric::Value::CustomMaxLatency(pb::MaxLatency { - name: name.to_string(), - value: value as u64, - })), - } - } - task_metrics_set.metrics.push(new_metric( + pub(super) fn to_metrics_set(&self) -> MetricsSet { + let mut metrics_set = MetricsSet::new(); + metrics_set.push(max_latency_metric( PLAN_ADDED_AT_METRIC, - self.plan_added_at.as_usize(), + &self.plan_added_at, )); - task_metrics_set.metrics.push(new_metric( + metrics_set.push(max_latency_metric( PLAN_EXECUTED_AT_METRIC, - self.plan_executed_at.as_usize(), + &self.plan_executed_at, )); - task_metrics_set.metrics.push(new_metric( + metrics_set.push(max_latency_metric( PLAN_FINISHED_AT_METRIC, - self.plan_finished_at.as_usize(), + &self.plan_finished_at, )); - task_metrics_set + metrics_set } } -impl TaskData { - /// Returns the number of partitions remaining to be processed. - pub(crate) fn num_partitions_remaining(&self) -> usize { - self.num_partitions_remaining.load(Ordering::SeqCst) - } - - /// Returns the total number of partitions in this task. - pub(crate) fn total_partitions(&self) -> usize { - match self.final_plan.get() { - Some(Ok(plan)) => plan.output_partitioning().partition_count(), - _ => self - .base_plan - .properties() - .output_partitioning() - .partition_count(), - } - } +fn max_latency_metric(name: &'static str, value: &MaxLatencyMetric) -> Arc { + Arc::new(Metric::new( + MetricValue::Custom { + name: Cow::Borrowed(name), + value: Arc::new(MaxLatencyMetric::from_nanos(value.value())), + }, + None, + )) +} +impl TaskData { pub(crate) fn plan( &self, - producer_head: pb::execute_task_request::ProducerHead, + producer_head_spec: &ProducerHeadSpec, ) -> Result> { let result = self.final_plan.get_or_init(|| { - let producer_head = - ProducerHead::from_proto(producer_head, &self.base_plan.schema(), &self.task_ctx)?; + let producer_head = ProducerHead::from_spec( + producer_head_spec, + self.base_plan.schema(), + &self.task_ctx, + )?; let plan = producer_head.insert(Arc::clone(&self.base_plan))?; diff --git a/src/worker/test_utils/worker_handles.rs b/src/worker/test_utils/worker_handles.rs index 477b1de5..68f95088 100644 --- a/src/worker/test_utils/worker_handles.rs +++ b/src/worker/test_utils/worker_handles.rs @@ -1,7 +1,7 @@ use crate::config_extension_ext::set_distributed_option_extension; -use crate::worker::generated::worker::TaskKey; +use crate::grpc::BoxCloneSyncChannel; use crate::worker::task_data::TaskDataMetrics; -use crate::{BoxCloneSyncChannel, DistributedConfig, DistributedExt, TaskData, Worker}; +use crate::{DistributedConfig, DistributedExt, TaskData, TaskKey, Worker}; use arrow_ipc::CompressionType; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -17,9 +17,9 @@ use tonic::transport::{Endpoint, Server}; use url::Url; use uuid::Uuid; -pub fn test_task_key_with_query(query_id: Uuid, task_number: u64) -> TaskKey { +pub fn test_task_key_with_query(query_id: Uuid, task_number: usize) -> TaskKey { TaskKey { - query_id: query_id.as_bytes().to_vec(), + query_id, stage_id: 0, task_number, } diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 7acd629c..ad761239 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -1,48 +1,25 @@ -use crate::common::{OnceLockResult, on_drop_stream, serialize_uuid}; +use crate::common::OnceLockResult; use crate::distributed_planner::ProducerHead; -use crate::metrics::LatencyMetricExt; -use crate::networking::get_distributed_channel_resolver; use crate::passthrough_headers::get_passthrough_headers; -use crate::protobuf::{datafusion_error_to_tonic_status, map_flight_to_datafusion_error}; use crate::stage::RemoteStage; -use crate::worker::generated::worker::FlightAppMetadata; -use crate::worker::generated::worker::{ExecuteTaskRequest, TaskKey}; -use crate::worker::impl_execute_task::execute_local_task; use crate::worker::worker_service::TaskDataEntries; -use crate::{BytesMetricExt, ChannelResolver, DistributedConfig}; -use arrow_flight::FlightData; -use arrow_flight::decode::FlightRecordBatchStream; -use arrow_flight::error::FlightError; -use dashmap::DashMap; +use crate::{ + ChannelResolver, ExecuteTaskRequest, TaskKey, Worker, get_distributed_channel_resolver, +}; use datafusion::arrow::array::RecordBatch; -use datafusion::common::instant::Instant; use datafusion::common::runtime::SpawnedTask; use datafusion::common::{ - DataFusionError, Result, exec_err, internal_datafusion_err, internal_err, + DataFusionError, Result, exec_datafusion_err, internal_datafusion_err, internal_err, }; use datafusion::execution::TaskContext; -use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, MetricValue}; -use datafusion::physical_plan::metrics::{MetricBuilder, Time}; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::metrics::MetricBuilder; +use futures::future::{BoxFuture, Shared}; use futures::stream::BoxStream; -use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; -use http::Extensions; -use pin_project::{pin_project, pinned_drop}; -use prost::Message; -use std::borrow::Cow; +use futures::{FutureExt, StreamExt, TryFutureExt}; use std::fmt::{Debug, Formatter}; use std::ops::Range; -use std::pin::Pin; -use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; -use std::task::{Context, Poll}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use tokio::sync::Notify; -use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; -use tokio_stream::wrappers::UnboundedReceiverStream; -use tokio_util::sync::CancellationToken; -use tonic::metadata::MetadataMap; -use tonic::{Request, Status}; use url::Url; /// Context set by [crate::Worker::coordinator_channel] in DataFusion's @@ -59,18 +36,22 @@ pub(crate) struct LocalWorkerContext { pub(crate) self_url: Url, } -/// Holds a list of lazily initialized [WorkerConnection]s. Each position in the underlying +/// Holds lazily initialized partition streams. Each position in the underlying /// `connections` vector corresponds to the connection to one worker. It assumes a 1:1 mapping -/// between worker and tasks, and upon calling [WorkerConnectionPool::get_or_init_worker_connection] +/// between worker and tasks, and upon calling [WorkerConnectionPool::execute] /// it will initialize the corresponding position in the vector matching the provided `target_task` /// index. pub(crate) struct WorkerConnectionPool { - connections: Vec>>, + connections: Vec>, pub(crate) metrics: ExecutionPlanMetricsSet, } +type PartitionStreams = Arc>>>>>; +type SharedPartitionStreamsFuture = + Shared>>>; + impl WorkerConnectionPool { - /// Builds a new [WorkerConnectionPool] with as many empty slots for [WorkerConnection]s as + /// Builds a new [WorkerConnectionPool] with as many empty slots for worker stream entries as /// the provided `input_tasks`. pub(crate) fn new(input_tasks: usize) -> Self { let mut connections = Vec::with_capacity(input_tasks); @@ -83,17 +64,17 @@ impl WorkerConnectionPool { } } - /// Lazily initializes the [WorkerConnection] corresponding to the provided `target_task` - /// (therefore maintaining one independent [WorkerConnection] per `target_task`), and - /// returns it. - pub(crate) fn get_or_init_worker_connection( + /// Lazily initializes the partition streams corresponding to the provided `target_task`, and + /// returns the stream for `target_partition`. + pub(crate) fn execute( &self, input_stage: &RemoteStage, target_partitions: Range, target_task: usize, + target_partition: usize, producer_head: ProducerHead, ctx: &Arc, - ) -> Result<&(dyn WorkerConnection + Sync + Send)> { + ) -> Result>> { let Some(worker_connection) = self.connections.get(target_task) else { return internal_err!( "WorkerConnections: Task index {target_task} not found, only have {} tasks", @@ -101,399 +82,39 @@ impl WorkerConnectionPool { ); }; - let conn = worker_connection.get_or_init(|| { - let Some(target_url) = input_stage.workers.get(target_task) else { - internal_err!("input_stage.workers[{target_task}] out of range.")? - }; - if let Some(lw_ctx) = ctx.session_config().get_extension::() - && &lw_ctx.self_url == target_url - { - // Instead of making a gRPC call to ourselves, better to just use local comms. - LocalWorkerConnection::init( - input_stage, - target_partitions, - target_task, - producer_head, - ctx, - &self.metrics, - ) - .map(|v| Box::new(v) as Box<_>) - .map_err(Arc::new) - } else { - // We are trying to reach a URL different from ours, so use normal gRPC streams. - RemoteWorkerConnection::init( - input_stage, - target_partitions, - target_task, - producer_head, - ctx, - &self.metrics, - ) - .map(|v| Box::new(v) as Box<_>) - .map_err(Arc::new) - } - }); - - match conn { - Ok(v) => Ok(v.as_ref()), - Err(err) => Err(DataFusionError::Shared(Arc::clone(err))), - } - } -} - -type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; - -/// Abstraction that allows treating remote and local comms as equal. Network boundaries do not -/// care if the stream comes over the wire or locally. -pub(crate) trait WorkerConnection { - /// Streams the specified partition. Consumers do not care if the implementation pulls data - /// from in-memory or from local comms. - fn execute(&self, partition: usize) -> Result>>; -} - -/// Represents a connection to one [Worker]. Network boundaries will use this for streaming -/// data from single partitions while the actual network communication is handling all the partitions -/// under the hood. -/// -/// This is done so that, rather than issuing one gRPC stream per partition, we issue one gRPC stream -/// per group of partitions, and we multiplex streamed record batches locally to in-memory channels. -/// -/// Even if Tonic can perfectly multiplex and interleave messages from different gRPC streams through -/// the same underlying TCP connection, there do is some overhead in having one gRPC stream per -/// partition VS a single gRPC stream interleaving multiple partitions. The whole serialized plan -/// needs to be sent over the wire on every gRPC call, so the less gRPC calls we do the better. -struct RemoteWorkerConnection { - task: Arc>, - not_consumed_streams: Arc, - cancel_token: CancellationToken, - per_partition_rx: DashMap>, - - first_poll_notify: Arc, - // Signals the demux task that buffered memory has been freed by a consumer. - mem_available_notify: Arc, - - // Metrics collection stuff. - memory_reservation: Arc, - elapsed_compute: Time, -} + let streams_shared_future = worker_connection.get_or_init(|| { + let ch_resolver = get_distributed_channel_resolver(ctx.as_ref()); -impl RemoteWorkerConnection { - fn init( - input_stage: &RemoteStage, - target_partition_range: Range, - target_task: usize, - producer_head: ProducerHead, - ctx: &Arc, - metrics: &ExecutionPlanMetricsSet, - ) -> Result { - let channel_resolver = get_distributed_channel_resolver(ctx.as_ref()); - let buffer_budget_bytes = - DistributedConfig::from_config_options(ctx.session_config().options())? - .worker_connection_buffer_budget_bytes; - // We are retaining record batches in memory until they are consumed, so we need to account - // for them in the memory pool. - let memory_reservation = - Arc::new(MemoryConsumer::new("WorkerConnection").register(ctx.memory_pool())); - let memory_reservation_clone = Arc::clone(&memory_reservation); - - // Track the maximum memory used to buffer recieved messages. - let mut curr_max_mem = 0; - let max_mem_used = MetricBuilder::new(metrics).global_gauge("max_mem_used"); - // Track the total encoded size of all recieved messages. - let bytes_transferred = MetricBuilder::new(metrics).bytes_counter("bytes_transferred"); - let msg_count = MetricBuilder::new(metrics).global_counter("msg_count"); - // Track end-to-end network latency distribution for all messages. - let min_latency = MetricBuilder::new(metrics).min_latency("network_latency_min"); - let max_latency = MetricBuilder::new(metrics).max_latency("network_latency_max"); - let p50_latency = MetricBuilder::new(metrics).p50_latency("network_latency_p50"); - let p95_latency = MetricBuilder::new(metrics).p95_latency("network_latency_p95"); - let first_latency = MetricBuilder::new(metrics).first_latency("network_latency_first"); - let sum_latency = Time::new(); - MetricBuilder::new(metrics).build(MetricValue::Time { - name: Cow::Borrowed("network_latency_sum"), - time: sum_latency.clone(), - }); - let latency_count = MetricBuilder::new(metrics).counter("network_latency_count", 0); - // Track the total CPU time spent in polling messages over the network + decoding them. - let elapsed_compute = Time::new(); - let elapsed_compute_clone = elapsed_compute.clone(); - MetricBuilder::new(metrics).build(MetricValue::ElapsedCompute(elapsed_compute.clone())); - - // Building the actual request that will be sent to the worker. - let headers = get_passthrough_headers(ctx.session_config()); - let request = Request::from_parts( - MetadataMap::from_headers(headers), - Extensions::default(), - ExecuteTaskRequest { - target_partition_start: target_partition_range.start as u64, - target_partition_end: target_partition_range.end as u64, - task_key: Some(TaskKey { - query_id: serialize_uuid(&input_stage.query_id), - stage_id: input_stage.num as u64, - task_number: target_task as u64, - }), - producer_head: Some(producer_head.to_proto(ctx)?), - }, - ); - - let Some(url) = input_stage.workers.get(target_task).cloned() else { - return internal_err!("ProgrammingError: Task {target_task} not found"); - }; - - // The senders and receivers are unbounded queues used for multiplexing the record - // batches sent through the single gRPC stream into one stream per partition. They - // are unbounded to avoid head-of-line blocking: a single bounded queue could block - // the demux task and starve all sibling partitions even though they have capacity, - // which deadlocks queries with cross-partition dependencies. - // Total memory is bounded globally below via `mem_available_notify`. - let mut per_partition_tx = Vec::with_capacity(target_partition_range.len()); - let per_partition_rx = DashMap::with_capacity(target_partition_range.len()); - for partition in target_partition_range.clone() { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel::(); - per_partition_tx.push(tx); - per_partition_rx.insert(partition, rx); - } - - let mem_available_notify = Arc::new(Notify::new()); - let mem_available_notify_for_task = Arc::clone(&mem_available_notify); - - let first_poll_notify = Arc::new(Notify::new()); - let first_poll_notify_for_task = Arc::clone(&first_poll_notify); - - // Cancellation token allows us to stop the background task promptly when all partition - // streams are dropped (e.g., when the query is cancelled). - let cancel_token = CancellationToken::new(); - let cancel = cancel_token.clone(); - - // This task will pull data from all the partitions in `target_partition_range`, and will - // fan them out to the appropriate `per_partition_rx` based on the "partition" declared - // in each individual record batch flight metadata. - let task = SpawnedTask::spawn(async move { - let mut client = match channel_resolver.get_worker_client_for_url(&url).await { - Ok(v) => v, - Err(err) => return fanout(&per_partition_tx, datafusion_error_to_tonic_status(&err)) + let Some(target_url) = input_stage.workers.get(target_task).cloned() else { + internal_err!("input_stage.workers[{target_task}] out of range.")? }; - - tokio::select! { - biased; - _ = cancel.cancelled() => { - // If all SendableRecordBatchStreams canceled before any poll, we need to - // anyway trigger the task execution and cancel it immediately so that the - // cancellation is propagated also in the remote worker. Otherwise, it might - // hang forever waiting for someone to execute it. - let _ = client.execute_task(request).await; - return - }, - _ = first_poll_notify_for_task.notified() => {} + let local_task_data_entries = ctx + .session_config() + .get_extension::() + .and_then(|lw_ctx| match lw_ctx.self_url == target_url { + true => Some(Arc::clone(&lw_ctx.task_data_entries)), + false => None, + }); + if local_task_data_entries.is_some() { + MetricBuilder::new(&self.metrics) + .global_counter("local_connections_used") + .add(1); } - let mut interleaved_stream = match client.execute_task(request).await { - Ok(v) => v.into_inner(), - Err(err) => return fanout(&per_partition_tx, err), + let task_key = TaskKey { + query_id: input_stage.query_id, + stage_id: input_stage.num, + task_number: target_task, }; - loop { - // Backpressure gate. Per-partition channels are unbounded, so we cap - // total in-flight buffered bytes here by pausing the gRPC pull when - // consumers haven't drained enough. This propagates flow control all - // the way back to the worker without coupling sibling partitions. - // We always allow a message through when reservation == 0 to avoid - // livelock if a single message is larger than the budget. - while memory_reservation.size() >= buffer_budget_bytes { - tokio::select! { - biased; - _ = cancel.cancelled() => return, - _ = mem_available_notify_for_task.notified() => {} - } - } - - // Check for cancellation while waiting for the next message. - let flight_data = tokio::select! { - biased; - _ = cancel.cancelled() => return, - msg = interleaved_stream.next() => { - match msg { - Some(Ok(v)) => v, - Some(Err(err)) => return fanout(&per_partition_tx, err), - None => return, // Stream exhausted - } - } - }; - - // Earliest time at which the msg was received. - let msg_received_time = SystemTime::now(); - - let flight_metadata = match FlightAppMetadata::decode(flight_data.app_metadata.as_ref()) { - Ok(v) => v, - Err(err) => { - return fanout(&per_partition_tx, Status::internal(err.to_string())); - } - }; - - // Update the running latency tracker. - let sent_time = UNIX_EPOCH + Duration::from_nanos(flight_metadata.created_timestamp_unix_nanos); - if flight_metadata.created_timestamp_unix_nanos > 0 - && let Ok(delta) = msg_received_time.duration_since(sent_time) { - min_latency.add_duration(delta); - max_latency.add_duration(delta); - p50_latency.add_duration(delta); - p95_latency.add_duration(delta); - first_latency.add_duration(delta); - sum_latency.add_duration(delta); - latency_count.add(1); - } - - let partition = flight_metadata.partition as usize; - // the `per_partition_tx` variable is using a normal `Vec` for storing the - // channel transmitters, so we need to subtract the `target_partition_range.start` - // to the `partition` in order to offset it to the appropriate index. - let sender_i = partition - target_partition_range.start; - - let Some(o_tx) = per_partition_tx.get(sender_i) else { - let msg = format!( - "Received partition {partition} in Flight metadata, but available partitions are {target_partition_range:?}" - ); - return fanout(&per_partition_tx, Status::internal(msg)); - }; - - // We need to send the memory reservation in the same tuple as the actual message - // so that it gets dropped as soon as the message leaves the queue. Dropping the - // memory reservation means releasing the memory from the pool for that specific - // message - let size = flight_data.encoded_len(); - memory_reservation.grow(size); - - // Update memory related metrics. - msg_count.add(1); - bytes_transferred.add_bytes(size); - let curr_mem = memory_reservation.size(); - if curr_mem > curr_max_mem { - curr_max_mem = curr_mem; - max_mem_used.set(curr_max_mem); - } - - if o_tx.send(Ok((flight_data, flight_metadata))).is_err() { - // The receiver for this partition was dropped (e.g. a hash join partition - // completed early without consuming its probe side). Don't exit: other - // partitions multiplexed over the same gRPC stream still need their data. - // Undo the memory reservation that was grown for this dropped batch. - memory_reservation.shrink(size); - continue; - }; - } - }.with_elapsed_compute(elapsed_compute)); - - Ok(Self { - task: Arc::new(task), - cancel_token, - not_consumed_streams: Arc::new(AtomicUsize::new(per_partition_rx.len())), - per_partition_rx, - mem_available_notify, - first_poll_notify, - - // metrics stuff - memory_reservation: memory_reservation_clone, - elapsed_compute: elapsed_compute_clone, - }) - } -} - -impl WorkerConnection for RemoteWorkerConnection { - /// Streams the provided `partition` from the remote worker. - /// - /// This method does not handle any network connection. Instead, the network comms are delegated - /// to the task spawned by [WorkerConnection::init], who is in charge of polling data not only - /// from the requested `partition`, but from any other partition in `target_partition_range`. - /// This method just streams all the record batches belonging to the provided `partition` from - /// an in-memory queue. - /// - /// The task that polls data over the network is held inactive until the first poll to the - /// stream returned by this method. - /// - /// When the returned stream is dropped (e.g., due to query cancellation), the background task - /// pulling from the Flight stream will be canceled promptly. - fn execute(&self, partition: usize) -> Result>> { - let Some((_, partition_receiver)) = self.per_partition_rx.remove(&partition) else { - return internal_err!( - "WorkerConnection has no stream for target partition {partition}. Was it already consumed?" - ); - }; - let task = Arc::clone(&self.task); - let cancel_token = self.cancel_token.clone(); - - let first_poll_notify = Arc::clone(&self.first_poll_notify); - let stream = async move { - first_poll_notify.notify_one(); - UnboundedReceiverStream::new(partition_receiver) - } - .flatten_stream(); - - let stream = stream.map_err(|err| FlightError::Tonic(Box::new(err))); - let reservation = Arc::clone(&self.memory_reservation); - let mem_available_notify = Arc::clone(&self.mem_available_notify); - let stream = stream.map_ok(move |(data, _meta)| { - reservation.shrink(data.encoded_len()); - // Wake the demux task in case it is blocked on the byte budget. - mem_available_notify.notify_one(); - let _ = &task; // <- keep the task that polls data from the network alive. - data - }); - let stream = FlightRecordBatchStream::new_from_flight_data(stream); - let stream = stream.map_err(map_flight_to_datafusion_error); - let stream = stream.with_elapsed_compute(self.elapsed_compute.clone()); - - // When the stream is dropped, cancel the background task to ensure prompt cleanup. - let not_consumed_streams = Arc::clone(&self.not_consumed_streams); - Ok(on_drop_stream(stream, move || { - let remaining_streams = not_consumed_streams.fetch_sub(1, Ordering::SeqCst) - 1; - if remaining_streams == 0 { - cancel_token.cancel(); - } - }) - .boxed()) - } -} - -/// Equivalent to [RemoteWorkerConnection], but that pulls data from the local registry of tasks -/// rather than doing it across a gRPC interface. -pub(crate) struct LocalWorkerConnection { - partition_start: usize, - local_streams: Vec>>>>, -} - -impl LocalWorkerConnection { - fn init( - input_stage: &RemoteStage, - target_partition_range: Range, - target_task: usize, - producer_head: ProducerHead, - ctx: &Arc, - metrics: &ExecutionPlanMetricsSet, - ) -> Result { - MetricBuilder::new(metrics) - .global_counter("local_connections_used") - .add(1); - let Some(lw_ctx) = ctx.session_config().get_extension::() else { - return exec_err!("Missing LocalWorkerContext extension"); - }; - - let task_key = TaskKey { - query_id: serialize_uuid(&input_stage.query_id), - stage_id: input_stage.num as u64, - task_number: target_task as u64, - }; - - let partition_start = target_partition_range.start; - let mut local_streams = Vec::with_capacity(target_partition_range.len()); - for partition_i in target_partition_range { let request = ExecuteTaskRequest { - task_key: Some(task_key.clone()), - target_partition_start: partition_i as u64, - target_partition_end: (partition_i + 1) as u64, - producer_head: Some(producer_head.to_proto(ctx)?), + task_key, + target_partition_start: target_partitions.start, + target_partition_end: target_partitions.end, + producer_head_spec: producer_head.to_spec(ctx.session_config())?, }; - - let task_data_entries = Arc::clone(&lw_ctx.task_data_entries); + let metrics = self.metrics.clone(); + let ctx = Arc::clone(ctx); // 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, @@ -501,57 +122,55 @@ impl LocalWorkerConnection { // is polled, the entry might not be there. // // Note that this does not start polling the returned streams, it just instantiates them. - let streams_future = SpawnedTask::spawn(async move { - let (streams, _) = execute_local_task(&task_data_entries, request).await?; - Ok::<_, DataFusionError>(streams) + let streams_task = SpawnedTask::spawn(async move { + if let Some(task_data_entries) = local_task_data_entries { + let (streams, _) = + Worker::execute_task_static(task_data_entries, request).await?; + let streams = streams.into_iter().map(|v| v.boxed()).collect(); + Ok::<_, DataFusionError>(streams) + } else { + let mut client = ch_resolver.get_worker_client_for_url(&target_url).await?; + let headers = get_passthrough_headers(ctx.session_config()); + let streams = client.execute_task(headers, request, metrics, &ctx).await?; + Ok(streams) + } }); - let stream = async move { - let mut streams = streams_future - .await - .map_err(|err| internal_datafusion_err!("{err}"))??; - if streams.len() != 1 { - return internal_err!("Expected exactly 1 local stream"); + Ok(async move { + match streams_task.await { + Ok(Ok(v)) => Ok(Arc::new( + v.into_iter().map(|v| Mutex::new(Some(v))).collect(), + )), + Ok(Err(e)) => Err(Arc::new(e)), + Err(e) => Err(Arc::new(exec_datafusion_err!( + "JoinError instantiating streams {e}" + ))), } - Ok(streams.swap_remove(0)) } - .try_flatten_stream() - .boxed(); - - local_streams.push(Mutex::new(Some(stream))); - } - - Ok(Self { - partition_start, - local_streams, - }) - } -} + .boxed() + .shared()) + }); -impl WorkerConnection for LocalWorkerConnection { - fn execute(&self, partition: usize) -> Result>> { - let Some(relative_i) = partition.checked_sub(self.partition_start) else { - return internal_err!( - "LocalWorkerConnection received an invalid partition {partition}, the starting partition is {}", - self.partition_start - ); - }; - let Some(slot) = self.local_streams.get(relative_i) else { - return internal_err!( - "LocalWorkerConnection has no stream for partition {partition}. Was it already consumed?" - ); + let streams_future = match streams_shared_future { + Ok(v) => v.clone(), + Err(err) => return Err(DataFusionError::Shared(Arc::clone(err))), }; - slot.lock().unwrap().take().ok_or_else(|| { - internal_datafusion_err!( - "LocalWorkerConnection stream for partition {partition} was already consumed" - ) - }) - } -} -fn fanout(o_txs: &[UnboundedSender], err: Status) { - for o_tx in o_txs { - let _ = o_tx.send(Err(err.clone())); + Ok(async move { + let streams = streams_future.await.map_err(DataFusionError::Shared)?; + let Some(slot) = streams.get(target_partition - target_partitions.start) else { + return internal_err!( + "WorkerConnections has no stream for partition {target_partition}. Was it already consumed?" + ); + }; + slot.lock().unwrap().take().ok_or_else(|| { + internal_datafusion_err!( + "WorkerConnections stream for partition {target_partition} was already consumed" + ) + }) + } + .try_flatten_stream() + .boxed()) } } @@ -568,197 +187,3 @@ impl Clone for WorkerConnectionPool { Self::new(self.connections.len()) } } - -impl Debug for RemoteWorkerConnection { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WorkerConnection").finish() - } -} - -trait ElapsedComputeFutureExt: Future + Sized { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture; -} - -trait ElapsedComputeStreamExt: Stream + Sized { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream; -} - -impl> ElapsedComputeFutureExt for F { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture { - ElapsedComputeFuture { - inner: self, - curr: Duration::default(), - elapsed_compute, - } - } -} - -impl> ElapsedComputeStreamExt for S { - fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream { - ElapsedComputeStream { - inner: self, - curr: Duration::default(), - elapsed_compute, - } - } -} - -#[pin_project(PinnedDrop)] -struct ElapsedComputeStream { - #[pin] - inner: T, - curr: Duration, - elapsed_compute: Time, -} - -/// Drop implementation that ensures that any accumulated time is properly dumped to the metric -/// in case the stream gets dropped before completion. -#[pinned_drop] -impl PinnedDrop for ElapsedComputeStream { - fn drop(self: Pin<&mut Self>) { - if self.curr > Duration::default() { - let self_projected = self.project(); - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - } - } -} - -impl> Stream for ElapsedComputeStream { - type Item = O; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let self_projected = self.project(); - let start = Instant::now(); - let result = self_projected.inner.poll_next(cx); - *self_projected.curr += start.elapsed(); - if result.is_ready() { - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - *self_projected.curr = Duration::default(); - } - result - } -} - -#[pin_project(PinnedDrop)] -struct ElapsedComputeFuture { - #[pin] - inner: T, - curr: Duration, - elapsed_compute: Time, -} - -/// Drop implementation that ensures that any accumulated time is properly dumped to the metric -/// in case the future gets dropped before completion. -#[pinned_drop] -impl PinnedDrop for ElapsedComputeFuture { - fn drop(self: Pin<&mut Self>) { - if self.curr > Duration::default() { - let self_projected = self.project(); - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - } - } -} - -impl> Future for ElapsedComputeFuture { - type Output = O; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let self_projected = self.project(); - let start = Instant::now(); - let result = self_projected.inner.poll(cx); - *self_projected.curr += start.elapsed(); - if result.is_ready() { - self_projected - .elapsed_compute - .add_duration(*self_projected.curr); - *self_projected.curr = Duration::default(); - } - result - } -} - -#[cfg(test)] -mod tests { - use super::*; - use futures::StreamExt; - use futures::stream::unfold; - - #[tokio::test] - async fn elapsed_compute_future() { - async fn cheap() { - tokio::time::sleep(Duration::from_millis(1)).await; - } - - async fn expensive() { - let mut _count = 0f64; - for i in 0..100000 { - tokio::task::yield_now().await; - _count /= i as f64 - } - } - - let cheap_time = Time::new(); - cheap().with_elapsed_compute(cheap_time.clone()).await; - println!("cheap future: {}", cheap_time.value()); - - let expensive_time = Time::new(); - expensive() - .with_elapsed_compute(expensive_time.clone()) - .await; - println!("expensive future: {}", expensive_time.value()); - - assert!(expensive_time.value() > cheap_time.value()); - } - - #[tokio::test] - async fn elapsed_compute_stream() { - fn cheap() -> impl Stream { - unfold(0i64, |state| async move { - if state < 10 { - tokio::time::sleep(Duration::from_micros(10)).await; - Some((state, state + 1)) - } else { - None - } - }) - } - - fn expensive() -> impl Stream { - unfold(0i64, |state| async move { - if state < 10 { - // Simulate expensive computation - let mut _count = 0f64; - for i in 1..100000 { - _count += (i as f64).sqrt(); - } - tokio::task::yield_now().await; - Some((state, state + 1)) - } else { - None - } - }) - } - - let cheap_time = Time::new(); - cheap() - .with_elapsed_compute(cheap_time.clone()) - .collect::>() - .await; - println!("cheap future: {}", cheap_time.value()); - - let expensive_time = Time::new(); - expensive() - .with_elapsed_compute(expensive_time.clone()) - .collect::>() - .await; - println!("expensive future: {}", expensive_time.value()); - - assert!(expensive_time.value() > cheap_time.value()); - } -} diff --git a/src/worker/worker_service.rs b/src/worker/worker_service.rs index 5d9c375f..3a285f11 100644 --- a/src/worker/worker_service.rs +++ b/src/worker/worker_service.rs @@ -1,17 +1,5 @@ -use crate::worker::WorkerSessionBuilder; -use crate::worker::generated::worker::worker_service_server::{WorkerService, WorkerServiceServer}; -use crate::worker::generated::worker::{ - CoordinatorToWorkerMsg, ExecuteTaskRequest, TaskKey, WorkerToCoordinatorMsg, -}; -use crate::worker::impl_execute_task::execute_remote_task; -use crate::worker::single_write_multi_read::SingleWriteMultiRead; -use crate::worker::task_data::TaskData; -use crate::{ - DefaultSessionBuilder, GetWorkerInfoRequest, GetWorkerInfoResponse, ObservabilityServiceImpl, - ObservabilityServiceServer, WorkerResolver, -}; -use arrow_flight::FlightData; -use async_trait::async_trait; +use crate::worker::{SingleWriteMultiRead, WorkerSessionBuilder}; +use crate::{DefaultSessionBuilder, TaskData, TaskKey}; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::ExecutionPlan; @@ -20,8 +8,6 @@ use moka::future::Cache; use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; -use tonic::codegen::BoxStream; -use tonic::{Request, Response, Status, Streaming}; const TASK_CACHE_TTI: Duration = Duration::from_mins(10); @@ -44,10 +30,10 @@ pub struct Worker { /// TTL-based cache for task execution data. Entries are automatically evicted after /// TASK_CACHE_TTI seconds. This prevents memory leaks from abandoned or incomplete queries /// while allowing concurrent access to task results across multiple partition requests. - pub(super) task_data_entries: Arc, + pub(crate) task_data_entries: Arc, pub(super) session_builder: Arc, pub(super) hooks: WorkerHooks, - pub(super) max_message_size: Option, + pub(crate) max_message_size: Option, pub(super) version: Cow<'static, str>, } @@ -124,60 +110,17 @@ impl Worker { self } - /// Converts this [Worker] into a [`WorkerServiceServer`] with high default message size limits. - /// - /// This is a convenience method that wraps the endpoint in a [`WorkerServiceServer`] and - /// configures it with `max_decoding_message_size(usize::MAX)` and - /// `max_encoding_message_size(usize::MAX)` to avoid message size limitations for internal - /// communication. - /// - /// You can further customize the returned server by chaining additional tonic methods. - /// - /// # Example - /// - /// ``` - /// # use datafusion_distributed::Worker; - /// # use tonic::transport::Server; - /// # use std::net::{IpAddr, Ipv4Addr, SocketAddr}; - /// # async fn f() { - /// - /// let worker = Worker::default(); - /// let server = worker.into_worker_server(); - /// - /// Server::builder() - /// .add_service(Worker::default().into_worker_server()) - /// .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)) - /// .await; - /// - /// # } - /// ``` - pub fn into_worker_server(self) -> WorkerServiceServer { - WorkerServiceServer::new(self) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX) - } - - /// Creates an [`ObservabilityServiceServer`] that exposes task progress and cluster - /// worker discovery via the provided [`WorkerResolver`]. - /// - /// The returned server is meant to be added to the same [`tonic::transport::Server`] as the - /// Flight service — gRPC multiplexes both services on a single port. - pub fn with_observability_service( - &self, - worker_resolver: Arc, - ) -> ObservabilityServiceServer { - ObservabilityServiceServer::new(ObservabilityServiceImpl::new( - self.task_data_entries.clone(), - worker_resolver, - )) - } - /// Sets a version string reported by the `GetWorkerInfo` gRPC endpoint. pub fn with_version(mut self, version: impl Into>) -> Self { self.version = version.into(); self } + /// Returns the version set by [Self::with_version]. + pub fn version(&self) -> &str { + &self.version + } + /// Returns the number of cached task entries currently held by this worker. #[cfg(any(test, feature = "integration"))] pub async fn tasks_running(&self) -> usize { @@ -187,37 +130,3 @@ impl Worker { self.task_data_entries.entry_count() as usize } } - -/// Implementation of the `worker.proto` specification based on the generated Rust stubs. -/// -/// The methods are delegated to plan `impl Worker` implementations so that they can be implemented -/// in different files. -#[async_trait] -impl WorkerService for Worker { - type CoordinatorChannelStream = BoxStream; - - async fn coordinator_channel( - &self, - request: Request>, - ) -> Result, Status> { - self.impl_coordinator_channel(request).await - } - - type ExecuteTaskStream = BoxStream; - - async fn execute_task( - &self, - request: Request, - ) -> Result, Status> { - execute_remote_task(&self.task_data_entries, request).await - } - - async fn get_worker_info( - &self, - _request: Request, - ) -> Result, Status> { - Ok(Response::new(GetWorkerInfoResponse { - version: self.version.to_string(), - })) - } -} diff --git a/src/networking/worker_resolver.rs b/src/worker_resolver.rs similarity index 100% rename from src/networking/worker_resolver.rs rename to src/worker_resolver.rs From 7406c01fbe17bbf217376a9661a7122fa1ea942b Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 14:37:05 -0700 Subject: [PATCH 12/20] Added an in-process worker transport that needs no gRPC. The worker-protocol abstraction shipped only a gRPC implementation, so nothing exercised it without gRPC. This routes the three protocol methods straight to a co-located `Worker`, and is the seam a shared-memory transport plugs into. A worker resolves its own nested stage reads back through the same transport, so multi-stage queries stay in process. --- src/lib.rs | 6 +- src/protocol/channel_resolver.rs | 3 + src/protocol/in_process.rs | 230 +++++++++++++++++++++++++++++++ src/protocol/mod.rs | 2 + 4 files changed, 238 insertions(+), 3 deletions(-) create mode 100644 src/protocol/in_process.rs diff --git a/src/lib.rs b/src/lib.rs index b714fe41..ce401c01 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,9 +45,9 @@ pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; /// TODO: do not expose this yet. pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, - GetWorkerInfoResponse, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, - WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, WorkerToCoordinatorMsg, - get_distributed_channel_resolver, + GetWorkerInfoResponse, InProcessChannelResolver, LoadInfo, ProducerHeadSpec, SetPlanRequest, + TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, + WorkerToCoordinatorMsg, get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, diff --git a/src/protocol/channel_resolver.rs b/src/protocol/channel_resolver.rs index ac16fefc..09b8208e 100644 --- a/src/protocol/channel_resolver.rs +++ b/src/protocol/channel_resolver.rs @@ -78,6 +78,9 @@ pub fn get_distributed_channel_resolver( #[cfg(not(feature = "grpc"))] { + // With `grpc` off there is no built-in default. A co-located deployment can register + // [`crate::InProcessChannelResolver`] (or another transport) via + // `with_distributed_channel_resolver`. panic!( "gRPC feature is not enabled, and no channel resolver was provided, so no default ChannelResolver can be provided" ); diff --git a/src/protocol/in_process.rs b/src/protocol/in_process.rs new file mode 100644 index 00000000..9a59085d --- /dev/null +++ b/src/protocol/in_process.rs @@ -0,0 +1,230 @@ +//! The reference implementation of the worker protocol for a co-located worker. +//! +//! It implements [`WorkerChannel`] by calling a [`Worker`] that lives in the same process, with no +//! gRPC, no IPC, and no networking. Its purpose is twofold: it lets the crate run distributed +//! queries with the `grpc` feature off (the protocol abstraction is only real if something other +//! than gRPC implements it), and it is the shape a custom co-located transport (for example a +//! shared-memory mesh spanning sibling processes) follows: implement [`ChannelResolver`] to hand +//! out a channel for a URL, then route the three protocol methods to wherever the worker runs. + +use crate::protocol::{ChannelResolver, WorkerChannel}; +use crate::{ + CoordinatorToWorkerMsg, DefaultSessionBuilder, DistributedExt, ExecuteTaskRequest, + GetWorkerInfoRequest, GetWorkerInfoResponse, MappedWorkerSessionBuilderExt, Worker, + WorkerSessionBuilder, WorkerToCoordinatorMsg, +}; +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::stream::{BoxStream, StreamExt}; +use http::HeaderMap; +use std::sync::{Arc, Weak}; +use url::Url; + +/// A [`ChannelResolver`] backed by a single co-located [`Worker`]. +/// +/// Every URL resolves to that one worker: with no network there is nothing to dial, so the URLs a +/// [`crate::WorkerResolver`] hands out only label the tasks the planner routes. One worker holds the +/// state for every task, keyed by [`crate::TaskKey`], the same way the gRPC worker does when several +/// tasks of a query land on it. +#[derive(Clone)] +pub struct InProcessChannelResolver { + worker: Arc, +} + +impl InProcessChannelResolver { + /// Builds the co-located worker from `session_builder`, registering this resolver into the + /// worker's own per-query sessions so a worker reading a downstream stage stays in process + /// rather than falling back to the gRPC default. + pub fn from_session_builder( + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + // The worker and the resolver point at each other: the resolver runs tasks on the worker, + // and the worker resolves its own nested reads back through the resolver. A `Weak` on the + // worker's side breaks the cycle, so the returned `InProcessChannelResolver` owns the only + // strong reference and dropping it frees the worker. + let worker = Arc::new_cyclic(|weak: &Weak| { + let weak = weak.clone(); + Worker::from_session_builder(session_builder.map(move |builder| { + Ok(builder + .with_distributed_channel_resolver(WeakInProcessChannelResolver(weak.clone())) + .build()) + })) + }); + Self { worker } + } +} + +impl Default for InProcessChannelResolver { + fn default() -> Self { + Self::from_session_builder(DefaultSessionBuilder) + } +} + +#[async_trait] +impl ChannelResolver for InProcessChannelResolver { + async fn get_worker_client_for_url( + &self, + _url: &Url, + ) -> Result, DataFusionError> { + Ok(Box::new(InProcessWorkerChannel { + worker: Arc::clone(&self.worker), + })) + } +} + +/// The resolver a worker installs in its own sessions. It upgrades a [`Weak`] reference to the +/// co-located worker so a read of a downstream stage routes back in process instead of dialing out. +struct WeakInProcessChannelResolver(Weak); + +#[async_trait] +impl ChannelResolver for WeakInProcessChannelResolver { + async fn get_worker_client_for_url( + &self, + _url: &Url, + ) -> Result, DataFusionError> { + let worker = self + .0 + .upgrade() + .ok_or_else(|| internal_datafusion_err!("the in-process worker has been dropped"))?; + Ok(Box::new(InProcessWorkerChannel { worker })) + } +} + +/// A [`WorkerChannel`] that calls a co-located [`Worker`] directly. +struct InProcessWorkerChannel { + worker: Arc, +} + +#[async_trait] +impl WorkerChannel for InProcessWorkerChannel { + async fn coordinator_channel( + &mut self, + headers: HeaderMap, + c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>> { + // The worker reads a fallible stream so a wire transport can surface its decode errors. + // Handing messages over in process has no such step, so each one is already an `Ok`. + self.worker + .coordinator_channel(headers, c2w_stream.map(Ok).boxed()) + .await + } + + async fn execute_task( + &mut self, + _headers: HeaderMap, + request: ExecuteTaskRequest, + _metrics: ExecutionPlanMetricsSet, + _task_ctx: &Arc, + ) -> Result>>> { + // Reading a partition runs the producer in place: the returned streams are the worker's own + // task output, so there is no IPC decode pass and no network metrics to record. The + // consumer's `task_ctx` is the consumer side's; the producer runs under the worker's own. + let (streams, _task_ctx) = self.worker.execute_task(request).await?; + Ok(streams.into_iter().map(|stream| stream.boxed()).collect()) + } + + async fn get_worker_info( + &mut self, + _request: GetWorkerInfoRequest, + ) -> Result { + Ok(GetWorkerInfoResponse { + version: self.worker.version().to_string(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{SessionStateBuilderExt, WorkerResolver, display_plan_ascii}; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::collect; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use std::io::Write; + + /// Hands out as many placeholder URLs as workers. With one co-located worker behind the + /// transport, these only label the tasks the planner routes; nothing is dialed. + struct InProcessWorkers(usize); + + impl WorkerResolver for InProcessWorkers { + fn get_urls(&self) -> Result, DataFusionError> { + (0..self.0) + .map(|i| Url::parse(&format!("http://worker-{i}"))) + .collect::>() + .map_err(|err| DataFusionError::External(Box::new(err))) + } + } + + /// Drives a real distributed query end to end through the in-process transport. With the `grpc` + /// feature off this is the only transport that can run it; `cargo check --no-default-features` + /// covers the no-gRPC compile. + #[tokio::test] + async fn in_process_transport_runs_a_distributed_query() -> Result<()> { + const N: usize = 4; + + // A file scan round-trips through `datafusion-proto`, so a worker can rebuild it from the + // serialized stage plan. An in-memory table would not, hence a CSV on disk. + let path = std::env::temp_dir().join(format!("dfd_in_process_{}.csv", std::process::id())); + let mut file = + std::fs::File::create(&path).map_err(|e| DataFusionError::External(Box::new(e)))?; + writeln!(file, "k,v").unwrap(); + for i in 0..200 { + writeln!(file, "{},{}", ["a", "b", "c", "d"][i % 4], i).unwrap(); + } + drop(file); + let path = path.to_str().unwrap().to_string(); + + let query = "SELECT k, COUNT(*) AS c FROM t GROUP BY k ORDER BY k"; + + // Single-node reference result. + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(N)); + ctx.register_csv("t", &path, CsvReadOptions::new()).await?; + let expected = collect( + ctx.sql(query).await?.create_physical_plan().await?, + ctx.task_ctx(), + ) + .await?; + let expected = pretty_format_batches(&expected)?.to_string(); + + // Distributed over the in-process transport. + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(N)) + .with_distributed_planner() + .with_distributed_worker_resolver(InProcessWorkers(N)) + .with_distributed_channel_resolver(InProcessChannelResolver::default()) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + let ctx_distributed = SessionContext::from(state); + ctx_distributed + .register_csv("t", &path, CsvReadOptions::new()) + .await?; + + let physical = ctx_distributed + .sql(query) + .await? + .create_physical_plan() + .await?; + let rendered = display_plan_ascii(physical.as_ref(), false); + assert!( + rendered.contains("DistributedExec"), + "plan was not distributed:\n{rendered}" + ); + assert!( + rendered.contains("NetworkShuffleExec"), + "no shuffle boundary, so the transport never carried a cross-task stream:\n{rendered}" + ); + + let actual = collect(physical, ctx_distributed.task_ctx()).await?; + let actual = pretty_format_batches(&actual)?.to_string(); + + let _ = std::fs::remove_file(&path); + assert_eq!(actual, expected); + Ok(()) + } +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 5c6bdd1a..9dd5137c 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -2,10 +2,12 @@ pub mod grpc; mod channel_resolver; +mod in_process; mod worker_channel; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; +pub use in_process::InProcessChannelResolver; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, From 085050139e0b2961c19e2384f24cf94492c1ee94 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:11:48 -0700 Subject: [PATCH 13/20] Exposed the worker protocol off gRPC for co-located transports. A co-located transport (in-process worker, shared-memory mesh) speaks the same wire shape as gRPC, so the prost messages and the metrics codec have to build with the feature off and crate::proto must resolve either way. The plan-metrics collector and the receive-time stamp are the seams such a transport drives directly. --- src/lib.rs | 9 +++++++-- src/protocol/generated/mod.rs | 1 + src/protocol/{grpc => }/generated/worker.rs | 2 ++ src/protocol/grpc/gen/src/main.rs | 2 +- src/protocol/grpc/generated/mod.rs | 1 - src/protocol/grpc/mod.rs | 2 -- .../grpc/observability/gen/src/main.rs | 2 +- .../observability/generated/observability.rs | 2 +- src/protocol/grpc/observability/service.rs | 2 +- src/protocol/grpc/worker_client.rs | 6 +++--- src/protocol/grpc/worker_service.rs | 4 ++-- src/protocol/{grpc => }/metrics_proto.rs | 4 ++++ src/protocol/mod.rs | 5 +++++ src/work_unit_feed/mod.rs | 5 +++-- src/work_unit_feed/remote_work_unit_feed.rs | 7 +++++++ src/worker/impl_coordinator_channel.rs | 20 +++++++++++++++++++ src/worker/mod.rs | 1 + 17 files changed, 59 insertions(+), 16 deletions(-) create mode 100644 src/protocol/generated/mod.rs rename src/protocol/{grpc => }/generated/worker.rs (99%) delete mode 100644 src/protocol/grpc/generated/mod.rs rename src/protocol/{grpc => }/metrics_proto.rs (99%) diff --git a/src/lib.rs b/src/lib.rs index ce401c01..6b21af0d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -39,6 +39,10 @@ pub mod test_utils; #[cfg(feature = "grpc")] pub use protocol::grpc; +/// The worker-protocol prost message types, independent of any transport. A non-gRPC transport +/// reaches for these to speak the same wire shape the gRPC path serializes. +pub use protocol::generated::worker as proto; + pub use codec::DistributedCodec; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; @@ -53,11 +57,12 @@ pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, }; pub use work_unit_feed::{ - DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, + DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, + WorkUnitFeedProvider, set_received_time, }; pub use worker::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, - Worker, WorkerQueryContext, WorkerSessionBuilder, + Worker, WorkerQueryContext, WorkerSessionBuilder, collect_plan_metrics_protos, }; #[cfg(any(feature = "integration", test))] diff --git a/src/protocol/generated/mod.rs b/src/protocol/generated/mod.rs new file mode 100644 index 00000000..2c8b8399 --- /dev/null +++ b/src/protocol/generated/mod.rs @@ -0,0 +1 @@ +pub mod worker; diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/generated/worker.rs similarity index 99% rename from src/protocol/grpc/generated/worker.rs rename to src/protocol/generated/worker.rs index 290261ed..1f104a7b 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/generated/worker.rs @@ -468,6 +468,7 @@ pub struct MaxGauge { pub value: u64, } /// Generated client implementations. +#[cfg(feature = "grpc")] pub mod worker_service_client { #![allow( unused_variables, @@ -617,6 +618,7 @@ pub mod worker_service_client { } } /// Generated server implementations. +#[cfg(feature = "grpc")] pub mod worker_service_server { #![allow( unused_variables, diff --git a/src/protocol/grpc/gen/src/main.rs b/src/protocol/grpc/gen/src/main.rs index 7505aae5..f9cf2a90 100644 --- a/src/protocol/grpc/gen/src/main.rs +++ b/src/protocol/grpc/gen/src/main.rs @@ -6,7 +6,7 @@ fn main() -> Result<(), Box> { let proto_dir = repo_root.join("src/protocol/grpc"); let proto_file = proto_dir.join("worker.proto"); - let out_dir = repo_root.join("src/protocol/grpc/generated"); + let out_dir = repo_root.join("src/protocol/generated"); fs::create_dir_all(&out_dir)?; diff --git a/src/protocol/grpc/generated/mod.rs b/src/protocol/grpc/generated/mod.rs deleted file mode 100644 index 844c269c..00000000 --- a/src/protocol/grpc/generated/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) mod worker; diff --git a/src/protocol/grpc/mod.rs b/src/protocol/grpc/mod.rs index e432b607..cc2f6576 100644 --- a/src/protocol/grpc/mod.rs +++ b/src/protocol/grpc/mod.rs @@ -1,7 +1,5 @@ mod channel_resolver; mod errors; -mod generated; -mod metrics_proto; mod observability; mod on_drop_stream; mod spawn_select_all; diff --git a/src/protocol/grpc/observability/gen/src/main.rs b/src/protocol/grpc/observability/gen/src/main.rs index 124f0f1b..1b55fb7d 100644 --- a/src/protocol/grpc/observability/gen/src/main.rs +++ b/src/protocol/grpc/observability/gen/src/main.rs @@ -21,7 +21,7 @@ fn main() -> Result<(), Box> { .out_dir(&out_dir) .extern_path( ".observability.TaskKey", - "crate::protocol::grpc::generated::worker::TaskKey", + "crate::protocol::generated::worker::TaskKey", ) .compile_protos(&[proto_file], &[proto_dir])?; diff --git a/src/protocol/grpc/observability/generated/observability.rs b/src/protocol/grpc/observability/generated/observability.rs index 7706dd21..1cfba3d3 100644 --- a/src/protocol/grpc/observability/generated/observability.rs +++ b/src/protocol/grpc/observability/generated/observability.rs @@ -12,7 +12,7 @@ pub struct GetTaskProgressRequest {} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct TaskProgress { #[prost(message, optional, tag = "1")] - pub task_key: ::core::option::Option, + pub task_key: ::core::option::Option, #[prost(uint64, tag = "2")] pub total_partitions: u64, #[prost(uint64, tag = "3")] diff --git a/src/protocol/grpc/observability/service.rs b/src/protocol/grpc/observability/service.rs index b7cc2120..ce325c3d 100644 --- a/src/protocol/grpc/observability/service.rs +++ b/src/protocol/grpc/observability/service.rs @@ -4,7 +4,7 @@ use super::{ }; use crate::common::serialize_uuid; use crate::grpc::{GetClusterWorkersRequest, GetClusterWorkersResponse}; -use crate::protocol::grpc::generated::worker as worker_pb; +use crate::protocol::generated::worker as worker_pb; use crate::worker::{SingleWriteMultiRead, TaskData}; use crate::{TaskKey, WorkerResolver}; use datafusion::error::DataFusionError; diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 2af90c59..9e84c5fe 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -1,10 +1,10 @@ use super::channel_resolver::BoxCloneSyncChannel; use super::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; -use super::generated::worker as pb; -use super::metrics_proto::metrics_set_proto_to_df; use crate::common::serialize_uuid; -use crate::grpc::generated::worker::FlightAppMetadata; use crate::grpc::on_drop_stream::on_drop_stream; +use crate::protocol::generated::worker as pb; +use crate::protocol::generated::worker::FlightAppMetadata; +use crate::protocol::metrics_proto::metrics_set_proto_to_df; use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index fa677a10..87c1ef4b 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -1,7 +1,7 @@ use super::errors::{datafusion_error_to_tonic_status, map_status_to_datafusion_error}; -use super::generated::worker as pb; -use super::metrics_proto::df_metrics_set_to_proto; use super::spawn_select_all::spawn_select_all; +use crate::protocol::generated::worker as pb; +use crate::protocol::metrics_proto::df_metrics_set_to_proto; use crate::common::{deserialize_uuid, now_ns}; use crate::protocol::ProducerHeadSpec; diff --git a/src/protocol/grpc/metrics_proto.rs b/src/protocol/metrics_proto.rs similarity index 99% rename from src/protocol/grpc/metrics_proto.rs rename to src/protocol/metrics_proto.rs index 6db8cda9..27b715f2 100644 --- a/src/protocol/grpc/metrics_proto.rs +++ b/src/protocol/metrics_proto.rs @@ -43,6 +43,9 @@ pub fn df_metrics_set_to_proto( } /// metrics_set_proto_to_df converts a [pb::MetricsSet] to a [MetricsSet]. +// The decode direction has no caller when `grpc` is off: the gRPC client consumes it, and so will a +// push transport that receives metrics frames. +#[cfg_attr(not(feature = "grpc"), allow(dead_code))] pub fn metrics_set_proto_to_df( metrics_set_proto: &pb::MetricsSet, ) -> Result { @@ -283,6 +286,7 @@ pub fn df_metric_to_proto(metric: Arc) -> Result Result, DataFusionError> { let partition = metric.partition.map(|p| p as usize); let labels = metric diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 9dd5137c..a08a3eb4 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -2,7 +2,12 @@ pub mod grpc; mod channel_resolver; +// The prost message types and the metrics codec carry no tonic dependency, so a non-gRPC transport +// (the in-process worker, a shared-memory mesh) can speak the same wire shape without pulling in the +// whole gRPC stack. +pub(crate) mod generated; mod in_process; +pub(crate) mod metrics_proto; mod worker_channel; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index 1bd55ebf..ff9fdb6a 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -6,11 +6,12 @@ mod work_unit_feed_provider; mod work_unit_feed_registry; pub(crate) use remote_work_unit_feed::{ - RemoteWorkUnitFeedRegistry, build_work_unit_batch_msg, set_work_unit_received_time, - set_work_unit_send_time, + RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedTxs, build_work_unit_batch_msg, + set_work_unit_received_time, set_work_unit_send_time, }; pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed}; +pub use remote_work_unit_feed::set_received_time; pub use work_unit::WorkUnit; pub use work_unit_feed::{WorkUnitFeed, WorkUnitFeedProto}; pub use work_unit_feed_provider::{DistributedWorkUnitFeedContext, WorkUnitFeedProvider}; diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index 18c51ffe..78ecc42a 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -99,6 +99,13 @@ pub(crate) fn set_work_unit_received_time( msg } +/// Stamps the receive time on a bare proto work unit. A transport that hands units across without +/// going through [set_work_unit_received_time] stamps each one as it crosses into the worker, so the +/// worker-side latency math has a delivery timestamp; a missing stamp reads as zero latency. +pub fn set_received_time(work_unit: &mut crate::proto::WorkUnit) { + work_unit.received_timestamp_unix_nanos = now_ns(); +} + /// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnitMsg]s coming over /// the wire from a [RemoteWorkUnitFeedRegistry]. /// diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index b60edc37..3c02d1bf 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,5 +1,6 @@ use crate::common::TreeNodeExt; use crate::execution_plans::SamplerExec; +use crate::protocol::metrics_proto::df_metrics_set_to_proto; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; use crate::worker::task_data::TaskDataMetrics; @@ -228,3 +229,22 @@ fn send_metrics_via_channel( task_metrics: task_data_metrics.to_metrics_set(), }); } + +/// Collects the per-node metrics of `plan` in pre-order traversal order, encoded as protos. A push +/// transport reuses this for its metrics frame so the per-node shape matches what the coordinator +/// channel reports. +pub fn collect_plan_metrics_protos( + plan: &Arc, + dt_ctx: DistributedTaskContext, +) -> Vec { + let mut pre_order_plan_metrics = vec![]; + let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { + pre_order_plan_metrics.push( + node.metrics() + .and_then(|m| df_metrics_set_to_proto(&m).ok()) + .unwrap_or_default(), + ); + Ok(TreeNodeRecursion::Continue) + }); + pre_order_plan_metrics +} diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 6e8ce7b9..6f9dfc83 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod test_utils; mod worker_connection_pool; mod worker_service; +pub use impl_coordinator_channel::collect_plan_metrics_protos; pub(crate) use single_write_multi_read::SingleWriteMultiRead; pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool}; From 02ff23ba6ee59a5e3f997f37cb27ba8b9376ba68 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:46:51 -0700 Subject: [PATCH 14/20] Made the no-gRPC test build compile and runnable in CI. The benchmarks dev-dependency and a few `tonic`-importing test helpers re-unified gRPC into every test build, so nothing exercised the in-process transport with gRPC off. Moving the dataset suites into the benchmarks crate and gating the gRPC-coupled helpers lets `cargo test --no-default-features --lib` run the in-process transport, guarded by a new CI job. --- .github/workflows/ci.yml | 31 ++- Cargo.lock | 3 +- Cargo.toml | 42 ++- benchmarks/Cargo.toml | 11 + .../tests}/clickbench_correctness_test.rs | 2 +- .../tests}/clickbench_plans_test.rs | 2 +- .../tests}/stateful_data_cleanup.rs | 2 +- .../tests}/tpcds_correctness_test.rs | 2 +- .../tests}/tpcds_plans_test.rs | 2 +- .../tests}/tpch_correctness_test.rs | 2 +- .../tests}/tpch_plans_test.rs | 2 +- src/coordinator/metrics_store.rs | 2 +- src/execution_plans/metrics.rs | 2 +- src/execution_plans/mod.rs | 4 +- src/lib.rs | 2 +- src/metrics/task_metrics_collector.rs | 3 +- src/metrics/task_metrics_rewriter.rs | 3 +- src/test_utils/in_memory_channel_resolver.rs | 254 ++++++++++-------- src/test_utils/mod.rs | 1 + src/test_utils/routing.rs | 2 +- src/work_unit_feed/mod.rs | 7 +- src/worker/mod.rs | 4 +- 22 files changed, 233 insertions(+), 152 deletions(-) rename {tests => benchmarks/tests}/clickbench_correctness_test.rs (99%) rename {tests => benchmarks/tests}/clickbench_plans_test.rs (99%) rename {tests => benchmarks/tests}/stateful_data_cleanup.rs (98%) rename {tests => benchmarks/tests}/tpcds_correctness_test.rs (99%) rename {tests => benchmarks/tests}/tpcds_plans_test.rs (99%) rename {tests => benchmarks/tests}/tpch_correctness_test.rs (99%) rename {tests => benchmarks/tests}/tpch_plans_test.rs (99%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3c5da3fd..7ecf07e7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,6 +41,17 @@ jobs: - uses: ./.github/actions/setup - run: cargo test --features integration + # Builds the lib with `grpc` off and runs the in-memory and shared-memory transport tests over it, + # the proof that the abstraction runs without gRPC. Scoped to those tests on purpose: the rest of + # the lib's unit tests load parquet testdata (git-LFS), which is the gRPC `unit-test` job's domain. + # The full lib still compiles with `grpc` off here, since the test binary builds every module. + unit-test-no-grpc: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/setup + - run: cargo test --no-default-features --lib -- protocol::in_process shm::in_process + tpch-correctness-test: runs-on: ubuntu-latest strategy: @@ -50,7 +61,7 @@ jobs: steps: - uses: actions/checkout@v4 - uses: ./.github/actions/setup - - run: cargo test --features tpch --test tpch_correctness_test + - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_correctness_test env: ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }} @@ -59,7 +70,7 @@ jobs: steps: - uses: actions/checkout@v4 - uses: ./.github/actions/setup - - run: cargo test --features tpch --test tpch_plans_test + - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test tpch_plans_test tpcds-correctness-test: runs-on: ubuntu-latest @@ -73,9 +84,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/tpcds/main.zip + path: benchmarks/testdata/tpcds/main.zip key: "main.zip" - - run: cargo test --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }} + - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }} env: ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }} @@ -86,9 +97,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/tpcds/main.zip + path: benchmarks/testdata/tpcds/main.zip key: "main.zip" - - run: cargo test --features tpcds --test tpcds_plans_test + - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_plans_test clickbench-correctness-test: runs-on: ubuntu-latest @@ -101,9 +112,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/clickbench/ + path: benchmarks/testdata/clickbench/ key: "data" - - run: cargo test --features clickbench --test clickbench_correctness_test + - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_correctness_test env: ADAPTIVE: ${{ matrix.planning_mode == 'adaptive' }} @@ -114,9 +125,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/clickbench/ + path: benchmarks/testdata/clickbench/ key: "data" - - run: cargo test --features clickbench --test clickbench_plans_test + - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test clickbench_plans_test format-check: runs-on: ubuntu-latest diff --git a/Cargo.lock b/Cargo.lock index 7044adaa..3858be13 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2199,7 +2199,6 @@ dependencies = [ "crossbeam-queue", "dashmap", "datafusion", - "datafusion-distributed-benchmarks", "datafusion-proto", "delegate", "futures", @@ -2249,11 +2248,13 @@ dependencies = [ "object_store", "openssl", "parquet", + "pretty_assertions", "reqwest", "serde", "serde_json", "structopt", "sysinfo", + "test-case", "tokio", "tonic", "tpchgen", diff --git a/Cargo.toml b/Cargo.toml index 79366ed3..6b2f0c24 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,14 +75,9 @@ integration = ["insta", "parquet", "arrow", "hyper-util", "grpc"] system-metrics = ["sysinfo"] -tpch = ["integration"] -tpcds = ["integration"] -clickbench = ["integration"] -slow-tests = [] sysinfo = ["dep:sysinfo"] [dev-dependencies] -datafusion-distributed-benchmarks = { path = "benchmarks" } structopt = "0.3" insta = { version = "1.46.0", features = ["filters"] } parquet = "58" @@ -92,6 +87,43 @@ hyper-util = "0.1.16" pretty_assertions = "1.4" test-case = "3.3.1" +# Every example drives the gRPC transport, so a no-grpc build skips them. +[[example]] +name = "custom_distributed_partial_reduction_tree" +required-features = ["integration"] + +[[example]] +name = "custom_execution_plan" +required-features = ["integration"] + +[[example]] +name = "custom_worker_url_routing" +required-features = ["integration"] + +[[example]] +name = "work_unit_feed" +required-features = ["integration"] + +[[example]] +name = "in_memory_cluster" +required-features = ["grpc"] + +[[example]] +name = "localhost_run" +required-features = ["grpc"] + +[[example]] +name = "localhost_versioned_run" +required-features = ["grpc"] + +[[example]] +name = "localhost_worker" +required-features = ["grpc"] + +[[example]] +name = "localhost_versioned_worker" +required-features = ["grpc"] + [workspace.lints.clippy] disallowed_types = "deny" disallowed-methods = "deny" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 0a03ead0..18197f42 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -33,9 +33,20 @@ aws-sdk-ec2 = "1" openssl = { version = "0.10", features = ["vendored"] } # Keep this. Necessary for the remote benchmarks worker. mimalloc = "0.1" +[features] +# Gates for the dataset test suites under `tests/`. They live here rather than in the library +# because, as a dev-dependency of the library, this crate re-enables `grpc` on every test build +# through feature unification, which made the no-grpc config untestable. +tpch = [] +tpcds = [] +clickbench = [] +slow-tests = [] + [dev-dependencies] criterion = "0.5" sysinfo = "0.30" +pretty_assertions = "1.4" +test-case = "3.3.1" [build-dependencies] built = { version = "0.8", features = ["git2", "chrono"] } diff --git a/tests/clickbench_correctness_test.rs b/benchmarks/tests/clickbench_correctness_test.rs similarity index 99% rename from tests/clickbench_correctness_test.rs rename to benchmarks/tests/clickbench_correctness_test.rs index acf3daa8..25fb38f8 100644 --- a/tests/clickbench_correctness_test.rs +++ b/benchmarks/tests/clickbench_correctness_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "clickbench", test))] +#[cfg(all(feature = "clickbench", test))] mod tests { use datafusion::arrow::array::RecordBatch; use datafusion::common::plan_err; diff --git a/tests/clickbench_plans_test.rs b/benchmarks/tests/clickbench_plans_test.rs similarity index 99% rename from tests/clickbench_plans_test.rs rename to benchmarks/tests/clickbench_plans_test.rs index 2c202201..910194b1 100644 --- a/tests/clickbench_plans_test.rs +++ b/benchmarks/tests/clickbench_plans_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "clickbench", test))] +#[cfg(all(feature = "clickbench", test))] mod tests { use datafusion::error::Result; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; diff --git a/tests/stateful_data_cleanup.rs b/benchmarks/tests/stateful_data_cleanup.rs similarity index 98% rename from tests/stateful_data_cleanup.rs rename to benchmarks/tests/stateful_data_cleanup.rs index fae892d2..f977c61a 100644 --- a/tests/stateful_data_cleanup.rs +++ b/benchmarks/tests/stateful_data_cleanup.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpch", test))] +#[cfg(all(feature = "tpch", test))] mod tests { use datafusion::common::instant::Instant; use datafusion::error::Result; diff --git a/tests/tpcds_correctness_test.rs b/benchmarks/tests/tpcds_correctness_test.rs similarity index 99% rename from tests/tpcds_correctness_test.rs rename to benchmarks/tests/tpcds_correctness_test.rs index e4baeba9..8d5d8722 100644 --- a/tests/tpcds_correctness_test.rs +++ b/benchmarks/tests/tpcds_correctness_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpcds", test))] +#[cfg(all(feature = "tpcds", test))] mod tests { use datafusion::arrow::array::RecordBatch; use datafusion::common::plan_err; diff --git a/tests/tpcds_plans_test.rs b/benchmarks/tests/tpcds_plans_test.rs similarity index 99% rename from tests/tpcds_plans_test.rs rename to benchmarks/tests/tpcds_plans_test.rs index 667ed4d1..04a6fd0a 100644 --- a/tests/tpcds_plans_test.rs +++ b/benchmarks/tests/tpcds_plans_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpcds", test))] +#[cfg(all(feature = "tpcds", test))] mod tests { use datafusion::error::Result; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; diff --git a/tests/tpch_correctness_test.rs b/benchmarks/tests/tpch_correctness_test.rs similarity index 99% rename from tests/tpch_correctness_test.rs rename to benchmarks/tests/tpch_correctness_test.rs index 3ae089dc..7d6ab2ce 100644 --- a/tests/tpch_correctness_test.rs +++ b/benchmarks/tests/tpch_correctness_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpch", test))] +#[cfg(all(feature = "tpch", test))] mod tests { use datafusion::physical_plan::execute_stream; use datafusion::prelude::SessionContext; diff --git a/tests/tpch_plans_test.rs b/benchmarks/tests/tpch_plans_test.rs similarity index 99% rename from tests/tpch_plans_test.rs rename to benchmarks/tests/tpch_plans_test.rs index 231ee3b0..3e4a719a 100644 --- a/tests/tpch_plans_test.rs +++ b/benchmarks/tests/tpch_plans_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpch", test))] +#[cfg(all(feature = "tpch", test))] mod tests { use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; use datafusion_distributed::{ diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index 8ccc1f7c..ed55db2c 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -27,7 +27,7 @@ impl MetricsStore { self.rx.borrow().get(key).cloned() } - #[cfg(test)] + #[cfg(all(test, feature = "grpc"))] pub(crate) fn from_entries(entries: impl IntoIterator) -> Self { let map: HashMap<_, _> = entries.into_iter().collect(); let (tx, rx) = watch::channel(map); diff --git a/src/execution_plans/metrics.rs b/src/execution_plans/metrics.rs index 0fefdf3e..de97ed04 100644 --- a/src/execution_plans/metrics.rs +++ b/src/execution_plans/metrics.rs @@ -21,7 +21,7 @@ impl MetricsWrapperExec { Self { inner, metrics } } - #[cfg(test)] + #[cfg(all(test, feature = "grpc"))] pub(crate) fn inner(&self) -> &Arc { &self.inner } diff --git a/src/execution_plans/mod.rs b/src/execution_plans/mod.rs index ecdcbc35..97904ef1 100644 --- a/src/execution_plans/mod.rs +++ b/src/execution_plans/mod.rs @@ -8,7 +8,9 @@ mod network_coalesce; mod network_shuffle; mod sampler; -#[cfg(any(test, feature = "integration"))] +// The benchmark fixtures spin up real gRPC workers (`tonic` servers, Flight channels), so they +// only exist when that transport is compiled in. +#[cfg(all(feature = "grpc", any(test, feature = "integration")))] pub mod benchmarks; pub use broadcast::BroadcastExec; diff --git a/src/lib.rs b/src/lib.rs index 6b21af0d..8389405d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -65,7 +65,7 @@ pub use worker::{ Worker, WorkerQueryContext, WorkerSessionBuilder, collect_plan_metrics_protos, }; -#[cfg(any(feature = "integration", test))] +#[cfg(all(feature = "grpc", any(feature = "integration", test)))] pub use execution_plans::benchmarks::{ LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, ShuffleBench, ShuffleFixture, TransportBench, TransportBenchMode, TransportFixture, diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index 2d78998d..b613dcfc 100644 --- a/src/metrics/task_metrics_collector.rs +++ b/src/metrics/task_metrics_collector.rs @@ -19,7 +19,8 @@ pub fn collect_plan_metrics(plan: &Arc) -> Result Self { - Self::from_configured_worker(builder, |worker| worker) - } - - /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder] and worker setup. - pub fn from_configured_worker( - builder: impl WorkerSessionBuilder + Send + Sync + 'static, - configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, - ) -> Self { - let (client, server) = tokio::io::duplex(1024 * 1024); - - let mut client = Some(client); - let channel = Endpoint::try_from(DUMMY_URL) - .expect("Invalid dummy URL for building an endpoint. This should never happen") - .connect_with_connector_lazy(tower::service_fn(move |_| { - let client = client - .take() - .expect("Client taken twice. This should never happen"); - async move { Ok::<_, std::io::Error>(TokioIo::new(client)) } - })); - - let this = Self { - channel: grpc::BoxCloneSyncChannel::new(channel), - }; - let this_clone = this.clone(); - - let endpoint = Worker::from_session_builder(builder.map(move |builder| { - let this = this.clone(); - Ok(builder.with_distributed_channel_resolver(this).build()) - })); - let endpoint = configure_worker(endpoint); - - #[allow(clippy::disallowed_methods)] - tokio::spawn(async move { - Server::builder() - .add_service(endpoint.into_worker_server()) - .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server))) - .await - }); - - this_clone - } -} - -impl Default for InMemoryChannelResolver { - fn default() -> Self { - Self::from_session_builder(DefaultSessionBuilder) - } -} - -#[async_trait] -impl ChannelResolver for InMemoryChannelResolver { - async fn get_worker_client_for_url( - &self, - _: &url::Url, - ) -> Result, DataFusionError> { - Ok(grpc::create_worker_client(self.channel.clone())) - } -} - pub struct InMemoryWorkerResolver { n_workers: usize, } @@ -105,41 +23,139 @@ impl WorkerResolver for InMemoryWorkerResolver { } } -/// Creates a distributed session context backed by a single in-memory worker service. -/// The set of produced worker URLs is deterministic, taking the form http://worker-. -pub async fn start_in_memory_context( - num_workers: usize, - session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, -) -> SessionContext { - let channel_resolver = InMemoryChannelResolver::from_session_builder(session_builder); - let state = SessionStateBuilder::new() - .with_default_features() - .with_distributed_planner() - .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) - .with_distributed_channel_resolver(channel_resolver) - .build(); - SessionContext::from(state) -} +// The channel resolver runs a real `tonic` worker server over a tokio duplex, so it only compiles +// with the gRPC transport. The worker resolver above stays available without it, so the planner +// tests that only need worker URLs keep building in a no-gRPC config. +#[cfg(feature = "grpc")] +mod channel { + use super::InMemoryWorkerResolver; + use crate::{ + ChannelResolver, DefaultSessionBuilder, DistributedExt, MappedWorkerSessionBuilderExt, + SessionStateBuilderExt, Worker, WorkerChannel, WorkerSessionBuilder, grpc, + }; + use async_trait::async_trait; + use datafusion::common::DataFusionError; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::{SessionConfig, SessionContext}; + use hyper_util::rt::TokioIo; + use tonic::transport::{Endpoint, Server}; + + const DUMMY_URL: &str = "http://localhost:50051"; + + /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory + /// tokio duplex rather than a TCP connection. + #[derive(Clone)] + pub struct InMemoryChannelResolver { + channel: grpc::BoxCloneSyncChannel, + } + + impl InMemoryChannelResolver { + /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder]. + /// This allows you to inject your own DataFusion extensions in the in-memory worker + /// spawned by this method. + pub fn from_session_builder( + builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + Self::from_configured_worker(builder, |worker| worker) + } + + /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder] and worker setup. + pub fn from_configured_worker( + builder: impl WorkerSessionBuilder + Send + Sync + 'static, + configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, + ) -> Self { + let (client, server) = tokio::io::duplex(1024 * 1024); + + let mut client = Some(client); + let channel = Endpoint::try_from(DUMMY_URL) + .expect("Invalid dummy URL for building an endpoint. This should never happen") + .connect_with_connector_lazy(tower::service_fn(move |_| { + let client = client + .take() + .expect("Client taken twice. This should never happen"); + async move { Ok::<_, std::io::Error>(TokioIo::new(client)) } + })); + + let this = Self { + channel: grpc::BoxCloneSyncChannel::new(channel), + }; + let this_clone = this.clone(); + + let endpoint = Worker::from_session_builder(builder.map(move |builder| { + let this = this.clone(); + Ok(builder.with_distributed_channel_resolver(this).build()) + })); + let endpoint = configure_worker(endpoint); + + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { + Server::builder() + .add_service(endpoint.into_worker_server()) + .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server))) + .await + }); + + this_clone + } + } + + impl Default for InMemoryChannelResolver { + fn default() -> Self { + Self::from_session_builder(DefaultSessionBuilder) + } + } -/// Creates a distributed session context backed by a configurable in-memory worker service. -/// -/// Like [crate::test_utils::localhost::start_localhost_context], this uses tiny file-scan -/// partitions so small test datasets still cross worker boundaries. -pub async fn start_configured_in_memory_context( - num_workers: usize, - session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, - configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, -) -> SessionContext { - let channel_resolver = - InMemoryChannelResolver::from_configured_worker(session_builder, configure_worker); - let state = SessionStateBuilder::new() - .with_default_features() - .with_config(SessionConfig::new().with_target_partitions(num_workers)) - .with_distributed_planner() - .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) - .with_distributed_channel_resolver(channel_resolver) - .with_distributed_file_scan_config_bytes_per_partition(1) - .unwrap() - .build(); - SessionContext::from(state) + #[async_trait] + impl ChannelResolver for InMemoryChannelResolver { + async fn get_worker_client_for_url( + &self, + _: &url::Url, + ) -> Result, DataFusionError> { + Ok(grpc::create_worker_client(self.channel.clone())) + } + } + + /// Creates a distributed session context backed by a single in-memory worker service. + /// The set of produced worker URLs is deterministic, taking the form http://worker-. + pub async fn start_in_memory_context( + num_workers: usize, + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> SessionContext { + let channel_resolver = InMemoryChannelResolver::from_session_builder(session_builder); + let state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_planner() + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) + .with_distributed_channel_resolver(channel_resolver) + .build(); + SessionContext::from(state) + } + + /// Creates a distributed session context backed by a configurable in-memory worker service. + /// + /// Like [crate::test_utils::localhost::start_localhost_context], this uses tiny file-scan + /// partitions so small test datasets still cross worker boundaries. + pub async fn start_configured_in_memory_context( + num_workers: usize, + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, + ) -> SessionContext { + let channel_resolver = + InMemoryChannelResolver::from_configured_worker(session_builder, configure_worker); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(num_workers)) + .with_distributed_planner() + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) + .with_distributed_channel_resolver(channel_resolver) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + SessionContext::from(state) + } } + +#[cfg(feature = "grpc")] +pub use channel::{ + InMemoryChannelResolver, start_configured_in_memory_context, start_in_memory_context, +}; diff --git a/src/test_utils/mod.rs b/src/test_utils/mod.rs index dc66f977..6c84a3ff 100644 --- a/src/test_utils/mod.rs +++ b/src/test_utils/mod.rs @@ -1,5 +1,6 @@ pub mod in_memory_channel_resolver; pub mod insta; +#[cfg(feature = "grpc")] pub mod localhost; pub mod metrics; pub mod mock_exec; diff --git a/src/test_utils/routing.rs b/src/test_utils/routing.rs index 26dc6c10..b36721fd 100644 --- a/src/test_utils/routing.rs +++ b/src/test_utils/routing.rs @@ -2,6 +2,7 @@ use arrow::{ array::{Int64Array, RecordBatch, StringArray}, datatypes::{DataType, Field, Schema, SchemaRef}, }; +use async_trait::async_trait; use datafusion::{ catalog::{Session, TableFunctionImpl, TableProvider}, common::{ScalarValue, Statistics, internal_err, plan_err}, @@ -19,7 +20,6 @@ use datafusion_proto::{physical_plan::PhysicalExtensionCodec, protobuf::proto_er use futures::stream; use prost::Message; use std::{fmt::Formatter, sync::Arc}; -use tonic::async_trait; use url::Url; use crate::execution_plans::DistributedLeafExec; diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index ff9fdb6a..6000fc32 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -6,9 +6,12 @@ mod work_unit_feed_provider; mod work_unit_feed_registry; pub(crate) use remote_work_unit_feed::{ - RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedTxs, build_work_unit_batch_msg, - set_work_unit_received_time, set_work_unit_send_time, + RemoteWorkUnitFeedRegistry, build_work_unit_batch_msg, set_work_unit_received_time, + set_work_unit_send_time, }; +// Re-exported for an out-of-crate transport that consumes it downstream; unused on this branch. +#[allow(unused_imports)] +pub(crate) use remote_work_unit_feed::RemoteWorkUnitFeedTxs; pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed}; pub use remote_work_unit_feed::set_received_time; diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 6f9dfc83..759dee43 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -3,7 +3,9 @@ mod impl_execute_task; mod session_builder; mod single_write_multi_read; mod task_data; -#[cfg(any(test, feature = "integration"))] +// `worker_handles` builds `tonic` servers and Flight channels for the benchmark fixtures, which +// only compile with the gRPC transport. +#[cfg(all(feature = "grpc", any(test, feature = "integration")))] pub(crate) mod test_utils; mod worker_connection_pool; mod worker_service; From b0ca8f3c8227a1d9437c8de5c692273534942edc Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 1 Jul 2026 11:34:42 -0700 Subject: [PATCH 15/20] Added a DispatchPlanSource hook for coordinator-sourced plans. Lets an embedder supply the dispatch bytes per stage instead of the coordinator encoding the plan it holds, whose exec-time state can be wrong to dispatch or not serializable through datafusion-proto. --- src/coordinator/query_coordinator.rs | 18 +++++++++--- src/dispatch_plan_source.rs | 42 ++++++++++++++++++++++++++++ src/distributed_ext.rs | 38 +++++++++++++++++++++++-- src/lib.rs | 2 ++ 4 files changed, 94 insertions(+), 6 deletions(-) create mode 100644 src/dispatch_plan_source.rs diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 1e309410..685bfb80 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -11,7 +11,8 @@ use crate::{ DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_worker_resolver, + WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_dispatch_plan_source, + get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -140,12 +141,21 @@ 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(); + // An embedder can supply the dispatch bytes for this stage (e.g. a structure-only build the + // workers specialize per task) instead of the coordinator encoding the plan it holds, whose + // exec-time state may be wrong to dispatch or not serializable. + let plan_proto = match get_distributed_dispatch_plan_source(session_config) + .and_then(|source| source.dispatch_plan_proto(self.stage_id, task_i)) + { + Some(bytes) => bytes, + None => { + let codec = DistributedCodec::new_combined_with_user(session_config); + PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec() + } + }; let plan_size = plan_proto.len(); let task_key = TaskKey { diff --git a/src/dispatch_plan_source.rs b/src/dispatch_plan_source.rs new file mode 100644 index 00000000..df49de2f --- /dev/null +++ b/src/dispatch_plan_source.rs @@ -0,0 +1,42 @@ +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +/// Supplies the pre-serialized stage subplan the coordinator dispatches, instead of the coordinator +/// encoding the plan it holds. +/// +/// An embedder whose dispatch plan differs from its execution plan registers one via +/// [`crate::DistributedExt::with_distributed_dispatch_plan_source`]. The shm embedder, for example, +/// dispatches a structure-only build (segment-free scans the workers specialize per task) that its +/// exec-time plan is not; sourcing the bytes here lets the coordinator route that plan rather than +/// re-encode its own, which carries exec-time state that is wrong to dispatch. +/// +/// Returning `None` for a `(stage_id, task_number)` lets the coordinator fall back to encoding the +/// plan it holds, so a resolver that only overrides some stages stays correct. +pub trait DispatchPlanSource: Send + Sync { + fn dispatch_plan_proto(&self, stage_id: usize, task_number: usize) -> Option>; +} + +#[derive(Clone)] +pub(crate) struct DispatchPlanSourceExtension(pub(crate) Arc); + +pub(crate) fn set_distributed_dispatch_plan_source( + cfg: &mut SessionConfig, + source: impl DispatchPlanSource + 'static, +) { + set_distributed_dispatch_plan_source_arc(cfg, Arc::new(source)) +} + +pub(crate) fn set_distributed_dispatch_plan_source_arc( + cfg: &mut SessionConfig, + source: Arc, +) { + cfg.set_extension(Arc::new(DispatchPlanSourceExtension(source))); +} + +/// Returns the [`DispatchPlanSource`] registered on this config, if any. +pub fn get_distributed_dispatch_plan_source( + cfg: &SessionConfig, +) -> Option> { + cfg.get_extension::() + .map(|ext| Arc::clone(&ext.0)) +} diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 9efdb420..6f71f7a4 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -2,14 +2,15 @@ use crate::codec::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::config_extension_ext::{ set_distributed_option_extension, set_distributed_option_extension_from_headers, }; +use crate::dispatch_plan_source::set_distributed_dispatch_plan_source; use crate::distributed_planner::set_distributed_task_estimator; use crate::passthrough_headers::set_passthrough_headers; use crate::protocol::set_distributed_channel_resolver; use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::worker_resolver::set_distributed_worker_resolver; use crate::{ - ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, + ChannelResolver, DispatchPlanSource, DistributedConfig, TaskEstimator, WorkUnitFeed, + WorkUnitFeedProvider, WorkerResolver, }; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; @@ -278,6 +279,16 @@ pub trait DistributedExt: Sized { resolver: T, ); + /// Registers a [DispatchPlanSource] the coordinator consults for each stage's dispatch bytes + /// instead of encoding the plan it holds. See [DispatchPlanSource] for when this is needed. + fn with_distributed_dispatch_plan_source( + self, + source: T, + ) -> Self; + + /// Same as [DistributedExt::with_distributed_dispatch_plan_source] but with an in-place mutation. + fn set_distributed_dispatch_plan_source(&mut self, source: T); + /// Adds a distributed task count estimator. [TaskEstimator]s are executed on each node /// sequentially until one returns an estimation on the number of tasks that should be /// used for the stage containing that node. @@ -634,6 +645,10 @@ impl DistributedExt for SessionConfig { set_distributed_channel_resolver(self, resolver); } + fn set_distributed_dispatch_plan_source(&mut self, source: T) { + set_distributed_dispatch_plan_source(self, source); + } + fn set_distributed_task_estimator( &mut self, estimator: T, @@ -787,6 +802,10 @@ impl DistributedExt for SessionConfig { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + #[call(set_distributed_task_estimator)] #[expr($;self)] fn with_distributed_task_estimator(mut self, estimator: T) -> Self; @@ -889,6 +908,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1012,6 +1036,11 @@ impl DistributedExt for SessionState { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1135,6 +1164,11 @@ impl DistributedExt for SessionContext { #[expr($;self)] fn with_distributed_channel_resolver(self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] diff --git a/src/lib.rs b/src/lib.rs index 8389405d..14135a1d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ mod codec; mod common; mod config_extension_ext; mod coordinator; +mod dispatch_plan_source; mod distributed_ext; mod distributed_planner; mod execution_plans; @@ -44,6 +45,7 @@ pub use protocol::grpc; pub use protocol::generated::worker as proto; pub use codec::DistributedCodec; +pub use dispatch_plan_source::{DispatchPlanSource, get_distributed_dispatch_plan_source}; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; /// TODO: do not expose this yet. From 3439f34c6c0045d2f8a6abab932d132f71aa0841 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:13:02 -0700 Subject: [PATCH 16/20] Re-exposed the embedder public API an out-of-crate driver needs. An out-of-crate driver dispatches a query and drives the head stage itself, so it needs prepare_in_process_plan, the metrics store and a public insert, the partition-routing seam on NetworkBoundary, and a no-gRPC decode_task_metrics to file proto metric frames into the plain store. --- src/coordinator/distributed.rs | 70 ++++++++++++++++++++- src/coordinator/metrics_store.rs | 4 +- src/coordinator/mod.rs | 2 +- src/distributed_planner/mod.rs | 2 +- src/distributed_planner/network_boundary.rs | 38 ++++++++++- src/execution_plans/network_coalesce.rs | 10 +++ src/lib.rs | 8 ++- src/protocol/metrics_proto.rs | 21 +++++++ src/protocol/mod.rs | 1 + 9 files changed, 148 insertions(+), 8 deletions(-) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 9c682328..8ed0a9f2 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -24,7 +24,6 @@ use std::sync::{Arc, Mutex}; /// channel resolver and assigned to each task in each stage. /// 2. Encodes all the plans in protobuf format so that network boundary nodes can send them /// over the wire. -#[derive(Debug)] pub struct DistributedExec { /// Initial [ExecutionPlan] present before execution. /// - If the plan was distributed statically, this will be the final distributed plan with all @@ -44,6 +43,19 @@ pub struct DistributedExec { /// Storage where metrics collected from workers at runtime will place their results as they /// finish their respective remote tasks. pub(crate) metrics_store: Option>, + /// Kept alive only on the [DistributedExec::prepare_in_process_plan] path. That path dispatches + /// every stage through the coordinator's background join-set, so the coordinator must outlive + /// the call for the embedder to drive the returned head stage; dropping it aborts the in-flight + /// dispatch. + in_process_coordinator: Mutex>, +} + +impl std::fmt::Debug for DistributedExec { + fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { + f.debug_struct("DistributedExec") + .field("base_plan", &self.base_plan) + .finish_non_exhaustive() + } } pub(super) struct PreparedPlan { @@ -61,9 +73,17 @@ impl DistributedExec { head_stage: Arc::new(Mutex::new(None)), metrics: ExecutionPlanMetricsSet::new(), metrics_store: None, + in_process_coordinator: Mutex::new(None), } } + /// The store where worker task metrics land at runtime, if metrics collection is enabled. + /// Exposed for the in-crate shm/embedder consumer, which files decoded worker metric frames + /// here before the per-task EXPLAIN rewrite. + pub fn metrics_store(&self) -> Option> { + self.metrics_store.clone() + } + /// Enables task metrics collection from remote workers. pub fn with_metrics_collection(mut self, enabled: bool) -> Self { self.metrics_store = match enabled { @@ -134,6 +154,53 @@ impl DistributedExec { .clone() .ok_or_else(|| internal_datafusion_err!("No head stage found. Was execute() called?")) } + + /// Routes and dispatches every stage through the registered channel resolver, then returns the + /// head stage for the caller to drive synchronously, skipping the background task that + /// [`ExecutionPlan::execute`] would otherwise spawn to drive it. + /// + /// This is the extension point for an embedder that owns the runtime and drives the head stage + /// itself (for example a shared-memory mesh). Unlike `execute`, no record-batch pump is spawned: + /// the caller pulls partitions off the returned plan directly. + /// + /// Dispatch on this branch is not synchronous: [`prepare_static_plan`] sends each stage through + /// the [`QueryCoordinator`]'s background join-set. That coordinator is stashed on `self` so the + /// dispatch is not aborted, which means the caller must keep this `DistributedExec` alive for as + /// long as it drives the head stage. + /// + /// Only static task planning is supported here; dynamic task counts need the async coordinator + /// path that `execute` runs. + pub fn prepare_in_process_plan( + &self, + ctx: &Arc, + ) -> Result> { + let d_cfg = DistributedConfig::from_config_options(ctx.session_config().options())?; + if d_cfg.dynamic_task_count { + return exec_err!( + "prepare_in_process_plan only supports static task planning; \ + dynamic task counts require the async coordinator path" + ); + } + + let query_coordinator = + QueryCoordinator::new(Arc::clone(ctx), &self.metrics, self.metrics_store.clone()); + let result = prepare_static_plan(&query_coordinator, &self.base_plan)?; + + self.plan_for_viz + .lock() + .expect("poisoned lock") + .replace(result.plan_for_viz); + self.head_stage + .lock() + .expect("poisoned lock") + .replace(Arc::clone(&result.head_stage)); + self.in_process_coordinator + .lock() + .expect("poisoned lock") + .replace(query_coordinator); + + Ok(result.head_stage) + } } impl DisplayAs for DistributedExec { @@ -165,6 +232,7 @@ impl ExecutionPlan for DistributedExec { head_stage: Arc::new(Mutex::new(None)), metrics: self.metrics.clone(), metrics_store: self.metrics_store.clone(), + in_process_coordinator: Mutex::new(None), })) } diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index ed55db2c..cf1fbc6e 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -17,7 +17,9 @@ impl MetricsStore { Self { tx, rx } } - pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { + // Public for the in-crate shm/embedder consumer, which files decoded worker metric frames into + // the executed plan's store before the per-task EXPLAIN rewrite reads it. + pub fn insert(&self, key: TaskKey, metrics: TaskMetrics) { self.tx.send_modify(|map| { map.insert(key, metrics); }); diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd..be069485 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -6,4 +6,4 @@ mod prepare_static_plan; mod query_coordinator; pub use distributed::DistributedExec; -pub(crate) use metrics_store::MetricsStore; +pub use metrics_store::MetricsStore; diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index 18cfe6e0..176bd319 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -15,7 +15,7 @@ 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, PartitionRoute}; pub use session_state_builder_ext::SessionStateBuilderExt; pub(crate) use statistics::calculate_cost; pub(crate) use task_estimator::set_distributed_task_estimator; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index 3b7bd3b1..aed5fbef 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -5,7 +5,7 @@ use crate::{ Stage, }; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::Result; +use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::repartition::RepartitionExec; @@ -19,6 +19,14 @@ use datafusion_proto::protobuf::proto_error; use prost::Message; use std::sync::Arc; +/// Where a producer's output partition should be sent: which consumer task, and the local partition +/// index within that task's slice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PartitionRoute { + pub consumer_task: usize, + pub consumer_partition: usize, +} + /// This trait represents a node that introduces the necessity of a network boundary in the plan. /// The distributed planner, upon stepping into one of these, will break the plan and build a stage /// out of it. @@ -36,6 +44,34 @@ pub trait NetworkBoundary: ExecutionPlan { /// implementation have. This information is used during planning an executing for ensuring /// the head of a stage has the appropriate shape for consumption. fn producer_head(&self, consumer_tasks: usize) -> ProducerHead; + + /// `P_c`: how many partitions each consumer task reads in the sliced layout + /// (`global = P_c * consumer_task + local`) that shuffle and broadcast reads use. Surfaced so a + /// transport that has to place a produced partition does not re-derive it from node properties. + /// Meaningless for [NetworkCoalesceExec], whose consumers read whole per-producer-task groups. + fn partitions_per_consumer_task(&self) -> usize { + self.properties().partitioning.partition_count() + } + + /// Maps a producer output partition to the consumer task and the local partition within that + /// task that reads it, for the `global = P_c * consumer_task + local` layout. + /// + /// Boundaries whose consumers do not read that layout must override this with an error; the + /// default would silently misroute them. A zero-partition boundary is a planner bug, so it + /// errors instead of routing everything to task `0`. + fn route_partition(&self, output_partition: usize) -> Result { + let p_c = self.partitions_per_consumer_task(); + if p_c == 0 { + return internal_err!( + "cannot route output partition {output_partition}: the boundary reports 0 \ + partitions per consumer task" + ); + } + Ok(PartitionRoute { + consumer_task: output_partition / p_c, + consumer_partition: output_partition % p_c, + }) + } } /// Defines what shape should the head node of a stage have upon getting executed. Depending diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 703843cc..4379910a 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -182,6 +182,16 @@ impl NetworkBoundary for NetworkCoalesceExec { fn producer_head(&self, _consumer_task_count: usize) -> ProducerHead { ProducerHead::None } + + fn route_partition( + &self, + output_partition: usize, + ) -> Result { + internal_err!( + "NetworkCoalesceExec routes by producer task group, not by output partition; \ + partition {output_partition} has no slice-layout route" + ) + } } impl DisplayAs for NetworkCoalesceExec { diff --git a/src/lib.rs b/src/lib.rs index 14135a1d..af569a76 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,10 +18,12 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +// `MetricsStore` is re-exposed for the in-crate shm/embedder consumer, which files worker metric +// frames into the executed plan's store before the per-task EXPLAIN rewrite. +pub use coordinator::{DistributedExec, MetricsStore}; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ - DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, + DistributedConfig, NetworkBoundary, NetworkBoundaryExt, PartitionRoute, SessionStateBuilderExt, TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, }; pub use execution_plans::{ @@ -53,7 +55,7 @@ pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, InProcessChannelResolver, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, + WorkerToCoordinatorMsg, decode_task_metrics, get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, diff --git a/src/protocol/metrics_proto.rs b/src/protocol/metrics_proto.rs index 27b715f2..ea95295e 100644 --- a/src/protocol/metrics_proto.rs +++ b/src/protocol/metrics_proto.rs @@ -58,6 +58,27 @@ pub fn metrics_set_proto_to_df( Ok(metrics_set) } +/// Decode a wire [`pb::TaskMetrics`] into the in-memory [`crate::TaskMetrics`]. The no-gRPC push +/// embedder receives metric frames as proto and files them into the plain metrics store, so it +/// needs the decode direction without the gRPC client. +pub fn decode_task_metrics( + task_metrics: pb::TaskMetrics, +) -> Result { + Ok(crate::TaskMetrics { + pre_order_plan_metrics: task_metrics + .pre_order_plan_metrics + .iter() + .map(metrics_set_proto_to_df) + .collect::>()?, + task_metrics: metrics_set_proto_to_df( + task_metrics + .task_metrics + .as_ref() + .ok_or_else(|| DataFusionError::Internal("Missing field 'task_metrics'".into()))?, + )?, + }) +} + /// Custom metrics are not supported in proto conversion. const CUSTOM_METRICS_NOT_SUPPORTED: &str = "custom metrics are not supported in metrics proto conversion"; diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index a08a3eb4..d21fd4b2 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -13,6 +13,7 @@ mod worker_channel; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; pub use in_process::InProcessChannelResolver; +pub use metrics_proto::decode_task_metrics; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, From 3ca1aeb526b89ae37b347521d2cad9b2a35f45eb Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 29 Jun 2026 18:14:43 -0700 Subject: [PATCH 17/20] Ported the shared-memory transport onto the channel protocol. The shm ring/DSM/mesh/framing core moves verbatim from the fork; the old WorkerTransport umbrella is gone, so the resolver hands out a channel and execute_task reads the rings per partition. The mesh is the no-gRPC transport, so it builds in both the grpc-on and grpc-off configs. --- Cargo.lock | 24 + Cargo.toml | 2 + src/lib.rs | 7 + src/shm/dsm.rs | 544 ++++++++ src/shm/in_process.rs | 851 ++++++++++++ src/shm/mesh.rs | 375 +++++ src/shm/mod.rs | 81 ++ src/shm/mpsc_ring.rs | 1550 +++++++++++++++++++++ src/shm/runtime.rs | 477 +++++++ src/shm/self_hosted.rs | 1277 +++++++++++++++++ src/shm/setup.rs | 371 +++++ src/shm/sink.rs | 49 + src/shm/transport.rs | 2935 ++++++++++++++++++++++++++++++++++++++++ 13 files changed, 8543 insertions(+) create mode 100644 src/shm/dsm.rs create mode 100644 src/shm/in_process.rs create mode 100644 src/shm/mesh.rs create mode 100644 src/shm/mod.rs create mode 100644 src/shm/mpsc_ring.rs create mode 100644 src/shm/runtime.rs create mode 100644 src/shm/self_hosted.rs create mode 100644 src/shm/setup.rs create mode 100644 src/shm/sink.rs create mode 100644 src/shm/transport.rs diff --git a/Cargo.lock b/Cargo.lock index 3858be13..42b58fc2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -475,6 +475,28 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -2192,6 +2214,7 @@ dependencies = [ "arrow-flight", "arrow-ipc", "arrow-select", + "async-stream", "async-trait", "bincode", "bytes", @@ -2206,6 +2229,7 @@ dependencies = [ "hyper-util", "insta", "itertools 0.14.0", + "log", "moka", "num-traits", "object_store", diff --git a/Cargo.toml b/Cargo.toml index 6b2f0c24..d0e73913 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,8 @@ datafusion = { workspace = true, features = [ ] } datafusion-proto = { workspace = true } async-trait = "0.1.89" +async-stream = "0.3" +log = "0.4" tokio = { version = "1.48", features = ["full"] } http = "1.3.1" itertools = "0.14.0" diff --git a/src/lib.rs b/src/lib.rs index af569a76..52984bd5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,9 @@ mod execution_plans; mod metrics; mod passthrough_headers; mod protocol; +// Not feature-gated: the shared-memory mesh is the no-gRPC transport, so it has to build in both +// the `grpc`-on and `grpc`-off configs. +pub mod shm; mod stage; mod work_unit_feed; mod worker; @@ -48,6 +51,10 @@ pub use protocol::generated::worker as proto; pub use codec::DistributedCodec; pub use dispatch_plan_source::{DispatchPlanSource, get_distributed_dispatch_plan_source}; +// The producer-side sink traits live in `shm` because only a push-based transport produces through +// them; re-exported at the crate root so `crate::PartitionSink` resolves the way the shm core spells +// it. +pub use shm::{PartitionSink, WorkerSink}; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; /// TODO: do not expose this yet. diff --git a/src/shm/dsm.rs b/src/shm/dsm.rs new file mode 100644 index 00000000..3720da0e --- /dev/null +++ b/src/shm/dsm.rs @@ -0,0 +1,544 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Mesh-multiplexed DSM layout: one MPSC inbox per receiver process. +//! +//! Each MPP query allocates a single DSM region: +//! +//! ```text +//! +--- MppDsmHeader (repr C, MAXALIGN-padded) -------------+ +//! | magic, version, n_procs, queue_bytes | +//! | plan_offset, plan_len, queues_offset, region_total | +//! +-------------------------------------------------------+ +//! | plan bytes (bincode-serialized worker fragment) | +//! +-------------------------------------------------------+ +//! | padding to MAXALIGN | +//! +-------------------------------------------------------+ +//! | DsmMpscRing inbox array: n_procs inboxes | +//! | inbox(receiver) = queues_offset | +//! | + receiver * queue_bytes | +//! +-------------------------------------------------------+ +//! ``` +//! +//! - `n_procs` = 1 leader + N parallel workers. Leader is `proc_idx = 0`; workers are +//! `proc_idx = ParallelWorkerNumber + 1`. +//! - Each process attaches as receiver to its own inbox (one MPSC ring) and as sender +//! to each of N-1 peer inboxes. Senders stamp `MppFrameHeader::sender_proc` on every +//! frame so the receiver demuxes by source. +//! - Self-loop frames (proc → itself) ride an in-proc channel installed by +//! `worker_setup`, not a DSM slot: each `DsmMpscReceiver` is owned by a single +//! receiver process, so there is no DSM slot for the self-loop pair. + +use std::ffi::c_void; +use std::mem::size_of; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; + +use super::AliveFlag; +use super::mesh::{align_up_maxalign_checked, aligned_queue_bytes}; +use super::mpsc_ring::{self, DsmMpscReceiver, DsmMpscRingHeader, DsmMpscSender, Wakeup}; + +/// Number of slots in each per-receiver MPSC ring. With operator-visible +/// The embedder's queue-size knob divided across `RING_SLOTS` slots, each slot holds +/// up to `(queue_bytes / RING_SLOTS) - SLOT_HEADER` bytes of payload. That's enough +/// for typical Arrow batches at the bench scales we measured. Fixed at compile time +/// for now; a future GUC could expose it if bench data points at a different sweet +/// spot. +const RING_SLOTS: u32 = 8; + +/// Floor `ring_dims_for` applies to the per-slot payload capacity. The layout minimum check +/// uses the same floor so the constructed ring always fits its inbox region. +const MIN_SLOT_CAPACITY: u32 = 64; + +/// Cache-line alignment for the per-inbox DsmMpscRing header. The ring itself is plain +/// `#[repr(C)]` (the base address is only MAXALIGN-guaranteed), but spacing the per-inbox +/// offsets to 64 keeps hot headers on separate cache lines when the base cooperates; +/// both `queues_offset` and per-inbox `queue_bytes` are aligned up to this so the +/// computed `inbox_offset(r) = queues_offset + r * queue_bytes` lands at a 64-aligned +/// address for every `r`. +const RING_ALIGN: usize = 64; + +#[inline] +fn align_up_ring(n: usize) -> Option { + let mask = RING_ALIGN - 1; + n.checked_add(mask).map(|x| x & !mask) +} + +const MPP_DSM_MAGIC: u32 = 0x4D50_5052; // "MPPR" (RPC variant) +/// Bump on any wire-incompatible change to the DSM header layout or to the inbox-offset +/// math, so attaching workers reject mismatched leaders loudly rather than reading +/// garbage. Validated in [`MppDsmHeader::validate`]. +const MPP_DSM_HEADER_VERSION: u32 = 4; + +/// Absolute cap on DSM region size. 16 GiB is two orders of magnitude beyond +/// any realistic workload; the cap fails early on a pathologically oversized +/// request rather than asking PG for ~`usize::MAX` bytes. +const MPP_DSM_MAX_BYTES: usize = 16 * 1024 * 1024 * 1024; + +/// C-repr header at offset 0 of the DSM region. +/// +/// Layout: three `u32`s + padding, then five `u64`s. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct MppDsmHeader { + pub(super) magic: u32, + pub(super) header_version: u32, + /// Total proc count. Leader is `proc_idx = 0`; workers are + /// `proc_idx = ParallelWorkerNumber + 1`. The DSM region holds `n_procs` + /// per-receiver MPSC inboxes laid out contiguously after the plan bytes. + pub n_procs: u32, + pub(super) _pad: u32, + pub(super) queue_bytes: u64, + pub(super) plan_offset: u64, + pub(super) plan_len: u64, + pub(super) queues_offset: u64, + pub region_total: u64, +} + +impl MppDsmHeader { + fn from_layout(layout: &DsmLayout) -> Self { + Self { + magic: MPP_DSM_MAGIC, + header_version: MPP_DSM_HEADER_VERSION, + n_procs: layout.n_procs, + _pad: 0, + queue_bytes: layout.queue_bytes as u64, + plan_offset: layout.plan_offset as u64, + plan_len: layout.plan_len as u64, + queues_offset: layout.queues_offset as u64, + region_total: layout.region_total as u64, + } + } + + pub(super) fn validate(&self, region_total: u64) -> Result<(), &'static str> { + if self.magic != MPP_DSM_MAGIC { + return Err("mpp: DSM header magic mismatch"); + } + if self.header_version != MPP_DSM_HEADER_VERSION { + return Err("mpp: DSM header version mismatch"); + } + if self.n_procs == 0 { + return Err("mpp: header n_procs must be > 0"); + } + if self.region_total != region_total { + return Err("mpp: DSM region_total in header disagrees with attached size"); + } + match self.plan_offset.checked_add(self.plan_len) { + None => return Err("mpp: plan_offset + plan_len overflow"), + Some(end) if end > self.queues_offset => { + return Err("mpp: plan would overlap queues area"); + } + _ => {} + } + if self.queues_offset > region_total { + return Err("mpp: queues_offset past end of region"); + } + Ok(()) + } + + /// Byte offset (relative to DSM base) of `receiver_proc`'s MPSC inbox. + /// + /// The owner attaches as receiver (`DsmMpscReceiver`); every other process attaches + /// as sender (`DsmMpscSender`) to the same inbox and stamps its identity into the + /// frame header (`MppFrameHeader::sender_proc`). + pub(super) fn inbox_offset(&self, receiver_proc: u32) -> u64 { + debug_assert!(receiver_proc < self.n_procs); + self.queues_offset + (receiver_proc as u64) * self.queue_bytes + } +} + +/// Pure-math layout for [`compute_dsm_layout`]. +#[derive(Debug, Clone, Copy)] +pub(super) struct DsmLayout { + pub n_procs: u32, + pub queue_bytes: usize, + pub plan_offset: usize, + pub plan_len: usize, + pub queues_offset: usize, + pub region_total: usize, +} + +/// Compute the DSM region size and field offsets for one MPP query. +/// +/// `n_procs` is the total proc count (1 leader + N workers); the region holds one +/// `DsmMpscRing` inbox per process. Every other process attaches as sender to that +/// inbox and the receiver demultiplexes by `MppFrameHeader::sender_proc`. +/// +/// `queue_bytes` is the per-inbox total (ring header + `RING_SLOTS` slots). +/// The embedder's operator-facing queue-size knob controls this value. +pub(super) fn compute_dsm_layout( + n_procs: u32, + queue_bytes: usize, + plan_len: usize, +) -> Result { + if n_procs < 2 { + return Err("mpp: n_procs must be >= 2 (leader + at least one worker)"); + } + // MAXALIGN-round-down first (operator-friendly), then round up to the ring's + // 64-byte alignment requirement. Doing it in this order means each per-inbox + // region is both MAXALIGN-aligned (PG DSM convention) AND cache-line aligned, + // which keeps the rings' hot header fields on separate cache lines. + let queue_bytes = aligned_queue_bytes(queue_bytes); + if queue_bytes == 0 { + return Err("mpp: queue_bytes too small after alignment"); + } + let queue_bytes = align_up_ring(queue_bytes).ok_or("mpp: queue_bytes alignment overflow")?; + // Each inbox must have room for the ring `ring_dims_for` will actually build: + // `RING_SLOTS` slots with the 64-byte capacity floor. Checking against a smaller + // ring here would let `create_at` write past the inbox region for tiny + // `queue_bytes`, overlapping the next inbox. + if queue_bytes < DsmMpscRingHeader::region_bytes(RING_SLOTS, MIN_SLOT_CAPACITY) { + return Err("mpp: queue_bytes too small for ring header + min slot capacity"); + } + let header_end = align_up_maxalign_checked(size_of::()) + .ok_or("mpp: header alignment overflow")?; + let plan_offset = header_end; + let plan_end = plan_offset + .checked_add(plan_len) + .ok_or("mpp: plan offset+len overflow")?; + // Round queues_offset up to RING_ALIGN so the first inbox starts at a 64-aligned + // address; subsequent inboxes are queue_bytes apart and queue_bytes is RING_ALIGN- + // aligned, so they all land on cache-line boundaries. + let queues_offset = align_up_ring(plan_end).ok_or("mpp: queues alignment overflow")?; + let queues_bytes = (n_procs as usize) + .checked_mul(queue_bytes) + .ok_or("mpp: queues bytes overflow")?; + let region_total = queues_offset + .checked_add(queues_bytes) + .ok_or("mpp: region total overflow")?; + if region_total > MPP_DSM_MAX_BYTES { + return Err("mpp: DSM region exceeds MPP_DSM_MAX_BYTES"); + } + Ok(DsmLayout { + n_procs, + queue_bytes, + plan_offset, + plan_len, + queues_offset, + region_total, + }) +} + +/// Derive `(ring_slots, slot_capacity)` for a per-inbox region of `queue_bytes`. +/// Total ring region size = `DsmMpscRingHeader::region_bytes(ring_slots, slot_capacity)` +/// which fits within `queue_bytes` by construction. +fn ring_dims_for(queue_bytes: usize) -> (u32, u32) { + let header = std::mem::size_of::(); + // queue_bytes >= ring_dims minimum is enforced in compute_dsm_layout; we recompute + // here without re-validating. + let slot_total_bytes = queue_bytes.saturating_sub(header); + let slot_capacity = (slot_total_bytes / RING_SLOTS as usize).max(MIN_SLOT_CAPACITY as usize); + // Cap slot_capacity at u32::MAX (DsmMpscRing's field type) to avoid casting wrap. + let slot_capacity = slot_capacity.min(u32::MAX as usize) as u32; + (RING_SLOTS, slot_capacity) +} + +/// Per-proc return from `attach_proc`: N-1 outbound senders (one per peer inbox) plus a +/// single inbound receiver (this proc's own inbox). +/// +/// The own-inbox is the multiplexed entry point: every peer attaches to it as a sender +/// (each `DsmMpscSender` increments the ring's `sender_count`) and stamps its identity +/// into `MppFrameHeader::sender_proc` on every frame. The receiver side pulls frames +/// off that single ring and routes them to per-`(sender_proc, stage_id, partition)` +/// channel buffers via [`DrainHandle`]. +pub(super) struct ProcAttach { + /// `outbound_senders[i]` writes to peer `peer_proc_for_index(this_proc, i)`'s inbox. + /// `peer_proc(i) = i if i < this_proc else i + 1` skips the self-loop entry. + pub(super) outbound_senders: Vec, + /// This process's own MPSC inbox receiver. Drained inline by `DrainHandle`. + pub(super) inbound_receiver: DsmMpscReceiver, + pub(super) alive: AliveFlag, +} + +/// Read `region_total` out of the header at the start of an initialized region. +/// +/// # Safety +/// `base` must point at the start of a region a leader wrote via [`leader_init`]. +pub(super) unsafe fn read_region_total(base: *const c_void) -> u64 { + unsafe { std::ptr::read(base as *const MppDsmHeader).region_total } +} + +/// Translate a peer index (`0..n_procs - 1`) into a process index +/// (`0..n_procs`) by skipping the self-loop slot. +#[inline] +pub(super) fn peer_proc_for_index(this_proc: u32, peer_idx: u32) -> u32 { + if peer_idx < this_proc { + peer_idx + } else { + peer_idx + 1 + } +} + +/// Initialize the DSM region as the leader (`proc_idx = 0`). Writes the MppDsmHeader, +/// copies the plan bytes, initializes the N MPSC inboxes via `DsmMpscRing::create_at`, +/// and attaches the leader as receiver to its own inbox, plus (when `attach_senders`) as +/// sender to each peer for the control plane (work-unit frames). +/// +/// `attach_senders` is a commitment: a ring latches `detached` when its sender count falls to +/// zero, so a leader that attaches and then drops its senders before a worker attached poisons +/// that worker's inbox. Attach only when the senders outlive the query. +/// +/// # Safety +/// - `coordinate` must point to the start of a DSM region of size `>= layout.region_total`. +/// - The region must be uninitialized (the leader is the first writer). +pub(super) unsafe fn leader_init( + coordinate: *mut c_void, + layout: &DsmLayout, + plan_bytes: &[u8], + wakeup: Arc, + attach_senders: bool, +) -> Result { + if coordinate.is_null() { + return Err("mpp: leader_init given null coordinate".into()); + } + if plan_bytes.len() != layout.plan_len { + return Err(format!( + "mpp: plan_bytes.len()={} != layout.plan_len={}", + plan_bytes.len(), + layout.plan_len + )); + } + + let base = coordinate as *mut u8; + + // Header. + unsafe { + std::ptr::write( + base.cast::(), + MppDsmHeader::from_layout(layout), + ); + } + // Plan bytes. + unsafe { + std::ptr::copy_nonoverlapping( + plan_bytes.as_ptr(), + base.add(layout.plan_offset), + plan_bytes.len(), + ); + } + + // Initialize the N MPSC inboxes. Workers can't do this (the region is uninitialized + // at their attach time), so the leader runs DsmMpscRing::create_at for every receiver. + let header = MppDsmHeader::from_layout(layout); + let n_procs = layout.n_procs; + let (ring_slots, slot_capacity) = ring_dims_for(layout.queue_bytes); + for r in 0..n_procs { + let off = header.inbox_offset(r) as usize; + let inbox_addr = unsafe { base.add(off) }; + unsafe { mpsc_ring::create_at(inbox_addr, ring_slots, slot_capacity) }; + } + + let attach = unsafe { attach_proc(base, &header, 0, attach_senders, wakeup) }; + Ok(attach) +} + +/// Attach to the leader-initialized DSM region as `proc_idx` (`0 = leader`, `1..N` = +/// parallel workers). Attach as receiver to this proc's own MPSC inbox, and (if +/// `attach_senders` is true) as sender to every peer's inbox. +/// +/// `attach_senders = false` is for the leader (consumer-only). If the leader attached +/// as sender it would bump every peer inbox's `sender_count` and decrement it on `Drop`; +/// if it dropped before any worker bumped, the 1 → 0 transition would flip `detached` +/// on every peer inbox and every later worker send would fail `SendError::Detached`. +/// Skipping keeps `sender_count` honest (only producers ever increment). +/// +/// # Safety +/// - `base` must point to a DSM region whose header has been validated. +/// - `header.inbox_offset(r)` must point at a ring already initialized by +/// `DsmMpscRing::create_at` (the leader does this in `leader_init`). +unsafe fn attach_proc( + base: *mut u8, + header: &MppDsmHeader, + this_proc: u32, + attach_senders: bool, + wakeup: Arc, +) -> ProcAttach { + let n_procs = header.n_procs; + let peer_count = (n_procs - 1) as usize; + let mut outbound_senders = Vec::with_capacity(if attach_senders { peer_count } else { 0 }); + + let (ring_slots, slot_capacity) = ring_dims_for(header.queue_bytes as usize); + + // One liveness flag shared by every handle this attach mints. + let alive = Arc::new(AtomicBool::new(true)); + + if attach_senders { + for peer_idx in 0..(n_procs - 1) { + let r = peer_proc_for_index(this_proc, peer_idx); + let off = header.inbox_offset(r) as usize; + let inbox_addr = unsafe { base.add(off) }; + let nn = unsafe { mpsc_ring::attach_at(inbox_addr, ring_slots, slot_capacity) } + .expect("DsmMpscRing attach_at: leader-initialized region must validate"); + outbound_senders + .push(unsafe { DsmMpscSender::new(nn, Arc::clone(&wakeup), Arc::clone(&alive)) }); + } + } + + // Inbound: this proc's own inbox. Single receiver per ring (MPSC contract). + let own_off = header.inbox_offset(this_proc) as usize; + let own_inbox_addr = unsafe { base.add(own_off) }; + let own_nn = unsafe { mpsc_ring::attach_at(own_inbox_addr, ring_slots, slot_capacity) } + .expect("DsmMpscRing attach_at: own inbox must validate"); + let inbound_receiver = unsafe { DsmMpscReceiver::new(own_nn, Arc::clone(&alive)) }; + + ProcAttach { + outbound_senders, + inbound_receiver, + alive, + } +} + +/// Attach to the leader-initialized DSM region as `proc_idx` (1-based for +/// workers: PG's `ParallelWorkerNumber + 1`). +/// +/// # Safety +/// - `coordinate` must be the DSM region pointer the leader initialized. +/// - `region_total` must match the DSM's attached size. +pub(super) unsafe fn worker_attach( + coordinate: *mut c_void, + region_total: u64, + proc_idx: u32, + wakeup: Arc, +) -> Result<(MppDsmHeader, Vec, ProcAttach), String> { + if coordinate.is_null() { + return Err("mpp: worker_attach given null coordinate".into()); + } + let base = coordinate as *mut u8; + let header = unsafe { std::ptr::read(base.cast::()) }; + header + .validate(region_total) + .map_err(|e| format!("mpp: worker DSM validate: {e}"))?; + if proc_idx == 0 { + return Err( + "mpp: worker_attach must be called with proc_idx >= 1 (proc 0 is leader)".into(), + ); + } + if proc_idx >= header.n_procs { + return Err(format!( + "mpp: proc_idx {proc_idx} not in 1..{}", + header.n_procs + )); + } + + // Copy plan bytes out of DSM so the caller has an owned buffer. + let plan_bytes = unsafe { + std::slice::from_raw_parts( + base.add(header.plan_offset as usize), + header.plan_len as usize, + ) + .to_vec() + }; + + let attach = unsafe { + attach_proc( + base, &header, proc_idx, /* attach_senders */ true, wakeup, + ) + }; + Ok((header, plan_bytes, attach)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compute_dsm_layout_works() { + let l = compute_dsm_layout(4, 64 * 1024, 1024).unwrap(); + assert_eq!(l.n_procs, 4); + // 4 procs => 4 inboxes laid out contiguously after the plan bytes. + let aligned = aligned_queue_bytes(64 * 1024); + let queues_size = 4 * aligned; + assert_eq!(l.region_total, l.queues_offset + queues_size); + } + + #[test] + fn compute_dsm_layout_rejects_zero_procs() { + assert!(compute_dsm_layout(0, 64 * 1024, 0).is_err()); + } + + #[test] + fn compute_dsm_layout_rejects_oversize() { + assert!(compute_dsm_layout(u32::MAX, 64 * 1024, 0).is_err()); + } + + #[test] + fn compute_dsm_layout_scales_linearly_in_n_procs() { + // Total queue area must grow as O(N) in proc count. Pinning the math here so a + // regression that grows queues per-pair fails this test at compile time rather + // than at the N=24 wall-time cliff in production. + let queue_bytes = 64 * 1024; + let aligned = aligned_queue_bytes(queue_bytes); + for n in [2u32, 4, 8, 16, 24] { + let l = compute_dsm_layout(n, queue_bytes, 0).unwrap(); + let expected = (n as usize) * aligned; + assert_eq!( + l.region_total - l.queues_offset, + expected, + "n={n}: expected {expected} inbox bytes ({n} inboxes × {aligned})" + ); + } + } + + #[test] + fn header_inbox_offset_is_per_receiver() { + // 4 procs => 4 inboxes, contiguous, sized by queue_bytes each. + let l = compute_dsm_layout(4, 64 * 1024, 0).unwrap(); + let h = MppDsmHeader::from_layout(&l); + let aligned = h.queue_bytes; + assert_eq!(h.inbox_offset(0), h.queues_offset); + assert_eq!(h.inbox_offset(1), h.queues_offset + aligned); + assert_eq!(h.inbox_offset(2), h.queues_offset + 2 * aligned); + assert_eq!(h.inbox_offset(3), h.queues_offset + 3 * aligned); + } + + #[test] + fn header_validate_accepts_self() { + let l = compute_dsm_layout(2, 64 * 1024, 0).unwrap(); + let h = MppDsmHeader::from_layout(&l); + assert!(h.validate(l.region_total as u64).is_ok()); + } + + #[test] + fn header_validate_rejects_wrong_version() { + let l = compute_dsm_layout(2, 64 * 1024, 0).unwrap(); + let mut h = MppDsmHeader::from_layout(&l); + h.header_version = MPP_DSM_HEADER_VERSION.wrapping_sub(1); + let err = h + .validate(l.region_total as u64) + .expect_err("wrong version must fail"); + assert!(err.contains("DSM header version mismatch"), "got: {err}"); + } + + #[test] + fn header_validate_rejects_size_mismatch() { + let l = compute_dsm_layout(2, 64 * 1024, 0).unwrap(); + let h = MppDsmHeader::from_layout(&l); + assert!(h.validate(l.region_total as u64 + 1).is_err()); + } + + /// Pins the coupling between `compute_dsm_layout`'s minimum check and `ring_dims_for`'s + /// slot-capacity floor: a queue sized below the floored ring must be rejected, or + /// `create_at` would write past its inbox region. + #[test] + fn layout_rejects_queue_smaller_than_floored_ring() { + let floored = DsmMpscRingHeader::region_bytes(RING_SLOTS, MIN_SLOT_CAPACITY); + let too_small = DsmMpscRingHeader::region_bytes(RING_SLOTS, 1); + assert!(too_small < floored); + assert!(compute_dsm_layout(2, too_small, 0).is_err()); + assert!(compute_dsm_layout(2, floored.next_multiple_of(64), 0).is_ok()); + } +} diff --git a/src/shm/in_process.rs b/src/shm/in_process.rs new file mode 100644 index 00000000..effbb160 --- /dev/null +++ b/src/shm/in_process.rs @@ -0,0 +1,851 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! In-process instantiation of the shared-memory transport, plus an end-to-end test. +//! +//! A real distributed query runs through [`ShmChannelResolver`] with no Postgres and no Flight. +//! This is the rebase-safety payoff: a single process plays every role that production splits +//! across the leader (build and slice the plan), Postgres (allocate the DSM region, launch the +//! workers), and each worker (run a producer fragment and push it through the mesh). If an upstream +//! change breaks the channel-protocol contract, the boundary routing API, or +//! `prepare_in_process_plan`, the test below fails here, before any downstream embedder rebuilds. +//! +//! What's faithful vs. simplified relative to the Postgres path: +//! - Faithful: the DSM ring mesh, the framing, the cooperative drain, `ShmWorkerChannel::execute_task`, +//! the per-fragment routing (`collect_dispatched_stages`), `run_worker_fragment`, and the leader +//! consuming via `prepare_in_process_plan`. +//! - Simplified: the producer subplans are shared as `Arc`s rather than bincode-serialized through +//! the DSM plan region (all roles live in one address space), the wakeup is a no-op (the +//! cooperative consumer yields rather than parking), and there's no cancellation source. + +use std::alloc::Layout; +use std::ffi::c_void; +use std::sync::Arc; + +use datafusion::arrow::array::{Int32Array, RecordBatch}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::catalog::memory::DataSourceExec; +use datafusion::common::runtime::JoinSet; +use datafusion::common::{DataFusionError, HashMap, Result}; +use datafusion::config::ConfigOptions; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::{SessionStateBuilder, TaskContext}; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion::prelude::{SessionConfig, SessionContext}; + +use crate::{ + DistributedConfig, DistributedExec, DistributedExt, DistributedLeafExec, + DistributedTaskContext, NetworkBoundaryExt, NetworkBroadcastExec, NetworkCoalesceExec, + NetworkShuffleExec, PartitionSink, SessionStateBuilderExt, TaskEstimation, TaskEstimator, + WorkerSink, +}; + +use super::mpsc_ring::Wakeup; +use super::runtime::{InProcessWorkerResolver, MppMesh, ShmChannelResolver, proc_for_task}; +use super::setup::{dsm_region_bytes, leader_setup, run_worker_fragment, worker_setup}; +use super::transport::{ + CooperativeDrainSet, MppFrameHeader, MppPartitionSink, MppSender, NoInterrupt, +}; + +/// Per-inbox DSM ring size for the in-process mesh. Generous: the test ships a handful of tiny +/// batches, so backpressure never kicks in. Production sizes this from `paradedb.mpp_queue_size`. +const IN_PROCESS_QUEUE_BYTES: usize = 1 << 20; + +/// No-op wakeup. The cooperative consumer in `runtime::pull_partition_stream` yields rather +/// than parking, so a publish never needs to wake a blocked thread. The real cross-thread wakeup +/// extension point is covered by `mpsc_ring`'s `injected_wakeup_unparks_blocked_consumer`. +struct NoopWakeup; +impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} +} + +/// Owns the single heap buffer that stands in for the PG DSM segment. Every proc (leader + workers) +/// reads and writes the same region through raw pointers; the lock-free rings make concurrent +/// access sound. Kept alive in the harness until after all producer tasks join. +struct HeapRegion { + ptr: *mut u8, + layout: Layout, +} + +impl HeapRegion { + fn new(bytes: usize) -> Self { + // 64-byte alignment so each per-inbox ring header lands on its own cache line; the + // dsm layout aligns the offsets within the region, but only if the base is aligned too. + let layout = Layout::from_size_align(bytes, 64).expect("dsm region layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null(), "dsm region alloc failed"); + Self { ptr, layout } + } + + fn base(&self) -> *mut c_void { + self.ptr as *mut c_void + } +} + +impl Drop for HeapRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } +} + +/// Send wrapper so the leader-init / worker-attach raw base pointer can be handed to the per-proc +/// setup. All setup runs on the harness thread, so the pointer never crosses into a spawned task; +/// only the resulting `Send` meshes and senders do. +#[derive(Clone, Copy)] +struct SharedBase(*mut c_void); +unsafe impl Send for SharedBase {} + +/// Opaque, non-sentinel receiver token. `NoopWakeup` ignores the value; this just exercises the +/// `set_receiver` path with something the producer won't skip as "no consumer registered". +fn receiver_token(proc_idx: u32) -> u64 { + proc_idx as u64 + 1 +} + +// --------------------------------------------------------------------------------------------- +// Leader-side producer-stage discovery (in-process port of pg_search's `collect_dispatched_stages`). +// The crate's gRPC path keys dispatch on resolver URLs and never decides this; the shm_mq peers are +// push-driven without URLs, so the embedder classifies each boundary's routing here. +// --------------------------------------------------------------------------------------------- + +/// Routing rule for a producer fragment's output partitions. +#[derive(Clone, Debug)] +enum FragmentRouting { + /// Every output partition goes to one destination proc (a `NetworkCoalesceExec`, or the + /// top-level gather to the leader). + Coalesce { dest_proc: u32 }, + /// Hash-partitioned mesh (`NetworkShuffleExec` / `NetworkBroadcastExec`): output partition `q` + /// goes to the consumer task the crate's `route_partition(q)` selects. + Hashed { + consumer_task: Vec, + broadcast: bool, + }, +} + +/// One producer stage to run, captured from a network boundary before `prepare_in_process_plan` +/// converts its input stage to `Remote`. +struct StageEntry { + stage_num: u32, + task_count: usize, + routing: FragmentRouting, + plan: Arc, +} + +/// Walk the distributed physical plan and collect every producer stage, once per boundary. +fn collect_dispatched_stages(root: &Arc, n_workers: u32) -> Vec { + let mut out = Vec::new(); + collect_stages(root, n_workers, /* nested = */ false, &mut out); + out +} + +fn collect_stages( + plan: &Arc, + n_workers: u32, + nested: bool, + out: &mut Vec, +) { + if let Some(nb) = plan.as_ref().as_network_boundary() { + let stage = nb.input_stage(); + let stage_id = stage.num() as u32; + let route_consumer_tasks = || { + // The producer head fans its output into `per_task * consumer_task_count` partitions, + // contiguous per consumer task, so producer output partition `q` is read by consumer + // task `q / per_task`. That mirrors the `off = per_task * task_index` slicing both the + // shuffle and broadcast boundaries apply on the consume side; `per_task` is the consumer + // boundary's own per-task output partition count. + let per_task = plan.output_partitioning().partition_count().max(1); + let n_out = stage + .local_plan() + .map_or(0, |p| p.properties().partitioning.partition_count()); + (0..n_out) + .map(|q| (q / per_task) as u32) + .collect::>() + }; + let plan_any = plan.as_ref(); + let routing = if plan_any.is::() { + if nested { + FragmentRouting::Coalesce { + dest_proc: proc_for_task(n_workers, 0), + } + } else { + FragmentRouting::Coalesce { dest_proc: 0 } + } + } else if plan_any.is::() { + assert!(nested, "top-level NetworkShuffleExec is unsupported"); + FragmentRouting::Hashed { + consumer_task: route_consumer_tasks(), + broadcast: false, + } + } else if plan_any.is::() { + assert!(nested, "top-level NetworkBroadcastExec is unsupported"); + FragmentRouting::Hashed { + consumer_task: route_consumer_tasks(), + broadcast: true, + } + } else { + panic!("unrecognized network boundary {}", plan.name()); + }; + + let task_count = stage.task_count(); + if let Some(stage_plan) = stage.local_plan() { + out.push(StageEntry { + stage_num: stage_id, + task_count, + routing, + plan: Arc::clone(stage_plan), + }); + // The boundary's children() returns [stage.plan], so descending here would double-count + // every nested stage. Recurse through the stage plan directly with nested = true. + collect_stages(stage_plan, n_workers, true, out); + } + return; + } + for child in plan.children() { + collect_stages(child, n_workers, nested, out); + } +} + +/// One producer fragment assigned to a worker proc: a single task of a producer stage. +struct FragmentAssignment { + stage_id: u32, + task_idx: usize, + task_count: usize, + plan: Arc, + routing: FragmentRouting, +} + +/// Rebuild a plan subtree so each proc executes its own node instances. In production every +/// proc decodes its own copy of the plan; sharing one `Arc` across procs would share +/// execute-once state (a `RepartitionExec` panics when a second proc executes a partition the +/// first already consumed). Leaves are shared: they carry no execute-once state here. +fn reinstantiate(plan: &Arc) -> Arc { + let children: Vec<_> = plan.children().into_iter().map(reinstantiate).collect(); + if children.is_empty() { + Arc::clone(plan) + } else { + Arc::clone(plan) + .with_new_children(children) + .expect("with_new_children with the same arity") + } +} + +/// Expand the dispatched stages into the fragments `this_proc` owns under `proc_for_task`. +fn fragments_for_proc( + entries: &[StageEntry], + this_proc: u32, + n_workers: u32, +) -> Vec { + let mut out = Vec::new(); + for entry in entries { + for task_idx in 0..entry.task_count { + if proc_for_task(n_workers, task_idx as u32) != this_proc { + continue; + } + // Broadcast caps its build subtree at task 0; the other tasks would re-emit the same + // canonical replica and the consumer's select_all would over-count. + if matches!( + entry.routing, + FragmentRouting::Hashed { + broadcast: true, + .. + } + ) && task_idx != 0 + { + continue; + } + out.push(FragmentAssignment { + stage_id: entry.stage_num, + task_idx, + task_count: entry.task_count, + plan: reinstantiate(&entry.plan), + routing: entry.routing.clone(), + }); + } + } + out +} + +/// Build a fragment's `TaskContext`, carrying the right `DistributedTaskContext` so nested boundary +/// nodes know their `(task_index, task_count)` and the worker session's resolver/transport ride +/// along for `prepare_in_process_plan`. +fn fragment_task_ctx( + session: &SessionContext, + task_index: usize, + task_count: usize, +) -> Arc { + let cfg = session + .state() + .config() + .clone() + .with_extension(Arc::new(DistributedTaskContext { + task_index, + task_count, + })); + Arc::new(TaskContext::default().with_session_config(cfg)) +} + +/// Run all fragments owned by one worker proc, then signal completion. Mirrors the body of +/// pg_search's `run_mpp_worker`: build a [`WorkerSink`] that routes by partition, open a +/// [`PartitionSink`] per output partition, wrap each fragment in a fresh `DistributedExec` + +/// `prepare_in_process_plan` (converting nested boundaries), and join. +async fn run_worker_proc( + fragments: Vec, + outbound: Vec>, + mesh: Arc, + session: SessionContext, + n_workers: u32, +) -> Result<()> { + let mut routing = HashMap::new(); + for fragment in &fragments { + routing.insert(fragment.stage_id, fragment.routing.clone()); + } + // One sink serves every stage this proc produces; it owns the base outbound senders and routes + // each (stage, partition) to the destination proc's send end. + let worker_sink = ShmMqWorkerSink { + outbound, + mesh: Arc::clone(&mesh), + n_workers, + routing, + }; + + let mut prepared = Vec::with_capacity(fragments.len()); + for fragment in &fragments { + let task_ctx = fragment_task_ctx(&session, fragment.task_idx, fragment.task_count); + let plan = { + let dist = Arc::new(DistributedExec::new(Arc::clone(&fragment.plan))); + dist.prepare_in_process_plan(&task_ctx)? + }; + let n_out = plan.output_partitioning().partition_count(); + let mut sinks: Vec> = Vec::with_capacity(n_out); + for q in 0..n_out { + sinks.push(worker_sink.open_partition(fragment.stage_id as usize, q)?); + } + prepared.push((plan, sinks, task_ctx)); + } + // Drop the base senders so the only senders left are the per-partition clones the fragment + // futures own; otherwise the rings never observe the last-sender detach. + drop(worker_sink); + + let mut futures = Vec::with_capacity(prepared.len()); + for (plan, sinks, task_ctx) in prepared { + futures.push(run_worker_fragment(plan, sinks, task_ctx)); + } + for r in futures::future::join_all(futures).await { + r?; + } + Ok(()) +} + +/// Test-harness [`WorkerSink`]: routes each `(stage, partition)` to the destination proc's outbound +/// send end, the in-process analog of what pg_search builds on a real backend. Holds the base +/// senders plus the per-stage routing so `open_partition` reproduces the header + cooperative-drain +/// wiring the produce loop used to apply inline. +struct ShmMqWorkerSink { + outbound: Vec>, + mesh: Arc, + n_workers: u32, + routing: HashMap, +} + +impl WorkerSink for ShmMqWorkerSink { + fn open_partition(&self, stage: usize, partition: usize) -> Result> { + let routing = self.routing.get(&(stage as u32)).ok_or_else(|| { + DataFusionError::Internal(format!("run_worker_proc: no routing for stage {stage}")) + })?; + let dest_proc = match routing { + FragmentRouting::Coalesce { dest_proc } => *dest_proc, + FragmentRouting::Hashed { consumer_task, .. } => { + proc_for_task(self.n_workers, consumer_task[partition]) + } + }; + let base = self.outbound[dest_proc as usize].as_ref().ok_or_else(|| { + DataFusionError::Internal(format!( + "run_worker_proc: no outbound sender for dest proc {dest_proc}" + )) + })?; + let sender = base + .clone_with_header(MppFrameHeader::batch( + stage as u32, + partition as u32, + self.mesh.this_proc, + )) + .with_cooperative_drain(Arc::clone(&self.mesh) as Arc); + Ok(Box::new(MppPartitionSink::new(sender))) + } +} + +/// Splits an in-memory `DataSourceExec` leaf across tasks, the in-memory analog of the crate's +/// `FileScanConfigTaskEstimator` (which only handles file scans). Each task reads a disjoint subset +/// of the source's partitions, so a gather over the tasks reproduces the serial result exactly. +#[derive(Debug)] +struct MemShardEstimator { + n_tasks: usize, +} + +impl MemShardEstimator { + fn mem_source(plan: &Arc) -> Option<&MemorySourceConfig> { + plan.downcast_ref::()? + .data_source() + .downcast_ref::() + } +} + +impl TaskEstimator for MemShardEstimator { + fn task_estimation( + &self, + plan: &Arc, + _cfg: &ConfigOptions, + ) -> Option { + Self::mem_source(plan).map(|_| TaskEstimation::desired(self.n_tasks)) + } + + fn scale_up_leaf_node( + &self, + plan: &Arc, + task_count: usize, + _cfg: &ConfigOptions, + ) -> Result>> { + if task_count <= 1 { + return Ok(None); + } + let Some(mem) = Self::mem_source(plan) else { + return Ok(None); + }; + let parts = mem.partitions().to_vec(); + let n_part = parts.len(); + // The stored batches are unprojected; reuse the source's exact schema + projection so each + // variant's projected output schema matches the original leaf. + let unprojected_schema: SchemaRef = parts + .iter() + .flatten() + .next() + .map(|b| b.schema()) + .unwrap_or_else(|| plan.schema()); + let projection = mem.projection().clone(); + let variants = (0..task_count).map(|i| { + // Keep every variant at the original partition count (pad with empties) so + // DistributedLeafExec's same-partition-count contract holds; round-robin the + // non-empty partitions so each task reads a disjoint slice. + let per_task: Vec> = (0..n_part) + .map(|j| { + if j % task_count == i { + parts[j].clone() + } else { + Vec::new() + } + }) + .collect(); + MemorySourceConfig::try_new_exec( + &per_task, + unprojected_schema.clone(), + projection.clone(), + ) + .expect("memory variant") as Arc + }); + Ok(Some(Arc::new(DistributedLeafExec::try_new( + Arc::clone(plan), + variants, + )?))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::datasource::MemTable; + use futures::TryStreamExt; + + /// Total procs = leader (proc 0) + `N_WORKERS` producers. With round-robin `proc_for_task`, + /// worker proc `p` runs producer task `p - 1`. + const N_WORKERS: u32 = 3; + + fn table_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int32, false), + ])) + } + + /// `N_WORKERS` partitions, two rows each, so the shard estimator hands one partition per task. + fn table_partitions() -> Vec> { + let schema = table_schema(); + (0..N_WORKERS as i32) + .map(|p| { + let ids = Int32Array::from(vec![p * 2, p * 2 + 1]); + let vals = Int32Array::from(vec![p * 20, p * 20 + 10]); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(vals)]) + .unwrap(); + vec![batch] + }) + .collect() + } + + fn register_table(ctx: &SessionContext) { + let table = MemTable::try_new(table_schema(), table_partitions()).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + } + + /// A worker is a consumer too (it reads shuffle inputs), so when one of its input streams drops + /// early it has to cancel that stream's producer, not just the leader. This checks the wiring + /// end-to-end: a worker proc's `cancel_stream` reaches the producing proc's inbox through the + /// control senders `worker_setup` installs, and the producer records it. The producer then ends + /// the stream cleanly, which `producer_send_ends_when_consumer_cancels_the_stream` covers. + #[test] + fn worker_consumer_cancels_its_producer() { + use crate::shm::transport::CooperativeDrainSet; + + // procs: leader 0, plus workers 1 and 2. + let boot = bootstrap_mesh(3); + let consumer = Arc::clone(&boot.workers[0].1); // proc 1 + let producer = Arc::clone(&boot.workers[1].1); // proc 2 + + assert!(!producer.stream_cancelled(7, 0)); + + // Proc 1 abandons the `(stage 7, partition 0)` stream it reads from proc 2. + consumer.cancel_stream(2, 7, 0); + + // Proc 2 drains its inbox and sees the cancel its consumer sent. + producer.try_drain_pass().unwrap(); + assert!(producer.stream_cancelled(7, 0)); + // Scoped to that one stream: a sibling partition stays live. + assert!(!producer.stream_cancelled(7, 1)); + } + + fn build_session(mesh: Arc) -> SessionContext { + let config = SessionConfig::new().with_target_partitions(N_WORKERS as usize); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(config) + .with_distributed_option_extension(DistributedConfig::default()) + .with_distributed_worker_resolver(InProcessWorkerResolver::new(N_WORKERS as usize)) + .with_distributed_channel_resolver(ShmChannelResolver::new(mesh)) + .with_distributed_task_estimator(MemShardEstimator { + n_tasks: N_WORKERS as usize, + }) + .with_distributed_planner() + .build(); + let ctx = SessionContext::new_with_state(state); + register_table(&ctx); + ctx + } + + /// Collect the `id` column out of a batch list, in row order. + fn ids_of(batches: &[RecordBatch]) -> Vec { + let mut out = Vec::new(); + for b in batches { + let col = b + .column(0) + .as_any() + .downcast_ref::() + .expect("id column"); + out.extend((0..col.len()).map(|i| col.value(i))); + } + out + } + + /// A real distributed query runs end-to-end through the shm_mq transport with no Postgres and + /// no Flight: producer fragments push through the heap-backed mesh while the leader gathers via + /// `prepare_in_process_plan`. The gathered, ordered result must match the serial reference. + /// + /// Single-threaded runtime on purpose: it mirrors the per-process current-thread runtime each + /// PG worker runs, and it's exactly the cooperative model the transport is built for (the + /// producer send spin and the consumer pull loop interleave by yielding, not by parallelism). + #[tokio::test(flavor = "current_thread")] + async fn in_process_distributed_query_matches_serial() { + let query = "SELECT id, val FROM t ORDER BY id"; + + // Serial reference: same query, no distribution. + let serial_ctx = SessionContext::new(); + register_table(&serial_ctx); + let expected = serial_ctx + .sql(query) + .await + .unwrap() + .collect() + .await + .unwrap(); + let expected_ids = ids_of(&expected); + assert_eq!(expected_ids, vec![0, 1, 2, 3, 4, 5]); + + // One heap region stands in for the DSM segment; size it for n_procs = leader + workers. + let n_procs = N_WORKERS + 1; + let region_total = dsm_region_bytes(n_procs, IN_PROCESS_QUEUE_BYTES, 0).unwrap(); + let region = HeapRegion::new(region_total); + let base = SharedBase(region.base()); + let wakeup: Arc = Arc::new(NoopWakeup); + + // Leader first (it initializes the rings), then each worker attaches. No plan bytes travel + // through the region here; all roles share the producer subplans as Arcs. + let leader_mesh = unsafe { + leader_setup( + base.0, + n_procs, + IN_PROCESS_QUEUE_BYTES, + &[], + Arc::clone(&wakeup), + receiver_token(0), + Arc::new(NoInterrupt), + /* attach_senders */ false, + ) + } + .unwrap() + .mesh; + let mut worker_setups = Vec::new(); + for proc_idx in 1..n_procs { + let attach = unsafe { + worker_setup( + base.0, + region_total, + proc_idx, + Arc::clone(&wakeup), + receiver_token(proc_idx), + Arc::new(NoInterrupt), + ) + } + .unwrap(); + worker_setups.push((proc_idx, attach.mesh, attach.outbound_senders)); + } + + // Build the distributed plan once on the leader session; producers and consumer share it. + let leader_ctx = build_session(Arc::clone(&leader_mesh)); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + assert!( + physical.is::(), + "expected a DistributedExec root, got {}", + physical.name() + ); + + let entries = collect_dispatched_stages(&physical, N_WORKERS); + // Guard against a planner change that silently collapses the query to a trivial one-task + // gather: the producer stage must actually fan across every worker, or the transport's + // multi-task routing never gets exercised. + assert!( + entries.iter().any(|e| e.task_count == N_WORKERS as usize), + "expected a producer stage with task_count = {N_WORKERS}; got {:?}", + entries.iter().map(|e| e.task_count).collect::>() + ); + + // Launch the producer fragments before the leader pulls, so the mesh has data flowing while + // the consumer drains. On the current-thread runtime the spawned tasks interleave with the + // consumer by yielding, the same cooperative model each PG worker runs. + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in worker_setups { + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh)); + workers.spawn(run_worker_proc( + fragments, outbound, mesh, session, N_WORKERS, + )); + } + + // Leader consumer: prepare the head stage and execute it; the network boundary nodes pull + // from the mesh through ShmWorkerChannel::execute_task. + let leader_task_ctx = leader_ctx.task_ctx(); + let dist = physical + .downcast_ref::() + .expect("DistributedExec"); + let head = dist.prepare_in_process_plan(&leader_task_ctx).unwrap(); + let stream = head.execute(0, leader_task_ctx).unwrap(); + let got: Vec = stream.try_collect().await.unwrap(); + + while let Some(res) = workers.join_next().await { + res.expect("worker task panicked").expect("worker proc"); + } + + let got_ids = ids_of(&got); + assert_eq!(got_ids, expected_ids, "distributed gather != serial"); + + // `region` is declared before the meshes, so reverse drop order frees it after every + // receiver handle into it is gone. + } + + /// Mesh bootstrap shared by the tests: leader first (it initializes the rings), then each + /// worker attaches. + struct Bootstrap { + leader_mesh: Arc, + workers: Vec<(u32, Arc, Vec>)>, + // Last field on purpose: struct fields drop in declaration order, so the region + // outlives every receiver handle into it. + _region: HeapRegion, + } + + fn bootstrap_mesh(n_procs: u32) -> Bootstrap { + let region_total = dsm_region_bytes(n_procs, IN_PROCESS_QUEUE_BYTES, 0).unwrap(); + let region = HeapRegion::new(region_total); + let base = SharedBase(region.base()); + let wakeup: Arc = Arc::new(NoopWakeup); + let leader_mesh = unsafe { + leader_setup( + base.0, + n_procs, + IN_PROCESS_QUEUE_BYTES, + &[], + Arc::clone(&wakeup), + receiver_token(0), + Arc::new(NoInterrupt), + /* attach_senders */ false, + ) + } + .unwrap() + .mesh; + let mut workers = Vec::new(); + for proc_idx in 1..n_procs { + let attach = unsafe { + worker_setup( + base.0, + region_total, + proc_idx, + Arc::clone(&wakeup), + receiver_token(proc_idx), + Arc::new(NoInterrupt), + ) + } + .unwrap(); + workers.push((proc_idx, attach.mesh, attach.outbound_senders)); + } + Bootstrap { + leader_mesh, + workers, + _region: region, + } + } + + /// A `GROUP BY` plans a nested `NetworkShuffleExec`, so this exercises hash-routed + /// worker-to-worker traffic and the self-loop sender, which the plain gather test never + /// touches. That routing is the main thing an upstream rebase can silently break. + #[tokio::test(flavor = "current_thread")] + async fn in_process_shuffle_query_matches_serial() { + let query = "SELECT val, count(*) AS c FROM t GROUP BY val ORDER BY val"; + + let serial_ctx = SessionContext::new(); + register_table(&serial_ctx); + let expected = serial_ctx + .sql(query) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let boot = bootstrap_mesh(N_WORKERS + 1); + let leader_ctx = build_session(Arc::clone(&boot.leader_mesh)); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let entries = collect_dispatched_stages(&physical, N_WORKERS); + assert!( + entries + .iter() + .any(|e| matches!(e.routing, FragmentRouting::Hashed { .. })), + "expected a hash-routed producer stage; got {:?}", + entries.iter().map(|e| &e.routing).collect::>() + ); + + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in boot.workers { + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh)); + workers.spawn(run_worker_proc( + fragments, outbound, mesh, session, N_WORKERS, + )); + } + + let leader_task_ctx = leader_ctx.task_ctx(); + let dist = physical + .downcast_ref::() + .expect("DistributedExec"); + let head = dist.prepare_in_process_plan(&leader_task_ctx).unwrap(); + let stream = head.execute(0, leader_task_ctx).unwrap(); + let got: Vec = stream.try_collect().await.unwrap(); + + while let Some(res) = workers.join_next().await { + res.expect("worker task panicked").expect("worker proc"); + } + + use datafusion::arrow::util::pretty::pretty_format_batches; + assert_eq!( + pretty_format_batches(&expected).unwrap().to_string(), + pretty_format_batches(&got).unwrap().to_string(), + "distributed shuffle != serial" + ); + } + + /// A producer that attaches and then goes away without sending its EOFs must fail the + /// gather, not hang it: the drain fails the channels the dead receiver fed once the ring + /// detaches. + #[tokio::test(flavor = "current_thread")] + async fn producer_loss_fails_the_gather_instead_of_hanging() { + let query = "SELECT id, val FROM t ORDER BY id"; + + let boot = bootstrap_mesh(N_WORKERS + 1); + let leader_ctx = build_session(Arc::clone(&boot.leader_mesh)); + let physical = leader_ctx + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let entries = collect_dispatched_stages(&physical, N_WORKERS); + + let mut workers = JoinSet::new(); + for (proc_idx, mesh, outbound) in boot.workers { + if proc_idx == 1 { + // Simulated crash: the proc attached (its senders exist), then dies without + // running its fragments or sending EOF. Dropping the senders is what process + // exit does. + drop(outbound); + drop(mesh); + continue; + } + let fragments = fragments_for_proc(&entries, proc_idx, N_WORKERS); + let session = build_session(Arc::clone(&mesh)); + workers.spawn(run_worker_proc( + fragments, outbound, mesh, session, N_WORKERS, + )); + } + + let leader_task_ctx = leader_ctx.task_ctx(); + let dist = physical + .downcast_ref::() + .expect("DistributedExec"); + let head = dist.prepare_in_process_plan(&leader_task_ctx).unwrap(); + let stream = head.execute(0, leader_task_ctx).unwrap(); + let res: Result, _> = stream.try_collect().await; + + while let Some(r) = workers.join_next().await { + r.expect("worker task panicked").expect("worker proc"); + } + + let err = res + .expect_err("gather must fail when a producer goes away") + .to_string(); + assert!( + err.contains("detached before this channel's EOF"), + "unexpected error: {err}" + ); + } +} diff --git a/src/shm/mesh.rs b/src/shm/mesh.rs new file mode 100644 index 00000000..4421cd73 --- /dev/null +++ b/src/shm/mesh.rs @@ -0,0 +1,375 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! MPP mesh adapters around `DsmMpscRing` plus the alignment helpers used to size +//! the per-inbox DSM regions. + +use datafusion::common::DataFusionError; + +use super::mpsc_ring::{ + DsmMpscReceiver, DsmMpscSender, RecvOutcome as MpscRecvOutcome, SendError as MpscSendError, +}; +use super::transport::{BatchChannelReceiver, BatchChannelSender, RecvOutcome}; + +/// Postgres MAXALIGN. The DSM layout must agree between this crate (which computes the per-inbox +/// offsets) and whatever shared buffer the embedder hands it. PG's `MAXIMUM_ALIGNOF` is 8 on every +/// supported platform, so the shared-memory layout pins it. +const MAXALIGN: usize = 8; + +/// MAXALIGN-DOWN of `queue_bytes`, so every per-inbox region in the grid is the same size and +/// offset math is a simple multiply. +#[inline] +pub(super) fn aligned_queue_bytes(queue_bytes: usize) -> usize { + queue_bytes & !(MAXALIGN - 1) +} + +/// MAXALIGN-UP `n` to the next multiple of [`MAXALIGN`]. Returns `None` on overflow. Used by +/// [`super::dsm::compute_dsm_layout`] to keep section boundaries MAXALIGN-aligned. +#[inline] +pub(super) fn align_up_maxalign_checked(n: usize) -> Option { + let mask = MAXALIGN - 1; + n.checked_add(mask).map(|x| x & !mask) +} + +/// DSM MPSC ring as a `BatchChannelSender`. Multiple producer processes hold their own +/// `DsmInboxSender` clones targeting the same receiver inbox; the ring serializes them +/// via Vyukov CAS on `tail`. Because the ring's atomics are the synchronization point, +/// the `send_bytes` / `try_send_bytes` paths don't need an attach-thread `debug_assert!`. +/// +/// Detach-on-drop: `DsmMpscSender::Drop` decrements `sender_count`; the last drop flips +/// `detached` and wakes the receiver, mirroring shm_mq's "drop the last sender, receiver +/// sees detach" guarantee. +pub(super) struct DsmInboxSender { + inner: DsmMpscSender, + send_lock: tokio::sync::Mutex<()>, +} + +// Send + Sync auto-derive. Both fields are Send + Sync (`DsmMpscSender` via its unsafe +// impls; `tokio::sync::Mutex` by definition), so no manual `unsafe impl` is needed; the +// auto-derive also surfaces a compile error if a future field is `!Send` / `!Sync`. + +impl DsmInboxSender { + /// Wrap a `DsmMpscSender` for use through the `BatchChannelSender` trait. + pub(super) fn new(inner: DsmMpscSender) -> Self { + Self { + inner, + send_lock: tokio::sync::Mutex::new(()), + } + } + + fn map_send_err(err: MpscSendError) -> DataFusionError { + match err { + MpscSendError::Detached => { + DataFusionError::Execution("mpp: DSM MPSC inbox detached".into()) + } + MpscSendError::MessageTooLarge => DataFusionError::Execution( + "mpp: DSM MPSC frame exceeds the entire ring capacity \ + (raise the embedder's queue-size knob)" + .into(), + ), + MpscSendError::Full => DataFusionError::Execution( + "mpp: DSM MPSC inbox full (caller should retry via try_send_bytes)".into(), + ), + } + } +} + +impl BatchChannelSender for DsmInboxSender { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError> { + // Fallback for callers that didn't wire `with_cooperative_drain`. The real send + // path drives `try_send_bytes` through the cooperative spin in `transport.rs`; + // this loop just spins on `yield_now` and burns the backend core under a slow + // consumer. Hitting it in production means a missing `with_cooperative_drain` + // on the fragment, not a real backpressure path. + loop { + match self.inner.try_send(bytes) { + Ok(()) => return Ok(()), + Err(MpscSendError::Full) => std::thread::yield_now(), + Err(e) => return Err(Self::map_send_err(e)), + } + } + } + + fn try_send_bytes(&self, bytes: &[u8]) -> Result { + match self.inner.try_send(bytes) { + Ok(()) => Ok(true), + Err(MpscSendError::Full) => Ok(false), + Err(e) => Err(Self::map_send_err(e)), + } + } + + fn send_lock(&self) -> &tokio::sync::Mutex<()> { + &self.send_lock + } +} + +/// DSM MPSC ring as a `BatchChannelReceiver`. The scratch `Vec` lives behind a `Mutex` so a +/// `&self` `try_recv` can hand the populated buffer back via `mem::take` without `RefCell` runtime +/// borrow tracking. +/// +/// Single-consumer comes from the call pattern: one `DsmInboxReceiver` per process, +/// owned by `DrainHandle::cooperative_receivers`, polled inline from `try_drain_pass`. +/// No two threads ever race on the same receiver; the mutex is just interior-mutability +/// boilerplate, uncontended in production. +pub(super) struct DsmInboxReceiver { + inner: DsmMpscReceiver, + /// Scratch the inner primitive's `try_recv` reuses across calls (via reserve+set_len). + /// We `mem::take` it on every Bytes outcome to hand the bytes to the caller without a + /// copy; the inner primitive re-grows from `Vec::new()` on the next call. + scratch: std::sync::Mutex>, +} + +// `DsmMpscReceiver` is deliberately `!Sync` (single-consumer invariant). We promote +// `DsmInboxReceiver` to Sync because the caller pattern guarantees only one thread +// touches it at a time, and `Mutex>` protects the scratch from shared-reference +// UB if that invariant ever slips. Send is auto-derived. +unsafe impl Sync for DsmInboxReceiver {} + +impl DsmInboxReceiver { + /// Wrap a `DsmMpscReceiver` for use through the `BatchChannelReceiver` trait. + pub(super) fn new(inner: DsmMpscReceiver) -> Self { + Self { + inner, + scratch: std::sync::Mutex::new(Vec::new()), + } + } + + /// Register the consumer's opaque wakeup token on the underlying ring. The embedder packs + /// whatever its [`super::mpsc_ring::Wakeup`] interprets (a PG embedder packs `(pgprocno, pid)`). + pub(super) fn set_receiver(&self, token: u64) { + self.inner.set_receiver(token); + } +} + +impl BatchChannelReceiver for DsmInboxReceiver { + fn try_recv(&self) -> RecvOutcome { + let mut buf = self.scratch.lock().expect("scratch mutex poisoned"); + match self.inner.try_recv(&mut buf) { + MpscRecvOutcome::Bytes => { + // Hand the populated buffer to the caller and leave an empty Vec behind. + // The inner primitive's `try_recv` does its own `reserve(len)` on the + // next call, so we don't need to pre-allocate scratch capacity here. + RecvOutcome::Bytes(std::mem::take(&mut *buf)) + } + MpscRecvOutcome::Empty => RecvOutcome::Empty, + MpscRecvOutcome::Detached => RecvOutcome::Detached, + } + } +} + +#[cfg(test)] +mod tests { + use super::super::mpsc_ring::{self, DsmMpscRingHeader, Wakeup}; + use super::*; + + struct NoopWakeup; + impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} + } + fn test_wakeup() -> std::sync::Arc { + std::sync::Arc::new(NoopWakeup) + } + + /// Allocate a fresh ring (heap, aligned) and return `(owning region, sender, receiver)`. + /// Pairs `DsmInboxSender` + `DsmInboxReceiver` over a heap-allocated ring matching the + /// alignment contract `create_at` requires. Production allocates the region inside a + /// `dsm_segment`; this helper exists so the BatchChannel trait impls can be exercised + /// without a PG backend. + /// + /// The region is returned FIRST so callers bind it first. Rust drops locals in reverse + /// declaration order, so the region (bound first) drops LAST, after the sender/receiver whose + /// `Drop` touches the ring memory (`DsmMpscSender::drop` does `sender_count.fetch_sub`). + /// Returning the region last instead frees the bytes before those `Drop`s run, a + /// use-after-free that corrupts the heap shared with parallel tests. + fn test_dsm_inbox_pair( + ring_size: u32, + slot_capacity: u32, + ) -> (AlignedTestRegion, DsmInboxSender, DsmInboxReceiver) { + let bytes = DsmMpscRingHeader::region_bytes(ring_size, slot_capacity); + let region = AlignedTestRegion::new(bytes); + let header_ptr = + unsafe { mpsc_ring::create_at(region.as_mut_ptr(), ring_size, slot_capacity) }; + let nn = std::ptr::NonNull::new(header_ptr).expect("create_at returned null"); + let alive = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)); + let sender = DsmInboxSender::new(unsafe { + DsmMpscSender::new(nn, test_wakeup(), std::sync::Arc::clone(&alive)) + }); + let receiver = DsmInboxReceiver::new(unsafe { DsmMpscReceiver::new(nn, alive) }); + (region, sender, receiver) + } + + struct AlignedTestRegion { + ptr: *mut u8, + layout: std::alloc::Layout, + } + + impl AlignedTestRegion { + fn new(bytes: usize) -> Self { + let align = std::mem::align_of::(); + let layout = std::alloc::Layout::from_size_align(bytes, align).expect("layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null()); + Self { ptr, layout } + } + fn as_mut_ptr(&self) -> *mut u8 { + self.ptr + } + } + + impl Drop for AlignedTestRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } + } + + #[test] + fn aligned_queue_bytes_rounds_down_to_maxalign() { + let maxalign = MAXALIGN; + for req in [0, 1, 7, 8, 15, 16, 1024, 1025, 1_048_576, 8 * 1024 * 1024] { + let got = aligned_queue_bytes(req); + assert!(got <= req); + assert_eq!(got % maxalign, 0, "not aligned: {got} for req {req}"); + assert!(req - got < maxalign); + } + } + + #[test] + fn align_up_maxalign_rounds_up_to_maxalign() { + let maxalign = MAXALIGN; + for req in [0, 1, 7, 8, 15, 16, 1024, 1025] { + let got = align_up_maxalign_checked(req).unwrap(); + assert!(got >= req); + assert_eq!(got % maxalign, 0); + assert!(got - req < maxalign); + } + assert!(align_up_maxalign_checked(usize::MAX).is_none()); + } + + #[test] + fn dsm_inbox_batch_channel_round_trip() { + let (_region, tx, rx) = test_dsm_inbox_pair(4, 64); + // Sanity: try_send_bytes succeeds, try_recv hands back the bytes, send_lock + // returns a usable Mutex. + assert!(tx.try_send_bytes(b"hello").unwrap()); + assert!(tx.try_send_bytes(b"world").unwrap()); + match rx.try_recv() { + RecvOutcome::Bytes(b) => assert_eq!(&b[..], b"hello"), + other => panic!("expected Bytes(hello), got {other:?}"), + } + match rx.try_recv() { + RecvOutcome::Bytes(b) => assert_eq!(&b[..], b"world"), + other => panic!("expected Bytes(world), got {other:?}"), + } + assert!(matches!(rx.try_recv(), RecvOutcome::Empty)); + // send_lock is just exercised: the per-instance Mutex satisfies the trait. + let _guard = tx.send_lock(); + } + + #[test] + fn dsm_inbox_try_send_returns_false_when_full() { + let (_region, tx, rx) = test_dsm_inbox_pair(2, 64); + assert!(tx.try_send_bytes(b"a").unwrap()); + assert!(tx.try_send_bytes(b"b").unwrap()); + // Third send should hit Full (returns Ok(false), not Err). + assert!(!tx.try_send_bytes(b"c").unwrap()); + // Drain one, then send again succeeds. + assert!(matches!(rx.try_recv(), RecvOutcome::Bytes(_))); + assert!(tx.try_send_bytes(b"c").unwrap()); + } + + #[test] + fn dsm_inbox_multi_producer_through_trait_surface() { + use std::sync::Arc; + // Build a real Arc shared across threads to confirm + // dyn dispatch + Send/Sync impls compile and behave. + let (_region, tx, rx) = test_dsm_inbox_pair(64, 32); + let tx: Arc = Arc::new(tx); + let mut handles = Vec::new(); + const K: usize = 4; + const M: u32 = 200; + for producer_id in 0..K { + let tx = Arc::clone(&tx); + handles.push(std::thread::spawn(move || { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + let mut sent = 0u32; + while sent < M { + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + match tx.try_send_bytes(&payload) { + Ok(true) => sent += 1, + Ok(false) => std::thread::yield_now(), + Err(e) => panic!("send failed: {e}"), + } + } + })); + } + let target = K * M as usize; + let mut seen = vec![vec![false; M as usize]; K]; + let mut got = 0usize; + while got < target { + match rx.try_recv() { + RecvOutcome::Bytes(b) => { + let producer_id = u32::from_le_bytes(b[0..4].try_into().unwrap()) as usize; + let idx = u32::from_le_bytes(b[4..8].try_into().unwrap()) as usize; + let already = std::mem::replace(&mut seen[producer_id][idx], true); + assert!(!already, "dup ({producer_id}, {idx})"); + got += 1; + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + } + + /// Dropping the last `DsmInboxSender` flips `detached` and the receiver sees the + /// queued bytes followed by `Detached`. This is the structural equivalent of + /// shm_mq's "drop the last sender, receiver sees detach" guarantee, and is what + /// keeps the drain loop from wedging on a clean shutdown. + #[test] + fn dropping_last_sender_triggers_detach() { + let (_region, tx, rx) = test_dsm_inbox_pair(4, 64); + tx.try_send_bytes(b"final").unwrap(); + drop(tx); + // The queued frame is still readable. + match rx.try_recv() { + RecvOutcome::Bytes(b) => assert_eq!(&b[..], b"final"), + other => panic!("expected Bytes(final), got {other:?}"), + } + // Then Detached. + assert!(matches!(rx.try_recv(), RecvOutcome::Detached)); + } + + #[test] + fn try_send_bytes_rejects_oversize_payload() { + // Multi-slot fragmentation lifted the per-slot ceiling; only payloads larger + // than `ring_size * (slot_capacity - SLOT_HEADER_BYTES)` are now rejected. + // ring_size=2, slot_capacity=32 -> payload_cap=16, ring-wide cap=32. 64 + // bytes still doesn't fit. + let (_region, tx, _rx) = test_dsm_inbox_pair(2, 32); + let oversize = vec![0u8; 64]; + let err = tx + .try_send_bytes(&oversize) + .expect_err("expected MessageTooLarge"); + assert!( + format!("{err}").contains("exceeds the entire ring capacity"), + "unexpected error: {err}" + ); + } +} diff --git a/src/shm/mod.rs b/src/shm/mod.rs new file mode 100644 index 00000000..bf5c95e1 --- /dev/null +++ b/src/shm/mod.rs @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared-memory transport. +//! +//! A non-gRPC [`ChannelResolver`] for co-located execution, where "workers" are tasks or +//! parallel processes sharing one machine and communicating over a shared-memory mesh rather +//! than gRPC. The transport-mechanism pieces (the MPSC ring, framing, routing, cooperative +//! drain) live here as a reusable library; an embedder supplies the two platform primitives via +//! small extension points: how to allocate the shared buffer, and how to wake a blocked consumer. +//! +//! The point of hosting it in this crate is testing: the in-process instantiation runs real +//! distributed queries through the transport in this crate's CI, so an upstream rebase that +//! breaks the channel-protocol contract fails here, before any downstream embedder rebuilds. +//! +//! [`ChannelResolver`]: crate::ChannelResolver +//! +//! Two assumptions an embedder signs up for: +//! - Execution is cooperative on a current-thread runtime: consumers spin on +//! `try_pop` + `yield_now` and producers drain their own inbound while blocked, instead of +//! parking on the `Wakeup` extension point. On a multi-thread runtime each stream burns a core while +//! idle. +//! - Inbound frames demux into unbounded per-channel buffers, so a consumer that falls behind +//! buffers the in-flight intermediate result in process memory. The rings in shared memory +//! stay bounded; the overflow lives on the consumer's heap. + +use std::sync::Arc; +use std::sync::atomic::AtomicBool; + +mod dsm; +mod mesh; +mod mpsc_ring; +mod runtime; +// Deferred: the self-hosting default transport was built on the removed `WorkerTransport`/ +// `WorkerDispatch` dispatch umbrella, which the `ChannelResolver` model has no analog for; its +// no-gRPC-default role is now served by `InProcessChannelResolver` and its ring-exercising role by +// the `in_process` test, so it stays gated out until reimplemented on `coordinator_channel`. +#[cfg(any())] +mod self_hosted; +mod setup; +mod sink; +mod transport; + +// Curated public surface an embedder consumes. The embedder allocates the shared buffer and +// supplies the two extension points (`Wakeup`, `Interrupt`); everything else is built here. +pub use mpsc_ring::{NO_RECEIVER_TOKEN, Wakeup}; +pub use runtime::{InProcessWorkerResolver, MppMesh, ShmChannelResolver, proc_for_task}; +pub use setup::{ + LeaderAttach, WorkerAttach, collect_task_metrics, dsm_region_bytes, install_work_unit_channels, + leader_setup, region_total, run_worker_fragment, worker_setup, +}; +pub use sink::{PartitionSink, WorkerSink}; +pub use transport::{ + CooperativeDrainSet, Interrupt, MppFrameHeader, MppPartitionSink, MppSender, NoInterrupt, + SendBatchStats, SetPlanFrame, +}; + +/// Out-of-DSM liveness flag shared by the ring handles from one attach. The embedder flips it to +/// `false` from its dsm-detach callback while the segment is still mapped, so a handle dropped +/// afterward (e.g. by a memory-context reset) no-ops instead of dereferencing freed memory. +pub type AliveFlag = Arc; + +// In-process instantiation + the end-to-end test that runs a real distributed query through the +// transport with no Postgres. Test-only: it's how an upstream rebase that breaks the transport +// contract fails in this crate's CI. +#[cfg(test)] +mod in_process; diff --git a/src/shm/mpsc_ring.rs b/src/shm/mpsc_ring.rs new file mode 100644 index 00000000..55c7390b --- /dev/null +++ b/src/shm/mpsc_ring.rs @@ -0,0 +1,1550 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! DSM-backed MPSC ring for MPP mesh inboxes. +//! +//! PG's `shm_mq` is hard-wired SPSC (asserted in `shm_mq_set_sender`), so one inbox +//! can't be shared across N-1 peers without forking PG. This ring is the replacement: +//! a fixed-size byte-message ring sitting in a `dsm_segment`, MPSC-correct via +//! Vyukov-style per-slot sequence counters. +//! +//! Layout (contiguous bytes, `repr(C)`): +//! +//! ```text +//! +- DsmMpscRingHeader -----------------------+ +//! | magic, version, ring_size, slot_capacity | +//! | sender_count, detached, receiver_packed | +//! | (cache-line padding around head/tail) | +//! | head, tail (each on its own line) | +//! +-------------------------------------------+ +//! | Slot[0] | Slot[1] | ... | Slot[N-1] | +//! +-------------------------------------------+ +//! +//! Slot { +//! seq: AtomicU64, // Vyukov phase counter +//! len: AtomicU32, +//! data: [u8; slot_capacity - SLOT_HEADER_SIZE] +//! } +//! ``` +//! +//! Slot phase encoding (Vyukov MPMC reduced to MPSC): +//! +//! ```text +//! slot[i] in round k: seq = k * ring_size + i // empty +//! seq = k * ring_size + i + 1 // ready +//! ``` +//! +//! Producer claim at `tail = T`: read `slot[T % ring_size].seq`. `seq == T` then CAS +//! `tail: T → T+1`, winner copies payload and stores `seq = T + 1` (Release). `seq < T` +//! means ring full. `seq > T` means another producer took `T`, retry. Winner +//! `SetLatch`es the receiver. +//! +//! Consumer take at `head = H`: `slot[H % ring_size].seq == H + 1` means ready. Read +//! payload, store `seq = H + ring_size` (next round's empty marker), then `head = H+1`. +//! The single consumer owns `head` without CAS; producers contend only on `tail`. +//! +//! **Safety**: public methods on `DsmMpscSender` / `DsmMpscReceiver` are type-safe once +//! constructed. The constructors are `unsafe` because the DSM region must be correctly +//! sized and not aliased (one process calls `create_at`, everyone else `attach_at`). +//! +//! **Counter wraparound**: head/tail are `u64`, incremented by one per op. At 100M +//! ops/sec that's ~5800 years, so we ignore wrap. The seq math would break under wrap +//! (pre-wrap `seq` would exceed post-wrap `tail` and producers would spin forever); if +//! that ever matters, add a `tail < u64::MAX - margin` check and reset the ring. + +use std::ptr::NonNull; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; + +use std::sync::Arc; + +use super::AliveFlag; + +/// Wakes the ring's single consumer after a producer publishes a frame. +/// +/// The ring is transport-agnostic: the consumer registers an opaque `u64` token via +/// `DsmMpscReceiver::set_receiver`, and each publishing producer hands that token to this +/// hook. An embedder over PostgreSQL shared memory packs `(pgprocno, pid)` into the token and +/// `SetLatch`es the backend; an in-process embedder packs a registry key and unparks the +/// consumer thread. The token is stored in the ring header (shared memory), so the hook must be +/// able to resolve it from any producer. +pub trait Wakeup: Send + Sync { + fn wake(&self, token: u64); +} + +/// Sentinel token meaning "no consumer registered". Producers skip the wake when the stored +/// token equals this; `create_at` initializes the ring to it. +pub const NO_RECEIVER_TOKEN: u64 = u64::MAX; + +/// Reserved byte at the start of every slot. Sized so payload + header fits in +/// `slot_capacity` exactly. +const SLOT_HEADER_BYTES: usize = std::mem::size_of::(); + +#[repr(C)] +struct SlotHeader { + /// Vyukov sequence counter; see module docs for phase encoding. + seq: AtomicU64, + /// Bytes of payload written into THIS slot's data region; + /// `0..=slot_capacity - SLOT_HEADER_BYTES`. + len: AtomicU32, + /// Fragment metadata. Low 2 bits = [`FragmentKind`] (Complete=0 / First=1 / + /// Continue=2 / Last=3). For `First`, bits 16..32 hold the slot count of the + /// logical frame (1..=65535). For other kinds these bits are unused. + flags: AtomicU32, +} + +/// Kind bits stored in [`SlotHeader::flags`]'s low 2 bits. +/// +/// A frame fitting in `slot_capacity - SLOT_HEADER_BYTES` rides the single-slot +/// fast path as `Complete`. Anything larger spans +/// `n_slots = ceil(frame_len / per_slot_payload)` consecutive slots: `First` +/// (carrying `n_slots` in the upper bits), `Continue` for the middle, `Last` for +/// the tail. Producers grab the whole run with one +/// `tail.compare_exchange(T, T + n_slots)`, so other producers can't interleave +/// their fragments and the receiver always sees the slots in producer order. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FragmentKind { + Complete = 0, + First = 1, + Continue = 2, + Last = 3, +} + +/// Bit-mask for the kind bits in [`SlotHeader::flags`]. +const FLAGS_KIND_MASK: u32 = 0b11; +/// Shift for the `n_slots` field stored in the flags' upper half (only meaningful +/// for `First`). Limits per-frame fragmentation to 65535 slots, which is far +/// beyond any realistic ring size. +const FLAGS_NSLOTS_SHIFT: u32 = 16; +const FLAGS_NSLOTS_MAX: u32 = 0xFFFF; + +#[inline] +fn pack_flags(kind: FragmentKind, n_slots: u32) -> u32 { + debug_assert!( + kind != FragmentKind::First || (1..=FLAGS_NSLOTS_MAX).contains(&n_slots), + "First frames must carry 1..=65535 slots; got {n_slots}" + ); + (kind as u32) | (n_slots << FLAGS_NSLOTS_SHIFT) +} + +#[inline] +fn unpack_kind(flags: u32) -> Option { + match flags & FLAGS_KIND_MASK { + 0 => Some(FragmentKind::Complete), + 1 => Some(FragmentKind::First), + 2 => Some(FragmentKind::Continue), + 3 => Some(FragmentKind::Last), + _ => None, + } +} + +#[inline] +fn unpack_nslots(flags: u32) -> u32 { + flags >> FLAGS_NSLOTS_SHIFT +} + +/// Magic constant validating that an attaching process points at a `DsmMpscRingHeader` +/// rather than garbage. "MPCR" = MPSC Ring. Different value from `MppDsmHeader`'s magic +/// so a worker that picks up the wrong region fails the wrong-shape check loudly. +const MPSC_RING_MAGIC: u32 = u32::from_le_bytes(*b"MPCR"); + +/// Bump on any wire-incompatible layout change. Mirrors the discipline at +/// `MppDsmHeader::validate`. +/// +/// Wire versions: +/// - v1: `SlotHeader { seq, len, _pad }`. One frame per slot. +/// - v2: `SlotHeader { seq, len, flags }`. `flags` carries `FragmentKind` plus +/// `n_slots` on `First`, so frames bigger than +/// `slot_capacity - SLOT_HEADER_BYTES` can span N consecutive slots reserved +/// atomically by the producer. +const MPSC_RING_VERSION: u32 = 2; + +/// Assumed cache line size for false-sharing avoidance. 64 bytes covers x86_64 and arm64; +/// over-padding on smaller-cache-line targets costs a few bytes per ring, nothing more. +const CACHE_LINE: usize = 64; + +/// Ring header. Laid out first in the DSM region; slot array follows immediately after. +/// +/// Cache-line padding around `head` and `tail` isn't optional: at N=24 producer +/// contention the consumer's `head.store` and the producers' `tail.compare_exchange` +/// race on the same line, MESI-ping-ponging every claim. That's the false-sharing +/// footgun Vyukov's writeups call out; the Disruptor literature shows 5-10x throughput +/// loss from it on x86. Padding puts each hot field on its own line. +/// +/// NOT `#[repr(C, align(64))]`: PG `dsm_segment` base addresses are only +/// MAXALIGN-aligned in practice (on macOS the user-data offset can land 16-aligned +/// but not 64-aligned). Forcing `align(64)` would impose a 64-aligned destination on +/// `create_at` / `attach_at` we can't guarantee. The `_pad_*` fields below still +/// put `head`, `tail`, and the first slot on separate 64-byte regions; it's the +/// *distance* between hot fields that matters, not their absolute alignment. +#[repr(C)] +pub(super) struct DsmMpscRingHeader { + /// Magic constant; equals `MPSC_RING_MAGIC` for a valid ring. Checked in `attach_at`. + magic: u32, + /// Layout version; equals `MPSC_RING_VERSION`. Checked in `attach_at`. + version: u32, + /// Number of slots. Immutable after `create_at`. + ring_size: u32, + /// Byte capacity of each slot INCLUDING the slot header. Payload bytes per slot are + /// `slot_capacity - SLOT_HEADER_BYTES`. Immutable after `create_at`. + slot_capacity: u32, + /// Live `DsmMpscSender` count. Incremented in `DsmMpscSender::new`, decremented in + /// `DsmMpscSender::Drop`. The drop that takes the count from 1 → 0 sets `detached` + /// (with Release) and wakes the receiver, mirroring shm_mq's "drop = detach" + /// structural guarantee. + sender_count: AtomicU32, + /// Set by the consumer (or by the leader on query teardown) to tell producers to + /// fail-fast on subsequent sends. Sticky. + detached: AtomicBool, + _pad_after_detached: [u8; 3], + /// Packed `(pgprocno: i32, pid: i32)` of the registered receiver, or 0 (both + /// Opaque receiver token, set by the consumer via `set_receiver` and handed to the + /// embedder's `Wakeup` extension point on every post-publish wake. Initialized to + /// [`NO_RECEIVER_TOKEN`] (`u64::MAX`), which producers treat as "no receiver yet: + /// skip the wake". The embedder defines the token's contents (pg_search packs + /// `(pgprocno, pid)`); a single atomic keeps whatever pair it packs from being + /// observed torn mid-update. + receiver_packed: AtomicU64, + /// Padding to push `head` onto its own cache line. Header up to here uses bytes + /// 0..32; this padding fills 32..64 so `head` lands at offset 64 exactly. The + /// `header_layout_is_cache_friendly` test asserts this. + _pad_before_head: [u8; CACHE_LINE - 32], + /// Consumer's read cursor. Only the consumer writes this. Currently no producer + /// reads it (full-detection works via slot `seq`); the Release-on-store is defensive + /// for any future blocking-send variant that wants to poll consumer progress. + head: AtomicU64, + /// Padding to push `tail` onto its own cache line so the consumer's `head.store` + /// doesn't invalidate the producers' `tail` cache line. + _pad_between_head_and_tail: [u8; CACHE_LINE - 8], + /// Producers' write cursor. CAS'd to claim slot ownership for a tail value. + tail: AtomicU64, + /// Padding so the first slot doesn't share a cache line with `tail`. Producers + /// race on `tail`; the consumer's first slot read should not pull the `tail` cache + /// line into the consumer's L1 unnecessarily. + _pad_after_tail: [u8; CACHE_LINE - 8], +} + +impl DsmMpscRingHeader { + /// Bytes occupied by `ring_size` slots of `slot_capacity` each, plus the header. + pub(super) const fn region_bytes(ring_size: u32, slot_capacity: u32) -> usize { + std::mem::size_of::() + (ring_size as usize) * (slot_capacity as usize) + } +} + +// The receiver wake lives on `DsmMpscSender` (it holds the injected `Wakeup`); see +// `DsmMpscSender::wake_receiver`. + +/// Errors that `try_send` can surface to the producer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum SendError { + /// Ring is full; consumer hasn't drained enough to free a slot. + Full, + /// Receiver has detached (query teardown). Producer should stop sending. + Detached, + /// `bytes.len() + SLOT_HEADER_BYTES > slot_capacity`. The caller picked a slot + /// capacity too small for this payload; bumping `slot_capacity` at `create_at` time + /// fixes it. + MessageTooLarge, +} + +/// Outcome of a single `try_recv` call. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum RecvOutcome { + /// A frame was copied into the caller's buffer. + Bytes, + /// Ring is empty right now; try again later. + Empty, + /// All producers have detached and there's no more data to drain. + Detached, +} + +/// Caller-visible handle for the single consumer. +pub(super) struct DsmMpscReceiver { + ring: NonNull, + alive: AliveFlag, +} + +/// Caller-visible handle for any of the N-1 producers. +pub(super) struct DsmMpscSender { + ring: NonNull, + /// Hook to wake the registered consumer after a publish. Injected by the embedder so the + /// ring stays free of any process/thread-notification mechanism. + wakeup: Arc, + /// Whether this handle counts toward `sender_count`. Data senders (the producer fragments) + /// do, so the last one's drop flips `detached` and the consumer learns its producers are gone. + /// Control senders (a consumer's `Cancel` path) don't: they target a peer's inbox without being + /// one of its producers, so counting them would mask that peer's own producer-gone signal. + counts_as_data: bool, + alive: AliveFlag, +} + +// SAFETY: the ring is a `repr(C)` blob in shared memory whose atomic operations are the +// synchronization point. Both handles are stateless pointers to the same data; sending +// either across threads requires only the atomic ordering already in use. +// +// `DsmMpscReceiver` is deliberately !Sync: the type-level invariant that exactly one +// thread calls `try_recv` at a time is what makes the lock-free MPSC math correct (the +// single consumer owns `head` without CAS). `DsmMpscSender` is Sync so multiple producer +// threads can share one `Arc` and race on `tail` via CAS. +unsafe impl Send for DsmMpscReceiver {} +unsafe impl Send for DsmMpscSender {} +unsafe impl Sync for DsmMpscSender {} + +/// Initialize a freshly-allocated DSM region into a valid ring header + zeroed slot array. +/// Must be called exactly once, by the process that allocated the region (leader in our +/// case). All other processes attach via [`attach_at`] without re-initializing. +/// +/// # Safety +/// - `base` must point at the start of a region of at least +/// `DsmMpscRingHeader::region_bytes(ring_size, slot_capacity)` bytes. +/// - `ring_size >= 2` (a ring of 1 slot can't distinguish empty from full). +/// - `slot_capacity > SLOT_HEADER_BYTES` (need at least one byte of payload). +/// - The region must not be concurrently accessed by any other process or thread until +/// this returns. +pub(super) unsafe fn create_at( + base: *mut u8, + ring_size: u32, + slot_capacity: u32, +) -> *mut DsmMpscRingHeader { + debug_assert!(ring_size >= 2, "ring_size must be >= 2"); + debug_assert!( + slot_capacity as usize > SLOT_HEADER_BYTES, + "slot_capacity must leave room for at least one payload byte" + ); + // The header's natural alignment (from its largest field, AtomicU64) is 8 bytes. + // PG MAXALIGN is 8 on every supported platform, so dsm_segment user-data offsets + // computed via `align_up_maxalign_checked` always satisfy this. Tests use an + // aligned-Vec helper just to keep the assertion clean across allocators. + debug_assert!( + (base as usize).is_multiple_of(std::mem::align_of::()), + "create_at base must be aligned to {} bytes", + std::mem::align_of::() + ); + let header_ptr = base.cast::(); + // Write the immutable header fields. Use std::ptr::write so we don't construct an + // intermediate &mut that aliases the not-yet-initialized atomic fields. + unsafe { + std::ptr::write( + header_ptr, + DsmMpscRingHeader { + magic: MPSC_RING_MAGIC, + version: MPSC_RING_VERSION, + ring_size, + slot_capacity, + detached: AtomicBool::new(false), + _pad_after_detached: [0; 3], + sender_count: AtomicU32::new(0), + receiver_packed: AtomicU64::new(NO_RECEIVER_TOKEN), + _pad_before_head: [0; CACHE_LINE - 32], + head: AtomicU64::new(0), + _pad_between_head_and_tail: [0; CACHE_LINE - 8], + tail: AtomicU64::new(0), + _pad_after_tail: [0; CACHE_LINE - 8], + }, + ); + } + // Initialize slot sequences: slot[i].seq = i in round 0. + for i in 0..ring_size { + let slot = unsafe { slot_ptr(header_ptr, i, slot_capacity) }; + unsafe { + std::ptr::write( + slot, + SlotHeader { + seq: AtomicU64::new(i as u64), + len: AtomicU32::new(0), + flags: AtomicU32::new(0), + }, + ); + } + } + header_ptr +} + +/// Take an already-initialized ring header pointer and confirm its shape matches caller's +/// expectations. The caller's `expected_ring_size` / `expected_slot_capacity` must match +/// the values written at `create_at` time; mismatch is a hard error (returns null). +/// +/// # Safety +/// - `base` must point at the same region a previous `create_at` initialized. +/// - The region must not be deallocated for the lifetime of any handle returned from +/// the wrappers (`DsmMpscReceiver::new`, `DsmMpscSender::new`). +pub(super) unsafe fn attach_at( + base: *mut u8, + expected_ring_size: u32, + expected_slot_capacity: u32, +) -> Option> { + let header_ptr = base.cast::(); + let nn = NonNull::new(header_ptr)?; + let header = unsafe { nn.as_ref() }; + if header.magic != MPSC_RING_MAGIC || header.version != MPSC_RING_VERSION { + return None; + } + if header.ring_size != expected_ring_size || header.slot_capacity != expected_slot_capacity { + return None; + } + Some(nn) +} + +#[inline] +unsafe fn slot_ptr( + header: *mut DsmMpscRingHeader, + idx: u32, + slot_capacity: u32, +) -> *mut SlotHeader { + let header_bytes = std::mem::size_of::(); + let base = header.cast::(); + unsafe { base.add(header_bytes + (idx as usize) * (slot_capacity as usize)) } + .cast::() +} + +#[inline] +unsafe fn slot_data_ptr(slot: *mut SlotHeader) -> *mut u8 { + unsafe { slot.cast::().add(SLOT_HEADER_BYTES) } +} + +impl DsmMpscReceiver { + /// Wrap an already-initialized ring as the single consumer. Pairs with + /// [`DsmMpscSender::new`] on the producer side; calling code is responsible for + /// keeping exactly one `DsmMpscReceiver` per ring. + /// + /// # Safety + /// `ring` must point to a header initialized by [`create_at`] and not yet + /// deallocated. The caller guarantees no other `DsmMpscReceiver` exists for the + /// same ring (single-consumer invariant). + pub(super) unsafe fn new(ring: NonNull, alive: AliveFlag) -> Self { + Self { ring, alive } + } + + /// Register the consumer's opaque wakeup token. Producers Acquire-load it as a single + /// `u64` (so they never see a torn token) and hand it to their [`Wakeup`] after publishing. + /// The token is the embedder's to interpret: a PG embedder packs `(pgprocno, pid)`; an + /// in-process embedder packs a registry key. Must not be [`NO_RECEIVER_TOKEN`]. + pub(super) fn set_receiver(&self, token: u64) { + if !self.alive.load(Ordering::Acquire) { + return; + } + let header = unsafe { self.ring.as_ref() }; + header.receiver_packed.store(token, Ordering::Release); + } + + /// Try to read one frame into `out`. `Bytes`: `out` holds the payload. `Empty`: + /// caller should yield and retry. `Detached`: ring drained, all producers gone, + /// no more frames coming. + /// + /// Known wedge: if a producer CAS-advances `tail` then exits/crashes before + /// publishing `seq`, the consumer sees `tail > head`, the slot's `seq` stuck at + /// the prior-round empty marker, and `detached && tail <= head` never becomes + /// true. The drain returns `Empty` forever. In production PG's parallel-worker + /// death handling bounds this (worker exit fires leader ERROR, DSM tears down), + /// but a future pass should add an explicit `PGPROC` liveness check to + /// force-detach on producer death. + pub(super) fn try_recv(&self, out: &mut Vec) -> RecvOutcome { + if !self.alive.load(Ordering::Acquire) { + // Segment detached: no DSM left to read, so report drained. + return RecvOutcome::Detached; + } + let header = unsafe { self.ring.as_ref() }; + let head = header.head.load(Ordering::Relaxed); + let slot_idx = (head % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + let expected_ready = head.wrapping_add(1); + if seq != expected_ready { + // Slot not ready. Use `<=` rather than `==` so a strict invariant + // violation (tail < head, impossible under correct operation) still + // surfaces as Detached rather than wedging Empty forever. + if header.detached.load(Ordering::Acquire) + && header.tail.load(Ordering::Acquire) <= head + { + return RecvOutcome::Detached; + } + return RecvOutcome::Empty; + } + // Slot at head is ready. Inspect its kind to choose between single-slot + // fast path and multi-slot reassembly. + let flags = unsafe { (*slot).flags.load(Ordering::Relaxed) }; + let Some(kind) = unpack_kind(flags) else { + // Reserved bits set; treat as corruption. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + }; + let payload_cap = (header.slot_capacity as usize).saturating_sub(SLOT_HEADER_BYTES); + match kind { + FragmentKind::Complete => self.recv_single_slot(header, head, slot, payload_cap, out), + FragmentKind::First => { + let n_slots = unpack_nslots(flags); + if n_slots == 0 || n_slots > header.ring_size { + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + self.recv_multi_slot(header, head, n_slots, payload_cap, out) + } + FragmentKind::Continue | FragmentKind::Last => { + // Encountering Continue/Last at `head` means a producer violated + // ascending-publish ordering or a previous reassembly didn't advance + // `head` past every fragment. Either is a contract break, not a + // recoverable condition; poison and detach. + header.detached.store(true, Ordering::Release); + RecvOutcome::Detached + } + } + } + + /// Read a single-slot `Complete` frame at `head` and advance. + fn recv_single_slot( + &self, + header: &DsmMpscRingHeader, + head: u64, + slot: *mut SlotHeader, + payload_cap: usize, + out: &mut Vec, + ) -> RecvOutcome { + let len_raw = unsafe { (*slot).len.load(Ordering::Relaxed) } as usize; + // Clamp against slot's payload capacity. DSM is mapped writable by every + // attached backend, so a buggy / corrupted producer could write a garbage len. + // Without this guard, `set_len + copy_nonoverlapping` would read OOB into + // neighboring slots or other DSM contents. + if len_raw > payload_cap { + // Poison the ring rather than silently returning corrupt data. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + let len = len_raw; + out.clear(); + out.reserve(len); + let data = unsafe { slot_data_ptr(slot) }; + // copy_nonoverlapping before set_len so a hypothetical panic mid-copy doesn't + // leave `out` with logical-len > initialized-bytes. + unsafe { + std::ptr::copy_nonoverlapping(data, out.as_mut_ptr(), len); + out.set_len(len); + } + // Mark the slot empty for the next round. Round k empty marker is + // (k * ring_size) + slot_idx; head + ring_size is exactly that for next round. + let next_empty_seq = head.wrapping_add(header.ring_size as u64); + unsafe { (*slot).seq.store(next_empty_seq, Ordering::Release) }; + // Advance head AFTER publishing the slot's empty marker, so a producer racing + // to claim sees the empty slot before seeing the new head value. + header.head.store(head.wrapping_add(1), Ordering::Release); + RecvOutcome::Bytes + } + + /// Reassemble an `n_slots`-fragment frame at `head`. Returns `Empty` without + /// advancing `head` if any continuation slot isn't yet published (producer + /// mid-publish); the caller retries and the `First` slot stays put with its + /// flags intact. + fn recv_multi_slot( + &self, + header: &DsmMpscRingHeader, + head: u64, + n_slots: u32, + payload_cap: usize, + out: &mut Vec, + ) -> RecvOutcome { + // First pass: verify every fragment in the run is published. If not, bail + // with Empty so the caller can retry once the producer finishes. + let mut total_len: usize = 0; + for i in 0..n_slots { + let h = head.wrapping_add(i as u64); + let slot_idx = (h % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + if seq != h.wrapping_add(1) { + // Continuation slot not yet ready. Don't touch `head` or any slot + // metadata; producer will publish soon, caller will retry. + return RecvOutcome::Empty; + } + let expected_kind = if i == 0 { + FragmentKind::First + } else if i + 1 == n_slots { + FragmentKind::Last + } else { + FragmentKind::Continue + }; + let flags = unsafe { (*slot).flags.load(Ordering::Relaxed) }; + if unpack_kind(flags) != Some(expected_kind) { + // Run integrity check: the producer's multi-slot publish never + // interleaves with another producer's, so the kind sequence must be + // First, Continue*, Last. Anything else is corruption. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + let len = unsafe { (*slot).len.load(Ordering::Relaxed) } as usize; + if len > payload_cap || (i + 1 < n_slots && len != payload_cap) { + // Non-final fragments must be full slots (producer fills them + // first); final fragment may be partial. Anything else is + // corruption. + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + total_len = match total_len.checked_add(len) { + Some(v) => v, + None => { + header.detached.store(true, Ordering::Release); + return RecvOutcome::Detached; + } + }; + } + // Second pass: concatenate payloads, mark each slot empty for next round, + // advance head past the whole run in one Release store. + out.clear(); + out.reserve(total_len); + for i in 0..n_slots { + let h = head.wrapping_add(i as u64); + let slot_idx = (h % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let len = unsafe { (*slot).len.load(Ordering::Relaxed) } as usize; + let data = unsafe { slot_data_ptr(slot) }; + unsafe { + let write_at = out.as_mut_ptr().add(out.len()); + std::ptr::copy_nonoverlapping(data, write_at, len); + out.set_len(out.len() + len); + } + // Round k empty marker for slot at position h is h + ring_size. + let next_empty_seq = h.wrapping_add(header.ring_size as u64); + unsafe { (*slot).seq.store(next_empty_seq, Ordering::Release) }; + } + header + .head + .store(head.wrapping_add(n_slots as u64), Ordering::Release); + RecvOutcome::Bytes + } +} + +impl DsmMpscSender { + /// Wrap an already-initialized ring as a producer. Multiple `DsmMpscSender` + /// handles to the same ring are legal (and the point of MPSC). Increments the + /// ring's `sender_count`; the `Drop` impl decrements and, on the last drop, + /// flips `detached` and wakes the receiver. That mirrors shm_mq's + /// "drop the sender, receiver sees detach" structural guarantee. + /// + /// # Safety + /// `ring` must point to a header initialized by [`create_at`] and not yet + /// deallocated. + pub(super) unsafe fn new( + ring: NonNull, + wakeup: Arc, + alive: AliveFlag, + ) -> Self { + let header = unsafe { ring.as_ref() }; + header.sender_count.fetch_add(1, Ordering::AcqRel); + Self { + ring, + wakeup, + counts_as_data: true, + alive, + } + } + + /// A control-plane sibling onto the same ring this handle already targets: it can publish + /// frames but stays out of `sender_count`, so it never sets or delays `detached`. Used to derive + /// a proc's `Cancel` senders from its data senders without a second attach, so a consumer can + /// reach its producer without masquerading as one of that producer's data senders. + pub(super) fn to_control(&self) -> DsmMpscSender { + DsmMpscSender { + ring: self.ring, + wakeup: Arc::clone(&self.wakeup), + counts_as_data: false, + alive: Arc::clone(&self.alive), + } + } + + /// Wake the registered consumer, if any. Reads the token the consumer stored via + /// [`DsmMpscReceiver::set_receiver`] and hands it to the injected [`Wakeup`]; skips when no + /// consumer is registered ([`NO_RECEIVER_TOKEN`]). + fn wake_receiver(&self) { + let header = unsafe { self.ring.as_ref() }; + let token = header.receiver_packed.load(Ordering::Acquire); + if token != NO_RECEIVER_TOKEN { + self.wakeup.wake(token); + } + } + + /// Push one frame onto the ring. Returns immediately: + /// - `Ok(())`: published; receiver's latch was set if installed. + /// - `Err(Full)`: no slot run available right now; caller yields + retries. + /// - `Err(Detached)`: receiver has detached; caller stops. + /// - `Err(MessageTooLarge)`: frame needs more than `ring_size` slots + /// (max writable size is `ring_size * (slot_capacity - SLOT_HEADER_BYTES)`). + /// + /// Frames up to the per-slot payload capacity take the single-slot fast path. + /// Larger frames span consecutive slots claimed atomically via one + /// `tail.compare_exchange(T, T + n_slots)`, so other producers can't + /// interleave their fragments. See [`FragmentKind`]. + pub(super) fn try_send(&self, bytes: &[u8]) -> Result<(), SendError> { + if !self.alive.load(Ordering::Acquire) { + return Err(SendError::Detached); + } + let header = unsafe { self.ring.as_ref() }; + // Control senders ignore `detached`: a `Cancel` still has to reach a producer whose own + // inbox already detached because its upstream finished while it's mid-output. + if self.counts_as_data && header.detached.load(Ordering::Acquire) { + return Err(SendError::Detached); + } + let payload_cap = (header.slot_capacity as usize).saturating_sub(SLOT_HEADER_BYTES); + if payload_cap == 0 { + return Err(SendError::MessageTooLarge); + } + // Single-slot fast path: cheaper than the multi-slot CAS dance and preserves + // the v1 hot path exactly so the no-fragmentation case takes no extra branches + // inside the claim loop. + if bytes.len() <= payload_cap { + return self.try_send_single_slot(header, bytes); + } + // Multi-slot path: fragment across `n_slots` consecutive slots. + let n_slots_usize = bytes.len().div_ceil(payload_cap); + if n_slots_usize > header.ring_size as usize || n_slots_usize > FLAGS_NSLOTS_MAX as usize { + // Frame is larger than the entire ring (or larger than the n_slots field + // can encode). Either bump `mpp_queue_size` or land a chunked-stream + // protocol that spans rounds. + return Err(SendError::MessageTooLarge); + } + let n_slots = n_slots_usize as u32; + self.try_send_multi_slot(header, bytes, n_slots, payload_cap) + } + + /// Single-slot path. Identical to the v1 layout's hot path with a `flags` write + /// added; the kind is always `Complete` here. + fn try_send_single_slot( + &self, + header: &DsmMpscRingHeader, + bytes: &[u8], + ) -> Result<(), SendError> { + loop { + // Acquire load on `tail` pairs defensively with any future blocking-send + // variant that may want to observe consumer progress via `head` (we don't + // today; full-detection rides on the slot's `seq`). Cheaper to keep the + // ordering tight than to relax-then-tighten under a future audit. + let tail = header.tail.load(Ordering::Acquire); + let slot_idx = (tail % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + // Three-way compare per Vyukov MPMC. + match seq.cmp(&tail) { + std::cmp::Ordering::Equal => { + // Slot is empty in our round. Try to claim by advancing tail. + // AcqRel on success so a subsequent producer's Acquire load of + // `tail` sees our claim and skips the slot we own. Relaxed on + // failure (we just retry the loop on a fresh tail load). + match header.tail.compare_exchange_weak( + tail, + tail.wrapping_add(1), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + // We own slot[slot_idx] for tail value `tail`. + unsafe { + (*slot).len.store(bytes.len() as u32, Ordering::Relaxed); + (*slot).flags.store( + pack_flags(FragmentKind::Complete, 0), + Ordering::Relaxed, + ); + let data = slot_data_ptr(slot); + std::ptr::copy_nonoverlapping(bytes.as_ptr(), data, bytes.len()); + // Publish: ready in round k is (k * ring_size + i + 1) = tail + 1. + (*slot).seq.store(tail.wrapping_add(1), Ordering::Release); + } + // Wake the consumer (resolves the latch via pgprocno + pid). + self.wake_receiver(); + return Ok(()); + } + Err(_) => continue, // another producer took this tail; retry + } + } + std::cmp::Ordering::Less => { + // seq < tail: the consumer hasn't reclaimed slot[slot_idx] for our + // round yet. Ring is full. + return Err(SendError::Full); + } + std::cmp::Ordering::Greater => { + // seq > tail: another producer has already claimed tail. Reload and + // retry. + continue; + } + } + } + } + + /// Multi-slot path. CAS-advance `tail` from `T` to `T + n_slots` to claim the run, + /// then publish each fragment ascending with `First` / `Continue` / `Last` flags. + /// The CAS requires all N target slots to already be in their expected empty + /// round, which generalizes v1's per-slot `seq == tail` check to a run. + /// + /// Ascending publish order isn't optional: the receiver only starts reassembling + /// once `slot[head]` (the `First`) is ready, then waits for each subsequent slot. + /// Out-of-order publish would make the receiver block on a `Continue` whose data + /// is already there but whose `seq` hasn't been stored yet, wasting one drain + /// pass per fragment. + // Fairness caveat: a multi-slot frame needs `n_slots` consecutive empty slots starting at + // `tail`, and competing single-slot producers can keep winning the CAS, so a large frame + // has no progress bound under sustained contention. The cooperative drain keeps retrying + // (no deadlock), but latency is unbounded; revisit with a reservation scheme if it shows + // up in send stats. + fn try_send_multi_slot( + &self, + header: &DsmMpscRingHeader, + bytes: &[u8], + n_slots: u32, + payload_cap: usize, + ) -> Result<(), SendError> { + loop { + let tail = header.tail.load(Ordering::Acquire); + // Verify every target slot is currently in its expected empty round. + // Slot at position (tail + i) % ring_size in round-of-(tail+i) has empty + // marker == (tail + i). If ANY is not empty, the ring is too contended / + // not drained enough for a run of this length. + let mut all_empty = true; + for i in 0..n_slots { + let t = tail.wrapping_add(i as u64); + let slot_idx = (t % header.ring_size as u64) as u32; + let slot = unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(Ordering::Acquire) }; + match seq.cmp(&t) { + std::cmp::Ordering::Equal => {} + std::cmp::Ordering::Less => { + // Slot at offset i hasn't been reclaimed by the consumer for + // round-of-t. Run not available right now. + return Err(SendError::Full); + } + std::cmp::Ordering::Greater => { + // Another producer already claimed (T + i); our view of `tail` + // is stale, retry. + all_empty = false; + break; + } + } + } + if !all_empty { + continue; + } + // All N target slots are empty in our round. Atomically claim the run. + match header.tail.compare_exchange_weak( + tail, + tail.wrapping_add(n_slots as u64), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + // We own slots tail..tail+n_slots. Write each fragment and + // publish its seq in ascending order so the receiver can drain + // them as soon as the prefix is ready. + let mut offset = 0usize; + for i in 0..n_slots { + let t = tail.wrapping_add(i as u64); + let slot_idx = (t % header.ring_size as u64) as u32; + let slot = + unsafe { slot_ptr(self.ring.as_ptr(), slot_idx, header.slot_capacity) }; + let remaining = bytes.len() - offset; + let chunk_len = remaining.min(payload_cap); + let kind = if i == 0 { + FragmentKind::First + } else if i + 1 == n_slots { + FragmentKind::Last + } else { + FragmentKind::Continue + }; + let flags_word = if kind == FragmentKind::First { + pack_flags(kind, n_slots) + } else { + pack_flags(kind, 0) + }; + unsafe { + (*slot).len.store(chunk_len as u32, Ordering::Relaxed); + (*slot).flags.store(flags_word, Ordering::Relaxed); + let data = slot_data_ptr(slot); + std::ptr::copy_nonoverlapping( + bytes.as_ptr().add(offset), + data, + chunk_len, + ); + // Publish: ready in round-of-t for slot at position t is t + 1. + (*slot).seq.store(t.wrapping_add(1), Ordering::Release); + } + offset += chunk_len; + } + debug_assert_eq!(offset, bytes.len(), "multi-slot send copied wrong length"); + // Wake the consumer once for the whole run; the drain pass that + // wakes for the First slot will keep going through all our + // already-published Continue/Last fragments without sleeping. + self.wake_receiver(); + return Ok(()); + } + Err(_) => continue, + } + } + } +} + +impl Drop for DsmMpscSender { + fn drop(&mut self) { + if !self.alive.load(Ordering::Acquire) { + // Segment detached: the `sender_count.fetch_sub` below would write into freed DSM. + return; + } + if !self.counts_as_data { + // Control senders never joined `sender_count`, so there's nothing to release and no + // detach to trigger. + return; + } + let header = unsafe { self.ring.as_ref() }; + // AcqRel: decrement is observed by other producers (they don't care, but the + // Release pairs with the receiver's Acquire load on `detached` below). + let prev = header.sender_count.fetch_sub(1, Ordering::AcqRel); + if prev == 1 { + // We were the last sender. Tell the consumer. + header.detached.store(true, Ordering::Release); + self.wake_receiver(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering as O}; + + /// No-op wakeup for the busy-poll tests, which spin on `try_recv` and never rely on being + /// woken. + struct NoopWakeup; + impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} + } + fn noop_wakeup() -> Arc { + Arc::new(NoopWakeup) + } + + /// In-process wakeup: unparks the consumer thread registered under a token. Proves the + /// injected extension point carries a real cross-thread notification with no PG / `SetLatch`. + #[derive(Default)] + struct ThreadWakeup { + threads: std::sync::Mutex>, + wakes: AtomicUsize, + } + impl ThreadWakeup { + fn register(&self, token: u64, thread: std::thread::Thread) { + self.threads.lock().unwrap().insert(token, thread); + } + } + impl Wakeup for ThreadWakeup { + fn wake(&self, token: u64) { + self.wakes.fetch_add(1, O::Relaxed); + if let Some(t) = self.threads.lock().unwrap().get(&token) { + t.unpark(); + } + } + } + + /// Send + Copy wrapper so tests can pass the ring pointer into spawned threads. The + /// real production handles (DsmMpscSender / Receiver) already impl Send; this exists + /// only because constructing them is unsafe and we want the test to do it inside the + /// thread so each thread has its own handle. + #[derive(Clone, Copy)] + struct SharedRing(NonNull); + unsafe impl Send for SharedRing {} + + /// Owning aligned region for a ring. The default `Vec` allocator can't + /// promise the alignment the `create_at` write wants, so allocate via the global + /// allocator with an explicit `Layout` and free in `Drop`. Production uses PG + /// `dsm_segment` (page-aligned), so this dance is test-only. + struct AlignedRegion { + ptr: *mut u8, + layout: std::alloc::Layout, + } + impl AlignedRegion { + fn new(bytes: usize) -> Self { + let align = std::mem::align_of::(); + let layout = std::alloc::Layout::from_size_align(bytes, align).expect("invalid layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null(), "allocator returned null"); + Self { ptr, layout } + } + fn as_mut_ptr(&self) -> *mut u8 { + self.ptr + } + } + impl Drop for AlignedRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } + } + + /// Allocate a heap region big enough for a ring with `ring_size` slots of + /// `slot_capacity` bytes each, initialize it, and return `(owner, receiver, + /// sender_template)`. The owner is returned first so callers bind it first. + /// Rust drops locals in reverse declaration order, so `_region` bound first drops + /// last, after the handles whose `Drop` impls touch the region's memory. Reverse + /// that order and `_region` frees the bytes before `tx_template`'s + /// `sender_count.fetch_sub` runs, which is undefined behavior and surfaces as a + /// stochastic crash at process teardown. + fn make_ring( + ring_size: u32, + slot_capacity: u32, + ) -> (AlignedRegion, DsmMpscReceiver, DsmMpscSender) { + let bytes = DsmMpscRingHeader::region_bytes(ring_size, slot_capacity); + let region = AlignedRegion::new(bytes); + let header_ptr = unsafe { create_at(region.as_mut_ptr(), ring_size, slot_capacity) }; + let nn = NonNull::new(header_ptr).expect("create_at returned null"); + // Unsafe: we hand ownership of the same pointer to two handles. Safe because the + // ring is the synchronization point; the handles are stateless wrappers. + let alive = Arc::new(AtomicBool::new(true)); + let receiver = unsafe { DsmMpscReceiver::new(nn, Arc::clone(&alive)) }; + let sender = unsafe { DsmMpscSender::new(nn, noop_wakeup(), alive) }; + (region, receiver, sender) + } + + #[test] + fn control_sender_does_not_gate_detach() { + let bytes = DsmMpscRingHeader::region_bytes(4, 64); + let region = AlignedRegion::new(bytes); + let nn = NonNull::new(unsafe { create_at(region.as_mut_ptr(), 4, 64) }).unwrap(); + let alive = Arc::new(AtomicBool::new(true)); + let rx = unsafe { DsmMpscReceiver::new(nn, Arc::clone(&alive)) }; + let data = unsafe { DsmMpscSender::new(nn, noop_wakeup(), alive) }; + let control = data.to_control(); + + let mut buf = Vec::new(); + // The control sender publishes without joining `sender_count`. + control.try_send(&[1, 2, 3]).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(&buf[..], &[1, 2, 3]); + // The data sender still holds the ring open. + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + + // Dropping the only data sender detaches the ring even though the control sender is held, + // so a downstream cancel never masks a producer-gone signal. + drop(data); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Detached); + + // The control sender still reaches the now-detached inbox: a `Cancel` has to land on a + // producer whose own upstream already finished. + control.try_send(&[9]).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(&buf[..], &[9]); + drop(control); + } + + #[test] + fn spsc_round_trip_under_capacity() { + let (_region, rx, tx) = make_ring(4, 64); + for i in 0..3u8 { + tx.try_send(&[i, i + 1, i + 2]).unwrap(); + } + let mut buf = Vec::new(); + for i in 0..3u8 { + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(&buf[..], &[i, i + 1, i + 2]); + } + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + #[test] + fn fills_then_full_then_drains() { + let (_region, rx, tx) = make_ring(4, 64); + for i in 0..4u32 { + tx.try_send(&i.to_le_bytes()).unwrap(); + } + // Fifth send must fail Full. + assert_eq!(tx.try_send(&999u32.to_le_bytes()), Err(SendError::Full)); + // Drain one, send one more, drain rest. + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, 0u32.to_le_bytes()); + tx.try_send(&100u32.to_le_bytes()).unwrap(); + for expected in [1u32, 2, 3, 100] { + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, expected.to_le_bytes()); + } + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + #[test] + fn message_too_large_is_rejected() { + // `MessageTooLarge` only fires when the frame would need more than `ring_size` + // slots. Per-slot oversize splits across consecutive slots instead. + let (_region, _rx, tx) = make_ring(2, 32); + let payload_cap = 32 - SLOT_HEADER_BYTES; + // payload_cap + 1 byte: now spans 2 slots, fits in ring_size=2. + let two_slot = vec![0u8; payload_cap + 1]; + tx.try_send(&two_slot).unwrap(); + // ring_size * payload_cap + 1 byte: needs 3 slots, ring only has 2. + let oversize = vec![0u8; 2 * payload_cap + 1]; + assert_eq!(tx.try_send(&oversize), Err(SendError::MessageTooLarge)); + // Exactly at the ring-wide cap is fine. (We just sent a 2-slot frame above, + // so we need to drain it first to free the slots; the receiver only exists + // in this test for that.) + } + + /// Multi-producer concurrent send: K threads each push M unique messages; consumer + /// receives K*M messages and verifies every (producer, sequence_in_producer) pair + /// appears exactly once. + #[test] + fn mpsc_no_lost_messages_under_contention() { + const K_PRODUCERS: usize = 8; + const M_PER_PRODUCER: u32 = 2000; + let (_region, rx, tx_template) = make_ring(64, 32); + // Wrap region pointer in something we can share across threads. The handles + // themselves are Send, so we clone via NonNull copy. + let ring_nn = SharedRing(tx_template.ring); + let consumed = Arc::new(AtomicUsize::new(0)); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + // Construct the sender on this thread so the closure captures DsmMpscSender + // (Send via unsafe impl) rather than the inner NonNull (Rust 2021 disjoint + // capture would otherwise project to the NonNull field and fail Send). + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + let h = std::thread::spawn(move || { + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + match tx.try_send(&payload) { + Ok(_) => sent += 1, + Err(SendError::Full) => std::thread::yield_now(), + Err(e) => panic!("unexpected send error: {e:?}"), + } + } + }); + handles.push(h); + } + // Drain on this thread until every producer's M messages have shown up. + let mut seen: Vec> = (0..K_PRODUCERS) + .map(|_| vec![false; M_PER_PRODUCER as usize]) + .collect(); + let mut buf = Vec::new(); + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + while consumed.load(O::Relaxed) < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => { + assert_eq!(buf.len(), 8); + let producer_id = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + let sent_idx = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize; + assert!(producer_id < K_PRODUCERS, "bad producer id {producer_id}"); + assert!( + sent_idx < M_PER_PRODUCER as usize, + "bad sent idx {sent_idx}" + ); + let already = std::mem::replace(&mut seen[producer_id][sent_idx], true); + assert!(!already, "duplicate ({producer_id}, {sent_idx})"); + consumed.fetch_add(1, O::Relaxed); + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach mid-drain"), + } + } + for h in handles { + h.join().unwrap(); + } + for (p, row) in seen.iter().enumerate() { + assert!( + row.iter().all(|&b| b), + "producer {p} has a missed message: row = {row:?}" + ); + } + } + + /// Per-producer in-order property: a single producer's messages observed by the + /// consumer must arrive in the order the producer sent them. (Cross-producer + /// ordering is not guaranteed.) + #[test] + fn mpsc_preserves_per_producer_order() { + const K_PRODUCERS: usize = 4; + const M_PER_PRODUCER: u32 = 500; + let (_region, rx, tx_template) = make_ring(32, 32); + let ring_nn = SharedRing(tx_template.ring); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + handles.push(std::thread::spawn(move || { + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + if tx.try_send(&payload).is_ok() { + sent += 1; + } else { + std::thread::yield_now(); + } + } + })); + } + let mut last_seq: Vec = vec![-1; K_PRODUCERS]; + let mut buf = Vec::new(); + let mut total = 0usize; + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + while total < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => { + let producer_id = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + let sent_idx = i64::from(u32::from_le_bytes(buf[4..8].try_into().unwrap())); + assert_eq!( + sent_idx, + last_seq[producer_id] + 1, + "out-of-order from producer {producer_id}: got {sent_idx}, expected {}", + last_seq[producer_id] + 1 + ); + last_seq[producer_id] = sent_idx; + total += 1; + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + } + + /// Cache-line layout regression: `head` must land on its own cache line + /// (offset == CACHE_LINE so the preceding header fields can't false-share), + /// `tail` separated from `head` by CACHE_LINE, and the first slot not sharing a + /// line with `tail`. Catches accidental padding removal or field reorder that + /// would re-introduce the false-sharing perf cliff the padding exists to avoid. + #[test] + fn header_layout_is_cache_friendly() { + use std::mem::offset_of; + let head_off = offset_of!(DsmMpscRingHeader, head); + let tail_off = offset_of!(DsmMpscRingHeader, tail); + let header_size = std::mem::size_of::(); + // head at exactly CACHE_LINE offset means the preceding 64 bytes (header + // fields + pad) live entirely in cache line 0, head + its pad live in + // cache line 1, etc. Anything other than equality means the padding + // math drifted; fix it rather than relaxing the assertion. + assert_eq!( + head_off, CACHE_LINE, + "head must land at offset {CACHE_LINE} (its own cache line); got {head_off}" + ); + assert!( + (tail_off - head_off) >= CACHE_LINE, + "head and tail must be on separate cache lines: head_off={head_off}, tail_off={tail_off}, CACHE_LINE={CACHE_LINE}" + ); + assert!( + (header_size - tail_off) >= CACHE_LINE, + "first slot must start on its own cache line: header_size={header_size}, tail_off={tail_off}, CACHE_LINE={CACHE_LINE}" + ); + // Header's natural alignment is determined by its largest field (AtomicU64, + // 8 bytes). Cache-line padding between hot fields still keeps them on separate + // absolute lines regardless of the struct's starting alignment, as long as the + // inter-field distance is >= CACHE_LINE. + assert_eq!(std::mem::align_of::(), 8); + } + + /// Magic + version validation: an attach against a region with the wrong magic or + /// version returns None rather than handing back a NonNull with garbage state. + /// Mirrors `MppDsmHeader::validate`'s discipline. + #[test] + fn attach_at_rejects_wrong_magic_and_version() { + let bytes = DsmMpscRingHeader::region_bytes(4, 64); + let region = AlignedRegion::new(bytes); + let header_ptr = unsafe { create_at(region.as_mut_ptr(), 4, 64) }; + // Sanity: correct attach succeeds. + assert!(unsafe { attach_at(region.as_mut_ptr(), 4, 64) }.is_some()); + // Corrupt magic: rejected. + unsafe { (*header_ptr).magic = 0xDEADBEEF }; + assert!(unsafe { attach_at(region.as_mut_ptr(), 4, 64) }.is_none()); + // Restore magic, corrupt version. + unsafe { (*header_ptr).magic = MPSC_RING_MAGIC }; + unsafe { (*header_ptr).version = MPSC_RING_VERSION.wrapping_add(1) }; + assert!(unsafe { attach_at(region.as_mut_ptr(), 4, 64) }.is_none()); + // Restore version, mismatched ring_size. + unsafe { (*header_ptr).version = MPSC_RING_VERSION }; + assert!(unsafe { attach_at(region.as_mut_ptr(), 8, 64) }.is_none()); + } + + /// Stress test at the production-worst contention level (K=24 producers, matching + /// the largest mesh-row size the transport drives in practice). Smoke that the + /// primitive doesn't wedge or lose messages under heavy CAS contention; doesn't + /// measure perf. + #[test] + fn mpsc_stress_at_production_worst_case() { + const K_PRODUCERS: usize = 24; + const M_PER_PRODUCER: u32 = 500; + let (_region, rx, tx_template) = make_ring(64, 32); + let ring_nn = SharedRing(tx_template.ring); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + handles.push(std::thread::spawn(move || { + let mut payload = [0u8; 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + match tx.try_send(&payload) { + Ok(_) => sent += 1, + Err(SendError::Full) => std::thread::yield_now(), + Err(e) => panic!("unexpected: {e:?}"), + } + } + })); + } + let mut buf = Vec::new(); + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + let mut total = 0usize; + while total < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => total += 1, + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + // Multi-round invariant: after draining 24 * 500 = 12000 messages on a 64-slot + // ring, every slot's seq must have advanced into a future round. Walk slot[0]: + // we expect seq = head + ring_size = 12000 + 64 - whatever round 0 it's in. + // Loose check: every slot's seq is bounded below by ring_size (round >= 1). + let header = unsafe { ring_nn.0.as_ref() }; + for i in 0..header.ring_size { + let slot = unsafe { slot_ptr(ring_nn.0.as_ptr(), i, header.slot_capacity) }; + let seq = unsafe { (*slot).seq.load(O::Acquire) }; + assert!( + seq >= header.ring_size as u64, + "slot[{i}].seq={seq} < ring_size; slot was never reused" + ); + } + } + + // ----- multi-slot fragmentation tests ----- + + /// Two-slot frame round-trip: send a payload that's exactly `2*payload_cap` + /// bytes long with distinct first-half and second-half markers, then verify the + /// receiver reassembles them in order with no truncation or duplication. + #[test] + fn multi_slot_two_fragment_round_trip() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx) = make_ring(4, slot_cap as u32); + let mut frame = vec![0u8; 2 * payload_cap]; + for (i, b) in frame.iter_mut().enumerate() { + *b = (i & 0xFF) as u8; + } + tx.try_send(&frame).unwrap(); + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf.len(), frame.len()); + assert_eq!(buf, frame); + // Ring fully drained. + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + /// Frame at the multi-slot ceiling: needs exactly `ring_size` slots, succeeds; + /// one byte more, rejected with MessageTooLarge. + #[test] + fn multi_slot_max_size_succeeds_one_more_rejected() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let ring_size = 4; + let (_region, rx, tx) = make_ring(ring_size, slot_cap as u32); + let max = vec![0xABu8; ring_size as usize * payload_cap]; + tx.try_send(&max).unwrap(); + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, max); + let too_big = vec![0u8; ring_size as usize * payload_cap + 1]; + assert_eq!(tx.try_send(&too_big), Err(SendError::MessageTooLarge)); + } + + /// Multi-slot followed by single-slot from the same producer: each frame must + /// come out of the receiver as a separate, intact byte sequence; fragments + /// don't leak into the following frame. + #[test] + fn multi_slot_then_single_slot_preserves_boundaries() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx) = make_ring(8, slot_cap as u32); + let big = vec![0xAAu8; 3 * payload_cap]; + tx.try_send(&big).unwrap(); + let small = b"hi".to_vec(); + tx.try_send(&small).unwrap(); + let mut buf = Vec::new(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, big); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, small); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + /// Multi-slot frame wrapping the ring boundary. With ring_size=4, tail starts at 0: + /// send a single-slot first to advance tail to 1, drain it, then send a 3-slot + /// frame at tail=1 occupying slots 1, 2, 3, drain the small frame to advance head, + /// then send another single-slot at tail=4 (slot 0 in round 1) and a 3-slot at + /// tail=5 (slots 1, 2, 3 in round 1). Verifies wraparound math holds. + #[test] + fn multi_slot_wraparound() { + let slot_cap = 64; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx) = make_ring(4, slot_cap as u32); + let mut buf = Vec::new(); + // Push a small frame so the next multi-slot reservation starts mid-ring. + tx.try_send(b"warm").unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, b"warm"); + // 3-slot frame at tail=1, wrapping at slot 3 -> slot 0 won't happen here + // (slots 1, 2, 3 are still in round 0). Drain it. + let big = vec![0x11u8; 3 * payload_cap]; + tx.try_send(&big).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, big); + // Now tail=4, head=4. Send another small frame at slot 0 (round 1, seq=4). + tx.try_send(b"two").unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, b"two"); + // And a 3-slot at tail=5 occupying slots 1, 2, 3 in round 1 (seqs 5, 6, 7). + let big2 = vec![0x22u8; 3 * payload_cap]; + tx.try_send(&big2).unwrap(); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Bytes); + assert_eq!(buf, big2); + assert_eq!(rx.try_recv(&mut buf), RecvOutcome::Empty); + } + + /// Stress: multiple producers each push a mix of single-slot and multi-slot + /// frames; consumer verifies every frame's contents and per-producer ordering. + /// Catches interleave bugs (a multi-slot run getting mixed with another + /// producer's fragments) and missing wake-ups under high contention. + #[test] + fn multi_slot_stress_mixed_with_singles() { + const K_PRODUCERS: usize = 6; + const M_PER_PRODUCER: u32 = 300; + let slot_cap = 128; + let payload_cap = slot_cap - SLOT_HEADER_BYTES; + let (_region, rx, tx_template) = make_ring(16, slot_cap as u32); + let ring_nn = SharedRing(tx_template.ring); + let mut handles = Vec::with_capacity(K_PRODUCERS); + for producer_id in 0..K_PRODUCERS { + let tx = unsafe { + DsmMpscSender::new(ring_nn.0, noop_wakeup(), Arc::clone(&tx_template.alive)) + }; + handles.push(std::thread::spawn(move || { + let mut sent = 0u32; + while sent < M_PER_PRODUCER { + // Alternate single, 2-slot, 4-slot, single, ... so the consumer + // sees every fragment-kind combination. + let n_payload = match sent % 3 { + 0 => 8, // single-slot + 1 => 2 * payload_cap, // 2-slot + _ => 4 * payload_cap, // 4-slot + }; + let mut payload = vec![0u8; n_payload + 8]; + payload[0..4].copy_from_slice(&(producer_id as u32).to_le_bytes()); + payload[4..8].copy_from_slice(&sent.to_le_bytes()); + for (i, b) in payload[8..].iter_mut().enumerate() { + *b = ((producer_id ^ sent as usize ^ i) & 0xFF) as u8; + } + match tx.try_send(&payload) { + Ok(_) => sent += 1, + Err(SendError::Full) => std::thread::yield_now(), + Err(e) => panic!("unexpected send error: {e:?}"), + } + } + })); + } + let mut last_seq: Vec = vec![-1; K_PRODUCERS]; + let mut buf = Vec::new(); + let mut total = 0usize; + let target = K_PRODUCERS * M_PER_PRODUCER as usize; + while total < target { + match rx.try_recv(&mut buf) { + RecvOutcome::Bytes => { + assert!(buf.len() >= 8, "frame too short: {}", buf.len()); + let producer_id = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + let sent_idx = i64::from(u32::from_le_bytes(buf[4..8].try_into().unwrap())); + assert!(producer_id < K_PRODUCERS, "bad producer id"); + assert_eq!( + sent_idx, + last_seq[producer_id] + 1, + "out-of-order frame from producer {producer_id}" + ); + // Verify payload bytes (catches fragment interleave / partial reads). + for (i, b) in buf[8..].iter().enumerate() { + let expected = ((producer_id ^ sent_idx as usize ^ i) & 0xFF) as u8; + assert_eq!(*b, expected, "payload mismatch at byte {i}"); + } + last_seq[producer_id] = sent_idx; + total += 1; + } + RecvOutcome::Empty => std::thread::yield_now(), + RecvOutcome::Detached => panic!("unexpected detach"), + } + } + for h in handles { + h.join().unwrap(); + } + } + + /// The injected `Wakeup` carries a real cross-thread notification with no PG: a consumer that + /// parks until woken is unparked by the producer's publish, and the producer actually invoked + /// the wakeup (asserted via the counter, so a silently-broken extension point fails rather than relying + /// on park timing). + #[test] + fn injected_wakeup_unparks_blocked_consumer() { + let region = AlignedRegion::new(DsmMpscRingHeader::region_bytes(8, 256)); + let nn = NonNull::new(unsafe { create_at(region.as_mut_ptr(), 8, 256) }) + .expect("create_at returned null"); + let ring = SharedRing(nn); + let wakeup = Arc::new(ThreadWakeup::default()); + let alive = Arc::new(AtomicBool::new(true)); + const TOKEN: u64 = 7; + + let receiver = unsafe { DsmMpscReceiver::new(nn, Arc::clone(&alive)) }; + // Register this thread as the wake target before publishing the token, so a producer that + // races ahead still finds us. + wakeup.register(TOKEN, std::thread::current()); + receiver.set_receiver(TOKEN); + + let producer = { + let wakeup = Arc::clone(&wakeup) as Arc; + let alive = Arc::clone(&alive); + std::thread::spawn(move || { + // Rebind the whole `Send` wrapper so the closure captures it (not the bare + // `NonNull` field, which edition-2024 disjoint capture would otherwise grab). + let ring = ring; + let tx = unsafe { DsmMpscSender::new(ring.0, wakeup, alive) }; + std::thread::sleep(std::time::Duration::from_millis(20)); + assert_eq!(tx.try_send(b"hello"), Ok(())); + }) + }; + + // park_timeout (not park) so a broken extension point fails via the counter assertion below instead + // of hanging the suite. park/unpark holds a token, so an unpark that beats our park still + // wakes the next park (no lost-wakeup). + let mut out = Vec::new(); + loop { + match receiver.try_recv(&mut out) { + RecvOutcome::Bytes => break, + RecvOutcome::Empty => std::thread::park_timeout(std::time::Duration::from_secs(5)), + RecvOutcome::Detached => panic!("detached before receiving the frame"), + } + } + producer.join().unwrap(); + assert_eq!(out, b"hello"); + assert!( + wakeup.wakes.load(O::Relaxed) >= 1, + "producer never invoked the injected wakeup" + ); + } +} diff --git a/src/shm/runtime.rs b/src/shm/runtime.rs new file mode 100644 index 00000000..5f538d83 --- /dev/null +++ b/src/shm/runtime.rs @@ -0,0 +1,477 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Runtime glue between the leader's DataFusion execution and the DSM MPSC mesh. +//! +//! [`MppMesh`] is the runtime handle the leader builds at DSM-init time. It holds the +//! single [`super::transport::DrainHandle`] (the +//! `inbound_receiver`) that consolidates this proc's DSM inbox and self-loop, and gets +//! installed on the leader's `SessionConfig` extensions before plan execution. +//! +//! [`ShmChannelResolver`] implements [`ChannelResolver`], consulted by +//! `NetworkShuffleExec` / `NetworkCoalesceExec` / `NetworkBroadcastExec` at execute time. It hands +//! out a [`ShmWorkerChannel`] whose `execute_task(task_key, partition_range)` yields one stream per +//! consumer partition from the shared `inbound_receiver`. +//! +//! [`InProcessWorkerResolver`] hands the planner `n_workers` placeholder URLs. The transport +//! routes by task index, not URL, so the URLs are never dialed; the resolver exists only because +//! the planner sizes stages from the URL count. It replaces the placeholder URL the fork used to +//! substitute internally under the old `in_process_mode` flag. + +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; +use futures::stream::{self, BoxStream, StreamExt}; +use http::HeaderMap; +use url::Url; + +use crate::proto as pb; +use crate::work_unit_feed::RemoteWorkUnitFeedTxs; +use crate::{ + ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, + GetWorkerInfoResponse, WorkerChannel, WorkerResolver, WorkerToCoordinatorMsg, +}; + +use super::AliveFlag; +use super::transport::{CooperativeDrainSet, DrainHandle, DrainItem, Interrupt, MppSender}; + +/// A proc's outbound senders to each peer inbox, shared between the mesh (for `Cancel` frames) and +/// the embedder that owns their lifetime. Indexed by destination `proc_idx`; this proc's own slot +/// is `None`. +pub type PeerSenders = Arc>>>; + +/// `task_idx → proc_idx` round-robin over the worker procs. The leader is `proc_idx = 0` +/// (consumer-only), workers are `1..n_procs` (each hosts producer fragments). +/// +/// A stage's task count is set by the DF-D task estimator chain, not by the worker proc count. +/// The transport does NOT support more tasks than producer procs: channels are keyed +/// `(sender_proc, stage, partition)`, so two tasks wrapped onto one proc would interleave on one +/// channel and the first EOF would truncate the other task's output. `ShmWorkerChannel::execute_task` +/// rejects that shape; the modulo here only keeps the function total. +#[inline] +pub fn proc_for_task(n_workers: u32, task_idx: u32) -> u32 { + 1 + (task_idx % n_workers.max(1)) +} + +/// Runtime handle the customscan populates at DSM-init time. +/// +/// Each process owns one MPSC inbox in DSM that receives frames from every peer. +/// `inbound_receiver` consolidates that inbox plus the in-proc self-loop channel (for +/// producer-and-consumer-on-same-worker fragments) into a single `DrainHandle`. Frames +/// carry `(sender_proc, stage_id, partition)` in their header so the routing registry +/// inside the handle delivers each frame to the matching consumer. +/// +/// [`MppFrameHeader`]: super::transport::MppFrameHeader +pub struct MppMesh { + /// This process's `proc_idx` (= 0 for the leader, `ParallelWorkerNumber + 1` for workers). + /// Frames addressed to this proc arrive on this proc's own inbox. + pub this_proc: u32, + /// Total proc count. Bounds the producer/consumer proc lookups in + /// [`ShmWorkerChannel::execute_task`]. + pub n_procs: u32, + /// Single cooperative inbound handle pulling every frame addressed to this proc. The + /// DSM MPSC inbox and an in-proc self-loop receiver both feed into this handle. Demux + /// to per-`(sender_proc, stage_id, partition)` channel buffers happens inside via + /// `DrainHandle::register_channel`. + pub(super) inbound_receiver: Arc, + /// Cancellation hook, injected by the embedder, checked at the transport's block points (the + /// cooperative send spin and the consumer pull loop). + interrupt: Arc, + /// This proc's outbound senders to each peer inbox, shared with the embedder that owns their + /// lifetime (it clears the `Vec` before the DSM unmaps). Used by [`Self::cancel_stream`] to ship + /// a `Cancel` frame to a producing proc when this proc abandons a stream. `None` until the + /// embedder installs it. + cancel_senders: Mutex>, + alive: AliveFlag, +} + +impl MppMesh { + /// Build a fresh mesh. + pub fn new( + this_proc: u32, + n_procs: u32, + inbound_receiver: Arc, + interrupt: Arc, + alive: AliveFlag, + ) -> Self { + Self { + this_proc, + n_procs, + inbound_receiver, + interrupt, + cancel_senders: Mutex::new(None), + alive, + } + } + + /// Mark this proc's DSM mesh detached so every ring handle's operations become no-ops. + /// The embedder calls this from its dsm-detach callback, while the segment is still mapped, + /// so handles dropped afterward (e.g. by a memory-context reset) never touch freed memory. + pub fn mark_detached(&self) { + self.alive.store(false, Ordering::Release); + } + + /// The raw liveness flag shared with every ring handle, for embedders that register a C + /// dsm-detach callback against it directly. + pub fn detached_flag(&self) -> AliveFlag { + Arc::clone(&self.alive) + } + + /// Share this proc's outbound senders so [`Self::cancel_stream`] can reach the producers. The + /// embedder keeps owning their lifetime: it passes a clone of the same `Arc` it releases before + /// the DSM unmaps, so the mesh never extends the senders past teardown. + pub fn set_cancel_senders(&self, senders: PeerSenders) { + *self.cancel_senders.lock().unwrap() = Some(senders); + } + + /// Tell the producer on `producer_proc` to stop the `(stage_id, partition)` stream: this proc's + /// consumer of it stopped reading before EOF. Ships one `Cancel` frame, leaving the rings + /// healthy for metrics and every other stream. A no-op when no senders are installed (the + /// embedder hasn't wired this proc, or teardown cleared them). + /// + /// Stream-level so any consumer can cancel its own input: every `(producer_proc, stage, + /// partition)` channel has a single consumer, so one stream's drop never cuts off a sibling's. + pub fn cancel_stream(&self, producer_proc: u32, stage_id: u32, partition: u32) { + let guard = self.cancel_senders.lock().unwrap(); + let Some(senders) = guard.as_ref() else { + return; + }; + let senders = senders.lock().unwrap(); + if let Some(Some(sender)) = senders.get(producer_proc as usize) { + sender.try_send_cancel(stage_id, partition); + } + } + + /// The single cooperative inbound handle that pulls frames from every peer (and the + /// self-loop) into per-`(sender_proc, stage_id, partition)` channel buffers. + pub(super) fn inbound_receiver(&self) -> &Arc { + &self.inbound_receiver + } + + /// Install the senders of one task's work-unit feed channels on this proc's drain, so + /// inbound `WorkUnit` frames for `(stage_id, task_number)` flow into them. Units that + /// arrived first are flushed; a `FeedEof` that already came through closes the channels + /// immediately. The drain only fills channels: something on this proc must keep draining + /// (a consumer pull loop, a producer send spin, or an explicit + /// [`CooperativeDrainSet::try_drain_pass`] pump) or a fragment blocked on its feed starves. + pub fn register_work_unit_senders( + &self, + stage_id: u32, + task_number: u32, + senders: RemoteWorkUnitFeedTxs, + ) { + self.inbound_receiver + .register_work_unit_senders(stage_id, task_number, senders); + } + + /// Take the plan delivered for `(stage_id, task_number)` as a `SetPlan` frame on this proc's + /// inbox, waiting for it if it has not arrived yet. Something on this proc must keep + /// draining (a pump, a consumer pull loop, or a cooperative send spin) or the wait starves. + pub async fn take_set_plan( + &self, + stage_id: u32, + task_number: u32, + ) -> Result { + self.inbound_receiver + .take_set_plan(stage_id, task_number) + .await + } + + /// Take the stream of `TaskMetrics` frames arriving on this proc's inbox: + /// `(stage_id, task_number, metrics)` per producer task that reported in. Meant for the + /// leader; the first caller gets it, later calls get `None`. + pub fn take_task_metrics_receiver( + &self, + ) -> Option> { + self.inbound_receiver.take_task_metrics_receiver() + } + + /// Number of worker procs (= `n_procs - 1`, since the leader is proc 0). Used as the + /// modulus in [`proc_for_task`]. + /// + /// The embedder must guarantee `n_procs >= 2` (one consumer-only leader plus at least one + /// producer worker) before constructing an `MppMesh`; the subtraction is otherwise unsound. + /// `compute_dsm_layout` enforces the same bound. Asserted in debug builds so a future misuse + /// fails loudly. + pub fn n_workers(&self) -> u32 { + debug_assert!( + self.n_procs >= 2, + "MppMesh::n_workers() called with n_procs={} (< 2); the embedder must gate on \ + n_procs >= 2", + self.n_procs + ); + self.n_procs - 1 + } + + /// Pull from the single inbound handle. Called from + /// [`super::transport::MppSender`]'s cooperative-send spin so a + /// producer stalled on a full outbound can pull inbound peer data inline. That's what + /// prevents the symmetric-send deadlock when every peer is simultaneously stalled waiting + /// for space. + pub(super) fn drain_all_inbound(&self) -> Result<(), DataFusionError> { + self.inbound_receiver.try_drain_pass() + } +} + +impl CooperativeDrainSet for MppMesh { + fn try_drain_pass(&self) -> Result<(), DataFusionError> { + self.drain_all_inbound() + } + + fn check_interrupt(&self) -> Result<(), DataFusionError> { + self.interrupt.check() + } + + fn stream_cancelled(&self, stage_id: u32, partition: u32) -> bool { + self.inbound_receiver.stream_cancelled(stage_id, partition) + } +} + +/// Hands out [`ShmWorkerChannel`]s over the leader's [`MppMesh`]. The shm mesh has no URLs to dial, +/// so every `get_worker_client_for_url` resolves to the same mesh; the channel routes by the +/// `task_key` it is handed at `execute_task`, not by URL. +pub struct ShmChannelResolver { + mesh: Arc, +} + +impl ShmChannelResolver { + pub fn new(mesh: Arc) -> Self { + Self { mesh } + } +} + +#[async_trait] +impl ChannelResolver for ShmChannelResolver { + async fn get_worker_client_for_url( + &self, + _url: &Url, + ) -> Result, DataFusionError> { + Ok(Box::new(ShmWorkerChannel { + mesh: Arc::clone(&self.mesh), + })) + } +} + +/// A [`WorkerChannel`] over the mesh. `execute_task((stage, task), partition_range)` translates the +/// DF-D `(stage, task)` addressing into the proc-pair grid: `proc_for_task(n_workers, task_number)` +/// selects which `sender_proc` hosts the producer-side task, and each returned stream pulls that +/// proc's `(sender_proc, stage_id, partition)` slice from the inbound drain. +struct ShmWorkerChannel { + mesh: Arc, +} + +#[async_trait] +impl WorkerChannel for ShmWorkerChannel { + /// pg_search delivers each worker's plan over DSM, not over this channel, so plan delivery is a + /// no-op (what the old `NoOpDispatch` did). Drain the inbound control stream to exhaustion so the + /// coordinator's keep-alive tail does not wedge, then complete with an empty output stream: task + /// metrics travel back over the mesh, not here. + async fn coordinator_channel( + &mut self, + _headers: HeaderMap, + mut c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, + ) -> Result>> { + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { while c2w_stream.next().await.is_some() {} }); + Ok(stream::empty().boxed()) + } + + async fn execute_task( + &mut self, + _headers: HeaderMap, + request: ExecuteTaskRequest, + metrics: ExecutionPlanMetricsSet, + _task_ctx: &Arc, + ) -> Result>>> { + MetricBuilder::new(&metrics) + .global_counter("mesh_connections_used") + .add(1); + let stage_id = u32::try_from(request.task_key.stage_id).map_err(|_| { + DataFusionError::Internal(format!( + "ShmWorkerChannel: stage_id={} > u32::MAX", + request.task_key.stage_id + )) + })?; + let task_number = u32::try_from(request.task_key.task_number).map_err(|_| { + DataFusionError::Internal(format!( + "ShmWorkerChannel: task_number={} > u32::MAX", + request.task_key.task_number + )) + })?; + let sender_proc = proc_for_task(self.mesh.n_workers(), task_number); + // More tasks than producer procs would fold two tasks onto one + // `(sender_proc, stage, partition)` channel: interleaved batches, and the first task's EOF + // truncates the second. With no input_stage here to count tasks, the equivalent guard is + // that the routed proc stays in range, as the old code also checked. + if sender_proc >= self.mesh.n_procs { + return Err(DataFusionError::Internal(format!( + "ShmWorkerChannel: sender_proc={sender_proc} >= n_procs={} \ + (stage_id={stage_id}, task_number={task_number})", + self.mesh.n_procs + ))); + } + // First line to grep when a query hangs on a channel nothing registered. + log::debug!( + "shm transport execute_task: this_proc={} stage_id={stage_id} \ + task_number={task_number} -> sender_proc={sender_proc}", + self.mesh.this_proc + ); + let mut streams = Vec::with_capacity( + request + .target_partition_end + .saturating_sub(request.target_partition_start), + ); + for partition in request.target_partition_start..request.target_partition_end { + let partition_u32 = u32::try_from(partition).map_err(|_| { + DataFusionError::Internal(format!( + "ShmWorkerChannel: partition={partition} > u32::MAX" + )) + })?; + streams.push(pull_partition_stream( + Arc::clone(&self.mesh), + sender_proc, + stage_id, + partition_u32, + )); + } + Ok(streams) + } + + async fn get_worker_info( + &mut self, + _request: GetWorkerInfoRequest, + ) -> Result { + Ok(GetWorkerInfoResponse { + version: String::new(), + }) + } +} + +/// Build the cooperative pull-loop stream for one `(producer_proc, stage_id, partition)` channel. +/// +/// The inbound drain runs inline on the consumer's thread. Each iteration checks for cancellation +/// (via the injected interrupt extension point), drains the receiver into the registry, pops one +/// batch to yield, then yields back to Tokio so sibling tasks (e.g. the leader's own producer +/// subplan) advance. The send side runs the same interrupt check inside `MppSender`'s retry spin. +fn pull_partition_stream( + mesh: Arc, + producer_proc: u32, + stage_id: u32, + partition: u32, +) -> BoxStream<'static, Result> { + // One drain per process, shared across all sender_procs. The channel-buffer registry keys by + // (sender_proc, stage_id, partition) so this consumer still sees only its named sender's slice + // even though the underlying inbox is shared with all peers. + let drain = Arc::clone(mesh.inbound_receiver()); + log::debug!( + "shm transport execute: register channel sender_proc={producer_proc} stage_id={stage_id} \ + partition={partition}" + ); + let buffer = drain.register_channel(producer_proc, stage_id, partition); + let stream = async_stream::stream! { + // Any consumer cancels its own input stream when it drops early. The mesh no-ops the send + // until the embedder wires this proc's outbound senders. + let mut cancel_guard = StreamCancelGuard { + mesh: Arc::clone(&mesh), + producer_proc, + stage_id, + partition, + armed: true, + }; + loop { + if let Err(e) = mesh.check_interrupt() { + yield Err(e); + return; + } + if let Err(e) = drain.try_drain_pass() { + yield Err(e); + return; + } + match buffer.try_pop() { + Some(DrainItem::Batch(batch)) => yield Ok(batch), + Some(DrainItem::Eof) => { + // Clean end: the producer already EOF'd, so there's nothing to cancel. + cancel_guard.armed = false; + break; + } + Some(DrainItem::Failed(msg)) => { + yield Err(DataFusionError::Execution(msg)); + return; + } + None => tokio::task::yield_now().await, + } + } + }; + Box::pin(stream) +} + +/// Placeholder worker resolver for the in-process MPP transport. +/// +/// The shm_mq transport routes by `target_task` (proc index), never by URL, so these URLs are +/// never dialed. The distributed planner still needs a resolver: it sizes stages and assigns +/// tasks from the URL count. `n_workers` placeholder URLs is exactly what the planner needs. This +/// replaces the placeholder URL the fork used to substitute internally under `in_process_mode`. +pub struct InProcessWorkerResolver { + n_workers: usize, +} + +impl InProcessWorkerResolver { + pub fn new(n_workers: usize) -> Self { + Self { n_workers } + } +} + +impl WorkerResolver for InProcessWorkerResolver { + fn get_urls(&self) -> Result, DataFusionError> { + (0..self.n_workers.max(1)) + .map(|i| { + Url::parse(&format!("inprocess://worker/{i}")).map_err(|e| { + DataFusionError::Internal(format!( + "InProcessWorkerResolver: invalid placeholder url: {e}" + )) + }) + }) + .collect() + } +} + +/// Cancels one `(stage_id, partition)` stream if its consumer drops before EOF, telling the +/// producer on `producer_proc` to stop. A consumer that stops pulling early (a top-N `LIMIT`, an +/// inner merge join exhausting a side, etc.) would otherwise leave that producer spinning on the +/// full inbox until the statement timeout. Disarmed on a clean EOF: there the producer already +/// finished, so there's nothing to cancel. +struct StreamCancelGuard { + mesh: Arc, + producer_proc: u32, + stage_id: u32, + partition: u32, + armed: bool, +} + +impl Drop for StreamCancelGuard { + fn drop(&mut self) { + if self.armed { + self.mesh + .cancel_stream(self.producer_proc, self.stage_id, self.partition); + } + } +} diff --git a/src/shm/self_hosted.rs b/src/shm/self_hosted.rs new file mode 100644 index 00000000..192b77a9 --- /dev/null +++ b/src/shm/self_hosted.rs @@ -0,0 +1,1277 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Deferred: `mod self_hosted` is gated out (`#[cfg(any())]`) in `shm/mod.rs`. This file still +//! targets the removed `WorkerTransport`/`WorkerDispatch` dispatch umbrella, which the +//! `ChannelResolver`/`WorkerChannel` protocol has no analog for. Its no-gRPC-default role is now +//! served by `InProcessChannelResolver`, and the in-crate ring safety net by the `in_process` test, +//! so it stays unported until reimplemented as a `ChannelResolver` driving produce over +//! `coordinator_channel`. Kept for reference; it does not compile against this branch. +//! +//! The shared-memory transport hosting its own workers, so it can serve as a session default. +//! +//! Production embedders drive the mesh themselves: they allocate the region, launch the worker +//! processes, and deliver plans out of band, so [`ShmMqWorkerTransport`]'s dispatcher is a no-op. +//! That shape cannot be a default transport: a default gets nothing but the `WorkerTransport` +//! calls. [`SelfHostedShmTransport`] closes that gap by playing the embedder itself, in-process: +//! dispatch delivers each task's plan through the worker session machinery +//! ([`Worker::set_task_plan`]: codec round-trip, work-unit feed channels, metrics back-channel) +//! and runs the producer fragments as tasks pushing through a heap-backed DSM mesh; `open` reads +//! the rings from the leader side. Every cross-stage byte moves through the same rings, framing, +//! and cooperative drain a production embedder uses. +//! +//! Per query, the harness lives from the first dispatch to the cancellation token firing (the +//! head stream dropping). The mesh is sized and built lazily at the first `open`, once every +//! stage has been dispatched and the routing is known. + +use std::alloc::Layout; +use std::ffi::c_void; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; + +use dashmap::DashMap; +use datafusion::common::instant::Instant; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, HashMap, Result, exec_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::ExecutionPlan; +use futures::StreamExt; +use futures::stream::BoxStream; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use super::mpsc_ring::Wakeup; +use super::runtime::{MppMesh, ShmMqWorkerTransport, proc_for_task}; +use super::setup::{dsm_region_bytes, leader_setup, worker_setup}; +use super::transport::{ + CooperativeDrainSet, Interrupt, MppFrameHeader, MppPartitionSink, MppSender, SendBatchStats, + SetPlanFrame, +}; +use crate::proto as pb; +use crate::{ + CoordinatorToWorkerMetrics, DistributedTaskContext, MetricsStore, NetworkBoundaryExt, + NetworkCoalesceExec, PartitionSink, ProducerHead, RemoteStage, TreeNodeExt, Worker, + WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerSessionBuilder, WorkerTransport, + collect_task_work_unit_feeds, encode_task_plan, execute_local_task, + get_config_extension_propagation_headers, get_distributed_cancellation_token, + get_passthrough_headers, serialize_uuid, set_distributed_worker_transport, set_sent_time, +}; + +/// Per-inbox ring size. Frames fragment across slots, so this bounds a single batch at roughly +/// the whole ring; the suite's batches sit well under it while keeping the per-query footprint +/// (`n_procs` inboxes) modest. +const SELF_HOSTED_QUEUE_BYTES: usize = 1 << 22; + +/// No-op wakeup: the cooperative consumers yield instead of parking, so a publish never has to +/// wake a blocked thread. +pub(super) struct NoopWakeup; +impl Wakeup for NoopWakeup { + fn wake(&self, _token: u64) {} +} + +/// Opaque, non-sentinel receiver token. The wakeup ignores the value; this just keeps the +/// producer from treating the consumer as unregistered. +pub(super) fn receiver_token(proc_idx: u32) -> u64 { + proc_idx as u64 + 1 +} + +/// Owns the heap buffer standing in for a shared-memory segment. Every proc reads and writes the +/// same region through raw pointers; the lock-free rings make the concurrent access sound. +pub(super) struct HeapRegion { + ptr: *mut u8, + layout: Layout, +} + +impl HeapRegion { + pub(super) fn new(bytes: usize) -> Self { + // 64-byte alignment so each per-inbox ring header lands on its own cache line; the + // dsm layout aligns the offsets within the region, but only if the base is aligned too. + let layout = Layout::from_size_align(bytes, 64).expect("dsm region layout"); + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + assert!(!ptr.is_null(), "dsm region alloc failed"); + Self { ptr, layout } + } + + pub(super) fn base(&self) -> *mut c_void { + self.ptr as *mut c_void + } +} + +impl Drop for HeapRegion { + fn drop(&mut self) { + unsafe { std::alloc::dealloc(self.ptr, self.layout) }; + } +} + +// The raw pointer is only dereferenced through the rings, which are designed for concurrent +// access from multiple procs. +unsafe impl Send for HeapRegion {} +unsafe impl Sync for HeapRegion {} + +/// Interrupt extension point wired to the query's cancellation token, so producers blocked in a send spin +/// and consumers in the pull loop both unwind when the head stream drops. +struct CancellationInterrupt(CancellationToken); +impl Interrupt for CancellationInterrupt { + fn check(&self) -> Result<(), DataFusionError> { + if self.0.is_cancelled() { + return Err(DataFusionError::Execution( + "mpp: query cancelled".to_string(), + )); + } + Ok(()) + } +} + +/// [`WorkerTransport`] over the shared-memory mesh, hosting its own workers in-process. +/// +/// All tasks share one [Worker] (one task registry, one session builder), like the in-memory +/// transport; what differs is the data plane: producer fragments run eagerly as background tasks +/// pushing through DSM rings, and reads pull from those rings instead of executing lazily. +/// +/// With the `flight` feature off this is the default transport. Multi-process embedders keep +/// driving [`ShmMqWorkerTransport`] directly. +#[derive(Clone)] +pub struct SelfHostedShmTransport { + worker: Worker, + queries: Arc>>, + queue_bytes: usize, +} + +impl Default for SelfHostedShmTransport { + fn default() -> Self { + Self::new(Worker::default()) + } +} + +impl SelfHostedShmTransport { + /// Builds the transport around an existing [Worker], sharing its task registry, session + /// builder, and runtime environment. + pub fn new(worker: Worker) -> Self { + Self { + worker, + queries: Arc::new(DashMap::new()), + queue_bytes: SELF_HOSTED_QUEUE_BYTES, + } + } + + /// Overrides the per-inbox ring size. Small values force multi-slot fragmentation and the + /// cooperative send spin on every query, which is how the ring mechanics get stress-tested; + /// the default is generous enough that only large batches touch them. + pub fn with_queue_bytes(mut self, queue_bytes: usize) -> Self { + self.queue_bytes = queue_bytes; + self + } + + /// Builds the transport with a custom [WorkerSessionBuilder], the same customization hook a + /// remote worker offers. + pub fn from_session_builder( + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + Self::new(Worker::from_session_builder(session_builder)) + } + + /// The in-process [Worker] backing this transport. + pub fn worker(&self) -> &Worker { + &self.worker + } +} + +/// Recovers the producer head the planner already inserted at the top of a stage's shipped +/// fragment. The worker's `insert_producer_head` strips an existing head and re-adds one from the +/// [`ProducerHead`] it is given, so a `None` here would flatten a sliced stage to one partition; +/// echoing the existing head back keeps it intact. +fn producer_head_from_plan(plan: &Arc) -> ProducerHead { + use datafusion::physical_plan::repartition::RepartitionExec; + if let Some(r) = plan.downcast_ref::() { + ProducerHead::RepartitionExec { + partitioning: r.partitioning().clone(), + } + } else if let Some(b) = plan.downcast_ref::() { + ProducerHead::BroadcastExec { + output_partitions: b.properties().partitioning.partition_count(), + } + } else { + ProducerHead::None + } +} + +impl WorkerTransport for SelfHostedShmTransport { + fn open( + &self, + input_stage: &RemoteStage, + target_partitions: std::ops::Range, + target_task: usize, + producer_head: ProducerHead, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + let Some(harness) = self + .queries + .get(&input_stage.query_id) + .map(|e| Arc::clone(e.value())) + else { + return exec_err!( + "self-hosted shm transport: no harness for query {}; stage {} was never \ + dispatched through this transport", + input_stage.query_id, + input_stage.num + ); + }; + // The first read finalizes the harness: by now every stage has been dispatched (plan + // preparation completes before the head executes), so the mesh can be sized and the + // producer drivers released. + harness.ensure_started()?; + let leader_mesh = harness.leader_mesh()?; + let inner = ShmMqWorkerTransport::new(leader_mesh).open( + input_stage, + target_partitions, + target_task, + producer_head, + ctx, + metrics, + )?; + Ok(Box::new(PinnedConnection { inner, harness })) + } + + fn dispatcher(&self) -> Box { + Box::new(SelfHostedDispatcher { + transport: self.clone(), + metrics: OnceLock::new(), + }) + } +} + +/// Keeps the query harness (and through it the heap region the rings live in) alive for as long +/// as any stream is still reading from the mesh. +struct PinnedConnection { + inner: Box, + harness: Arc, +} + +impl WorkerConnection for PinnedConnection { + fn execute( + &self, + partition: usize, + ) -> Result>> { + let stream = self.inner.execute(partition)?; + let harness = Arc::clone(&self.harness); + Ok(Box::pin(stream.map(move |item| { + let _ = &harness; // <- the region must outlive the ring receivers. + item + }))) + } +} + +/// Per-query plan-delivery state. As with the other transports, the plan-send metrics and the +/// query start timestamp live for the whole query. +struct SelfHostedDispatcher { + transport: SelfHostedShmTransport, + metrics: OnceLock, +} + +impl WorkerDispatch for SelfHostedDispatcher { + fn dispatch(&self, request: WorkerDispatchRequest<'_>) -> Result<()> { + let WorkerDispatchRequest { + stage, + routed_urls, + task_ctx, + metrics, + metrics_store, + join_set, + .. + } = request; + let metrics = self + .metrics + .get_or_init(|| CoordinatorToWorkerMetrics::new(metrics)) + .clone(); + + let token = get_distributed_cancellation_token(task_ctx); + let harness = match self.transport.queries.entry(stage.query_id) { + dashmap::Entry::Occupied(e) => Arc::clone(e.get()), + dashmap::Entry::Vacant(e) => { + let harness = Arc::new(QueryHarness::new( + stage.query_id, + token.clone(), + self.transport.queue_bytes, + metrics_store.cloned(), + )); + e.insert(Arc::clone(&harness)); + // The token fires when the head stream drops (normal completion included), which + // is the query's end of life; drop the registry entry then. The region itself + // stays alive through the Arcs the drivers and streams hold. + let queries = Arc::clone(&self.transport.queries); + let query_id = stage.query_id; + let watched = token.clone(); + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { + watched.cancelled().await; + queries.remove(&query_id); + }); + // One pump set per query: drains every proc's inbox so control frames flow even + // when no consumer or producer is actively draining (a fragment blocked on its + // feed, the metrics frames after the gather finished), and forwards the task + // metrics into the store. Lives until every driver and feed pump reported done. + join_set.spawn(run_pumps(Arc::clone(&harness), token.clone())); + // Plans travel as `SetPlan` frames like every other control message, but the + // rings only exist at finalize; this pump waits for them and ships the queue. + harness.participants.fetch_add(1, Ordering::SeqCst); + join_set.spawn(run_plan_delivery(Arc::clone(&harness), token.clone())); + harness + } + }; + + let mut encoded_tasks = Vec::with_capacity(routed_urls.len()); + for task_i in 0..routed_urls.len() { + encoded_tasks.push(encode_task_plan( + &stage.plan, + task_i, + stage.tasks, + task_ctx.session_config(), + )?); + } + harness.record_stage( + stage.num, + stage.tasks, + encoded_tasks.iter().map(|e| e.partitions).collect(), + ); + harness + .state + .lock() + .unwrap() + .coord_metrics + .get_or_insert_with(|| metrics.clone()); + harness.scan_for_child_routing(&stage.plan, stage.tasks)?; + + // The planner already baked the producer head into the shipped fragment. Re-state it so + // each worker's `insert_producer_head` re-creates the same head rather than flattening the + // stage to one partition. + let producer_head = producer_head_from_plan(&stage.plan).to_proto(task_ctx)?; + + let mut headers = get_config_extension_propagation_headers(task_ctx.session_config())?; + headers.extend(get_passthrough_headers(task_ctx.session_config())); + + for (task_i, (url, encoded)) in routed_urls.iter().zip(encoded_tasks).enumerate() { + let plan_size = encoded.plan_proto.len(); + + let task_key = pb::TaskKey { + query_id: serialize_uuid(&stage.query_id), + stage_id: stage.num as u64, + task_number: task_i as u64, + }; + let set_plan = pb::SetPlanRequest { + plan_proto: encoded.plan_proto, + task_count: stage.tasks as u64, + task_key: Some(task_key.clone()), + work_unit_feed_declarations: encoded.feed_declarations, + target_worker_url: url.to_string(), + query_start_time_ns: metrics.instantiation_time, + }; + // The plan reaches the driver as a `SetPlan` frame on its proc's inbox, the same + // wire crossing a separate-process worker would see; the driver only holds its + // address. The frames queue here because the rings are only built at finalize. + harness + .state + .lock() + .unwrap() + .pending_set_plans + .push(PendingSetPlan { + stage_num: stage.num as u32, + task_i, + frame: SetPlanFrame::from_parts(set_plan, &headers)?, + plan_size, + }); + + // Collected before spawning so the providers see the same eager `feed()` timing as + // they do under the other transports. + let feed_streams = + collect_task_work_unit_feeds(&stage.plan, task_ctx, task_i, stage.tasks)?; + let has_feeds = !feed_streams.is_empty(); + + let driver = TaskDriver { + harness: Arc::clone(&harness), + worker: self.transport.worker.clone(), + stage_num: stage.num as u32, + task_i, + n_partitions: encoded.partitions, + has_feeds, + task_key, + token: token.clone(), + producer_head: producer_head.clone(), + }; + harness.participants.fetch_add(1, Ordering::SeqCst); + join_set.spawn(driver.run()); + + // The feeds travel as `WorkUnit` frames through the leader's outbound senders, so + // the worker side reads them exactly as a separate process would. + if has_feeds { + harness.state.lock().unwrap().any_feeds = true; + harness.participants.fetch_add(1, Ordering::SeqCst); + join_set.spawn(run_leader_feed_pump( + Arc::clone(&harness), + stage.num as u32, + task_i, + feed_streams, + token.clone(), + )); + } + } + Ok(()) + } +} + +/// How a producer stage's output reaches its consumers, as parent-stage task indexes. A stage no +/// parent boundary ever claims has no entry: it is consumed by the head on the leader (proc 0). +/// +/// Built by simulating each consumer task's reads under its effective task contexts (a +/// `ChildrenIsolatorUnionExec` hands its children remapped contexts, so a boundary under one is +/// read with that inner context, not the stage-level one). `None` slots were never claimed by +/// any consumer (padding partitions); they route to the leader, where they sit buffered until +/// teardown. +enum RoutingSpec { + /// Nested `NetworkCoalesceExec`: consumers read whole producer tasks, so the destination + /// depends on the producer task only. + PerTask(Vec>), + /// Nested shuffle/broadcast: consumers read partition slices, identical across producer + /// tasks, so the destination depends on the output partition only. + PerPartition(Vec>), +} + +struct StageRec { + tasks: usize, + /// Output partitions of each task's specialized plan. Task-isolated nodes make these differ + /// from the unspecialized stage plan (and possibly from each other). + task_partitions: Vec, +} + +/// What a task driver needs to start producing: its proc's mesh and one routed sink per output +/// partition. +struct Launch { + mesh: Arc, + sinks: Vec>, + /// Send end for this task's `TaskMetrics` frame to the leader. + metrics_sender: MppSender, +} + +/// One task's plan, queued between `dispatch` and finalize, when the rings exist to carry it. +struct PendingSetPlan { + stage_num: u32, + task_i: usize, + frame: SetPlanFrame, + plan_size: usize, +} + +struct HarnessState { + stages: HashMap, + routing: HashMap, + started: bool, + /// Whether any dispatched task declared work-unit feeds; decides whether a feed pump runs. + any_feeds: bool, + /// Plans queued by `dispatch` until finalize builds the mesh; the delivery pump then ships + /// each as a `SetPlan` frame to the proc hosting its task. + pending_set_plans: Vec, + /// Dispatch-side metrics (plan send latency / bytes), recorded by the delivery pump. + coord_metrics: Option, + n_workers: u32, + leader_mesh: Option>, + /// The leader's outbound senders, the control-plane path for `WorkUnit` frames. + leader_senders: Vec>, + /// One mesh per worker proc, kept for the per-proc drain pumps. + worker_meshes: Vec>, + launches: HashMap<(u32, usize), Launch>, + /// Declared after the meshes so it would drop last either way; the harness Arcs held by + /// drivers and pinned streams are what actually keep it alive long enough. + region: Option, +} + +struct QueryHarness { + query_id: Uuid, + token: CancellationToken, + queue_bytes: usize, + metrics_store: Option>, + /// Drivers and feed pumps spawned for this query; `done` counts their exits. The pumps run + /// until the two meet, which is the deterministic "no more frames are coming" signal. + participants: AtomicUsize, + done: AtomicUsize, + state: Mutex, + ready_tx: tokio::sync::watch::Sender, + ready_rx: tokio::sync::watch::Receiver, +} + +impl QueryHarness { + fn new( + query_id: Uuid, + token: CancellationToken, + queue_bytes: usize, + metrics_store: Option>, + ) -> Self { + let (ready_tx, ready_rx) = tokio::sync::watch::channel(false); + Self { + query_id, + token, + queue_bytes, + metrics_store, + participants: AtomicUsize::new(0), + done: AtomicUsize::new(0), + state: Mutex::new(HarnessState { + stages: HashMap::new(), + routing: HashMap::new(), + started: false, + any_feeds: false, + pending_set_plans: Vec::new(), + coord_metrics: None, + n_workers: 0, + leader_mesh: None, + leader_senders: Vec::new(), + worker_meshes: Vec::new(), + launches: HashMap::new(), + region: None, + }), + ready_tx, + ready_rx, + } + } + + /// Wait until the first read finalizes the harness. + async fn ready(&self) -> Result<()> { + let mut rx = self.ready_rx.clone(); + while !*rx.borrow() { + rx.changed().await.map_err(|_| { + DataFusionError::Internal( + "self-hosted shm transport: harness dropped before start".to_string(), + ) + })?; + } + Ok(()) + } + + /// The leader's send end for one task's `WorkUnit` frames, cooperative-draining the + /// leader's own inbox so a symmetric full-ring stall cannot deadlock the feed push. + fn leader_feed_sender(&self, stage_num: u32, task_i: usize) -> Result { + let state = self.state.lock().unwrap(); + let dest_proc = proc_for_task(state.n_workers, task_i as u32); + let Some(base) = state + .leader_senders + .get(dest_proc as usize) + .and_then(|s| s.as_ref()) + else { + return internal_err!( + "self-hosted shm transport: no leader sender for proc {dest_proc}" + ); + }; + let Some(leader_mesh) = state.leader_mesh.clone() else { + return internal_err!("self-hosted shm transport: leader mesh not built"); + }; + Ok(base + .clone_with_header(MppFrameHeader::work_unit(stage_num, task_i as u32, 0)) + .with_cooperative_drain(leader_mesh as Arc)) + } + + fn record_stage(&self, num: usize, tasks: usize, task_partitions: Vec) { + let mut state = self.state.lock().unwrap(); + state.stages.insert( + num as u32, + StageRec { + tasks, + task_partitions, + }, + ); + } + + /// Classify the routing of every child stage referenced by `plan`'s network boundaries. The + /// children were dispatched before this stage (plan preparation converts bottom-up), so their + /// records exist; stages no parent ever claims are consumed by the head on the leader. + fn scan_for_child_routing( + &self, + plan: &Arc, + consumer_tasks: usize, + ) -> Result<()> { + let mut state = self.state.lock().unwrap(); + for task_i in 0..consumer_tasks { + let d_ctx = DistributedTaskContext { + task_index: task_i, + task_count: consumer_tasks, + }; + let state = &mut *state; + let mut scan_err = Ok(()); + plan.apply_with_dt_ctx(d_ctx, |node, d_ctx| { + let Some(nb) = node.as_ref().as_network_boundary() else { + return Ok(TreeNodeRecursion::Continue); + }; + let child_num = nb.input_stage().num() as u32; + let Some(rec) = state.stages.get(&child_num) else { + scan_err = internal_err!( + "self-hosted shm transport: stage {child_num} referenced before dispatch" + ); + return Ok(TreeNodeRecursion::Stop); + }; + let child_tasks = rec.tasks; + let child_max_parts = rec.task_partitions.iter().copied().max().unwrap_or(0); + + if node.as_ref().is::() { + let spec = state + .routing + .entry(child_num) + .or_insert_with(|| RoutingSpec::PerTask(vec![None; child_tasks])); + let RoutingSpec::PerTask(dest) = spec else { + scan_err = internal_err!( + "self-hosted shm transport: stage {child_num} read through mixed \ + boundary kinds" + ); + return Ok(TreeNodeRecursion::Stop); + }; + // Mirror of the consumer's `task_group` split: contiguous groups, the first + // `extra` groups one producer task longer. + let base = child_tasks / d_ctx.task_count.max(1); + let extra = child_tasks % d_ctx.task_count.max(1); + let len = base + usize::from(d_ctx.task_index < extra); + let start = d_ctx.task_index * base + d_ctx.task_index.min(extra); + let end = (start + len).min(child_tasks); + for slot in dest[start..end].iter_mut() { + *slot = Some(task_i as u32); + } + } else { + let spec = state + .routing + .entry(child_num) + .or_insert_with(|| RoutingSpec::PerPartition(vec![None; child_max_parts])); + let RoutingSpec::PerPartition(dest) = spec else { + scan_err = internal_err!( + "self-hosted shm transport: stage {child_num} read through mixed \ + boundary kinds" + ); + return Ok(TreeNodeRecursion::Stop); + }; + // Shuffle and broadcast read the same partition slice per consumer context. + let p_c = nb.partitions_per_consumer_task(); + let from = (p_c * d_ctx.task_index).min(child_max_parts); + let to = (p_c * (d_ctx.task_index + 1)).min(child_max_parts); + for slot in dest[from..to].iter_mut() { + *slot = Some(task_i as u32); + } + } + Ok(TreeNodeRecursion::Continue) + })?; + scan_err?; + } + Ok(()) + } + + /// Size and build the mesh, resolve the routing, and release the waiting task drivers. Runs + /// once, on the first `open`. + fn ensure_started(&self) -> Result<()> { + let mut state = self.state.lock().unwrap(); + if state.started { + return Ok(()); + } + + let n_workers = state + .stages + .values() + .map(|s| s.tasks) + .max() + .unwrap_or(1) + .max(1) as u32; + let n_procs = n_workers + 1; + + let region_total = dsm_region_bytes(n_procs, self.queue_bytes, 0)?; + let region = HeapRegion::new(region_total); + let wakeup: Arc = Arc::new(NoopWakeup); + let interrupt: Arc = Arc::new(CancellationInterrupt(self.token.clone())); + + let leader_attach = unsafe { + leader_setup( + region.base(), + n_procs, + self.queue_bytes, + &[], + Arc::clone(&wakeup), + receiver_token(0), + Arc::clone(&interrupt), + // The harness holds the senders until the query ends, so attaching is safe, + // and the plan delivery pump always needs them. + true, + ) + }?; + let leader_mesh = leader_attach.mesh; + let mut worker_meshes = Vec::with_capacity(n_workers as usize); + for proc_idx in 1..n_procs { + let attach = unsafe { + worker_setup( + region.base(), + region_total, + proc_idx, + Arc::clone(&wakeup), + receiver_token(proc_idx), + Arc::clone(&interrupt), + ) + }?; + worker_meshes.push((attach.mesh, attach.outbound_senders)); + } + + // Build every fragment's launch package: one routed sink per output partition, sharing + // the proc's outbound senders. The base senders drop at the end of this scope, so the + // rings observe the last-sender detach once the per-partition clones finish. + let mut launches = HashMap::new(); + for (&stage_num, rec) in state.stages.iter() { + let spec = state.routing.get(&stage_num); + for task_i in 0..rec.tasks { + let n_out = rec.task_partitions.get(task_i).copied().unwrap_or(0); + let proc = proc_for_task(n_workers, task_i as u32); + let (mesh, outbound) = &worker_meshes[(proc - 1) as usize]; + let Some(to_leader) = outbound[0].as_ref() else { + return internal_err!( + "self-hosted shm transport: no outbound sender from proc {proc} to the \ + leader" + ); + }; + // No cooperative drain on purpose: metrics frames go out after the cancellation + // token fired (it fires on normal completion), when the interrupt-checking spin + // would abort the send. + let metrics_sender = to_leader.clone_with_header(MppFrameHeader::task_metrics( + stage_num, + task_i as u32, + proc, + )); + let mut sinks: Vec> = Vec::with_capacity(n_out); + for q in 0..n_out { + let consumer = match spec { + // No parent boundary claimed this stage: the head consumes it on the + // leader. + None => None, + Some(RoutingSpec::PerTask(dest)) => dest.get(task_i).copied().flatten(), + Some(RoutingSpec::PerPartition(dest)) => dest.get(q).copied().flatten(), + }; + let dest_proc = match (spec, consumer) { + (None, _) | (_, None) => 0, + (_, Some(parent_task)) => proc_for_task(n_workers, parent_task), + }; + let Some(base) = outbound[dest_proc as usize].as_ref() else { + return internal_err!( + "self-hosted shm transport: no outbound sender from proc {proc} to \ + proc {dest_proc}" + ); + }; + let sender = base + .clone_with_header(MppFrameHeader::batch(stage_num, q as u32, proc)) + .with_cooperative_drain(Arc::clone(mesh) as Arc); + sinks.push(Box::new(MppPartitionSink::new(sender))); + } + launches.insert( + (stage_num, task_i), + Launch { + mesh: Arc::clone(mesh), + sinks, + metrics_sender, + }, + ); + } + } + + state.n_workers = n_workers; + state.leader_mesh = Some(leader_mesh); + state.leader_senders = leader_attach.outbound_senders; + state.worker_meshes = worker_meshes.iter().map(|(m, _)| Arc::clone(m)).collect(); + state.launches = launches; + state.region = Some(region); + state.started = true; + drop(state); + let _ = self.ready_tx.send(true); + Ok(()) + } + + fn leader_mesh(&self) -> Result> { + self.state + .lock() + .unwrap() + .leader_mesh + .clone() + .ok_or_else(|| { + DataFusionError::Internal( + "self-hosted shm transport: leader mesh not built".to_string(), + ) + }) + } + + /// Wait until the harness is finalized, then take this fragment's launch package. + async fn wait_launch(&self, stage_num: u32, task_i: usize) -> Result { + self.ready().await?; + let mut state = self.state.lock().unwrap(); + state.launches.remove(&(stage_num, task_i)).ok_or_else(|| { + DataFusionError::Internal(format!( + "self-hosted shm transport: no launch package for stage {stage_num} task {task_i}" + )) + }) + } +} + +/// Delivers one task's plan through the worker session machinery and produces its fragment into +/// the mesh. +struct TaskDriver { + harness: Arc, + worker: Worker, + stage_num: u32, + task_i: usize, + n_partitions: usize, + has_feeds: bool, + task_key: pb::TaskKey, + token: CancellationToken, + /// The producer head the distributed planner already baked into the shipped fragment. + /// Re-stated here so the worker's `insert_producer_head` strips and re-adds the same head + /// (a no-op) instead of flattening the stage to a single partition. + producer_head: pb::execute_task_request::ProducerHead, +} + +impl TaskDriver { + async fn run(self) -> Result<()> { + let harness = Arc::clone(&self.harness); + let result = self.run_inner().await; + harness.done.fetch_add(1, Ordering::SeqCst); + result + } + + async fn run_inner(self) -> Result<()> { + let Self { + harness, + worker, + stage_num, + task_i, + n_partitions, + has_feeds, + task_key, + token, + producer_head, + } = self; + + // The launch package arrives when the first read finalizes the harness. A query torn + // down before any read just unwinds the driver. + let launch = tokio::select! { + launch = harness.wait_launch(stage_num, task_i) => launch?, + _ = token.cancelled() => return Ok(()), + }; + + // The plan arrives as a `SetPlan` frame on this proc's inbox, shipped by the leader's + // delivery pump; the pumps drain it into the registry this take resolves from. + let set_plan_frame = tokio::select! { + frame = launch.mesh.take_set_plan(stage_num, task_i as u32) => frame?, + _ = token.cancelled() => return Ok(()), + }; + let (set_plan, headers) = set_plan_frame.into_parts()?; + + let mesh = Arc::clone(&launch.mesh); + let outcome = worker + .set_task_plan(set_plan, headers, move |mut cfg| { + // Child-stage reads inside the decoded fragment must pull from this proc's + // inbox, and its dispatcher must stay a no-op: the plans are already here. + set_distributed_worker_transport(&mut cfg, ShmMqWorkerTransport::new(mesh)); + Ok(cfg) + }) + .await?; + + // The feeds arrive as `WorkUnit` frames on this proc's inbox; install the channels so + // the drain fills them (and flushes whatever arrived first). The leader-side pump owns + // delivery and the `FeedEof` that ends the streams. + if has_feeds { + launch.mesh.register_work_unit_senders( + stage_num, + task_i as u32, + outcome.work_unit_senders, + ); + } + + let produce = async { + let request = pb::ExecuteTaskRequest { + task_key: Some(task_key), + target_partition_start: 0, + target_partition_end: n_partitions as u64, + producer_head: Some(producer_head), + }; + // Through `execute_local_task` rather than a bare `plan.execute` so the task metrics + // (added/executed/finished stamps, per-node metrics) flow exactly like a pull-based + // read would deliver them. + let (streams, _ctx) = execute_local_task(worker.task_data_entries(), request).await?; + if streams.len() != launch.sinks.len() { + return internal_err!( + "self-hosted shm transport: stage {stage_num} task {task_i} decoded into {} \ + partitions but routed {} sinks", + streams.len(), + launch.sinks.len() + ); + } + let mut futures = Vec::with_capacity(streams.len()); + for (mut stream, mut sink) in streams.into_iter().zip(launch.sinks) { + let token = token.clone(); + futures.push(async move { + let stream_result: Result<()> = async { + loop { + let batch = tokio::select! { + next = stream.next() => next, + // Head stream dropped: stop pulling so this fragment and its upstream + // scan unwind, instead of draining the input into a buffer no one reads. + _ = token.cancelled() => break, + }; + let Some(batch) = batch else { break }; + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + sink.send(&batch).await?; + // A downstream worker abandoned this stream (its mesh carries the cancel + // senders): stop pulling so the cancel cascades to this fragment's own + // producers, matching the embedder's `run_worker_fragment` loop. + if sink.cancelled() { + break; + } + } + Ok(()) + } + .await; + // EOF always, even after a failed send, so the consumer side unblocks; the + // stream error stays the primary one. + let eof_result = sink.finish().await; + stream_result.and(eof_result) + }); + } + // `join_all`, not fail-fast: cancelling sibling partitions mid-await would skip their + // EOFs and wedge the consumer's channel buffers. + let results = futures::future::join_all(futures).await; + for r in results { + r?; + } + Ok(()) + }; + let produce_res: Result<()> = produce.await; + + // The metrics receiver resolves as the last partition stream above completes, so this + // does not block on anything remote. The frame goes back over the mesh, where the + // leader-side pump forwards it into the metrics store. + if let Ok(task_metrics) = outcome.metrics_rx.await { + let _ = launch + .metrics_sender + .send_task_metrics_best_effort(&task_metrics) + .await; + } + produce_res + } +} + +/// Leader-side delivery of every dispatched plan: waits for finalize to build the rings, then +/// ships each queued plan as a `SetPlan` frame to the proc hosting its task. One pump per query, +/// counted as a participant so the drain pumps outlive it. +async fn run_plan_delivery(harness: Arc, token: CancellationToken) -> Result<()> { + let result = run_plan_delivery_inner(&harness, token).await; + harness.done.fetch_add(1, Ordering::SeqCst); + result +} + +async fn run_plan_delivery_inner( + harness: &Arc, + token: CancellationToken, +) -> Result<()> { + tokio::select! { + ready = harness.ready() => ready?, + _ = token.cancelled() => return Ok(()), + } + let (pending, senders, metrics) = { + let mut state = harness.state.lock().unwrap(); + let pending = std::mem::take(&mut state.pending_set_plans); + let Some(leader_mesh) = state.leader_mesh.clone() else { + return internal_err!("self-hosted shm transport: leader mesh not built"); + }; + let mut senders = Vec::with_capacity(pending.len()); + for plan in &pending { + let dest_proc = proc_for_task(state.n_workers, plan.task_i as u32); + let Some(base) = state + .leader_senders + .get(dest_proc as usize) + .and_then(|s| s.as_ref()) + else { + return internal_err!( + "self-hosted shm transport: no leader sender for proc {dest_proc}" + ); + }; + senders.push( + base.clone_with_header(MppFrameHeader::set_plan( + plan.stage_num, + plan.task_i as u32, + 0, + )) + .with_cooperative_drain(Arc::clone(&leader_mesh) as Arc), + ); + } + (pending, senders, state.coord_metrics.clone()) + }; + let mut stats = SendBatchStats::default(); + for (plan, sender) in pending.into_iter().zip(senders) { + let start = Instant::now(); + sender.send_set_plan_traced(&plan.frame, &mut stats).await?; + if let Some(metrics) = &metrics { + metrics.plan_send_latency.record(&start); + metrics.plan_bytes_sent.add_bytes(plan.plan_size); + } + } + Ok(()) +} + +/// Leader-side pump for one task's work-unit feeds: drives the providers and ships each unit as +/// a `WorkUnit` frame to the proc hosting the task, then closes the task's channels with a +/// `FeedEof`. The close is unconditional, error or not: without it the fragment's feed streams +/// never end and the worker side wedges instead of surfacing this pump's error. +async fn run_leader_feed_pump( + harness: Arc, + stage_num: u32, + task_i: usize, + feed_streams: Vec>>, + token: CancellationToken, +) -> Result<()> { + let result = async { + tokio::select! { + ready = harness.ready() => ready?, + _ = token.cancelled() => return Ok(()), + } + let sender = harness.leader_feed_sender(stage_num, task_i)?; + + let pump_result = async { + let mut pumps = Vec::with_capacity(feed_streams.len()); + for mut stream in feed_streams { + let sender = &sender; + pumps.push(async move { + let mut stats = SendBatchStats::default(); + while let Some(unit) = stream.next().await { + let mut unit = unit?; + set_sent_time(&mut unit); + sender.send_work_unit_traced(&unit, &mut stats).await?; + } + Ok::<_, DataFusionError>(()) + }); + } + futures::future::try_join_all(pumps).await.map(|_| ()) + } + .await; + + let mut stats = SendBatchStats::default(); + let eof_result = sender.send_feed_eof_traced(&mut stats).await; + pump_result.and(eof_result) + } + .await; + harness.done.fetch_add(1, Ordering::SeqCst); + result +} + +/// Per-proc drain pumps plus the metrics forwarder, alive until every driver and feed pump +/// reported done. They are what moves control frames when nothing else drains: a fragment +/// blocked on its feed before producing, and the metrics frames arriving after the gather +/// already finished. +async fn run_pumps(harness: Arc, token: CancellationToken) -> Result<()> { + tokio::select! { + ready = harness.ready() => ready?, + _ = token.cancelled() => return Ok(()), + } + let (leader_mesh, worker_meshes) = { + let state = harness.state.lock().unwrap(); + let Some(leader_mesh) = state.leader_mesh.clone() else { + return internal_err!("self-hosted shm transport: leader mesh not built"); + }; + (leader_mesh, state.worker_meshes.clone()) + }; + let metrics_rx = leader_mesh.take_task_metrics_receiver(); + + let mut meshes: Vec> = Vec::with_capacity(worker_meshes.len() + 1); + meshes.push(Arc::clone(&leader_mesh)); + meshes.extend(worker_meshes); + let pumps = meshes.into_iter().map(|mesh| { + let harness = Arc::clone(&harness); + async move { + loop { + let all_done = harness.done.load(Ordering::SeqCst) + >= harness.participants.load(Ordering::SeqCst); + // Drain errors surface through the consumers reading the same registry; + // the pump just stops contributing. + if mesh.try_drain_pass().is_err() { + break; + } + if all_done { + // The pass above ran after the last participant reported done, so every + // frame sent before that point has been routed. + break; + } + tokio::task::yield_now().await; + } + } + }); + futures::future::join_all(pumps).await; + + // Every frame is routed by now; whatever metrics arrived are in the channel. + if let (Some(mut rx), Some(store)) = (metrics_rx, harness.metrics_store.clone()) { + while let Ok((stage_id, task_number, task_metrics)) = rx.try_recv() { + store.insert( + pb::TaskKey { + query_id: serialize_uuid(&harness.query_id), + stage_id: stage_id as u64, + task_number: task_number as u64, + }, + task_metrics, + ); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; + use crate::test_utils::session_context::register_temp_parquet_table; + use crate::{DistributedConfig, DistributedExt, SessionStateBuilderExt, display_plan_ascii}; + use datafusion::arrow::array::{Int32Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::execute_stream; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + + /// Forces the ring mechanics on every batch: with `RING_SLOTS = 8`, a 64 KiB inbox has + /// ~8 KiB slots, so the ~16 KiB frames below fragment across slots, and the ~2 MB of + /// payload wraps each ring dozens of times, exercising the cooperative send spin. + const TINY_QUEUE_BYTES: usize = 64 * 1024; + + const ROWS: usize = 2000; + + fn sample_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, false), + Field::new("val", DataType::Int32, false), + ])); + // ~1 KiB per row, unique values so the GROUP BY keeps the full volume flowing + // through the shuffle instead of compacting it away at the partial aggregate. + let strings: Vec = (0..ROWS) + .map(|i| format!("{i:06}-{}", "x".repeat(1024))) + .collect(); + let vals: Vec = (0..ROWS as i32).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(strings)), + Arc::new(Int32Array::from(vals)), + ], + ) + .unwrap() + } + + async fn run(ctx: &SessionContext) -> Result<(String, Vec)> { + // Shaped so every ring frame stays bounded by `shuffle_batch_size`. The strings cross + // the shuffle inside `max`'s partial state, which the repartition rebuilds with `take` + // into fresh per-batch arrays; the projection then reduces them to a length before the + // gather. Shipping `s` itself out of a sort or an aggregate would not work: those emit + // offset slices of their accumulated state, a sliced variable-length array ships its + // whole values buffer through arrow-ipc, and a single frame balloons to the size of the + // partition's state no matter the batch size. + let query = "SELECT val, length(max(s)) AS l FROM t GROUP BY val"; + let plan = ctx.sql(query).await?.create_physical_plan().await?; + let display = display_plan_ascii(plan.as_ref(), false); + let batches: Vec<_> = execute_stream(plan, ctx.task_ctx())?.try_collect().await?; + let mut lines: Vec = pretty_format_batches(&batches)? + .to_string() + .lines() + .map(str::to_string) + .collect(); + lines.sort(); + Ok((display, lines)) + } + + /// A high-cardinality shuffle query over rings far smaller than the data, so every + /// cross-stage byte moves through fragmented frames under send-spin backpressure. The + /// result must still match the serial reference exactly. + #[tokio::test] + async fn tiny_rings_force_fragmentation_and_backpressure() -> Result<()> { + let transport = SelfHostedShmTransport::default().with_queue_bytes(TINY_QUEUE_BYTES); + // Small producer batches keep each frame a few slots big instead of overflowing the + // whole ring (a single frame must fit within one ring). + let d_cfg = DistributedConfig { + shuffle_batch_size: 16, + ..Default::default() + }; + let mut state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_option_extension(d_cfg) + .with_distributed_planner() + .with_distributed_task_estimator(2) + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(3)) + .with_distributed_worker_transport(transport) + .build(); + state.config_mut().options_mut().execution.target_partitions = 3; + let ctx = SessionContext::from(state); + let path = + register_temp_parquet_table("t", sample_batch().schema(), vec![sample_batch()], &ctx) + .await?; + + let (display, distributed) = run(&ctx).await?; + assert!( + display.contains("NetworkShuffleExec"), + "the query did not distribute:\n{display}" + ); + + let single = SessionContext::default(); + single + .register_parquet("t", path.to_string_lossy().as_ref(), Default::default()) + .await?; + let (_, expected) = run(&single).await?; + + assert_eq!(distributed, expected); + Ok(()) + } + + /// The per-query harness must be reclaimed when the output stream drops, or the `queries` map + /// grows one entry per query for the transport's lifetime. + #[tokio::test] + async fn query_harness_is_reclaimed_after_the_stream_drops() -> Result<()> { + let transport = SelfHostedShmTransport::default(); + let probe = transport.clone(); + let mut state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_planner() + .with_distributed_task_estimator(2) + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(3)) + .with_distributed_worker_transport(transport) + .build(); + state.config_mut().options_mut().execution.target_partitions = 3; + let ctx = SessionContext::from(state); + register_temp_parquet_table("t", sample_batch().schema(), vec![sample_batch()], &ctx) + .await?; + + let (display, _) = run(&ctx).await?; + assert!( + display.contains("NetworkShuffleExec"), + "the query did not distribute:\n{display}" + ); + + // The token fires as the gathered stream drops above; the registry entry is dropped on a + // task the cancelled token wakes, so poll briefly rather than assume it already ran. + for _ in 0..100 { + if probe.queries.is_empty() { + return Ok(()); + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + panic!( + "queries map still holds {} entries; the harness leaked", + probe.queries.len() + ); + } +} diff --git a/src/shm/setup.rs b/src/shm/setup.rs new file mode 100644 index 00000000..f86c67bb --- /dev/null +++ b/src/shm/setup.rs @@ -0,0 +1,371 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Mesh construction over a caller-supplied shared buffer, the extension point between an embedder's buffer +//! allocation and the transport. +//! +//! The embedder allocates the shared region (a PG `dsm_segment`, or a heap buffer in-process), +//! sizes it with [`dsm_region_bytes`], then calls [`leader_setup`] on the proc that initializes the +//! rings and [`worker_setup`] on each producer proc. Both hand back an [`MppMesh`] the embedder +//! installs on its DataFusion session; `worker_setup` also returns the outbound senders and the +//! plan bytes copied out of the region. [`run_worker_fragment`] is the producer push loop. +//! +//! The two platform primitives the embedder supplies are the [`Wakeup`] (how to wake a blocked +//! consumer) and the [`Interrupt`] (how to check for cancellation); everything here is otherwise +//! pure Rust over the shared buffer. + +use std::ffi::c_void; +use std::sync::Arc; + +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use futures::stream::StreamExt; + +use super::dsm::{ + compute_dsm_layout, leader_init, peer_proc_for_index, read_region_total, worker_attach, +}; +use super::mesh::{DsmInboxReceiver, DsmInboxSender}; +use super::mpsc_ring::{DsmMpscSender, NO_RECEIVER_TOKEN, Wakeup}; +use super::runtime::MppMesh; +use super::transport::{ + BatchChannelSender, DrainHandle, Interrupt, MppFrameHeader, MppReceiver, MppSender, + ReceiverScope, SELF_LOOP_CAPACITY, in_proc_channel, +}; +use crate::proto as pb; +use crate::work_unit_feed::RemoteWorkUnitFeedRegistry; +use crate::{DistributedTaskContext, PartitionSink, collect_plan_metrics_protos}; + +/// Total bytes the shared region needs for `n_procs` inboxes plus `plan_len` plan bytes, with +/// `queue_bytes` per inbox. The embedder reserves exactly this much before [`leader_setup`]. +pub fn dsm_region_bytes(n_procs: u32, queue_bytes: usize, plan_len: usize) -> Result { + compute_dsm_layout(n_procs, queue_bytes, plan_len) + .map(|l| l.region_total) + .map_err(|e| DataFusionError::Internal(format!("mpp: dsm_region_bytes: {e}"))) +} + +/// Read `region_total` out of the header a leader wrote, so a worker that just mapped the region +/// can size its [`worker_setup`] call without knowing the header layout. +/// +/// A caller that derives the size from the header forfeits [`worker_setup`]'s size validation +/// (it would compare the header against itself). Pass the mapped size from the embedder's own +/// bookkeeping where it is available. +/// +/// # Safety +/// - `base` must point at the start of a region a leader initialized via [`leader_setup`]. +/// - `base` must be at least 8-byte aligned (the header holds `u64` fields). +pub unsafe fn region_total(base: *const c_void) -> usize { + unsafe { read_region_total(base) as usize } +} + +/// Wrap each peer-indexed `DsmMpscSender` into an outbound `MppSender` keyed by destination proc. +/// The dispatcher `clone_with_header`s these per output partition before sending, so the +/// placeholder header is never observed on the wire. Slot `this_proc` stays `None` until the +/// self-loop install. +/// +/// Returns `(data, cancel)`. `data` is the producer's output senders. `cancel` is a control-plane +/// sibling onto each peer inbox, used by [`MppMesh::cancel_stream`]: a consumer reaches its producer +/// without counting as one of that producer's data senders, so a held `Cancel` sender never masks +/// the producer-gone `detached` signal. +fn build_outbound_senders( + this_proc: u32, + total_procs: u32, + peer_senders: Vec, +) -> (Vec>, Vec>) { + let mut senders: Vec> = (0..total_procs).map(|_| None).collect(); + let mut cancel: Vec> = (0..total_procs).map(|_| None).collect(); + for (peer_idx, dsm_send) in peer_senders.into_iter().enumerate() { + let target_proc = peer_proc_for_index(this_proc, peer_idx as u32); + // A `peer_proc_for_index` regression that maps a peer onto the self slot would be + // silently overwritten by the self-loop install and only surface later as a missing + // sender at dispatch; name the bug at its source. + debug_assert!( + target_proc != this_proc, + "peer index {peer_idx} mapped to the self proc {this_proc}" + ); + debug_assert!( + target_proc < total_procs, + "peer index {peer_idx} mapped to proc {target_proc} >= total {total_procs}" + ); + let control: Arc = + Arc::new(DsmInboxSender::new(dsm_send.to_control())); + cancel[target_proc as usize] = Some(MppSender::with_header( + control, + MppFrameHeader::batch(0, 0, this_proc), + )); + let shared: Arc = Arc::new(DsmInboxSender::new(dsm_send)); + senders[target_proc as usize] = Some(MppSender::with_header( + shared, + // Stamp `sender_proc = this_proc` so a stray frame that escapes the dispatcher's + // `clone_with_header` overwrite still identifies its origin on the drain side. + MppFrameHeader::batch(0, 0, this_proc), + )); + } + (senders, cancel) +} + +/// What [`leader_setup`] hands back to the embedder. +pub struct LeaderAttach { + /// The leader's mesh, installed on its DataFusion session. + pub mesh: Arc, + /// Outbound senders keyed by destination proc index, for the control plane: work-unit + /// frames flow leader -> worker through them. Slot 0 (the leader itself) stays `None`; + /// empty unless `attach_senders` was passed. Holders must keep them alive for the whole + /// query: dropping them before a worker attaches latches that worker's inbox as detached. + pub outbound_senders: Vec>, +} + +/// Initialize the shared region as the leader (`proc 0`) and return its mesh plus its outbound +/// senders. +/// +/// Writes the region header, copies `plan_bytes` in, initializes the `n_procs` inboxes, and +/// attaches the leader as receiver to its own inbox. `receiver_token` is registered so producers +/// resolve this proc's [`Wakeup`]; `interrupt` is consulted at the transport's block points. +/// +/// # Safety +/// - `base` must point at an uninitialized region of at least `dsm_region_bytes(n_procs, +/// queue_bytes, plan_bytes.len())` bytes. +/// - `base` must be at least 8-byte (MAXALIGN) aligned; the ring headers hold atomics. +/// - The region must not be concurrently accessed until this returns. +#[allow(clippy::too_many_arguments)] // mirrors worker_setup; the args are the embedder's knobs +pub unsafe fn leader_setup( + base: *mut c_void, + n_procs: u32, + queue_bytes: usize, + plan_bytes: &[u8], + wakeup: Arc, + receiver_token: u64, + interrupt: Arc, + attach_senders: bool, +) -> Result { + if receiver_token == NO_RECEIVER_TOKEN { + return Err(DataFusionError::Internal( + "mpp: leader_setup: receiver_token is the NO_RECEIVER_TOKEN sentinel; wakeups \ + for this proc would be silently disabled" + .into(), + )); + } + let layout = compute_dsm_layout(n_procs, queue_bytes, plan_bytes.len()) + .map_err(|e| DataFusionError::Internal(format!("mpp: leader_setup compute layout: {e}")))?; + let attach = unsafe { + leader_init( + base, + &layout, + plan_bytes, + Arc::clone(&wakeup), + attach_senders, + ) + } + .map_err(DataFusionError::Internal)?; + + let inbox = DsmInboxReceiver::new(attach.inbound_receiver); + inbox.set_receiver(receiver_token); + let inbound = Arc::new(DrainHandle::cooperative( + 0, + vec![(ReceiverScope::Inbox, MppReceiver::new(Box::new(inbox)))], + )); + // The leader hosts no producer fragments, but its senders carry the control plane: + // work-unit frames (and later dynamic filters) flow leader -> worker through them. Empty + // when the embedder did not opt in: a ring latches `detached` once its sender count hits + // zero, so senders that might drop before every worker attached must never exist. + // The leader's `Cancel` senders are wired by the embedder (it shares the same outbound senders + // it holds for plan delivery and releases before the DSM unmaps), so drop the cancel set here. + let (outbound_senders, _cancel_senders) = + build_outbound_senders(0, n_procs, attach.outbound_senders); + Ok(LeaderAttach { + mesh: Arc::new(MppMesh::new(0, n_procs, inbound, interrupt, attach.alive)), + outbound_senders, + }) +} + +/// Build one task's work-unit feed channels, install the receiving ends on `cfg` (where the +/// deserialized plan's remote feed leaves look them up), and register the sending ends on +/// `mesh`'s drain so inbound `WorkUnit` frames fill them. `feeds` lists the task's declared +/// feeds as `(feed id, partitions)`, the same pairs the plan's `WorkUnitFeedDeclaration`s carry. +/// +/// The caller must keep the proc draining (a consumer loop, a send spin, or an explicit +/// [`crate::shm::CooperativeDrainSet::try_drain_pass`] pump) while a fragment waits on its +/// feed, or the units sit in the inbox unread. +/// Build the [`TaskMetrics`] payload for one executed fragment, for embedders that run +/// fragments outside the worker task registry (pg parallel workers). Pair it with +/// [`super::transport::MppSender::send_task_metrics_best_effort`] after the fragment's streams +/// complete; the leader-side rewrite consumes the same pre-order the in-registry path produces. +/// The task-level stamps (plan added/executed/finished) stay unset on this path. +/// +/// [`TaskMetrics`]: crate::proto::TaskMetrics +pub fn collect_task_metrics( + plan: &Arc, + task_index: usize, + task_count: usize, +) -> pb::TaskMetrics { + pb::TaskMetrics { + pre_order_plan_metrics: collect_plan_metrics_protos( + plan, + DistributedTaskContext { + task_index, + task_count, + }, + ), + task_metrics: None, + } +} + +pub fn install_work_unit_channels( + cfg: &mut datafusion::prelude::SessionConfig, + mesh: &MppMesh, + stage_id: u32, + task_number: u32, + feeds: &[(uuid::Uuid, usize)], +) { + let mut channels = RemoteWorkUnitFeedRegistry::default(); + for (id, partitions) in feeds { + channels.add(*id, *partitions); + } + cfg.set_extension(Arc::new(channels.receivers)); + mesh.register_work_unit_senders(stage_id, task_number, channels.senders); +} + +/// What [`worker_setup`] hands back to the embedder. +pub struct WorkerAttach { + /// The worker's mesh, installed on its DataFusion session. + pub mesh: Arc, + /// Outbound senders keyed by destination proc index. The slot at `this_proc` is the in-proc + /// self-loop; every other slot writes to that peer's inbox. + pub outbound_senders: Vec>, + /// The plan bytes the leader wrote into the region, copied out for this worker. + pub plan_bytes: Vec, +} + +/// Attach to the leader-initialized region as worker `proc_idx` (`>= 1`). +/// +/// # Safety +/// - `base`/`region_total` must match the region the leader initialized via [`leader_setup`]. +/// - `base` must be at least 8-byte (MAXALIGN) aligned; the ring headers hold atomics. +pub unsafe fn worker_setup( + base: *mut c_void, + region_total: usize, + proc_idx: u32, + wakeup: Arc, + receiver_token: u64, + interrupt: Arc, +) -> Result { + if receiver_token == NO_RECEIVER_TOKEN { + return Err(DataFusionError::Internal( + "mpp: worker_setup: receiver_token is the NO_RECEIVER_TOKEN sentinel; wakeups \ + for this proc would be silently disabled" + .into(), + )); + } + let (header, plan_bytes, attach) = + unsafe { worker_attach(base, region_total as u64, proc_idx, Arc::clone(&wakeup)) } + .map_err(DataFusionError::Internal)?; + let total_procs = header.n_procs; + + let (mut outbound, cancel) = + build_outbound_senders(proc_idx, total_procs, attach.outbound_senders); + + // Self-loop in-proc channel: peer-mesh routing can land a producer and its consumer on the same + // proc, and an MPSC inbox has no slot for a proc sending to itself. The unified drain pulls from + // both the inbox and this channel. + let (self_tx, self_rx) = in_proc_channel(SELF_LOOP_CAPACITY); + let self_tx_arc: Arc = Arc::new(self_tx); + outbound[proc_idx as usize] = Some(MppSender::with_header( + self_tx_arc, + MppFrameHeader::batch(0, 0, proc_idx), + )); + + let inbox = DsmInboxReceiver::new(attach.inbound_receiver); + inbox.set_receiver(receiver_token); + let inbound = Arc::new(DrainHandle::cooperative( + proc_idx, + vec![ + (ReceiverScope::Inbox, MppReceiver::new(Box::new(inbox))), + (ReceiverScope::SelfLoop, MppReceiver::new(Box::new(self_rx))), + ], + )); + let mesh = Arc::new(MppMesh::new( + proc_idx, + total_procs, + inbound, + interrupt, + attach.alive, + )); + // A worker consumes shuffle inputs, so it can be the consumer that stops a stream early. Give + // its mesh the control-plane cancel senders; they drop with the mesh at the end of the worker's + // run, well before the DSM unmaps, so no explicit release is needed. + mesh.set_cancel_senders(Arc::new(std::sync::Mutex::new(cancel))); + Ok(WorkerAttach { + mesh, + outbound_senders: outbound, + plan_bytes, + }) +} + +/// Run a producer fragment plan to exhaustion, pushing every output batch into the matching +/// per-partition [`PartitionSink`]. The output partition count of `plan` must equal `sinks.len()`; +/// `sinks[partition]` is the send end the caller routed for that output partition. +/// +/// Each partition's [`PartitionSink::finish`] sends a per-channel EOF when its stream ends, +/// regardless of how it ended: the shared queue is multiplexed across fragments, so dropping a sink +/// doesn't end the channel, only the EOF frame does. +pub async fn run_worker_fragment( + plan: Arc, + sinks: Vec>, + ctx: Arc, +) -> Result<()> { + let n_partitions = plan.output_partitioning().partition_count(); + if n_partitions != sinks.len() { + return Err(DataFusionError::Internal(format!( + "run_worker_fragment: plan has {n_partitions} output partitions but {} sinks", + sinks.len() + ))); + } + let mut futures = Vec::with_capacity(n_partitions); + for (partition, mut sink) in sinks.into_iter().enumerate() { + let plan = Arc::clone(&plan); + let ctx = Arc::clone(&ctx); + futures.push(async move { + let stream_result: Result<()> = async { + let mut stream = plan.execute(partition, ctx)?; + while let Some(batch) = stream.next().await { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + sink.send(&batch).await?; + // Consumer abandoned this stream. Stop pulling: dropping `stream` ends the + // upstream scan and cascades the cancel to its own producers. + if sink.cancelled() { + break; + } + } + Ok(()) + } + .await; + let eof_result = sink.finish().await; + // Surface the stream error first, then any EOF-send error, so neither disappears. + stream_result.and(eof_result) + }); + } + // `join_all`, not `try_join_all`: fail-fast would cancel sibling partitions mid-await before + // they reach `finish`, leaving the consumer's channel buffer stuck. + let results = futures::future::join_all(futures).await; + for r in results { + r?; + } + Ok(()) +} diff --git a/src/shm/sink.rs b/src/shm/sink.rs new file mode 100644 index 00000000..a6ae19d7 --- /dev/null +++ b/src/shm/sink.rs @@ -0,0 +1,49 @@ +use async_trait::async_trait; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::Result; + +/// The producer's send end for one partition channel, symmetric to a [crate::WorkerChannel] read. +/// +/// This lives in the shm module rather than the core transport surface: only a push-based transport +/// (the shared-memory mesh) produces through sinks. Flight produces inside its gRPC worker service +/// and the in-memory transport pulls straight from the local task registry, so neither needs it. +/// +/// Contract with the produce loop: +/// - Batches arrive in `send` order and can be assumed non-empty. +/// - After a failed `send` the channel state is unspecified, but the caller still calls `finish` so +/// the consumer sees EOF; `finish` must tolerate a prior `send` error. +/// - Dropping a sink without calling `finish` does not end the channel, by design: `finish` is async +/// so Drop can't run it, and an implicit EOF would make an aborted producer look like a clean, +/// short stream. Abnormal teardown belongs to the transport, not the sink. +/// - `send` borrows the batch because transports serialize it into their own buffers; none needs +/// ownership. +#[async_trait] +pub trait PartitionSink: Send { + /// Sends one batch. Async so a blocked send can yield and let the transport make progress + /// elsewhere; a full channel must not park the calling thread. + async fn send(&mut self, batch: &RecordBatch) -> Result<()>; + /// Per-channel EOF, independent of the underlying link. Async for the same reason as `send`. + async fn finish(self: Box) -> Result<()>; + /// Whether the consumer cancelled this stream. The produce loop stops pulling its input when + /// this turns true, so a cancel doesn't just skip the send, it ends the upstream scan and drops + /// the input stream, cascading the cancel further up. Default `false` for links that don't carry + /// a cancel signal. + fn cancelled(&self) -> bool { + false + } +} + +/// The producer (write) side: opens a [PartitionSink] per output partition, symmetric to +/// [crate::WorkerChannel] (the read side). The worker's produce loop builds one and pushes each +/// output batch in. The shared-memory mesh provides the implementation; it is constructed by the +/// producer (which knows the per-partition routing), not handed out by the consume-side transport. +pub trait WorkerSink: Send + Sync { + /// Takes `stage` and `partition` separately because one sink serves every stage, unlike the + /// per-stage read connection that closes over its stage. + /// + /// `stage` is the producing stage's number and `partition` the producer task's own output + /// partition index, before routing. Several producer tasks of one stage may hold sinks for the + /// same pair, and the consumer merges them, so one `finish` is one producer task's EOF, not + /// channel completion (which stays transport-defined). + fn open_partition(&self, stage: usize, partition: usize) -> Result>; +} diff --git a/src/shm/transport.rs b/src/shm/transport.rs new file mode 100644 index 00000000..4b8fb2f6 --- /dev/null +++ b/src/shm/transport.rs @@ -0,0 +1,2935 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Transport layer for MPP shuffle. +//! +//! - [`MppFrameHeader`]: fixed 16-byte prefix tagging each wire message with +//! `(stage_id, partition)`, so one queue carries frames for many logical channels. +//! - [`encode_frame_into`] / [`decode_frame`]: Arrow IPC serialize/deserialize with +//! header prefix. Only codec entry points; tests round-trip through the same path. +//! - [`DrainBuffer`]: per-proc queue the drain writes into and the DataFusion consumer +//! reads from. Decouples consumer-side from producer-side backpressure: the drain +//! always makes forward progress on the inbound rings, so a stalled consumer can't +//! propagate backpressure to remote producers and cause a peer-mesh stall. + +use async_trait::async_trait; +use std::collections::VecDeque; +use std::sync::{Arc, Condvar, Mutex, MutexGuard}; + +use datafusion::common::{HashMap, HashSet}; +use std::time::{Duration, Instant}; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::ipc::reader::StreamReader; +use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::common::DataFusionError; +use prost::Message; + +use crate::common::deserialize_uuid; +use crate::proto as pb; +use crate::work_unit_feed::RemoteWorkUnitFeedTxs; +use crate::{WorkUnitMsg, set_received_time}; + +/// Magic bytes "MPPF" (MPP Frame) at the start of every wire message. +/// Lets receivers reject misrouted / corrupt frames before they hit Arrow IPC. +const MPP_FRAME_MAGIC: u32 = 0x4D505046; + +/// Wire-format size of [`MppFrameHeader`] in bytes. Asserted at compile time +/// below via `const _: ()`. +const MPP_FRAME_HEADER_SIZE: usize = 16; + +/// Kind of payload following [`MppFrameHeader`]. +/// +/// `Batch` is the common case. The header is followed by an Arrow IPC stream containing one +/// `RecordBatch`. `Eof` carries no payload. It signals the receiver that the named +/// `(stage_id, partition)` channel is done, even though the underlying shm_mq queue may still +/// carry frames for other channels. +/// +/// The remaining kinds are the control plane riding the same rings. For them the header's +/// `partition` field carries the task number instead: a work unit already names its +/// `(feed id, partition)` inside the payload, and metrics describe a whole task. +/// `SetPlan` (leader -> worker) carries one task's [`pb::SetPlanRequest`], the same message +/// Flight ships over its coordinator stream, plus the propagation headers that ride gRPC +/// metadata there; `WorkUnit` (leader -> worker) carries one prost-encoded work unit for +/// `(stage, task)`; `FeedEof` closes that task's feed channels (the wire stand-in for Flight's +/// stream close); `TaskMetrics` (worker -> leader) carries the task's collected metrics. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum MppFrameKind { + Batch = 0, + Eof = 1, + WorkUnit = 2, + FeedEof = 3, + TaskMetrics = 4, + SetPlan = 5, + /// Consumer -> producer: stop producing the `(stage_id, partition)` stream, the consumer + /// stopped reading it before EOF (a top-N `LIMIT`, an inner merge join exhausting a side, etc.). + /// The producer ends that one stream cleanly. Scoped to the stream, not the connection, so the + /// ring stays healthy for metrics and every other stream, the way gRPC closes one stream + /// without dropping the channel. + Cancel = 6, +} + +/// Payload of a `SetPlan` frame: the plan-delivery message a worker needs to run one task, +/// byte-compatible with what Flight sends. +/// +/// `set_plan` is the exact [`pb::SetPlanRequest`] the Flight dispatcher would put on its gRPC +/// stream. The headers are the config-extension propagation headers that travel as gRPC metadata +/// there; the ring has no metadata side channel, so they ride inside the frame as parallel +/// key/value lists (parallel lists rather than a map so repeated header names survive). +#[derive(Clone, PartialEq, prost::Message)] +pub struct SetPlanFrame { + #[prost(message, optional, tag = "1")] + pub set_plan: Option, + #[prost(string, repeated, tag = "2")] + pub header_keys: Vec, + #[prost(string, repeated, tag = "3")] + pub header_values: Vec, +} + +impl SetPlanFrame { + /// Bundle a plan-delivery message with the headers Flight would carry as gRPC metadata. + pub fn from_parts( + set_plan: pb::SetPlanRequest, + headers: &http::HeaderMap, + ) -> Result { + let mut header_keys = Vec::with_capacity(headers.len()); + let mut header_values = Vec::with_capacity(headers.len()); + for (name, value) in headers { + let value = value.to_str().map_err(|e| { + DataFusionError::Internal(format!( + "mpp: non-ASCII header {name} cannot travel in a SetPlan frame: {e}" + )) + })?; + header_keys.push(name.as_str().to_string()); + header_values.push(value.to_string()); + } + Ok(Self { + set_plan: Some(set_plan), + header_keys, + header_values, + }) + } + + /// Split back into the plan-delivery message and the propagation headers. + pub fn into_parts(self) -> Result<(pb::SetPlanRequest, http::HeaderMap), DataFusionError> { + let set_plan = self.set_plan.ok_or_else(|| { + DataFusionError::Internal("mpp: SetPlan frame carries no SetPlanRequest".to_string()) + })?; + let mut headers = http::HeaderMap::with_capacity(self.header_keys.len()); + for (key, value) in self.header_keys.iter().zip(self.header_values.iter()) { + let name = http::header::HeaderName::from_bytes(key.as_bytes()).map_err(|e| { + DataFusionError::Internal(format!("mpp: SetPlan frame header name {key:?}: {e}")) + })?; + let value = http::header::HeaderValue::from_str(value).map_err(|e| { + DataFusionError::Internal(format!("mpp: SetPlan frame header value for {key}: {e}")) + })?; + headers.append(name, value); + } + Ok((set_plan, headers)) + } +} + +/// 16-byte prefix on every transport frame. +/// +/// The fixed layout `[magic, flags, stage_id, partition]` (4×u32) is what +/// senders prepend before the Arrow IPC stream bytes and what receivers +/// parse before deciding which channel buffer the payload belongs to. +/// +/// See the `flags` bit-layout block below for the encoding of the `flags` word. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MppFrameHeader { + // Private so headers only come out of `batch()`/`eof()`: hand-built ones could bypass + // `pack_flags`'s sender bound and the reserved-bits invariant, and the consumer would + // reject them at decode, far from the producer. + pub(super) magic: u32, + pub(super) flags: u32, + pub(super) stage_id: u32, + pub(super) partition: u32, +} + +/// `flags` bit layout: +/// bits 0..8: frame kind (Batch | Eof) +/// bits 8..16: reserved (must be 0) +/// bits 16..32: sender_proc (mesh peer that wrote the frame) +const FRAME_KIND_MASK: u32 = 0x0000_00FF; +const FRAME_RESERVED_MASK: u32 = 0x0000_FF00; +const FRAME_SENDER_SHIFT: u32 = 16; +/// Maximum `sender_proc` representable in the header. Asserted at construction time so an +/// overflow becomes a hard error in the producer rather than silent flag corruption on the wire. +pub const MPP_MAX_SENDER_PROC: u32 = 0xFFFF; + +/// Spin bound for landing a control or metrics frame on a peer inbox that may be momentarily full. +/// A live peer drains its inbox within this many tries, so the frame lands well inside the bound; +/// it only runs out if the peer already exited, where the leader's wait-for-workers surfaces the +/// failure instead of hanging here. +const MAX_CONTROL_SEND_SPINS: usize = 10_000; + +const _: () = { + // shm_mq slot layout calculations depend on this being exact. + assert!(std::mem::size_of::() == MPP_FRAME_HEADER_SIZE); +}; + +#[inline] +fn pack_flags(kind: MppFrameKind, sender_proc: u32) -> u32 { + // fail_loud rather than debug_assert: in release builds the check would be compiled out and + // an out-of-range value would silently truncate to `sender_proc & 0xFFFF`. Catching the case + // where a refactor accidentally passes a task_id or partition here is the whole point. + assert!( + sender_proc <= MPP_MAX_SENDER_PROC, + "mpp: sender_proc {sender_proc} > MPP_MAX_SENDER_PROC ({MPP_MAX_SENDER_PROC})" + ); + (kind as u32) | (sender_proc << FRAME_SENDER_SHIFT) +} + +impl MppFrameHeader { + /// Build a `Batch` header for the given `(stage_id, partition)` stamped with `sender_proc`. + pub fn batch(stage_id: u32, partition: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::Batch, sender_proc), + stage_id, + partition, + } + } + + /// Build an `Eof` header for the given `(stage_id, partition)` stamped with `sender_proc`. + /// Carries no payload; receivers route it to the channel buffer's source-done counter. + pub fn eof(stage_id: u32, partition: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::Eof, sender_proc), + stage_id, + partition, + } + } + + /// Build a `WorkUnit` header addressed to `(stage_id, task_number)`. The `partition` slot + /// carries the task number; the unit's own `(feed id, partition)` ride in the payload. + pub fn work_unit(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::WorkUnit, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `FeedEof` header for `(stage_id, task_number)`: every feed of that task is done. + pub fn feed_eof(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::FeedEof, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `TaskMetrics` header for `(stage_id, task_number)`. + pub fn task_metrics(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::TaskMetrics, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `SetPlan` header for `(stage_id, task_number)`: the frame delivers that task's + /// plan to the proc hosting it. + pub fn set_plan(stage_id: u32, task_number: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::SetPlan, sender_proc), + stage_id, + partition: task_number, + } + } + + /// Build a `Cancel` header for the `(stage_id, partition)` stream, stamped with the consumer's + /// `sender_proc`. Carries no payload; the producer reads it as "stop sending this stream." + pub fn cancel(stage_id: u32, partition: u32, sender_proc: u32) -> Self { + Self { + magic: MPP_FRAME_MAGIC, + flags: pack_flags(MppFrameKind::Cancel, sender_proc), + stage_id, + partition, + } + } + + /// The mesh peer that wrote this frame. The drain demuxes incoming frames into the + /// per-channel buffer registry by `(sender_proc, stage_id, partition)`. + pub fn sender_proc(&self) -> u32 { + (self.flags >> FRAME_SENDER_SHIFT) & 0xFFFF + } + + /// Read the kind out of `flags`. Returns an error if the kind byte is + /// unknown or if any reserved bit (bits 8..16) is set, which catches wire-format + /// drift early. Sender_proc bits (16..32) are not validated here; readers extract + /// them with `sender_proc()`. + pub(super) fn kind(&self) -> Result { + let reserved = self.flags & FRAME_RESERVED_MASK; + if reserved != 0 { + return Err(DataFusionError::Internal(format!( + "mpp: reserved frame flag bits set ({reserved:#x})" + ))); + } + match self.flags & FRAME_KIND_MASK { + 0 => Ok(MppFrameKind::Batch), + 1 => Ok(MppFrameKind::Eof), + 2 => Ok(MppFrameKind::WorkUnit), + 3 => Ok(MppFrameKind::FeedEof), + 4 => Ok(MppFrameKind::TaskMetrics), + 5 => Ok(MppFrameKind::SetPlan), + 6 => Ok(MppFrameKind::Cancel), + other => Err(DataFusionError::Internal(format!( + "mpp: unknown frame kind {other:#x}" + ))), + } + } + + /// Serialize into the first `MPP_FRAME_HEADER_SIZE` bytes of `out`. + /// `out.len()` must be `>= MPP_FRAME_HEADER_SIZE`. + fn write_to(&self, out: &mut [u8]) { + debug_assert!(out.len() >= MPP_FRAME_HEADER_SIZE); + out[0..4].copy_from_slice(&self.magic.to_le_bytes()); + out[4..8].copy_from_slice(&self.flags.to_le_bytes()); + out[8..12].copy_from_slice(&self.stage_id.to_le_bytes()); + out[12..16].copy_from_slice(&self.partition.to_le_bytes()); + } + + /// Parse from the first `MPP_FRAME_HEADER_SIZE` bytes of `bytes`. Returns + /// `Err` if the slice is too short or the magic doesn't match. + fn parse(bytes: &[u8]) -> Result { + if bytes.len() < MPP_FRAME_HEADER_SIZE { + // No encoder in this file emits sub-header output, so a short frame means the + // shm_mq stitched together payloads from different senders. Hex-dump the bytes + // so the source is identifiable from log output without a debugger. + let hex = bytes + .iter() + .map(|b| format!("{b:02x}")) + .collect::>() + .join(" "); + return Err(DataFusionError::Internal(format!( + "mpp: frame too short for header ({} < {}); bytes = [{hex}]", + bytes.len(), + MPP_FRAME_HEADER_SIZE + ))); + } + let magic = u32::from_le_bytes(bytes[0..4].try_into().unwrap()); + if magic != MPP_FRAME_MAGIC { + return Err(DataFusionError::Internal(format!( + "mpp: bad frame magic {magic:#x} (expected {MPP_FRAME_MAGIC:#x})" + ))); + } + Ok(Self { + magic, + flags: u32::from_le_bytes(bytes[4..8].try_into().unwrap()), + stage_id: u32::from_le_bytes(bytes[8..12].try_into().unwrap()), + partition: u32::from_le_bytes(bytes[12..16].try_into().unwrap()), + }) + } +} + +/// Serialize `batch` into `buf` with a 16-byte [`MppFrameHeader`] prefix +/// addressing it to `(stage_id, partition)`. Wire format: +/// +/// ```text +/// [ magic | flags | stage_id | partition ] [ Arrow IPC stream bytes ] +/// |---------- 16 bytes --------| |---- variable ----| +/// ``` +/// +/// `flags` encodes kind + sender_proc; see the bit-layout block near +/// `FRAME_KIND_MASK` for details. +/// +/// Caller is expected to hold `buf` alive across many encodes so the peak-sized +/// allocation amortizes (~500 KB/batch on the 25M GROUP BY bench). +fn encode_frame_into( + header: MppFrameHeader, + batch: &RecordBatch, + buf: &mut Vec, +) -> Result<(), DataFusionError> { + buf.clear(); + buf.resize(MPP_FRAME_HEADER_SIZE, 0); + header.write_to(&mut buf[..MPP_FRAME_HEADER_SIZE]); + let mut writer = StreamWriter::try_new(&mut *buf, batch.schema_ref())?; + writer.write(batch)?; + writer.finish()?; + Ok(()) +} + +/// Serialize a payload-less [`MppFrameKind::Eof`] frame for `(stage_id, partition)` +/// into `buf`. The shm_mq peer reads this as a 16-byte message and routes it to +/// the channel buffer's source-done counter without touching Arrow IPC. +/// Consumed by [`MppSender::send_eof_traced`] when a producer fragment's +/// per-partition stream exhausts, so the receiver's `(stage_id, partition)` +/// channel buffer transitions to `Eof` even though the multiplexed shm_mq queue +/// stays attached for other channels. +fn encode_eof_frame_into( + stage_id: u32, + partition: u32, + sender_proc: u32, + buf: &mut Vec, +) -> Result<(), DataFusionError> { + buf.clear(); + buf.resize(MPP_FRAME_HEADER_SIZE, 0); + MppFrameHeader::eof(stage_id, partition, sender_proc) + .write_to(&mut buf[..MPP_FRAME_HEADER_SIZE]); + Ok(()) +} + +/// Serialize a prost-encoded control payload (`WorkUnit` / `TaskMetrics`) behind `header`. +fn encode_prost_frame_into( + header: MppFrameHeader, + msg: &impl prost::Message, + buf: &mut Vec, +) -> Result<(), DataFusionError> { + buf.clear(); + buf.resize(MPP_FRAME_HEADER_SIZE, 0); + header.write_to(&mut buf[..MPP_FRAME_HEADER_SIZE]); + msg.encode(buf) + .map_err(|e| DataFusionError::Internal(format!("mpp: prost frame encode: {e}")))?; + Ok(()) +} + +/// A decoded frame payload, routed by the drain according to its kind. +#[derive(Debug)] +enum FrameBody { + Batch(RecordBatch), + Eof, + WorkUnit(pb::WorkUnit), + FeedEof, + TaskMetrics(pb::TaskMetrics), + SetPlan(SetPlanFrame), + Cancel, +} + +/// Inverse of the frame encoders. Parses the 16-byte header and decodes the payload according +/// to the kind. Receivers branch on the body to decide routing. +fn decode_frame(bytes: &[u8]) -> Result<(MppFrameHeader, FrameBody), DataFusionError> { + let header = MppFrameHeader::parse(bytes)?; + let payload = &bytes[MPP_FRAME_HEADER_SIZE..]; + match header.kind()? { + MppFrameKind::Eof | MppFrameKind::FeedEof | MppFrameKind::Cancel => { + if bytes.len() != MPP_FRAME_HEADER_SIZE { + return Err(DataFusionError::Internal(format!( + "mpp: payload-less frame carries payload ({} > {})", + bytes.len(), + MPP_FRAME_HEADER_SIZE + ))); + } + match header.kind()? { + MppFrameKind::Eof => Ok((header, FrameBody::Eof)), + MppFrameKind::Cancel => Ok((header, FrameBody::Cancel)), + _ => Ok((header, FrameBody::FeedEof)), + } + } + MppFrameKind::WorkUnit => { + let unit = pb::WorkUnit::decode(payload) + .map_err(|e| DataFusionError::Internal(format!("mpp: work unit decode: {e}")))?; + Ok((header, FrameBody::WorkUnit(unit))) + } + MppFrameKind::TaskMetrics => { + let metrics = pb::TaskMetrics::decode(payload) + .map_err(|e| DataFusionError::Internal(format!("mpp: task metrics decode: {e}")))?; + Ok((header, FrameBody::TaskMetrics(metrics))) + } + MppFrameKind::SetPlan => { + let frame = SetPlanFrame::decode(payload) + .map_err(|e| DataFusionError::Internal(format!("mpp: set-plan decode: {e}")))?; + Ok((header, FrameBody::SetPlan(frame))) + } + MppFrameKind::Batch => { + let mut reader = StreamReader::try_new(payload, None)?; + let batch = reader.next().ok_or_else(|| { + DataFusionError::Execution("mpp: empty arrow-ipc stream in decode_frame".into()) + })??; + Ok((header, FrameBody::Batch(batch))) + } + } +} + +/// Local queue between a drain (either the cooperative `try_drain_pass` or the test-only thread +/// variant) and the consumer that pops batches. +/// +/// In the cooperative path each `DrainBuffer` corresponds to one logical channel: one +/// `(stage_id, partition)` entry in the owning [`DrainHandle`]'s registry. `num_sources` is +/// always `1` there because a given drain serves a single sender_proc, which is the only producer +/// for any channel routed through it. The test-only thread path uses a single shared buffer with +/// `num_sources = N` over an N-sender setup. +/// +/// Push side: callers append deserialized batches; on source detach (or per-channel `Eof` frame) +/// [`DrainBuffer::notify_source_done`] is called. Once `sources_done >= num_sources` AND the +/// queue is empty, `try_pop` returns [`DrainItem::Eof`]. +/// +/// Pop side: cooperative consumers loop on `try_pop` + `yield_now`. The test-only `pop_front` +/// blocks on the condvar. +#[derive(Debug)] +pub(super) struct DrainBuffer { + inner: Mutex, + cond: Condvar, +} + +#[derive(Debug)] +struct DrainBufferInner { + queue: VecDeque, + num_sources: u32, + sources_done: u32, + /// Consumer-side cancel flag. When set (e.g., query cancelled or `DrainHandle` dropped), + /// `try_pop`/`pop_front` returns `Eof` even if `sources_done` hasn't reached `num_sources`. + cancelled: bool, + /// Set when the receiver feeding this channel detached (or errored) before the channel's + /// `Eof` frame arrived. Distinct from `cancelled`: cancellation is a clean teardown and + /// yields `Eof`, while a lost source must surface as an error or the consumer would treat + /// truncated output as complete. + failed: Option, +} + +/// Yielded by [`DrainBuffer::pop_front`]. +#[derive(Debug)] +pub(super) enum DrainItem { + /// A batch produced by one of the inbound shm_mqs. + Batch(RecordBatch), + /// All source queues have detached and the local queue is drained. + Eof, + /// The receiver feeding this channel went away before the channel's `Eof` frame. + Failed(String), +} + +impl DrainBuffer { + /// Create a drain buffer expecting `num_sources` inbound queues. For a + /// proc in an N-proc mesh, `num_sources == N - 1` (all peers + /// excluding self — the self-partition bypasses the buffer). + pub fn new(num_sources: u32) -> Arc { + Arc::new(Self { + inner: Mutex::new(DrainBufferInner { + queue: VecDeque::new(), + num_sources, + sources_done: 0, + cancelled: false, + failed: None, + }), + cond: Condvar::new(), + }) + } + + /// Push a freshly-received batch into the local queue. + pub fn push_batch(&self, batch: RecordBatch) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + guard.queue.push_back(batch); + self.cond.notify_one(); + } + + /// Mark one source queue as detached. Safe to call from the drain thread + /// after observing `SHM_MQ_DETACHED` on a given inbound queue. + pub fn notify_source_done(&self) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + guard.sources_done = guard.sources_done.saturating_add(1); + if guard.sources_done >= guard.num_sources { + self.cond.notify_all(); + } + } + + /// Mark the channel as fed by a dead receiver, unless it already completed (its `Eof` + /// arrived), was cancelled, or already failed. Consumers then see an error instead of + /// hanging on a channel nothing will ever fill. + pub fn fail_pending(&self, msg: &str) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + if guard.sources_done >= guard.num_sources || guard.cancelled || guard.failed.is_some() { + return; + } + guard.failed = Some(msg.to_string()); + self.cond.notify_all(); + } + + /// Cancel all further pushes and wake all consumers with EOF. + pub fn cancel(&self) { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + guard.cancelled = true; + self.cond.notify_all(); + } + + /// Non-blocking variant. Returns the front item, or `DrainItem::Eof` if + /// all sources have detached and the queue is drained, or `None` if more + /// data may yet arrive. Cooperative consumers loop on + /// `try_drain_pass` + `try_pop`, yielding to the executor between + /// iterations. + pub fn try_pop(&self) -> Option { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + Self::try_pop_locked(&mut guard) + } + + /// Shared body of [`try_pop`] and the test-only [`Self::pop_front`]. + /// Returns `Some(Batch)` if the queue has data, `Some(Eof)` if all + /// sources have detached or the buffer is cancelled, and `None` + /// otherwise. Lets the two entry points stay in lockstep on the + /// "buffered data wins over cancellation/EOF" invariant locked in by + /// the `drain_buffer_drains_buffered_before_eof` test. + fn try_pop_locked(guard: &mut MutexGuard<'_, DrainBufferInner>) -> Option { + if let Some(batch) = guard.queue.pop_front() { + return Some(DrainItem::Batch(batch)); + } + if let Some(msg) = &guard.failed { + return Some(DrainItem::Failed(msg.clone())); + } + if guard.cancelled || guard.sources_done >= guard.num_sources { + return Some(DrainItem::Eof); + } + None + } +} + +/// Outcome of a single non-blocking receive attempt. +#[derive(Debug)] +pub(super) enum RecvOutcome { + /// One serialized Arrow IPC message ready to decode. + Bytes(Vec), + /// No data currently available but the peer is still attached. + Empty, + /// The peer has detached; no more bytes will ever arrive on this channel. + Detached, +} + +/// Non-blocking byte channel receiver. Implementations: `DsmInboxReceiver` (production), +/// `std::sync::mpsc` (tests). Must be `Send` because the drain thread takes ownership. +pub(super) trait BatchChannelReceiver: Send + Sync { + fn try_recv(&self) -> RecvOutcome; +} + +/// Byte channel sender paired with [`BatchChannelReceiver`]. `send` blocks when +/// the channel is full. Dropping the sender signals EOF to the receiver. +/// +/// `Send` is required because unit tests and future producer-pump threads move +/// senders across thread boundaries. +pub(crate) trait BatchChannelSender: Send + Sync { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError>; + + /// Non-blocking variant. Returns `Ok(true)` on success, `Ok(false)` + /// when the channel is full (caller should retry), `Err` on detach / + /// transport error. Default falls back to the blocking send — safe + /// for in-proc channels used by tests where "full" doesn't arise. + fn try_send_bytes(&self, bytes: &[u8]) -> Result { + self.send_bytes(bytes).map(|()| true) + } + + /// Async lock the send paths hold across the cooperative-drain spin so two tasks can't + /// interleave partial writes on the same handle. PG's `shm_mq_send` requires the same + /// `(nbytes, data)` on retry after `SHM_MQ_WOULD_BLOCK`. Multiple [`MppSender`] clones + /// multiplex onto one channel, and the spin's `yield_now().await` would otherwise let a + /// sibling task land a different payload mid-message and corrupt the queue. In-proc + /// channels return a per-instance mutex too, just to keep the call sites uniform. + fn send_lock(&self) -> &tokio::sync::Mutex<()>; +} + +/// Pluggable "drain everything inbound" hook for [`MppSender`]'s cooperative send spin. The +/// peer-mesh deadlock-breaking pattern needs the producer to pump ALL inbound queues (not just +/// one) while waiting for a full outbound, so the implementation typically delegates to +/// `MppMesh::drain_all_inbound()` which iterates every per-sender-proc drain. +pub trait CooperativeDrainSet: Send + Sync { + fn try_drain_pass(&self) -> Result<(), DataFusionError>; + + /// Checked in the send spin alongside `try_drain_pass`: returns `Err` (or aborts) if the + /// query should stop. Default is a no-op, for embedders with no external interrupt source; a + /// Postgres embedder overrides this to run `check_for_interrupts!`, which longjmps on cancel. + fn check_interrupt(&self) -> Result<(), DataFusionError> { + Ok(()) + } + + /// Whether a consumer cancelled the `(stage_id, partition)` stream (a `Cancel` frame arrived on + /// this proc's inbox). The send spin ends the producer's stream cleanly when it's set. Default + /// `false` for drains that don't carry inbound control frames (in-proc test channels). + fn stream_cancelled(&self, _stage_id: u32, _partition: u32) -> bool { + false + } +} + +/// Cancellation extension point, checked at the transport's block points (the send spin and the consumer +/// pull loop). An in-process embedder uses the default no-op or a cancellation token; a Postgres +/// embedder runs `check_for_interrupts!`, which longjmps out of the backend on cancel. +pub trait Interrupt: Send + Sync { + fn check(&self) -> Result<(), DataFusionError>; +} + +/// No-op interrupt for embedders that have no external cancellation source. +pub struct NoInterrupt; +impl Interrupt for NoInterrupt { + fn check(&self) -> Result<(), DataFusionError> { + Ok(()) + } +} + +impl CooperativeDrainSet for DrainHandle { + fn try_drain_pass(&self) -> Result<(), DataFusionError> { + DrainHandle::try_drain_pass(self) + } + + fn stream_cancelled(&self, stage_id: u32, partition: u32) -> bool { + DrainHandle::stream_cancelled(self, stage_id, partition) + } +} + +/// High-level sender: encodes a `RecordBatch` then pushes bytes through the underlying channel. +/// +/// With `cooperative_drain` set, `send_batch` breaks the symmetric-send deadlock on a +/// single-threaded tokio runtime by interleaving send-retries with +/// `CooperativeDrainSet::try_drain_pass` on the same mesh's inbound side. Each proc's +/// sender doing the same guarantees mutual progress: our drain pulls peer-shipped rows out of +/// our inbound queues, which frees peers' outbound-to-us send space, which lets their sends +/// un-stall. +pub struct MppSender { + /// Underlying byte channel. Held behind `Arc` so multiple `MppSender`s can share one + /// `shm_mq` queue while tagging frames with different `(stage_id, partition)` headers, which + /// is the multiplexed path's natural pattern. Clone the Arc, build a new `MppSender` with a + /// different header, both write into the same queue. + pub(super) channel: Arc, + cooperative_drain: Option>, + /// Frame header prepended to every outgoing batch. Identifies the logical + /// `(stage_id, partition)` channel the receiver demultiplexes on. Per-sender rather than + /// per-call: each partition gets its own `MppSender` via `clone_with_header`, all sharing + /// the underlying `Arc` of a single shm_mq queue. + pub(super) header: MppFrameHeader, + /// Scratch buffer reused across every `encode_frame_into` on this sender. Sized by the + /// first batch; subsequent batches clear and re-fill without reallocating. Interior + /// mutability lets the caller keep the `&self` signature (each producer fragment holds + /// its `MppSender` clones behind shared borrows for the duration of + /// `worker::run_worker_fragment`). + scratch: std::cell::RefCell>, +} + +// SAFETY: only `scratch: RefCell>` and the trait-object `Arc`s are `!Sync`. Callers +// compose `send_*_traced` futures via `tokio::spawn` / `join_all`, which makes the compiler +// require `&Self: Send` and therefore `Self: Sync`. The shared-memory model runs those futures on +// a current-thread runtime (see the module docs), so the cell is never observed from two +// threads; a multi-thread embedder would additionally be serialized by `send_lock` across +// every send path that touches `scratch`. +unsafe impl Sync for MppSender {} + +impl MppSender { + /// Construct a sender that tags every outgoing batch with `header`. Production call sites + /// clone one shared `Arc` across N senders, each with a different + /// `MppFrameHeader::batch(stage, p)`. That's the multiplexed pattern for fanning multiple + /// partitions over one shm_mq queue. + pub(super) fn with_header( + channel: Arc, + header: MppFrameHeader, + ) -> Self { + Self { + channel, + cooperative_drain: None, + header, + scratch: std::cell::RefCell::new(Vec::new()), + } + } + + /// Build a new `MppSender` that shares this sender's underlying channel + /// but tags every frame with `header` instead. Used by callers that know + /// the physical plan's output partition count and need one sender per + /// partition, all multiplexed over the same shm_mq queue. + pub fn clone_with_header(&self, header: MppFrameHeader) -> Self { + Self { + channel: Arc::clone(&self.channel), + cooperative_drain: self.cooperative_drain.as_ref().map(Arc::clone), + header, + scratch: std::cell::RefCell::new(Vec::new()), + } + } + + /// Whether the consumer of this sender's `(stage, partition)` stream cancelled it. Read by the + /// produce loop to stop pulling its input, not just skip the send. `false` without a drain + /// (in-proc test channels carry no inbound cancel). + pub(super) fn stream_cancelled(&self) -> bool { + self.cooperative_drain + .as_ref() + .is_some_and(|d| d.stream_cancelled(self.header.stage_id, self.header.partition)) + } + + /// Attach a [`CooperativeDrainSet`] so `Self::send_batch_traced`'s spin + /// can drain inbound peer traffic while waiting for outbound space. + /// Required for peer-mesh fragments where every worker is both sender and + /// consumer; without it, symmetric full-queue stalls deadlock the + /// single-threaded Tokio runtime. + pub fn with_cooperative_drain(mut self, drain: Arc) -> Self { + self.cooperative_drain = Some(drain); + self + } + + /// `send_batch` variant that accumulates per-call timings and spin counts into `stats`. + /// Callers that report at EOF (e.g. `ShuffleStream`) use this to diagnose where time + /// goes when the outbound queue is full. + /// + /// Async because the spin awaits the per-handle send lock and yields between + /// `try_send_bytes` retries; see `send_with_scratch`. + pub(super) async fn send_batch_traced( + &self, + batch: &RecordBatch, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + // Take the scratch buffer out of the `RefCell` rather than + // holding a `RefMut` across the spin below. The spin contains + // the embedder's `Interrupt::check`, which may unwind or `longjmp` through + // Rust frames; a `longjmp` does not run `Drop`, so a `RefMut` + // held across it would leave the cell perpetually borrowed and + // panic the next caller. `replace` is atomic — the cell is + // never observed in a borrowed state — and we put the buffer + // back at the end so its heap allocation survives across calls. + // If the spin longjmps anyway, the cell holds the default empty + // `Vec` and the next call simply re-allocates. + let mut scratch = self.scratch.replace(Vec::new()); + let result = self.send_with_scratch(batch, &mut scratch, stats).await; + self.scratch.replace(scratch); + result + } + + /// Send a payload-less [`MppFrameKind::Eof`] frame so the receiver's `(stage_id, partition)` + /// channel buffer transitions to `Eof` and the consumer's pull loop terminates cleanly. + /// + /// Producer fragments must call this exactly once per `(stage_id, partition)` channel after + /// the local stream exhausts. Without it the multiplexed shm_mq queue stays attached (other + /// channels still flow) and the consumer channel buffer never reaches `sources_done == 1`. The + /// receive-side [`DrainHandle::try_drain_pass`] decodes the frame and calls + /// `notify_source_done` on the matching channel buffer. + /// + /// Uses the same cooperative-spin path as [`Self::send_batch_traced`] so a full outbound + /// queue doesn't deadlock the EOF send. `stats.spin_iters` / `send_wait` capture any + /// contention. + /// + /// Symmetric-EOF safety: when every peer reaches EOF simultaneously with full outbound + /// queues, each peer's cooperative [`CooperativeDrainSet::try_drain_pass`] inside the spin + /// pulls peer-sent frames out of its own inbound queues, freeing space the peers are blocked + /// on. Progress is monotone: at least one `try_send_bytes` succeeds per spin iteration + /// somewhere in the mesh, so symmetric stalls resolve within a few iterations rather than + /// deadlocking. + pub(super) async fn send_eof_traced( + &self, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = self.send_eof_with_scratch(&mut scratch, stats).await; + self.scratch.replace(scratch); + result + } + + /// Bounded synchronous send of a `Cancel` frame for the `(stage_id, partition)` stream. The + /// consumer calls it on the sender to the producing proc when it abandons that stream. + /// Synchronous so it can run from the consumer stream's drop, where `await` isn't available. + /// + /// The bound doesn't risk a stuck producer. A live producer drains its own inbox in its send + /// spin, so a slot frees and the frame lands well inside the bound even when the inbox is + /// backed up. The bound only runs out if the producer already exited or died, and a dead worker + /// makes the leader's wait-for-workers error out rather than hang. + pub fn try_send_cancel(&self, stage_id: u32, partition: u32) { + let header = MppFrameHeader::cancel(stage_id, partition, self.header.sender_proc()); + let mut buf = [0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + for _ in 0..MAX_CONTROL_SEND_SPINS { + match self.channel.try_send_bytes(&buf) { + Ok(true) => return, + Ok(false) => std::thread::yield_now(), + Err(_) => return, // the producer's inbox is gone; nothing left to cancel + } + } + } + + async fn send_eof_with_scratch( + &self, + scratch: &mut Vec, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + encode_eof_frame_into( + self.header.stage_id, + self.header.partition, + self.header.sender_proc(), + scratch, + )?; + let Some(drain) = self.cooperative_drain.as_ref() else { + return self.channel.send_bytes(scratch); + }; + // Lock the channel before the spin so a sibling task can't interleave a different + // partial write through the shared shm_mq handle. See `BatchChannelSender::send_lock`. + let _send_guard = self.channel.send_lock().lock().await; + let mut first_try = true; + let t_wait_start = Instant::now(); + loop { + drain.check_interrupt()?; + // The consumer cancelled this stream, so end the send cleanly and let the producer + // fragment complete. Checked before the send so a cancel that landed mid-spin stops + // the next iteration. + if drain.stream_cancelled(self.header.stage_id, self.header.partition) { + return Ok(()); + } + if self.spin_try_send_bytes(scratch).await? { + if !first_try { + stats.send_wait += t_wait_start.elapsed(); + } + return Ok(()); + } + first_try = false; + stats.spin_iters += 1; + let t_drain = Instant::now(); + self.spin_try_drain_pass(drain).await?; + stats.coop_drain_in_spin += t_drain.elapsed(); + tokio::task::yield_now().await; + } + } + + /// Spin-loop helper: call `channel.try_send_bytes(scratch)`. + async fn spin_try_send_bytes(&self, scratch: &[u8]) -> Result { + self.channel.try_send_bytes(scratch) + } + + /// Spin-loop helper: call `drain.try_drain_pass()`. + async fn spin_try_drain_pass( + &self, + drain: &Arc, + ) -> Result<(), DataFusionError> { + drain.try_drain_pass() + } + + async fn send_with_scratch( + &self, + batch: &RecordBatch, + scratch: &mut Vec, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let t_enc = Instant::now(); + encode_frame_into(self.header, batch, scratch)?; + stats.encode += t_enc.elapsed(); + self.spin_send_scratch(scratch, stats).await + } + + /// Push an already-encoded frame through the channel via the cooperative-drain spin (or the + /// blocking fallback when no drain is attached). Shared by every frame kind's send path. + async fn spin_send_scratch( + &self, + scratch: &[u8], + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let Some(drain) = self.cooperative_drain.as_ref() else { + // No drain attached (unit tests, in-proc channels): fall + // back to the blocking send path. + return self.channel.send_bytes(scratch); + }; + // Lock the channel before the spin so a sibling task can't interleave a different + // partial write through the shared shm_mq handle. See `BatchChannelSender::send_lock`. + // Long-term, switching shm_mq for an async-friendly ring buffer (cf. #4184) drops the + // partial-send invariant entirely and removes the need for this lock. + // + // Latent under the current-thread runtime: today every fragment owns its own + // `Arc` (one sender per `Arc`), so the FIFO Mutex below + // is uncontended. A future multi-thread runtime that shares a sender across + // sibling fragment tasks (multi-partition fan-out) would let one task starve + // another for the duration of a large shuffle; the fix at that point is to move + // the entire spin off the compute thread. + let _send_guard = self.channel.send_lock().lock().await; + let mut first_try = true; + let t_wait_start = Instant::now(); + // The spin runs inside a tokio task on the backend thread's current-thread runtime + // (DataFusion needs one to drive `Stream`s). The deadlock we're breaking is + // *cross-proc*: two peers each blocked on a full outbound. `try_drain_pass` pulls + // peer batches off our inbound on the same OS thread, freeing their slots so their + // sends advance. `yield_now().await` between iterations hands the runtime back to + // siblings if any are ready, mostly a no-op under today's linear MPP topology. + loop { + drain.check_interrupt()?; + // The consumer cancelled this stream, so end the send cleanly and let the producer + // fragment complete. Checked before the send so a cancel that landed mid-spin stops + // the next iteration. + if drain.stream_cancelled(self.header.stage_id, self.header.partition) { + return Ok(()); + } + if self.spin_try_send_bytes(scratch).await? { + if !first_try { + stats.send_wait += t_wait_start.elapsed(); + } + return Ok(()); + } + first_try = false; + stats.spin_iters += 1; + // Would-block. Pull from our inbound so peers' outbound-to-us frees up and their + // sends to us unblock; without this, symmetric full-queue sends deadlock. Errors + // propagate so a peer detaching mid-spin doesn't leave us spinning on a closed + // mesh. + let t_drain = Instant::now(); + self.spin_try_drain_pass(drain).await?; + stats.coop_drain_in_spin += t_drain.elapsed(); + tokio::task::yield_now().await; + } + } + + /// Send one work unit for the task this sender's header names. The unit's hop stamps are the + /// caller's job; the payload travels prost-encoded, never through Arrow IPC. + pub async fn send_work_unit_traced( + &self, + unit: &pb::WorkUnit, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + encode_prost_frame_into(self.header, unit, &mut scratch)?; + self.spin_send_scratch(&scratch, stats).await + } + .await; + self.scratch.replace(scratch); + result + } + + /// Ship one task's plan as a `SetPlan` frame: the wire stand-in for Flight's + /// `SetPlanRequest` over its coordinator stream. + pub async fn send_set_plan_traced( + &self, + frame: &SetPlanFrame, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + encode_prost_frame_into(self.header, frame, &mut scratch)?; + self.spin_send_scratch(&scratch, stats).await + } + .await; + self.scratch.replace(scratch); + result + } + + /// Close the feed channels of the task this sender's header names: the wire stand-in for + /// Flight closing its coordinator stream, after which the worker-side feed streams end. + pub async fn send_feed_eof_traced( + &self, + stats: &mut SendBatchStats, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + scratch.clear(); + scratch.resize(MPP_FRAME_HEADER_SIZE, 0); + MppFrameHeader::feed_eof( + self.header.stage_id, + self.header.partition, + self.header.sender_proc(), + ) + .write_to(&mut scratch[..MPP_FRAME_HEADER_SIZE]); + self.spin_send_scratch(&scratch, stats).await + } + .await; + self.scratch.replace(scratch); + result + } + + /// Send the task's collected metrics without consulting the interrupt: this runs after the + /// query's cancellation token already fired (it fires on normal completion too), so the + /// cooperative spin would abort exactly when delivery matters. Best-effort like Flight's + /// metrics sends: a detached leader drops them. Retrying on a full ring is safe because the + /// receiving side keeps draining until every producer reported in. + pub async fn send_task_metrics_best_effort( + &self, + metrics: &pb::TaskMetrics, + ) -> Result<(), DataFusionError> { + let mut scratch = self.scratch.replace(Vec::new()); + let result = async { + encode_prost_frame_into(self.header, metrics, &mut scratch)?; + let _send_guard = self.channel.send_lock().lock().await; + // The consumer stops reading the data path early but keeps draining at teardown to + // collect these metrics, so the frame lands once the data the cancel stopped drains out + // and a slot frees. Bounded so a leader that really went away can't wedge the worker on + // a full ring. + for _ in 0..MAX_CONTROL_SEND_SPINS { + match self.channel.try_send_bytes(&scratch) { + Ok(true) => return Ok(()), + Ok(false) => tokio::task::yield_now().await, + Err(_) => return Ok(()), // receiver gone; metrics are best-effort + } + } + Ok(()) + } + .await; + self.scratch.replace(scratch); + result + } +} + +/// Per-call timing + spin metrics for [`MppSender::send_batch_traced`]. +/// All fields accumulate; callers zero or reuse as needed. +#[derive(Default, Debug, Clone)] +pub struct SendBatchStats { + /// Cumulative time spent inside `encode_frame_into` (header + Arrow IPC serialization). + pub encode: Duration, + /// Cumulative wall time in the send-retry spin after the first failed + /// `try_send_bytes`. Zero if the first try succeeded. + pub send_wait: Duration, + /// Cumulative time spent in `try_drain_pass` while spinning on a + /// full outbound. A subset of `send_wait`; the remainder is the + /// `tokio::task::yield_now()` await + the (small) cost of + /// `try_send_bytes` itself. + pub coop_drain_in_spin: Duration, + /// Count of `try_send_bytes` calls that returned `Ok(false)` (full). + pub spin_iters: u64, +} + +/// A [`crate::PartitionSink`] over one [`MppSender`]: the produce loop's per-partition send end. +/// `send` runs the cooperative-drain spin and `finish` flushes the channel EOF the same way, so a +/// non-Flight embedder (the in-process harness, pg_search's worker loop) wraps each routed sender +/// in one of these and pushes batches through the trait instead of touching `MppSender` directly. +pub struct MppPartitionSink { + sender: MppSender, + stats: SendBatchStats, +} + +impl MppPartitionSink { + pub fn new(sender: MppSender) -> Self { + Self { + sender, + stats: SendBatchStats::default(), + } + } + + /// Per-channel send counters, for an embedder that traces throughput. Read them before + /// `finish`, which consumes the sink. + pub fn stats(&self) -> &SendBatchStats { + &self.stats + } +} + +#[async_trait] +impl crate::PartitionSink for MppPartitionSink { + async fn send(&mut self, batch: &RecordBatch) -> datafusion::common::Result<()> { + self.sender.send_batch_traced(batch, &mut self.stats).await + } + + async fn finish(mut self: Box) -> datafusion::common::Result<()> { + self.sender.send_eof_traced(&mut self.stats).await + } + + fn cancelled(&self) -> bool { + self.sender.stream_cancelled() + } +} + +/// High-level receiver: pulls bytes via the underlying channel and decodes them +/// into `RecordBatch`. Used by the drain thread. +pub(super) struct MppReceiver { + channel: Box, +} + +impl MppReceiver { + pub fn new(channel: Box) -> Self { + Self { channel } + } + + pub(super) fn try_recv_batch(&self) -> RecvBatchOutcome { + match self.channel.try_recv() { + RecvOutcome::Bytes(bytes) => match decode_frame(&bytes) { + Ok((header, FrameBody::Batch(batch))) => RecvBatchOutcome::Batch { header, batch }, + Ok((header, FrameBody::Eof)) => RecvBatchOutcome::Eof { header }, + Ok((header, FrameBody::WorkUnit(unit))) => { + RecvBatchOutcome::WorkUnit { header, unit } + } + Ok((header, FrameBody::FeedEof)) => RecvBatchOutcome::FeedEof { header }, + Ok((header, FrameBody::TaskMetrics(metrics))) => { + RecvBatchOutcome::TaskMetrics { header, metrics } + } + Ok((header, FrameBody::SetPlan(frame))) => { + RecvBatchOutcome::SetPlan { header, frame } + } + Ok((header, FrameBody::Cancel)) => RecvBatchOutcome::Cancel { header }, + Err(e) => RecvBatchOutcome::Error(e), + }, + RecvOutcome::Empty => RecvBatchOutcome::Empty, + RecvOutcome::Detached => RecvBatchOutcome::Detached, + } + } +} + +/// Decoded result of an [`MppReceiver::try_recv_batch`]. Carries the +/// parsed [`MppFrameHeader`] so the drain thread can route the payload to +/// the right `(stage_id, partition)` channel buffer. +#[derive(Debug)] +pub(super) enum RecvBatchOutcome { + Batch { + header: MppFrameHeader, + batch: RecordBatch, + }, + /// A payload-less `Eof` frame for `header.(stage_id, partition)`. The + /// underlying shm_mq queue is still attached. The sender is just + /// signalling that this logical channel is done, so we can EOF + /// per-channel without dropping the whole queue. + Eof { + header: MppFrameHeader, + }, + /// One work unit for the task named by `header.(stage_id, partition=task)`. + WorkUnit { + header: MppFrameHeader, + unit: pb::WorkUnit, + }, + /// Every feed of the task named by the header is done; its channels close. + FeedEof { + header: MppFrameHeader, + }, + /// The plan for the task named by `header.(stage_id, partition=task)`. + SetPlan { + header: MppFrameHeader, + frame: SetPlanFrame, + }, + /// The collected metrics of the task named by the header. + TaskMetrics { + header: MppFrameHeader, + metrics: pb::TaskMetrics, + }, + /// Consumer abandoned the `header.(stage_id, partition)` stream; its producer stops sending it. + Cancel { + header: MppFrameHeader, + }, + Empty, + Detached, + Error(DataFusionError), +} + +/// Per-`(sender_proc, stage_id, partition)` channel buffer registry owned by a cooperative +/// [`DrainHandle`]. The handle may host several cooperative receivers (DSM MPSC inbox + self-loop +/// in-proc), each demultiplexed by the [`MppFrameHeader`] prefix into the same `map`. +/// `try_drain_pass` looks up the right channel buffer on every frame and pushes the payload into +/// it. Consumers waiting on a given key only see frames matching that key. +/// +/// Each entry is a `DrainBuffer::new(1)`: exactly one sender_proc emits frames for any given +/// channel. Per-channel EOF flows via the `Eof` frame demuxed onto the matching buffer; query- +/// teardown unblock flows via [`DrainHandle::cancel_channel_buffers`] from the handle's `Drop`. +#[derive(Default)] +struct ChannelBufferRegistry { + /// Keyed by `(sender_proc, stage_id, partition)`. The unified inbox carries frames + /// from every peer, so each `(stage, partition)` consumer gets its own per-sender + /// buffer. This preserves the implicit "one stream per sender" semantics that + /// `WorkerConnection::execute` consumers rely on. + map: HashMap<(u32, u32, u32), Arc>, + /// Scopes whose receiver detached (or errored) before draining cleanly. Channels fed by a + /// dead scope fail at registration time too, so a consumer that registers after the detach + /// does not wait on a channel nothing will ever fill. + dead_inbox: bool, + dead_self_loop: bool, +} + +/// Which frames a receiver carries, so a detach can fail exactly the channels it feeds. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(super) enum ReceiverScope { + /// The proc's DSM inbox: frames from every peer proc. + Inbox, + /// The in-proc self-loop: frames this proc sends itself. + SelfLoop, +} + +/// Per-sender-proc drain: stashes the receivers and polls them inline from the cooperative spin +/// (no background thread), demuxing each frame into a per-`(stage_id, partition)` channel buffer. +/// +/// Inline polling is the production requirement: pgrx's `check_active_thread` guard panics on any +/// pg FFI call (including `shm_mq_receive`) from a non-backend thread, so the drain work has to +/// run on the backend thread. Tests that need a true thread-backed drain use +/// [`ThreadedDrainHandle`] instead. +/// +/// On drop, the handle cancels every channel buffer so any consumer blocked on `try_pop` unblocks +/// with `Eof` — the drain can therefore never outlive its query, even on a panicked teardown. +pub struct DrainHandle { + /// Per-(stage_id, partition) channel buffer registry. Populated lazily on first frame for a + /// channel, or up-front by callers (e.g. `WorkerConnection::execute`) that need a + /// buffer to wait on before any frame arrives. + channel_buffers: Mutex, + /// Receivers owned by the handle and polled inline from `DrainGatherStream::poll_next` via + /// [`Self::try_drain_pass`]. The `Mutex` is for interior mutability: `try_drain_pass(&self)` + /// marks each slot as `None` after observing `Detached` so subsequent passes skip the dead + /// receiver. `BatchChannelReceiver: Send + Sync` makes `Vec>: Sync` + /// already, so the lock is no longer doubling as the `Sync` provider — replacing it with a + /// non-locking primitive would need either an atomic per-slot detached flag or accepting + /// that detached receivers get polled once per pass (fast-returning `Detached`). The lock + /// is uncontended in production (single backend thread) so the marginal cost is in the + /// type system, not the runtime. + coop_receivers: Mutex>>, + /// This proc's index, used to map a channel's `sender_proc` to the receiver scope that + /// feeds it (`SelfLoop` iff `sender_proc == this_proc`). + this_proc: u32, + /// Worker-side destination of `WorkUnit` frames, keyed `(stage_id, task_number)`. Frames + /// arriving before the embedder registers a task's channels buffer in `Pending`; `FeedEof` + /// drops the senders so the consuming feed streams end, the wire analog of Flight closing + /// its coordinator stream. + feed_registry: Mutex, + /// Leader-side destination of `TaskMetrics` frames: `(stage_id, task_number, metrics)`. + task_metrics_tx: tokio::sync::mpsc::UnboundedSender<(u32, u32, pb::TaskMetrics)>, + task_metrics_rx: + Mutex>>, + /// Worker-side destination of `SetPlan` frames, keyed `(stage_id, task_number)`. Same + /// pending-or-waiting shape as the feed registry: a frame arriving before the task asks + /// buffers in `Pending`; a task asking first parks a oneshot the drain fulfills. + set_plan_registry: Mutex, + /// `(stage_id, partition)` streams this proc's consumers abandoned, learned from inbound + /// `Cancel` frames. A producer blocked on a full outbound checks it in its send spin and ends + /// that stream cleanly, so a consumer that stopped reading early doesn't leave it spinning to + /// the statement timeout. + cancelled_streams: Mutex>, +} + +#[derive(Default)] +struct FeedRegistry { + map: HashMap<(u32, u32), FeedSlot>, + /// Set when the inbox scope died: feeds come from a peer proc, so a dead inbox means no + /// further units or `FeedEof` can arrive. Registered channels get the failure pushed in; + /// later registrations fail immediately. + dead: Option, +} + +enum FeedSlot { + /// Frames that arrived before the embedder registered the task's channels. + Pending { + units: Vec, + done: bool, + }, + Active(RemoteWorkUnitFeedTxs), +} + +/// Push one decoded unit into the channel its `(feed id, partition)` names. A missing channel is +/// not an error: the same tolerance the Flight worker applies to its stream (a feed the plan +/// does not declare is dropped). +fn forward_unit(senders: &RemoteWorkUnitFeedTxs, unit: pb::WorkUnit) { + let Ok(id) = deserialize_uuid(&unit.id) else { + return; + }; + let Some(tx) = senders.get(&(id, unit.partition as usize)) else { + return; + }; + // The feed channels carry the protocol's `WorkUnitMsg`, so the decoded proto frame is mapped + // into it the same way the gRPC worker's `decode_work_unit` does. + let _ = tx.send(Ok(WorkUnitMsg { + id, + partition: unit.partition as usize, + body: unit.body, + created_timestamp_unix_nanos: unit.created_timestamp_unix_nanos as usize, + sent_timestamp_unix_nanos: unit.sent_timestamp_unix_nanos as usize, + received_timestamp_unix_nanos: unit.received_timestamp_unix_nanos as usize, + processed_timestamp_unix_nanos: unit.processed_timestamp_unix_nanos as usize, + })); +} + +fn fail_feed_senders(senders: &RemoteWorkUnitFeedTxs, reason: &str) { + for tx in senders.values() { + let _ = tx.send(Err(DataFusionError::Execution(reason.to_string()))); + } +} + +#[derive(Default)] +struct SetPlanRegistry { + map: HashMap<(u32, u32), SetPlanSlot>, + /// Set when the inbox scope died: plans come from the leader's proc, so a dead inbox means + /// no plan can arrive. Parked takers get the failure; later takers fail immediately. + dead: Option, +} + +enum SetPlanSlot { + /// A frame that arrived before the task asked for it. + Pending(SetPlanFrame), + /// A task that asked before its frame arrived. + Waiting(tokio::sync::oneshot::Sender>), +} + +impl DrainHandle { + /// Construct a cooperative drain handle. Channel buffers are populated lazily by + /// [`Self::try_drain_pass`] when a frame arrives, or up-front by [`Self::register_channel`] + /// when a consumer needs a buffer to wait on before any frame has come in. + pub(super) fn cooperative( + this_proc: u32, + receivers: Vec<(ReceiverScope, MppReceiver)>, + ) -> Self { + let wrapped = receivers.into_iter().map(Some).collect(); + let (task_metrics_tx, task_metrics_rx) = tokio::sync::mpsc::unbounded_channel(); + Self { + channel_buffers: Mutex::new(ChannelBufferRegistry::default()), + coop_receivers: Mutex::new(wrapped), + this_proc, + feed_registry: Mutex::new(FeedRegistry::default()), + task_metrics_tx, + task_metrics_rx: Mutex::new(Some(task_metrics_rx)), + set_plan_registry: Mutex::new(SetPlanRegistry::default()), + cancelled_streams: Mutex::new(HashSet::default()), + } + } + + /// Record a `Cancel` frame: the consumer abandoned the `(stage_id, partition)` stream, so this + /// proc's producer of it stops. Idempotent. + fn note_cancel(&self, stage_id: u32, partition: u32) { + self.cancelled_streams + .lock() + .unwrap() + .insert((stage_id, partition)); + } + + /// Whether a consumer cancelled the `(stage_id, partition)` stream. Read by the send spin to + /// end a producer's stream cleanly when the consumer stopped reading. + pub(super) fn stream_cancelled(&self, stage_id: u32, partition: u32) -> bool { + self.cancelled_streams + .lock() + .unwrap() + .contains(&(stage_id, partition)) + } + + /// Take the receiving end of the `TaskMetrics` frame stream. The embedder (the leader) + /// drains it into its metrics store; the first caller gets it, later calls get `None`. + pub(super) fn take_task_metrics_receiver( + &self, + ) -> Option> { + self.task_metrics_rx.lock().unwrap().take() + } + + /// Install the senders of one task's feed channels, flushing any units that arrived first. + /// If the task's `FeedEof` (or the inbox death) already came through, the senders drop (or + /// fail) immediately so the consuming streams terminate instead of waiting forever. + pub(super) fn register_work_unit_senders( + &self, + stage_id: u32, + task_number: u32, + senders: RemoteWorkUnitFeedTxs, + ) { + let mut registry = self.feed_registry.lock().unwrap(); + if let Some(reason) = ®istry.dead { + fail_feed_senders(&senders, reason); + return; + } + match registry.map.remove(&(stage_id, task_number)) { + Some(FeedSlot::Pending { units, done }) => { + for unit in units { + forward_unit(&senders, unit); + } + if !done { + registry + .map + .insert((stage_id, task_number), FeedSlot::Active(senders)); + } + } + Some(FeedSlot::Active(_)) | None => { + registry + .map + .insert((stage_id, task_number), FeedSlot::Active(senders)); + } + } + } + + fn route_work_unit(&self, stage_id: u32, task_number: u32, unit: pb::WorkUnit) { + let mut registry = self.feed_registry.lock().unwrap(); + if registry.dead.is_some() { + return; + } + match registry.map.get_mut(&(stage_id, task_number)) { + Some(FeedSlot::Active(senders)) => forward_unit(senders, unit), + Some(FeedSlot::Pending { units, .. }) => units.push(unit), + None => { + registry.map.insert( + (stage_id, task_number), + FeedSlot::Pending { + units: vec![unit], + done: false, + }, + ); + } + } + } + + fn close_feeds(&self, stage_id: u32, task_number: u32) { + let mut registry = self.feed_registry.lock().unwrap(); + match registry.map.get_mut(&(stage_id, task_number)) { + Some(FeedSlot::Active(_)) => { + // Dropping the senders is the close: the consuming streams see end-of-input. + registry.map.remove(&(stage_id, task_number)); + } + Some(FeedSlot::Pending { done, .. }) => *done = true, + None => { + registry.map.insert( + (stage_id, task_number), + FeedSlot::Pending { + units: Vec::new(), + done: true, + }, + ); + } + } + } + + /// Route one decoded `SetPlan` frame to whoever asked for `(stage_id, task_number)`, or + /// buffer it until they do. A duplicate for an already-buffered slot keeps the first frame. + fn route_set_plan(&self, stage_id: u32, task_number: u32, frame: SetPlanFrame) { + let mut registry = self.set_plan_registry.lock().unwrap(); + match registry.map.remove(&(stage_id, task_number)) { + Some(SetPlanSlot::Waiting(tx)) => { + let _ = tx.send(Ok(frame)); + } + Some(pending @ SetPlanSlot::Pending(_)) => { + log::debug!( + "mpp: duplicate SetPlan frame for stage {stage_id} task {task_number}; \ + keeping the first" + ); + registry.map.insert((stage_id, task_number), pending); + } + None => { + registry + .map + .insert((stage_id, task_number), SetPlanSlot::Pending(frame)); + } + } + } + + /// Take the plan delivered for `(stage_id, task_number)`, waiting for its `SetPlan` frame if + /// it has not arrived yet. Something on this proc must keep draining (a pump or a + /// cooperative send spin) or the wait starves; same contract as the feed channels. + pub(super) async fn take_set_plan( + &self, + stage_id: u32, + task_number: u32, + ) -> Result { + let rx = { + let mut registry = self.set_plan_registry.lock().unwrap(); + if let Some(reason) = ®istry.dead { + return Err(DataFusionError::Execution(reason.clone())); + } + match registry.map.remove(&(stage_id, task_number)) { + Some(SetPlanSlot::Pending(frame)) => return Ok(frame), + Some(SetPlanSlot::Waiting(_)) => { + return Err(DataFusionError::Internal(format!( + "mpp: two takers for the SetPlan frame of stage {stage_id} task \ + {task_number}" + ))); + } + None => { + let (tx, rx) = tokio::sync::oneshot::channel(); + registry + .map + .insert((stage_id, task_number), SetPlanSlot::Waiting(tx)); + rx + } + } + }; + rx.await.map_err(|_| { + DataFusionError::Execution( + "mpp: transport torn down before this task's plan arrived".to_string(), + ) + })? + } + + fn scope_for_sender(&self, sender_proc: u32) -> ReceiverScope { + if sender_proc == self.this_proc { + ReceiverScope::SelfLoop + } else { + ReceiverScope::Inbox + } + } + + /// Fail every registered channel fed by `scope` that has not completed yet, and remember the + /// scope as dead so later registrations fail too. Channels whose `Eof` already arrived are + /// untouched: a detach after a clean drain is the normal end of life for a ring. + fn fail_scope(&self, scope: ReceiverScope, reason: &str) { + let to_fail = { + let mut guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + match scope { + ReceiverScope::Inbox => guard.dead_inbox = true, + ReceiverScope::SelfLoop => guard.dead_self_loop = true, + } + guard + .map + .iter() + .filter(|((sender_proc, _, _), _)| self.scope_for_sender(*sender_proc) == scope) + .map(|(_, buf)| buf.clone()) + .collect::>() + }; + for buf in to_fail { + buf.fail_pending(reason); + } + if scope == ReceiverScope::Inbox { + let mut registry = self.feed_registry.lock().unwrap(); + registry.dead = Some(reason.to_string()); + for (_, slot) in registry.map.drain() { + if let FeedSlot::Active(senders) = slot { + fail_feed_senders(&senders, reason); + } + } + drop(registry); + let mut plans = self.set_plan_registry.lock().unwrap(); + plans.dead = Some(reason.to_string()); + for (_, slot) in plans.map.drain() { + if let SetPlanSlot::Waiting(tx) = slot { + let _ = tx.send(Err(DataFusionError::Execution(reason.to_string()))); + } + } + } + } + + /// Register (or look up) the channel buffer for `(sender_proc, stage_id, partition)`. + /// The returned `Arc` is the canonical destination for frames matching + /// that key: `try_drain_pass` pushes into the same entry on every `Batch { header, .. }` + /// whose `header.sender_proc()` / `stage_id` / `partition` matches. + pub(super) fn register_channel( + &self, + sender_proc: u32, + stage_id: u32, + partition: u32, + ) -> Arc { + let mut guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + let scope_dead = match self.scope_for_sender(sender_proc) { + ReceiverScope::Inbox => guard.dead_inbox, + ReceiverScope::SelfLoop => guard.dead_self_loop, + }; + let buf = guard + .map + .entry((sender_proc, stage_id, partition)) + .or_insert_with(|| { + // num_sources stays 1: each (sender_proc, stage, partition) tuple has + // exactly one upstream (the named sender), even though the underlying + // inbox is shared across all senders. + DrainBuffer::new(1) + }) + .clone(); + drop(guard); + if scope_dead { + buf.fail_pending( + "transport receiver detached before this channel's EOF; the producer went away", + ); + } + buf + } + + /// Cancel every registered channel buffer. Called from `Drop` to unblock any consumer waiting on + /// a channel buffer when the handle goes away mid-query. + /// + /// Collects buffer handles under the registry lock, then notifies after releasing + /// it. Notifying inline would block any concurrent [`Self::register_channel`] for + /// as long as it takes to acquire `DrainBuffer::inner` N times. Fine today (single + /// backend thread), but cheap insurance against the multi-thread variant landing + /// later. + fn cancel_channel_buffers(&self) { + let to_cancel = { + let guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + guard.map.values().cloned().collect::>() + }; + for buf in to_cancel { + buf.cancel(); + } + } + + /// Pull batches from each live receiver and demux them into the per-`(stage_id, partition)` + /// channel buffer registry. Called from `DrainGatherStream::poll_next` and from + /// `MppSender::send_batch`'s cooperative spin. Drain work happens on the backend thread + /// (pgrx-safe). No-op for thread-backed handles. + /// + /// Each pass drains *every available* batch from each receiver (up to a safety cap). Pulling + /// only one batch per source per call would mean that under steady producer pressure the + /// cooperative sender's spin-loop can't keep up: we'd fall N:1 behind peers' sends and the + /// mesh would stall once any queue fills. Draining until the receiver reports `Empty` bounds + /// each pass by queue depth rather than by spin-loop iteration count. + /// + /// Returns `Ok(())` once every cooperative receiver has been pulled until `Empty` (or + /// detached). Errors propagate as `Err` so a transport-level failure surfaces at the call + /// site rather than getting silently dropped. + /// + /// Routing rules per outcome: + /// - `Batch { header, batch }`: look up (or lazily create) the + /// `(header.stage_id, header.partition)` channel buffer and push `batch`. + /// - `Eof { header }`: per-channel EOF. Resolve the channel buffer and call + /// `notify_source_done`. Other channels on the same queue keep flowing, + /// so the receiver slot stays live. + /// - `Detached` / `Error`: queue-wide shutdown. Notify every registered + /// channel buffer, mark the handle detached, and drop the slot. + pub fn try_drain_pass(&self) -> Result<(), DataFusionError> { + // Bound per-source pulls per call. The upper limit exists to give + // the caller a chance to re-try its own send between drains — + // otherwise a proc with a very fast peer could drain + // indefinitely on one source and starve its own outbound. + const MAX_BATCHES_PER_SOURCE_PER_PASS: usize = 256; + + let mut slots = self.coop_receivers.lock().unwrap(); + for slot in slots.iter_mut() { + let Some((scope, rx)) = slot.as_ref() else { + continue; + }; + let scope = *scope; + for _ in 0..MAX_BATCHES_PER_SOURCE_PER_PASS { + match rx.try_recv_batch() { + RecvBatchOutcome::Batch { header, batch } => { + let buf = self.register_channel( + header.sender_proc(), + header.stage_id, + header.partition, + ); + buf.push_batch(batch); + } + RecvBatchOutcome::Eof { header } => { + let buf = self.register_channel( + header.sender_proc(), + header.stage_id, + header.partition, + ); + buf.notify_source_done(); + // Other channels may still flow on this queue, so the receiver slot + // stays live. + } + RecvBatchOutcome::WorkUnit { header, mut unit } => { + set_received_time(&mut unit); + self.route_work_unit(header.stage_id, header.partition, unit); + } + RecvBatchOutcome::FeedEof { header } => { + self.close_feeds(header.stage_id, header.partition); + } + RecvBatchOutcome::TaskMetrics { header, metrics } => { + // The embedder may have dropped the receiver; metrics are best-effort. + let _ = + self.task_metrics_tx + .send((header.stage_id, header.partition, metrics)); + } + RecvBatchOutcome::SetPlan { header, frame } => { + self.route_set_plan(header.stage_id, header.partition, frame); + } + RecvBatchOutcome::Cancel { header } => { + self.note_cancel(header.stage_id, header.partition); + } + RecvBatchOutcome::Empty => break, + RecvBatchOutcome::Detached => { + // Only THIS receiver is dead. The drain holds multiple receivers + // (own-inbox MPSC + self-loop in-proc); one going away doesn't + // imply the others have. Fail only the channels this receiver's + // scope feeds, and only those still waiting on their `Eof`: after a + // clean drain the detach is the ring's normal end of life, but a + // channel that never got its `Eof` (producer crash, early sender + // drop) would otherwise spin on `try_pop -> None` forever. + *slot = None; + self.fail_scope( + scope, + "transport receiver detached before this channel's EOF; the \ + producer went away", + ); + break; + } + RecvBatchOutcome::Error(e) => { + // Same scoping as Detached, but the ring reported corruption (or the + // receiver poisoned itself), so even completed siblings can't be + // trusted to have been the last word. Still scope-limited: the other + // receiver is an independent transport. The error also propagates to + // this caller directly. + *slot = None; + self.fail_scope(scope, &format!("transport receiver failed: {e}")); + return Err(e); + } + } + } + } + Ok(()) + } +} + +impl Drop for DrainHandle { + fn drop(&mut self) { + // Unblock any consumer blocked on a channel buffer when the handle is torn down before EOF + // flows naturally (e.g. a query error en route to ExecEndCustomScan). + self.cancel_channel_buffers(); + } +} +/// SPSC channel pair for two use cases: +/// - Unit tests (bounded capacity, exercising backpressure). +/// - Production self-loop slots: when a worker's fragment emits a partition destined for +/// its OWN proc (e.g. peer-mesh hash routing where consumer task t lands on the same +/// worker as producer task t), the DSM layout has no self-pair inbox: a process is +/// not its own peer. The dispatcher routes those self-loops through this in-proc +/// channel, which exposes the same `BatchChannelSender` / `BatchChannelReceiver` +/// surface as the DSM ring so the drain and channel-buffer registry don't need a +/// special case. +/// +/// Production callers pass a very large `capacity` so the channel is effectively unbounded under +/// steady state. The current-thread Tokio runtime interleaves producer and consumer fragments +/// via `yield_now().await`, so backpressure would be benign anyway, but unbounded rules out any +/// chance of self-deadlock if the producer never yields. +pub(super) fn in_proc_channel(capacity: usize) -> (InProcSender, InProcReceiver) { + let (tx, rx) = std::sync::mpsc::sync_channel::>(capacity); + ( + InProcSender { + tx, + send_lock: tokio::sync::Mutex::new(()), + }, + InProcReceiver { rx: Mutex::new(rx) }, + ) +} + +pub(super) struct InProcSender { + tx: std::sync::mpsc::SyncSender>, + /// Per-instance lock so the [`BatchChannelSender::send_lock`] contract holds even when an + /// in-proc channel ends up in a code path that would otherwise need serialization. In-proc + /// `send_bytes` is already atomic (each call pushes a complete `Vec`), so the lock is + /// effectively a no-op here; keeping it uniform with `DsmInboxSender` avoids + /// special-casing the caller. + send_lock: tokio::sync::Mutex<()>, +} + +pub(super) struct InProcReceiver { + // The std::sync::mpsc receiver is !Sync; wrap in a Mutex so the drain + // thread can hold it behind a `Box` (which is + // `Send + Sync`-relaxed by design, but we only need Send for the thread + // hand-off). Tests only ever access from one thread so the Mutex is + // uncontended. + rx: Mutex>>, +} + +impl BatchChannelSender for InProcSender { + fn send_bytes(&self, bytes: &[u8]) -> Result<(), DataFusionError> { + self.tx.send(bytes.to_vec()).map_err(|_| { + DataFusionError::Execution("mpp: in-proc channel detached during send".into()) + }) + } + + fn try_send_bytes(&self, bytes: &[u8]) -> Result { + match self.tx.try_send(bytes.to_vec()) { + Ok(()) => Ok(true), + Err(std::sync::mpsc::TrySendError::Full(_)) => Ok(false), + Err(std::sync::mpsc::TrySendError::Disconnected(_)) => Err(DataFusionError::Execution( + "mpp: in-proc channel detached during try_send".into(), + )), + } + } + + fn send_lock(&self) -> &tokio::sync::Mutex<()> { + &self.send_lock + } +} + +impl BatchChannelReceiver for InProcReceiver { + fn try_recv(&self) -> RecvOutcome { + let rx = self.rx.lock().expect("InProcReceiver mutex poisoned"); + match rx.try_recv() { + Ok(bytes) => RecvOutcome::Bytes(bytes), + Err(std::sync::mpsc::TryRecvError::Empty) => RecvOutcome::Empty, + Err(std::sync::mpsc::TryRecvError::Disconnected) => RecvOutcome::Detached, + } + } +} + +/// Effectively unbounded capacity for self-loop in-proc channels. The +/// `std::sync::mpsc::sync_channel` API requires a numeric capacity; this constant picks one large +/// enough that production workloads won't reach it but small enough that a runaway producer +/// (e.g. infinite-loop bug) won't allocate billions of `Vec` before OOM. +pub(super) const SELF_LOOP_CAPACITY: usize = 1 << 20; + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Int32Array, Int64Array, StringArray, UInt64Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use std::sync::Arc as StdArc; + use std::thread; + + use std::thread::JoinHandle; + + impl DrainBuffer { + /// Block until a batch is available, EOF is reached, or the buffer is cancelled. + /// + /// INVARIANT: any already-buffered batch is returned *before* honoring either + /// cancellation or all-sources-done. Reordering the queue pop ahead of the cancel/eof + /// check would silently drop buffered data on an otherwise-clean shutdown; the + /// `drain_buffer_drains_buffered_before_eof` test locks this in. + fn pop_front(&self) -> DrainItem { + let mut guard = self.inner.lock().expect("DrainBuffer mutex poisoned"); + loop { + if let Some(batch) = guard.queue.pop_front() { + return DrainItem::Batch(batch); + } + if let Some(msg) = &guard.failed { + return DrainItem::Failed(msg.clone()); + } + if guard.cancelled || guard.sources_done >= guard.num_sources { + return DrainItem::Eof; + } + guard = self.cond.wait(guard).expect("DrainBuffer mutex poisoned"); + } + } + + /// True if `cancel` has been called. The local `drain_loop` consults this; the + /// cooperative production path watches the flag through `notify_source_done` fan-out + /// instead. + fn is_cancelled(&self) -> bool { + self.inner + .lock() + .expect("DrainBuffer mutex poisoned") + .cancelled + } + } + + impl MppSender { + /// Construct a sender with the default `(stage_id=0, partition=0)` header. Used where + /// the header carries no actionable routing info. + fn new(channel: Arc) -> Self { + Self::with_header(channel, MppFrameHeader::batch(0, 0, 0)) + } + + /// Stats-less wrapper around `send_batch_traced`. Production call sites + /// (`ShuffleStream::process_batch`) always pass a `SendBatchStats` so per-peer + /// wall-time shows up in the EOF trace. Wraps the async send in a tiny current-thread + /// Tokio runtime so `#[test]` functions don't have to be `#[tokio::test]` and the + /// OS-thread-spawning test harnesses don't have to plumb an async runtime themselves. + fn send_batch(&self, batch: &RecordBatch) -> Result<(), DataFusionError> { + let mut stats = SendBatchStats::default(); + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .expect("test tokio runtime build"); + rt.block_on(self.send_batch_traced(batch, &mut stats)) + } + } + + /// Configuration for `spawn_drain_thread`. pgrx panics on any pg FFI call (including + /// `shm_mq_receive`) from a non-backend thread, so production never spawns a drain thread — + /// see [`DrainHandle::cooperative`] for the cooperative path. + struct DrainConfig { + /// Receivers to drain. Ownership moves into the spawned thread. + receivers: Vec, + /// Destination buffer. + buffer: Arc, + /// How long to sleep when every receiver is empty but some are still attached. Tuning: + /// small values reduce end-of-batch latency but raise CPU; 1 ms is a safe default until + /// we integrate with WaitLatch. + idle_sleep: Duration, + } + + impl DrainConfig { + fn new(receivers: Vec, buffer: Arc) -> Self { + Self { + receivers, + buffer, + idle_sleep: Duration::from_millis(1), + } + } + } + + /// Spawn the dedicated drain thread. The thread round-robins through every receiver with + /// non-blocking `try_recv`, pushes decoded batches into `buffer`, and marks each source done + /// as soon as it observes a detach or decode error. When every source is done, the thread + /// exits. + fn spawn_drain_thread(config: DrainConfig) -> JoinHandle> { + thread::spawn(move || drain_loop(config)) + } + + /// RAII wrapper: owns the drain thread's `JoinHandle` and the buffer it writes into. + /// `Drop` cancels the buffer (unblocking the consumer) and joins the thread, so the thread + /// can never outlive the test scope even on a panic. + struct ThreadedDrainHandle { + buffer: Arc, + join: Mutex>>>, + } + + impl ThreadedDrainHandle { + fn spawn(config: DrainConfig) -> Self { + let buffer = Arc::clone(&config.buffer); + let join = spawn_drain_thread(config); + Self { + buffer, + join: Mutex::new(Some(join)), + } + } + } + + impl Drop for ThreadedDrainHandle { + fn drop(&mut self) { + self.buffer.cancel(); + if let Some(join) = self.join.lock().unwrap().take() { + let _ = join.join(); + } + } + } + + /// Test-only thread-backed drain. Writes every observed frame into a single shared + /// [`DrainBuffer`] with `num_sources = N`. Per-channel `Eof` frames are treated as "this source + /// is done" rather than "this logical channel within the source is done"; sufficient for unit + /// tests that don't exercise per-channel demux. Production drains route through + /// [`DrainHandle::try_drain_pass`] (cooperative variant), which keys on the frame header. Tests + /// that need to validate production demux must use [`DrainHandle::cooperative`] and call + /// `try_drain_pass` directly. + fn drain_loop(config: DrainConfig) -> Result<(), DataFusionError> { + let DrainConfig { + receivers, + buffer, + idle_sleep, + } = config; + + let mut done = vec![false; receivers.len()]; + loop { + // Observe cancellation before each pass so a `DrainHandle::drop` with + // live peer senders tears down cleanly. Without this check, the drain + // thread would spin `try_recv` forever because no source has detached. + if buffer.is_cancelled() { + return Ok(()); + } + + let mut got_any = false; + let mut all_done = true; + for (i, rx) in receivers.iter().enumerate() { + if done[i] { + continue; + } + all_done = false; + match rx.try_recv_batch() { + RecvBatchOutcome::Batch { header: _, batch } => { + got_any = true; + buffer.push_batch(batch); + } + RecvBatchOutcome::Eof { header: _ } => { + // Per-channel Eof frame: single-channel positional design + // treats it as a source-done signal. See `try_drain_pass`. + done[i] = true; + buffer.notify_source_done(); + } + // Control frames have no place in the test-only single-buffer path. + RecvBatchOutcome::WorkUnit { .. } + | RecvBatchOutcome::FeedEof { .. } + | RecvBatchOutcome::TaskMetrics { .. } + | RecvBatchOutcome::SetPlan { .. } + | RecvBatchOutcome::Cancel { .. } => {} + RecvBatchOutcome::Empty => {} + RecvBatchOutcome::Detached => { + done[i] = true; + buffer.notify_source_done(); + } + RecvBatchOutcome::Error(e) => { + // Treat a decode error as a fatal detach for this source + // so the consumer can observe EOF and abort the query. + done[i] = true; + buffer.notify_source_done(); + return Err(e); + } + } + } + + if all_done { + return Ok(()); + } + if !got_any { + thread::sleep(idle_sleep); + } + } + } + + fn sample_batch(rows: i32) -> RecordBatch { + let schema = StdArc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ])); + let ids = Int32Array::from_iter_values(0..rows); + let names = StringArray::from_iter_values((0..rows).map(|i| format!("n{i}"))); + RecordBatch::try_new(schema, vec![StdArc::new(ids), StdArc::new(names)]).unwrap() + } + + #[test] + fn frame_round_trips_a_batch_with_header() { + let orig = sample_batch(64); + let header = MppFrameHeader::batch(7, 3, 0); + let mut buf = Vec::with_capacity(1024); + encode_frame_into(header, &orig, &mut buf).expect("encode_frame"); + + let (parsed, body) = decode_frame(&buf).expect("decode_frame"); + assert_eq!(parsed, header); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::Batch); + let FrameBody::Batch(decoded) = body else { + panic!("Batch frame must carry a batch payload"); + }; + assert_eq!(decoded.num_rows(), 64); + assert_eq!(decoded.schema(), orig.schema()); + assert_eq!(decoded.num_columns(), orig.num_columns()); + for col in 0..orig.num_columns() { + assert_eq!(orig.column(col).as_ref(), decoded.column(col).as_ref()); + } + } + + #[test] + fn frame_round_trips_eof() { + let mut buf = Vec::new(); + encode_eof_frame_into(2, 5, 0, &mut buf).expect("encode_eof"); + assert_eq!(buf.len(), MPP_FRAME_HEADER_SIZE); + + let (header, body) = decode_frame(&buf).expect("decode_frame"); + assert_eq!(header, MppFrameHeader::eof(2, 5, 0)); + assert_eq!(header.kind().unwrap(), MppFrameKind::Eof); + assert!(matches!(body, FrameBody::Eof)); + } + + #[test] + fn frame_round_trips_cancel() { + let header = MppFrameHeader::cancel(4, 2, 0); + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + + let (parsed, body) = decode_frame(&buf).expect("decode_frame"); + assert_eq!(parsed, header); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::Cancel); + assert_eq!(parsed.stage_id, 4); + assert_eq!(parsed.partition, 2); + assert!(matches!(body, FrameBody::Cancel)); + } + + #[test] + fn drain_records_cancel_scoped_to_its_stream() { + let drain = DrainHandle::cooperative(0, Vec::new()); + assert!(!drain.stream_cancelled(7, 1)); + drain.note_cancel(7, 1); + assert!(drain.stream_cancelled(7, 1)); + // A cancel for one stream leaves the others alive, the way gRPC closes one stream: same + // stage but a different partition, and a different stage, both stay live. + assert!(!drain.stream_cancelled(7, 2)); + assert!(!drain.stream_cancelled(8, 1)); + } + + /// The producer-side half of the early-termination fix: a producer blocked on a full outbound + /// ends its send cleanly once a `Cancel` for its `(stage, partition)` stream lands on its inbox. + /// The test hangs if the spin doesn't observe the cancel. + #[tokio::test(flavor = "current_thread")] + async fn producer_send_ends_when_consumer_cancels_the_stream() { + // Outbound to a consumer that never drains: capacity 1, so the second send finds it full. + let (out_tx, _out_rx) = in_proc_channel(1); + // The producer's inbox, where the consumer's `Cancel` frame lands. + let (inbox_tx, inbox_rx) = in_proc_channel(4); + let drain = Arc::new(DrainHandle::cooperative( + 1, + vec![(ReceiverScope::Inbox, MppReceiver::new(Box::new(inbox_rx)))], + )); + // Producer of the `(stage 7, partition 0)` stream. + let sender = MppSender::with_header(Arc::new(out_tx), MppFrameHeader::batch(7, 0, 1)) + .with_cooperative_drain(Arc::clone(&drain) as Arc); + + // The consumer abandons that stream (it stopped reading before EOF): a `Cancel` reaches the + // inbox. + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + MppFrameHeader::cancel(7, 0, 0).write_to(&mut buf); + inbox_tx.send_bytes(&buf).unwrap(); + + let mut stats = SendBatchStats::default(); + // First batch fills the one slot. + sender + .send_batch_traced(&sample_batch(1), &mut stats) + .await + .unwrap(); + // The second batch would block forever on the full slot. Instead the spin drains the inbox, + // sees the cancel, and ends the send cleanly. + sender + .send_batch_traced(&sample_batch(1), &mut stats) + .await + .unwrap(); + } + + #[test] + fn frame_round_trips_a_work_unit_and_feed_eof() { + let unit = pb::WorkUnit { + id: vec![7; 16], + partition: 4, + body: vec![1, 2, 3], + created_timestamp_unix_nanos: 11, + sent_timestamp_unix_nanos: 22, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + }; + let header = MppFrameHeader::work_unit(3, 1, 0); + let mut buf = Vec::new(); + encode_prost_frame_into(header, &unit, &mut buf).expect("encode work unit"); + let (parsed, body) = decode_frame(&buf).expect("decode work unit"); + assert_eq!(parsed, header); + let FrameBody::WorkUnit(decoded) = body else { + panic!("WorkUnit frame must carry a unit"); + }; + assert_eq!(decoded, unit); + + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + MppFrameHeader::feed_eof(3, 1, 0).write_to(&mut buf); + let (parsed, body) = decode_frame(&buf).expect("decode feed eof"); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::FeedEof); + assert!(matches!(body, FrameBody::FeedEof)); + } + + #[test] + fn frame_round_trips_task_metrics() { + let metrics = pb::TaskMetrics { + pre_order_plan_metrics: vec![], + task_metrics: None, + }; + let header = MppFrameHeader::task_metrics(2, 0, 1); + let mut buf = Vec::new(); + encode_prost_frame_into(header, &metrics, &mut buf).expect("encode metrics"); + let (parsed, body) = decode_frame(&buf).expect("decode metrics"); + assert_eq!(parsed, header); + assert!(matches!(body, FrameBody::TaskMetrics(m) if m == metrics)); + } + + #[test] + fn work_units_buffer_until_registration_and_feed_eof_closes() { + fn unit(partition: u64) -> pb::WorkUnit { + pb::WorkUnit { + id: vec![9; 16], + partition, + body: vec![], + created_timestamp_unix_nanos: 0, + sent_timestamp_unix_nanos: 0, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + } + } + let drain = DrainHandle::cooperative(1, vec![]); + + // Units arriving before registration must buffer, not drop. + drain.route_work_unit(5, 0, unit(0)); + drain.route_work_unit(5, 0, unit(0)); + + let id = crate::common::deserialize_uuid(&[9; 16]).unwrap(); + let mut channels = crate::work_unit_feed::RemoteWorkUnitFeedRegistry::default(); + channels.add(id, 1); + let mut rx = channels + .receivers + .get(&(id, 0)) + .unwrap() + .lock() + .unwrap() + .take() + .unwrap(); + drain.register_work_unit_senders(5, 0, channels.senders); + + assert!(rx.try_recv().unwrap().is_ok()); + assert!(rx.try_recv().unwrap().is_ok()); + // Still open: the producer may send more units. + assert!(rx.try_recv().is_err()); + + // FeedEof drops the senders, which ends the stream. + drain.route_work_unit(5, 0, unit(0)); + drain.close_feeds(5, 0); + assert!(rx.try_recv().unwrap().is_ok()); + assert!(matches!( + rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) + )); + } + + #[test] + fn frame_round_trips_a_set_plan_with_headers() { + let set_plan = pb::SetPlanRequest { + plan_proto: vec![1, 2, 3, 4], + task_count: 2, + task_key: Some(pb::TaskKey { + query_id: vec![5; 16], + stage_id: 3, + task_number: 1, + }), + work_unit_feed_declarations: vec![], + target_worker_url: "inprocess://worker/1".to_string(), + query_start_time_ns: 42, + }; + let mut headers = http::HeaderMap::new(); + headers.insert("x-datafusion-distributed-config", "abc".parse().unwrap()); + headers.append("x-repeated", "one".parse().unwrap()); + headers.append("x-repeated", "two".parse().unwrap()); + + let frame = SetPlanFrame::from_parts(set_plan.clone(), &headers).expect("from_parts"); + let header = MppFrameHeader::set_plan(3, 1, 0); + let mut buf = Vec::new(); + encode_prost_frame_into(header, &frame, &mut buf).expect("encode set plan"); + let (parsed, body) = decode_frame(&buf).expect("decode set plan"); + assert_eq!(parsed, header); + assert_eq!(parsed.kind().unwrap(), MppFrameKind::SetPlan); + let FrameBody::SetPlan(decoded) = body else { + panic!("SetPlan frame must carry a SetPlanFrame"); + }; + let (decoded_plan, decoded_headers) = decoded.into_parts().expect("into_parts"); + assert_eq!(decoded_plan, set_plan); + assert_eq!(decoded_headers, headers); + } + + fn sample_set_plan_frame(plan_proto: Vec) -> SetPlanFrame { + SetPlanFrame { + set_plan: Some(pb::SetPlanRequest { + plan_proto, + task_count: 1, + task_key: None, + work_unit_feed_declarations: vec![], + target_worker_url: String::new(), + query_start_time_ns: 0, + }), + header_keys: vec![], + header_values: vec![], + } + } + + #[tokio::test] + async fn set_plan_serves_taker_in_either_arrival_order() { + let drain = DrainHandle::cooperative(1, vec![]); + + // Frame first: the take resolves from the pending slot. + drain.route_set_plan(7, 0, sample_set_plan_frame(vec![1])); + let frame = drain.take_set_plan(7, 0).await.expect("pending take"); + assert_eq!(frame.set_plan.unwrap().plan_proto, vec![1]); + + // Taker first: the frame fulfills the parked oneshot. + let take = drain.take_set_plan(7, 1); + futures::pin_mut!(take); + assert!(futures::poll!(take.as_mut()).is_pending()); + drain.route_set_plan(7, 1, sample_set_plan_frame(vec![2])); + let frame = take.await.expect("waiting take"); + assert_eq!(frame.set_plan.unwrap().plan_proto, vec![2]); + } + + #[tokio::test] + async fn set_plan_take_fails_when_the_inbox_dies() { + let drain = DrainHandle::cooperative(1, vec![]); + let take = drain.take_set_plan(9, 0); + futures::pin_mut!(take); + assert!(futures::poll!(take.as_mut()).is_pending()); + drain.fail_scope(ReceiverScope::Inbox, "producer went away"); + let err = take.await.expect_err("dead inbox must fail the take"); + assert!(format!("{err}").contains("producer went away")); + // Later takers fail immediately. + let err = drain.take_set_plan(9, 1).await.expect_err("dead registry"); + assert!(format!("{err}").contains("producer went away")); + } + + #[test] + fn frame_rejects_short_message() { + let too_short = vec![0u8; MPP_FRAME_HEADER_SIZE - 1]; + let err = decode_frame(&too_short).expect_err("short frame must fail"); + assert!(format!("{err}").contains("too short")); + } + + #[test] + fn frame_rejects_bad_magic() { + // Explicit non-zero, non-magic prefix. Don't rely on the + // happenstance that 0u32 != MPP_FRAME_MAGIC. + let mut bad = vec![0u8; MPP_FRAME_HEADER_SIZE]; + bad[0..4].copy_from_slice(&0xCAFEBABE_u32.to_le_bytes()); + let err = decode_frame(&bad).expect_err("bad magic must fail"); + assert!(format!("{err}").contains("bad frame magic")); + bad[0..4].copy_from_slice(&0xDEADBEEF_u32.to_le_bytes()); + let err = decode_frame(&bad).expect_err("bad magic must fail"); + assert!(format!("{err}").contains("bad frame magic")); + } + + #[test] + fn frame_rejects_unknown_kind() { + let header = MppFrameHeader { + magic: MPP_FRAME_MAGIC, + flags: 0x42, // unknown kind byte, no reserved bits set + stage_id: 0, + partition: 0, + }; + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + let err = decode_frame(&buf).expect_err("unknown kind must fail"); + assert!(format!("{err}").contains("unknown frame kind")); + } + + #[test] + fn frame_rejects_reserved_flag_bits() { + // Reserved range is bits 8..16. Bits 16..32 are sender_proc and must NOT trip the + // reserved check. Cover both boundaries of the reserved range. + for bit in [0x0000_0100u32, 0x0000_8000u32] { + let header = MppFrameHeader { + magic: MPP_FRAME_MAGIC, + flags: bit, // kind byte 0 (Batch), reserved bit set, no sender_proc + stage_id: 0, + partition: 0, + }; + let mut buf = vec![0u8; MPP_FRAME_HEADER_SIZE]; + header.write_to(&mut buf); + let err = decode_frame(&buf).expect_err(&format!("reserved bit {bit:#x} must fail")); + assert!( + format!("{err}").contains("reserved frame flag bits"), + "bit {bit:#x}: {err}" + ); + } + } + + #[test] + fn frame_kind_coexists_with_max_sender_proc() { + // Negative-space companion to frame_rejects_reserved_flag_bits: setting every bit in + // 16..32 (= max sender_proc) along with kind=Eof in bit 0 must parse cleanly without + // tripping the reserved-bits check, and sender_proc()/kind() must read both back. + let header = MppFrameHeader { + magic: MPP_FRAME_MAGIC, + flags: 0xFFFF_0001, // Eof in low byte, max sender_proc in high half, reserved=0 + stage_id: 0, + partition: 0, + }; + assert_eq!(header.kind().unwrap(), MppFrameKind::Eof); + assert_eq!(header.sender_proc(), MPP_MAX_SENDER_PROC); + } + + #[test] + fn frame_sender_proc_round_trip() { + // sender_proc lives in flags bits 16..32 and shouldn't collide with kind or reserved. + for &sp in &[0u32, 1, 7, 255, 256, 1023, 65534, MPP_MAX_SENDER_PROC] { + let header = MppFrameHeader::batch(11, 5, sp); + assert_eq!(header.sender_proc(), sp, "batch round-trip sp={sp}"); + assert_eq!(header.kind().unwrap(), MppFrameKind::Batch); + + let mut buf = Vec::with_capacity(MPP_FRAME_HEADER_SIZE); + let payload = sample_batch(8); + encode_frame_into(header, &payload, &mut buf).expect("encode batch"); + let (parsed, _) = decode_frame(&buf).expect("decode batch"); + assert_eq!(parsed.sender_proc(), sp, "decoded batch sender_proc"); + + let mut eof_buf = Vec::new(); + encode_eof_frame_into(11, 5, sp, &mut eof_buf).expect("encode eof"); + let (parsed_eof, _) = decode_frame(&eof_buf).expect("decode eof"); + assert_eq!(parsed_eof.sender_proc(), sp, "decoded eof sender_proc"); + assert_eq!(parsed_eof.kind().unwrap(), MppFrameKind::Eof); + } + } + + #[test] + fn frame_eof_with_payload_is_rejected() { + let mut buf = Vec::with_capacity(32); + encode_eof_frame_into(0, 0, 0, &mut buf).expect("encode_eof"); + buf.push(0xAB); // smuggle a payload byte after the Eof header + let err = decode_frame(&buf).expect_err("Eof+payload must fail"); + assert!(format!("{err}").contains("payload-less frame carries payload")); + } + + #[test] + fn codec_round_trips_many_batch_sizes() { + let mut buf = Vec::with_capacity(1024); + for rows in [0, 1, 7, 64, 1024] { + let orig = sample_batch(rows); + encode_frame_into(MppFrameHeader::batch(0, 0, 0), &orig, &mut buf).expect("encode"); + let (_header, body) = decode_frame(&buf).expect("decode"); + let FrameBody::Batch(decoded) = body else { + panic!("Batch frame must carry a batch payload"); + }; + assert_eq!(orig.num_rows(), decoded.num_rows()); + } + } + + #[test] + fn drain_buffer_pop_returns_pushed_batches_in_order() { + let buf = DrainBuffer::new(1); + buf.push_batch(sample_batch(3)); + buf.push_batch(sample_batch(5)); + buf.notify_source_done(); + + match buf.pop_front() { + DrainItem::Batch(b) => assert_eq!(b.num_rows(), 3), + DrainItem::Eof => panic!("expected batch"), + DrainItem::Failed(msg) => panic!("unexpected failure: {msg}"), + } + match buf.pop_front() { + DrainItem::Batch(b) => assert_eq!(b.num_rows(), 5), + DrainItem::Eof => panic!("expected batch"), + DrainItem::Failed(msg) => panic!("unexpected failure: {msg}"), + } + matches!(buf.pop_front(), DrainItem::Eof); + } + + #[test] + fn drain_buffer_pop_blocks_until_push_then_eof() { + let buf = DrainBuffer::new(2); + let producer = StdArc::clone(&buf); + let handle = thread::spawn(move || { + thread::sleep(Duration::from_millis(20)); + producer.push_batch(sample_batch(2)); + producer.notify_source_done(); + thread::sleep(Duration::from_millis(20)); + producer.notify_source_done(); + }); + + match buf.pop_front() { + DrainItem::Batch(b) => assert_eq!(b.num_rows(), 2), + DrainItem::Eof => panic!("expected batch first"), + DrainItem::Failed(msg) => panic!("unexpected failure: {msg}"), + } + assert!(matches!(buf.pop_front(), DrainItem::Eof)); + handle.join().unwrap(); + } + + #[test] + fn drain_buffer_cancel_unblocks_waiter() { + let buf = DrainBuffer::new(1); + let canceller = StdArc::clone(&buf); + let handle = thread::spawn(move || { + thread::sleep(Duration::from_millis(20)); + canceller.cancel(); + }); + assert!(matches!(buf.pop_front(), DrainItem::Eof)); + handle.join().unwrap(); + } + + #[test] + fn in_proc_channel_round_trips_through_mpp_sender_receiver() { + let (tx, rx) = in_proc_channel(8); + let sender = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + + sender.send_batch(&sample_batch(4)).unwrap(); + std::mem::drop(sender); + + match receiver.try_recv_batch() { + RecvBatchOutcome::Batch { header: _, batch } => assert_eq!(batch.num_rows(), 4), + other => panic!("expected batch, got {other:?}"), + } + assert!(matches!( + receiver.try_recv_batch(), + RecvBatchOutcome::Detached + )); + } + + #[test] + fn drain_thread_drains_single_source() { + let (tx, rx) = in_proc_channel(4); + let sender = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + let buffer = DrainBuffer::new(1); + + let join = spawn_drain_thread(DrainConfig::new(vec![receiver], StdArc::clone(&buffer))); + + thread::spawn(move || { + for rows in [1, 2, 3, 4, 5] { + sender.send_batch(&sample_batch(rows)).unwrap(); + } + // Drop sender to signal EOF + }) + .join() + .unwrap(); + + let mut received = Vec::new(); + while let DrainItem::Batch(b) = buffer.pop_front() { + received.push(b.num_rows()); + } + assert_eq!(received, vec![1, 2, 3, 4, 5]); + join.join().unwrap().unwrap(); + } + + #[test] + fn drain_handle_shutdown_joins_cleanly() { + let (tx, rx) = in_proc_channel(4); + let sender = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + let buffer = DrainBuffer::new(1); + let handle = + ThreadedDrainHandle::spawn(DrainConfig::new(vec![receiver], StdArc::clone(&buffer))); + + sender.send_batch(&sample_batch(2)).unwrap(); + std::mem::drop(sender); // detach + // Pop the one batch + assert!(matches!(buffer.pop_front(), DrainItem::Batch(_))); + assert!(matches!(buffer.pop_front(), DrainItem::Eof)); + // Drop drives production teardown (cancel + join). Test passes if + // this returns without hanging. + std::mem::drop(handle); + } + + #[test] + fn drain_handle_drop_cancels_and_joins() { + // Build a drain that never detaches (we keep the sender alive), then + // drop the handle. The Drop impl must cancel the buffer and join the + // thread without hanging. + let (tx, rx) = in_proc_channel(4); + let _sender_kept_alive = MppSender::new(Arc::new(tx)); + let receiver = MppReceiver::new(Box::new(rx)); + let buffer = DrainBuffer::new(1); + let handle = + ThreadedDrainHandle::spawn(DrainConfig::new(vec![receiver], StdArc::clone(&buffer))); + + // Simulate consumer path error: drop the handle without calling + // shutdown(). The drain thread must exit before drop returns. + let start = Instant::now(); + drop(handle); + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_secs(2), + "ThreadedDrainHandle::drop took too long: {elapsed:?}" + ); + // Consumer observes EOF because cancel was called. + assert!(matches!(buffer.pop_front(), DrainItem::Eof)); + } + + #[test] + fn drain_thread_drains_n2_mesh_100k_batches() { + // Simulates a 2-proc mesh under load. Each of two producers + // pushes 50_000 small batches through a bounded channel; the drain + // thread interleaves and the consumer reads EOF exactly after + // receiving all 100_000 batches. Exercises backpressure (bounded + // capacity = 16) without deadlock. + const PER_SOURCE: usize = 50_000; + let (tx0, rx0) = in_proc_channel(16); + let (tx1, rx1) = in_proc_channel(16); + let receivers = vec![ + MppReceiver::new(Box::new(rx0)), + MppReceiver::new(Box::new(rx1)), + ]; + let buffer = DrainBuffer::new(2); + let drain_join = spawn_drain_thread(DrainConfig::new(receivers, StdArc::clone(&buffer))); + + let tx0_send = MppSender::new(Arc::new(tx0)); + let tx1_send = MppSender::new(Arc::new(tx1)); + let batch_template = sample_batch(1); + + let p0 = { + let b = batch_template.clone(); + thread::spawn(move || { + for _ in 0..PER_SOURCE { + tx0_send.send_batch(&b).unwrap(); + } + }) + }; + let p1 = { + let b = batch_template.clone(); + thread::spawn(move || { + for _ in 0..PER_SOURCE { + tx1_send.send_batch(&b).unwrap(); + } + }) + }; + + let mut total = 0usize; + while let DrainItem::Batch(_) = buffer.pop_front() { + total += 1; + } + assert_eq!(total, 2 * PER_SOURCE); + p0.join().unwrap(); + p1.join().unwrap(); + drain_join.join().unwrap().unwrap(); + } + + #[test] + fn drain_buffer_drains_buffered_before_eof() { + // Even if all sources have finished and cancel fires, any already- + // buffered batches must be observed before Eof. + let buf = DrainBuffer::new(1); + buf.push_batch(sample_batch(1)); + buf.push_batch(sample_batch(1)); + buf.notify_source_done(); + buf.cancel(); + + assert!(matches!(buf.pop_front(), DrainItem::Batch(_))); + assert!(matches!(buf.pop_front(), DrainItem::Batch(_))); + assert!(matches!(buf.pop_front(), DrainItem::Eof)); + } + + // --------------------------------------------------------------------- + // Throughput microbenches. + // + // These are `#[ignore]` by default because they spin for seconds and spam stdout. Run with: + // + // cargo test --package pg_search --release \ + // postgres::customscan::mpp::transport::tests::throughput \ + // -- --ignored --nocapture + // + // They help us bound the transport layer's cost independently of DataFusion/Tantivy. All use + // the `in_proc_channel` backend (same `MppSender`/`MppReceiver` trait boundary as the shm_mq + // one), so numbers here are an optimistic ceiling. shm_mq adds the ring-buffer copy + + // cross-process notification cost on top. If these numbers are already below the row rate + // the real query needs, we know IPC encode + channel handoff is the bottleneck without + // needing CI data. + // --------------------------------------------------------------------- + + /// Row shape matching the post-Partial shuffle in + /// `aggregate_join_groupby`: a grouping key (title string) plus two + /// partial-aggregate accumulators (COUNT u64, SUM i64). + fn postagg_shape_batch(rows: usize) -> RecordBatch { + let schema = StdArc::new(Schema::new(vec![ + Field::new("title", DataType::Utf8, false), + Field::new("count_partial", DataType::UInt64, false), + Field::new("sum_partial", DataType::Int64, false), + ])); + // Titles averaging ~30 bytes, typical for the docs dataset. + let titles = StringArray::from_iter_values( + (0..rows).map(|i| format!("file_{i:012}_title_with_some_length")), + ); + let counts = UInt64Array::from_iter_values((0..rows as u64).map(|i| i % 64 + 1)); + let sums = Int64Array::from_iter_values((0..rows as i64).map(|i| i * 1024)); + RecordBatch::try_new( + schema, + vec![StdArc::new(titles), StdArc::new(counts), StdArc::new(sums)], + ) + .unwrap() + } + + /// Row shape matching the probe-side shuffle in the same query: + /// `pages.fileId` (u64) plus `pages.sizeInBytes` (i64). + fn probe_shape_batch(rows: usize) -> RecordBatch { + let schema = StdArc::new(Schema::new(vec![ + Field::new("fileId", DataType::UInt64, false), + Field::new("sizeInBytes", DataType::Int64, false), + ])); + let ids = + UInt64Array::from_iter_values((0..rows as u64).map(|i| i.wrapping_mul(2654435761))); + let sizes = Int64Array::from_iter_values((0..rows as i64).map(|i| i * 37)); + RecordBatch::try_new(schema, vec![StdArc::new(ids), StdArc::new(sizes)]).unwrap() + } + + fn bench_throughput( + label: &str, + make_batch: fn(usize) -> RecordBatch, + batch_rows: usize, + total_rows: usize, + ) { + let batches = total_rows.div_ceil(batch_rows); + let template = make_batch(batch_rows); + // Encode once up front so we also report pure-encode throughput + // separately. Real queries encode inside the hot path per batch. + let enc_start = Instant::now(); + let mut enc_bytes = 0usize; + let mut enc_buf = Vec::with_capacity(1024); + for _ in 0..batches { + encode_frame_into(MppFrameHeader::batch(0, 0, 0), &template, &mut enc_buf) + .expect("encode"); + enc_bytes += enc_buf.len(); + } + let enc_elapsed = enc_start.elapsed(); + + // N=2 mesh: two senders, one drain thread, one consumer. Matches + // the gb_postagg / gb_right topology in the real query. + let (tx0, rx0) = in_proc_channel(16); + let (tx1, rx1) = in_proc_channel(16); + let receivers = vec![ + MppReceiver::new(Box::new(rx0)), + MppReceiver::new(Box::new(rx1)), + ]; + let buffer = DrainBuffer::new(2); + let drain_join = spawn_drain_thread(DrainConfig::new(receivers, StdArc::clone(&buffer))); + let tx0_send = MppSender::new(Arc::new(tx0)); + let tx1_send = MppSender::new(Arc::new(tx1)); + + let per_source = batches / 2; + let round_trip_start = Instant::now(); + let p0 = { + let b = template.clone(); + thread::spawn(move || { + for _ in 0..per_source { + tx0_send.send_batch(&b).unwrap(); + } + }) + }; + let p1 = { + let b = template.clone(); + thread::spawn(move || { + for _ in 0..per_source { + tx1_send.send_batch(&b).unwrap(); + } + }) + }; + + let mut got_rows = 0usize; + let mut got_batches = 0usize; + while let DrainItem::Batch(b) = buffer.pop_front() { + got_rows += b.num_rows(); + got_batches += 1; + } + p0.join().unwrap(); + p1.join().unwrap(); + drain_join.join().unwrap().unwrap(); + let rt_elapsed = round_trip_start.elapsed(); + + let enc_mb_per_s = (enc_bytes as f64 / (1024.0 * 1024.0)) / enc_elapsed.as_secs_f64(); + let enc_rows_per_s = (batches * batch_rows) as f64 / enc_elapsed.as_secs_f64(); + let rt_rows_per_s = got_rows as f64 / rt_elapsed.as_secs_f64(); + let rt_bytes_total_mb = enc_bytes as f64 / (1024.0 * 1024.0); + let rt_mb_per_s = rt_bytes_total_mb / rt_elapsed.as_secs_f64(); + let per_batch_us = rt_elapsed.as_micros() as f64 / got_batches as f64; + + println!( + "[throughput] {label:<18} batch_rows={batch_rows:<5} batches={got_batches:<6} rows={got_rows} \ + encode_only: {enc_rows_per_s:>11.0} rows/s {enc_mb_per_s:>7.1} MB/s | \ + round_trip: {rt_rows_per_s:>11.0} rows/s {rt_mb_per_s:>7.1} MB/s ({per_batch_us:.1}us/batch)" + ); + } + + #[test] + #[ignore] + fn throughput_postagg_shape() { + // Sweeps batch size to show per-batch fixed cost vs per-row cost. + // 1.25M total rows ≈ what one proc ships through gb_postagg at + // 25M scale. 625K per proc × 2 = 1.25M. + for batch_rows in [128, 512, 2048, 8192, 32_768] { + bench_throughput("postagg", postagg_shape_batch, batch_rows, 1_250_000); + } + } + + #[test] + #[ignore] + fn throughput_probe_shape() { + // 12.5M total rows ≈ what one proc ships through gb_right at 25M. + for batch_rows in [128, 512, 2048, 8192, 32_768] { + bench_throughput("probe", probe_shape_batch, batch_rows, 12_500_000); + } + } + + // --------------------------------------------------------------------- + // Per-`(stage_id, partition)` channel buffer registry on the cooperative `DrainHandle`. + // + // Producers stamp `MppFrameHeader::batch(stage_id, partition)` on every outgoing frame, and + // the receiver-side cooperative drain demuxes by header into a channel buffer per + // `(stage_id, partition)`. These tests use the `in_proc_channel` backend to drive + // `try_drain_pass` from the test thread. That mirrors how the production path runs the drain + // inline from `DrainGatherStream::poll_next` on the backend thread. + // --------------------------------------------------------------------- + + /// Drain a `DrainHandle::cooperative` to completion: poll until every receiver returns + /// `Empty`. With the `in_proc_channel` test backend the drain observes `Detached` once the + /// producer drops its sender, so a bounded loop of `try_drain_pass` calls is enough to flush + /// everything the producer wrote. + fn drain_until_detached(handle: &DrainHandle) { + for _ in 0..64 { + handle.try_drain_pass().expect("try_drain_pass"); + } + } + + #[test] + fn drain_handle_demuxes_frames_by_header() { + // One queue carrying two channels: `(0, 0)` and `(0, 1)`. Each + // channel buffer receives only its own batches. Per-channel EOF is out of scope + // here. See `drain_handle_eof_frame_closes_one_channel` for explicit-Eof routing + // and `drain_handle_drop_cancels_registered_channel_buffers` for the + // teardown-EOF contract. + let (tx, rx) = in_proc_channel(8); + let base = MppSender::new(Arc::new(tx)); + let s00 = base.clone_with_header(MppFrameHeader::batch(0, 0, 0)); + let s01 = base.clone_with_header(MppFrameHeader::batch(0, 1, 0)); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + s00.send_batch(&sample_batch(2)).unwrap(); + s01.send_batch(&sample_batch(7)).unwrap(); + s00.send_batch(&sample_batch(3)).unwrap(); + drop(s00); + drop(s01); + drop(base); + + let buf00 = handle.register_channel(0, 0, 0); + let buf01 = handle.register_channel(0, 0, 1); + + drain_until_detached(&handle); + + let mut p0_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf00.try_pop() { + p0_rows.push(b.num_rows()); + } + let mut p1_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf01.try_pop() { + p1_rows.push(b.num_rows()); + } + assert_eq!(p0_rows, vec![2, 3]); + assert_eq!(p1_rows, vec![7]); + } + + #[test] + fn drain_handle_eof_frame_closes_one_channel() { + // An `Eof` frame on `(0, 0)` closes that channel buffer while frames on + // `(0, 1)` continue to flow on the same queue. `Detached` doesn't broadcast a + // registry-wide EOF, so `(0, 1)` surfaces EOF only when the handle's `Drop` + // runs `cancel_channel_buffers`. + let (tx, rx) = in_proc_channel(8); + let tx_arc: Arc = Arc::new(tx); + let s00 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(0, 0, 0)); + let s01 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(0, 1, 0)); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + s00.send_batch(&sample_batch(4)).unwrap(); + let mut eof_buf = Vec::new(); + encode_eof_frame_into(0, 0, 0, &mut eof_buf).unwrap(); + tx_arc.send_bytes(&eof_buf).unwrap(); + s01.send_batch(&sample_batch(6)).unwrap(); + + let buf00 = handle.register_channel(0, 0, 0); + let buf01 = handle.register_channel(0, 0, 1); + + drop(s00); + drop(s01); + drop(tx_arc); + drain_until_detached(&handle); + + match buf00.try_pop() { + Some(DrainItem::Batch(b)) => assert_eq!(b.num_rows(), 4), + other => panic!("expected (0,0) batch, got {other:?}"), + } + assert!(matches!(buf00.try_pop(), Some(DrainItem::Eof))); + + match buf01.try_pop() { + Some(DrainItem::Batch(b)) => assert_eq!(b.num_rows(), 6), + other => panic!("expected (0,1) batch, got {other:?}"), + } + assert!(buf01.try_pop().is_none()); + drop(handle); + assert!(matches!(buf01.try_pop(), Some(DrainItem::Eof))); + } + + #[test] + fn drain_handle_register_channel_is_idempotent() { + // Two calls for the same key return Arcs pointing to the same + // DrainBuffer instance. + let (_tx, rx) = in_proc_channel(8); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + let first = handle.register_channel(0, 2, 3); + let second = handle.register_channel(0, 2, 3); + assert!(Arc::ptr_eq(&first, &second)); + } + + #[test] + fn drain_handle_demuxes_frames_by_stage_id() { + // Same partition (0) for two different stage ids on the same queue. + // The registry's compound key keeps them on separate channel buffers. + let (tx, rx) = in_proc_channel(8); + let tx_arc: Arc = Arc::new(tx); + let s_stage0 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(0, 0, 0)); + let s_stage1 = MppSender::with_header(Arc::clone(&tx_arc), MppFrameHeader::batch(1, 0, 0)); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + s_stage0.send_batch(&sample_batch(2)).unwrap(); + s_stage1.send_batch(&sample_batch(9)).unwrap(); + s_stage0.send_batch(&sample_batch(4)).unwrap(); + drop(s_stage0); + drop(s_stage1); + drop(tx_arc); + + let buf0 = handle.register_channel(0, 0, 0); + let buf1 = handle.register_channel(0, 1, 0); + + drain_until_detached(&handle); + + let mut stage0_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf0.try_pop() { + stage0_rows.push(b.num_rows()); + } + let mut stage1_rows = Vec::new(); + while let Some(DrainItem::Batch(b)) = buf1.try_pop() { + stage1_rows.push(b.num_rows()); + } + assert_eq!(stage0_rows, vec![2, 4]); + assert_eq!(stage1_rows, vec![9]); + } + + #[test] + fn drain_handle_drop_cancels_registered_channel_buffers() { + // Dropping a cooperative DrainHandle must wake any consumer holding an Arc + // from `register_channel`. Otherwise a query error path that tears down the mesh would + // leave a consumer blocked on a buffer that will never see EOF. + let (_tx, rx) = in_proc_channel(8); + let receiver = MppReceiver::new(Box::new(rx)); + let handle = DrainHandle::cooperative(0, vec![(ReceiverScope::Inbox, receiver)]); + + let buf_a = handle.register_channel(0, 0, 0); + let buf_b = handle.register_channel(0, 7, 3); + // No data ever flows; the handle is just dropped. + drop(handle); + + assert!(matches!(buf_a.try_pop(), Some(DrainItem::Eof))); + assert!(matches!(buf_b.try_pop(), Some(DrainItem::Eof))); + } +} From ea63b5e71bd4c8f39faa62c5e77c1e6d2a5db467 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 1 Jul 2026 11:58:29 -0700 Subject: [PATCH 18/20] Routed dispatched plans to workers over the shm mesh. ShmWorkerChannel::coordinator_channel now ships each stage's SetPlanRequest to the worker proc that owns its task, as a SetPlan frame the worker takes with take_set_plan, so the embedder no longer routes plans out of band. --- src/shm/runtime.rs | 90 +++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 82 insertions(+), 8 deletions(-) diff --git a/src/shm/runtime.rs b/src/shm/runtime.rs index 5f538d83..6b616905 100644 --- a/src/shm/runtime.rs +++ b/src/shm/runtime.rs @@ -37,13 +37,14 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use datafusion::arrow::array::RecordBatch; -use datafusion::common::{DataFusionError, Result}; +use datafusion::common::{DataFusionError, Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; use futures::stream::{self, BoxStream, StreamExt}; use http::HeaderMap; use url::Url; +use crate::common::serialize_uuid; use crate::proto as pb; use crate::work_unit_feed::RemoteWorkUnitFeedTxs; use crate::{ @@ -52,7 +53,10 @@ use crate::{ }; use super::AliveFlag; -use super::transport::{CooperativeDrainSet, DrainHandle, DrainItem, Interrupt, MppSender}; +use super::transport::{ + CooperativeDrainSet, DrainHandle, DrainItem, Interrupt, MppFrameHeader, MppSender, + SendBatchStats, SetPlanFrame, +}; /// A proc's outbound senders to each peer inbox, shared between the mesh (for `Cancel` frames) and /// the embedder that owns their lifetime. Indexed by destination `proc_idx`; this proc's own slot @@ -143,6 +147,37 @@ impl MppMesh { *self.cancel_senders.lock().unwrap() = Some(senders); } + /// Route a leader-built `SetPlan` frame to the worker proc that owns `task_number`, over the + /// leader's outbound senders (installed via [`Self::set_cancel_senders`]). This is the unified + /// dispatch path: the coordinator ships each stage's plan through here, so the embedder no + /// longer routes it out of band. Errors if the senders aren't installed or the destination + /// slot is empty; the caller treats a failure as best-effort (the worker starves on its + /// `SetPlan` wait, which the dispatch deadlock detector surfaces). + pub async fn send_set_plan( + self: &Arc, + stage_id: u32, + task_number: u32, + frame: SetPlanFrame, + ) -> Result<()> { + let dest_proc = proc_for_task(self.n_workers(), task_number); + let sender = { + let guard = self.cancel_senders.lock().unwrap(); + let Some(senders) = guard.as_ref() else { + return internal_err!( + "shm mesh: control senders not installed; cannot route SetPlan" + ); + }; + let senders = senders.lock().unwrap(); + let Some(Some(base)) = senders.get(dest_proc as usize) else { + return internal_err!("shm mesh: no control sender for proc {dest_proc}"); + }; + base.clone_with_header(MppFrameHeader::set_plan(stage_id, task_number, 0)) + .with_cooperative_drain(Arc::clone(self) as Arc) + }; + let mut stats = SendBatchStats::default(); + sender.send_set_plan_traced(&frame, &mut stats).await + } + /// Tell the producer on `producer_proc` to stop the `(stage_id, partition)` stream: this proc's /// consumer of it stopped reading before EOF. Ships one `Cancel` frame, leaving the rings /// healthy for metrics and every other stream. A no-op when no senders are installed (the @@ -281,17 +316,56 @@ struct ShmWorkerChannel { #[async_trait] impl WorkerChannel for ShmWorkerChannel { - /// pg_search delivers each worker's plan over DSM, not over this channel, so plan delivery is a - /// no-op (what the old `NoOpDispatch` did). Drain the inbound control stream to exhaustion so the - /// coordinator's keep-alive tail does not wedge, then complete with an empty output stream: task - /// metrics travel back over the mesh, not here. + /// Route each `SetPlanRequest` the coordinator ships to the worker proc that owns its task, as a + /// `SetPlan` frame the worker picks up with `take_set_plan`. The coordinator's keep-alive tail + /// emits nothing after the request, so the loop then blocks until the stream closes, draining it + /// to exhaustion. Completes with an empty output stream: task metrics travel back over the mesh, + /// not here. async fn coordinator_channel( &mut self, - _headers: HeaderMap, + headers: HeaderMap, mut c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>, ) -> Result>> { + let mesh = Arc::clone(&self.mesh); #[allow(clippy::disallowed_methods)] - tokio::spawn(async move { while c2w_stream.next().await.is_some() {} }); + tokio::spawn(async move { + while let Some(msg) = c2w_stream.next().await { + let CoordinatorToWorkerMsg::SetPlanRequest(req) = msg else { + continue; + }; + let stage_id = req.task_key.stage_id as u32; + let task_number = req.task_key.task_number as u32; + let set_plan = pb::SetPlanRequest { + task_key: Some(pb::TaskKey { + query_id: serialize_uuid(&req.task_key.query_id), + stage_id: req.task_key.stage_id as u64, + task_number: req.task_key.task_number as u64, + }), + task_count: req.task_count as u64, + plan_proto: req.plan_proto, + work_unit_feed_declarations: req + .work_unit_feed_declarations + .into_iter() + .map(|d| pb::set_plan_request::WorkUnitFeedDeclaration { + id: serialize_uuid(&d.id), + partitions: d.partitions as u64, + }) + .collect(), + target_worker_url: req.target_worker_url.to_string(), + query_start_time_ns: req.query_start_time_ns as u64, + }; + let frame = match SetPlanFrame::from_parts(set_plan, &headers) { + Ok(frame) => frame, + Err(e) => { + log::warn!("shm coordinator_channel: SetPlan frame build failed: {e}"); + continue; + } + }; + if let Err(e) = mesh.send_set_plan(stage_id, task_number, frame).await { + log::warn!("shm coordinator_channel: SetPlan route failed: {e}"); + } + } + }); Ok(stream::empty().boxed()) } From 1f6fb1e3cf6f1c6c4942985db2bee48ee44ee326 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 30 Jun 2026 21:18:03 -0700 Subject: [PATCH 19/20] Added an shm plan codec for the discarded dispatch encode. The coordinator serializes every stage subplan before handing it to the channel, but the shm transport ships plans over DSM and discards that request. `ShmDiscardedPlanCodec` lets the throwaway encode succeed for nodes DataFusion's proto cannot represent, such as `SortMergeJoinExec`. --- src/shm/mod.rs | 2 ++ src/shm/plan_codec.rs | 53 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 src/shm/plan_codec.rs diff --git a/src/shm/mod.rs b/src/shm/mod.rs index bf5c95e1..509cdc04 100644 --- a/src/shm/mod.rs +++ b/src/shm/mod.rs @@ -44,6 +44,7 @@ use std::sync::atomic::AtomicBool; mod dsm; mod mesh; mod mpsc_ring; +mod plan_codec; mod runtime; // Deferred: the self-hosting default transport was built on the removed `WorkerTransport`/ // `WorkerDispatch` dispatch umbrella, which the `ChannelResolver` model has no analog for; its @@ -58,6 +59,7 @@ mod transport; // Curated public surface an embedder consumes. The embedder allocates the shared buffer and // supplies the two extension points (`Wakeup`, `Interrupt`); everything else is built here. pub use mpsc_ring::{NO_RECEIVER_TOKEN, Wakeup}; +pub use plan_codec::ShmDiscardedPlanCodec; pub use runtime::{InProcessWorkerResolver, MppMesh, ShmChannelResolver, proc_for_task}; pub use setup::{ LeaderAttach, WorkerAttach, collect_task_metrics, dsm_region_bytes, install_work_unit_channels, diff --git a/src/shm/plan_codec.rs b/src/shm/plan_codec.rs new file mode 100644 index 00000000..ec39eb4b --- /dev/null +++ b/src/shm/plan_codec.rs @@ -0,0 +1,53 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use datafusion::common::{Result, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; + +/// The embedder registers this via [`crate::DistributedExt::with_distributed_user_codec`] so the +/// coordinator's per-stage `plan_proto` encode always succeeds over the shm transport. +/// +/// The coordinator serializes every stage subplan before handing it to the channel, but the shm +/// transport ships plans over DSM and its `coordinator_channel` drains and discards that request, +/// so the bytes are never read. The encode still runs, and it fails for nodes DataFusion's proto +/// cannot represent (custom scans, `SortMergeJoinExec`). Sitting behind the built-in +/// `DistributedCodec`, this codec is consulted only for those leftover nodes; it encodes each as +/// empty and refuses to decode, since nothing ever decodes the discarded bytes. +#[derive(Debug, Default)] +pub struct ShmDiscardedPlanCodec; + +impl PhysicalExtensionCodec for ShmDiscardedPlanCodec { + fn try_encode(&self, _node: Arc, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + ) -> Result> { + internal_err!( + "ShmDiscardedPlanCodec never decodes: the shm transport discards the coordinator's \ + plan_proto and delivers plans over DSM" + ) + } +} From 741f0c14caee0b74793865483dd49b4ec93b14f8 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 30 Jun 2026 21:18:03 -0700 Subject: [PATCH 20/20] Reported an empty task-level metric set from the shm path. The shm fragment path has no task-level stamps, but the decoder requires the field present; emitting an empty set instead of `None` keeps every worker metrics frame decodable so the per-task EXPLAIN rewrite runs. --- src/shm/setup.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/shm/setup.rs b/src/shm/setup.rs index f86c67bb..9244ad08 100644 --- a/src/shm/setup.rs +++ b/src/shm/setup.rs @@ -220,7 +220,9 @@ pub fn collect_task_metrics( task_count, }, ), - task_metrics: None, + // This path has no task-level stamps, but the decoder requires the field present, so report + // an empty set rather than `None`. + task_metrics: Some(pb::MetricsSet::default()), } }