diff --git a/differential-dataflow/src/trace/implementations/ord_neu.rs b/differential-dataflow/src/trace/implementations/ord_neu.rs index cf0c22a5f..03322643f 100644 --- a/differential-dataflow/src/trace/implementations/ord_neu.rs +++ b/differential-dataflow/src/trace/implementations/ord_neu.rs @@ -358,6 +358,17 @@ pub mod val_batch { staging: UpdsBuilder, } + impl OrdValMerger { + /// The partially assembled merge output. + /// + /// The containers are allocated at their full merged capacity when the merge begins, so + /// this storage accounts for the merge's whole memory cost from the outset, even though + /// the number of updates in it grows as the merge proceeds. + pub fn result(&self) -> &OrdValStorage { + &self.result + } + } + impl Merger> for OrdValMerger where OrdValBatch: Batch>, @@ -860,6 +871,17 @@ pub mod key_batch { staging: UpdsBuilder, } + impl OrdKeyMerger { + /// The partially assembled merge output. + /// + /// The containers are allocated at their full merged capacity when the merge begins, so + /// this storage accounts for the merge's whole memory cost from the outset, even though + /// the number of updates in it grows as the merge proceeds. + pub fn result(&self) -> &OrdKeyStorage { + &self.result + } + } + impl>> Merger> for OrdKeyMerger where OrdKeyBatch: Batch>, diff --git a/differential-dataflow/src/trace/implementations/spine_fueled.rs b/differential-dataflow/src/trace/implementations/spine_fueled.rs index f4e6ad036..b5e826a78 100644 --- a/differential-dataflow/src/trace/implementations/spine_fueled.rs +++ b/differential-dataflow/src/trace/implementations/spine_fueled.rs @@ -223,6 +223,23 @@ impl TraceReader for Spine { } } +impl Spine { + /// Applies `f` to the merger of each layer that is currently mid-merge. + /// + /// A merger owns the partially assembled output of its merge. That storage is memory the + /// spine holds but which [`TraceReader::map_batches`] does not reach: it presents the two + /// input batches of a merge and stops there. A caller that accounts for the spine's memory + /// footprint must visit both, and must re-read a merger on each observation rather than + /// cache what it learns, because a merger's contents change as the merge proceeds. + pub fn map_mergers::Merger)>(&self, mut f: F) { + for state in self.merging.iter().rev() { + if let MergeState::Double(MergeVariant::InProgress(_, _, merger)) = state { + f(merger); + } + } + } +} + // A trace implementation for any key type that can be borrowed from or converted into `Key`. // TODO: Almost all this implementation seems to be generic with respect to the trace and batch types. impl Trace for Spine { diff --git a/differential-dataflow/src/trace/mod.rs b/differential-dataflow/src/trace/mod.rs index c6f864f53..1f6600411 100644 --- a/differential-dataflow/src/trace/mod.rs +++ b/differential-dataflow/src/trace/mod.rs @@ -430,6 +430,11 @@ pub mod rc_blanket_impls { /// Wrapper type for merging reference counted batches. pub struct RcMerger { merger: B::Merger } + impl RcMerger { + /// The wrapped batch's merger, which owns the merge's partially assembled output. + pub fn inner(&self) -> &B::Merger { &self.merger } + } + /// Represents a merge in progress. impl Merger> for RcMerger { fn new(source1: &Rc, source2: &Rc, compaction_frontier: AntichainRef) -> Self { RcMerger { merger: B::begin_merge(source1, source2, compaction_frontier) } }