dd hydration gate and tests for arrangement-size collection

* Gate the first snapshot on every compute object being hydrated, so
  early history rows don't record partial sizes. Sticky flag.
* Extract pruner cutoff and hydration check into pure helpers, each
  covered by new unit tests.
* New sqllogictest for the 10 MiB quantization formula.
* New testdrive for live + history behavior, plus a 2s cadence
  default in the testdrive harness so it fits the timeout.

Author: leedqin

src/adapter/src/coord.rs

@@ -1931,6 +1931,11 @@ pub struct Coordinator {
     /// The interval at which to collect storage usage information.
     storage_usage_collection_interval: Duration,
 
+    /// Set once all compute objects have been observed as hydrated, gating
+    /// the first write into `mz_object_arrangement_size_history`. Sticky:
+    /// later partial re-hydrations (e.g. replica restart) don't re-arm it.
+    arrangement_sizes_hydration_observed: bool,
+
     /// Segment analytics client.
     #[derivative(Debug = "ignore")]
     segment_client: Option<mz_segment::Client>,
@@ -4274,27 +4279,11 @@ impl Coordinator {
             differential_dataflow::consolidation::consolidate(&mut current_contents);
 
             let cutoff_ts = u128::from(read_ts).saturating_sub(retention_period.as_millis());
-            let mut expired = Vec::new();
-            for (row, diff) in current_contents {
-                assert_eq!(
-                    diff, 1,
-                    "consolidated contents should not contain retractions: ({row:#?}, {diff:#?})"
-                );
-                // Column 3 is `collection_timestamp`.
-                let collection_timestamp = row
-                    .unpack()
-                    .get(3)
-                    .expect("definition of mz_object_arrangement_size_history changed")
-                    .unwrap_timestamptz();
-                let collection_timestamp = collection_timestamp.timestamp_millis();
-                let collection_timestamp: u128 = collection_timestamp
-                    .try_into()
-                    .expect("all collections happen after Jan 1 1970");
-                if collection_timestamp < cutoff_ts {
-                    let builtin_update = BuiltinTableUpdate::row(item_id, row, Diff::MINUS_ONE);
-                    expired.push(builtin_update);
-                }
-            }
+            let expired = arrangement_sizes_expired_retractions(
+                current_contents,
+                cutoff_ts,
+                item_id,
+            );
 
             // TODO(arrangement-sizes): when the writeable-catalog-server
             // plumbing in https://github.com/MaterializeInc/materialize/pull/35436
@@ -4322,6 +4311,40 @@ impl Coordinator {
     }
 }
 
+/// Returns retraction updates for rows in a consolidated
+/// `mz_object_arrangement_size_history` snapshot whose `collection_timestamp`
+/// (column 3) is strictly before `cutoff_ts`.
+///
+/// Panics if any input row has `diff != 1`: the caller must consolidate first,
+/// and a consolidated history table should never contain retractions because
+/// the only source of retractions is this function itself.
+fn arrangement_sizes_expired_retractions(
+    rows: impl IntoIterator<Item = (mz_repr::Row, i64)>,
+    cutoff_ts: u128,
+    item_id: CatalogItemId,
+) -> Vec<BuiltinTableUpdate> {
+    let mut expired = Vec::new();
+    for (row, diff) in rows {
+        assert_eq!(
+            diff, 1,
+            "consolidated contents should not contain retractions: ({row:#?}, {diff:#?})"
+        );
+        let collection_timestamp = row
+            .unpack()
+            .get(3)
+            .expect("definition of mz_object_arrangement_size_history changed")
+            .unwrap_timestamptz()
+            .timestamp_millis();
+        let collection_timestamp: u128 = collection_timestamp
+            .try_into()
+            .expect("all collections happen after Jan 1 1970");
+        if collection_timestamp < cutoff_ts {
+            expired.push(BuiltinTableUpdate::row(item_id, row, Diff::MINUS_ONE));
+        }
+    }
+    expired
+}
+
 #[cfg(test)]
 impl Coordinator {
     #[allow(dead_code)]
@@ -4783,6 +4806,7 @@ pub fn serve(
                     cloud_resource_controller,
                     storage_usage_client,
                     storage_usage_collection_interval,
+                    arrangement_sizes_hydration_observed: false,
                     segment_client,
                     metrics,
                     optimizer_metrics,
@@ -5248,3 +5272,56 @@ mod id_pool_tests {
         pool.refill(10, 5);
     }
 }
+
+#[cfg(test)]
+mod arrangement_sizes_pruner_tests {
+    use mz_repr::catalog_item_id::CatalogItemId;
+    use mz_repr::{Datum, Row};
+
+    use super::arrangement_sizes_expired_retractions;
+
+    // Pack a row shaped like `mz_object_arrangement_size_history`: the pruner
+    // only cares about column 3 (`collection_timestamp`), but we stuff the
+    // other three columns with realistic values so shape changes would fail.
+    fn history_row(ts_ms: i64) -> Row {
+        let dt = mz_ore::now::to_datetime(ts_ms.try_into().expect("non-negative"));
+        Row::pack_slice(&[
+            Datum::String("r1"),
+            Datum::String("u1"),
+            Datum::Int64(123),
+            Datum::TimestampTz(dt.try_into().expect("fits in TimestampTz")),
+        ])
+    }
+
+    fn item_id() -> CatalogItemId {
+        // Any CatalogItemId will do; tests don't dispatch on it.
+        CatalogItemId::User(42)
+    }
+
+    #[mz_ore::test]
+    fn empty_input_produces_no_retractions() {
+        let out = arrangement_sizes_expired_retractions(Vec::new(), 1_000, item_id());
+        assert!(out.is_empty());
+    }
+
+    #[mz_ore::test]
+    fn retracts_only_rows_strictly_before_cutoff() {
+        // Mixes both sides of the filter and includes a row at exactly
+        // the cutoff timestamp to pin down the strict-less-than boundary.
+        let rows = vec![
+            (history_row(100), 1),
+            (history_row(500), 1),
+            (history_row(1_000), 1), // at cutoff: kept (strict <)
+            (history_row(5_000), 1),
+        ];
+        let out = arrangement_sizes_expired_retractions(rows, 1_000, item_id());
+        assert_eq!(out.len(), 2);
+    }
+
+    #[mz_ore::test]
+    #[should_panic(expected = "consolidated contents should not contain retractions")]
+    fn retraction_in_input_panics() {
+        let rows = vec![(history_row(100), -1)];
+        let _ = arrangement_sizes_expired_retractions(rows, 1_000, item_id());
+    }
+}

src/adapter/src/coord/introspection.rs

@@ -566,11 +566,18 @@ const SUBSCRIBES: &[SubscribeSpec] = &[
             GROUP BY export_id, lir_id
         )",
     },
-    // Per-object arrangement sizes via unified introspection subscribe.
-    // Quantization strategy: objects <10MB show real size for precise monitoring;
-    // objects ≥10MB round to nearest 10MB to suppress byte-level differential churn.
-    // Uses mz_dataflow_addresses.address[1] to extract dataflow_id, avoiding an extra join.
-    // Per-worker raw tables aggregate bytes across all workers for total replica memory per object.
+    // Per-object arrangement sizes, one row per `(object_id, replica)`.
+    //
+    // `mz_arrangement_heap_size_raw` and `mz_arrangement_batcher_size_raw` are
+    // differential logs where each `+1` row represents one byte of heap delta;
+    // after consolidation, `COUNT(*)` is the current arrangement size in bytes.
+    //
+    // Objects smaller than 10 MiB report exact bytes; larger ones are rounded
+    // to the nearest 10 MiB to suppress byte-level churn in the collection.
+    //
+    // `mz_dataflow_addresses.address[1]` is the root of each operator's address
+    // tree, which equals the owning `dataflow_id` — so we can go addresses →
+    // operator → dataflow without joining `mz_dataflow_operator_dataflows`.
     SubscribeSpec {
         introspection_type: IntrospectionType::ComputeObjectArrangementSizes,
         sql: "SUBSCRIBE (

src/adapter/src/coord/message_handler.rs

@@ -411,6 +411,19 @@ impl Coordinator {
             return;
         }
 
+        // Delay writes until every compute object has hydrated, so we
+        // don't record partial sizes from still-building dataflows.
+        // Sticky: later snapshots skip this check. Gate retries run at
+        // the full collection interval, so the first snapshot can lag
+        // hydration by up to one interval.
+        if !self.arrangement_sizes_hydration_observed {
+            if !self.check_arrangement_sizes_hydration().await {
+                self.schedule_arrangement_sizes_collection().await;
+                return;
+            }
+            self.arrangement_sizes_hydration_observed = true;
+        }
+
         let collection_timer = self
             .metrics
             .arrangement_sizes_collection_time_seconds
@@ -547,6 +560,32 @@ impl Coordinator {
         });
     }
 
