compute: sort snapshot arrangements once instead of merging them

src/compute/src/render/context.rs

@@ -34,9 +34,7 @@ use mz_repr::{DatumVec, DatumVecBorrow, Diff, GlobalId, Row, RowArena, SharedRow
 use mz_storage_types::controller::CollectionMetadata;
 use mz_timely_util::columnar::batcher;
 use mz_timely_util::columnar::builder::ColumnBuilder;
-use mz_timely_util::columnar::{
-    Col2ValBatcher, Col2ValColBatcher, Col2ValPagedBatcher, columnar_exchange,
-};
+use mz_timely_util::columnar::{Col2ValColBatcher, Col2ValPagedBatcher, columnar_exchange};
 use mz_timely_util::columnation::ColumnationChunker;
 use timely::ContainerBuilder;
 use timely::container::{CapacityContainerBuilder, PushInto};
@@ -1284,8 +1282,8 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> {
             >(exchange, name),
             ArrangementBatcher::Columnation => ok_stream.mz_arrange_core::<
                 _,
-                batcher::Chunker<_>,
-                Col2ValBatcher<_, _, _, _>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
+                mz_row_spine::RowRowBatcher<_, _>,
                 RowRowBuilder<_, _>,
                 RowRowSpine<_, _>,
             >(exchange, name),

src/compute/src/render/join/linear_join.rs

@@ -28,9 +28,7 @@ use mz_repr::fixed_length::ExtendDatums;
 use mz_repr::{DatumVec, Diff, Row, RowArena, SharedRow};
 use mz_timely_util::columnar::batcher;
 use mz_timely_util::columnar::builder::ColumnBuilder;
-use mz_timely_util::columnar::{
-    Col2ValBatcher, Col2ValColBatcher, Col2ValPagedBatcher, columnar_exchange,
-};
+use mz_timely_util::columnar::{Col2ValColBatcher, Col2ValPagedBatcher, columnar_exchange};
 use mz_timely_util::operator::{CollectionExt, StreamExt};
 use timely::dataflow::Scope;
 use timely::dataflow::channels::pact::{ExchangeCore, Pipeline};
@@ -407,8 +405,8 @@ where
                 >(exchange, "JoinStage"),
                 ArrangementBatcher::Columnation => keyed.mz_arrange_core::<
                     _,
-                    batcher::Chunker<_>,
-                    Col2ValBatcher<_, _, _, _>,
+                    mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
+                    mz_row_spine::RowRowBatcher<_, _>,
                     RowRowBuilder<_, _>,
                     RowRowSpine<_, _>,
                 >(exchange, "JoinStage"),

src/compute/src/render/reduce.rs

@@ -54,12 +54,12 @@ use crate::render::errors::MaybeValidatingRow;
 use crate::render::reduce::monoids::{ReductionMonoid, get_monoid};
 use crate::render::{ArrangementFlavor, Pairer, RenderTimestamp};
 use crate::typedefs::{
-    ErrBatcher, ErrBuilder, KeyBatcher, RowErrBuilder, RowErrSpine, RowRowAgent, RowRowArrangement,
+    ErrBatcher, ErrBuilder, RowErrBuilder, RowErrSpine, RowRowAgent, RowRowArrangement,
     RowRowSpine, RowSpine, RowValSpine,
 };
 use mz_row_spine::{
-    DatumContainer, DatumSeq, RowBatcher, RowBuilder, RowRowBatcher, RowRowBuilder, RowValBatcher,
-    RowValBuilder,
+    DatumContainer, DatumSeq, KeyBatcher, RowBatcher, RowBuilder, RowRowBatcher, RowRowBuilder,
+    RowValBatcher, RowValBuilder,
 };
 
 /// Key container of trace `Tr`'s batch cursor.
@@ -301,7 +301,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
 
         let arranged = collection
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowRowBatcher<_, _>,
                 RowRowBuilder<_, _>,
                 RowRowSpine<_, _>,
@@ -411,7 +411,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
 
         let arranged = differential_dataflow::collection::concatenate(input.scope(), to_collect)
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowValBatcher<_, _, _>,
                 RowValBuilder<_, _, _>,
                 RowValSpine<_, _, _>,
@@ -577,7 +577,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
         };
         let arranged = partial
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowRowBatcher<_, _>,
                 RowRowBuilder<_, _>,
                 RowRowSpine<_, _>,
@@ -806,7 +806,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
 
         let input: KeyCollection<_, _, _> = input.into();
         let arranged = input.mz_arrange::<
-            ColumnationChunker<_>,
+            mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
             RowBatcher<_, _>,
             RowBuilder<_, _>,
             RowSpine<_, _>,
@@ -941,7 +941,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
                 // view mz_introspection.mz_expected_group_size_advice.
                 let arranged = partial
                     .mz_arrange::<
-                        ColumnationChunker<_>,
+                        mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                         RowRowBatcher<_, _>,
                         RowRowBuilder<_, _>,
                         RowRowSpine<_, _>,
@@ -1152,7 +1152,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
         // view mz_introspection.mz_expected_group_size_advice.
         let arranged_input = input
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowRowBatcher<_, _>,
                 RowRowBuilder<_, _>,
                 RowRowSpine<_, _>,
@@ -1285,7 +1285,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
         let partial: KeyCollection<_, _, _> = partial.into();
         let arranged = partial
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowBatcher<_, _>,
                 RowBuilder<_, _>,
                 RowSpine<_, Vec<ReductionMonoid>>,
@@ -1441,7 +1441,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
                     (pairer.merge(&key, std::iter::once(value)), ())
                 })
                 .mz_arrange::<
