From 7e87024da152c21c968d5dce88e6160603bd423b Mon Sep 17 00:00:00 2001
From: Frank McSherry <fmcsherry@me.com>
Date: Thu, 10 Sep 2026 13:39:44 -0400
Subject: [PATCH] Advance Corgi timestamps in runtime-width rows

---
 interactive/src/corgi/chunk.rs              |  53 ++++++++
 interactive/src/corgi/col_times.rs          |  10 +-
 interactive/src/corgi/time_advance_tests.rs | 127 ++++++++++++++++++
 interactive/src/corgi/time_rows.rs          | 138 ++++++++++++++++++++
 4 files changed, 325 insertions(+), 3 deletions(-)
 create mode 100644 interactive/src/corgi/time_advance_tests.rs
 create mode 100644 interactive/src/corgi/time_rows.rs

diff --git a/interactive/src/corgi/chunk.rs b/interactive/src/corgi/chunk.rs
index 683d4872c..3d8a0f6e3 100644
--- a/interactive/src/corgi/chunk.rs
+++ b/interactive/src/corgi/chunk.rs
@@ -327,6 +327,20 @@ where
             input.push_front(Self::from_kv(gather(&ckv, &idx), ct, cdiffs[end..].to_vec()));
         }
 
+        // Dispatch once per buffer for DDIR's numeric product lattice. The general
+        // timestamp path below retains its own Lattice::advance_by semantics.
+        use std::any::Any;
+        use crate::ir::Time;
+        if let Some(times) = (&ctimes as &dyn Any).downcast_ref::<ColTimes<Time>>() {
+            let frontier: Antichain<Time> = frontier.iter()
+                .map(|t| (t as &dyn Any).downcast_ref::<Time>().expect("checked time type").clone())
+                .collect();
+            let target = (out as &mut dyn Any).downcast_mut::<VecDeque<CorgiChunk<Time, R>>>()
+                .expect("checked time type");
+            CorgiChunk::<Time, R>::advance_rows(&ckv, times, &cdiffs, &bounds, end, frontier.borrow(), target);
+            return;
+        }
+
         // Advance + consolidate each complete group; emit `TARGET`-sized chunks. All rows of a group
         // share `(key, val)`, so one representative offset materializes each output row's kv. Times are
         // materialized here (owned `T`) because `advance_by` mutates and the tiebreak re-sort is a Rust
@@ -401,6 +415,41 @@ where
     }
 }
 
+impl<R: Semigroup + Clone + 'static> CorgiChunk<crate::ir::Time, R> {
+    /// Advance complete key/value groups without constructing an owned timestamp per record.
+    fn advance_rows(
+        kv: &CValue, times: &ColTimes<crate::ir::Time>, diffs: &[R], bounds: &[usize], end: usize,
+        frontier: AntichainRef<crate::ir::Time>, out: &mut VecDeque<Self>,
+    ) {
+        let times = crate::corgi::col_times::TimeRows::advance(times, end, frontier);
+        let (mut tags, mut offs, mut order) = (Vec::new(), Vec::new(), Vec::new());
+        let (mut otimes, mut odiffs) = (ColTimes::new(), Vec::new());
+        let mut start = 0;
+        for &stop in bounds {
+            if stop > end { break; }
+            order.extend(start..stop);
+            order.sort_by(|&a, &b| times.cmp(a, b));
+            let mut run = order.drain(..).peekable();
+            while let Some(row) = run.next() {
+                let mut diff = diffs[row].clone();
+                while run.peek().is_some_and(|&other| times.cmp(row, other).is_eq()) {
+                    diff.plus_equals(&diffs[run.next().unwrap()]);
+                }
+                if !diff.is_zero() {
+                    // A complete group shares one key/value; keep its representative offset.
+                    tags.push(0); offs.push(start); times.push_to(row, &mut otimes); odiffs.push(diff);
+                    if otimes.len() >= TARGET {
+                        Self::emit(&[Some(kv)], &tags, &offs, std::mem::replace(&mut otimes, ColTimes::new()), std::mem::take(&mut odiffs), out);
+                        tags.clear(); offs.clear();
+                    }
+                }
+            }
+            start = stop;
+        }
+        if !otimes.is_empty() { Self::emit(&[Some(kv)], &tags, &offs, otimes, odiffs, out); }
+    }
+}
+
 /// Sort parallel columns by `(key, val, time)` and consolidate exact `(key, val, time)` triples
 /// (summing diffs, dropping zeros). Returns a sorted+consolidated `(keys, vals, times, diffs)`.
 ///