+    /// Returns `true` when every row in `mz_compute_hydration_times` has a
+    /// non-null `time_ns` (i.e. every compute object on every replica has
+    /// finished its initial hydration). An empty collection also returns
+    /// `true`. On snapshot failure, returns `false` so the caller retries.
+    async fn check_arrangement_sizes_hydration(&self) -> bool {
+        let item_id = self
+            .catalog()
+            .resolve_builtin_storage_collection(&mz_catalog::builtin::MZ_COMPUTE_HYDRATION_TIMES);
+        let global_id = self.catalog.get_entry(&item_id).latest_global_id();
+        let read_ts = self.get_local_read_ts().await;
+        let mut snapshot = match self
+            .controller
+            .storage_collections
+            .snapshot(global_id, read_ts)
+            .await
+        {
+            Ok(s) => s,
+            Err(e) => {
+                tracing::warn!("arrangement-sizes hydration gate snapshot failed: {e:?}");
+                return false;
+            }
+        };
+        differential_dataflow::consolidation::consolidate(&mut snapshot);
+        arrangement_sizes_all_hydrated(&snapshot)
+    }
+
     #[mz_ore::instrument(level = "debug")]
     async fn message_command(&mut self, cmd: Command) {
         self.handle_command(cmd).await;
@@ -1028,3 +1067,88 @@ impl Coordinator {
         }
     }
 }