-                    ColumnationChunker<_>,
+                    mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                     RowBatcher<_, _>,
                     RowBuilder<_, _>,
                     RowSpine<_, _>,
@@ -1484,7 +1484,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
         let err_full_aggrs = full_aggrs.clone();
         let arranged = collection
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowBatcher<_, _>,
                 RowBuilder<_, _>,
                 RowSpine<_, (Vec<Accum>, Diff)>,

src/compute/src/render/top_k.rs

@@ -48,10 +48,10 @@ use crate::render::Pairer;
 use crate::render::context::{ArrangementFlavor, CollectionBundle, Context};
 use crate::render::errors::DataflowErrorSer;
 use crate::render::errors::MaybeValidatingRow;
-use crate::typedefs::{ErrBatcher, ErrBuilder, KeyBatcher, MzTimestamp, RowRowSpine, RowSpine};
+use crate::typedefs::{ErrBatcher, ErrBuilder, MzTimestamp, RowRowSpine, RowSpine};
 use mz_row_spine::{
-    DatumContainer, DatumSeq, RowBatcher, RowBuilder, RowRowBatcher, RowRowBuilder, RowValBuilder,
-    RowValSpine,
+    DatumContainer, DatumSeq, KeyBatcher, RowBatcher, RowBuilder, RowRowBatcher, RowRowBuilder,
+    RowValBuilder, RowValSpine,
 };
 
 // The implementation requires integer timestamps to be able to delay feedback for monotonic inputs.
@@ -581,7 +581,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim
             .into();
         let result = partial
             .mz_arrange::<
-                ColumnationChunker<_>,
+                mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
                 RowBatcher<_, _>,
                 RowBuilder<_, _>,
                 RowSpine<_, _>,
@@ -647,7 +647,7 @@ where
     let arranged = input
         .clone()
         .mz_arrange::<
-            ColumnationChunker<_>,
+            mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
             RowRowBatcher<_, _>,
             RowRowBuilder<_, _>,
             RowRowSpine<_, _>,

src/repr/src/row.rs

@@ -753,6 +753,22 @@ impl RowRef {
         &self.0
     }
 
