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..923675b4bf612 100644 --- a/src/compute/src/render/columnar.rs +++ b/src/compute/src/render/columnar.rs @@ -9,28 +9,18 @@ //! 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 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`, +//! `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; @@ -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,216 +47,136 @@ 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(), - } - } +/// 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>; - /// 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), - } - } +/// 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) +} - /// 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)), +/// 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, + ); + } + }); } - } - - /// 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::>(), - )) - } + }); + (ok_stream, err_stream) +} - /// 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)>, - ) +/// 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. +/// +/// Negation stays checked: `Neg for Overflowing` runs `overflowing_neg` and +/// reports overflow, which `-Diff::MIN` triggers. +/// +/// 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 - DCB: ContainerBuilder, - L: for<'a> FnMut( - &'a mut DatumVecBorrow<'_>, - T, - Diff, - &mut Session, - &mut Session>, - ) -> usize - + 'static, + D: Len + Index + 'a, { - 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) - } + let mut negated = ::Container::default(); + for index in 0..diffs.len() { + negated.push(-diffs.get(index)); } + negated } - /// 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)), + 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`]. -/// -/// Rows and times are pushed from their borrowed forms. Only the diff is -/// materialized, and it is `Copy`. -/// -/// 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. pub fn columnar_negate<'scope, T>( collection: ColumnarCollection<'scope, T, Row, Diff>, ) -> ColumnarCollection<'scope, T, Row, Diff> @@ -276,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); }); } }, @@ -296,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 @@ -339,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, @@ -363,8 +268,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>, @@ -393,10 +301,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> @@ -487,7 +397,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 +410,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 +440,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 +463,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 +505,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 +523,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..9fcb2eaee8852 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 = @@ -1579,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() @@ -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 { @@ -1901,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(); @@ -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..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; @@ -29,7 +27,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 +49,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 +60,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) @@ -90,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 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. +/// The fueled FlatMap operator. /// -/// 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, @@ -144,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, @@ -165,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, @@ -185,7 +146,7 @@ where &mut err_session, &mut budget, ); - }); + } if budget == 0 { activator.activate(); break; @@ -368,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); @@ -426,56 +389,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 +420,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 @@ -547,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 b5b9e545387ea..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::{CollectionEdge, 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}; @@ -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,7 @@ 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(oks, as_of, source_relation, initial_closure); }; match bundle.arrangement(&source_key) { Some(ArrangementFlavor::Local(oks, _errs)) => { @@ -876,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, @@ -892,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 8a52e4d14c1a6..a9fefcf57ee83 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; @@ -279,13 +279,11 @@ 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) = joined - .into_vec() + let (j, errs) = columnar_to_vec(joined) .flat_map_fallible::, CB<_>, _, _, _, _>(name, { // Reuseable allocation for unpacking. let mut datums = DatumVec::new(); @@ -333,10 +331,9 @@ 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) => edge.into_vec(), + JoinedFlavor::Edge(edge) => columnar_to_vec(edge), JoinedFlavor::Collection(collection) => collection, _ => panic!("Unexpectedly arranged join output"), }; @@ -359,21 +356,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"), } }; @@ -567,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, @@ -583,7 +574,7 @@ where { stream.unary_fallible::, _, _, _>( Pipeline, - "LinearJoinKeyPreparation", + "LinearJoinAccumulatorKeyPreparation", |_, _| { Box::new(move |input, ok, errs| { let mut temp_storage = RowArena::new(); @@ -659,14 +650,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 +666,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) } @@ -776,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() @@ -829,122 +810,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..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, 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; @@ -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 }; @@ -158,24 +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. 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| { - 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()); } } @@ -614,16 +628,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 +644,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 +704,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 +1267,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 +1323,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 +1407,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()