+
+/// Returns `true` when every `+1` row in a consolidated snapshot of
+/// `mz_compute_hydration_times` has a non-null `time_ns` (column 2), i.e.
+/// every compute object on every replica has finished hydrating at least
+/// once. Empty input returns `true`. Unexpected row shape returns `false`
+/// so the caller re-polls rather than crashing on schema drift.
+fn arrangement_sizes_all_hydrated(snapshot: &[(mz_repr::Row, i64)]) -> bool {
+    const HYDRATION_COL_TIME_NS: usize = 2;
+    for (row, diff) in snapshot {
+        if *diff != 1 {
+            continue;
+        }
+        let datums = row.unpack();
+        let Some(time_ns) = datums.get(HYDRATION_COL_TIME_NS) else {
+            return false;
+        };
+        if time_ns.is_null() {
+            return false;
+        }
+    }
+    true
+}
+
+#[cfg(test)]
+mod arrangement_sizes_hydration_tests {
+    use mz_repr::{Datum, Row};
+
+    use super::arrangement_sizes_all_hydrated;
+
+    // Columns: (replica_id, object_id, time_ns). `time_ns` is the value
+    // under test; the other columns are packed with realistic defaults so
+    // shape drift would surface.
+    fn hydration_row(time_ns: Option<u64>) -> Row {
+        Row::pack_slice(&[
+            Datum::String("r1"),
+            Datum::String("u1"),
+            match time_ns {
+                Some(ns) => Datum::UInt64(ns),
+                None => Datum::Null,
+            },
+        ])
+    }
+
+    #[mz_ore::test]
+    fn empty_snapshot_is_hydrated() {
+        assert!(arrangement_sizes_all_hydrated(&[]));
+    }
+
+    #[mz_ore::test]
+    fn all_non_null_is_hydrated() {
+        let rows = vec![
+            (hydration_row(Some(100)), 1),
+            (hydration_row(Some(200)), 1),
+        ];
+        assert!(arrangement_sizes_all_hydrated(&rows));
+    }
+
+    #[mz_ore::test]
+    fn any_null_blocks_hydration() {
+        let rows = vec![
+            (hydration_row(Some(100)), 1),
+            (hydration_row(None), 1),
+            (hydration_row(Some(300)), 1),
+        ];
+        assert!(!arrangement_sizes_all_hydrated(&rows));
+    }
+
+    #[mz_ore::test]
+    fn retractions_are_ignored() {
+        // A -1 row represents stale state that consolidation would remove.
+        // We skip it so a retracted null doesn't veto hydration.
+        let rows = vec![
+            (hydration_row(None), -1),
+            (hydration_row(Some(100)), 1),
+        ];
+        assert!(arrangement_sizes_all_hydrated(&rows));
+    }
+
+    #[mz_ore::test]
+    fn malformed_row_blocks_hydration() {
+        let malformed = Row::pack_slice(&[Datum::String("r1"), Datum::String("u1")]);
+        let rows = vec![(malformed, 1)];
+        assert!(!arrangement_sizes_all_hydrated(&rows));
+    }
+}