@@ -831,3 +880,7 @@ mod test {
         }
     }
 }
+
+#[cfg(test)]
+#[path = "time_advance_tests.rs"]
+mod time_advance_tests;
diff --git a/interactive/src/corgi/col_times.rs b/interactive/src/corgi/col_times.rs
index f9d00a678..b0d488307 100644
--- a/interactive/src/corgi/col_times.rs
+++ b/interactive/src/corgi/col_times.rs
@@ -8,9 +8,9 @@
 //!
 //! Times are compared IN PLACE via the container's derived `Ord` on `Ref` (both `Product` and
 //! `PointStamp` carry `#[columnar(derive(Ord, PartialOrd))]`), so merge/sort never materialize a
-//! `T`. An owned `T` is reconstructed (`get`) only where a `Lattice` op is unavoidable — `join` in
-//! the join cross-product, `advance_by` in compaction — or at the emit boundary handing `T` back to
-//! DD. Range copies (`emit`/`concat`) push `Ref`s straight across (`push_ref`), also no `T`.
+//! `T`. DDIR's numeric product times can also advance in bulk through `TimeRows`. Other lattice
+//! operations and the emit boundary reconstruct owned times through `get`. Range copies
+//! (`emit`/`concat`) push `Ref`s straight across (`push_ref`), also no `T`.
 //!
 //! This is the O(data) time store; DD's `Chunk` boundary only ever sees whole chunks + frontier
 //! antichains (control complexity), so this stays entirely inside the backend — no DD change.
@@ -22,6 +22,10 @@ use columnar::{Borrow, Clear, Columnar, Container, Index, Len, Push};
 use differential_dataflow::lattice::Lattice;
 use timely::progress::Timestamp;
 
+#[path = "time_rows.rs"]
+mod time_rows;
+pub(crate) use time_rows::TimeRows;
+
 /// A timestamp usable as a columnar time column: `Timestamp + Lattice` (DD's algebra) plus
 /// `Columnar` with an *ordered* `Ref` (so times compare in their SoA form). Our
 /// `Product<u64, PointStamp<u64>>` satisfies it — every layer derives `Columnar` with
