sinks: drop the arrangement's trace reader, pass only batches

SinkRender::render_sink now takes a Stream<Vec<SinkBatch>> instead of an
Arranged<SinkTrace>. Sinks only ever walk incoming batches (via
for_each_diff_pair) — they never use the TraceAgent for random access — so the
reader handle is dead weight.

arrange_sink_input still calls arrange_named, but immediately extracts
arranged.stream and lets the surrounding Arranged (and its TraceAgent) drop.
With no reader holding compaction frontiers, the arrange operator's spine can
compact to the empty antichain as batches flow, releasing historical batch
state instead of accumulating it.

Mirrors the pattern used by DD's consolidate_named, which builds an
arrangement only to call as_collection on it and drop the trace.

Author: DAlperin

src/storage/src/render/sinks.rs

@@ -12,7 +12,8 @@
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
-use differential_dataflow::operators::arrange::{Arrange, Arranged, TraceAgent};
+use differential_dataflow::operators::arrange::{Arrange, TraceAgent};
+use differential_dataflow::trace::TraceReader;
 use differential_dataflow::trace::implementations::ord_neu::{
     OrdValBatcher, OrdValSpine, RcOrdValBuilder,
 };
@@ -31,10 +32,21 @@ use tracing::warn;
 use crate::healthcheck::HealthStatusMessage;
 use crate::storage_state::StorageState;
 
-/// The concrete trace type used to hand arranged sink input to
-/// [`SinkRender::render_sink`].
+/// The concrete trace type produced internally when arranging a sink's input.
+/// The sink never sees this directly — only the batches flowing through it —
+/// but it's the anchor for the batch type in [`SinkBatchStream`].
 pub(crate) type SinkTrace = TraceAgent<OrdValSpine<Option<Row>, Row, Timestamp, Diff>>;
 
+/// Stream of arrangement batches handed to [`SinkRender::render_sink`].
+///
+/// This is `Arranged::stream` with the trace reader dropped: sinks only need
+/// batch-level access (no random-access reads via a cursor), so we don't keep
+/// a `TraceAgent` alive. Dropping the reader lets the spine's compaction
+/// frontiers advance to the empty antichain, so the arrange operator can
+/// aggressively compact / release batch state as updates flow through.
+pub(crate) type SinkBatchStream<'scope> =
+    StreamVec<'scope, Timestamp, <SinkTrace as TraceReader>::Batch>;
+
 /// _Renders_ complete _differential_ collections
 /// that represent the sink and its errors as requested
 /// by the original `CREATE SINK` statement.
@@ -79,15 +91,15 @@ pub(crate) fn render_sink<'scope>(
         );
         tokens.extend(persist_tokens);
 
-        let arranged = arrange_sink_input(&*sink_render, ok_collection.as_collection());
+        let batches = arrange_sink_input(&*sink_render, ok_collection.as_collection());
         let key_is_synthetic = sink_render.get_key_indices().is_none()
             && sink_render.get_relation_key_indices().is_none();
 
         let (health, sink_tokens) = sink_render.render_sink(
             storage_state,
             sink,
             sink_id,
-            arranged,
+            batches,
             key_is_synthetic,
             err_collection.as_collection(),
         );
@@ -96,17 +108,22 @@ pub(crate) fn render_sink<'scope>(
     })
 }
 
-/// Extract the sink's key column(s) from each row and arrange the resulting
-/// `(Option<Row>, Row)` collection by key.
+/// Extract the sink's key column(s) from each row, arrange the resulting
+/// `(Option<Row>, Row)` collection by key, and return just the stream of
+/// batches — dropping the trace reader.
 ///
 /// Prefers the user-specified sink key, falling back to any natural key of the
 /// underlying relation. When neither exists, a synthetic per-row hash is used
 /// purely to distribute work across workers — in that case the sink should
 /// treat the key as absent (`key_is_synthetic`).
+///
+/// Partial-moving `arranged.stream` lets the surrounding `Arranged` (and the
+/// `TraceAgent` it holds) drop, releasing the spine's compaction holds so the
+/// arrange operator can compact batch state as it's emitted.
 fn arrange_sink_input<'scope>(
     sink_render: &dyn SinkRender<'scope>,
     collection: VecCollection<'scope, Timestamp, Row, Diff>,
-) -> Arranged<'scope, SinkTrace> {
+) -> SinkBatchStream<'scope> {
     let key_indices = sink_render
         .get_key_indices()
         .or_else(|| sink_render.get_relation_key_indices())
@@ -131,7 +148,8 @@ fn arrange_sink_input<'scope>(
     // Allow access to `arrange_named` because we cannot access Mz's wrapper
     // from here. TODO(database-issues#5046): Revisit with cluster unification.
     #[allow(clippy::disallowed_methods)]
-    keyed.arrange_named::<OrdValBatcher<_, _, _, _>, RcOrdValBuilder<_, _, _, _>, OrdValSpine<_, _, _, _>>("Arrange Sink")
+    let arranged = keyed.arrange_named::<OrdValBatcher<_, _, _, _>, RcOrdValBuilder<_, _, _, _>, OrdValSpine<_, _, _, _>>("Arrange Sink");
+    arranged.stream
 }
 
 /// Rate-limited detector for primary-key uniqueness violations as a sink's
