Skip to content

One Batcher trait, shaped like the one half_join had to write #844

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
frankmcsherry merged 1 commit into from
Aug 26, 2026
Merged

One Batcher trait, shaped like the one half_join had to write #844

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions differential-dataflow/examples/columnar/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,17 +129,17 @@ mod reachability {

let edges_arr = arrange_core::<_, _,
ValChunker<(Node, Node, IterTime, Diff)>,
ValBatcher<Node, Node, IterTime, Diff>,
_,
ValBuilder<Node, Node, IterTime, Diff>,
ValSpine<Node, Node, IterTime, Diff>,
>(edges_inner.inner, edges_pact, "Edges");
>(edges_inner.inner, edges_pact, "Edges", ValBatcher::new);

let reach_arr = arrange_core::<_, _,
ValChunker<(Node, (), IterTime, Diff)>,
ValBatcher<Node, (), IterTime, Diff>,
_,
ValBuilder<Node, (), IterTime, Diff>,
ValSpine<Node, (), IterTime, Diff>,
>(reach.inner, reach_pact, "Reach");
>(reach.inner, reach_pact, "Reach", ValBatcher::new);

// join_traces with ValColBuilder: produces Stream<_, RecordedUpdates<...>>.
let proposed =
Expand All @@ -162,10 +162,10 @@ mod reachability {
let combined_pact = ValPact { hashfunc: |k: columnar::Ref<'_, Node>| *k as u64 };
let combined_arr = arrange_core::<_, _,
ValChunker<(Node, (), IterTime, Diff)>,
ValBatcher<Node, (), IterTime, Diff>,
_,
ValBuilder<Node, (), IterTime, Diff>,
ValSpine<Node, (), IterTime, Diff>,
>(combined.inner, combined_pact, "Combined");
>(combined.inner, combined_pact, "Combined", ValBatcher::new);

// reduce_abelian on the columnar arrangement.
let result = combined_arr.reduce_abelian::<_,
Expand Down
4 changes: 2 additions & 2 deletions differential-dataflow/examples/columnar_spill.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,10 +239,10 @@ fn run_timely_dataflow(times: u64, keys_per_time: u64, workers: usize, sample_se
let arranged = arrange_core::<
_, _,
ValChunker<(u64, u64, u64, i64)>,
ValBatcher<u64, u64, u64, i64>,
_,
ValBuilder<u64, u64, u64, i64>,
ValSpine<u64, u64, u64, i64>,
>(stream, Pipeline, "ColumnarSpillArrange");
>(stream, Pipeline, "ColumnarSpillArrange", ValBatcher::new);
arranged.stream.probe_with(&mut probe);
});

Expand Down
16 changes: 8 additions & 8 deletions differential-dataflow/examples/spines.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,8 @@ fn main() {
use differential_dataflow::trace::implementations::ord_neu::{OrdKeyBatcher, VecOrdKeyBuilder, OrdKeySpine};
let (data_input, data) = scope.new_collection::<String, isize>();
let (keys_input, keys) = scope.new_collection::<String, isize>();
let data = data.map(|k| (k, ())).arrange::<OrdKeyBatcher<String,_,isize>, VecOrdKeyBuilder<String,_,isize>, OrdKeySpine<String,_,isize>>();
let keys = keys.map(|k| (k, ())).arrange::<OrdKeyBatcher<String,_,isize>, VecOrdKeyBuilder<String,_,isize>, OrdKeySpine<String,_,isize>>();
let data = data.map(|k| (k, ())).arrange::<_, VecOrdKeyBuilder<String,_,isize>, OrdKeySpine<String,_,isize>>(OrdKeyBatcher::new);
let keys = keys.map(|k| (k, ())).arrange::<_, VecOrdKeyBuilder<String,_,isize>, OrdKeySpine<String,_,isize>>(OrdKeyBatcher::new);
keys.join_core(data, |_k, &(), &()| Option::<()>::None)
.probe_with(&mut probe);
Workload { data_input, keys_input }
Expand All @@ -68,8 +68,8 @@ fn main() {
use differential_dataflow::trace::implementations::ord_neu::{OrdValBatcher, VecOrdValBuilder, OrdValSpine};
let (data_input, data) = scope.new_collection::<String, isize>();
let (keys_input, keys) = scope.new_collection::<String, isize>();
let data = data.map(|x| (x, ())).arrange::<OrdValBatcher<String,(),_,isize>, VecOrdValBuilder<String,(),_,isize>, OrdValSpine<String,(),_,isize>>();
let keys = keys.map(|x| (x, ())).arrange::<OrdValBatcher<String,(),_,isize>, VecOrdValBuilder<String,(),_,isize>, OrdValSpine<String,(),_,isize>>();
let data = data.map(|x| (x, ())).arrange::<_, VecOrdValBuilder<String,(),_,isize>, OrdValSpine<String,(),_,isize>>(OrdValBatcher::new);
let keys = keys.map(|x| (x, ())).arrange::<_, VecOrdValBuilder<String,(),_,isize>, OrdValSpine<String,(),_,isize>>(OrdValBatcher::new);
keys.join_core(data, |_k, &(), &()| Option::<()>::None)
.probe_with(&mut probe);
Workload { data_input, keys_input }
Expand All @@ -94,8 +94,8 @@ fn main() {
type Sp = Spine<String, (), u64, isize>;
type Chu = ContainerChunker<ColChunk<(String, (), u64, isize)>>;
let exchange = || Exchange::new(|u: &((String, ()), u64, isize)| (u.0).0.hashed().into());
let data = arrange_core::<_, _, Chu, Ba, Bu, Sp>(data.inner, exchange(), "DataArrange");
let keys = arrange_core::<_, _, Chu, Ba, Bu, Sp>(keys.inner, exchange(), "KeysArrange");
let data = arrange_core::<_, _, Chu, Ba, Bu, Sp>(data.inner, exchange(), "DataArrange", Ba::new);
let keys = arrange_core::<_, _, Chu, Ba, Bu, Sp>(keys.inner, exchange(), "KeysArrange", Ba::new);
// `ColChunk`'s cursor yields `Val = columnar::Ref<()> = ()`, not `&()`.
keys.join_core(data, |_k, _, _| Option::<()>::None)
.probe_with(&mut probe);
Expand All @@ -121,8 +121,8 @@ fn main() {
type Sp = ChunkSpine<String, (), u64, isize>;
type Chu = ContainerChunker<VecChunk<String, (), u64, isize>>;
let exchange = || Exchange::new(|u: &((String, ()), u64, isize)| (u.0).0.hashed().into());
let data = arrange_core::<_, _, Chu, Ba, Bu, Sp>(data.inner, exchange(), "DataArrange");
let keys = arrange_core::<_, _, Chu, Ba, Bu, Sp>(keys.inner, exchange(), "KeysArrange");
let data = arrange_core::<_, _, Chu, Ba, Bu, Sp>(data.inner, exchange(), "DataArrange", Ba::new);
let keys = arrange_core::<_, _, Chu, Ba, Bu, Sp>(keys.inner, exchange(), "KeysArrange", Ba::new);
keys.join_core(data, |_k, &(), &()| Option::<()>::None)
.probe_with(&mut probe);
Workload { data_input, keys_input }
Expand Down
36 changes: 18 additions & 18 deletions differential-dataflow/src/collection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -958,21 +958,21 @@ pub mod vec {
/// ```
pub fn consolidate(self) -> Self {
use crate::trace::implementations::{KeyBatcher, KeyBuilder, KeySpine};
self.consolidate_named::<KeyBatcher<_, _, _>,KeyBuilder<_,_,_>, KeySpine<_,_,_>,_>("Consolidate", |key,&()| key.clone())
self.consolidate_named::<_,KeyBuilder<_,_,_>, KeySpine<_,_,_>,_>("Consolidate", KeyBatcher::new, |key,&()| key.clone())
}

/// As `consolidate` but with the ability to name the operator, specify the trace type,
/// and provide the function `reify` to produce owned keys and values..
pub fn consolidate_named<Ba, Bu, Tr, F>(self, name: &str, reify: F) -> Self
pub fn consolidate_named<Ba, Bu, Tr, F>(self, name: &str, batcher: impl FnOnce(Option<crate::logging::Logger>, usize) -> Ba, reify: F) -> Self
where
Ba: crate::trace::Batcher<Output=Vec<((D, ()), T, R)>, Time=T> + 'static,
Ba: crate::trace::Batcher<T, Vec<((D, ()), T, R)>, Vec<Vec<((D, ()), T, R)>>> + 'static,
Tr: crate::trace::Trace<Batch: Navigable, Time=T>+'static,
for<'a> BatchCursor<Tr>: Cursor<Time=Tr::Time, Diff=R>,
Bu: crate::trace::Builder<Time=Tr::Time, Input=Vec<((D, ()), T, R)>, Output: Into<Tr::Batch>>,
F: Fn(BatchKey<'_, Tr>, BatchVal<'_, Tr>) -> D + 'static,
{
self.map(|k| (k, ()))
.arrange_named::<Ba, Bu, Tr>(name)
.arrange_named::<Ba, Bu, Tr>(name, batcher)
.as_collection(reify)
}

Expand Down Expand Up @@ -1032,28 +1032,28 @@ pub mod vec {
{
/// Arranges updates into a shared trace, exchanged by a hash of the key.
///
/// The batcher's output container must equal the stream container; the default chunker
/// only consolidates same-type containers. For chunker setups that convert between
/// container types (e.g. columnar layouts), call
/// [`arrange_core`](crate::operators::arrange::arrangement::arrange_core) directly.
pub fn arrange<Ba, Bu, Tr>(self) -> Arranged<'scope, TraceAgent<Tr>>
/// The batcher must accept the stream container; the default chunker only consolidates
/// same-type containers. For chunker setups that convert between container types (e.g.
/// columnar layouts), call [`arrange_core`](crate::operators::arrange::arrangement::arrange_core)
/// directly.
pub fn arrange<Ba, Bu, Tr>(self, batcher: impl FnOnce(Option<crate::logging::Logger>, usize) -> Ba) -> Arranged<'scope, TraceAgent<Tr>>
where
Ba: crate::trace::Batcher<Output=Vec<((K, V), T, R)>, Time=T> + 'static,
Bu: crate::trace::Builder<Time=T, Input=Vec<((K, V), T, R)>, Output: Into<Tr::Batch>>,
Ba: crate::trace::Batcher<T, Vec<((K, V), T, R)>, Vec<Bu::Input>> + 'static,
Bu: crate::trace::Builder<Time=T, Output: Into<Tr::Batch>>,
Tr: crate::trace::Trace<Time=T> + 'static,
{
self.arrange_named::<Ba, Bu, Tr>("Arrange")
self.arrange_named::<Ba, Bu, Tr>("Arrange", batcher)
}

/// As [`Collection::arrange`] but with the ability to name the operator.
pub fn arrange_named<Ba, Bu, Tr>(self, name: &str) -> Arranged<'scope, TraceAgent<Tr>>
pub fn arrange_named<Ba, Bu, Tr>(self, name: &str, batcher: impl FnOnce(Option<crate::logging::Logger>, usize) -> Ba) -> Arranged<'scope, TraceAgent<Tr>>
where
Ba: crate::trace::Batcher<Output=Vec<((K, V), T, R)>, Time=T> + 'static,
Bu: crate::trace::Builder<Time=T, Input=Vec<((K, V), T, R)>, Output: Into<Tr::Batch>>,
Ba: crate::trace::Batcher<T, Vec<((K, V), T, R)>, Vec<Bu::Input>> + 'static,
Bu: crate::trace::Builder<Time=T, Output: Into<Tr::Batch>>,
Tr: crate::trace::Trace<Time=T> + 'static,
{
let exchange = timely::dataflow::channels::pact::Exchange::new(move |update: &((K,V),T,R)| (update.0).0.hashed().into());
crate::operators::arrange::arrangement::arrange_core::<_, _, ContainerChunker<Vec<((K, V), T, R)>>, Ba, Bu, _>(self.inner, exchange, name)
crate::operators::arrange::arrangement::arrange_core::<_, _, ContainerChunker<Vec<((K, V), T, R)>>, Ba, Bu, _>(self.inner, exchange, name, batcher)
}
}

Expand All @@ -1072,7 +1072,7 @@ pub mod vec {

/// As `arrange_by_key` but with the ability to name the arrangement.
pub fn arrange_by_key_named(self, name: &str) -> Arranged<'scope, TraceAgent<ValSpine<K, V, T, R>>> {
self.arrange_named::<ValBatcher<_,_,_,_>,ValBuilder<_,_,_,_>,_>(name)
self.arrange_named::<_,ValBuilder<_,_,_,_>,_>(name, ValBatcher::new)
}
}

Expand All @@ -1092,7 +1092,7 @@ pub mod vec {
/// As `arrange_by_self` but with the ability to name the arrangement.
pub fn arrange_by_self_named(self, name: &str) -> Arranged<'scope, TraceAgent<KeySpine<K, T, R>>> {
self.map(|k| (k, ()))
.arrange_named::<KeyBatcher<_,_,_>,KeyBuilder<_,_,_>,_>(name)
.arrange_named::<_,KeyBuilder<_,_,_>,_>(name, KeyBatcher::new)
}
}

Expand Down
45 changes: 29 additions & 16 deletions differential-dataflow/src/operators/arrange/arrangement.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,8 @@
use crate::{Data, VecCollection, AsCollection};
use crate::difference::Semigroup;
use crate::lattice::Lattice;
use crate::trace::{self, SpanOf, Trace, TraceReader, Navigable, Batcher, Builder, Cursor, BatchCursor, BatchDiff, BatchKey, BatchVal, BatchValOwn};
use crate::logging::Logger;
use crate::trace::{self, Description, SpanOf, Trace, TraceReader, Navigable, Batcher, Builder, Cursor, BatchCursor, BatchDiff, BatchKey, BatchVal, BatchValOwn};

use trace::wrappers::enter::{TraceEnter, enter_span};

Expand Down Expand Up @@ -184,7 +185,7 @@
while let Some(key) = cursor.get_key(batch) {
while let Some(val) = cursor.get_val(batch) {
for datum in logic(key, val) {
cursor.map_times(batch, |time, diff| {

Check warning on line 188 in differential-dataflow/src/operators/arrange/arrangement.rs

View workflow job for this annotation

GitHub Actions / Cargo clippy 

`time` shadows a previous, unrelated binding
session.give((datum.clone(), <BatchCursor<Tr> as Cursor>::owned_time(time), <BatchCursor<Tr> as Cursor>::owned_diff(diff)));
});
}
Expand Down Expand Up @@ -302,13 +303,18 @@
/// This operator arranges a stream of values into a shared trace, whose contents it maintains.
/// It uses the supplied parallelization contract to distribute the data, which does not need to
/// be consistently by key (though this is the most common).
pub fn arrange_core<'scope, P, C, Chu, Ba, Bu, Tr>(stream: Stream<'scope, Tr::Time, C>, pact: P, name: &str) -> Arranged<'scope, TraceAgent<Tr>>
pub fn arrange_core<'scope, P, C, Chu, Ba, Bu, Tr>(
stream: Stream<'scope, Tr::Time, C>,
pact: P,
name: &str,
batcher: impl FnOnce(Option<Logger>, usize) -> Ba,
) -> Arranged<'scope, TraceAgent<Tr>>
where
C: Container + Clone + 'static,
P: ParallelizationContract<Tr::Time, C>,
Chu: ContainerBuilder<Container=Ba::Output> + for<'a> PushInto<&'a mut C> + 'static,
Ba: Batcher<Time=Tr::Time> + 'static,
Bu: Builder<Time=Tr::Time, Input=Ba::Output, Output: Into<Tr::Batch>>,
Chu: ContainerBuilder + for<'a> PushInto<&'a mut C> + 'static,
Ba: Batcher<Tr::Time, Chu::Container, Vec<Bu::Input>> + 'static,
Bu: Builder<Time=Tr::Time, Output: Into<Tr::Batch>>,
Tr: Trace+'static,
{
// The `Arrange` operator is tasked with reacting to an advancing input
Expand Down Expand Up @@ -338,7 +344,7 @@
let logger = scope.worker().logger_for::<crate::logging::DifferentialEventBuilder>("differential/arrange").map(Into::into);

// Where we will deposit received updates, and from which we extract batches.
let mut batcher = Ba::new(logger.clone(), info.global_id);
let mut batcher = batcher(logger.clone(), info.global_id);

// Capabilities for the lower envelope of updates in `batcher`.
let mut capabilities = Antichain::<Capability<Tr::Time>>::new();
Expand Down Expand Up @@ -371,7 +377,7 @@
}
chunker.push_into(data);
while let Some(chunk) = chunker.extract() {
batcher.push_into(std::mem::take(chunk));
batcher.insert(chunk);
}
});

Expand All @@ -391,7 +397,7 @@
// seal. The batcher only sees chunks the chunker has emitted; without this drain
// a partial final chunk would never reach the batcher.
while let Some(chunk) = chunker.finish() {
batcher.push_into(std::mem::take(chunk));
batcher.insert(chunk);
}

// There are two cases to handle with some care:
Expand All @@ -409,7 +415,7 @@
if capabilities.elements().iter().any(|c| !frontier.less_equal(c.time())) {

// The capabilities to retire: those not in advance of the input frontier.
// Each update sealed below is greater or equal to one of them, as updates
// Each update extracted below is greater or equal to one of them, as updates
// supported only by the remaining capabilities are in advance of the input
// frontier and remain in the batcher.
let retired = capabilities
Expand All @@ -420,8 +426,16 @@
.collect::<CapabilitySet<_>>();

// Extract all updates not in advance of the input frontier, as one batch.
let (mut chain, description) = batcher.seal(frontier.frontier().to_owned());
let batch = trace::Span::new(description, Bu::seal(&mut chain).map(Into::into));
// The batch spans the interval from the previously reported frontier to the
// current one, which is exactly the interval the batcher carves out.
let description = Description::new(
prev_frontier.clone(),
frontier.frontier().to_owned(),
Antichain::from_elem(Tr::Time::minimum()),
);
let (chain, retained) = batcher.extract(frontier.frontier());

let batch = trace::Span::new(description, chain.and_then(|mut chain| Bu::seal(&mut chain)).map(Into::into));

let stamp = retired.iter().map(|c| c.time().clone()).collect::<Stamp<_>>();
writer.insert(batch.clone(), stamp);
Expand All @@ -436,7 +450,7 @@
// in messages with new capabilities.

let mut new_capabilities = Antichain::new();
for time in batcher.frontier().iter() {
for time in retained.iter() {
if let Some(capability) = capabilities.elements().iter().find(|c| c.time().less_equal(time)) {
new_capabilities.insert(capability.delayed(time));
}
Expand All @@ -448,10 +462,9 @@
capabilities = new_capabilities;
}
else {
// Announce progress updates, even without data. We seal the batcher to
// advance its lower bound and frontier, but discard the readied updates
// rather than building a batch we would immediately drop.
let _ = batcher.seal(frontier.frontier().to_owned());
// Announce progress updates, even without data. No held capability precedes
// the input frontier, so no update does either, and the batcher has nothing
// to extract.
writer.seal(frontier.frontier().to_owned());
}

Expand Down
Loading
Loading