diff --git a/src/compute/src/render/top_k.rs b/src/compute/src/render/top_k.rs index 26112ad7d3708..3e162ea8e2db1 100644 --- a/src/compute/src/render/top_k.rs +++ b/src/compute/src/render/top_k.rs @@ -36,6 +36,7 @@ use mz_ore::cast::CastFrom; use mz_ore::soft_assert_or_log; use mz_repr::fixed_length::ExtendDatums; use mz_repr::{Datum, DatumVec, Diff, ReprScalarType, Row, SharedRow}; +use mz_timely_util::columnar::builder::ColumnBuilder; use mz_timely_util::columnation::ColumnationChunker; use mz_timely_util::operator::CollectionExt; use timely::Container; @@ -293,10 +294,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim "requested no validation, but received error collection" ); - CollectionBundle::from_collections( - result.map(|(_key_hash, row)| row), - err_collection, - ) + CollectionBundle::from_edge(topk_result_to_columnar(result), err_collection) } TopKPlan::Basic(BasicTopKPlan { group_key, @@ -319,7 +317,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim ok_input, group_key, order_key, offset, limit, arity, buckets, ); err_collection = err_collection.concat(errs); - CollectionBundle::from_collections(oks, err_collection) + CollectionBundle::from_edge(oks, err_collection) } }; @@ -341,7 +339,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim arity: usize, buckets: Vec, ) -> ( - VecCollection<'s, T, Row, Diff>, + CollectionEdge<'s, T>, VecCollection<'s, T, DataflowErrorSer, Diff>, ) { let pairer = Pairer::new(1); @@ -417,7 +415,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim err_collection = errs; } ( - collection.map(|(_key_hash, row)| row), + topk_result_to_columnar(collection), err_collection.expect("at least one stage validated its inputs"), ) } @@ -663,6 +661,36 @@ where } } +/// Drops the hash-key pairing from a consolidated `(hash_key, row)` TopK result. +/// +/// The hash key is a function of the row and the input is consolidated, so dropping the +/// key is injective and the output has no within-batch duplicates for a consolidating +/// builder to fold. Rows are pushed borrowed. +/// +/// TODO: TopK renders its stages over `Vec` containers, so this encode sits at the very +/// end of the plan. Pushing columnar containers down through `build_topk` and the +/// monotonic path would remove it. +fn topk_result_to_columnar<'s, T>( + collection: VecCollection<'s, T, (Row, Row), Diff>, +) -> CollectionEdge<'s, T> +where + T: crate::render::RenderTimestamp, +{ + let stream = collection + .inner + .unary::, _, _, _>(Pipeline, "TopKUnkey", |_cap, _info| { + move |input, output| { + input.for_each(|time, data| { + let mut session = output.session_with_builder(&time); + for ((_key_hash, row), t, d) in data.drain(..) { + session.give((&row, &t, &d)); + } + }); + } + }); + CollectionEdge::Columnar(stream.as_collection()) +} + /// 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. @@ -1328,4 +1356,57 @@ mod tests { assert_eq!(key_datums[1], value_datums[0]); } } + + #[mz_ore::test] + fn topk_result_to_columnar_drops_key() { + let key = Row::pack_slice(&[Datum::Int64(7)]); + let rows = vec![ + ( + (key.clone(), Row::pack_slice(&[Datum::Int32(1)])), + 0u64, + Diff::ONE, + ), + ( + (key.clone(), Row::pack_slice(&[Datum::Int32(2)])), + 1u64, + Diff::ONE, + ), + // Retracts at a `(row, time)` with no insertion, so the `InputSession`'s + // pre-send consolidation does not cancel it out. + ( + (key.clone(), Row::pack_slice(&[Datum::Int32(1)])), + 2u64, + -Diff::ONE, + ), + ]; + let mut expected: Vec<(Row, Timestamp, Diff)> = rows + .iter() + .map(|((_, v), t, d)| (v.clone(), Timestamp::from(*t), *d)) + .collect(); + expected.sort(); + + let (is_columnar, captured) = timely::execute_directly(move |worker| { + worker.dataflow::(|scope| { + let (mut handle, collection) = scope.new_collection(); + let edge = topk_result_to_columnar(collection); + let is_columnar = matches!(edge, CollectionEdge::Columnar(_)); + let captured = edge.into_vec().inner.capture(); + for (kv, time, diff) in rows { + handle.update_at(kv, Timestamp::from(time), diff); + } + handle.advance_to(Timestamp::from(3u64)); + handle.flush(); + (is_columnar, captured) + }) + }); + assert!(is_columnar, "the TopK output must be a columnar edge"); + + let mut got: Vec<(Row, Timestamp, Diff)> = captured + .extract() + .into_iter() + .flat_map(|(_, data)| data) + .collect(); + got.sort(); + assert_eq!(got, expected); + } }