diff --git a/src/compute/src/render/join/linear_join.rs b/src/compute/src/render/join/linear_join.rs index 84de916c96c47..0e637da72a06e 100644 --- a/src/compute/src/render/join/linear_join.rs +++ b/src/compute/src/render/join/linear_join.rs @@ -101,6 +101,7 @@ impl LinearJoinSpec { arranged1: Arranged<'s, Tr1>, arranged2: Arranged<'s, Tr2>, result: L, + consolidate_output: bool, ) -> VecCollection<'s, T, I::Item, Diff> where T: Lattice + timely::progress::Timestamp, @@ -122,19 +123,23 @@ impl LinearJoinSpec { (Materialize, Some(work_limit), Some(time_limit)) => { let yield_fn = move |start: Instant, work| work >= work_limit || start.elapsed() >= time_limit; - mz_join_core(arranged1, arranged2, result, yield_fn).as_collection() + mz_join_core(arranged1, arranged2, result, yield_fn, consolidate_output) + .as_collection() } (Materialize, Some(work_limit), None) => { let yield_fn = move |_start, work| work >= work_limit; - mz_join_core(arranged1, arranged2, result, yield_fn).as_collection() + mz_join_core(arranged1, arranged2, result, yield_fn, consolidate_output) + .as_collection() } (Materialize, None, Some(time_limit)) => { let yield_fn = move |start: Instant, _work| start.elapsed() >= time_limit; - mz_join_core(arranged1, arranged2, result, yield_fn).as_collection() + mz_join_core(arranged1, arranged2, result, yield_fn, consolidate_output) + .as_collection() } (Materialize, None, None) => { let yield_fn = |_start, _work| false; - mz_join_core(arranged1, arranged2, result, yield_fn).as_collection() + mz_join_core(arranged1, arranged2, result, yield_fn, consolidate_output) + .as_collection() } } } @@ -498,24 +503,61 @@ where // Reuseable allocation for unpacking. let mut datums = DatumVec::new(); + if closure.is_identity() { + // The output row is the key, the stream value and the lookup value in that order, + // so when all three are row-encoded the bytes can be concatenated without a trip + // through datums. Distinct input pairs give distinct outputs, so the join core need + // not consolidate them. + let oks = self.linear_join_spec.render( + prev_keyed, + next_input, + move |key, old, new| match (key.as_row_ref(), old.as_row_ref(), new.as_row_ref()) { + (Some(key), Some(old), Some(new)) => { + let mut row = + Row::with_capacity(key.byte_len() + old.byte_len() + new.byte_len()); + let mut packer = row.packer(); + packer.extend_by_row_ref(key); + packer.extend_by_row_ref(old); + packer.extend_by_row_ref(new); + Some(row) + } + _ => { + let temp_storage = RowArena::new(); + let mut datums_local = datums.borrow(); + key.extend_datums(&temp_storage, &mut datums_local, None); + old.extend_datums(&temp_storage, &mut datums_local, None); + new.extend_datums(&temp_storage, &mut datums_local, None); + Some(Row::pack(datums_local.iter())) + } + }, + false, + ); + return (oks, None); + } + if closure.could_error() { let (oks, err) = self .linear_join_spec - .render(prev_keyed, next_input, move |key, old, new| { - let mut row_builder = SharedRow::get(); - let temp_storage = RowArena::new(); + .render( + prev_keyed, + next_input, + move |key, old, new| { + let mut row_builder = SharedRow::get(); + let temp_storage = RowArena::new(); - let mut datums_local = datums.borrow(); - key.extend_datums(&temp_storage, &mut datums_local, None); - old.extend_datums(&temp_storage, &mut datums_local, None); - new.extend_datums(&temp_storage, &mut datums_local, None); + let mut datums_local = datums.borrow(); + key.extend_datums(&temp_storage, &mut datums_local, None); + old.extend_datums(&temp_storage, &mut datums_local, None); + new.extend_datums(&temp_storage, &mut datums_local, None); - closure - .apply(&mut datums_local, &temp_storage, &mut row_builder) - .map(|row| row.cloned()) - .map_err(DataflowErrorSer::from) - .transpose() - }) + closure + .apply(&mut datums_local, &temp_storage, &mut row_builder) + .map(|row| row.cloned()) + .map_err(DataflowErrorSer::from) + .transpose() + }, + true, + ) .inner .ok_err(|(x, t, d)| { // TODO(mcsherry): consider `ok_err()` for `Collection`. @@ -527,9 +569,10 @@ where (oks.as_collection(), Some(err.as_collection())) } else { - let oks = self - .linear_join_spec - .render(prev_keyed, next_input, move |key, old, new| { + let oks = self.linear_join_spec.render( + prev_keyed, + next_input, + move |key, old, new| { let mut row_builder = SharedRow::get(); let temp_storage = RowArena::new(); @@ -542,7 +585,9 @@ where .apply(&mut datums_local, &temp_storage, &mut row_builder) .expect("Closure claimed to never error") .cloned() - }); + }, + true, + ); (oks, None) } diff --git a/src/compute/src/render/join/mz_join_core.rs b/src/compute/src/render/join/mz_join_core.rs index 3ef61aa63ef22..9832a44b16da9 100644 --- a/src/compute/src/render/join/mz_join_core.rs +++ b/src/compute/src/render/join/mz_join_core.rs @@ -59,6 +59,7 @@ pub(super) fn mz_join_core<'scope, T, Tr1, Tr2, L, I, YFn, C>( arranged2: Arranged<'scope, Tr2>, result: L, yield_fn: YFn, + consolidate_output: bool, ) -> Stream<'scope, T, C> where T: timely::progress::Timestamp + Lattice, @@ -107,9 +108,12 @@ where // cursor against trace2's merged cursor, and vice versa for input 2. let mut todo1 = Work::, CursorList>, _, _>::new( Rc::clone(&result_fn), + consolidate_output, + ); + let mut todo2 = Work::>, BatchCursor, _, _>::new( + result_fn, + consolidate_output, ); - let mut todo2 = - Work::>, BatchCursor, _, _>::new(result_fn); // We'll unload the initial batches here, to put ourselves in a less non-deterministic state to start. trace1.map_batches(|batch1| { @@ -519,6 +523,10 @@ where /// /// Used with `yield_fn` to inform when `Work::process` should yield. produced: Rc>, + /// Whether to consolidate each chunk of produced results before emitting it. Worth it when + /// the closure can map distinct input pairs to equal results, as a projection does; wasted + /// sorting when it cannot, as with an identity closure over consolidated inputs. + consolidate_output: bool, _cursors: PhantomData<(C1, C2)>, } @@ -531,12 +539,13 @@ where L: FnMut(C1::Key<'_>, C1::Val<'_>, C2::Val<'_>) -> I + 'static, I: IntoIterator + 'static, { - fn new(result_fn: Rc>) -> Self { + fn new(result_fn: Rc>, consolidate_output: bool) -> Self { Self { todo: Default::default(), result_fn, output: Default::default(), produced: Default::default(), + consolidate_output, _cursors: PhantomData, } } @@ -597,13 +606,12 @@ where // Drain the produced join results. let mut output_buf = self.output.borrow_mut(); - - // Consolidating here is important when the join closure produces data that - // consolidates well, for example when projecting columns. - let old_len = output_buf.len(); - consolidate_updates(&mut output_buf); - let recovered = old_len - output_buf.len(); - self.produced.update(|x| x - recovered); + if self.consolidate_output { + let old_len = output_buf.len(); + consolidate_updates(&mut output_buf); + let recovered = old_len - output_buf.len(); + self.produced.update(|x| x - recovered); + } output.session(&cap).give_iterator(output_buf.drain(..)); diff --git a/src/repr/src/fixed_length.rs b/src/repr/src/fixed_length.rs index dcccaba4340d9..8e39f090a2f3e 100644 --- a/src/repr/src/fixed_length.rs +++ b/src/repr/src/fixed_length.rs @@ -11,7 +11,7 @@ //! `Row` is the most obvious implementor, but other trace types that may use more advanced //! representations only need to commit to implementing this trait. -use crate::{Datum, Row, RowArena}; +use crate::{Datum, Row, RowArena, RowRef}; /// A helper trait for types that can append their datums to a `Vec`. pub trait ExtendDatums { @@ -32,6 +32,15 @@ pub trait ExtendDatums { target: &mut Vec>, max: Option, ); + + /// The datums as one row encoding, when they are stored as such. + /// + /// `None` means the datums are not available as contiguous row-encoded bytes (for example + /// a dictionary-coded representation), and the caller must go through `extend_datums`. + #[inline] + fn as_row_ref(&self) -> Option<&RowRef> { + None + } } impl ExtendDatums for &T { @@ -45,6 +54,11 @@ impl ExtendDatums for &T { // Forward to T's impl so an override isn't lost behind a reference. (**self).extend_datums(arena, target, max) } + + #[inline] + fn as_row_ref(&self) -> Option<&RowRef> { + (**self).as_row_ref() + } } // Identity implementation for Row, whose datums borrow directly from its bytes. @@ -61,4 +75,9 @@ impl ExtendDatums for Row { None => target.extend(self.iter()), } } + + #[inline] + fn as_row_ref(&self) -> Option<&RowRef> { + Some(Row::as_row_ref(self)) + } } diff --git a/src/row-spine/src/lib.rs b/src/row-spine/src/lib.rs index ad84b647029af..dc248c990eab1 100644 --- a/src/row-spine/src/lib.rs +++ b/src/row-spine/src/lib.rs @@ -1778,6 +1778,17 @@ mod dictionary { } } } + + #[inline] + fn as_row_ref(&self) -> Option<&RowRef> { + if self.iter.index.is_none() { + // SAFETY: without a codec the iterator's bytes are exactly one row encoding, + // which is what `extend_datums` reads them as above. + Some(unsafe { RowRef::from_slice(self.iter.data) }) + } else { + None + } + } } }