From 3599ddc17f430cc9dfef9755d67bdac0a5e60235 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Tue, 21 Jul 2026 21:26:32 +0200 Subject: [PATCH 1/5] compute: Collapse CollectionEdge to the columnar collection Every producer emits the columnar edge and no producer constructs a `Vec` edge, so the `CollectionEdge` enum is now a single-variant wrapper. Replace it with a type alias for the columnar collection and delete the `Vec` variant. Deleting the variant is the compiler-enforced completeness check for the migration: nothing failed to compile, so no live dependency on the row-based edge remained. De-match the carrier operations, which no longer branch on a variant: `scope`, `enter_region`, and `leave_region` are the collection's own methods; `negate`, `concat_many`, `consolidate_named`, and `flat_map_datums` become single columnar free functions (the existing `columnar_negate` / `columnar_consolidate`, and new `concat_many` / `flat_map_datums`). The row-forming consumers (`arrange_collection`, `map_topk_key`, linear-join key preparation, FlatMap input) drop their `Vec` arms. The leaf conversions stay: `.into_vec()` method calls become `columnar_to_vec(...)` at the sanctioned leaves (sink, LetRec, temporal bucketing, TopK fallible-limit, linear-join initial closure, delta-join raw source), and `vec_to_columnar` remains the leaf encode. Pure type collapse with no runtime change; everything was already columnar. The cross-arm unit tests become single-path columnar correctness tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/compute/src/render.rs | 89 ++-- src/compute/src/render/columnar.rs | 385 +++++---------- src/compute/src/render/context.rs | 536 ++++++++------------- src/compute/src/render/flat_map.rs | 95 ++-- src/compute/src/render/join/delta_join.rs | 11 +- src/compute/src/render/join/linear_join.rs | 238 ++++----- src/compute/src/render/sinks.rs | 5 +- src/compute/src/render/top_k.rs | 177 +++---- 8 files changed, 558 insertions(+), 978 deletions(-) diff --git a/src/compute/src/render.rs b/src/compute/src/render.rs index 6a7351cc09bb4..1f6552f49ed10 100644 --- a/src/compute/src/render.rs +++ b/src/compute/src/render.rs @@ -166,7 +166,9 @@ use crate::extensions::temporal_bucket::TemporalBucketing; use crate::logging::compute::{ ComputeEvent, DataflowGlobal, LirMapping, LirMetadata, LogDataflowErrors, OperatorHydration, }; -use crate::render::columnar::{CollectionEdge, vec_to_columnar}; +use crate::render::columnar::{ + columnar_consolidate, columnar_negate, columnar_to_vec, concat_many, vec_to_columnar, +}; use crate::render::context::{ArrangementFlavor, Context}; use crate::render::errors::DataflowErrorSer; use crate::typedefs::{ErrBatcher, ErrBuilder, ErrSpine, KeyBatcher, MzTimestamp}; @@ -379,7 +381,7 @@ pub fn build_compute_dataflow( // columnar edge at the boundary. The batches are already // consolidated, so this leaf-encode is non-consolidating. let bundle = crate::render::CollectionBundle::from_edge( - CollectionEdge::Columnar(vec_to_columnar(oks.enter(region))), + vec_to_columnar(oks.enter(region)), errs.enter(region), ); // Associate collection bundle with the source identifier. @@ -482,7 +484,7 @@ pub fn build_compute_dataflow( // columnar edge at the boundary. The batches are already // consolidated, so this leaf-encode is non-consolidating. let bundle = crate::render::CollectionBundle::from_edge( - CollectionEdge::Columnar(vec_to_columnar(oks.enter_region(region))), + vec_to_columnar(oks.enter_region(region)), errs.enter_region(region), ); // Associate collection bundle with the source identifier. @@ -677,10 +679,7 @@ where // The filtered index collection is row-shaped. Encode it to // the columnar edge at the boundary. It is already // consolidated, so this leaf-encode is non-consolidating. - CollectionBundle::from_edge( - CollectionEdge::Columnar(vec_to_columnar(oks)), - errs, - ) + CollectionBundle::from_edge(vec_to_columnar(oks), errs) } }; self.update_id(Id::Global(idx.on_id), bundle); @@ -960,10 +959,7 @@ impl<'scope> Context<'scope, Product>> { // iterative frontier or fixpoint behavior. self.insert_id( Id::Local(*id), - CollectionBundle::from_edge( - CollectionEdge::Columnar(vec_to_columnar(oks_collection)), - err_collection, - ), + CollectionBundle::from_edge(vec_to_columnar(oks_collection), err_collection), ); variables.insert(Id::Local(*id), (oks_v, err_v)); } @@ -979,7 +975,7 @@ impl<'scope> Context<'scope, Product>> { // We need to ensure that the raw collection exists, but do not have enough information // here to cause that to happen. let (oks, mut err) = bundle.collection.clone().unwrap(); - let oks = oks.into_vec(); + let oks = columnar_to_vec(oks); decoded_oks.insert(id, oks.clone()); // Collapses what forward reads see. `err_v` below feeds reads rendered before this // binding and is collapsed separately; without this, a `Get` in a later rec binding @@ -1051,7 +1047,7 @@ impl<'scope> Context<'scope, Product>> { self.insert_id( Id::Local(id), CollectionBundle::from_edge( - CollectionEdge::Columnar(vec_to_columnar(oks.leave_dynamic(level + 1))), + vec_to_columnar(oks.leave_dynamic(level + 1)), err.leave_dynamic(level + 1), ), ); @@ -1256,25 +1252,24 @@ impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> { // left distinct may become duplicates here; a // `ConsolidatingColumnBuilder` folds those within the batch (the // rows are already owned, so the give is a move into staging). - let ok_collection = CollectionEdge::Columnar( - rows.into_iter() - .filter_map(move |(row, mut time, diff)| { - time.advance_by(as_of_frontier.borrow()); - if !until.less_equal(&time) { - Some(( - row.0, - >::to_inner(time), - diff, - )) - } else { - None - } - }) - .to_stream_with_builder::<_, ConsolidatingColumnBuilder>( - self.scope, - ) - .as_collection(), - ); + let ok_collection = rows + .into_iter() + .filter_map(move |(row, mut time, diff)| { + time.advance_by(as_of_frontier.borrow()); + if !until.less_equal(&time) { + Some(( + row.0, + >::to_inner(time), + diff, + )) + } else { + None + } + }) + .to_stream_with_builder::<_, ConsolidatingColumnBuilder>( + self.scope, + ) + .as_collection(); let mut error_time: mz_repr::Timestamp = Timestamp::minimum(); error_time.advance_by(self.as_of_frontier.borrow()); @@ -1397,7 +1392,7 @@ impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> { .collection .clone() .expect("Negate input must be an unarranged collection"); - CollectionBundle::from_edge(oks.negate(), errs) + CollectionBundle::from_edge(columnar_negate(oks), errs) } Threshold { input, @@ -1431,13 +1426,11 @@ impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> { // Temporal bucketing operates on `Vec`: decode the edge // into it, then re-encode the result so this Union input // is a columnar edge like every other. - let os = os.into_vec(); - CollectionEdge::Columnar(vec_to_columnar( - T::maybe_apply_temporal_bucketing( - os.inner, - self.as_of_frontier.clone(), - summary, - ), + let os = columnar_to_vec(os); + vec_to_columnar(T::maybe_apply_temporal_bucketing( + os.inner, + self.as_of_frontier.clone(), + summary, )) } else { os @@ -1445,9 +1438,9 @@ impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> { oks.push(os); errs.push(es); } - let oks = CollectionEdge::concat_many(self.scope, oks); + let oks = concat_many(self.scope, oks); let oks = if consolidate_output { - oks.consolidate_named("UnionConsolidation") + columnar_consolidate(oks, "UnionConsolidation") } else { oks }; @@ -1528,16 +1521,8 @@ impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> { .collection .as_mut() .expect("CollectionBundle invariant"); - match oks { - CollectionEdge::Vec(c) => { - let stream = self.log_operator_hydration_inner(c.inner.clone(), lir_id); - *c = stream.as_collection(); - } - CollectionEdge::Columnar(c) => { - let stream = self.log_operator_hydration_inner(c.inner.clone(), lir_id); - *c = stream.as_collection(); - } - } + let stream = self.log_operator_hydration_inner(oks.inner.clone(), lir_id); + *oks = stream.as_collection(); } } } diff --git a/src/compute/src/render/columnar.rs b/src/compute/src/render/columnar.rs index 935c9105e833d..9dffb52346fba 100644 --- a/src/compute/src/render/columnar.rs +++ b/src/compute/src/render/columnar.rs @@ -9,26 +9,16 @@ //! Columnar dataflow edge support. //! -//! Defines [`CollectionEdge`], a wrapper that lets dataflow edges between Plan -//! nodes carry either row-based ([`VecCollection`]) or columnar -//! ([`ColumnarCollection`]) batches of `(D, T, R)` updates. +//! Defines [`CollectionEdge`], the columnar batch representation that dataflow +//! edges between Plan nodes carry. Every producer emits this representation. //! -//! # Migration model -//! -//! The migration is consumer-first: every Plan-node consumer learns to accept -//! both variants before any producer emits the columnar variant. Producers can -//! then flip to columnar one at a time. -//! -//! Within a Plan node, operators may freely materialize Vec collections; only -//! the inter-node edge format is constrained. A decode from columnar to Vec at -//! a consumer's input is acceptable only when the consumer would have decoded -//! `Row` to [`mz_repr::Datum`] anyway. Pure passthrough consumers (Negate, -//! Union) round-trip the columnar variant without decoding. -//! -//! Consumers that have not yet learned the columnar form fall back to -//! [`CollectionEdge::into_vec`], which decodes through the named -//! `ColumnarToVec` operator. Repack seams therefore stay visible in dataflow -//! introspection, so they can be found and retired. +//! Within a Plan node, operators may freely materialize `Vec` collections. Only +//! the inter-node edge format is constrained. A node that produces a row-based +//! collection re-encodes it to the columnar edge at its output leaf via +//! [`vec_to_columnar`]. A node that must consume rows decodes at its input leaf +//! via [`columnar_to_vec`]. Both are named operators (`VecToColumnar`, +//! `ColumnarToVec`), so those leaf seams stay visible in dataflow +//! introspection. use columnar::{Columnar, Index}; use differential_dataflow::{AsCollection, Collection, VecCollection}; @@ -38,7 +28,7 @@ use mz_timely_util::columnar::batcher::ColumnChunker; use mz_timely_util::columnar::builder::ColumnBuilder; use mz_timely_util::columnar::columnar_consolidate_exchange; use mz_timely_util::columnar::merge_batcher::ColumnMergeBatcher; -use mz_timely_util::operator::{CollectionExt, consolidate_pact}; +use mz_timely_util::operator::consolidate_pact; use timely::ContainerBuilder; use timely::container::CapacityContainerBuilder; use timely::dataflow::channels::pact::{ExchangeCore, Pipeline}; @@ -49,7 +39,6 @@ use timely::dataflow::{Scope, Stream, StreamVec}; use crate::render::RenderTimestamp; use crate::render::context::{ECB, Session}; use crate::render::errors::DataflowErrorSer; -use crate::typedefs::KeyBatcher; /// A columnar collection of `(D, T, R)` updates traveling on a compute /// dataflow edge. @@ -58,204 +47,84 @@ use crate::typedefs::KeyBatcher; /// container is [`Column<(D, T, R)>`] instead of `Vec<(D, T, R)>`. pub type ColumnarCollection<'scope, T, D, R> = Collection<'scope, T, Column<(D, T, R)>>; -/// A dataflow edge carrying records as either a row-based [`VecCollection`] or -/// a [`ColumnarCollection`]. -/// -/// Producers choose a variant; consumers must accept either. Variant-mixing -/// `concat`s repack the row-based inputs and produce the columnar variant. -#[derive(Clone)] -pub enum CollectionEdge<'scope, T: RenderTimestamp> { - /// Row-formatted collection. No producer constructs this after the - /// migration; the variant and its remaining match arms are removed when the - /// enum collapses to a columnar alias. - #[allow(dead_code)] - Vec(VecCollection<'scope, T, Row, Diff>), - /// Columnar collection. Currently unused by any producer; reserved for the - /// producer flip at the end of the migration. - Columnar(ColumnarCollection<'scope, T, Row, Diff>), -} - -impl<'scope, T: RenderTimestamp> CollectionEdge<'scope, T> { - /// The scope containing this edge. - pub fn scope(&self) -> Scope<'scope, T> { - match self { - CollectionEdge::Vec(c) => c.inner.scope(), - CollectionEdge::Columnar(c) => c.inner.scope(), - } - } - - /// Brings the edge into a sub-region of its current scope. - pub fn enter_region<'inner>(self, region: Scope<'inner, T>) -> CollectionEdge<'inner, T> { - match self { - CollectionEdge::Vec(c) => CollectionEdge::Vec(c.enter_region(region)), - CollectionEdge::Columnar(c) => CollectionEdge::Columnar(c.enter_region(region)), - } - } - - /// Leaves a sub-region back to the outer scope. - pub fn leave_region<'outer>(self, outer: Scope<'outer, T>) -> CollectionEdge<'outer, T> { - match self { - CollectionEdge::Vec(c) => CollectionEdge::Vec(c.leave_region(outer)), - CollectionEdge::Columnar(c) => CollectionEdge::Columnar(c.leave_region(outer)), - } - } - - /// The edge as a row-based [`VecCollection`]. - /// - /// The Vec arm is returned as is. The columnar arm decodes through - /// [`columnar_to_vec`], which allocates an owned [`Row`] per record. - /// Consumers that can work on the columnar form directly should do so - /// instead of calling this. - pub fn into_vec(self) -> VecCollection<'scope, T, Row, Diff> { - match self { - CollectionEdge::Vec(c) => c, - CollectionEdge::Columnar(c) => columnar_to_vec(c), - } - } - - /// Negates the diff on every record in this edge. - /// - /// Preserves variant. The columnar arm uses [`columnar_negate`], which - /// negates diffs without decoding rows. - pub fn negate(self) -> Self { - match self { - CollectionEdge::Vec(c) => CollectionEdge::Vec(c.negate()), - CollectionEdge::Columnar(c) => CollectionEdge::Columnar(columnar_negate(c)), - } - } - - /// Concatenates a collection of edges. - /// - /// The inputs are all columnar, so they concatenate natively into the - /// columnar variant. - pub fn concat_many(scope: Scope<'scope, T>, edges: I) -> Self - where - I: IntoIterator, - { - let cols = edges.into_iter().map(|edge| match edge { - CollectionEdge::Columnar(c) => c, - // No producer emits `Vec`, so a `Vec` input cannot reach here. - CollectionEdge::Vec(_) => unreachable!("no producer emits a `Vec` edge"), - }); - CollectionEdge::Columnar(differential_dataflow::collection::concatenate( - scope, - cols.collect::>(), - )) - } +/// A dataflow edge between Plan nodes: a columnar collection of `(Row, Diff)` +/// updates. Every producer emits this representation; a row-based producer +/// re-encodes to it at its output leaf via [`vec_to_columnar`]. +pub type CollectionEdge<'scope, T> = ColumnarCollection<'scope, T, Row, Diff>; - /// Applies `logic` to each record in this edge, exposing the record as a - /// borrowed [`DatumVecBorrow`] and giving it ok and err output sessions. - /// - /// `max_demand` bounds the number of columns decoded per row; pass - /// `usize::MAX` to decode all columns. - /// - /// This is the canonical unified entry point for "decoding consumers" - /// (operators that read [`mz_repr::Datum`]s from each row anyway). The - /// Vec arm uses [`DatumVec::borrow_with_limit`] on each [`Row`]; the - /// Columnar arm iterates the columnar batch directly without going - /// through an owned [`Row`]. - pub fn flat_map_datums( - self, - max_demand: usize, - mut logic: L, - ) -> ( - Stream<'scope, T, DCB::Container>, - StreamVec<'scope, T, (DataflowErrorSer, T, Diff)>, - ) - where - DCB: ContainerBuilder, - L: for<'a> FnMut( - &'a mut DatumVecBorrow<'_>, - T, - Diff, - &mut Session, - &mut Session>, - ) -> usize - + 'static, - { - match self { - CollectionEdge::Vec(c) => { - let scope = c.inner.scope(); - let mut builder = OperatorBuilder::new("CollectionFlatMap".to_string(), scope); - let (ok_output, ok_stream) = builder.new_output(); - let mut ok_output = OutputBuilder::<_, DCB>::from(ok_output); - let (err_output, err_stream) = builder.new_output(); - let mut err_output = OutputBuilder::<_, ECB>::from(err_output); - let mut input = builder.new_input(c.inner, Pipeline); - builder.build(move |_capabilities| { - let mut datums = DatumVec::new(); - move |_frontiers| { - let mut ok_output = ok_output.activate(); - let mut err_output = err_output.activate(); - input.for_each(|time, data| { - // Retain the input capability to derive a `Capability` for each output; - // the `Session` type alias is fixed to `Capability`. - let ok_cap = time.retain(0); - let err_cap = time.retain(1); - let mut ok_session = ok_output.session_with_builder(&ok_cap); - let mut err_session = err_output.session_with_builder(&err_cap); - for (v, t, d) in data.drain(..) { - logic( - &mut datums.borrow_with_limit(&v, max_demand), - t, - d, - &mut ok_session, - &mut err_session, - ); - } - }); - } - }); - (ok_stream, err_stream) - } - CollectionEdge::Columnar(c) => { - let scope = c.inner.scope(); - let mut builder = OperatorBuilder::new("CollectionFlatMap".to_string(), scope); - let (ok_output, ok_stream) = builder.new_output(); - let mut ok_output = OutputBuilder::<_, DCB>::from(ok_output); - let (err_output, err_stream) = builder.new_output(); - let mut err_output = OutputBuilder::<_, ECB>::from(err_output); - let mut input = builder.new_input(c.inner, Pipeline); - builder.build(move |_capabilities| { - let mut datums = DatumVec::new(); - move |_frontiers| { - let mut ok_output = ok_output.activate(); - let mut err_output = err_output.activate(); - input.for_each(|time, data| { - // Retain the input capability to derive a `Capability` for each output; - // the `Session` type alias is fixed to `Capability`. - let ok_cap = time.retain(0); - let err_cap = time.retain(1); - let mut ok_session = ok_output.session_with_builder(&ok_cap); - let mut err_session = err_output.session_with_builder(&err_cap); - // Rows are read from the borrowed column, never - // materialized as owned `Row`s. - for (v, t, d) in data.borrow().into_index_iter() { - logic( - &mut datums.borrow_with_limit(v, max_demand), - Columnar::into_owned(t), - Columnar::into_owned(d), - &mut ok_session, - &mut err_session, - ); - } - }); - } - }); - (ok_stream, err_stream) - } - } - } +/// Concatenates a collection of columnar edges. +pub fn concat_many<'scope, T, I>(scope: Scope<'scope, T>, edges: I) -> CollectionEdge<'scope, T> +where + T: RenderTimestamp, + I: IntoIterator>, +{ + let cols: Vec<_> = edges.into_iter().collect(); + differential_dataflow::collection::concatenate(scope, cols) +} - /// Consolidates updates in the edge, preserving variant. - pub fn consolidate_named(self, name: &str) -> Self { - match self { - CollectionEdge::Vec(c) => CollectionEdge::Vec(CollectionExt::consolidate_named::< - KeyBatcher<_, _, _>, - >(c, name)), - CollectionEdge::Columnar(c) => CollectionEdge::Columnar(columnar_consolidate(c, name)), +/// Applies `logic` to each record in `edge`, exposing the record as a borrowed +/// [`DatumVecBorrow`] and giving it ok and err output sessions. +/// +/// `max_demand` bounds the number of columns decoded per row; pass `usize::MAX` +/// to decode all columns. +/// +/// This is the canonical entry point for "decoding consumers" (operators that +/// read [`mz_repr::Datum`]s from each row anyway). It iterates the columnar +/// batch directly without going through an owned [`Row`]. +pub fn flat_map_datums<'scope, T, DCB, L>( + edge: CollectionEdge<'scope, T>, + max_demand: usize, + mut logic: L, +) -> ( + Stream<'scope, T, DCB::Container>, + StreamVec<'scope, T, (DataflowErrorSer, T, Diff)>, +) +where + T: RenderTimestamp, + DCB: ContainerBuilder, + L: for<'a> FnMut( + &'a mut DatumVecBorrow<'_>, + T, + Diff, + &mut Session, + &mut Session>, + ) -> usize + + 'static, +{ + let scope = edge.inner.scope(); + let mut builder = OperatorBuilder::new("CollectionFlatMap".to_string(), scope); + let (ok_output, ok_stream) = builder.new_output(); + let mut ok_output = OutputBuilder::<_, DCB>::from(ok_output); + let (err_output, err_stream) = builder.new_output(); + let mut err_output = OutputBuilder::<_, ECB>::from(err_output); + let mut input = builder.new_input(edge.inner, Pipeline); + builder.build(move |_capabilities| { + let mut datums = DatumVec::new(); + move |_frontiers| { + let mut ok_output = ok_output.activate(); + let mut err_output = err_output.activate(); + input.for_each(|time, data| { + // Retain the input capability to derive a `Capability` for each output; + // the `Session` type alias is fixed to `Capability`. + let ok_cap = time.retain(0); + let err_cap = time.retain(1); + let mut ok_session = ok_output.session_with_builder(&ok_cap); + let mut err_session = err_output.session_with_builder(&err_cap); + // Rows are read from the borrowed column, never materialized as + // owned `Row`s. + for (v, t, d) in data.borrow().into_index_iter() { + logic( + &mut datums.borrow_with_limit(v, max_demand), + Columnar::into_owned(t), + Columnar::into_owned(d), + &mut ok_session, + &mut err_session, + ); + } + }); } - } + }); + (ok_stream, err_stream) } /// Negates the diff of every record in a [`ColumnarCollection`]. @@ -487,7 +356,7 @@ mod tests { } #[mz_ore::test] - fn negate_flips_diffs_on_columnar_arm() { + fn columnar_negate_flips_diffs() { let rows = test_rows(); let expected: Vec<_> = { let mut updates: Vec<_> = rows @@ -500,9 +369,8 @@ mod tests { let captured = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut input, collection) = scope.new_collection(); - let edge = CollectionEdge::Columnar(vec_to_columnar(collection)).negate(); - assert!(matches!(edge, CollectionEdge::Columnar(_))); - let captured = edge.into_vec().inner.capture(); + let edge = columnar_negate(vec_to_columnar(collection)); + let captured = columnar_to_vec(edge).inner.capture(); for row in rows { input.update(row, Diff::ONE); } @@ -531,15 +399,11 @@ mod tests { worker.dataflow::(|scope| { let (mut input1, collection1) = scope.new_collection(); let (mut input2, collection2) = scope.new_collection(); - let edge = CollectionEdge::concat_many( + let edge = concat_many( scope, - [ - CollectionEdge::Columnar(vec_to_columnar(collection1)), - CollectionEdge::Columnar(vec_to_columnar(collection2)), - ], + [vec_to_columnar(collection1), vec_to_columnar(collection2)], ); - assert!(matches!(edge, CollectionEdge::Columnar(_))); - let captured = edge.into_vec().inner.capture(); + let captured = columnar_to_vec(edge).inner.capture(); let (first, rest) = rows.split_first().unwrap(); input1.update(first.clone(), Diff::ONE); input2.update(first.clone(), Diff::ONE); @@ -558,44 +422,36 @@ mod tests { #[mz_ore::test] fn flat_map_datums_arms_agree() { - // Project the first datum of each row, exercising `max_demand` on both - // arms. The two captures must extract identical updates. + // Project the first datum of each row, exercising `max_demand`. let rows = test_rows(); - let (vec_captured, col_captured) = timely::execute_directly(move |worker| { + let captured = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut input, collection) = scope.new_collection(); - let mut captures = Vec::new(); - for edge in [ - CollectionEdge::Vec(collection.clone()), - CollectionEdge::Columnar(vec_to_columnar(collection)), - ] { - let (oks, _errs) = edge.flat_map_datums::( - 1, - |datums, t, d, ok_session, _err_session| { - ok_session.give((Row::pack(datums.iter()), t, d)); - 1 - }, - ); - captures.push(oks.capture()); - } - let col = captures.pop().unwrap(); - let vec = captures.pop().unwrap(); + let (oks, _errs) = flat_map_datums::<_, RowBuilder, _>( + vec_to_columnar(collection), + 1, + |datums, t, d, ok_session, _err_session| { + ok_session.give((Row::pack(datums.iter()), t, d)); + 1 + }, + ); + let captured = oks.capture(); for row in rows { input.update(row, Diff::ONE); } input.advance_to(Timestamp::from(1_u64)); input.flush(); - (vec, col) + captured }) }); - let vec_updates = extract_sorted(vec_captured); - assert_eq!(vec_updates, extract_sorted(col_captured)); + let updates = extract_sorted(captured); + assert!(!updates.is_empty()); // Each output row retains at most the first datum of its input. - assert!(vec_updates.iter().all(|(r, _, _)| r.iter().count() <= 1)); + assert!(updates.iter().all(|(r, _, _)| r.iter().count() <= 1)); } #[mz_ore::test] - fn consolidate_named_preserves_columnar() { + fn columnar_consolidate_accumulates_and_cancels() { let row1 = Row::pack_slice(&[Datum::Int32(1)]); let row2 = Row::pack_slice(&[Datum::Int32(2)]); let row3 = Row::pack_slice(&[Datum::Int32(3)]); @@ -608,26 +464,11 @@ mod tests { (row1.clone(), Timestamp::from(1_u64), Diff::ONE), ]; - // The columnar arm keeps the `Columnar` variant. No-ColumnarToVec is a - // by-inspection property: `consolidate_named`'s columnar arm calls - // `columnar_consolidate` (native `ColumnMergeBatcher` merge), never - // `columnar_to_vec`. The `into_vec` below is the capture harness - // decoding for the test only, not part of the consolidate. - let (vec_captured, col_captured) = timely::execute_directly(move |worker| { + let captured = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut input, collection) = scope.new_collection(); - let mut captures = Vec::new(); - for edge in [ - CollectionEdge::Vec(collection.clone()), - CollectionEdge::Columnar(vec_to_columnar(collection)), - ] { - let is_columnar = matches!(edge, CollectionEdge::Columnar(_)); - let edge = edge.consolidate_named("Test"); - assert_eq!(matches!(edge, CollectionEdge::Columnar(_)), is_columnar); - captures.push(edge.into_vec().inner.capture()); - } - let col = captures.pop().unwrap(); - let vec = captures.pop().unwrap(); + let edge = columnar_consolidate(vec_to_columnar(collection), "Test"); + let captured = columnar_to_vec(edge).inner.capture(); // t=0: row1 accumulates (+1, +1), row2 cancels (+1, -1). input.advance_to(Timestamp::from(0_u64)); input.update(row1.clone(), Diff::ONE); @@ -641,11 +482,9 @@ mod tests { input.update(row3, -Diff::ONE); input.advance_to(Timestamp::from(2_u64)); input.flush(); - (vec, col) + captured }) }); - let vec_updates = extract_sorted(vec_captured); - assert_eq!(vec_updates, expected); - assert_eq!(extract_sorted(col_captured), vec_updates); + assert_eq!(extract_sorted(captured), expected); } } diff --git a/src/compute/src/render/context.rs b/src/compute/src/render/context.rs index e81703169ab0e..36092a2f8c1fe 100644 --- a/src/compute/src/render/context.rs +++ b/src/compute/src/render/context.rs @@ -51,7 +51,7 @@ use timely::progress::{Antichain, Timestamp}; use crate::compute_state::ComputeState; use crate::extensions::arrange::{ArrangementBatcher, KeyCollection, MzArrange, MzArrangeCore}; use crate::extensions::reduce::MzReduce; -use crate::render::columnar::{CollectionEdge, vec_to_columnar}; +use crate::render::columnar::{CollectionEdge, columnar_to_vec, flat_map_datums, vec_to_columnar}; use crate::render::errors::{DataflowErrorSer, ErrorLogger}; use crate::render::{LinearJoinSpec, MaybeBucketByTime, RenderTimestamp}; use crate::typedefs::{ @@ -639,7 +639,7 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> { 1 }, ); - (CollectionEdge::Columnar(ok.as_collection()), err) + (ok.as_collection(), err) } } } @@ -691,7 +691,7 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> { .collection .clone() .expect("Invariant violated: CollectionBundle contains no collection."); - let (ok_stream, err_stream) = oks.flat_map_datums::(max_demand, logic); + let (ok_stream, err_stream) = flat_map_datums::<_, DCB, _>(oks, max_demand, logic); let errs = errs.concat(err_stream.as_collection()); (ok_stream, errs) } @@ -1048,7 +1048,7 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> { }, ); - (CollectionEdge::Columnar(stream.as_collection()), errors) + (stream.as_collection(), errors) } pub fn ensure_collections( mut self, @@ -1118,11 +1118,11 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> { // produces a `Vec` stream. Decode the edge into it, then re-encode // the `Vec` result to columnar at the boundary so the bucketed // output edge stays columnar like every other producer. - CollectionEdge::Columnar(vec_to_columnar(T::maybe_apply_temporal_bucketing( - oks.into_vec().inner, + vec_to_columnar(T::maybe_apply_temporal_bucketing( + columnar_to_vec(oks).inner, as_of.clone(), summary, - ))) + )) } else { oks }; @@ -1146,26 +1146,27 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> { } else { strategy }; - let oks = - if matches!(effective_strategy, ArrangementStrategy::TemporalBucketing) - && ENABLE_COMPUTE_TEMPORAL_BUCKETING.get(config_set) - { - let summary: mz_repr::Timestamp = TEMPORAL_BUCKETING_SUMMARY - .get(config_set) - .try_into() - .expect("must fit"); - bucketed = true; - // Temporal bucketing is `Vec`-internal: it consumes - // and produces a `Vec` stream. Decode the edge into it, then - // re-encode the `Vec` result to columnar at the boundary so the - // bucketed output edge stays columnar like every other producer. - let oks = oks.into_vec(); - CollectionEdge::Columnar(vec_to_columnar( - T::maybe_apply_temporal_bucketing(oks.inner, as_of.clone(), summary), - )) - } else { - oks - }; + let oks = if matches!(effective_strategy, ArrangementStrategy::TemporalBucketing) + && ENABLE_COMPUTE_TEMPORAL_BUCKETING.get(config_set) + { + let summary: mz_repr::Timestamp = TEMPORAL_BUCKETING_SUMMARY + .get(config_set) + .try_into() + .expect("must fit"); + bucketed = true; + // Temporal bucketing is `Vec`-internal: it consumes + // and produces a `Vec` stream. Decode the edge into it, then + // re-encode the `Vec` result to columnar at the boundary so the + // bucketed output edge stays columnar like every other producer. + let oks = columnar_to_vec(oks); + vec_to_columnar(T::maybe_apply_temporal_bucketing( + oks.inner, + as_of.clone(), + summary, + )) + } else { + oks + }; let batcher = ArrangementBatcher::from_config(config_set); let (oks, errs_keyed, passthrough) = Self::arrange_collection(&name, oks, key.clone(), thinning.clone(), batcher); @@ -1213,125 +1214,66 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> { // references, which is what we need to push into columnar streams. Instead, we use a // bespoke operator that also optimizes reuse of allocations across individual updates. // - // The two arms differ only in how records are read from the input container and in the - // variant of the passthrough output, which preserves the input variant. The ok output is - // columnar in both arms, since the arrangement key/value is always columnar. - let (ok_stream, err_stream, passthrough) = match oks { - CollectionEdge::Vec(oks) => { - let mut builder = - OperatorBuilder::new("FormArrangementKey".to_string(), oks.inner.scope()); - let (ok_output, ok_stream) = builder.new_output(); - let mut ok_output = - OutputBuilder::<_, ColumnBuilder<((Row, Row), T, Diff)>>::from(ok_output); - let (err_output, err_stream) = builder.new_output(); - let mut err_output = OutputBuilder::from(err_output); - let (passthrough_output, passthrough_stream) = builder.new_output(); - let mut passthrough_output = OutputBuilder::from(passthrough_output); - let mut input = builder.new_input(oks.inner, Pipeline); - builder.set_notify_for(0, FrontierInterest::Never); - builder.build(move |_capabilities| { - let mut key_buf = Row::default(); - let mut val_buf = Row::default(); - let mut datums = DatumVec::new(); - move |_frontiers| { - // Scoped to the activation so the arena's retained capacity does not - // outlive a single scheduling invocation; cleared per row to reuse it - // within the batch. - let mut temp_storage = RowArena::new(); - let mut ok_output = ok_output.activate(); - let mut err_output = err_output.activate(); - let mut passthrough_output = passthrough_output.activate(); - input.for_each(|time, data| { - let mut ok_session = ok_output.session_with_builder(&time); - let mut err_session = err_output.session(&time); - for (row, time, diff) in data.iter() { - temp_storage.clear(); - let datums = datums.borrow_with(row); - let key_iter = key.iter().map(|k| k.eval(&datums, &temp_storage)); - match key_buf.packer().try_extend(key_iter) { - Ok(()) => { - let val_datum_iter = thinning.iter().map(|c| datums[*c]); - val_buf.packer().extend(val_datum_iter); - ok_session.give(((&*key_buf, &*val_buf), time, diff)); - } - Err(e) => { - err_session.give((e.into(), time.clone(), *diff)); - } + // The passthrough output forwards the input `Column` unchanged, so downstream consumers + // can reuse the collection without teeing. + let (ok_stream, err_stream, passthrough) = { + let mut builder = + OperatorBuilder::new("FormArrangementKey".to_string(), oks.inner.scope()); + let (ok_output, ok_stream) = builder.new_output(); + let mut ok_output = + OutputBuilder::<_, ColumnBuilder<((Row, Row), T, Diff)>>::from(ok_output); + let (err_output, err_stream) = builder.new_output(); + let mut err_output = OutputBuilder::from(err_output); + let (passthrough_output, passthrough_stream) = builder.new_output(); + // The passthrough forwards the input `Column` unchanged; its builder's container + // type must match the input so `give_container` can hand the batch through. + let mut passthrough_output = OutputBuilder::< + _, + CapacityContainerBuilder>, + >::from(passthrough_output); + let mut input = builder.new_input(oks.inner, Pipeline); + builder.set_notify_for(0, FrontierInterest::Never); + builder.build(move |_capabilities| { + let mut key_buf = Row::default(); + let mut val_buf = Row::default(); + let mut datums = DatumVec::new(); + move |_frontiers| { + // Scoped to the activation so the arena's retained capacity does not + // outlive a single scheduling invocation; cleared per row to reuse it + // within the batch. + let mut temp_storage = RowArena::new(); + let mut ok_output = ok_output.activate(); + let mut err_output = err_output.activate(); + let mut passthrough_output = passthrough_output.activate(); + input.for_each(|time, data| { + let mut ok_session = ok_output.session_with_builder(&time); + let mut err_session = err_output.session(&time); + // Rows are read from the borrowed column, never materialized as + // owned `Row`s. Times and diffs are owned only on the error path. + for (row, t, d) in data.borrow().into_index_iter() { + temp_storage.clear(); + let datums = datums.borrow_with(row); + let key_iter = key.iter().map(|k| k.eval(&datums, &temp_storage)); + match key_buf.packer().try_extend(key_iter) { + Ok(()) => { + let val_datum_iter = thinning.iter().map(|c| datums[*c]); + val_buf.packer().extend(val_datum_iter); + ok_session.give(((&*key_buf, &*val_buf), t, d)); } - } - passthrough_output.session(&time).give_container(data); - }); - } - }); - ( - ok_stream, - err_stream, - CollectionEdge::Vec(passthrough_stream.as_collection()), - ) - } - CollectionEdge::Columnar(oks) => { - let mut builder = - OperatorBuilder::new("FormArrangementKey".to_string(), oks.inner.scope()); - let (ok_output, ok_stream) = builder.new_output(); - let mut ok_output = - OutputBuilder::<_, ColumnBuilder<((Row, Row), T, Diff)>>::from(ok_output); - let (err_output, err_stream) = builder.new_output(); - let mut err_output = OutputBuilder::from(err_output); - let (passthrough_output, passthrough_stream) = builder.new_output(); - // The passthrough forwards the input `Column` unchanged; its builder's container - // type must match the input so `give_container` can hand the batch through. - let mut passthrough_output = OutputBuilder::< - _, - CapacityContainerBuilder>, - >::from(passthrough_output); - let mut input = builder.new_input(oks.inner, Pipeline); - builder.set_notify_for(0, FrontierInterest::Never); - builder.build(move |_capabilities| { - let mut key_buf = Row::default(); - let mut val_buf = Row::default(); - let mut datums = DatumVec::new(); - move |_frontiers| { - // Scoped to the activation so the arena's retained capacity does not - // outlive a single scheduling invocation; cleared per row to reuse it - // within the batch. - let mut temp_storage = RowArena::new(); - let mut ok_output = ok_output.activate(); - let mut err_output = err_output.activate(); - let mut passthrough_output = passthrough_output.activate(); - input.for_each(|time, data| { - let mut ok_session = ok_output.session_with_builder(&time); - let mut err_session = err_output.session(&time); - // Rows are read from the borrowed column, never materialized as - // owned `Row`s. Times and diffs are owned only on the error path. - for (row, t, d) in data.borrow().into_index_iter() { - temp_storage.clear(); - let datums = datums.borrow_with(row); - let key_iter = key.iter().map(|k| k.eval(&datums, &temp_storage)); - match key_buf.packer().try_extend(key_iter) { - Ok(()) => { - let val_datum_iter = thinning.iter().map(|c| datums[*c]); - val_buf.packer().extend(val_datum_iter); - ok_session.give(((&*key_buf, &*val_buf), t, d)); - } - Err(e) => { - err_session.give(( - e.into(), - Columnar::into_owned(t), - Columnar::into_owned(d), - )); - } + Err(e) => { + err_session.give(( + e.into(), + Columnar::into_owned(t), + Columnar::into_owned(d), + )); } } - passthrough_output.session(&time).give_container(data); - }); - } - }); - ( - ok_stream, - err_stream, - CollectionEdge::Columnar(passthrough_stream.as_collection()), - ) - } + } + passthrough_output.session(&time).give_container(data); + }); + } + }); + (ok_stream, err_stream, passthrough_stream.as_collection()) }; let exchange = @@ -1592,76 +1534,46 @@ mod tests { updates } - /// Arranges `rows` (each stamped with its own time) as both a `Vec` edge and - /// the equivalent columnar edge, keyed by `key`, and returns the sorted ok - /// and err outputs of each arm plus whether the columnar arm kept the - /// columnar passthrough variant. - fn arrange_both_arms( + /// Arranges `rows` (each stamped with its own time) as a columnar edge, keyed + /// by `key` with the full row as the value, and returns the sorted ok and err + /// outputs read back from the arrangement. + fn arrange_columnar( rows: Vec<(Row, u64)>, key: Vec, - ) -> ( - Vec, - Vec, - Vec<(String, Timestamp, Diff)>, - Vec<(String, Timestamp, Diff)>, - bool, - ) { + ) -> (Vec, Vec<(String, Timestamp, Diff)>) { let thinning = vec![0, 1]; - let (ok_vec, ok_col, err_vec, err_col, col_is_columnar) = - timely::execute_directly(move |worker| { - worker.dataflow::(|scope| { - let (mut input, collection) = scope.new_collection(); - - let (vec_arranged, vec_errs, _vec_passthrough) = - CollectionBundle::::arrange_collection( - &"vec".to_string(), - CollectionEdge::Vec(collection.clone()), - key.clone(), - thinning.clone(), - ArrangementBatcher::Columnation, - ); - let (col_arranged, col_errs, col_passthrough) = - CollectionBundle::::arrange_collection( - &"col".to_string(), - CollectionEdge::Columnar(vec_to_columnar(collection)), - key, - thinning, - ArrangementBatcher::Columnation, - ); - let col_is_columnar = matches!(col_passthrough, CollectionEdge::Columnar(_)); - - let ok_vec = vec_arranged - .as_collection(|k, v| (k.to_row(), v.to_row())) - .inner - .capture(); - let ok_col = col_arranged - .as_collection(|k, v| (k.to_row(), v.to_row())) - .inner - .capture(); - let err_vec = vec_errs.inner.capture(); - let err_col = col_errs.inner.capture(); - - let max_time = rows.iter().map(|(_, t)| *t).max().unwrap_or(0); - for (row, time) in rows { - input.update_at(row, Timestamp::from(time), Diff::ONE); - } - input.advance_to(Timestamp::from(max_time + 1)); - input.flush(); - (ok_vec, ok_col, err_vec, err_col, col_is_columnar) - }) - }); - ( - extract_ok(ok_vec), - extract_ok(ok_col), - extract_err(err_vec), - extract_err(err_col), - col_is_columnar, - ) + let (ok, err) = timely::execute_directly(move |worker| { + worker.dataflow::(|scope| { + let (mut input, collection) = scope.new_collection(); + let (arranged, errs, _passthrough) = + CollectionBundle::::arrange_collection( + &"col".to_string(), + vec_to_columnar(collection), + key, + thinning, + ArrangementBatcher::Columnation, + ); + let ok = arranged + .as_collection(|k, v| (k.to_row(), v.to_row())) + .inner + .capture(); + let err = errs.inner.capture(); + + let max_time = rows.iter().map(|(_, t)| *t).max().unwrap_or(0); + for (row, time) in rows { + input.update_at(row, Timestamp::from(time), Diff::ONE); + } + input.advance_to(Timestamp::from(max_time + 1)); + input.flush(); + (ok, err) + }) + }); + (extract_ok(ok), extract_err(err)) } // Uniform two-column rows so `Column(0)` keys and full-row thinning are in - // bounds for every record. Times span three distinct values so the columnar - // arm's per-record time handling is exercised, not just t=0. + // bounds for every record. Times span three distinct values so the per-record + // time handling is exercised, not just t=0. fn test_rows() -> Vec<(Row, u64)> { vec![ (Row::pack_slice(&[Datum::Int32(1), Datum::String("a")]), 0), @@ -1671,48 +1583,39 @@ mod tests { ] } - /// The columnar arm of the arrange input produces the same arranged contents - /// as the `Vec` arm and keeps the columnar passthrough variant. - /// - /// What this proves: correctness of the columnar key-forming path (a mangled - /// key or dropped record would diverge from the `Vec` arm) and that the - /// columnar arm actually ran (the passthrough stays `Columnar`). - /// - /// What it does NOT prove: absence of a silent decode. A hypothetical - /// `columnar_to_vec` on the ok path would yield identical contents and still - /// return a `Columnar` passthrough. The no-decode property holds by code - /// inspection: the columnar arm reads records via `into_index_iter` on the - /// borrowed column and never calls `into_vec`. + /// The columnar arrange input forms the arrangement keyed by column 0 with the + /// full row as the value, producing exactly the input records. A mangled key + /// or a dropped record would diverge from the expected set. #[mz_ore::test] - fn arrange_collection_arms_agree() { - let (ok_vec, ok_col, err_vec, err_col, col_is_columnar) = - arrange_both_arms(test_rows(), vec![LirScalarExpr::column(0)]); + fn arrange_collection_keys_correctly() { + let rows = test_rows(); + let mut expected: Vec = rows + .iter() + .map(|(row, t)| { + let key = Row::pack_slice(&[row.iter().next().unwrap()]); + ((key, row.clone()), Timestamp::from(*t), Diff::ONE) + }) + .collect(); + expected.sort(); - assert!( - col_is_columnar, - "columnar arrange input must keep the columnar passthrough variant" - ); - assert!(!ok_vec.is_empty()); - assert_eq!(ok_vec, ok_col); - assert!(err_vec.is_empty() && err_col.is_empty()); + let (ok, err) = arrange_columnar(rows, vec![LirScalarExpr::column(0)]); + assert_eq!(ok, expected); + assert!(err.is_empty()); } /// A key expression that always errors drives every record onto the error /// path, exercising the columnar arm's `into_owned` reconstruction of the - /// error's `(time, diff)`. The two arms must agree on the errors, and the ok - /// output must be empty on both. + /// error's `(time, diff)`. The ok output must be empty and the errors present. #[mz_ore::test] - fn arrange_collection_arms_agree_on_error_path() { + fn arrange_collection_error_path() { let key = vec![LirScalarExpr::literal( Err(EvalError::DivisionByZero), ReprScalarType::Int32, )]; - let (ok_vec, ok_col, err_vec, err_col, _col_is_columnar) = - arrange_both_arms(test_rows(), key); + let (ok, err) = arrange_columnar(test_rows(), key); - assert!(ok_vec.is_empty() && ok_col.is_empty()); - assert!(!err_vec.is_empty()); - assert_eq!(err_vec, err_col); + assert!(ok.is_empty()); + assert!(!err.is_empty()); } fn extract_row_updates( @@ -1727,16 +1630,12 @@ mod tests { updates } - /// A `Get -> ArrangeBy` chain carries the columnar arm end to end. A - /// non-identity MFP drives `as_collection_core` down its columnar producer - /// path, and feeding that edge into the arrange input keeps the columnar - /// passthrough, so no `ColumnarToVec` sits on the arrange path. - /// - /// The producer output is checked against the projected input. Arrange - /// correctness itself is covered by `arrange_collection_arms_agree`; here we - /// only assert the variant survives the hand-off. + /// A `Get -> ArrangeBy` chain: a non-identity MFP drives `as_collection_core` + /// down its columnar producer path, and feeding that edge into the arrange + /// input produces the projected rows. Arrange correctness itself is covered by + /// `arrange_collection_keys_correctly`; here we check the producer output. #[mz_ore::test] - fn get_arrange_by_carries_columnar_end_to_end() { + fn get_arrange_by_produces_projected_rows() { let rows = vec![ (Row::pack_slice(&[Datum::Int64(1), Datum::Int64(10)]), 0u64), (Row::pack_slice(&[Datum::Int64(2), Datum::Int64(20)]), 1), @@ -1758,81 +1657,67 @@ mod tests { .collect(); expected.sort(); - let (producer_is_columnar, passthrough_is_columnar, produced) = - timely::execute_directly(move |worker| { - worker.dataflow::(|scope| { - let (mut input, collection) = scope.new_collection(); - let (_err_input, errs) = scope.new_collection::(); - let bundle = CollectionBundle::from_edge(CollectionEdge::Vec(collection), errs); - let (edge, _errs) = bundle.as_collection_core(mfp, None, Antichain::new()); - let producer_is_columnar = matches!(edge, CollectionEdge::Columnar(_)); - // Tee the producer output for a content check, then feed the - // original edge into the arrange input. - let produced = edge.clone().into_vec().inner.capture(); - let (_arranged, _arrange_errs, passthrough) = - CollectionBundle::::arrange_collection( - &"arrange".to_string(), - edge, - vec![LirScalarExpr::column(0)], - vec![0], - ArrangementBatcher::Columnation, - ); - let passthrough_is_columnar = - matches!(passthrough, CollectionEdge::Columnar(_)); - - let max_time = rows.iter().map(|(_, t)| *t).max().unwrap(); - for (row, time) in rows { - input.update_at(row, Timestamp::from(time), Diff::ONE); - } - input.advance_to(Timestamp::from(max_time + 1)); - input.flush(); - (producer_is_columnar, passthrough_is_columnar, produced) - }) - }); + let produced = timely::execute_directly(move |worker| { + worker.dataflow::(|scope| { + let (mut input, collection) = scope.new_collection(); + let (_err_input, errs) = scope.new_collection::(); + let bundle = CollectionBundle::from_edge(vec_to_columnar(collection), errs); + let (edge, _errs) = bundle.as_collection_core(mfp, None, Antichain::new()); + // Tee the producer output for a content check, then feed the + // original edge into the arrange input. + let produced = columnar_to_vec(edge.clone()).inner.capture(); + let (_arranged, _arrange_errs, _passthrough) = + CollectionBundle::::arrange_collection( + &"arrange".to_string(), + edge, + vec![LirScalarExpr::column(0)], + vec![0], + ArrangementBatcher::Columnation, + ); + + let max_time = rows.iter().map(|(_, t)| *t).max().unwrap(); + for (row, time) in rows { + input.update_at(row, Timestamp::from(time), Diff::ONE); + } + input.advance_to(Timestamp::from(max_time + 1)); + input.flush(); + produced + }) + }); - assert!( - producer_is_columnar, - "a non-identity MFP must produce a columnar edge" - ); - assert!( - passthrough_is_columnar, - "the arrange input must keep the columnar passthrough (no ColumnarToVec)" - ); assert_eq!(extract_row_updates(produced), expected); } - /// The reworked identity fast-path returns the unarranged input edge - /// unchanged, so a columnar producer stays columnar and a `Vec` producer - /// stays `Vec` with no repack in either direction. + /// The identity fast-path returns the unarranged input edge unchanged, with no + /// repack, so its contents pass straight through. #[mz_ore::test] fn as_collection_core_identity_passes_edge_through() { - for columnar_input in [false, true] { - let is_columnar = timely::execute_directly(move |worker| { - worker.dataflow::(|scope| { - let (mut input, collection) = scope.new_collection::(); - let (_err_input, errs) = scope.new_collection::(); - let edge = if columnar_input { - CollectionEdge::Columnar(vec_to_columnar(collection)) - } else { - CollectionEdge::Vec(collection) - }; - let bundle = CollectionBundle::from_edge(edge, errs); - let identity = MapFilterProject::::new(1) - .into_plan() - .expect("identity mfp"); - let (out, _errs) = bundle.as_collection_core(identity, None, Antichain::new()); - let is_columnar = matches!(out, CollectionEdge::Columnar(_)); - input.update(Row::pack_slice(&[Datum::Int64(1)]), Diff::ONE); - input.advance_to(Timestamp::from(1u64)); - input.flush(); - is_columnar - }) - }); - assert_eq!( - is_columnar, columnar_input, - "the identity fast-path must preserve the input edge variant" - ); - } + let expected = vec![( + Row::pack_slice(&[Datum::Int64(1)]), + Timestamp::from(0u64), + Diff::ONE, + )]; + let captured = timely::execute_directly(move |worker| { + worker.dataflow::(|scope| { + let (mut input, collection) = scope.new_collection::(); + let (_err_input, errs) = scope.new_collection::(); + let bundle = CollectionBundle::from_edge(vec_to_columnar(collection), errs); + let identity = MapFilterProject::::new(1) + .into_plan() + .expect("identity mfp"); + let (out, _errs) = bundle.as_collection_core(identity, None, Antichain::new()); + let captured = columnar_to_vec(out).inner.capture(); + input.update_at( + Row::pack_slice(&[Datum::Int64(1)]), + Timestamp::from(0u64), + Diff::ONE, + ); + input.advance_to(Timestamp::from(1u64)); + input.flush(); + captured + }) + }); + assert_eq!(extract_row_updates(captured), expected); } /// The columnar producer folds within-batch duplicates: input rows that @@ -1870,13 +1755,9 @@ mod tests { worker.dataflow::(|scope| { let (mut input, collection) = scope.new_collection(); let (_err_input, errs) = scope.new_collection::(); - let bundle = CollectionBundle::from_edge(CollectionEdge::Vec(collection), errs); + let bundle = CollectionBundle::from_edge(vec_to_columnar(collection), errs); let (edge, _errs) = bundle.as_collection_core(mfp, None, Antichain::new()); - assert!( - matches!(edge, CollectionEdge::Columnar(_)), - "a non-identity MFP must produce a columnar edge" - ); - let captured = edge.into_vec().inner.capture(); + let captured = columnar_to_vec(edge).inner.capture(); // Feed all rows at the same time in one batch so the fold is // within-batch, not a downstream re-consolidation. for row in rows { @@ -1912,13 +1793,13 @@ mod tests { .collect(); expected.sort(); - let (is_columnar, captured) = timely::execute_directly(move |worker| { + let captured = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut input, collection) = scope.new_collection(); let (arranged, arr_errs, _passthrough) = CollectionBundle::::arrange_collection( &"agg".to_string(), - CollectionEdge::Vec(collection), + vec_to_columnar(collection), key.clone(), vec![1], ArrangementBatcher::Columnation, @@ -1938,8 +1819,7 @@ mod tests { ArrangementFlavor::Local(arranged, err_arranged), ); let (edge, _errs) = bundle.as_specific_collection(Some(&key)); - let is_columnar = matches!(edge, CollectionEdge::Columnar(_)); - let captured = edge.into_vec().inner.capture(); + let captured = columnar_to_vec(edge).inner.capture(); let max_time = rows.iter().map(|(_, t)| *t).max().unwrap(); for (row, time) in rows { @@ -1947,14 +1827,10 @@ mod tests { } input.advance_to(Timestamp::from(max_time + 1)); input.flush(); - (is_columnar, captured) + captured }) }); - assert!( - is_columnar, - "as_specific_collection must materialize the arrangement as a columnar edge" - ); assert_eq!(extract_row_updates(captured), expected); } } diff --git a/src/compute/src/render/flat_map.rs b/src/compute/src/render/flat_map.rs index 787006d11691e..0a2bc7f3f0359 100644 --- a/src/compute/src/render/flat_map.rs +++ b/src/compute/src/render/flat_map.rs @@ -29,7 +29,6 @@ use timely::dataflow::{Scope, Stream}; use timely::progress::Antichain; use crate::render::RenderTimestamp; -use crate::render::columnar::CollectionEdge; use crate::render::context::{CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; @@ -52,10 +51,9 @@ impl<'scope, T: crate::render::RenderTimestamp> Context<'scope, T> { // a batch. A `generate_series` can still cause unavailability if it generates many rows. let budget = COMPUTE_FLAT_MAP_FUEL.get(&self.config_set); - // The unarranged path (no key) reads the input `CollectionEdge` directly, - // so the columnar arm never decodes rows at the input. The keyed path - // materializes an existing arrangement, which `as_specific_collection` - // presents as a columnar edge. + // The unarranged path (no key) reads the input edge directly. The keyed + // path materializes an existing arrangement, which + // `as_specific_collection` presents as a columnar edge. let (edge, err_collection) = match input_key.as_deref() { None => input .collection @@ -64,17 +62,10 @@ impl<'scope, T: crate::render::RenderTimestamp> Context<'scope, T> { Some(key) => input.as_specific_collection(Some(key)), }; - let (oks, errs) = match edge { - CollectionEdge::Vec(c) => { - flat_map_stage(c.inner, scope, exprs, func, mfp_plan, until, budget) - } - CollectionEdge::Columnar(c) => { - flat_map_stage(c.inner, scope, exprs, func, mfp_plan, until, budget) - } - }; + let (oks, errs) = flat_map_stage(edge.inner, scope, exprs, func, mfp_plan, until, budget); use differential_dataflow::AsCollection; - let ok_collection = CollectionEdge::Columnar(oks.as_collection()); + let ok_collection = oks.as_collection(); let new_err_collection = errs.as_collection(); let err_collection = err_collection.concat(new_err_collection); CollectionBundle::from_edge(ok_collection, err_collection) @@ -426,56 +417,27 @@ mod tests { } #[mz_ore::test] - fn flat_map_arms_agree() { - // The `Vec` and `Column` input arms must produce identical `(row, time, - // diff)` output. Multiple timestamps and a retraction exercise time - // handling and negative diffs. The columnar arm's `Columnar::into_owned` - // for time and diff runs on every ok record here (the input decode has - // no fallible/try_extend path, so the standing fallible-key rule does - // not apply), and this comparison proves it decodes correctly. - let (vec_captured, col_captured) = timely::execute_directly(move |worker| { + fn flat_map_reads_columnar_input() { + // Reads the columnar input edge, expands the table function, and emits the + // `(row, time, diff)` output. Multiple timestamps and a retraction exercise + // time handling and negative diffs; the columnar input's + // `Columnar::into_owned` for time and diff runs on every ok record. + let captured = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut input, collection) = scope.new_collection(); - let mut captures = Vec::new(); - for columnar in [false, true] { - let (exprs, func, mfp) = flat_map_args(); - let edge = if columnar { - CollectionEdge::Columnar(vec_to_columnar(collection.clone())) - } else { - CollectionEdge::Vec(collection.clone()) - }; - let (oks, _errs) = match edge { - CollectionEdge::Vec(c) => { - let stream = c.inner; - let scope = stream.scope(); - flat_map_stage( - stream, - scope, - exprs, - func, - mfp, - Antichain::new(), - usize::MAX, - ) - } - CollectionEdge::Columnar(c) => { - let stream = c.inner; - let scope = stream.scope(); - flat_map_stage( - stream, - scope, - exprs, - func, - mfp, - Antichain::new(), - usize::MAX, - ) - } - }; - captures.push(oks.capture()); - } - let col = captures.pop().unwrap(); - let vec = captures.pop().unwrap(); + let (exprs, func, mfp) = flat_map_args(); + let stream = vec_to_columnar(collection).inner; + let scope = stream.scope(); + let (oks, _errs) = flat_map_stage( + stream, + scope, + exprs, + func, + mfp, + Antichain::new(), + usize::MAX, + ); + let captured = oks.capture(); // t=0: generate_series(1, 2); t=1: generate_series(1, 3); // t=2: retract the t=0 row. input.advance_to(Timestamp::from(0_u64)); @@ -486,17 +448,16 @@ mod tests { input.update(input_row(2), -Diff::ONE); input.advance_to(Timestamp::from(3_u64)); input.flush(); - (vec, col) + captured }) }); - let vec_updates = extract_sorted_columns(vec_captured); - assert!(!vec_updates.is_empty()); + let updates = extract_sorted_columns(captured); + assert!(!updates.is_empty()); assert!( - vec_updates.iter().any(|(_, _, d)| *d < Diff::ZERO), + updates.iter().any(|(_, _, d)| *d < Diff::ZERO), "the retraction must survive as a negative diff" ); - assert_eq!(vec_updates, extract_sorted_columns(col_captured)); } /// Decodes a capture of the columnar FlatMap output into sorted owned diff --git a/src/compute/src/render/join/delta_join.rs b/src/compute/src/render/join/delta_join.rs index b5b9e545387ea..21adabe8a1b33 100644 --- a/src/compute/src/render/join/delta_join.rs +++ b/src/compute/src/render/join/delta_join.rs @@ -39,7 +39,7 @@ use timely::dataflow::operators::vec::Map; use timely::progress::Antichain; use crate::render::RenderTimestamp; -use crate::render::columnar::{CollectionEdge, vec_to_columnar}; +use crate::render::columnar::{columnar_to_vec, vec_to_columnar}; use crate::render::context::{ArrangementFlavor, CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; use crate::typedefs::{RowRowAgent, RowRowEnter}; @@ -254,7 +254,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> { // This is the sanctioned leaf-encode; a columnar `half_join`/algorithm is // a differential-side follow-up. Non-consolidating: the per-path // finalization already consolidated whatever it consolidates. - CollectionBundle::from_edge(CollectionEdge::Columnar(vec_to_columnar(oks)), errs) + CollectionBundle::from_edge(vec_to_columnar(oks), errs) } } @@ -731,7 +731,12 @@ where .collection .clone() .expect("The unarranged collection doesn't exist."); - return build_update_stream_stream(oks.into_vec(), as_of, source_relation, initial_closure); + return build_update_stream_stream( + columnar_to_vec(oks), + as_of, + source_relation, + initial_closure, + ); }; match bundle.arrangement(&source_key) { Some(ArrangementFlavor::Local(oks, _errs)) => { diff --git a/src/compute/src/render/join/linear_join.rs b/src/compute/src/render/join/linear_join.rs index 8a52e4d14c1a6..d64843f5186d5 100644 --- a/src/compute/src/render/join/linear_join.rs +++ b/src/compute/src/render/join/linear_join.rs @@ -42,7 +42,7 @@ use timely::dataflow::{Scope, Stream}; use crate::extensions::arrange::{ArrangementBatcher, MzArrangeCore}; use crate::render::RenderTimestamp; -use crate::render::columnar::{CollectionEdge, vec_to_columnar}; +use crate::render::columnar::{CollectionEdge, columnar_to_vec, vec_to_columnar}; use crate::render::context::{ArrangementFlavor, CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; use crate::render::join::mz_join_core::mz_join_core; @@ -284,8 +284,7 @@ where // but this branch is never taken in current lowering. let name = "LinearJoinInitialization"; type CB = ConsolidatingContainerBuilder; - let (j, errs) = joined - .into_vec() + let (j, errs) = columnar_to_vec(joined) .flat_map_fallible::, CB<_>, _, _, _, _>(name, { // Reuseable allocation for unpacking. let mut datums = DatumVec::new(); @@ -336,7 +335,7 @@ where // source edge is decoded to `Vec` first (`into_vec` is the identity on the // `Vec` arm); the accumulator is already a `VecCollection`. let input = match joined { - JoinedFlavor::Edge(edge) => edge.into_vec(), + JoinedFlavor::Edge(edge) => columnar_to_vec(edge), JoinedFlavor::Collection(collection) => collection, _ => panic!("Unexpectedly arranged join output"), }; @@ -359,21 +358,15 @@ where } }); errors.push(errs); - CollectionEdge::Columnar(updates) + updates } else { - // Identity finalization: the raw output is the result, encoded to the - // columnar edge via the sanctioned leaf-encode, non-consolidating to match - // the raw output. A columnar source (single-input join) is already an edge - // and passes through with no round-trip; the `Vec` accumulator and a `Vec` - // source encode via `vec_to_columnar`. + // Identity finalization: the raw output is the result. The source edge + // (single-input join) is already columnar and passes through; the `Vec` + // accumulator encodes via `vec_to_columnar`, non-consolidating to match + // the raw output. match joined { - JoinedFlavor::Edge(CollectionEdge::Columnar(c)) => CollectionEdge::Columnar(c), - JoinedFlavor::Edge(CollectionEdge::Vec(s)) => { - CollectionEdge::Columnar(vec_to_columnar(s)) - } - JoinedFlavor::Collection(collection) => { - CollectionEdge::Columnar(vec_to_columnar(collection)) - } + JoinedFlavor::Edge(edge) => edge, + JoinedFlavor::Collection(collection) => vec_to_columnar(collection), _ => panic!("Unexpectedly arranged join output"), } }; @@ -659,14 +652,10 @@ where /// Forms the source arrangement for a streamed join input off a collection edge. /// -/// Both arms build the same `((key, value), t, d)` columnar updates and push the -/// key and value borrowed into a `ColumnBuilder`, which the `Col2Val` batcher -/// consumes. This is the zero-allocation pattern: no owned `Row` per record on -/// the ok path. The arms differ only in how records are read, the `Vec` arm from -/// the owned container and the columnar arm from the borrowed column, and in the -/// error path, which owns time and diff. The columnar arm runs whenever the input -/// edge is columnar, which the source-key path (via `as_specific_collection`) and -/// upstream producers emit. +/// Reads records from the borrowed column and pushes the key and value borrowed +/// into a `ColumnBuilder`, which the `Col2Val` batcher consumes: a zero-allocation +/// pattern with no owned `Row` per record on the ok path. Only the error path +/// owns time and diff. fn arrange_join_input<'s, T>( edge: CollectionEdge<'s, T>, stream_key: Vec, @@ -679,54 +668,49 @@ fn arrange_join_input<'s, T>( where T: Lattice + RenderTimestamp, { - let (keyed, errs) = match edge { - CollectionEdge::Vec(stream) => { - key_join_input_vec(stream.inner, stream_key, stream_thinning) - } - CollectionEdge::Columnar(stream) => stream - .inner - .unary_fallible::, _, _, _>( - Pipeline, - "LinearJoinKeyPreparation", - |_, _| { - Box::new(move |input, ok, errs| { - let mut temp_storage = RowArena::new(); - let mut key_buf = Row::default(); - let mut val_buf = Row::default(); - let mut datums = DatumVec::new(); - input.for_each(|time, data| { - let mut ok_session = ok.session_with_builder(&time); - let mut err_session = errs.session(&time); - // Rows are read from the borrowed column; the key and - // value are pushed borrowed. Time and diff are owned - // only on the error path. - for (row, time, diff) in data.borrow().into_index_iter() { - temp_storage.clear(); - let datums_local = datums.borrow_with(row); - let datums = stream_key - .iter() - .map(|e| e.eval(&datums_local, &temp_storage)); - match key_buf.packer().try_extend(datums) { - Ok(()) => { - val_buf.packer().extend( - stream_thinning.iter().map(|e| datums_local[*e]), - ); - ok_session.give(((&key_buf, &val_buf), time, diff)); - } - Err(e) => { - err_session.give(( - e.into(), - Columnar::into_owned(time), - Columnar::into_owned(diff), - )); - } + let (keyed, errs) = edge + .inner + .unary_fallible::, _, _, _>( + Pipeline, + "LinearJoinKeyPreparation", + |_, _| { + Box::new(move |input, ok, errs| { + let mut temp_storage = RowArena::new(); + let mut key_buf = Row::default(); + let mut val_buf = Row::default(); + let mut datums = DatumVec::new(); + input.for_each(|time, data| { + let mut ok_session = ok.session_with_builder(&time); + let mut err_session = errs.session(&time); + // Rows are read from the borrowed column; the key and + // value are pushed borrowed. Time and diff are owned + // only on the error path. + for (row, time, diff) in data.borrow().into_index_iter() { + temp_storage.clear(); + let datums_local = datums.borrow_with(row); + let datums = stream_key + .iter() + .map(|e| e.eval(&datums_local, &temp_storage)); + match key_buf.packer().try_extend(datums) { + Ok(()) => { + val_buf + .packer() + .extend(stream_thinning.iter().map(|e| datums_local[*e])); + ok_session.give(((&key_buf, &val_buf), time, diff)); + } + Err(e) => { + err_session.give(( + e.into(), + Columnar::into_owned(time), + Columnar::into_owned(diff), + )); } } - }); - }) - }, - ), - }; + } + }); + }) + }, + ); arrange_keyed_join_input(keyed, errs, batcher) } @@ -829,122 +813,88 @@ mod tests { ] } - /// Runs `arrange_join_input` against the same input fed once as a `Vec` edge - /// and once as a columnar edge, keying by `key` with column 1 as the value. - /// Returns the sorted ok updates (read back from each arrangement) and the - /// sorted err updates of each arm. - #[allow(clippy::type_complexity)] - fn run_both_arms( + /// Runs `arrange_join_input` against the columnar input edge, keying by `key` + /// with column 1 as the value. Returns the sorted ok updates (read back from + /// the arrangement) and the sorted err updates. + fn run_columnar( input: Vec<(Row, u64, Diff)>, key: Vec, - ) -> ( - Vec, - Vec, - Vec<(String, Timestamp, Diff)>, - Vec<(String, Timestamp, Diff)>, - ) { - let (ok_vec, ok_col, err_vec, err_col) = timely::execute_directly(move |worker| { + ) -> (Vec, Vec<(String, Timestamp, Diff)>) { + let (ok, err) = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut handle, collection) = scope.new_collection(); - let mut ok_caps = Vec::new(); - let mut err_caps = Vec::new(); - for edge in [ - CollectionEdge::Vec(collection.clone()), - CollectionEdge::Columnar(vec_to_columnar(collection)), - ] { - let (arranged, errs) = arrange_join_input( - edge, - key.clone(), - vec![1], - ArrangementBatcher::Columnation, - ); - let keyed = arranged.as_collection(|k, v| (k.to_row(), v.to_row())); - ok_caps.push(keyed.inner.capture()); - err_caps.push(errs.inner.capture()); - } - let err_col = err_caps.pop().unwrap(); - let err_vec = err_caps.pop().unwrap(); - let ok_col = ok_caps.pop().unwrap(); - let ok_vec = ok_caps.pop().unwrap(); + let (arranged, errs) = arrange_join_input( + vec_to_columnar(collection), + key, + vec![1], + ArrangementBatcher::Columnation, + ); + let keyed = arranged.as_collection(|k, v| (k.to_row(), v.to_row())); + let ok = keyed.inner.capture(); + let err = errs.inner.capture(); for (row, time, diff) in input { handle.update_at(row, Timestamp::from(time), diff); } handle.advance_to(Timestamp::from(3_u64)); handle.flush(); - (ok_vec, ok_col, err_vec, err_col) + (ok, err) }) }); - ( - extract_sorted(ok_vec), - extract_sorted(ok_col), - extract_err(err_vec), - extract_err(err_col), - ) + (extract_sorted(ok), extract_err(err)) } - /// The columnar arm of `arrange_join_input` forms the same keyed arrangement - /// as the `Vec` arm, across several distinct timestamps and a retraction. - /// - /// This proves the columnar key-forming is correct and that the columnar arm - /// ran (the input is fed through `vec_to_columnar`). It does not prove the - /// absence of a silent `ColumnarToVec` decode on the ok path: such a decode - /// would yield identical contents. No-decode holds by code inspection, the - /// columnar arm reads via `into_index_iter` and pushes the key and value - /// borrowed into the `ColumnBuilder`, never calling `into_vec`. + /// `arrange_join_input` forms the keyed arrangement from the columnar input + /// across several distinct timestamps and a retraction. /// /// The stream key here is infallible column projection, so `try_extend` never /// fails and the ok path pushes the diff borrowed. The retraction records /// exercise a negative diff on that borrowed ok path. `Columnar::into_owned` - /// runs only on the error path, which `arrange_join_input_arms_agree_on_error_path` - /// covers. + /// runs only on the error path, which `arrange_join_input_error_path` covers. #[mz_ore::test] - fn arrange_join_input_arms_agree() { - let (ok_vec, ok_col, err_vec, err_col) = - run_both_arms(test_input(), vec![LirScalarExpr::column(0)]); - assert!(!ok_vec.is_empty()); - assert_eq!(ok_vec, ok_col); - assert!(err_vec.is_empty() && err_col.is_empty()); + fn arrange_join_input_keys_correctly() { + let (ok, err) = run_columnar(test_input(), vec![LirScalarExpr::column(0)]); + assert!(!ok.is_empty()); + assert!(err.is_empty()); // A retraction survives into the arrangement, so the ok path handled a // negative (borrowed) diff. - assert!(ok_vec.iter().any(|(_, _, d)| *d < Diff::ZERO)); + assert!(ok.iter().any(|(_, _, d)| *d < Diff::ZERO)); // Key is column 0, value is column 1 (the thinning), so both are single // datums. - for ((key, value), _t, _d) in &ok_vec { + for ((key, value), _t, _d) in &ok { assert_eq!(key.iter().count(), 1); assert_eq!(value.iter().count(), 1); } } /// A key expression that always errors drives every record onto the - /// `try_extend` Err branch, exercising the columnar arm's `Columnar::into_owned` - /// reconstruction of each error's `(time, diff)`. Both arms must agree on the - /// errors, and the ok output must be empty on both. + /// `try_extend` Err branch, exercising `Columnar::into_owned` reconstruction + /// of each error's `(time, diff)`. The ok output must be empty. #[mz_ore::test] - fn arrange_join_input_arms_agree_on_error_path() { + fn arrange_join_input_error_path() { let key = vec![LirScalarExpr::literal( Err(EvalError::DivisionByZero), ReprScalarType::Int32, )]; - let (ok_vec, ok_col, err_vec, err_col) = run_both_arms(test_input(), key); - assert!(ok_vec.is_empty() && ok_col.is_empty()); - assert!(!err_vec.is_empty()); - assert_eq!(err_vec, err_col); + let (ok, err) = run_columnar(test_input(), key); + assert!(ok.is_empty()); + assert!(!err.is_empty()); } /// The bare-`VecCollection` accumulator path (`arrange_join_collection`, used - /// for join stages after the first) forms the same keyed arrangement as - /// feeding the identical input through `arrange_join_input`'s `Vec` edge arm. - /// Both share `key_join_input_vec`, so this guards the accumulator wiring, - /// not the keying. + /// for join stages after the first) forms the same keyed arrangement as the + /// columnar source edge path (`arrange_join_input`). The two use different + /// keying implementations (`arrange_join_input` keys inline off the borrowed + /// column, `arrange_join_collection` keys via `key_join_input_vec`), so this + /// cross-checks the two keying paths against each other. #[mz_ore::test] - fn arrange_join_collection_matches_vec_edge() { + fn arrange_join_collection_matches_edge() { let key = vec![LirScalarExpr::column(0)]; let input = test_input(); let (edge_ok, acc_ok) = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut handle, collection) = scope.new_collection(); let (edge_arr, _edge_errs) = arrange_join_input( - CollectionEdge::Vec(collection.clone()), + vec_to_columnar(collection.clone()), key.clone(), vec![1], ArrangementBatcher::Columnation, diff --git a/src/compute/src/render/sinks.rs b/src/compute/src/render/sinks.rs index 35518b34d11d4..fd6c6bdf8a54d 100644 --- a/src/compute/src/render/sinks.rs +++ b/src/compute/src/render/sinks.rs @@ -30,6 +30,7 @@ use timely::progress::Antichain; use crate::compute_state::SinkToken; use crate::logging::compute::LogDataflowErrors; +use crate::render::columnar::columnar_to_vec; use crate::render::context::{Context, distinct_errs_collection}; use crate::render::errors::DataflowErrorSer; use crate::render::{RenderTimestamp, StartSignal}; @@ -68,7 +69,7 @@ impl<'g, T: RenderTimestamp> Context<'g, T> { .lookup_id(mz_expr::Id::Global(sink.from)) .expect("Sink source collection not loaded"); let (ok_collection, mut err_collection) = if let Some((oks, errs)) = &bundle.collection { - (oks.clone().into_vec(), errs.clone()) + (columnar_to_vec(oks.clone()), errs.clone()) } else { let (key, _arrangement) = bundle .arranged @@ -85,7 +86,7 @@ impl<'g, T: RenderTimestamp> Context<'g, T> { // above. let (oks, errs) = bundle.as_collection_core(mfp_plan, Some((key.clone(), None)), self.until.clone()); - (oks.into_vec(), errs) + (columnar_to_vec(oks), errs) }; // Attach logging of dataflow errors. diff --git a/src/compute/src/render/top_k.rs b/src/compute/src/render/top_k.rs index a8d5657032a01..f80bd98baa96e 100644 --- a/src/compute/src/render/top_k.rs +++ b/src/compute/src/render/top_k.rs @@ -49,7 +49,7 @@ use timely::dataflow::operators::generic::builder_rc::OperatorBuilder; use crate::extensions::arrange::{ArrangementSize, KeyCollection, MzArrange}; use crate::extensions::reduce::{ClearContainer, MzReduce}; use crate::render::Pairer; -use crate::render::columnar::{CollectionEdge, vec_to_columnar}; +use crate::render::columnar::{CollectionEdge, columnar_to_vec, vec_to_columnar}; use crate::render::context::{ArrangementFlavor, CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; use crate::render::errors::MaybeValidatingRow; @@ -125,11 +125,11 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim .get(&self.config_set) .try_into() .expect("must fit"); - CollectionEdge::Columnar(vec_to_columnar(T::maybe_apply_temporal_bucketing( - ok_input.into_vec().inner, + vec_to_columnar(T::maybe_apply_temporal_bucketing( + columnar_to_vec(ok_input).inner, self.as_of_frontier.clone(), summary, - ))) + )) } else { ok_input }; @@ -161,10 +161,9 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim let mut datum_vec = mz_repr::DatumVec::new(); // A literal, non-negative limit skips this branch entirely, so this // per-row evaluation only runs for column or otherwise fallible - // limits. On the `Vec` edge `into_vec` is the identity, so the - // `Vec` path is unchanged. On a columnar edge it is a narrow - // sanctioned decode confined to this rare path. - let errors = ok_input.clone().into_vec().flat_map(move |row| { + // limits. The columnar decode is a narrow sanctioned leaf confined + // to this rare path. + let errors = columnar_to_vec(ok_input.clone()).flat_map(move |row| { let temp_storage = mz_repr::RowArena::new(); let datums = datum_vec.borrow_with(&row); match expr.eval(&datums[..], &temp_storage) { @@ -614,16 +613,13 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim /// every TopK stage carries the row through to its output. /// /// The output is a `VecCollection<(Row, Row)>` because the TopK downstream -/// (`KeyBatcher`, `MzReduce`) is `Vec`-based. The columnar arm therefore still -/// decodes the value `Row` of every record inline via `Columnar::into_owned`, -/// which costs the same as `columnar_to_vec`'s body. This is not a zero-copy -/// borrowed-push: it removes the separate `ColumnarToVec` decode operator and the -/// intermediate `Vec<(Row, T, Diff)>` container it would produce, but it does not +/// (`KeyBatcher`, `MzReduce`) is `Vec`-based, so the value `Row` of every record +/// is decoded inline via `Columnar::into_owned` (the same per-record cost as +/// `columnar_to_vec`'s body). This removes the separate `ColumnarToVec` decode +/// operator and its intermediate `Vec<(Row, T, Diff)>` container, but does not /// avoid the per-record decode, because there is no columnar batcher to push -/// borrowed rows into here. The `Vec` arm maps the collection directly and reuses -/// the owned input row, so it is behaviorally identical to consuming the -/// collection directly. The columnar arm runs whenever the TopK input edge is -/// columnar. +/// borrowed rows into here. The key is formed from the borrowed datums. Time and +/// diff are owned per record. fn map_topk_key<'s, T, L>( edge: CollectionEdge<'s, T>, name: &str, @@ -633,54 +629,38 @@ where T: crate::render::RenderTimestamp, L: FnMut(&[Datum], &Row) -> Row + 'static, { - match edge { - CollectionEdge::Vec(oks) => { - let mut datum_vec = mz_repr::DatumVec::new(); - oks.map(move |row| { - let key_row = { - let datums = datum_vec.borrow_with(&row); - key(&datums, &row) - }; - (key_row, row) - }) - } - CollectionEdge::Columnar(oks) => { - let mut builder = OperatorBuilder::new(name.to_string(), oks.inner.scope()); - let (output, stream) = builder.new_output(); - let mut output = - OutputBuilder::<_, CapacityContainerBuilder>>::from( - output, - ); - let mut input = builder.new_input(oks.inner, Pipeline); - builder.build(move |_capabilities| { - let mut datum_vec = mz_repr::DatumVec::new(); - move |_frontiers| { - let mut output = output.activate(); - input.for_each(|time, data| { - let mut session = output.session_with_builder(&time); - // The value row is decoded to an owned `Row` here because the - // output is `Vec`-based. This is the same per-record cost as - // `columnar_to_vec`, just without the separate decode operator - // and its intermediate container. The key is formed from the - // borrowed datums. Time and diff are owned per record. - for (row, t, d) in data.borrow().into_index_iter() { - let value_row: Row = Columnar::into_owned(row); - let key_row = { - let datums = datum_vec.borrow_with(&value_row); - key(&datums, &value_row) - }; - session.give(( - (key_row, value_row), - Columnar::into_owned(t), - Columnar::into_owned(d), - )); - } - }); + let mut builder = OperatorBuilder::new(name.to_string(), edge.inner.scope()); + let (output, stream) = builder.new_output(); + let mut output = + OutputBuilder::<_, CapacityContainerBuilder>>::from(output); + let mut input = builder.new_input(edge.inner, Pipeline); + builder.build(move |_capabilities| { + let mut datum_vec = mz_repr::DatumVec::new(); + move |_frontiers| { + let mut output = output.activate(); + input.for_each(|time, data| { + let mut session = output.session_with_builder(&time); + // The value row is decoded to an owned `Row` here because the + // output is `Vec`-based. This is the same per-record cost as + // `columnar_to_vec`, just without the separate decode operator + // and its intermediate container. The key is formed from the + // borrowed datums. Time and diff are owned per record. + for (row, t, d) in data.borrow().into_index_iter() { + let value_row: Row = Columnar::into_owned(row); + let key_row = { + let datums = datum_vec.borrow_with(&value_row); + key(&datums, &value_row) + }; + session.give(( + (key_row, value_row), + Columnar::into_owned(t), + Columnar::into_owned(d), + )); } }); - stream.as_collection() } - } + }); + stream.as_collection() } /// Drops the hash-key pairing from a consolidated `(hash_key, row)` TopK result, @@ -709,7 +689,7 @@ where }); } }); - CollectionEdge::Columnar(stream.as_collection()) + stream.as_collection() } /// Build a stage of a topk reduction. Maintains the _retractions_ of the output instead of emitted @@ -1272,7 +1252,7 @@ mod tests { use timely::dataflow::operators::capture::{Event, Extract}; use super::*; - use crate::render::columnar::vec_to_columnar; + use crate::render::columnar::{columnar_to_vec, vec_to_columnar}; type KeyedUpdate = ((Row, Row), Timestamp, Diff); type Captured = std::sync::mpsc::Receiver>>; @@ -1328,65 +1308,50 @@ mod tests { ] } - /// Runs `map_topk_key` against the same input fed once as a `Vec` edge and - /// once as a columnar edge, forming a hash-and-group key exactly as - /// `build_topk` does. Returns the sorted `(key, value)` updates of each arm. - fn run_both_arms(input: Vec<(Row, u64, Diff)>) -> (Vec, Vec) { - let (vec, col) = timely::execute_directly(move |worker| { + /// Runs `map_topk_key` against the columnar input edge, forming a + /// hash-and-group key exactly as `build_topk` does. Returns the sorted + /// `(key, value)` updates. + fn run_columnar(input: Vec<(Row, u64, Diff)>) -> Vec { + let captured = timely::execute_directly(move |worker| { worker.dataflow::(|scope| { let (mut handle, collection) = scope.new_collection(); - let mut captures = Vec::new(); - for edge in [ - CollectionEdge::Vec(collection.clone()), - CollectionEdge::Columnar(vec_to_columnar(collection)), - ] { - let pairer = Pairer::new(1); - let group_key = [0usize]; - let keyed = map_topk_key(edge, "test", move |datums, row| { + let pairer = Pairer::new(1); + let group_key = [0usize]; + let keyed = + map_topk_key(vec_to_columnar(collection), "test", move |datums, row| { let hash = row.hashed(); let iterator = group_key.iter().map(|i| datums[*i]); pairer.merge(std::iter::once(Datum::from(hash)), iterator) }); - captures.push(keyed.inner.capture()); - } - let col = captures.pop().unwrap(); - let vec = captures.pop().unwrap(); + let captured = keyed.inner.capture(); for (row, time, diff) in input { handle.update_at(row, Timestamp::from(time), diff); } handle.advance_to(Timestamp::from(3_u64)); handle.flush(); - (vec, col) + captured }) }); - (extract_sorted(vec), extract_sorted(col)) + extract_sorted(captured) } - /// The columnar arm of `map_topk_key` forms the same `(key, value)` updates - /// as the `Vec` arm, across several distinct timestamps. - /// - /// This proves the columnar key-forming is correct and that the columnar arm - /// ran (the input is fed through `vec_to_columnar`). It does not prove the - /// absence of a silent `ColumnarToVec` decode on the ok path: such a decode - /// would yield identical contents. No-decode holds by code inspection, the - /// columnar arm reads via `into_index_iter` and never calls `into_vec`. + /// `map_topk_key` forms the `(key, value)` updates from the columnar input + /// across several distinct timestamps. /// - /// The key closure here is infallible (projection plus hash plus pack), so - /// there is no fallible-key path to test, unlike the arrange operator's key. - /// `Columnar::into_owned` for time and diff runs on the ok path for every + /// The key closure is infallible (projection plus hash plus pack), so there is + /// no fallible-key path. `Columnar::into_owned` for time and diff runs on every /// record, so the multi-timestamp, mixed-sign input exercises it on both /// positive and negative diffs. #[mz_ore::test] - fn map_topk_key_arms_agree() { - let (vec_updates, col_updates) = run_both_arms(test_input()); - assert!(!vec_updates.is_empty()); - assert_eq!(vec_updates, col_updates); - // Retractions reach the operator, so the columnar arm decoded a negative - // diff via `Columnar::into_owned`. - assert!(vec_updates.iter().any(|(_, _, d)| *d < Diff::ZERO)); + fn map_topk_key_forms_key() { + let updates = run_columnar(test_input()); + assert!(!updates.is_empty()); + // Retractions reach the operator, so a negative diff was decoded via + // `Columnar::into_owned`. + assert!(updates.iter().any(|(_, _, d)| *d < Diff::ZERO)); // The value is the full input row. The key is `(hash, group_column)`, so // the group component mirrors column 0 of the value row. - for ((key, value), _t, _d) in &vec_updates { + for ((key, value), _t, _d) in &updates { let key_datums: Vec<_> = key.iter().collect(); let value_datums: Vec<_> = value.iter().collect(); assert_eq!(key_datums.len(), 2); @@ -1427,21 +1392,19 @@ mod tests { .collect(); expected.sort(); - let (is_columnar, captured) = timely::execute_directly(move |worker| { + let 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(); + let captured = columnar_to_vec(edge).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) + captured }) }); - assert!(is_columnar, "the TopK output must be a columnar edge"); let mut got: Vec<(Row, Timestamp, Diff)> = captured .extract() From d53b433aa397cc2520224e998d98939dc4fb441f Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Wed, 22 Jul 2026 10:40:49 +0200 Subject: [PATCH 2/5] doc: use plain code span for consolidate_named cross-reference After the enum collapse the Vec arm is gone, so CollectionExt is no longer in scope and the intra-doc link cannot resolve. Reference it as a plain code span, matching the sibling comment. Co-Authored-By: Claude Opus 4.8 --- src/compute/src/render/columnar.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compute/src/render/columnar.rs b/src/compute/src/render/columnar.rs index 9dffb52346fba..8a3908cf04cb5 100644 --- a/src/compute/src/render/columnar.rs +++ b/src/compute/src/render/columnar.rs @@ -165,7 +165,7 @@ where /// Consolidates a [`ColumnarCollection`] natively, without a row round-trip. /// -/// Mirrors the `Vec` arm's [`CollectionExt::consolidate_named`], but keeps the +/// Mirrors the `Vec` arm's `CollectionExt::consolidate_named`, but keeps the /// data columnar throughout: a [`ColumnChunker`] sorts and consolidates the /// input columns and a [`ColumnMergeBatcher`] merges them under a /// [`columnar_exchange_data`] pact. Both hold their data in [`Column`], and the From 836a1a3732e8f26d84a9668ea04d9d8e15fae468 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Wed, 22 Jul 2026 18:23:24 +0200 Subject: [PATCH 3/5] doc: describe leaf encode/decode as permanent sanctioned boundary Reword the vec_to_columnar/columnar_to_vec docs to describe them as the permanent leaf encode/decode for row-serializing leaves rather than a transitional seam. Split the relocated map_topk_key doc semicolon. Comment-only, no behavior change. Co-Authored-By: Claude Opus 4.8 --- src/compute/src/render/columnar.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/compute/src/render/columnar.rs b/src/compute/src/render/columnar.rs index 8a3908cf04cb5..b68b5c45e6b10 100644 --- a/src/compute/src/render/columnar.rs +++ b/src/compute/src/render/columnar.rs @@ -232,8 +232,11 @@ where /// Repacks a row-based collection into columnar batches. /// -/// A transitional seam-healer, visible in rendered dataflows as a -/// `VecToColumnar` operator. Repacking copies row bytes but allocates no +/// The sanctioned leaf encode from `Vec` to the columnar edge, visible in +/// rendered dataflows as a `VecToColumnar` operator. Row-serializing leaves +/// (sinks, `LetRec`, temporal bucketing, join internals, TopK fallible-limit) +/// stay `Vec`-shaped internally and encode to the columnar edge at their +/// boundary through this pass. Repacking copies row bytes but allocates no /// per-record `Row`s. pub fn vec_to_columnar<'scope, T>( collection: VecCollection<'scope, T, Row, Diff>, @@ -262,10 +265,12 @@ where /// Decodes columnar batches into a row-based collection. /// -/// A transitional seam-healer, visible in rendered dataflows as a -/// `ColumnarToVec` operator. Decoding allocates an owned [`Row`] per record, -/// so it should only guard consumers that have not yet learned the columnar -/// form. +/// The sanctioned leaf decode from the columnar edge to `Vec`, visible in +/// rendered dataflows as a `ColumnarToVec` operator. Row-serializing leaves +/// (sinks, `LetRec`, temporal bucketing, join internals, TopK fallible-limit) +/// decode the columnar edge at their boundary through this pass. Decoding +/// allocates an owned [`Row`] per record, so it stays confined to those leaf +/// boundaries. pub fn columnar_to_vec<'scope, T>( collection: ColumnarCollection<'scope, T, Row, Diff>, ) -> VecCollection<'scope, T, Row, Diff> From 2f94b33055aebe199b49afbdda426d38d068e32b Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Wed, 22 Jul 2026 19:48:42 +0200 Subject: [PATCH 4/5] doc: de-plan columnar module doc and split comment semicolons Comment-only, no behavior change. Co-Authored-By: Claude Opus 4.8 --- src/compute/src/render/columnar.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/compute/src/render/columnar.rs b/src/compute/src/render/columnar.rs index b68b5c45e6b10..9c97758afeab7 100644 --- a/src/compute/src/render/columnar.rs +++ b/src/compute/src/render/columnar.rs @@ -13,7 +13,7 @@ //! edges between Plan nodes carry. Every producer emits this representation. //! //! Within a Plan node, operators may freely materialize `Vec` collections. Only -//! the inter-node edge format is constrained. A node that produces a row-based +//! the collection edge format is constrained. A node that produces a row-based //! collection re-encodes it to the columnar edge at its output leaf via //! [`vec_to_columnar`]. A node that must consume rows decodes at its input leaf //! via [`columnar_to_vec`]. Both are named operators (`VecToColumnar`, @@ -48,7 +48,7 @@ use crate::render::errors::DataflowErrorSer; pub type ColumnarCollection<'scope, T, D, R> = Collection<'scope, T, Column<(D, T, R)>>; /// A dataflow edge between Plan nodes: a columnar collection of `(Row, Diff)` -/// updates. Every producer emits this representation; a row-based producer +/// updates. Every producer emits this representation. A row-based producer /// re-encodes to it at its output leaf via [`vec_to_columnar`]. pub type CollectionEdge<'scope, T> = ColumnarCollection<'scope, T, Row, Diff>; @@ -65,7 +65,7 @@ where /// Applies `logic` to each record in `edge`, exposing the record as a borrowed /// [`DatumVecBorrow`] and giving it ok and err output sessions. /// -/// `max_demand` bounds the number of columns decoded per row; pass `usize::MAX` +/// `max_demand` bounds the number of columns decoded per row. Pass `usize::MAX` /// to decode all columns. /// /// This is the canonical entry point for "decoding consumers" (operators that @@ -104,8 +104,8 @@ where let mut ok_output = ok_output.activate(); let mut err_output = err_output.activate(); input.for_each(|time, data| { - // Retain the input capability to derive a `Capability` for each output; - // the `Session` type alias is fixed to `Capability`. + // Retain the input capability to derive a `Capability` for each + // output. The `Session` type alias is fixed to `Capability`. let ok_cap = time.retain(0); let err_cap = time.retain(1); let mut ok_session = ok_output.session_with_builder(&ok_cap); From c4f0a5df450fa7943fe401d3f0d27be71a8eda73 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Wed, 19 Aug 2026 15:08:11 +0200 Subject: [PATCH 5/5] compute: finish the collapse's end state Six leftovers from collapsing the edge enum, none of which a per-commit review could see, since each only becomes wrong once the enum is gone. Read the delta-join raw source and the TopK fallible limit from the borrowed column. Both consumers only read datums, so neither needs the owned rows that decoding to `Vec` built. The limit path is rarer, but it tees off the hot input, so the decode doubled that input's downstream work on those plans. Negate the diff column rather than re-encoding every record. Rows and times are handed over untouched for a typed input and copied in bulk for a serialized one, where before every row and time was pushed through a builder. Negation stays checked, so `-Diff::MIN` still reports overflow. Delete the FlatMap batch trait. Its `Vec` impl lost its caller here, and the remaining arm needs no abstraction. Its fuel test was passing a `Vec` stream, so the "fueling preserved" claim was being discharged against the dead arm; both tests now feed the columnar edge the operator is given in production. Give the accumulator's keying stage its own name. It shared "LinearJoinKeyPreparation" with the source edge's stage, so introspection could not tell the two apart, and its doc still claimed a caller it no longer has. Correct a doc that referred to `into_vec`, which no longer exists. --- src/compute/src/render/columnar.rs | 76 ++++++++++++++++------ src/compute/src/render/context.rs | 6 +- src/compute/src/render/flat_map.rs | 74 +++++++-------------- src/compute/src/render/join/delta_join.rs | 49 +++++++++----- src/compute/src/render/join/linear_join.rs | 21 +++--- src/compute/src/render/top_k.rs | 41 ++++++++---- 6 files changed, 151 insertions(+), 116 deletions(-) diff --git a/src/compute/src/render/columnar.rs b/src/compute/src/render/columnar.rs index 9c97758afeab7..923675b4bf612 100644 --- a/src/compute/src/render/columnar.rs +++ b/src/compute/src/render/columnar.rs @@ -20,7 +20,7 @@ //! `ColumnarToVec`), so those leaf seams stay visible in dataflow //! introspection. -use columnar::{Columnar, Index}; +use columnar::{Borrow, Columnar, Container, Index, Len, Push}; use differential_dataflow::{AsCollection, Collection, VecCollection}; use mz_repr::{DatumVec, DatumVecBorrow, Diff, Row}; use mz_timely_util::columnar::Column; @@ -127,16 +127,56 @@ where (ok_stream, err_stream) } -/// Negates the diff of every record in a [`ColumnarCollection`]. +/// Negates the diff of every record in `column`, rebuilding only the diff column. +/// +/// A `Typed` input hands its row and time columns over untouched, so only the +/// diffs are rebuilt, which is 8 bytes per record. A serialized input keeps all +/// three columns in one buffer, so its rows and times are copied in bulk. /// -/// Rows and times are pushed from their borrowed forms. Only the diff is -/// materialized, and it is `Copy`. +/// Negation stays checked: `Neg for Overflowing` runs `overflowing_neg` and +/// reports overflow, which `-Diff::MIN` triggers. /// -/// TODO: Rebuild only the diff column. Borrow the input column, build one owned -/// negated diff column from the borrowed diffs, and re-encode using the borrowed -/// row and time columns directly, so row and time bytes are copied once rather -/// than pushed per record. The serialized (`Align` / `Bytes`) input case needs -/// care, since all columns share a single buffer. +/// TODO: Negate a `Typed` input's diffs in place rather than rebuilding them. +/// This cannot go through `columnar::IndexMut`: `Overflows` stores the raw +/// integer and materializes `Overflowing` on read, so there is no wrapper in +/// memory to borrow mutably, and handing out `&mut` to the raw integer would let +/// writes bypass the checked arithmetic. It needs a checked bulk operation on the +/// container instead, keeping the overflow check inside. +fn negate_column(column: Column<(Row, T, Diff)>) -> Column<(Row, T, Diff)> +where + T: RenderTimestamp, +{ + /// Collects the negation of every diff in `diffs` into a fresh column. + fn negated_diffs<'a, D>(diffs: &'a D) -> ::Container + where + D: Len + Index + 'a, + { + let mut negated = ::Container::default(); + for index in 0..diffs.len() { + negated.push(-diffs.get(index)); + } + negated + } + + match column { + Column::Typed((rows, times, diffs)) => { + let negated = negated_diffs(&diffs.borrow()); + Column::Typed((rows, times, negated)) + } + column => { + let view = column.borrow(); + let len = view.len(); + let mut negated = <(Row, T, Diff) as Columnar>::Container::default(); + let (rows, times, diffs) = &mut negated; + rows.extend_from_self(view.0, 0..len); + times.extend_from_self(view.1, 0..len); + *diffs = negated_diffs(&view.2); + Column::Typed(negated) + } + } +} + +/// Negates the diff of every record in a [`ColumnarCollection`]. pub fn columnar_negate<'scope, T>( collection: ColumnarCollection<'scope, T, Row, Diff>, ) -> ColumnarCollection<'scope, T, Row, Diff> @@ -145,17 +185,14 @@ where { collection .inner - .unary::, _, _, _>( + .unary::>, _, _, _>( Pipeline, "ColumnarNegate", |_cap, _info| { move |input, output| { input.for_each(|time, data| { - let mut session = output.session_with_builder(&time); - for (v, t, d) in data.borrow().into_index_iter() { - let d = -Diff::into_owned(d); - session.give((v, t, &d)); - } + let mut negated = negate_column(std::mem::take(data)); + output.session(&time).give_container(&mut negated); }); } }, @@ -165,7 +202,7 @@ where /// Consolidates a [`ColumnarCollection`] natively, without a row round-trip. /// -/// Mirrors the `Vec` arm's `CollectionExt::consolidate_named`, but keeps the +/// Mirrors the row-based `CollectionExt::consolidate_named`, but keeps the /// data columnar throughout: a [`ColumnChunker`] sorts and consolidates the /// input columns and a [`ColumnMergeBatcher`] merges them under a /// [`columnar_exchange_data`] pact. Both hold their data in [`Column`], and the @@ -208,10 +245,9 @@ where // // TODO: This ships a whole sealed snapshot in one activation, an un-fueled // burst hazard on large consolidations. It is the same behavior as the - // `Vec` arm's `consolidate_named` unpack (see - // `mz_timely_util::operator::consolidate_named`), not new here. A future - // fuel fix should cover both arms, so the burst is not fixed on one and - // left on the other. + // row-based `mz_timely_util::operator::consolidate_named` unpack, not new + // here. A future fuel fix should cover both, so the burst is not fixed on + // one and left on the other. consolidated .unary::>, _, _, _>( Pipeline, diff --git a/src/compute/src/render/context.rs b/src/compute/src/render/context.rs index 36092a2f8c1fe..9fcb2eaee8852 100644 --- a/src/compute/src/render/context.rs +++ b/src/compute/src/render/context.rs @@ -1521,8 +1521,8 @@ mod tests { // `DataflowErrorSer` is not `Ord`, so project the error to its debug string // to get a stable, comparable ordering. The time and diff still ride along, - // so this verifies the columnar arm's `into_owned` on the error path - // reconstructs the same `(time, diff)` as the `Vec` arm. + // so an assertion on the result also covers `into_owned`'s reconstruction of + // them on the error path. fn extract_err(captured: Captured) -> Vec<(String, Timestamp, Diff)> { let mut updates: Vec<_> = captured .extract() @@ -1782,7 +1782,7 @@ mod tests { /// by column 0 and thinning the value to column 1 reconstructs the original /// two-column row. No-decode is a by-inspection property: the fueled path /// builds a `ColumnBuilder` via `flat_map_ok` and never calls - /// `columnar_to_vec`; the `into_vec` below is the capture harness only. + /// `columnar_to_vec`; the decode below is the capture harness only. #[mz_ore::test] fn as_specific_collection_materializes_columnar() { let rows = test_rows(); diff --git a/src/compute/src/render/flat_map.rs b/src/compute/src/render/flat_map.rs index 0a2bc7f3f0359..89729286bfab1 100644 --- a/src/compute/src/render/flat_map.rs +++ b/src/compute/src/render/flat_map.rs @@ -20,8 +20,6 @@ use mz_repr::{Diff, Row, RowRef, Timestamp}; use mz_timely_util::columnar::Column; use mz_timely_util::columnar::consolidate::ConsolidatingColumnBuilder; use mz_timely_util::operator::StreamExt; -use timely::Container; -use timely::container::DrainContainer; use timely::dataflow::channels::pact::Pipeline; use timely::dataflow::operators::Capability; use timely::dataflow::operators::generic::Session; @@ -81,48 +79,16 @@ type FlatMapOk = ConsolidatingColumnBuilder; /// Output err-session container builder for [`flat_map_stage`]. type FlatMapErr = ConsolidatingContainerBuilder>; -/// Yields the `(row, time, diff)` records of one queued FlatMap input batch. +/// The fueled FlatMap operator. /// -/// The two edge arms differ only in how records are read from a queued batch: -/// the `Vec` arm drains owned rows, the `Column` arm iterates the borrowed -/// column and never materializes an owned [`Row`]. Everything else in the -/// FlatMap operator (the fuel queue, budget, and re-activation) is shared -/// through [`flat_map_stage`]. -trait FlatMapBatch { - /// Calls `logic` once per record, presenting the row as a borrowed - /// [`RowRef`] and the time and diff by reference. - fn for_each_record(&mut self, logic: impl FnMut(&RowRef, &T, &Diff)); -} - -impl FlatMapBatch for Vec<(Row, T, Diff)> { - fn for_each_record(&mut self, mut logic: impl FnMut(&RowRef, &T, &Diff)) { - for (row, time, diff) in self.drain(..) { - logic(&row, &time, &diff); - } - } -} - -impl FlatMapBatch for Column<(Row, T, Diff)> { - fn for_each_record(&mut self, mut logic: impl FnMut(&RowRef, &T, &Diff)) { - // Rows are read from the borrowed column, never materialized as owned - // `Row`s. Times and diffs are owned only to hand `logic` a reference. - for (row, t, d) in self.borrow().into_index_iter() { - logic(row, &Columnar::into_owned(t), &Columnar::into_owned(d)); - } - } -} - -/// The fueled FlatMap operator, generic over the input edge arm. -/// -/// This is the sole owner of the fuel machinery, so the `Vec` and `Column` -/// arms cannot drift apart. Incoming batches are queued; each activation -/// processes queued batches and drains each record's table-function expansion -/// through the mfp, decrementing a per-activation `budget`. When the budget is -/// exhausted the operator re-activates itself and stops, deferring the rest of -/// the queue to a later activation. This bounds the work a single -/// `generate_series` can do before yielding the worker. -fn flat_map_stage<'scope, T, C>( - stream: Stream<'scope, T, C>, +/// Incoming batches are queued; each activation processes queued batches and +/// drains each record's table-function expansion through the mfp, decrementing +/// a per-activation `budget`. When the budget is exhausted the operator +/// re-activates itself and stops, deferring the rest of the queue to a later +/// activation. This bounds the work a single `generate_series` can do before +/// yielding the worker. +fn flat_map_stage<'scope, T>( + stream: Stream<'scope, T, Column<(Row, T, Diff)>>, scope: Scope<'scope, T>, exprs: Vec, func: TableFunc, @@ -135,7 +101,6 @@ fn flat_map_stage<'scope, T, C>( ) where T: RenderTimestamp, - C: Container + DrainContainer + Clone + Default + FlatMapBatch + 'static, { stream.unary_fallible::, FlatMapErr, _, _>( Pipeline, @@ -156,15 +121,20 @@ where queue.push_back((cap.retain(0), cap.retain(1), std::mem::take(data))) }); - while let Some((ok_cap, err_cap, mut data)) = queue.pop_front() { + while let Some((ok_cap, err_cap, data)) = queue.pop_front() { let mut ok_session = ok_output.session_with_builder(&ok_cap); let mut err_session = err_output.session_with_builder(&err_cap); - data.for_each_record(|input_row, time, diff| { + // Rows are read from the borrowed column, never materialized + // as owned `Row`s. Times and diffs are owned only to pass + // them by reference. + for (input_row, t, d) in data.borrow().into_index_iter() { + let time = Columnar::into_owned(t); + let diff = Columnar::into_owned(d); process_flat_map_row( input_row, - time, - diff, + &time, + &diff, &exprs, &func, &mfp_plan, @@ -176,7 +146,7 @@ where &mut err_session, &mut budget, ); - }); + } if budget == 0 { activator.activate(); break; @@ -359,7 +329,9 @@ mod tests { let mut input = worker.dataflow::(|scope| { let (input, collection) = scope.new_collection(); let (exprs, func, mfp) = flat_map_args(); - let stream = collection.inner; + // Feed the operator the columnar edge it is given in production, + // so the fuel assertions below cover the shipped path. + let stream = vec_to_columnar(collection).inner; let scope = stream.scope(); let (oks, _errs) = flat_map_stage(stream, scope, exprs, func, mfp, Antichain::new(), budget); @@ -508,7 +480,7 @@ mod tests { .project(vec![2]) .into_plan() .expect("project mfp"); - let stream = collection.inner; + let stream = vec_to_columnar(collection).inner; let scope = stream.scope(); let (oks, _errs) = flat_map_stage( stream, diff --git a/src/compute/src/render/join/delta_join.rs b/src/compute/src/render/join/delta_join.rs index 21adabe8a1b33..399d1dbb9e4f0 100644 --- a/src/compute/src/render/join/delta_join.rs +++ b/src/compute/src/render/join/delta_join.rs @@ -39,7 +39,7 @@ use timely::dataflow::operators::vec::Map; use timely::progress::Antichain; use crate::render::RenderTimestamp; -use crate::render::columnar::{columnar_to_vec, vec_to_columnar}; +use crate::render::columnar::{CollectionEdge, flat_map_datums, vec_to_columnar}; use crate::render::context::{ArrangementFlavor, CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; use crate::typedefs::{RowRowAgent, RowRowEnter}; @@ -731,12 +731,7 @@ where .collection .clone() .expect("The unarranged collection doesn't exist."); - return build_update_stream_stream( - columnar_to_vec(oks), - as_of, - source_relation, - initial_closure, - ); + return build_update_stream_stream(oks, as_of, source_relation, initial_closure); }; match bundle.arrangement(&source_key) { Some(ArrangementFlavor::Local(oks, _errs)) => { @@ -881,7 +876,7 @@ where /// first relation can be seeded from a raw collection, since the as-of filtering that the other /// paths rely on is only available from an arrangement's times. We assert that here. fn build_update_stream_stream<'scope, T>( - stream: VecCollection<'scope, T, Row, Diff>, + edge: CollectionEdge<'scope, T>, _as_of: Antichain, source_relation: usize, initial_closure: JoinClosure, @@ -897,18 +892,38 @@ where assert_eq!(source_relation, 0); type CB = ConsolidatingContainerBuilder; - stream.flat_map_fallible::, CB<_>, _, _, _, _>("UpdateStream", { - // Reuseable allocation for unpacking. - let mut datums = DatumVec::new(); - move |row| { + // The closure reads datums and builds a fresh row, so the input row is only + // ever borrowed. Reading it from the column directly keeps this path from + // materializing an owned `Row` per record. + let (oks, errs) = flat_map_datums::<_, CB>, _>(edge, usize::MAX, { + let mut datum_vec = DatumVec::new(); + move |row_datums, time, diff, ok_session, err_session| { let mut row_builder = SharedRow::get(); let temp_storage = RowArena::new(); - let mut datums_local = datums.borrow_with(&row); - initial_closure - .apply(&mut datums_local, &temp_storage, &mut row_builder) + // `JoinClosure::apply` unifies the lifetimes of `&self`, the datums, + // and the arena. Copying the datums into a local vec lets that + // lifetime shrink to this call. The copy moves datum references, not + // row data. + let mut datums = datum_vec.borrow(); + datums.extend(row_datums.iter()); + // `cloned` detaches the result from `temp_storage` and the shared row + // builder, both of which drop at the end of this call. + match initial_closure + .apply(&mut datums, &temp_storage, &mut row_builder) .map(|row| row.cloned()) - .map_err(DataflowErrorSer::from) .transpose() + { + Some(Ok(row)) => { + ok_session.give((row, time, diff)); + 1 + } + None => 0, + Some(Err(e)) => { + err_session.give((DataflowErrorSer::from(e), time, diff)); + 1 + } + } } - }) + }); + (oks.as_collection(), errs.as_collection()) } diff --git a/src/compute/src/render/join/linear_join.rs b/src/compute/src/render/join/linear_join.rs index d64843f5186d5..a9fefcf57ee83 100644 --- a/src/compute/src/render/join/linear_join.rs +++ b/src/compute/src/render/join/linear_join.rs @@ -279,9 +279,8 @@ where // If there is no starting arrangement, then we can run filters // directly on the starting collection. // If there is only one input, we are done joining, so run filters. - // `into_vec` is the identity on the `Vec` arm, so this is - // unchanged for `Vec` sources; a columnar source decodes here, - // but this branch is never taken in current lowering. + // The closure is `Vec`-internal, so the edge decodes here. + // This branch is never taken in current lowering. let name = "LinearJoinInitialization"; type CB = ConsolidatingContainerBuilder; let (j, errs) = columnar_to_vec(joined) @@ -332,8 +331,7 @@ where // The finalization closure computes fresh output rows, so build them into // a `ConsolidatingColumnBuilder` (owned give), matching the prior // `ConsolidatingContainerBuilder` and folding within-batch duplicates. A - // source edge is decoded to `Vec` first (`into_vec` is the identity on the - // `Vec` arm); the accumulator is already a `VecCollection`. + // A source edge decodes to `Vec` first; the accumulator already is one. let input = match joined { JoinedFlavor::Edge(edge) => columnar_to_vec(edge), JoinedFlavor::Collection(collection) => collection, @@ -560,9 +558,9 @@ where /// /// The key and value are pushed borrowed into a `ColumnBuilder`, so the ok path /// materializes no owned `Row` per record. The error path owns time and diff. -/// Shared by the `Vec` arm of [`arrange_join_input`] -/// (source edge) and by [`arrange_join_collection`] (the intra-operator -/// accumulator), both of which key a `Vec`-formatted stream. +/// Called by [`arrange_join_collection`] for the intra-operator accumulator, +/// which is row-formatted. [`arrange_join_input`] does the same job for the +/// columnar source edge, reading records from the borrowed column instead. fn key_join_input_vec<'s, T>( stream: Stream<'s, T, Vec<(Row, T, Diff)>>, stream_key: Vec, @@ -576,7 +574,7 @@ where { stream.unary_fallible::, _, _, _>( Pipeline, - "LinearJoinKeyPreparation", + "LinearJoinAccumulatorKeyPreparation", |_, _| { Box::new(move |input, ok, errs| { let mut temp_storage = RowArena::new(); @@ -760,9 +758,8 @@ mod tests { } // `DataflowErrorSer` is not `Ord`, so project the error to its debug string - // for a stable ordering. The time and diff ride along, so this verifies the - // columnar arm's `into_owned` on the error path reconstructs the same - // `(time, diff)` as the `Vec` arm. + // for a stable ordering. The time and diff ride along, so an assertion on the + // result also covers `into_owned`'s reconstruction of them on the error path. fn extract_err(captured: Captured) -> Vec<(String, Timestamp, Diff)> { let mut updates: Vec<_> = captured .extract() diff --git a/src/compute/src/render/top_k.rs b/src/compute/src/render/top_k.rs index f80bd98baa96e..e86c6e09dfe13 100644 --- a/src/compute/src/render/top_k.rs +++ b/src/compute/src/render/top_k.rs @@ -49,7 +49,7 @@ use timely::dataflow::operators::generic::builder_rc::OperatorBuilder; use crate::extensions::arrange::{ArrangementSize, KeyCollection, MzArrange}; use crate::extensions::reduce::{ClearContainer, MzReduce}; use crate::render::Pairer; -use crate::render::columnar::{CollectionEdge, columnar_to_vec, vec_to_columnar}; +use crate::render::columnar::{CollectionEdge, columnar_to_vec, flat_map_datums, vec_to_columnar}; use crate::render::context::{ArrangementFlavor, CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; use crate::render::errors::MaybeValidatingRow; @@ -158,23 +158,38 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim // the expression might still return a negative limit and // thus needs to be checked. let expr = expr.clone(); - let mut datum_vec = mz_repr::DatumVec::new(); // A literal, non-negative limit skips this branch entirely, so this // per-row evaluation only runs for column or otherwise fallible - // limits. The columnar decode is a narrow sanctioned leaf confined - // to this rare path. - let errors = columnar_to_vec(ok_input.clone()).flat_map(move |row| { - let temp_storage = mz_repr::RowArena::new(); - let datums = datum_vec.borrow_with(&row); - match expr.eval(&datums[..], &temp_storage) { - Ok(l) if l != Datum::Null && l.unwrap_int64() < 0 => { - Some(EvalError::NegLimit.into()) + // limits. The limit expression only reads datums, so they come + // from the borrowed column and no owned `Row` is built. This + // stage emits errors only; its ok output stays empty. + let (_, errors) = flat_map_datums::< + _, + CapacityContainerBuilder>, + _, + >(ok_input.clone(), usize::MAX, { + let mut datum_vec = mz_repr::DatumVec::new(); + move |row_datums, time, diff, _ok_session, err_session| { + let temp_storage = mz_repr::RowArena::new(); + // `eval` unifies the lifetimes of the expression, the + // datums, and the arena. Copying the datums into a local + // vec lets that lifetime shrink to this call. + let mut datums = datum_vec.borrow(); + datums.extend(row_datums.iter()); + match expr.eval(&datums[..], &temp_storage) { + Ok(l) if l != Datum::Null && l.unwrap_int64() < 0 => { + err_session.give((EvalError::NegLimit.into(), time, diff)); + 1 + } + Ok(_) => 0, + Err(e) => { + err_session.give((e.into(), time, diff)); + 1 + } } - Ok(_) => None, - Err(e) => Some(e.into()), } }); - err_collection = err_collection.concat(errors); + err_collection = err_collection.concat(errors.as_collection()); } }