use std::{collections::BTreeMap, fmt::Debug, hash::Hash};
use reifydb_codec::{
key::encoded::{EncodedKey, IntoEncodedKey},
row::operator::state::StateCodec,
};
use reifydb_core::{
interface::{catalog::flow::OperatorId, flow::OperatorCapability},
metrics::heap::{HeapSize, OperatorSample},
};
#[cfg(reifydb_assertions)]
use reifydb_flow::operator::state::reaper::queued;
use reifydb_flow::{
operator::state::{
reaper::{drain, enqueue},
seal::{coord::Coord, domain::SealDomain, rule::is_sealed},
},
window::{
accumulator::WindowAccumulator,
engine::{
AccumulatorEvent, EmitKind,
tumbling::{TumblingBuckets, TumblingEngine},
},
span::{Slot, SlotCoord, WindowAnchor, WindowSpan},
},
};
use reifydb_value::{
config::Config,
reifydb_assertions,
value::{diff_type::DiffType, duration::Duration, row_number::RowNumber},
};
use tracing::debug;
use crate::{
error::Result,
flow::operator::{
GuestOperator, OperatorMetadata,
column::{
batch::{InsertBatch, RemoveBatch, UpdateBatch},
operator::OperatorColumn,
row::Row,
},
context::GuestContext,
timer::Timer,
view::{ChangeView, ColumnsView, DiffView, RowView},
windowed::{
group_of, guest_as_host::GuestAsHost, intern_window_groups, observe_batch, seal_frontier,
timer_frontier, window_engine_config,
},
},
};
const SEAL_REAP_BATCH: usize = 256;
type AccumulatorContribution<A> = <<A as TumblingOperator>::Accumulator as WindowAccumulator>::Contribution;
type AccumulatorValue<A> = <<A as TumblingOperator>::Accumulator as WindowAccumulator>::Output;
type Anchor<A> = SlotCoord<<A as TumblingOperator>::WindowSlot>;
type Lateness<A> = <Anchor<A> as SealDomain>::Lateness;
type Buckets<A> = TumblingBuckets<<A as TumblingOperator>::GroupKey, Anchor<A>, AccumulatorContribution<A>>;
type WindowOrder<A> = Vec<(<A as TumblingOperator>::GroupKey, WindowSpan<Anchor<A>>)>;
pub trait TumblingOperator {
type GroupKey: Clone + Eq + Ord + Hash + Debug + StateCodec;
type WindowSlot: Slot<Coord: WindowAnchor + SealDomain + Hash + StateCodec + HeapSize + Send + Sync>
+ Hash
+ StateCodec
+ HeapSize;
type Accumulator: WindowAccumulator;
type Output: Clone + Debug + PartialEq;
fn coord(&self, row: &impl RowView) -> Option<Self::WindowSlot>;
fn extract(
&self,
ctx: &mut impl GuestContext,
row: &impl RowView,
) -> Option<(Self::GroupKey, AccumulatorContribution<Self>)>;
fn window_for(&self, coord: SlotCoord<Self::WindowSlot>) -> WindowSpan<SlotCoord<Self::WindowSlot>>;
fn lateness(&self) -> Option<<SlotCoord<Self::WindowSlot> as SealDomain>::Lateness> {
None
}
fn build_output(
&self,
group: &Self::GroupKey,
span: WindowSpan<SlotCoord<Self::WindowSlot>>,
value: AccumulatorValue<Self>,
) -> Option<Self::Output>;
fn new_accumulator(&self) -> Self::Accumulator {
Self::Accumulator::default()
}
}
pub trait TumblingRegistration: TumblingOperator + Sized
where
Self::Output: Row,
for<'a> &'a Self::GroupKey: IntoEncodedKey,
{
const NAME: &'static str;
const VERSION: &'static str;
const DESCRIPTION: &'static str;
const INPUT_COLUMNS: &'static [OperatorColumn];
const OUTPUT_COLUMNS: &'static [OperatorColumn];
const CAPABILITIES: &'static [OperatorCapability];
fn from_config(operator_id: OperatorId, config: &Config) -> Result<Self>;
fn encode_row_key(&self, group: &Self::GroupKey, window_start: SlotCoord<Self::WindowSlot>) -> EncodedKey;
}
pub struct TumblingDriver<A>
where
A: TumblingRegistration,
A::Output: Row,
for<'a> &'a A::GroupKey: IntoEncodedKey,
{
aggregator: A,
engine: TumblingEngine<A::GroupKey, Anchor<A>, A::Accumulator>,
reap_queue_empty: bool,
}
impl<A> TumblingDriver<A>
where
A: TumblingRegistration,
A::Output: Row,
for<'a> &'a A::GroupKey: IntoEncodedKey,
{
fn route(&self, ctx: &mut impl GuestContext, change: &impl ChangeView) -> Buckets<A> {
let mut buckets: Buckets<A> = BTreeMap::new();
for di in 0..change.diff_count() {
let Some(diff) = change.diff(di) else {
continue;
};
match diff.kind() {
DiffType::Insert => {
if let Some(cols) = diff.post() {
self.push_all(ctx, &cols, &mut buckets, true);
}
}
DiffType::Update => {
if let (Some(pre), Some(post)) = (diff.pre(), diff.post()) {
self.push_all(ctx, &pre, &mut buckets, false);
self.push_all(ctx, &post, &mut buckets, true);
}
}
DiffType::Remove => {
if let Some(cols) = diff.pre() {
self.push_all(ctx, &cols, &mut buckets, false);
}
}
}
}
buckets
}
fn push_all<C: ColumnsView>(
&self,
ctx: &mut impl GuestContext,
cols: &C,
buckets: &mut Buckets<A>,
is_add: bool,
) {
for i in 0..cols.row_count() {
let Some(row) = cols.row(i) else {
continue;
};
let Some(slot) = self.aggregator.coord(&row) else {
continue;
};
let Some((group, contribution)) = self.aggregator.extract(ctx, &row) else {
continue;
};
let span = self.aggregator.window_for(slot.order_key());
let event = if is_add {
AccumulatorEvent::Add(contribution)
} else {
AccumulatorEvent::Remove(contribution)
};
buckets.entry((group, span)).or_default().push(event);
}
}
}
impl<A> TumblingDriver<A>
where
A: TumblingRegistration + Send + Sync + 'static,
A::Output: Row,
A::GroupKey: Send + Sync,
A::WindowSlot: Send + Sync,
A::Accumulator: Send + Sync,
AccumulatorContribution<A>: Send + Sync,
for<'a> &'a A::GroupKey: IntoEncodedKey,
{
fn expire_through<C: GuestContext>(
engine: &mut TumblingEngine<A::GroupKey, Anchor<A>, A::Accumulator>,
reap_queue_empty: &mut bool,
store: &mut GuestAsHost<'_, C>,
frontier: Anchor<A>,
lateness: Lateness<A>,
) -> Result<()> {
let horizon = <Anchor<A> as SealDomain>::horizon(frontier, lateness);
if horizon <= <Anchor<A> as Coord>::from_order(0) {
return Ok(());
}
for window in engine.expire(store, horizon.to_order().saturating_sub(1))? {
enqueue(store, window.group_id)?;
*reap_queue_empty = false;
}
engine.expire_meta(store, horizon.to_order())?;
if *reap_queue_empty {
reifydb_assertions! {
let pending = queued(store, 1)?;
assert!(
pending.groups.is_empty(),
"the reap queue still holds {:?} while the driver believes it drained the \
queue; skipping the drain here leaves that group's state behind for good",
pending.groups
);
}
return Ok(());
}
let drained = drain(store, engine, SEAL_REAP_BATCH)?;
*reap_queue_empty = drained.queue_is_empty();
if !*reap_queue_empty {
observe_batch(store, frontier, lateness)?;
}
Ok(())
}
#[inline]
fn emit_batches(
&self,
ctx: &mut impl GuestContext,
inserts: &[(RowNumber, A::Output)],
updates: &[(RowNumber, A::Output)],
removes: &[(RowNumber, A::Output)],
) -> Result<()> {
if !inserts.is_empty() {
let mut batch = InsertBatch::<A::Output, _>::new(ctx, inserts.len())?;
for (rn, data) in inserts {
batch.push(*rn, data)?;
}
batch.finish()?;
}
if !updates.is_empty() {
let mut batch = UpdateBatch::<A::Output, _>::new(ctx, updates.len())?;
for (rn, data) in updates {
batch.push(*rn, data, data)?;
}
batch.finish()?;
}
if !removes.is_empty() {
let mut batch = RemoveBatch::<A::Output, _>::new(ctx, removes.len())?;
for (rn, data) in removes {
batch.push(*rn, data)?;
}
batch.finish()?;
}
Ok(())
}
}
impl<A> OperatorMetadata for TumblingDriver<A>
where
A: TumblingRegistration + 'static,
A::Output: Row,
for<'a> &'a A::GroupKey: IntoEncodedKey,
{
const NAME: &'static str = A::NAME;
const VERSION: &'static str = A::VERSION;
const DESCRIPTION: &'static str = A::DESCRIPTION;
const INPUT_COLUMNS: &'static [OperatorColumn] = A::INPUT_COLUMNS;
const OUTPUT_COLUMNS: &'static [OperatorColumn] = A::OUTPUT_COLUMNS;
const CAPABILITIES: &'static [OperatorCapability] = A::CAPABILITIES;
}
impl<A> GuestOperator for TumblingDriver<A>
where
A: TumblingRegistration + Send + Sync + 'static,
A::Output: Row,
A::GroupKey: Send + Sync,
A::WindowSlot: Send + Sync,
A::Accumulator: Send + Sync + HeapSize,
AccumulatorContribution<A>: Send + Sync,
for<'a> &'a A::GroupKey: IntoEncodedKey,
{
fn sample(&self) -> Option<OperatorSample> {
None
}
fn create(operator_id: OperatorId, config: &Config) -> Result<Self> {
let aggregator = A::from_config(operator_id, config)?;
let engine_config = window_engine_config(config);
Ok(Self {
aggregator,
engine: TumblingEngine::new(engine_config),
reap_queue_empty: false,
})
}
fn on_timer(&mut self, ctx: &mut impl GuestContext, timer: Timer<'_>) -> Result<()> {
let Some(lateness) = self.aggregator.lateness() else {
return Ok(());
};
let Self {
engine,
reap_queue_empty,
..
} = &mut *self;
let mut store = GuestAsHost(ctx);
let Some(frontier) = timer_frontier::<Anchor<A>>(&mut store, timer)? else {
return Ok(());
};
Self::expire_through(engine, reap_queue_empty, &mut store, frontier, lateness)
}
fn lateness(&self) -> Option<Duration> {
self.aggregator.lateness().and_then(<Anchor<A> as SealDomain>::lateness_duration)
}
fn apply(&mut self, ctx: &mut impl GuestContext, change: impl ChangeView) -> Result<()> {
let mut buckets = self.route(ctx, &change);
if buckets.is_empty() {
return Ok(());
}
let lateness = self.aggregator.lateness();
if let Some(lateness) = lateness {
let Self {
engine,
reap_queue_empty,
..
} = &mut *self;
let mut store = GuestAsHost(ctx);
let newest = buckets.keys().map(|(_, span)| span.start).max();
if let Some(newest) = newest {
observe_batch(&mut store, newest, lateness)?;
}
let watermark = seal_frontier::<Anchor<A>>(&mut store)?;
let horizon = <Anchor<A> as SealDomain>::horizon(watermark, lateness);
let mut dropped = 0u64;
buckets.retain(|(_, span), events| {
if is_sealed(span.start, horizon) {
dropped += events.len() as u64;
false
} else {
true
}
});
if dropped > 0 {
debug!(operator = A::NAME, dropped, "mutations targeting sealed windows were dropped");
}
Self::expire_through(engine, reap_queue_empty, &mut store, watermark, lateness)?;
if buckets.is_empty() {
return Ok(());
}
}
let groups = intern_window_groups(buckets.keys().map(|(group, span)| {
((group.clone(), span.start), self.aggregator.encode_row_key(group, span.start))
}));
let results = {
let Self {
aggregator,
engine,
..
} = &mut *self;
let mut store = GuestAsHost(ctx);
let order: WindowOrder<A> = buckets.keys().cloned().collect();
engine.apply(
&mut store,
buckets,
&order,
|group, window_start| {
(
group_of(&groups, group, window_start),
aggregator.encode_row_key(group, window_start),
)
},
|| aggregator.new_accumulator(),
)?
};
if lateness.is_some() {
let mut store = GuestAsHost(ctx);
for r in &results {
if r.kind == EmitKind::Insert {
let group = group_of(&groups, &r.group, r.span.start);
self.engine.reindex_window(
&mut store,
&r.group,
r.span.start,
group,
&self.aggregator.encode_row_key(&r.group, r.span.start),
None,
Some(r.span.start.to_order()),
)?;
}
}
}
let mut inserts: Vec<(RowNumber, A::Output)> = Vec::new();
let mut updates: Vec<(RowNumber, A::Output)> = Vec::new();
let mut removes: Vec<(RowNumber, A::Output)> = Vec::new();
for r in results {
let Some(out) = self.aggregator.build_output(&r.group, r.span, r.value) else {
continue;
};
match r.kind {
EmitKind::Insert => inserts.push((r.row_number, out)),
EmitKind::Update => updates.push((r.row_number, out)),
EmitKind::Remove => removes.push((r.row_number, out)),
}
}
self.emit_batches(ctx, &inserts, &updates, &removes)?;
Ok(())
}
}