Pre-arrange Differential introspection sources

In a quick test, this reduces the number of rows feeding into the output
as follows:
* Batches: 70%
* Records: 26%

Tested with introspection debugging turned on.

Signed-off-by: Moritz Hoffmann <mh@materialize.com>

Author: antiguru

src/compute/src/logging/differential.rs

@@ -115,33 +115,46 @@ pub(super) fn construct<A: Allocate>(
         });
 
         // Encode the contents of each logging stream into its expected `Row` format.
-        let arrangement_batches = batches.as_collection().map(move |op| {
-            Row::pack_slice(&[
-                Datum::UInt64(u64::cast_from(op)),
-                Datum::UInt64(u64::cast_from(worker_id)),
-            ])
-        });
-        let arrangement_records = records.as_collection().map(move |op| {
-            Row::pack_slice(&[
-                Datum::UInt64(u64::cast_from(op)),
-                Datum::UInt64(u64::cast_from(worker_id)),
-            ])
-        });
+        let arrangement_batches = batches
+            .as_collection()
+            .mz_arrange_core::<_, RowSpine<_, _, _, _>>(
+                Exchange::new(move |_| u64::cast_from(worker_id)),
+                "PreArrange Differential batches",
+                compute_state.enable_arrangement_size_logging,
+            )
+            .as_collection(move |op, ()| {
+                Row::pack_slice(&[
+                    Datum::UInt64(u64::cast_from(*op)),
+                    Datum::UInt64(u64::cast_from(worker_id)),
+                ])
+            });
+        let arrangement_records = records
+            .as_collection()
+            .mz_arrange_core::<_, RowSpine<_, _, _, _>>(
+                Exchange::new(move |_| u64::cast_from(worker_id)),
+                "PreArrange Differential records",
+                compute_state.enable_arrangement_size_logging,
+            )
+            .as_collection(move |op, ()| {
+                Row::pack_slice(&[
+                    Datum::UInt64(u64::cast_from(*op)),
+                    Datum::UInt64(u64::cast_from(worker_id)),
+                ])
+            });
 
         let sharing = sharing
             .as_collection()
             .mz_arrange_core::<_, RowSpine<_, _, _, _>>(
                 Exchange::new(move |_| u64::cast_from(worker_id)),
                 "PreArrange Differential sharing",
                 compute_state.enable_arrangement_size_logging,
-            );
-
-        let sharing = sharing.as_collection(move |op, ()| {
-            Row::pack_slice(&[
-                Datum::UInt64(u64::cast_from(*op)),
-                Datum::UInt64(u64::cast_from(worker_id)),
-            ])
-        });
+            )
+            .as_collection(move |op, ()| {
+                Row::pack_slice(&[
+                    Datum::UInt64(u64::cast_from(*op)),
+                    Datum::UInt64(u64::cast_from(worker_id)),
+                ])
+            });
 
         use DifferentialLog::*;
         let logs = [
@@ -194,8 +207,8 @@ type OutputBuffer<'a, 'b, D> = ConsolidateBuffer<'a, 'b, Timestamp, D, Diff, Pus
 
 /// Bundled output buffers used by the demux operator.
 struct DemuxOutput<'a, 'b> {
-    batches: OutputBuffer<'a, 'b, usize>,
-    records: OutputBuffer<'a, 'b, usize>,
+    batches: OutputBuffer<'a, 'b, (usize, ())>,
+    records: OutputBuffer<'a, 'b, (usize, ())>,
     sharing: OutputBuffer<'a, 'b, (usize, ())>,
 }
 
@@ -248,10 +261,10 @@ impl DemuxHandler<'_, '_, '_> {
     fn handle_batch(&mut self, event: BatchEvent) {
         let ts = self.ts();
         let op = event.operator;
-        self.output.batches.give(self.cap, (op, ts, 1));
+        self.output.batches.give(self.cap, ((op, ()), ts, 1));
 
         let diff = Diff::try_from(event.length).expect("must fit");
-        self.output.records.give(self.cap, (op, ts, diff));
+        self.output.records.give(self.cap, ((op, ()), ts, diff));
         self.notify_arrangement_size(op);
     }
 
@@ -260,24 +273,24 @@ impl DemuxHandler<'_, '_, '_> {
 
         let ts = self.ts();
         let op = event.operator;
-        self.output.batches.give(self.cap, (op, ts, -1));
+        self.output.batches.give(self.cap, ((op, ()), ts, -1));
 
         let diff = Diff::try_from(done).expect("must fit")
             - Diff::try_from(event.length1 + event.length2).expect("must fit");
         if diff != 0 {
-            self.output.records.give(self.cap, (op, ts, diff));
+            self.output.records.give(self.cap, ((op, ()), ts, diff));
         }
         self.notify_arrangement_size(op);
     }
 
     fn handle_drop(&mut self, event: DropEvent) {
         let ts = self.ts();
         let op = event.operator;
-        self.output.batches.give(self.cap, (op, ts, -1));
+        self.output.batches.give(self.cap, ((op, ()), ts, -1));
 
         let diff = -Diff::try_from(event.length).expect("must fit");
         if diff != 0 {
-            self.output.records.give(self.cap, (op, ts, diff));
+            self.output.records.give(self.cap, ((op, ()), ts, diff));
         }
         self.notify_arrangement_size(op);
     }