diff --git a/benchmarks/cdk/bin/datafusion-bench.ts b/benchmarks/cdk/bin/datafusion-bench.ts index 67d54449..44d414d3 100644 --- a/benchmarks/cdk/bin/datafusion-bench.ts +++ b/benchmarks/cdk/bin/datafusion-bench.ts @@ -25,7 +25,7 @@ async function main() { .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('--dynamic-bytes-per-partition ', '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') @@ -49,7 +49,7 @@ async function main() { 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 dynamicBytesPerPartition = parseInt(options.dynamicBytesPerPartition) const debug = options.debug === true || options.debug === 'true' || options.debug === 1 const warmup = options.warmup === true || options.warmup === 'true' || options.warmup === 1 @@ -64,7 +64,7 @@ async function main() { broadcastJoins, partialReduce, dynamicTaskCount, - bytesPerPartitionPerSecond, + dynamicBytesPerPartition, maxTasksPerStage, repartitionFileMinSize, targetPartitions @@ -107,7 +107,7 @@ class DataFusionRunner implements BenchmarkRunner { broadcastJoins: boolean; partialReduce: boolean; dynamicTaskCount: boolean; - bytesPerPartitionPerSecond: number; + dynamicBytesPerPartition: number; maxTasksPerStage: number; repartitionFileMinSize: number; targetPartitions: number; @@ -195,7 +195,7 @@ class DataFusionRunner implements BenchmarkRunner { 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.dynamic_bytes_per_partition=${this.options.dynamicBytesPerPartition}; 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/src/common/recursion.rs b/src/common/recursion.rs index 5c16116e..91506d8f 100644 --- a/src/common/recursion.rs +++ b/src/common/recursion.rs @@ -3,6 +3,7 @@ use crate::{DistributedTaskContext, NetworkBoundaryExt}; use datafusion::common::Result; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion}; use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::joins::{CrossJoinExec, HashJoinExec, NestedLoopJoinExec}; use std::cell::RefCell; use std::sync::Arc; @@ -26,12 +27,41 @@ pub(crate) trait TreeNodeExt { f: F, ) -> Result; + /// Applies `f` top-down (pre-order) to the nodes on the plan's *driver path*: the operators + /// whose ongoing data production feeds, batch for batch, into the root's output. The + /// [`TreeNodeRecursion`] returned by `f` steers it as in [`TreeNode::apply`]. + /// + /// It differs from a full traversal only at the pipeline-breaking joins ([`HashJoinExec`], + /// [`NestedLoopJoinExec`], [`CrossJoinExec`]): it follows the probe side (`right`) and skips the + /// build side (`left`), since a join fully materializes its build side before emitting any + /// output row, so build-side rows are setup work rather than output progress. Every other + /// operator is transparent. + /// + /// Used to estimate a running stage's completion: the total rows to pull come from the leaves + /// on the driver path (see `estimated_driver_path_leaf_rows`). + /// + /// ```text + /// ┌────────────┐ + /// │ HashJoin │ visited + /// └─────┬──────┘ + /// build │ probe + /// (left) │ │ (right) + /// ┌────────┘ └────────┐ + /// ▼ ▼ + /// ┌────────┐ ┌────────┐ + /// │ Scan B │ SKIPPED │ Scan P │ visited + /// └────────┘ └────────┘ + /// ``` + fn apply_driver_path Result>( + &self, + f: F, + ) -> Result; + /// Recursively rewrite the tree using `f` in a top-down (pre-order) fashion, propagating /// the appropriate [DistributedTaskContext] based on the presence of nodes that can isolate /// tasks, like [ChildrenIsolatorUnionExec]. /// /// `f` is applied to the node first, and then its children. - #[allow(dead_code)] // Used in follow up work. fn transform_down_with_dt_ctx< F: FnMut(Self, DistributedTaskContext) -> Result>, >( @@ -62,7 +92,6 @@ pub(crate) trait TreeNodeExt { /// count. /// /// `f` is applied to the node first, and then its children. - #[allow(dead_code)] // Used in follow up work. fn transform_down_with_task_count Result>>( self, task_count: usize, @@ -104,6 +133,32 @@ impl TreeNodeExt for Arc { recurse(self, ctx, &mut f) } + fn apply_driver_path Result>( + &self, + mut f: F, + ) -> Result { + fn recurse) -> Result>( + plan: &Arc, + f: &mut F, + ) -> Result { + f(plan)?.visit_children(|| { + if let Some(hash_join) = plan.downcast_ref::() { + recurse(hash_join.right(), f) + } else if let Some(nested_loop_join) = plan.downcast_ref::() { + recurse(nested_loop_join.right(), f) + } else if let Some(cross_join) = plan.downcast_ref::() { + recurse(cross_join.right(), f) + } else { + plan.children() + .into_iter() + .apply_until_stop(|child| recurse(child, f)) + } + }) + } + + recurse(self, &mut f) + } + fn transform_down_with_dt_ctx< F: FnMut(Self, DistributedTaskContext) -> Result>, >( @@ -213,9 +268,13 @@ mod tests { use crate::execution_plans::ChildWeight; use crate::stage::RemoteStage; use crate::{NetworkCoalesceExec, Stage}; - use datafusion::arrow::datatypes::Schema; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{JoinType, NullEquality}; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::joins::PartitionMode; use datafusion::physical_plan::union::UnionExec; use insta::assert_snapshot; // ── apply_with_dt_ctx ──────────────────────────────────────────────────────── @@ -348,6 +407,81 @@ mod tests { "); } + // ── apply_driver_path ──────────────────────────────────────────────────── + + #[test] + fn driver_path_leaf() { + let plan = leaf(); + assert_snapshot!(trace_driver_path(&plan), @"Leaf"); + } + + #[test] + fn driver_path_top_down_order() { + let plan = union(vec![single(leaf()), leaf()]); + assert_snapshot!(trace_driver_path(&plan), @r" + Union + Single + Leaf + Leaf + "); + } + + #[test] + fn driver_path_jump_skips_subtree() { + let child = single(leaf()); + let plan = single(Arc::clone(&child)); + assert_snapshot!( + trace_driver_path_with(&plan, |p| { + if Arc::ptr_eq(p, &child) { TreeNodeRecursion::Jump } else { TreeNodeRecursion::Continue } + }), + @r" + Single + Single [->jump] + "); + } + + #[test] + fn driver_path_hash_join_uses_probe_side() { + let plan = hash_join( + single(leaf_i32("l")), + union(vec![leaf_i32("r"), leaf_i32("r")]), + ); + assert_snapshot!(trace_driver_path(&plan), @r" + HashJoin + Union + Leaf + Leaf + "); + } + + #[test] + fn driver_path_nested_loop_join_uses_probe_side() { + let plan = nested_loop_join( + single(leaf_i32("l")), + union(vec![leaf_i32("r"), leaf_i32("r")]), + ); + assert_snapshot!(trace_driver_path(&plan), @r" + NestedLoopJoin + Union + Leaf + Leaf + "); + } + + #[test] + fn driver_path_cross_join_uses_probe_side() { + let plan = cross_join( + single(leaf_i32("l")), + union(vec![leaf_i32("r"), leaf_i32("r")]), + ); + assert_snapshot!(trace_driver_path(&plan), @r" + CrossJoin + Union + Leaf + Leaf + "); + } + // ── transform_down_with_dt_ctx ──────────────────────────────────────────── #[test] @@ -566,6 +700,14 @@ mod tests { Arc::new(EmptyExec::new(Arc::new(Schema::empty()))) } + fn leaf_i32(name: &str) -> Arc { + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + name, + DataType::Int32, + true, + )])))) + } + fn single(child: Arc) -> Arc { Arc::new(CoalescePartitionsExec::new(child)) } @@ -581,6 +723,47 @@ mod tests { Arc::new(NetworkCoalesceExec::try_new(input, producer_tasks, 1).unwrap()) } + fn hash_join( + left: Arc, + right: Arc, + ) -> Arc { + Arc::new( + HashJoinExec::try_new( + left, + right, + join_on(), + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) + } + + fn nested_loop_join( + left: Arc, + right: Arc, + ) -> Arc { + Arc::new(NestedLoopJoinExec::try_new(left, right, None, &JoinType::Inner, None).unwrap()) + } + + fn cross_join( + left: Arc, + right: Arc, + ) -> Arc { + Arc::new(CrossJoinExec::new(left, right)) + } + + fn join_on() -> Vec<(Arc, Arc)> { + vec![( + Arc::new(Column::new("l", 0)) as Arc, + Arc::new(Column::new("r", 0)) as Arc, + )] + } + fn remote_network_boundary() -> Arc { network_boundary(leaf(), 1) .as_network_boundary() @@ -630,6 +813,12 @@ mod tests { "CIU" } else if p.is::() { "Network" + } else if p.is::() { + "HashJoin" + } else if p.is::() { + "NestedLoopJoin" + } else if p.is::() { + "CrossJoin" } else { "?" } @@ -664,6 +853,29 @@ mod tests { lines.join("\n") } + fn trace_driver_path(root: &Arc) -> String { + trace_driver_path_with(root, |_| TreeNodeRecursion::Continue) + } + + fn trace_driver_path_with) -> TreeNodeRecursion>( + root: &Arc, + mut decide: F, + ) -> String { + let mut lines = vec![]; + root.apply_driver_path(|p| { + let rec = decide(p); + let suffix = match rec { + TreeNodeRecursion::Continue => "", + TreeNodeRecursion::Jump => " [->jump]", + TreeNodeRecursion::Stop => " [->stop]", + }; + lines.push(format!("{}{suffix}", plan_label(p))); + Ok(rec) + }) + .unwrap(); + lines.join("\n") + } + fn trace_dt_ctx_down(root: Arc, dt_ctx: DistributedTaskContext) -> String { trace_dt_ctx_down_with(root, dt_ctx, |_| TreeNodeRecursion::Continue) } diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index dfc062b2..11f46a74 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -8,7 +8,9 @@ use crate::distributed_planner::{ }; use crate::execution_plans::SamplerExec; use crate::stage::{LocalStage, RemoteStage}; -use crate::{BytesCounterMetric, LoadInfo, NetworkBoundaryExt, NetworkCoalesceExec, Stage}; +use crate::{ + BytesCounterMetric, LoadInfo, MaxGaugeMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage, +}; use dashmap::DashMap; use datafusion::common::stats::Precision; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -60,7 +62,7 @@ pub(super) async fn prepare_dynamic_plan( .cpu .get_value() .unwrap_or(&0) - .div_ceil(nb_ctx.d_cfg.bytes_per_partition_per_second.max(1)) + .div_ceil(nb_ctx.d_cfg.dynamic_bytes_per_partition.max(1)) .div_ceil(input_stage.plan.output_partitioning().partition_count()) .clamp(1, nb_ctx.max_tasks()?); let task_count = nb_ctx @@ -111,12 +113,14 @@ pub(super) async fn prepare_dynamic_plan( 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 (stats, new_metrics) = + 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, )); + metrics.extend(new_metrics); // 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. @@ -205,11 +209,17 @@ impl PlanReconstructor { async fn gather_runtime_statistics( per_task_load_info_stream: Vec + Unpin>, plan: &Arc, -) -> Result { - const ESTIMATED_QUERY_TIME_S: usize = 10; +) -> Result<(Statistics, MetricsSet)> { + /// How many LoadInfo samples with non-zero sampled bytes should be gathered before considering + /// that there's enough samples collected to make a final decision. The higher this value is, + /// the more execution is delayed because of sample collection, but better sampling precision. const BYTES_READY_SAMPLE_PERCENTAGE: f32 = 0.2; - const BYTES_PER_SECOND_SAMPLE_PERCENTAGE: f32 = 0.2; + /// If the leaf nodes in the provided `plan` do not provide any row count estimation, there's + /// no way for us here to estimate the % of completion of the stage during sampling. This is + /// the default % of completion for those cases. + const FALLBACK_PCT_COMPLETION: f32 = 0.2; + let mut new_metrics = MetricsSet::new(); let Some(sampler) = find_sampler(plan) else { return plan_err!("Mising SamplerExec while gathering load report"); }; @@ -223,34 +233,28 @@ async fn gather_runtime_statistics( 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 partitions_reached_eos = 0; let mut rows_ready = 0; - let mut rows_per_second = 0; + let mut rows_pulled_from_leafs = 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; 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, - )?; + rows_pulled_from_leafs += load_info.rows_pulled_from_leaf; 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_reached_eos += load_info.reached_eos as usize; partitions_done += 1; // Short circuit if we collected enough bytes_ready measurements. @@ -260,13 +264,6 @@ async fn gather_runtime_statistics( 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; @@ -274,31 +271,45 @@ async fn gather_runtime_statistics( } if partitions_done == 0 { - return Ok(zero_stats(plan.schema().fields.len())); + return Ok((zero_stats(plan.schema().fields.len()), new_metrics)); } 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 rows_pulled_from_leafs = rows_pulled_from_leafs * total_partitions / partitions_done; + + let estimated_pct_sampled = if partitions_reached_eos == partitions_done { + // Every sampled partition's stream reached end-of-stream, so `rows_ready` / + // `per_col_bytes_ready` are the partitions' final output rather than a partial snapshot — + // the stage is fully sampled. This is the reliable "done" signal, and it correctly covers + // legitimately-empty stages (which would otherwise report `rows_pulled_from_leafs == 0` and + // make the completion fraction 0, blowing up the `ready / fraction` extrapolation below). + 1.0 + } else if let Some(estimated_driver_path_leaf_rows) = estimated_driver_path_leaf_rows(plan) { + // The stage is still producing. Estimate how far along it is from the fraction of the + // driver-path leaf rows consumed so far. + (rows_pulled_from_leafs as f32 / estimated_driver_path_leaf_rows as f32).min(1.0) + } else { + // We can't measure progress (no leaf-row estimate, or nothing pulled from the leaves + // yet even though we're not at EOS): fall back rather than dividing by ~0. + FALLBACK_PCT_COMPLETION + }; + + new_metrics.push(MaxGaugeMetric::new_metric( + "estimated_pct_sampled", + (estimated_pct_sampled * 100.) as usize, + )); - let total_num_rows = rows_ready + rows_per_second * ESTIMATED_QUERY_TIME_S; + let total_num_rows = (rows_ready as f32 / estimated_pct_sampled) as usize; if total_num_rows == 0 { - return Ok(zero_stats(n_cols)); + return Ok((zero_stats(n_cols), new_metrics)); } - let per_col_byte_size = element_wise_sum( - per_col_bytes_ready, - &vec_mul(per_col_bytes_per_second, ESTIMATED_QUERY_TIME_S), - )?; + let per_col_byte_size = vec_mul(per_col_bytes_ready, 1. / estimated_pct_sampled); let total_byte_size: usize = per_col_byte_size.iter().sum(); let ndv_pct = vec_avg_reduce(ndv_pct)?; @@ -310,7 +321,7 @@ async fn gather_runtime_statistics( return plan_err!("Expected {n_cols} null values, but got {}", null_pct.len()); } - Ok(Statistics { + let stats = Statistics { num_rows: Precision::Inexact(total_num_rows), total_byte_size: Precision::Inexact(total_byte_size), column_statistics: ndv_pct @@ -326,7 +337,27 @@ async fn gather_runtime_statistics( sum_value: Precision::Absent, }) .collect(), - }) + }; + + Ok((stats, new_metrics)) +} + +fn estimated_driver_path_leaf_rows(plan: &Arc) -> Option { + let mut total_rows = None; + let _ = plan.apply_driver_path(|plan| { + if plan.children().is_empty() { + let stats = plan.partition_statistics(None)?; + if let Some(num_rows) = stats.num_rows.get_value() { + if let Some(total_rows) = &mut total_rows { + *total_rows += *num_rows; + } else { + total_rows = Some(*num_rows); + }; + } + } + Ok(TreeNodeRecursion::Continue) + }); + total_rows } fn find_sampler(plan: &Arc) -> Option<&SamplerExec> { diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index a6f9e032..daa26eaa 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -589,16 +589,16 @@ pub trait DistributedExt: Sized { /// 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( + fn with_distributed_dynamic_bytes_per_partition( self, - bytes_per_partition_per_second: usize, + dynamic_bytes_per_partition: usize, ) -> Result; - /// Same as [DistributedExt::with_distributed_bytes_per_partition_per_second] but with an + /// Same as [DistributedExt::with_distributed_dynamic_bytes_per_partition] but with an /// in-place mutation. - fn set_distributed_bytes_per_partition_per_second( + fn set_distributed_dynamic_bytes_per_partition( &mut self, - bytes_per_partition_per_second: usize, + dynamic_bytes_per_partition: usize, ) -> Result<(), DataFusionError>; } @@ -759,12 +759,12 @@ impl DistributedExt for SessionConfig { Ok(()) } - fn set_distributed_bytes_per_partition_per_second( + fn set_distributed_dynamic_bytes_per_partition( &mut self, - bytes_per_partition_per_second: usize, + dynamic_bytes_per_partition: 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; + d_cfg.dynamic_bytes_per_partition = dynamic_bytes_per_partition; Ok(()) } @@ -856,9 +856,9 @@ impl DistributedExt for SessionConfig { #[expr($?;Ok(self))] fn with_distributed_dynamic_task_count(mut self, enabled: bool) -> Result; - #[call(set_distributed_bytes_per_partition_per_second)] + #[call(set_distributed_dynamic_bytes_per_partition)] #[expr($?;Ok(self))] - fn with_distributed_bytes_per_partition_per_second(mut self, bytes_per_partition_per_second: usize) -> Result; + fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result; } } } @@ -983,10 +983,10 @@ impl DistributedExt for SessionStateBuilder { #[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)] + fn set_distributed_dynamic_bytes_per_partition(&mut self, dynamic_bytes_per_partition: usize) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_bytes_per_partition)] #[expr($?;Ok(self))] - fn with_distributed_bytes_per_partition_per_second(mut self, bytes_per_partition_per_second: usize) -> Result; + fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result; } } } @@ -1113,10 +1113,10 @@ impl DistributedExt for SessionState { #[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)] + fn set_distributed_dynamic_bytes_per_partition(&mut self, dynamic_bytes_per_partition: usize) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_bytes_per_partition)] #[expr($?;Ok(self))] - fn with_distributed_bytes_per_partition_per_second(mut self, bytes_per_partition_per_second: usize) -> Result; + fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result; } } } @@ -1236,10 +1236,10 @@ impl DistributedExt for SessionContext { #[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)] + fn set_distributed_dynamic_bytes_per_partition(&mut self, dynamic_bytes_per_partition: usize) -> Result<(), DataFusionError>; + #[call(set_distributed_dynamic_bytes_per_partition)] #[expr($?;Ok(self))] - fn with_distributed_bytes_per_partition_per_second(self, bytes_per_partition_per_second: usize) -> Result; + fn with_distributed_dynamic_bytes_per_partition(self, dynamic_bytes_per_partition: usize) -> Result; } } } diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index 03dd156b..88b541e2 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -68,9 +68,9 @@ extensions_options! { /// 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 + /// If `dynamic_task_count` is enabled, this value is the amount of bytes each /// partition is expected to handle. Lower values will result in greater parallelism. - pub bytes_per_partition_per_second: usize, default = 16 * 1024 * 1024 + pub dynamic_bytes_per_partition: usize, default = 16 * 1024 * 1024 } } diff --git a/src/execution_plans/sampler.rs b/src/execution_plans/sampler.rs index 836d9986..8eb5c364 100644 --- a/src/execution_plans/sampler.rs +++ b/src/execution_plans/sampler.rs @@ -1,4 +1,4 @@ -use crate::common::{require_one_child, vec_cast}; +use crate::common::{TreeNodeExt, require_one_child, vec_cast}; use crate::{ BytesCounterMetric, BytesMetricExt, GaugeMetricExt, LatencyMetricExt, LoadInfo, MaxGaugeMetric, MaxLatencyMetric, P50LatencyMetric, @@ -15,8 +15,11 @@ 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 datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; +use futures::stream::FusedStream; +use futures::{Stream, StreamExt, TryFutureExt}; use std::collections::VecDeque; use std::fmt::{Debug, Formatter}; use std::pin::Pin; @@ -28,15 +31,7 @@ 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; +const READY_CHUNK_LIMIT: usize = 1024 * 1024; /// Maximum number of rows sampled from the peek queue when estimating per-column NDV. const NDV_MAX_ROWS_SAMPLE: usize = 1000; @@ -69,8 +64,6 @@ pub(crate) struct SamplerExecMetrics { 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. @@ -89,7 +82,6 @@ impl SamplerExecMetrics { 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(); @@ -226,10 +218,10 @@ impl PartitionSampler { let first_batch_at = Arc::clone(&self.first_batch_at); let n_cols = self.input.schema().fields.len(); - let reporter = LoadInfoDropHandler { + let mut reporter = LoadInfoDropHandler { + plan: Arc::clone(&self.input), 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), @@ -255,57 +247,28 @@ impl PartitionSampler { // 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. + let Some(batches) = (&mut input_stream) + .ready_chunks(READY_CHUNK_LIMIT) + .next() + .await + else { + // The input produced nothing at all: the stream is exhausted, so this partition + // reached EOS with zero rows. + reporter.load_info.reached_eos = true; return Ok(peek.chain(input_stream).boxed()); }; - let _elapsed_compute_timer = elapsed_compute.timer(); + let _guard = elapsed_compute.timer(); + + peek.reserve(batches.len()); + first_batch_at.get_or_init(Instant::now); 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); + // If the input stream is already terminated, every batch it will ever produce is now in + // `peek`, so this partition has reached EOS and its `rows_ready` / `bytes_ready` are + // final rather than a partial snapshot. + reporter.load_info.reached_eos = input_stream.is_terminated(); + reporter.report(&peek); Ok(peek.chain(input_stream).boxed()) }); @@ -330,16 +293,16 @@ impl PartitionSampler { /// Emitting on drop ensures that it's always emitted. struct LoadInfoDropHandler { omit: Arc, + plan: Arc, load_info: 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) { + fn report(mut self, peek: &RecordBatchPeek) { if self.omit.load(Ordering::Relaxed) { return; } @@ -348,8 +311,6 @@ impl LoadInfoDropHandler { 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) { @@ -357,30 +318,10 @@ impl LoadInfoDropHandler { 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 = &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) / total_ready) - .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; } - fn set_rows_per_second(&mut self, rows_per_second: usize) { - self.load_info.rows_per_second = rows_per_second; - } - fn set_per_col_ndv(&mut self, per_column_ndv: Vec) { self.load_info.per_column_ndv_percentage = per_column_ndv; } @@ -388,6 +329,32 @@ impl LoadInfoDropHandler { fn set_per_col_null_pct(&mut self, per_column_null_pct: Vec) { self.load_info.per_column_null_percentage = per_column_null_pct; } + + fn set_rows_pulled_from_leaf(&mut self) { + let mut rows_pulled_from_leaf = 0; + let _ = self.plan.apply_driver_path(|plan| { + if plan.children().is_empty() + && let Some(metrics) = plan.metrics() + && let Some(output_rows) = metrics.output_rows() + { + let partition_count = self.plan.output_partitioning().partition_count(); + // Here, we need to divide by `partition_count` because the `output_rows` returned + // by the leaf node metrics is not per-partition, is global to the whole node. + // + // In a perfect scenario, we might be able to map the current SamplerExec partition + // in scope to the specific partition of the leaf node, but several things can + // happen in between these two nodes that mangle the partitions. For example, the + // presence of a RepartitionExec or an InterleaveExec will mess with partition + // mapping from SamplerExec to leaf node. + // + // Just dividing by `partition_count` gives a good approximation, but it's still + // not perfect as it will not account for data skews. + rows_pulled_from_leaf += output_rows / partition_count; + } + Ok(TreeNodeRecursion::Continue) + }); + self.load_info.rows_pulled_from_leaf = rows_pulled_from_leaf; + } } impl Drop for LoadInfoDropHandler { @@ -395,6 +362,7 @@ impl Drop for LoadInfoDropHandler { if self.omit.load(Ordering::Relaxed) { return; } + self.set_rows_pulled_from_leaf(); 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()); @@ -405,12 +373,12 @@ impl Drop for LoadInfoDropHandler { 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], 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], + rows_pulled_from_leaf: 0, + reached_eos: false, } } @@ -424,6 +392,10 @@ struct RecordBatchPeek { } impl RecordBatchPeek { + fn reserve(&mut self, size: usize) { + self.peek.reserve(size) + } + fn push(&mut self, batch: RecordBatch) { let batch_size = record_batch_size(&batch); if self.peek.is_empty() { diff --git a/src/metrics/max_gauge_metric.rs b/src/metrics/max_gauge_metric.rs index 9091204b..cfc77165 100644 --- a/src/metrics/max_gauge_metric.rs +++ b/src/metrics/max_gauge_metric.rs @@ -1,3 +1,4 @@ +use datafusion::physical_plan::Metric; use datafusion::physical_plan::metrics::{CustomMetricValue, MetricBuilder, MetricValue}; use std::sync::atomic::Ordering::Relaxed; use std::{ @@ -40,6 +41,16 @@ impl Default for MaxGaugeMetric { } impl MaxGaugeMetric { + pub fn new_metric(name: impl Into>, value: usize) -> Arc { + Arc::new(Metric::new( + MetricValue::Custom { + name: name.into(), + value: Arc::new(MaxGaugeMetric::from_value(value)), + }, + None, + )) + } + pub fn from_value(bytes: usize) -> Self { Self { value: Arc::new(AtomicUsize::new(bytes)), diff --git a/src/protocol/grpc/generated/worker.rs b/src/protocol/grpc/generated/worker.rs index 290261ed..d1f0fb4a 100644 --- a/src/protocol/grpc/generated/worker.rs +++ b/src/protocol/grpc/generated/worker.rs @@ -68,23 +68,22 @@ pub struct LoadInfo { /// 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, + /// The amount of rows that were pulled from leaf nodes while this partition was sampling data. + #[prost(uint64, tag = "8")] + pub rows_pulled_from_leaf: u64, + /// Whether the sampled partition stream reached end-of-stream by the time this LoadInfo was + /// captured. + #[prost(bool, tag = "9")] + pub reached_eos: bool, } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct GetWorkerInfoRequest {} diff --git a/src/protocol/grpc/worker.proto b/src/protocol/grpc/worker.proto index ef691b59..24f5c987 100644 --- a/src/protocol/grpc/worker.proto +++ b/src/protocol/grpc/worker.proto @@ -60,18 +60,19 @@ message LoadInfo { 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; + reserved 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; + reserved 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; + // The amount of rows that were pulled from leaf nodes while this partition was sampling data. + uint64 rows_pulled_from_leaf = 8; + // Whether the sampled partition stream reached end-of-stream (i.e. the partition finished + // producing all of its output) by the time this LoadInfo was captured. + bool reached_eos = 9; } message GetWorkerInfoRequest {} diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 9cc94de2..56612671 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -510,19 +510,15 @@ 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, + rows_pulled_from_leaf: load_info.rows_pulled_from_leaf as usize, + reached_eos: load_info.reached_eos, } } diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index fa677a10..6fba14ee 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -285,19 +285,15 @@ 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, + rows_pulled_from_leaf: load_info.rows_pulled_from_leaf as u64, + reached_eos: load_info.reached_eos, } } diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 1baa23fe..3ad3816c 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -145,18 +145,20 @@ pub struct LoadInfo { 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, + /// The amount of rows that were pulling from leaf nodes while the partition to which this + /// LoadInfo belongs to was sampling data. Used for estimating how much data is left by + /// comparing this value to the estimated total rows pulled from leaf nodes. + pub rows_pulled_from_leaf: usize, + /// Whether the sampled partition stream reached end-of-stream (i.e. the partition finished + /// producing all of its output) by the time this LoadInfo was captured. When true, `rows_ready` + /// and `per_column_bytes_ready` are final rather than a partial snapshot. + pub reached_eos: bool, } pub struct ExecuteTaskRequest { diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index ee6e66b7..9aa87e03 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -81,6 +81,7 @@ impl WorkerConnectionPool { task_number: target_task, }; let output_bytes = MetricBuilder::new(&self.metrics).output_bytes(target_partition); + let output_rows = MetricBuilder::new(&self.metrics).output_rows(target_partition); // If we are physically in the same worker, short circuit into a local connection without // going through the `WorkerChannel`. @@ -94,6 +95,7 @@ impl WorkerConnectionPool { return Ok(result .inspect_ok(move |batch| { output_bytes.add(logical_record_batch_size(batch)); + output_rows.add(batch.num_rows()); }) .boxed()); } @@ -163,6 +165,7 @@ impl WorkerConnectionPool { .try_flatten_stream() .inspect_ok(move |batch| { output_bytes.add(logical_record_batch_size(batch)); + output_rows.add(batch.num_rows()); }) .boxed()) }