diff --git a/interactive/src/corgi/time_advance_tests.rs b/interactive/src/corgi/time_advance_tests.rs
new file mode 100644
index 000000000..0d5d4c087
--- /dev/null
+++ b/interactive/src/corgi/time_advance_tests.rs
@@ -0,0 +1,127 @@
+use super::*;
+use std::collections::BTreeMap;
+#[test]
+fn advance_owned_and_shared_nested_times_matches_reference() {
+    use differential_dataflow::dynamic::pointstamp::PointStamp;
+    use differential_dataflow::lattice::Lattice;
+    use timely::order::Product;
+    type T = Product<u64, PointStamp<u64>>;
+    for depth in [0, 1, 2, 3, 4, 5, 8, 16] {
+        let time = |outer, coords: &[u64]| {
+            T::new(
+                outer,
+                PointStamp::new(coords.iter().copied().cycle().take(depth).collect()),
+            )
+        };
+        let times = [
+            time(0, &[]),
+            time(1, &[2, 3, 1]),
+            time(2, &[1, 4, 2]),
+            time(2, &[3, 1, 3]),
+            time(3, &[2]),
+        ];
+        let rows: Vec<_> = (0..5)
+            .flat_map(|k| (0..2).flat_map(move |v| (0..5).map(move |t| (k, v, t))))
+            .map(|(k, v, i)| ((k, v), times[i].clone(), if i % 2 == 0 { 1i64 } else { -1 }))
+            .collect();
+        for frontier in [
+            Antichain::new(),
+            Antichain::from_elem(time(3, &[2, 2, 1])),
+            Antichain::from(vec![time(1, &[4, 1, 2]), time(3, &[1, 2, 3])]),
+            Antichain::from_elem(T::new(
+                3,
+                PointStamp::new([3, 2, 1, 2].into_iter().collect()),
+            )),
+        ] {
+            let mut expected = BTreeMap::new();
+            for (kv, t, d) in &rows {
+                let mut t = t.clone();
+                t.advance_by(frontier.borrow());
+                *expected.entry((*kv, t)).or_insert(0i64) += d;
+            }
+            expected.retain(|_, d| *d != 0);
+            for size in [1, 3, rows.len()] {
+                for shared in [false, true] {
+                    let chunks: Vec<_> = rows
+                        .chunks(size)
+                        .map(|rows| {
+                            CorgiChunk::from_columns(
+                                CValue::u64(rows.iter().map(|r| r.0 .0).collect()),
+                                CValue::u64(rows.iter().map(|r| r.0 .1).collect()),
+                                rows.iter().map(|r| r.1.clone()).collect(),
+                                rows.iter().map(|r| r.2).collect(),
+                            )
+                        })
+                        .collect();
+                    let retained = if shared { chunks.clone() } else { Vec::new() };
+                    let retained_times: Vec<_> =
+                        retained.iter().map(|c| c.times().to_vec()).collect();
+                    let (mut input, mut output) = (VecDeque::new(), VecDeque::new());
+                    for chunk in chunks {
+                        input.push_back(chunk);
+                        CorgiChunk::advance(&mut input, frontier.borrow(), false, &mut output);
+                    }
+                    CorgiChunk::advance(&mut input, frontier.borrow(), true, &mut output);
+                    assert!(input.is_empty());
+                    let mut actual = BTreeMap::new();
+                    let mut previous = None;
+                    for chunk in output {
+                        let keys = corgi::arrange::leaf_slice(chunk.keys()).unwrap();
+                        let vals = corgi::arrange::leaf_slice(chunk.vals()).unwrap();
+                        for i in 0..chunk.len_() {
+                            let key = ((keys[i], vals[i]), chunk.times().get(i));
+                            assert!(previous.as_ref().is_none_or(|p| p < &key));
+                            previous = Some(key.clone());
+                            assert!(actual.insert(key, chunk.diffs()[i]).is_none());
+                        }
+                    }
+                    assert_eq!(
+                        actual, expected,
+                        "size={size}, shared={shared}, frontier={frontier:?}"
+                    );
+                    for (chunk, original) in retained.iter().zip(&retained_times) {
+                        assert_eq!(
+                            chunk.times().to_vec(),
+                            *original,
+                            "shared input was modified"
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[test]
+fn advance_carries_a_large_group_then_flushes_within_it() {
+    use crate::ir::Time;
+    let count = TARGET + 1;
+    let initial = CorgiChunk::from_parts(
+        CValue::Unit(count),
+        CValue::Unit(count),
+        (0..count)
+            .map(|i| Time::new(i as u64, Default::default()))
+            .collect(),
+        vec![1i64; count],
+    );
+    let frontier = Antichain::from_elem(Time::new(0, Default::default()));
+    let mut input = VecDeque::from([initial]);
+    let mut output = VecDeque::new();
+    // One key/value group is not complete until the input is done.
+    CorgiChunk::advance(&mut input, frontier.borrow(), false, &mut output);
+    assert_eq!(input.len(), 1);
+    assert!(output.is_empty());
+    CorgiChunk::advance(&mut input, frontier.borrow(), true, &mut output);
+    assert!(input.is_empty());
+    assert_eq!(output.len(), 2);
+    let mut row = 0;
+    for chunk in output {
+        assert!(chunk.len() <= TARGET);
+        for i in 0..chunk.len() {
+            assert_eq!(chunk.times().get(i), Time::new(row, Default::default()));
+            assert_eq!(chunk.diffs()[i], 1);
+            row += 1;
+        }
+    }
+    assert_eq!(row, count as u64);
+}
diff --git a/interactive/src/corgi/time_rows.rs b/interactive/src/corgi/time_rows.rs
new file mode 100644
index 000000000..0a1fe3405
--- /dev/null
+++ b/interactive/src/corgi/time_rows.rs
@@ -0,0 +1,138 @@
+//! Temporary, fixed-width rows for advancing DDIR's numeric product timestamps.
+use super::ColTimes;
+use crate::ir::Time;
+use columnar::{Borrow, Index};
+use std::cmp::Ordering;
+use timely::progress::frontier::AntichainRef;
+
+/// One flat coordinate buffer; the width is determined at runtime, with no arity cutoff.
+/// Rows include the outer epoch and zero-pad PointStamp's omitted trailing coordinates.
+pub(crate) struct TimeRows {
+    data: Vec<u64>,
+    width: usize,
+}
+
+impl TimeRows {
+    /// Read and advance the first `rows` timestamps. Scratch space is `rows * width` coordinates.
+    pub(crate) fn advance(
+        times: &ColTimes<Time>,
+        rows: usize,
+        frontier: AntichainRef<Time>,
+    ) -> Self {
+        let source = times.store.borrow();
+        let width = (0..rows)
+            .map(|r| source.get(r).inner.vector.len() + 1)
+            .chain(frontier.iter().map(|t| t.inner.len() + 1))
+            .max()
+            .unwrap_or(1);
+        let mut data = vec![0; rows.checked_mul(width).expect("timestamp storage overflow")];
+        for (index, row) in data.chunks_exact_mut(width).enumerate() {
+            let time = source.get(index);
+            row[0] = *time.outer;
+            for (value, coordinate) in row[1..].iter_mut().zip(time.inner.vector.into_iter()) {
+                *value = *coordinate;
+            }
+        }
+
+        // Numeric product lattices distribute: meet_f(t join f) = t join meet_f(f).
+        // Thus each coordinate advances to max(t_i, min_f(f_i)), with frontier minima
+        // computed once for the buffer. This is specific to this lattice, not arbitrary T.
+        // An empty frontier leaves times unchanged, matching Lattice::advance_by.
+        if !frontier.is_empty() {
+            let minima: Vec<u64> = (0..width)
+                .map(|c| {
+                    frontier
+                        .iter()
+                        .map(|t| {
+                            if c == 0 {
+                                t.outer
+                            } else {
+                                t.inner.get(c - 1).copied().unwrap_or(0)
+                            }
+                        })
+                        .min()
+                        .unwrap()
+                })
+                .collect();
+            for row in data.chunks_exact_mut(width) {
+                for (value, minimum) in row.iter_mut().zip(&minima) {
+                    *value = (*value).max(*minimum);
+                }
+            }
+        }
+        Self { data, width }
+    }
+
+    fn row(&self, index: usize) -> &[u64] {
+        &self.data[index * self.width..(index + 1) * self.width]
+    }
+
+    pub(crate) fn cmp(&self, left: usize, right: usize) -> Ordering {
+        self.row(left).cmp(self.row(right))
+    }
+
+    /// Write coordinates directly to the column, omitting trailing zeros as PointStamp requires.
+    pub(crate) fn push_to(&self, index: usize, output: &mut ColTimes<Time>) {
+        let row = self.row(index);
+        let inner = &row[1..];
+        let end = inner.iter().rposition(|v| *v != 0).map_or(0, |i| i + 1);
+        output.store.outer.push(row[0]);
+        output
+            .store
+            .inner
+            .vector
+            .push_iter(inner[..end].iter().copied());
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use differential_dataflow::{dynamic::pointstamp::PointStamp, lattice::Lattice};
+    use timely::progress::Antichain;
+
+    #[test]
+    fn rows_preserve_advancement_order_and_canonical_times() {
+        let time = |outer, coordinates: &[u64]| {
+            Time::new(
+                outer,
+                PointStamp::new(coordinates.iter().copied().collect()),
+            )
+        };
+        let times = [
+            time(0, &[]),
+            time(1, &[2]),
+            time(2, &[0, 4, 1]),
+            time(3, &[u64::MAX, 0, 7]),
+            time(u64::MAX, &[0; 16]),
+            time(4, &[3; 16]),
+        ];
+        let source: ColTimes<Time> = times.iter().cloned().collect();
+        for frontier in [
+            Antichain::new(),
+            Antichain::from_elem(time(2, &[1])),
+            Antichain::from(vec![time(4, &[2, 0, 9]), time(2, &[0, 4, 1])]),
+            Antichain::from_elem(time(0, &[2; 17])),
+        ] {
+            for count in [0, 1, 3, times.len()] {
+                let rows = TimeRows::advance(&source, count, frontier.borrow());
+                let mut output = ColTimes::new();
+                let expected: Vec<_> = times[..count]
+                    .iter()
+                    .cloned()
+                    .map(|mut t| {
+                        t.advance_by(frontier.borrow());
+                        t
+                    })
+                    .collect();
+                for i in 0..count {
+                    rows.push_to(i, &mut output);
+                    assert_eq!(output.get(i), expected[i]);
+                    for j in 0..count {
+                        assert_eq!(rows.cmp(i, j), expected[i].cmp(&expected[j]));
+                    }
+                }
+            }
+        }
+    }
+}