+    /// A `u64` whose order is consistent with this type's `Ord`: for rows `a < b`,
+    /// `a.sort_prefix() <= b.sort_prefix()`. Equal prefixes decide nothing, and a caller
+    /// sorting by prefix must fall back to a full comparison on ties.
+    ///
+    /// Sorting a compact array of `(prefix, position)` pairs and comparing rows only on equal
+    /// prefixes is much cheaper than comparing rows throughout. The prefix is the length,
+    /// saturated at `u16::MAX`, followed by the first six bytes, mirroring the length-first
+    /// order of `Ord` below. Two saturated lengths compare equal, which keeps the agreement.
+    pub fn sort_prefix(&self) -> u64 {
+        let len = u64::cast_from(self.0.len().min(usize::from(u16::MAX))) << 48;
+        let mut lead = [0u8; 8];
+        let n = self.0.len().min(6);
+        lead[2..2 + n].copy_from_slice(&self.0[..n]);
+        len | u64::from_be_bytes(lead)
+    }
+
     /// True iff there is no data in this [`RowRef`].
     pub fn is_empty(&self) -> bool {
         self.0.is_empty()

src/row-spine/src/lib.rs

@@ -18,12 +18,14 @@ pub use self::dictionary::DatumContainer;
 pub use self::dictionary::DatumSeq;
 pub use self::offset_opt::OffsetOptimized;
 pub use self::spines::{
-    ArcOrdKeyBuilder, ArcOrdKeySpine, ArcOrdValBuilder, ArcOrdValSpine, RowBatcher, RowBuilder,
-    RowRowBatcher, RowRowBuilder, RowRowColPagedBuilder, RowRowSpine, RowSpine, RowValBatcher,
-    RowValBuilder, RowValSpine, ValRowBatcher, ValRowBuilder, ValRowColPagedBuilder, ValRowSpine,
+    ArcOrdKeyBuilder, ArcOrdKeySpine, ArcOrdValBuilder, ArcOrdValSpine, KeyBatcher, RowBatcher,
+    RowBuilder, RowRowBatcher, RowRowBuilder, RowRowColPagedBuilder, RowRowSpine, RowSpine,
+    RowValBatcher, RowValBuilder, RowValSpine, ValRowBatcher, ValRowBuilder, ValRowColPagedBuilder,
+    ValRowSpine,
 };
 
 mod arc_batch;
+pub mod snapshot_batcher;
 
 use differential_dataflow::trace::implementations::OffsetList;
 
@@ -51,10 +53,11 @@ mod spines {
     /// Batcher matching `mz_compute::typedefs::KeyValBatcher`, redeclared
     /// locally so this crate does not need to depend on `mz_compute`.
     type KeyValBatcher<K, V, T, D> = MergeBatcher<ColInternalMerger<(K, V), T, D>>;
-    type KeyBatcher<K, T, D> = KeyValBatcher<K, (), T, D>;
+    /// Batcher for collections consolidated by their data, for data led by a `Row`.
+    pub type KeyBatcher<K, T, R> = crate::snapshot_batcher::SnapshotBatcher<(K, ()), T, R>;
 
     pub type RowRowSpine<T, R> = Spine<ArcBatch<OrdValBatch<RowRowLayout<((Row, Row), T, R)>>>>;
-    pub type RowRowBatcher<T, R> = KeyValBatcher<Row, Row, T, R>;
+    pub type RowRowBatcher<T, R> = crate::snapshot_batcher::SnapshotBatcher<(Row, Row), T, R>;
     pub type RowRowBuilder<T, R> = ArcBuilder<crate::dictionary::builders::RowRowBuilder<T, R>>;
 
     /// `RowRowBuilder` variant that consumes [`Column`] chunks. Pairs with any
@@ -72,12 +75,12 @@ mod spines {
         ArcBuilder<crate::dictionary::builders::RowRowColPagedBuilder<T, R>>;
 
     pub type RowValSpine<V, T, R> = Spine<ArcBatch<OrdValBatch<RowValLayout<((Row, V), T, R)>>>>;
-    pub type RowValBatcher<V, T, R> = KeyValBatcher<Row, V, T, R>;
+    pub type RowValBatcher<V, T, R> = crate::snapshot_batcher::SnapshotBatcher<(Row, V), T, R>;
     pub type RowValBuilder<V, T, R> =
         ArcBuilder<crate::dictionary::builders::RowValBuilder<V, T, R>>;
 
     pub type RowSpine<T, R> = Spine<ArcBatch<OrdKeyBatch<RowLayout<((Row, ()), T, R)>>>>;
-    pub type RowBatcher<T, R> = KeyBatcher<Row, T, R>;
+    pub type RowBatcher<T, R> = crate::snapshot_batcher::SnapshotBatcher<(Row, ()), T, R>;
     pub type RowBuilder<T, R> = ArcBuilder<crate::dictionary::builders::RowBuilder<T, R>>;
 
     pub type ValRowSpine<K, T, R> = Spine<ArcBatch<OrdValBatch<ValRowLayout<((K, Row), T, R)>>>>;