diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 5e91e9825e69d..fec06c90f8a5c 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -97,6 +97,8 @@ def get_minimal_system_parameters( "enable_coalesce_case_transform": "true", "enable_columnation_lgalloc": "false", "enable_compute_correction_v2": "true", + # Off in production; on here so hinted TopK plans exercise kept-row stages. + "enable_compute_topk_retained_stages": "true", "enable_compute_logical_backpressure": "true", "enable_connection_validation_syntax": "true", "enable_create_table_from_source": "true", diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 6a008fc235834..669c18ea6a258 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -3139,6 +3139,9 @@ def __init__( BOOLEAN_FLAG_VALUES ) self.flags_with_values["enable_compute_sync_mv_sink"] = BOOLEAN_FLAG_VALUES + self.flags_with_values["enable_compute_topk_retained_stages"] = ( + BOOLEAN_FLAG_VALUES + ) self.flags_with_values["enable_column_paged_batcher"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_columnar_merge_batcher"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_column_paged_batcher_spill"] = ( diff --git a/src/compute-types/src/dyncfgs.rs b/src/compute-types/src/dyncfgs.rs index 01cf2f5292d9c..9dd8174de2735 100644 --- a/src/compute-types/src/dyncfgs.rs +++ b/src/compute-types/src/dyncfgs.rs @@ -262,6 +262,16 @@ pub const ENABLE_SYNC_MV_SINK: Config = Config::new( ParameterScope::Environment, ); +/// Whether a TopK stage may emit the rows it keeps instead of negations of the rows it drops, +/// when the expected group size says the kept rows are fewer. Read at dataflow rendering time. +pub const ENABLE_COMPUTE_TOPK_RETAINED_STAGES: Config = Config::new( + "enable_compute_topk_retained_stages", + false, + "Whether a TopK stage may emit the rows it keeps instead of negations of the rows it drops, \ + when the expected group size says the kept rows are fewer.", + ParameterScope::Replica, +); + /// Whether rendering should use the new MV sink correction buffer implementation. pub const ENABLE_CORRECTION_V2: Config = Config::new( "enable_compute_correction_v2", @@ -793,6 +803,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&ENABLE_ERROR_DISTINCT) .add(&ENABLE_MZ_JOIN_CORE) .add(&ENABLE_SYNC_MV_SINK) + .add(&ENABLE_COMPUTE_TOPK_RETAINED_STAGES) .add(&ENABLE_CORRECTION_V2) .add(&CORRECTION_V2_CHAIN_PROPORTIONALITY) .add(&CORRECTION_V2_CHUNK_SIZE) diff --git a/src/compute-types/src/plan/top_k.rs b/src/compute-types/src/plan/top_k.rs index 417b0acddc937..4fb1b969e1fc4 100644 --- a/src/compute-types/src/plan/top_k.rs +++ b/src/compute-types/src/plan/top_k.rs @@ -227,5 +227,8 @@ pub struct BasicTopKPlan { /// The number of columns in the input and output. pub arity: usize, /// Bucket sizes for hierarchical stages of TopK. Should be decreasing. + /// + /// Empty when the query's group size hint is at most sixteen, the fan-in of one stage; the + /// renderer reads that as a promise that groups are small. pub buckets: Vec, } diff --git a/src/compute/src/render/top_k.rs b/src/compute/src/render/top_k.rs index 9810115be1bad..224c0e5f310fe 100644 --- a/src/compute/src/render/top_k.rs +++ b/src/compute/src/render/top_k.rs @@ -23,7 +23,10 @@ use differential_dataflow::operators::iterate::Variable as SemigroupVariable; use differential_dataflow::trace::cursor::{BatchCursor, BatchValOwn}; use differential_dataflow::trace::{Builder, Cursor, Navigable, Trace}; use differential_dataflow::{Data, VecCollection}; -use mz_compute_types::dyncfgs::{ENABLE_COMPUTE_TEMPORAL_BUCKETING, TEMPORAL_BUCKETING_SUMMARY}; +use mz_compute_types::dyncfgs::{ + ENABLE_COMPUTE_TEMPORAL_BUCKETING, ENABLE_COMPUTE_TOPK_RETAINED_STAGES, + TEMPORAL_BUCKETING_SUMMARY, +}; use mz_compute_types::plan::ArrangementStrategy; use mz_compute_types::plan::scalar::LirScalarExpr; use mz_compute_types::plan::top_k::{ @@ -268,8 +271,16 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim // intra-ts thinning. The maximum number of records per timestamp is // (num_workers * limit), which we expect to be a small number and so we render // a single topk stage. - let (result, errs) = - self.build_topk_stage(thinned, order_key, 1u64, 0, limit, arity, false); + let (result, errs) = self.build_topk_stage( + thinned, + order_key, + Some(1), + 0, + limit, + arity, + false, + StageOutput::Dropped, + ); // Consolidate the output of `build_topk_stage` because it's not guaranteed to be. let result = CollectionExt::consolidate_named::>( result, @@ -334,18 +345,69 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim ) { let pairer = Pairer::new(1); let mut datum_vec = mz_repr::DatumVec::new(); + // Key each row for the first stage directly: its bucket of the row hash, or zero when + // the only stage is the final one and the hash would be discarded anyway. + let first_modulus = buckets.first().copied().unwrap_or(1); + // A group key that is the row's leading columns is a byte prefix of the row. + let prefix_len = group_key + .iter() + .enumerate() + .all(|(index, column)| *column == index) + .then_some(group_key.len()); let mut collection = collection.map({ move |row| { - let group_row = { - let row_hash = row.hashed(); - let datums = datum_vec.borrow_with(&row); - let iterator = group_key.iter().map(|i| datums[*i]); - pairer.merge(std::iter::once(Datum::from(row_hash)), iterator) + let bucket = if first_modulus == 1 { + 0 + } else { + row.hashed() % first_modulus + }; + let group_row = match prefix_len { + Some(prefix_len) => { + let (group, _) = row.split_at_datum(prefix_len); + let mut group_row = Row::default(); + let mut packer = group_row.packer(); + packer.push(Datum::from(bucket)); + packer.extend_by_row_ref(group); + group_row + } + None => { + let datums = datum_vec.borrow_with(&row); + let iterator = group_key.iter().map(|i| datums[*i]); + pairer.merge(std::iter::once(Datum::from(bucket)), iterator) + } }; (group_row, row) } }); + // Rows a stage keeps per key: the literal limit plus the offset it must see past. + let kept_per_key = limit + .as_ref() + .and_then(|l| l.as_literal_int64()) + .and_then(|l| u64::try_from(l).ok()) + .map(|l| l.saturating_add(u64::cast_from(offset))); + let retain_allowed = ENABLE_COMPUTE_TOPK_RETAINED_STAGES.get(&self.config_set); + // Whether a stage emits kept rows or negated dropped rows. Kept rows cost less state when + // a key's group holds at least twice what the stage keeps. The plan carries the group + // size hint only as its bucket list: an empty list means the hint was at most sixteen, + // the fan-in of one stage, so the lone final stage keeps rows when twice its kept count + // fits in sixteen. With bucket stages the hint is unknown here (or absent, in which case + // the planner assumed billions), and every stage emits dropped rows, since a kept-row + // stage whose keys hold a single row each would copy its whole input into its output + // arrangement. + let single_stage = buckets.is_empty(); + let stage_output = |_modulus: u64| -> StageOutput { + let Some(kept) = kept_per_key else { + return StageOutput::Dropped; + }; + if retain_allowed && single_stage && kept.saturating_mul(2) <= 16 { + StageOutput::Kept + } else { + StageOutput::Dropped + } + }; + // The input arrives keyed for the first stage; later stages rekey it. + let mut first = true; let mut validating = true; let mut err_collection: Option> = None; @@ -384,12 +446,14 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim let (oks, errs) = self.build_topk_stage( collection, order_key.clone(), - bucket, + (!first).then_some(bucket), 0, Some(limit.clone()), arity, validating, + stage_output(bucket), ); + first = false; collection = oks; if validating { err_collection = errs; @@ -401,13 +465,26 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim // We do a final step, both to make sure that we complete the reduction, and to correctly // apply `offset` to the final group, as we have not yet been applying it to the partially // formed groups. + let final_output = stage_output(1); let (oks, errs) = self.build_topk_stage( - collection, order_key, 1u64, offset, limit, arity, validating, + collection, + order_key, + (!first).then_some(1), + offset, + limit, + arity, + validating, + final_output, ); - // Consolidate the output of `build_topk_stage` because it's not guaranteed to be. - let oks = - CollectionExt::consolidate_named::>(oks, "TopK final consolidate"); - collection = oks; + // A stage that emits its input plus negations leaves cancellations behind; one that + // emits kept rows comes out of a reduce, already consolidated. + collection = match final_output { + StageOutput::Dropped => CollectionExt::consolidate_named::>( + oks, + "TopK final consolidate", + ), + StageOutput::Kept => oks, + }; if validating { err_collection = errs; } @@ -453,32 +530,36 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim &self, collection: VecCollection<'s, T, (Row, Row), Diff>, order_key: Vec, - modulus: u64, + modulus: Option, offset: usize, limit: Option, arity: usize, validating: bool, + output: StageOutput, ) -> ( VecCollection<'s, T, (Row, Row), Diff>, Option>, ) { - // Form appropriate input by updating the `hash` column (first datum in `hash_key`) by - // applying `modulus`. - let input = collection.map(move |(hash_key, row)| { - let mut hash_key_iter = hash_key.iter(); - let hash = hash_key_iter.next().unwrap().unwrap_uint64() % modulus; - let hash_key = SharedRow::pack(std::iter::once(hash.into()).chain(hash_key_iter)); - (hash_key, row) - }); + // Rekey the input by applying `modulus` to the hash column (first datum in `hash_key`), + // unless the input already arrives keyed for this stage. + let input = match modulus { + Some(modulus) => collection.map(move |(hash_key, row)| { + let mut hash_key_iter = hash_key.iter(); + let hash = hash_key_iter.next().unwrap().unwrap_uint64() % modulus; + let hash_key = SharedRow::pack(std::iter::once(hash.into()).chain(hash_key_iter)); + (hash_key, row) + }), + None => collection, + }; // If validating: demux errors, otherwise we cannot produce errors. let (input, oks, errs) = if validating { // Build topk stage, produce errors for invalid multiplicities. - let (input, stage) = build_topk_negated_stage::< + let (input, stage) = build_topk_stage_arrangements::< T, RowValBuilder<_, _, _>, RowValSpine, _, _>, - >(&input, order_key, offset, limit, arity); + >(&input, order_key, offset, limit, arity, output); let stage = stage.as_collection(|k, v| (k.to_row(), v.clone())); // Demux oks and errors. @@ -502,16 +583,21 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim } else { // Build non-validating topk stage. let (input, stage) = - build_topk_negated_stage::, RowRowSpine<_, _>>( - &input, order_key, offset, limit, arity, + build_topk_stage_arrangements::, RowRowSpine<_, _>>( + &input, order_key, offset, limit, arity, output, ); // Turn arrangement into collection. let stage = stage.as_collection(|k, v| (k.to_row(), v.to_row())); (input, stage, None) }; - let input = input.as_collection(|k, v| (k.to_row(), v.to_row())); - (oks.concat(input), errs) + match output { + StageOutput::Kept => (oks, errs), + StageOutput::Dropped => { + let input = input.as_collection(|k, v| (k.to_row(), v.to_row())); + (oks.concat(input), errs) + } + } } fn render_top1_monotonic<'s>( @@ -604,19 +690,29 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim } } -/// Build a stage of a topk reduction. Maintains the _retractions_ of the output instead of emitted -/// rows. This has the benefit that we have to maintain state proportionally to size of the output -/// instead of the size of the input. +/// What a TopK stage's reduce emits. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum StageOutput { + /// The rows the stage keeps. The stage's result is the reduce output alone. + Kept, + /// Negations of the rows the stage drops. The stage's result is its input plus the + /// negations, so the maintained state is proportional to the rows dropped rather than kept, + /// which is smaller when groups are barely larger than the limit. + Dropped, +} + +/// Build a stage of a topk reduction. /// /// Returns two arrangements: /// * The arranged input data without modifications, and -/// * the maintained negated output data. -fn build_topk_negated_stage<'s, T, Bu, Tr>( +/// * the maintained output data, as `output` selects. +fn build_topk_stage_arrangements<'s, T, Bu, Tr>( input: &VecCollection<'s, T, (Row, Row), Diff>, order_key: Vec, offset: usize, limit: Option, arity: usize, + output: StageOutput, ) -> ( Arranged<'s, TraceAgent>>, Arranged<'s, TraceAgent>, @@ -702,16 +798,32 @@ where || limit .map(|l| source.iter().map(|(_, d)| *d).sum::() > l) .unwrap_or(false); - if !must_shrink { - return; - } - - // First go ahead and emit all records. Note that we ensure target - // has the capacity to hold at least these records, and avoid any - // dependencies on the user-provided (potentially unbounded) limit. - target.reserve(source.len()); - for (datums, diff) in source.iter() { - target.push((BatchValOwn::::ok((*datums).to_row()), -diff)); + match output { + StageOutput::Dropped => { + if !must_shrink { + return; + } + // First go ahead and emit all records negated; the kept rows are emitted + // again below and cancel. Note that we ensure target has the capacity to + // hold at least these records, and avoid any dependencies on the + // user-provided (potentially unbounded) limit. + target.reserve(source.len()); + for (datums, diff) in source.iter() { + target.push((BatchValOwn::::ok((*datums).to_row()), -diff)); + } + } + StageOutput::Kept => { + if !must_shrink { + // Every row is kept. + target.reserve(source.len()); + for (datums, diff) in source.iter() { + if diff.is_positive() { + target.push((BatchValOwn::::ok((*datums).to_row()), *diff)); + } + } + return; + } + } } // local copies that may count down to zero. let mut offset = offset; @@ -757,10 +869,9 @@ where diff = std::cmp::min(diff, Diff::from(*limit)); *limit -= diff; } - // Output the indicated number of rows. + // Output the indicated number of rows: the kept rows themselves, or, on + // the dropped path, the positive counterparts that cancel their negations. if diff.is_positive() { - // Emit retractions for the elements actually part of - // the set of TopK elements. target.push((BatchValOwn::::ok(datums.to_row()), diff)); } } diff --git a/src/repr/src/row.rs b/src/repr/src/row.rs index c6dc97f6c7e8d..ef8140c947891 100644 --- a/src/repr/src/row.rs +++ b/src/repr/src/row.rs @@ -753,6 +753,32 @@ impl RowRef { &self.0 } + /// Splits the row after its first `n` datums into two rows: the first `n` datums and the + /// rest. Datums are encoded back to back, so both halves are complete row encodings and no + /// datum is re-encoded. Only the first `n` datums are decoded, to find the boundary. + /// + /// # Panics + /// + /// Panics if the row has fewer than `n` datums. + pub fn split_at_datum(&self, n: usize) -> (&RowRef, &RowRef) { + let mut rest: &[u8] = &self.0; + for _ in 0..n { + assert!(!rest.is_empty(), "row has fewer than {n} datums"); + // SAFETY: `rest` is a suffix of a valid row encoding at a datum boundary. + unsafe { + read_datum(&mut rest); + } + } + let boundary = self.0.len() - rest.len(); + // SAFETY: both halves start and end at datum boundaries of a valid row encoding. + unsafe { + ( + RowRef::from_slice(&self.0[..boundary]), + RowRef::from_slice(&self.0[boundary..]), + ) + } + } + /// True iff there is no data in this [`RowRef`]. pub fn is_empty(&self) -> bool { self.0.is_empty() @@ -4750,4 +4776,31 @@ mod tests { assert_eq!(map_1.cmp(&map_null), Ordering::Less); assert_eq!(map_null.cmp(&map_1), Ordering::Greater); } + + #[mz_ore::test] + fn split_at_datum_yields_whole_rows() { + let datums = [ + Datum::Int64(7), + Datum::Null, + Datum::String("a longer string that is not inline"), + Datum::False, + ]; + let row = Row::pack_slice(&datums); + for n in 0..=datums.len() { + let (head, tail) = row.split_at_datum(n); + assert_eq!(head.iter().collect::>(), &datums[..n]); + assert_eq!(tail.iter().collect::>(), &datums[n..]); + assert_eq!( + head.byte_len() + tail.byte_len(), + row.as_row_ref().byte_len() + ); + } + } + + #[mz_ore::test] + #[should_panic(expected = "fewer than")] + fn split_at_datum_past_the_end_panics() { + let row = Row::pack_slice(&[Datum::Int64(1)]); + let _ = row.split_at_datum(2); + } }