@@ -213,8 +231,8 @@ pub(crate) trait SinkRender<'scope> {
 
     /// Renders the sink's dataflow.
     ///
-    /// The sink receives the input as an arrangement keyed on `Option<Row>`.
-    /// The sink is responsible for walking the arrangement (typically via
+    /// The sink receives a stream of arrangement batches keyed on `Option<Row>`.
+    /// The sink is responsible for walking each batch (typically via
     /// [`mz_interchange::envelopes::for_each_diff_pair`]) and handling any
     /// envelope-specific diff-pair construction. When `key_is_synthetic` is
     /// true the arrangement's key is a per-row hash used only for worker
@@ -225,7 +243,7 @@ pub(crate) trait SinkRender<'scope> {
         storage_state: &mut StorageState,
         sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
         sink_id: GlobalId,
-        arranged: Arranged<'scope, SinkTrace>,
+        batches: SinkBatchStream<'scope>,
         key_is_synthetic: bool,
         err_collection: VecCollection<'scope, Timestamp, DataflowError, Diff>,
     ) -> (

src/storage/src/sink/iceberg.rs

@@ -10,17 +10,19 @@
 //! Iceberg sink implementation.
 //!
 //! This code renders a [`IcebergSinkConnection`] into a dataflow that writes
-//! data to an Iceberg table. `SinkRender::render_sink` hands the sink an
-//! arranged collection keyed on the sink key; a small `walk_sink_arrangement`
-//! operator at the input walks the arrangement and emits one `DiffPair` per
-//! `(key, timestamp)` update into the pipeline below.
+//! data to an Iceberg table. `SinkRender::render_sink` hands the sink a stream
+//! of arrangement batches keyed on the sink key (the upstream arrangement's
+//! trace reader is already dropped, so the spine is free to compact as
+//! batches flow). A small `walk_sink_arrangement` operator consumes that
+//! stream and emits one `DiffPair` per `(key, timestamp)` update into the
+//! pipeline below.
 //!
 //! ```text
 //!        ┏━━━━━━━━━━━━━━┓
 //!        ┃   persist    ┃
 //!        ┃    source    ┃
 //!        ┗━━━━━━┯━━━━━━━┛
-//!               │ arranged (Option<Row>, Row) keyed by sink key
+//!               │ stream of arrangement batches (trace reader dropped)
 //!               │
 //!        ┏━━━━━━v━━━━━━━┓
 //!        ┃    walk      ┃
@@ -91,7 +93,6 @@ use anyhow::{Context, anyhow};
 use arrow::array::{ArrayRef, Int32Array, Int64Array, RecordBatch};
 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
 use differential_dataflow::lattice::Lattice;
-use differential_dataflow::operators::arrange::Arranged;
 use differential_dataflow::{AsCollection, Hashable, VecCollection};
 use futures::StreamExt;
 use iceberg::ErrorKind;
@@ -153,7 +154,7 @@ use tracing::debug;
 
 use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace};
 use crate::metrics::sink::iceberg::IcebergSinkMetrics;
-use crate::render::sinks::{PkViolationWarner, SinkRender, SinkTrace};
+use crate::render::sinks::{PkViolationWarner, SinkBatchStream, SinkRender};
 use crate::statistics::SinkStatistics;
 use crate::storage_state::StorageState;
 
