use timely::order::TotalOrder;
use timely::dataflow::*;
use timely::dataflow::operators::Operator;
use timely::dataflow::channels::pact::Pipeline;
use crate::lattice::Lattice;
use crate::{ExchangeData, Collection};
use crate::difference::{Semigroup, Abelian};
use crate::hashable::Hashable;
use crate::collection::AsCollection;
use crate::operators::arrange::{Arranged, ArrangeBySelf};
use crate::trace::{BatchReader, Cursor, TraceReader};
pub trait ThresholdTotal<G: Scope, K: ExchangeData, R: ExchangeData+Semigroup> where G::Timestamp: TotalOrder+Lattice+Ord {
fn threshold_semigroup<R2, F>(&self, thresh: F) -> Collection<G, K, R2>
where
R2: Semigroup,
F: FnMut(&K,&R,Option<&R>)->Option<R2>+'static,
;
fn threshold_total<R2: Abelian, F: FnMut(&K,&R)->R2+'static>(&self, mut thresh: F) -> Collection<G, K, R2> {
self.threshold_semigroup(move |key, new, old| {
let mut new = thresh(key, new);
if let Some(old) = old { new.plus_equals(&thresh(key, old).negate()); }
if !new.is_zero() { Some(new) } else { None }
})
}
fn distinct_total(&self) -> Collection<G, K, isize> {
self.distinct_total_core()
}
fn distinct_total_core<R2: Abelian+From<i8>>(&self) -> Collection<G, K, R2> {
self.threshold_total(|_,_| R2::from(1i8))
}
}
impl<G: Scope, K: ExchangeData+Hashable, R: ExchangeData+Semigroup> ThresholdTotal<G, K, R> for Collection<G, K, R>
where G::Timestamp: TotalOrder+Lattice+Ord {
fn threshold_semigroup<R2, F>(&self, thresh: F) -> Collection<G, K, R2>
where
R2: Semigroup,
F: FnMut(&K,&R,Option<&R>)->Option<R2>+'static,
{
self.arrange_by_self_named("Arrange: ThresholdTotal")
.threshold_semigroup(thresh)
}
}
impl<G: Scope, K, T1> ThresholdTotal<G, K, T1::Diff> for Arranged<G, T1>
where
G::Timestamp: TotalOrder+Lattice+Ord,
T1: for<'a> TraceReader<Key<'a>=&'a K, Val<'a>=&'a (), Time=G::Timestamp>+Clone+'static,
K: ExchangeData,
T1::Diff: ExchangeData+Semigroup,
{
fn threshold_semigroup<R2, F>(&self, mut thresh: F) -> Collection<G, K, R2>
where
R2: Semigroup,
F: for<'a> FnMut(T1::Key<'a>,&T1::Diff,Option<&T1::Diff>)->Option<R2>+'static,
{
let mut trace = self.trace.clone();
let mut buffer = Vec::new();
self.stream.unary_frontier(Pipeline, "ThresholdTotal", move |_,_| {
let mut upper_limit = timely::progress::frontier::Antichain::from_elem(<G::Timestamp as timely::progress::Timestamp>::minimum());
move |input, output| {
input.for_each(|capability, batches| {
batches.swap(&mut buffer);
let mut session = output.session(&capability);
for batch in buffer.drain(..) {
let mut batch_cursor = batch.cursor();
let (mut trace_cursor, trace_storage) = trace.cursor_through(batch.lower().borrow()).unwrap();
upper_limit.clone_from(batch.upper());
while let Some(key) = batch_cursor.get_key(&batch) {
let mut count: Option<T1::Diff> = None;
trace_cursor.seek_key(&trace_storage, key);
if trace_cursor.get_key(&trace_storage) == Some(key) {
trace_cursor.map_times(&trace_storage, |_, diff| {
count.as_mut().map(|c| c.plus_equals(diff));
if count.is_none() { count = Some(diff.clone()); }
});
}
batch_cursor.map_times(&batch, |time, diff| {
let difference =
match &count {
Some(old) => {
let mut temp = old.clone();
temp.plus_equals(diff);
thresh(key, &temp, Some(old))
},
None => { thresh(key, diff, None) },
};
if let Some(count) = &mut count {
count.plus_equals(diff);
}
else {
count = Some(diff.clone());
}
if let Some(difference) = difference {
if !difference.is_zero() {
session.give((key.clone(), time.clone(), difference));
}
}
});
batch_cursor.step_key(&batch);
}
}
});
trace.advance_upper(&mut upper_limit);
trace.set_logical_compaction(upper_limit.borrow());
trace.set_physical_compaction(upper_limit.borrow());
}
})
.as_collection()
}
}