compute: sort snapshot arrangements once instead of merging them

While a dataflow hydrates, every update reaching an arrangement carries the
snapshot time. The merge batcher still sorts each 64 KiB chunk and folds it
into a geometric ladder of two-way merges, copying every row once per level
(about thirteen times at 10M rows) and comparing full rows at every step,
although nothing can cancel across times and the sealed chain is simply
every update in `(data, time)` order, consolidated.

`SnapshotBatcher` holds incoming chunks untouched while they all share one
time. At `seal` it sorts a compact `(key prefix, chunk, position)` index,
comparing full rows only on equal prefixes, and emits the sorted and
consolidated chain in one copy pass. The first chunk with a second time hands
everything held to a `MergeBatcher` and the batcher stays on that path, so
steady-state behaviour is unchanged. `RowRef::sort_prefix` provides the
prefix next to the `Ord` impl whose order it must agree with.

Row-keyed arrangement sites use it through the `RowRowBatcher`,
`RowValBatcher` and `RowBatcher` aliases and the two direct
`Col2ValBatcher` sites in `context.rs` and `linear_join.rs`.

Measured on one worker at 10M rows, five repetitions each, against the
same tree without this change: a plain index on a table 3.53s to 2.81s
(the arrangement operator 2.00s to 1.27s), an index on a distinct view
5.40s to 4.76s (2.25s to 1.68s), a fact-to-dimension join 8.3 to 11.5s
down to 7.25s (input arrangement 2.0s to 1.6s, output 1.7s to 1.1s), and a
count-by-key aggregate 7.36s to 6.96s.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: frankmcsherry

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};
@@ -1285,7 +1283,7 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> {
             ArrangementBatcher::Columnation => ok_stream.mz_arrange_core::<
                 _,
                 batcher::Chunker<_>,
-                Col2ValBatcher<_, _, _, _>,
+                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};
@@ -408,7 +406,7 @@ where
                 ArrangementBatcher::Columnation => keyed.mz_arrange_core::<
                     _,
                     batcher::Chunker<_>,
-                    Col2ValBatcher<_, _, _, _>,
+                    mz_row_spine::RowRowBatcher<_, _>,
                     RowRowBuilder<_, _>,
                     RowRowSpine<_, _>,
                 >(exchange, "JoinStage"),

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

@@ -24,6 +24,7 @@ pub use self::spines::{
 };
 
 mod arc_batch;
+pub mod snapshot_batcher;
 
 use differential_dataflow::trace::implementations::OffsetList;
 
@@ -51,10 +52,12 @@ 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>>;
+    #[allow(dead_code)]
     type KeyBatcher<K, T, D> = KeyValBatcher<K, (), T, D>;
 
     pub type RowRowSpine<T, R> = Spine<ArcBatch<OrdValBatch<RowRowLayout<((Row, Row), T, R)>>>>;
-    pub type RowRowBatcher<T, R> = KeyValBatcher<Row, Row, T, R>;
+    // PROTOTYPE: snapshot batcher in place of the merge batcher at Row-keyed sites.
+    pub type RowRowBatcher<T, R> = crate::snapshot_batcher::SnapshotBatcher<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<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<(), 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)>>>>;