@@ -2251,18 +2252,18 @@ impl<'scope> SinkRender<'scope> for IcebergSinkConnection {
         storage_state: &mut StorageState,
         sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
         sink_id: GlobalId,
-        arranged: Arranged<'scope, SinkTrace>,
+        batches: SinkBatchStream<'scope>,
         key_is_synthetic: bool,
         _err_collection: VecCollection<'scope, Timestamp, DataflowError, Diff>,
     ) -> (
         StreamVec<'scope, Timestamp, HealthStatusMessage>,
         Vec<PressOnDropButton>,
     ) {
-        let scope = arranged.stream.scope();
+        let scope = batches.scope();
 
         let (input, walker_button) = walk_sink_arrangement(
             format!("{sink_id}-iceberg-walker"),
-            arranged,
+            batches,
             sink_id,
             sink.from,
             key_is_synthetic,
@@ -2413,25 +2414,25 @@ impl<'scope> SinkRender<'scope> for IcebergSinkConnection {
     }
 }
 
-/// Walks `arranged` and emits a stream of individual `(key, DiffPair)` records
-/// that feeds the rest of the Iceberg sink pipeline.
+/// Walks each arrangement batch and emits a stream of individual
+/// `(key, DiffPair)` records that feeds the rest of the Iceberg sink pipeline.
 ///
 /// Tracks per-`(key, timestamp)` group sizes and rate-limits a warning when a
 /// non-synthetic key has more than one `DiffPair`. When `key_is_synthetic` the
 /// arrangement's hash-based key is stripped before emission.
 fn walk_sink_arrangement<'scope>(
     name: String,
-    arranged: Arranged<'scope, SinkTrace>,
+    batches: SinkBatchStream<'scope>,
     sink_id: GlobalId,
     from_id: GlobalId,
     key_is_synthetic: bool,
 ) -> (
     VecCollection<'scope, Timestamp, (Option<Row>, DiffPair<Row>), Diff>,
     PressOnDropButton,
 ) {
-    let mut builder = OperatorBuilder::new(name, arranged.stream.scope());
+    let mut builder = OperatorBuilder::new(name, batches.scope());
     let (output, stream) = builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
-    let mut input = builder.new_input_for(arranged.stream, Pipeline, &output);
+    let mut input = builder.new_input_for(batches, Pipeline, &output);
 
     let button = builder.build(move |_caps| async move {
         let mut pk_warner = (!key_is_synthetic).then(|| PkViolationWarner::new(sink_id, from_id));

src/storage/src/sink/kafka.rs

@@ -15,10 +15,10 @@
 //!        ┃   persist    ┃
 //!        ┃    source    ┃
 //!        ┗━━━━━━┯━━━━━━━┛
-//!               │ arranged (Option<Row>, Row) keyed by sink key
+//!               │ stream of arrangement batches (trace reader dropped)
 //!               │
 //!        ┏━━━━━━v━━━━━━┓
-//!        ┃    row      ┃ walks the arrangement and emits one
+//!        ┃    row      ┃ walks each batch's cursor and emits one
 //!        ┃   encoder   ┃ encoded `KafkaMessage` per DiffPair
 //!        ┗━━━━━━┯━━━━━━┛
 //!               │ encoded data
@@ -82,11 +82,10 @@ use std::time::Duration;
 
 use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace};
 use crate::metrics::sink::kafka::KafkaSinkMetrics;
-use crate::render::sinks::{PkViolationWarner, SinkRender, SinkTrace};
+use crate::render::sinks::{PkViolationWarner, SinkBatchStream, SinkRender};
 use crate::statistics::SinkStatistics;
 use crate::storage_state::StorageState;
 use anyhow::{Context, anyhow, bail};
-use differential_dataflow::operators::arrange::Arranged;
 use differential_dataflow::{AsCollection, Hashable, VecCollection};
 use futures::StreamExt;
 use maplit::btreemap;
@@ -161,7 +160,7 @@ impl<'scope> SinkRender<'scope> for KafkaSinkConnection {
         storage_state: &mut StorageState,
         sink: &StorageSinkDesc<CollectionMetadata, Timestamp>,
         sink_id: GlobalId,
-        arranged: Arranged<'scope, SinkTrace>,
+        batches: SinkBatchStream<'scope>,
         key_is_synthetic: bool,
         // TODO(benesch): errors should stream out through the sink,
         // if we figure out a protocol for that.
@@ -170,7 +169,7 @@ impl<'scope> SinkRender<'scope> for KafkaSinkConnection {
         StreamVec<'scope, Timestamp, HealthStatusMessage>,
         Vec<PressOnDropButton>,
     ) {
-        let scope = arranged.stream.scope();
+        let scope = batches.scope();
 
         let write_handle = {
             let persist = Arc::clone(&storage_state.persist_clients);
@@ -196,7 +195,7 @@ impl<'scope> SinkRender<'scope> for KafkaSinkConnection {
 
         let (encoded, encode_status, encode_token) = encode_collection(
             format!("kafka-{sink_id}-{}-encode", self.format.get_format_name()),
-            arranged,
+            batches,
             sink.envelope,
             self.clone(),
             storage_state.storage_configuration.clone(),
@@ -1365,15 +1364,14 @@ async fn fetch_partition_count_loop<F>(
     }
 }
 
-/// Walks the sink's arrangement and emits encoded Kafka messages, one per
+/// Walks each arrangement batch and emits encoded Kafka messages, one per
 /// `DiffPair` observed at each `(key, timestamp)`.
 ///
-/// When `key_is_synthetic`, the arrangement's key is a per-row hash used only
-/// for worker distribution; the emitted `KafkaMessage` uses no key in that
-/// case.
+/// When `key_is_synthetic`, the batch keys are per-row hashes used only for
+/// worker distribution; the emitted `KafkaMessage` uses no key in that case.
 fn encode_collection<'scope>(
     name: String,
-    arranged: Arranged<'scope, SinkTrace>,
+    batches: SinkBatchStream<'scope>,
     envelope: SinkEnvelope,
     connection: KafkaSinkConnection,
     storage_configuration: StorageConfiguration,
@@ -1385,10 +1383,10 @@ fn encode_collection<'scope>(
     StreamVec<'scope, Timestamp, HealthStatusMessage>,
     PressOnDropButton,
 ) {
-    let mut builder = AsyncOperatorBuilder::new(name, arranged.stream.scope());
+    let mut builder = AsyncOperatorBuilder::new(name, batches.scope());
 
     let (output, stream) = builder.new_output::<CapacityContainerBuilder<Vec<_>>>();
-    let mut input = builder.new_input_for(arranged.stream, Pipeline, &output);
+    let mut input = builder.new_input_for(batches, Pipeline, &output);
 
     let (button, errors) = builder.build_fallible(move |caps| {
         Box::pin(async move {