src/catalog/src/builtin.rs

@@ -8684,6 +8684,9 @@ pub static MZ_OBJECT_ARRANGEMENT_SIZES_IND: LazyLock<BuiltinIndex> =
         is_retained_metrics_object: true,
     });
 
+/// Identifies [`MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY`] for the schema-migration
+/// guard in `builtin_schema_migration.rs`, which forbids migrating this table
+/// because its startup pruner assumes it is the only source of retractions.
 pub static MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY_DESCRIPTION: LazyLock<SystemObjectDescription> =
     LazyLock::new(|| SystemObjectDescription {
         schema_name: MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY.schema.to_string(),

test/sqllogictest/arrangement_sizes_quantization.slt

@@ -0,0 +1,53 @@
+# Copyright Materialize, Inc. and contributors. All rights reserved.
+#
+# Use of this software is governed by the Business Source License
+# included in the LICENSE file at the root of this repository.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0.
+
+# Exercises the `size` quantization formula used by the introspection subscribe
+# that populates `mz_internal.mz_object_arrangement_sizes`. The formula lives in
+# `src/adapter/src/coord/introspection.rs`.
+#
+# Semantics:
+#   - byte counts below 10 MiB (10485760) are reported exactly
+#   - byte counts >= 10 MiB round to the nearest 10 MiB boundary using
+#     `((n + 5 MiB) / 10 MiB) * 10 MiB`
+
+mode cockroach
+
+# Applies the quantization formula to a spread of raw byte counts.
+statement ok
+CREATE VIEW arrangement_size_quantize AS
+SELECT
+    n,
+    CASE
+        WHEN n < 10485760 THEN n::int8
+        ELSE ((n + 5242880) / 10485760 * 10485760)::int8
+    END AS size
+FROM (VALUES
+    (0),               -- zero bytes
+    (1),               -- tiny arrangement
+    (10485759),        -- one byte below the threshold: reported exactly
+    (10485760),        -- exactly 10 MiB: already on a boundary
+    (10485761),        -- one byte past threshold: rounds down to 10 MiB
+    (15728640),        -- 15 MiB (boundary midpoint): rounds up to 20 MiB
+    (15728639),        -- one byte below midpoint: rounds down to 10 MiB
+    (20971520),        -- 20 MiB: already on a boundary
+    (104857600)        -- 100 MiB: already on a boundary
+) AS v(n);
+
+query II
+SELECT n, size FROM arrangement_size_quantize ORDER BY n
+----
+0  0
+1  1
+10485759  10485759
+10485760  10485760
+10485761  10485760
+15728639  10485760
+15728640  20971520
+20971520  20971520
+104857600  104857600