compute: concatenate row bytes for identity join closures

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)
         }

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::<BatchCursor<Tr1>, CursorList<BatchCursor<Tr2>>, _, _>::new(
                 Rc::clone(&result_fn),
+                consolidate_output,
+            );
+            let mut todo2 = Work::<CursorList<BatchCursor<Tr1>>, BatchCursor<Tr2>, _, _>::new(
+                result_fn,
+                consolidate_output,
             );
-            let mut todo2 =
-                Work::<CursorList<BatchCursor<Tr1>>, BatchCursor<Tr2>, _, _>::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<Cell<usize>>,
+    /// 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<Item = D> + 'static,
 {
-    fn new(result_fn: Rc<RefCell<L>>) -> Self {
+    fn new(result_fn: Rc<RefCell<L>>, 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(..));
 

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<Datum>`.
 pub trait ExtendDatums {
@@ -32,6 +32,15 @@ pub trait ExtendDatums {
         target: &mut Vec<Datum<'a>>,
         max: Option<usize>,
     );
+
+    /// 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<T: ExtendDatums + ?Sized> ExtendDatums for &T {
@@ -45,6 +54,11 @@ impl<T: ExtendDatums + ?Sized> 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))
+    }
 }

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
+            }
+        }
     }
 }