From 94eec9c0aca69125689b1649ced43c8dc0c869a5 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 15 Jul 2026 19:48:56 +0200 Subject: [PATCH] Add stats estimation accuracy to benchmarks --- Cargo.lock | 1 + benchmarks/Cargo.toml | 1 + benchmarks/cdk/bin/@bench-common.ts | 16 ++++- benchmarks/cdk/bin/@results.ts | 46 ++++++++++++- benchmarks/cdk/bin/datafusion-bench.ts | 13 +++- benchmarks/cdk/bin/worker.rs | 6 ++ benchmarks/src/datasets/common.rs | 13 ++-- benchmarks/src/lib.rs | 1 + benchmarks/src/results.rs | 71 ++++++++++++++++---- benchmarks/src/run.rs | 72 ++++++++++++-------- benchmarks/src/stats.rs | 93 ++++++++++++++++++++++++++ src/worker/worker_connection_pool.rs | 21 +++++- 12 files changed, 298 insertions(+), 56 deletions(-) create mode 100644 benchmarks/src/stats.rs diff --git a/Cargo.lock b/Cargo.lock index 7044adaa..dbe6096c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2252,6 +2252,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "sketches-ddsketch", "structopt", "sysinfo", "tokio", diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 0a03ead0..531a8aba 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -32,6 +32,7 @@ aws-config = "1" aws-sdk-ec2 = "1" openssl = { version = "0.10", features = ["vendored"] } # Keep this. Necessary for the remote benchmarks worker. mimalloc = "0.1" +sketches-ddsketch = "0.3" [dev-dependencies] criterion = "0.5" diff --git a/benchmarks/cdk/bin/@bench-common.ts b/benchmarks/cdk/bin/@bench-common.ts index 014987c3..5f4b59ac 100644 --- a/benchmarks/cdk/bin/@bench-common.ts +++ b/benchmarks/cdk/bin/@bench-common.ts @@ -68,6 +68,8 @@ export interface ExecuteQueryResult { plan: string elapsed: number tasks: number + statsQErrorP50?: number + statsQErrorP95?: number } export interface BenchmarkRunner { @@ -193,11 +195,19 @@ export async function runBenchmark( rowCount: response.rowCount, plan: response.plan, tasks: response.tasks, + statsQErrorP50: response.statsQErrorP50, + statsQErrorP95: response.statsQErrorP95, }) - console.log( - `Query ${id} iteration ${i} took ${Math.round(response.elapsed)} ms and returned ${response.rowCount} rows` - ); + if (response.statsQErrorP50 !== undefined && response.statsQErrorP95 !== undefined) { + console.log( + `Query ${id} iteration ${i} took ${Math.round(response.elapsed)} ms, stats q-error P50 ${response.statsQErrorP50.toFixed(2)}x, P95 ${response.statsQErrorP95.toFixed(2)}x and returned ${response.rowCount} rows` + ); + } else { + console.log( + `Query ${id} iteration ${i} took ${Math.round(response.elapsed)} ms and returned ${response.rowCount} rows` + ); + } } console.log(`Query ${id} p50 time: ${result.p50()} ms`); diff --git a/benchmarks/cdk/bin/@results.ts b/benchmarks/cdk/bin/@results.ts index 9b39d654..542eb161 100644 --- a/benchmarks/cdk/bin/@results.ts +++ b/benchmarks/cdk/bin/@results.ts @@ -12,6 +12,8 @@ export interface QueryIter { rowCount: number; elapsed: number; // Duration in milliseconds tasks: number; + statsQErrorP50?: number; + statsQErrorP95?: number; error?: string; } @@ -80,6 +82,10 @@ export class BenchmarkRun { console.log(`=== Comparing ${this.dataset} results from engine '${other.engine}' [prev] with '${this.engine}' [new] ===`); let totalTimePrev = 0 let totalTimeNew = 0 + const statsQErrorP50Prev: number[] = [] + const statsQErrorP50New: number[] = [] + const statsQErrorP95Prev: number[] = [] + const statsQErrorP95New: number[] = [] for (const query of this.results) { const prevQuery = other.results.find(v => v.id === query.id); if (!prevQuery) { @@ -87,9 +93,21 @@ export class BenchmarkRun { } const timePrev = prevQuery.representativeTime() const timeNew = query.representativeTime() - if (timePrev && timeNew) { + if (timePrev !== undefined && timeNew !== undefined) { totalTimePrev += timePrev totalTimeNew += timeNew + statsQErrorP50Prev.push(...prevQuery.iterations.flatMap(iter => + iter.statsQErrorP50 === undefined ? [] : [iter.statsQErrorP50] + )) + statsQErrorP50New.push(...query.iterations.flatMap(iter => + iter.statsQErrorP50 === undefined ? [] : [iter.statsQErrorP50] + )) + statsQErrorP95Prev.push(...prevQuery.iterations.flatMap(iter => + iter.statsQErrorP95 === undefined ? [] : [iter.statsQErrorP95] + )) + statsQErrorP95New.push(...query.iterations.flatMap(iter => + iter.statsQErrorP95 === undefined ? [] : [iter.statsQErrorP95] + )) } query.compare(prevQuery); @@ -108,6 +126,9 @@ export class BenchmarkRun { console.log( `${"TOTAL".padStart(8)}: prev=${totalTimePrev.toString()} ms, new=${totalTimeNew.toString()} ms, diff=${f.toFixed(2)} ${tag} ${emoji}` ); + + printQErrorComparison("QERR P50", statsQErrorP50Prev, statsQErrorP50New) + printQErrorComparison("QERR P95", statsQErrorP95Prev, statsQErrorP95New) } compareWithPrevious(): void { @@ -238,6 +259,8 @@ export class BenchResult { error: z.string().optional(), plan: z.string(), tasks: z.number().default(0), + statsQErrorP50: z.number().optional(), + statsQErrorP95: z.number().optional(), }).array(), }) const data = fs.readFileSync(filePath, 'utf-8'); @@ -286,3 +309,24 @@ export class BenchResult { function numericId(queryName: string): number { return parseInt([...queryName.matchAll(/(\d+)/g)][0][0]) } + +function printQErrorComparison(label: string, prev: number[], next: number[]): void { + const prevValue = median(prev) + const nextValue = median(next) + if (prevValue !== undefined && nextValue !== undefined) { + console.log(`${label.padStart(8)}: prev=${prevValue.toFixed(2)}x, new=${nextValue.toFixed(2)}x`) + } else if (prevValue !== undefined) { + console.log(`${label.padStart(8)}: prev=${prevValue.toFixed(2)}x, new=n/a`) + } else if (nextValue !== undefined) { + console.log(`${label.padStart(8)}: prev=n/a, new=${nextValue.toFixed(2)}x`) + } +} + +function median(values: number[]): number | undefined { + if (values.length === 0) { + return undefined + } + const sorted = [...values].sort((a, b) => a - b) + const mid = Math.floor(sorted.length / 2) + return sorted.length % 2 ? sorted[mid] : (sorted[mid - 1] + sorted[mid]) / 2 +} diff --git a/benchmarks/cdk/bin/datafusion-bench.ts b/benchmarks/cdk/bin/datafusion-bench.ts index f32af829..67d54449 100644 --- a/benchmarks/cdk/bin/datafusion-bench.ts +++ b/benchmarks/cdk/bin/datafusion-bench.ts @@ -87,7 +87,9 @@ const QueryResponse = z.object({ count: z.number(), plan: z.string(), elapsed_ms: z.number(), - tasks: z.number() + tasks: z.number(), + stats_q_error_p50: z.number().nullable(), + stats_q_error_p95: z.number().nullable() }) type QueryResponse = z.infer @@ -142,7 +144,14 @@ class DataFusionRunner implements BenchmarkRunner { response = await this.query(sql) } - return { rowCount: response.count, plan: response.plan, elapsed: response.elapsed_ms, tasks: response.tasks }; + return { + rowCount: response.count, + plan: response.plan, + elapsed: response.elapsed_ms, + tasks: response.tasks, + statsQErrorP50: response.stats_q_error_p50 ?? undefined, + statsQErrorP95: response.stats_q_error_p95 ?? undefined + }; } private async query(sql: string): Promise { diff --git a/benchmarks/cdk/bin/worker.rs b/benchmarks/cdk/bin/worker.rs index 54d0e8ea..04c08fa0 100644 --- a/benchmarks/cdk/bin/worker.rs +++ b/benchmarks/cdk/bin/worker.rs @@ -20,6 +20,7 @@ use datafusion_distributed::{ get_distributed_channel_resolver, get_distributed_worker_resolver, rewrite_distributed_plan_with_metrics, }; +use datafusion_distributed_benchmarks::stats::stats_estimation_q_error; use futures::{StreamExt, TryFutureExt}; use log::{error, info, warn}; use object_store::aws::AmazonS3Builder; @@ -48,6 +49,8 @@ struct QueryResult { count: usize, elapsed_ms: f64, tasks: usize, + stats_q_error_p50: Option, + stats_q_error_p95: Option, } #[derive(Serialize)] @@ -208,6 +211,7 @@ async fn main() -> Result<(), Box> { ) .await .map_err(err)?; + let stats_q_error = stats_estimation_q_error(&physical); let plan = display_plan_ascii(physical.as_ref(), true); drop(task); @@ -233,6 +237,8 @@ async fn main() -> Result<(), Box> { plan, elapsed_ms: ms, tasks: task_count, + stats_q_error_p50: stats_q_error.map(|q_error| q_error.p50), + stats_q_error_p95: stats_q_error.map(|q_error| q_error.p95), })) } .inspect_err(|(_, msg)| { diff --git a/benchmarks/src/datasets/common.rs b/benchmarks/src/datasets/common.rs index b461be11..d8cc57a7 100644 --- a/benchmarks/src/datasets/common.rs +++ b/benchmarks/src/datasets/common.rs @@ -77,12 +77,13 @@ pub async fn register_tables( let path = entry?.path(); if path.is_dir() { let table_name = path.file_name().unwrap().to_str().unwrap(); - ctx.register_parquet( - table_name, - path.to_str().unwrap(), - ParquetReadOptions::default(), - ) - .await?; + let _ = ctx + .register_parquet( + table_name, + path.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await; } } Ok(()) diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 169cd1f1..93b3378d 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -1 +1,2 @@ pub mod datasets; +pub mod stats; diff --git a/benchmarks/src/results.rs b/benchmarks/src/results.rs index 59da88ea..1302308f 100644 --- a/benchmarks/src/results.rs +++ b/benchmarks/src/results.rs @@ -1,6 +1,7 @@ use crate::{DATA_PATH, RESULTS_DIR}; use datafusion::common::utils::get_available_parallelism; use datafusion::common::{Result, internal_datafusion_err}; +use datafusion_distributed_benchmarks::stats::median; use serde::ser::SerializeSeq; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use std::fs; @@ -13,6 +14,12 @@ use std::time::{Duration, SystemTime}; pub struct QueryIter { pub row_count: usize, pub n_tasks: usize, + /// P50 q-error of the byte estimates made at dynamic stage boundaries. + /// `None` when dynamic planning was disabled or no sampled boundary was present. + pub stats_q_error_p50: Option, + /// P95 q-error of the byte estimates made at dynamic stage boundaries. + /// `None` when dynamic planning was disabled or no sampled boundary was present. + pub stats_q_error_p95: Option, #[serde( serialize_with = "serialize_elapsed", deserialize_with = "deserialize_elapsed" @@ -279,6 +286,10 @@ impl BenchResult { pub fn print_comparison_total(base: &[BenchResult], new: &[BenchResult]) { let mut total_prev: u128 = 0; let mut total_new: u128 = 0; + let mut stats_q_error_p50_prev = vec![]; + let mut stats_q_error_p50_new = vec![]; + let mut stats_q_error_p95_prev = vec![]; + let mut stats_q_error_p95_new = vec![]; for query in new { let Some(prev) = base.iter().find(|v| v.id == query.id) else { continue; @@ -286,22 +297,56 @@ pub fn print_comparison_total(base: &[BenchResult], new: &[BenchResult]) { if let (Some(p), Some(n)) = (prev.representative_time(), query.representative_time()) { total_prev += p; total_new += n; + stats_q_error_p50_prev.extend( + prev.iterations + .iter() + .filter_map(|iteration| iteration.stats_q_error_p50), + ); + stats_q_error_p50_new.extend( + query + .iterations + .iter() + .filter_map(|iteration| iteration.stats_q_error_p50), + ); + stats_q_error_p95_prev.extend( + prev.iterations + .iter() + .filter_map(|iteration| iteration.stats_q_error_p95), + ); + stats_q_error_p95_new.extend( + query + .iterations + .iter() + .filter_map(|iteration| iteration.stats_q_error_p95), + ); } } - if total_prev == 0 && total_new == 0 { - return; + + if total_prev != 0 || total_new != 0 { + let (f, tag, emoji) = if total_new < total_prev { + let f = total_prev as f64 / total_new as f64; + (f, "faster", if f > 1.2 { "✅" } else { "✔" }) + } else { + let f = total_new as f64 / total_prev.max(1) as f64; + (f, "slower", if f > 1.2 { "❌" } else { "✖" }) + }; + println!( + "{:>8}: prev={total_prev} ms, new={total_new} ms, diff={f:.2} {tag} {emoji}", + "TOTAL" + ); + } + + print_q_error_comparison("QERR P50", stats_q_error_p50_prev, stats_q_error_p50_new); + print_q_error_comparison("QERR P95", stats_q_error_p95_prev, stats_q_error_p95_new); +} + +fn print_q_error_comparison(label: &str, prev: Vec, new: Vec) { + match (median(prev), median(new)) { + (Some(prev), Some(new)) => println!("{label:>8}: prev={prev:.2}x, new={new:.2}x"), + (Some(prev), None) => println!("{label:>8}: prev={prev:.2}x, new=n/a"), + (None, Some(new)) => println!("{label:>8}: prev=n/a, new={new:.2}x"), + (None, None) => {} } - let (f, tag, emoji) = if total_new < total_prev { - let f = total_prev as f64 / total_new as f64; - (f, "faster", if f > 1.2 { "✅" } else { "✔" }) - } else { - let f = total_new as f64 / total_prev.max(1) as f64; - (f, "slower", if f > 1.2 { "❌" } else { "✖" }) - }; - println!( - "{:>8}: prev={total_prev} ms, new={total_new} ms, diff={f:.2} {tag} {emoji}", - "TOTAL" - ); } fn serialize_bench_results( diff --git a/benchmarks/src/run.rs b/benchmarks/src/run.rs index 69aff7a6..62a81644 100644 --- a/benchmarks/src/run.rs +++ b/benchmarks/src/run.rs @@ -25,7 +25,7 @@ use datafusion::common::{config_err, exec_err, not_impl_err}; use datafusion::datasource::source::DataSourceExec; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::SessionStateBuilder; -use datafusion::physical_plan::collect; +use datafusion::physical_plan::{ExecutionPlan, collect}; use datafusion::prelude::*; use datafusion_distributed::test_utils::localhost::LocalHostWorkerResolver; use datafusion_distributed::test_utils::work_unit_file_scan::{ @@ -37,6 +37,7 @@ use datafusion_distributed::{ display_plan_ascii, rewrite_distributed_plan_with_metrics, }; use datafusion_distributed_benchmarks::datasets::{clickbench, register_tables, tpcds, tpch}; +use datafusion_distributed_benchmarks::stats::stats_estimation_q_error; use std::error::Error; use std::fs; use std::path::PathBuf; @@ -88,14 +89,6 @@ pub struct RunOpt { #[structopt(long)] cardinality_task_sf: Option, - /// Use children isolator UNIONs for distributing UNION operations. - #[structopt(long)] - children_isolator_unions: bool, - - /// Turns on broadcast joins. - #[structopt(long = "broadcast-joins")] - broadcast_joins: bool, - /// Collects metrics across network boundaries #[structopt(long)] collect_metrics: bool, @@ -216,9 +209,11 @@ impl RunOpt { "none" => None, v => return config_err!("Unknown compression type {v}"), })? - .with_distributed_children_isolator_unions(self.children_isolator_unions)? - .with_distributed_broadcast_joins(self.broadcast_joins)? - .with_distributed_metrics_collection(self.collect_metrics || self.debug)? + .with_distributed_children_isolator_unions(true)? + .with_distributed_broadcast_joins(true)? + .with_distributed_metrics_collection( + self.collect_metrics || self.debug || self.dynamic, + )? .with_distributed_max_tasks_per_stage(self.max_tasks_per_stage)? .with_distributed_user_codec(WorkUnitFileScanCodec) .with_distributed_task_estimator(WorkUnitFileScanTaskEstimator) @@ -290,18 +285,46 @@ impl RunOpt { } match self.execute_query(ctx, query).await { - Ok((result, n_tasks)) => { + Ok((result, n_tasks, physical_plan)) => { let elapsed = start.elapsed(); let ms = elapsed.as_secs_f64() * 1000.0; let row_count = result.iter().map(|b| b.num_rows()).sum(); - println!( - "Query {id} iteration {i} took {ms:.1} ms and returned {row_count} rows" - ); + let physical_plan = if self.dynamic || self.debug { + rewrite_distributed_plan_with_metrics( + physical_plan, + DistributedMetricsFormat::PerTask, + ) + .await? + } else { + physical_plan + }; + let stats_q_error = match self.dynamic { + true => stats_estimation_q_error(&physical_plan), + false => None, + }; + if let Some(q_error) = stats_q_error { + println!( + "Query {id} iteration {i} took {ms:.1} ms, stats q-error P50 {:.2}x, P95 {:.2}x and returned {row_count} rows", + q_error.p50, q_error.p95 + ); + } else { + println!( + "Query {id} iteration {i} took {ms:.1} ms and returned {row_count} rows" + ); + } + if self.debug { + println!( + "=== Physical plan with metrics ===\n{}\n", + display_plan_ascii(physical_plan.as_ref(), true) + ); + } bench_query.iterations.push(QueryIter { elapsed, row_count, n_tasks, + stats_q_error_p50: stats_q_error.map(|q_error| q_error.p50), + stats_q_error_p95: stats_q_error.map(|q_error| q_error.p95), error: None, }); } @@ -311,6 +334,8 @@ impl RunOpt { elapsed: Duration::from_millis(0), row_count: 0, n_tasks: 0, + stats_q_error_p50: None, + stats_q_error_p95: None, error: Some(err.to_string()), }); continue 'outer; @@ -327,7 +352,7 @@ impl RunOpt { &self, ctx: &SessionContext, sql: &str, - ) -> Result<(Vec, usize)> { + ) -> Result<(Vec, usize, Arc)> { let plan = ctx.sql(sql).await?; let (state, plan) = plan.into_parts(); @@ -341,18 +366,7 @@ impl RunOpt { Ok(Transformed::no(node)) })?; let result = collect(physical_plan.clone(), state.task_ctx()).await?; - if self.debug { - let physical_plan = rewrite_distributed_plan_with_metrics( - physical_plan.clone(), - DistributedMetricsFormat::PerTask, - ) - .await?; - println!( - "=== Physical plan with metrics ===\n{}\n", - display_plan_ascii(physical_plan.as_ref(), true) - ); - } - Ok((result, n_tasks)) + Ok((result, n_tasks, physical_plan)) } fn get_path(&self) -> Result { diff --git a/benchmarks/src/stats.rs b/benchmarks/src/stats.rs new file mode 100644 index 00000000..764800ba --- /dev/null +++ b/benchmarks/src/stats.rs @@ -0,0 +1,93 @@ +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion_distributed::{NetworkBoundaryExt, Stage}; +use sketches_ddsketch::{Config, DDSketch}; +use std::sync::Arc; + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct StatsEstimationQError { + pub p50: f64, + pub p95: f64, +} + +/// Computes P50 and P95 q-error between the sampled byte estimate and the actual output bytes for +/// every dynamically sampled stage boundary in `plan`. +pub fn stats_estimation_q_error(plan: &Arc) -> Option { + let mut boundary_q_errors = DDSketch::new(Config::defaults()); + + let _ = plan.apply(|node| { + if let Some(boundary) = node.as_network_boundary() + && let Stage::Local(input_stage) = boundary.input_stage() + && let Some(sampled_bytes) = metric_total(&input_stage.metrics_set, "sampled_bytes") + && let Some(actual_bytes) = node + .metrics() + .and_then(|metrics| metric_total(&metrics, "output_bytes")) + { + boundary_q_errors.add(q_error(sampled_bytes, actual_bytes)); + } + Ok(TreeNodeRecursion::Continue) + }); + + q_error_percentiles(&boundary_q_errors) +} + +fn q_error_percentiles(q_errors: &DDSketch) -> Option { + Some(StatsEstimationQError { + p50: q_errors.quantile(0.50).ok().flatten()?, + p95: q_errors.quantile(0.95).ok().flatten()?, + }) +} + +fn metric_total(metrics: &MetricsSet, name: &str) -> Option { + metrics + .sum(|metric| metric.value().name() == name) + .map(|value| value.as_usize()) +} + +/// Q-error is the standard cardinality-estimation metric because it treats equal-factor over- and +/// underestimates symmetrically. See https://www.vldb.org/pvldb/vol2/vldb09-657.pdf and +/// https://vldb.org/pvldb/vol9/p204-leis.pdf. +fn q_error(estimated: usize, actual: usize) -> f64 { + let estimated = estimated.max(1) as f64; + let actual = actual.max(1) as f64; + (estimated / actual).max(actual / estimated) +} + +pub fn median(mut values: Vec) -> Option { + if values.is_empty() { + return None; + } + values.sort_unstable_by(f64::total_cmp); + let mid = values.len() / 2; + Some(if values.len().is_multiple_of(2) { + (values[mid - 1] + values[mid]) / 2.0 + } else { + values[mid] + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn q_error_percentiles_returns_none_for_an_empty_sketch() { + assert_eq!( + q_error_percentiles(&DDSketch::new(Config::defaults())), + None + ); + } + + #[test] + fn q_error_percentiles_reports_regular_and_tail_cases() { + let mut sketch = DDSketch::new(Config::defaults()); + for value in 1..=100 { + sketch.add(value as f64); + } + + let percentiles = q_error_percentiles(&sketch).unwrap(); + assert!((49.0..=51.0).contains(&percentiles.p50)); + assert!((94.0..=96.0).contains(&percentiles.p95)); + } +} diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 4ccba411..ee6e66b7 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -13,7 +13,7 @@ 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, StreamExt, TryFutureExt}; +use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt}; use std::fmt::{Debug, Formatter}; use std::ops::Range; use std::sync::{Arc, Mutex, OnceLock}; @@ -80,6 +80,7 @@ impl WorkerConnectionPool { stage_id: input_stage.num, task_number: target_task, }; + let output_bytes = MetricBuilder::new(&self.metrics).output_bytes(target_partition); // If we are physically in the same worker, short circuit into a local connection without // going through the `WorkerChannel`. @@ -90,7 +91,11 @@ impl WorkerConnectionPool { &producer_head, ctx, )? { - return Ok(result); + return Ok(result + .inspect_ok(move |batch| { + output_bytes.add(logical_record_batch_size(batch)); + }) + .boxed()); } // Otherwise, we need to reach the remote worker through the `WorkerChannel`. Unlike local @@ -156,6 +161,9 @@ impl WorkerConnectionPool { }) } .try_flatten_stream() + .inspect_ok(move |batch| { + output_bytes.add(logical_record_batch_size(batch)); + }) .boxed()) } @@ -216,6 +224,15 @@ impl WorkerConnectionPool { } } +/// Returns the logical size of a batch's slices, excluding unused backing-buffer capacity. +fn logical_record_batch_size(batch: &RecordBatch) -> usize { + batch + .columns() + .iter() + .map(|column| column.to_data().get_slice_memory_size().unwrap_or(0)) + .sum() +} + impl Debug for WorkerConnectionPool { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("WorkerConnections")