From 69046c77ff19bd9c42360289d59c5825792de163 Mon Sep 17 00:00:00 2001 From: Moritz Hoffmann Date: Mon, 20 Jul 2026 17:06:11 +0200 Subject: [PATCH] compute: columnar linear-join input (C4) Rework the linear-join source-key path to consume its input as a `CollectionEdge` rather than decoding to `Vec` via `as_specific_collection`. The `JoinedFlavor::Collection` variant now carries an edge, and the source arrangement key is formed off it. Pattern: C1-reuse (borrowed-push, zero-alloc), not the C3 fallback. The join's source arrange already builds a `ColumnBuilder<((Row,Row),T,Diff)>` and arranges via the `Col2Val` batcher family, so the keying operator pushes the key and value borrowed with no owned `Row` per record. The keying-plus-arrange step is factored into a new free fn `arrange_join_input` with a `Vec` arm (behaviorally identical to before, same operator names) and a columnar arm reading via `into_index_iter`, owning time and diff only on the error path. The no-source-key path reads `input.collection` directly. The source-key path still reads from an existing arrangement (already columnar internally) and wraps its decoded rows as a `Vec` edge. Stage outputs and the final join output stay `Vec` (the join output producer is a later migration node); `into_vec` is the identity on the `Vec` arm, so Wave 1 is unchanged and the columnar arm is dead in production. Adds `arrange_join_input_arms_agree` (cross-arm agreement across several timestamps plus a retraction) and `arrange_join_input_arms_agree_on_error_path` (a fallible key drives every record onto the `try_extend` Err branch, covering the columnar arm's `Columnar::into_owned` for time and diff and asserting both arms produce identical errors). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/compute/src/render/join/linear_join.rs | 453 +++++++++++++++++---- 1 file changed, 366 insertions(+), 87 deletions(-) diff --git a/src/compute/src/render/join/linear_join.rs b/src/compute/src/render/join/linear_join.rs index 13d6d910d0d91..e62ec14dacb93 100644 --- a/src/compute/src/render/join/linear_join.rs +++ b/src/compute/src/render/join/linear_join.rs @@ -13,6 +13,7 @@ use std::time::{Duration, Instant}; +use columnar::{Columnar, Index}; use differential_dataflow::consolidation::ConsolidatingContainerBuilder; use differential_dataflow::lattice::Lattice; use differential_dataflow::operators::arrange::arrangement::Arranged; @@ -24,6 +25,7 @@ use mz_compute_types::dyncfgs::{ }; use mz_compute_types::plan::join::JoinClosure; use mz_compute_types::plan::join::linear_join::{LinearJoinPlan, LinearStagePlan}; +use mz_compute_types::plan::scalar::LirScalarExpr; use mz_dyncfg::ConfigSet; use mz_expr::Eval; use mz_repr::fixed_length::ExtendDatums; @@ -38,6 +40,7 @@ use timely::dataflow::operators::OkErr; use crate::extensions::arrange::MzArrangeCore; use crate::render::RenderTimestamp; +use crate::render::columnar::CollectionEdge; use crate::render::context::{ArrangementFlavor, CollectionBundle, Context}; use crate::render::errors::DataflowErrorSer; use crate::render::join::mz_join_core::mz_join_core; @@ -189,8 +192,10 @@ impl YieldSpec { /// Different forms the streamed data might take. enum JoinedFlavor<'scope, T: RenderTimestamp> { - /// Streamed data as a collection. - Collection(VecCollection<'scope, T, Row, Diff>), + /// Streamed data as a collection edge. `differential_join` forms its + /// arrangement key off the edge, so a columnar source flows in without a + /// `ColumnarToVec` decode. + Collection(CollectionEdge<'scope, T>), /// A dataflow-local arrangement. Local(Arranged<'scope, RowRowAgent>), /// An imported arrangement. @@ -241,39 +246,60 @@ where (_, initial_closure) => { // TODO: extract closure from the first stage in the join plan, should it exist. // TODO: apply that closure in `flat_map_ref` rather than calling `.collection`. - let (joined, errs) = inputs[linear_plan.source_relation] - .as_specific_collection(linear_plan.source_key.as_deref(), &self.config_set); + let (joined, errs) = match linear_plan.source_key.as_deref() { + // No source key: consume the input edge directly rather than + // decoding it to `Vec`. `differential_join` forms the source + // arrangement key off this edge below (C1 borrowed-push + // pattern), so a columnar source has no `ColumnarToVec` hop. + None => inputs[linear_plan.source_relation] + .collection + .clone() + .expect("The unarranged collection doesn't exist."), + // With a source key the input is read from an existing + // arrangement, which is already columnar internally; wrap its + // decoded rows as a `Vec` edge. + Some(key) => { + let (oks, errs) = inputs[linear_plan.source_relation] + .as_specific_collection(Some(key), &self.config_set); + (CollectionEdge::Vec(oks), errs) + } + }; errors.push(errs.enter_region(inner)); - let mut joined = joined.enter_region(inner); + let joined = joined.enter_region(inner); // In the current code this should always be `None`, but we have this here should // we change that and want to know what we should be doing. if let Some(closure) = initial_closure { // 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 + // 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. let name = "LinearJoinInitialization"; type CB = ConsolidatingContainerBuilder; - let (j, errs) = joined.flat_map_fallible::, CB<_>, _, _, _, _>(name, { - // Reuseable allocation for unpacking. - let mut datums = DatumVec::new(); - move |row| { - let mut row_builder = SharedRow::get(); - let temp_storage = RowArena::new(); - let mut datums_local = datums.borrow_with(&row); - // TODO(mcsherry): re-use `row` allocation. - closure - .apply(&mut datums_local, &temp_storage, &mut row_builder) - .map(|row| row.cloned()) - .map_err(DataflowErrorSer::from) - .transpose() - } - }); - joined = j; + let (j, errs) = joined + .into_vec() + .flat_map_fallible::, CB<_>, _, _, _, _>(name, { + // Reuseable allocation for unpacking. + let mut datums = DatumVec::new(); + move |row| { + let mut row_builder = SharedRow::get(); + let temp_storage = RowArena::new(); + let mut datums_local = datums.borrow_with(&row); + // TODO(mcsherry): re-use `row` allocation. + closure + .apply(&mut datums_local, &temp_storage, &mut row_builder) + .map(|row| row.cloned()) + .map_err(DataflowErrorSer::from) + .transpose() + } + }); errors.push(errs); + JoinedFlavor::Collection(CollectionEdge::Vec(j)) + } else { + JoinedFlavor::Collection(joined) } - - JoinedFlavor::Collection(joined) } }; @@ -287,14 +313,18 @@ where stage_plan, &mut errors, ); - // Update joined results and capture any errors. - joined = JoinedFlavor::Collection(stream); + // Update joined results and capture any errors. Stage output is a + // `Vec` collection (join output producer is a later migration node). + joined = JoinedFlavor::Collection(CollectionEdge::Vec(stream)); } // We have completed the join building, but may have work remaining. // For example, we may have expressions not pushed down (e.g. literals) // and projections that could not be applied (e.g. column repetition). - let bundle = if let JoinedFlavor::Collection(mut joined) = joined { + let bundle = if let JoinedFlavor::Collection(joined) = joined { + // The join output is a `Vec` collection until the output producer is + // flipped in a later node; `into_vec` is the identity on the `Vec` arm. + let mut joined = joined.into_vec(); if let Some(closure) = linear_plan.final_closure { let name = "LinearJoinFinalization"; type CB = ConsolidatingContainerBuilder; @@ -346,67 +376,13 @@ where ) -> VecCollection<'s, T, Row, Diff> { // If we have only a streamed collection, we must first form an arrangement. if let JoinedFlavor::Collection(stream) = joined { - let name = "LinearJoinKeyPreparation"; - let (keyed, errs) = stream - .inner - .unary_fallible::, _, _, _>( - Pipeline, - name, - |_, _| { - 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); - for (row, time, diff) in data.iter() { - temp_storage.clear(); - let datums_local = datums.borrow_with(row); - let datums = stream_key - .iter() - .map(|e| e.eval(&datums_local, &temp_storage)); - let result = key_buf.packer().try_extend(datums); - match result { - 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(), time.clone(), *diff)); - } - } - } - }); - }) - }, - ); - - errors.push(errs.as_collection()); - - let exchange = ExchangeCore::, _>::new_core( - columnar_exchange::, + let (arranged, errs) = arrange_join_input( + stream, + stream_key, + stream_thinning, + ENABLE_COLUMN_PAGED_BATCHER.get(&self.config_set), ); - let arranged = if ENABLE_COLUMN_PAGED_BATCHER.get(&self.config_set) { - keyed.mz_arrange_core::< - _, - batcher::ColumnChunker<_>, - Col2ValPagedBatcher<_, _, _, _>, - RowRowColPagedBuilder<_, _>, - RowRowSpine<_, _>, - >(exchange, "JoinStage") - } else { - keyed.mz_arrange_core::< - _, - batcher::Chunker<_>, - Col2ValBatcher<_, _, _, _>, - RowRowBuilder<_, _>, - RowRowSpine<_, _>, - >(exchange, "JoinStage") - }; + errors.push(errs); joined = JoinedFlavor::Local(arranged); } @@ -542,3 +518,306 @@ 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 C1's 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. Producers still emit `Vec` in Wave 1, so +/// the columnar arm is dead in production. +fn arrange_join_input<'s, T>( + edge: CollectionEdge<'s, T>, + stream_key: Vec, + stream_thinning: Vec, + use_paged_path: bool, +) -> ( + Arranged<'s, RowRowAgent>, + VecCollection<'s, T, DataflowErrorSer, Diff>, +) +where + T: Lattice + RenderTimestamp, +{ + let name = "LinearJoinKeyPreparation"; + let (keyed, errs) = match edge { + CollectionEdge::Vec(stream) => stream + .inner + .unary_fallible::, _, _, _>( + Pipeline, + name, + |_, _| { + 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); + for (row, time, diff) in data.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(), time.clone(), *diff)); + } + } + } + }); + }) + }, + ), + CollectionEdge::Columnar(stream) => stream + .inner + .unary_fallible::, _, _, _>( + Pipeline, + name, + |_, _| { + 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 exchange = + ExchangeCore::, _>::new_core(columnar_exchange::); + let arranged = if use_paged_path { + keyed.mz_arrange_core::< + _, + batcher::ColumnChunker<_>, + Col2ValPagedBatcher<_, _, _, _>, + RowRowColPagedBuilder<_, _>, + RowRowSpine<_, _>, + >(exchange, "JoinStage") + } else { + keyed.mz_arrange_core::< + _, + batcher::Chunker<_>, + Col2ValBatcher<_, _, _, _>, + RowRowBuilder<_, _>, + RowRowSpine<_, _>, + >(exchange, "JoinStage") + }; + (arranged, errs.as_collection()) +} + +#[cfg(test)] +mod tests { + use differential_dataflow::input::Input; + use mz_expr::EvalError; + use mz_repr::{Datum, ReprScalarType, Timestamp}; + use timely::dataflow::operators::Capture; + use timely::dataflow::operators::capture::{Event, Extract}; + + use super::*; + use crate::render::columnar::vec_to_columnar; + + type KeyedUpdate = ((Row, Row), Timestamp, Diff); + type ErrUpdate = (DataflowErrorSer, Timestamp, Diff); + type Captured = std::sync::mpsc::Receiver>>; + + fn extract_sorted(captured: Captured) -> Vec { + let mut updates: Vec<_> = captured + .extract() + .into_iter() + .flat_map(|(_, data)| data) + .collect(); + updates.sort(); + updates + } + + // `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. + fn extract_err(captured: Captured) -> Vec<(String, Timestamp, Diff)> { + let mut updates: Vec<_> = captured + .extract() + .into_iter() + .flat_map(|(_, data)| data) + .map(|(e, t, d)| (format!("{e:?}"), t, d)) + .collect(); + updates.sort(); + updates + } + + /// Input rows tagged with distinct timestamps and mixed-sign diffs. The two + /// `-1` records retract at a `(row, time)` with no matching insertion, so they + /// survive the `InputSession`'s pre-send consolidation and exercise a negative + /// diff on the ok path (pushed borrowed, not owned). + fn test_input() -> Vec<(Row, u64, Diff)> { + vec![ + ( + Row::pack_slice(&[Datum::Int32(1), Datum::String("a")]), + 0, + Diff::ONE, + ), + ( + Row::pack_slice(&[Datum::Int32(2), Datum::String("b")]), + 1, + Diff::ONE, + ), + ( + Row::pack_slice(&[Datum::Int32(1), Datum::String("c")]), + 2, + Diff::ONE, + ), + ( + Row::pack_slice(&[Datum::Int32(3), Datum::Null]), + 2, + Diff::ONE, + ), + ( + Row::pack_slice(&[Datum::Int32(2), Datum::String("b")]), + 2, + -Diff::ONE, + ), + ( + Row::pack_slice(&[Datum::Int32(4), Datum::String("d")]), + 1, + -Diff::ONE, + ), + ] + } + + /// 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( + 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| { + 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], false); + 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(); + 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) + }) + }); + ( + extract_sorted(ok_vec), + extract_sorted(ok_col), + extract_err(err_vec), + extract_err(err_col), + ) + } + + /// 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`. + /// + /// 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. + #[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()); + // A retraction survives into the arrangement, so the ok path handled a + // negative (borrowed) diff. + assert!(ok_vec.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 { + 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. + #[mz_ore::test] + fn arrange_join_input_arms_agree_on_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); + } +}