avantgardnerio commented on code in PR #2094: URL: https://github.com/apache/datafusion-ballista/pull/2094#discussion_r3614456395
########## ballista/core/src/execution_plans/runtime_stats.rs: ########## @@ -0,0 +1,687 @@ +// 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. + +//! Passthrough operator that accumulates runtime statistics on the data +//! flowing through it. Two accessor families: +//! +//! - **Row count** — always tracked, per input partition. Written via +//! `AtomicUsize::fetch_add` on the hot path (no cross-partition +//! contention, no lock overhead). +//! - **Quantile sketch** — optional (only when `order_by` is set at +//! construction). Per-partition `Mutex<TDigest>` keeps writes off any +//! shared lock. +//! +//! Timing is decoupled from correctness: both accessors are readable at +//! any point. Callers get whatever has flowed through so far — a +//! mid-stream snapshot for callers that decide the sample is accurate +//! enough, a post-drain snapshot for callers that want the full state +//! (typical after a blocking downstream like `SortExec`). +//! +//! The `order_by` field accepts the full `Vec<PhysicalSortExpr>` so +//! multi-key `ORDER BY` survives serde (tie-breakers get preserved for +//! downstream `SortExec` / `BoundedWindowAggExec` even though only the +//! first key drives the sketch today). +//! +//! TODO: swap `TDigest` for a generic-over-`Ord` KLL sketch. TDigest is +//! `Float64`-only and single-column; a KLL implementation would sketch +//! the full `Vec<PhysicalSortExpr>` (composite keys, non-numeric types) +//! and let the operator drop the "first expression, `Float64` only" +//! restriction on the routing key. +//! +//! This PR lands the tap in isolation: nothing wires it into a plan yet, +//! and the executor doesn't yet ship the accumulated state back to the +//! scheduler. Those pieces arrive with the range-repartition operator +//! (which is the first consumer). + +use std::fmt::{self, Debug, Formatter}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use datafusion::arrow::array::Float64Array; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalSortExpr; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use datafusion_functions_aggregate_common::tdigest::TDigest; +use futures::stream::StreamExt; +use log::info; + +/// T-Digest centroid budget. 100 is DataFusion's default and gives ~1% +/// quantile error, plenty of margin over the sub-partition counts we +/// expect at bin-pack time. +const TDIGEST_MAX_SIZE: usize = 100; + +/// Streaming runtime-stats operator. See module-level docs. +pub struct RuntimeStatsExec { + input: Arc<dyn ExecutionPlan>, + /// Lexicographic ORDER BY carried through from the wrapping window + /// operator when the caller wants quantile sketching. `None` when + /// only row counting is needed. + order_by: Option<Vec<PhysicalSortExpr>>, + /// Always tracked, per input partition. Written via + /// `AtomicUsize::fetch_add` on the hot path — no cross-partition + /// contention, no lock overhead. + row_counts: Arc<[AtomicUsize]>, + /// Only allocated when `order_by` is `Some`. Sketches over the first + /// ORDER BY expression's `Float64` values. `Mutex`-per-partition to + /// keep writes off any shared lock. + sketches: Option<Arc<[Mutex<TDigest>]>>, + properties: Arc<PlanProperties>, +} + +impl RuntimeStatsExec { + /// Wrap `input`. If `order_by` is provided, its first entry drives + /// the per-partition T-Digest; the full slice is preserved for serde + /// and for downstream operators (`SortExec`, `BoundedWindowAggExec`) + /// that need it. When `Some`, at least one expression is required — + /// nothing to sketch on with an empty slice. + pub fn try_new( + input: Arc<dyn ExecutionPlan>, + order_by: Option<Vec<PhysicalSortExpr>>, + ) -> Result<Self> { + if let Some(exprs) = &order_by { + let [_first, ..] = exprs.as_slice() else { + return internal_err!( + "RuntimeStatsExec: order_by is Some but empty; pass None to skip sketching" + ); + }; + } + let partition_count = input.output_partitioning().partition_count(); + let row_counts: Arc<[AtomicUsize]> = (0..partition_count) + .map(|_| AtomicUsize::new(0)) + .collect::<Vec<_>>() + .into(); + let sketches: Option<Arc<[Mutex<TDigest>]>> = order_by.as_ref().map(|_| { + (0..partition_count) + .map(|_| Mutex::new(TDigest::new(TDIGEST_MAX_SIZE))) + .collect::<Vec<_>>() + .into() + }); + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Ok(Self { + input, + order_by, + row_counts, + sketches, + properties, + }) + } + + /// Full ORDER BY carried through, or `None` if the operator was + /// built in row-count-only mode. + pub fn order_by(&self) -> Option<&[PhysicalSortExpr]> { + self.order_by.as_deref() + } + + /// Rows observed on `partition` so far. Cheap `Relaxed` load — the + /// value is a running counter, monotonically non-decreasing. + /// + /// Errors on out-of-range partition. Callers pass a partition id + /// they've already used with `execute`. + pub fn row_count(&self, partition: usize) -> Result<usize> { + let counter = self.row_counts.get(partition).ok_or_else(|| { + internal_datafusion_err!( + "RuntimeStatsExec: partition {} out of range (have {})", + partition, + self.row_counts.len() + ) + })?; + Ok(counter.load(Ordering::Relaxed)) + } + + /// Number of partition slots this operator was built with (matches + /// its input's declared partition count). Every slot has its own + /// row counter and — in sketch mode — its own sketch; a given task + /// only fills the slot(s) it actually executes. + pub fn partition_count(&self) -> usize { + self.row_counts.len() + } + + /// Rows observed across all partitions so far. + pub fn total_row_count(&self) -> usize { + self.row_counts + .iter() + .map(|c| c.load(Ordering::Relaxed)) + .sum() + } + + /// Snapshot of one partition's running quantile sketch. Returns + /// `None` when the operator was built in row-count-only mode (no + /// `order_by`). Cheap clone (a `Vec<Centroid>` of size + /// ≤ `TDIGEST_MAX_SIZE`). + /// + /// Errors if `partition` ≥ input's partition count — callers pass a + /// partition id they've already used with `execute`. + pub fn quantile_sketch(&self, partition: usize) -> Result<Option<TDigest>> { + let Some(sketches) = &self.sketches else { + return Ok(None); + }; + let slot = sketches.get(partition).ok_or_else(|| { + internal_datafusion_err!( + "RuntimeStatsExec: partition {} out of range (have {})", + partition, + sketches.len() + ) + })?; + Ok(Some( + slot.lock() + .expect("RuntimeStats sketch mutex poisoned") + .clone(), + )) + } + + /// All partitions merged into one sketch. `None` in row-count-only + /// mode. + pub fn merged_quantile_sketch(&self) -> Option<TDigest> { + let sketches = self.sketches.as_ref()?; + let snapshots: Vec<TDigest> = sketches + .iter() + .map(|m| { + m.lock() + .expect("RuntimeStats sketch mutex poisoned") + .clone() + }) + .collect(); + Some(TDigest::merge_digests(snapshots.iter())) + } +} + +impl Debug for RuntimeStatsExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("RuntimeStatsExec") + .field("order_by", &self.order_by) + .finish() + } +} + +impl DisplayAs for RuntimeStatsExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + match &self.order_by { + Some(exprs) => { + let routing = &exprs[0]; + write!( + f, + "RuntimeStatsExec: rows + sketch(routing={} {})", + routing.expr, + if routing.options.descending { + "desc" + } else { + "asc" + } + ) + } + None => write!(f, "RuntimeStatsExec: rows"), + } + } +} + +impl ExecutionPlan for RuntimeStatsExec { Review Comment: I implemented all 6 of the algebraic ones. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
