use std::{
collections::{BTreeMap, BTreeSet, HashMap},
fmt::{self, Debug, Formatter},
hash::Hash,
marker::PhantomData,
};
use reifydb_value::{Result, reifydb_assertions, value::row_number::RowNumber};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use crate::{
encoded::key::{EncodedKey, IntoEncodedKey},
window::{
accumulator::WindowAccumulator,
engine::{AccumulatorEvent, GroupMeta, LatePolicy, MetaKey, meta_key_for, rolling::RollingBuckets},
span::Slot,
state::StateCache,
store::WindowStore,
},
};
pub type MultiRollingBuffer<C, Accumulator> = BTreeMap<C, Accumulator>;
pub type MultiRollingEmit<SK, Output> = BTreeMap<SK, Output>;
pub enum MultiEmit<Output> {
Insert {
row_number: RowNumber,
value: Output,
},
Update {
row_number: RowNumber,
prior: Output,
value: Output,
},
Remove {
row_number: RowNumber,
value: Output,
},
}
#[derive(Serialize, Deserialize)]
#[serde(bound(
serialize = "C: Serialize + Ord, Accumulator: Serialize, SK: Serialize + Ord, Output: Serialize",
deserialize = "C: serde::de::DeserializeOwned + Ord, Accumulator: serde::de::DeserializeOwned, \
SK: serde::de::DeserializeOwned + Ord, Output: serde::de::DeserializeOwned"
))]
struct GroupState<C, Accumulator, SK, Output> {
buffer: MultiRollingBuffer<C, Accumulator>,
last_emit: MultiRollingEmit<SK, Output>,
}
impl<C: Ord, Accumulator, SK: Ord, Output> Default for GroupState<C, Accumulator, SK, Output> {
fn default() -> Self {
Self {
buffer: BTreeMap::new(),
last_emit: BTreeMap::new(),
}
}
}
impl<C: Ord + Clone, Accumulator: Clone, SK: Ord + Clone, Output: Clone> Clone
for GroupState<C, Accumulator, SK, Output>
{
fn clone(&self) -> Self {
Self {
buffer: self.buffer.clone(),
last_emit: self.last_emit.clone(),
}
}
}
impl<C, Accumulator, SK, Output> Debug for GroupState<C, Accumulator, SK, Output> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_struct("GroupState")
.field("buffer_len", &self.buffer.len())
.field("last_emit_len", &self.last_emit.len())
.finish()
}
}
type MetaLoaded<G, C> = HashMap<G, GroupMeta<C>>;
type StateRows<G> = HashMap<G, RowNumber>;
struct GroupSlot<C, Accumulator, SK, Output> {
state_row_number: RowNumber,
buffer: MultiRollingBuffer<C, Accumulator>,
prior_emit: MultiRollingEmit<SK, Output>,
buffer_changed: bool,
}
pub struct MultiRollingEngine<G, C, Accumulator, SK, Output> {
groups: StateCache<RowNumber, GroupState<C, Accumulator, SK, Output>>,
meta: StateCache<MetaKey, GroupMeta<C>>,
late_policy: LatePolicy,
_pd: PhantomData<G>,
}
impl<G, C, Accumulator, SK, Output> Default for MultiRollingEngine<G, C, Accumulator, SK, Output>
where
G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
C: Slot + Hash + Serialize + DeserializeOwned,
Accumulator: WindowAccumulator,
SK: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
Output: Clone + Debug + PartialEq + Serialize + DeserializeOwned,
for<'a> &'a G: IntoEncodedKey,
{
fn default() -> Self {
Self::new()
}
}
impl<G, C, Accumulator, SK, Output> MultiRollingEngine<G, C, Accumulator, SK, Output>
where
G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
C: Slot + Hash + Serialize + DeserializeOwned,
Accumulator: WindowAccumulator,
SK: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
Output: Clone + Debug + PartialEq + Serialize + DeserializeOwned,
for<'a> &'a G: IntoEncodedKey,
{
pub fn new() -> Self {
Self::with_late_policy(LatePolicy::Drop)
}
pub fn with_late_policy(late_policy: LatePolicy) -> Self {
Self {
groups: StateCache::<RowNumber, GroupState<C, Accumulator, SK, Output>>::new(8),
meta: StateCache::<MetaKey, GroupMeta<C>>::new_internal(64),
late_policy,
_pd: PhantomData,
}
}
pub fn apply<S, SKF, RKF, CB>(
&mut self,
store: &mut S,
buckets: RollingBuckets<G, C, Accumulator::Contribution>,
capacity: usize,
state_key: SKF,
row_key: RKF,
combine: CB,
) -> Result<Vec<MultiEmit<Output>>>
where
S: WindowStore,
SKF: Fn(&G) -> EncodedKey,
RKF: Fn(&G, &SK) -> EncodedKey,
CB: Fn(&G, &MultiRollingBuffer<C, Accumulator>) -> MultiRollingEmit<SK, Output>,
{
if buckets.is_empty() {
return Ok(Vec::new());
}
let mut meta_loaded = self.warm_and_load_meta(store, &buckets)?;
let state_rows = self.resolve_state_rows(store, &buckets, &meta_loaded, &state_key)?;
let group_slots = self.apply_events_into_buffers(
store,
buckets,
&mut meta_loaded,
&state_rows,
&state_key,
capacity,
)?;
let emits = self.diff_emits(store, group_slots, &row_key, &combine)?;
self.persist_meta(store, meta_loaded)?;
Ok(emits)
}
pub fn flush<S: WindowStore>(&mut self, store: &mut S) -> Result<()> {
self.groups.flush(store)?;
self.meta.flush(store)?;
Ok(())
}
fn warm_and_load_meta<S: WindowStore>(
&mut self,
store: &mut S,
buckets: &RollingBuckets<G, C, Accumulator::Contribution>,
) -> Result<MetaLoaded<G, C>> {
let meta_keys: Vec<MetaKey> = buckets
.keys()
.map(|(group, _)| group)
.collect::<BTreeSet<_>>()
.into_iter()
.map(meta_key_for)
.collect();
self.meta.warm(store, &meta_keys)?;
let mut meta_loaded: MetaLoaded<G, C> = HashMap::new();
for (group, _) in buckets.keys() {
if !meta_loaded.contains_key(group) {
let m = self.meta.get(store, &meta_key_for(group))?.unwrap_or_default();
meta_loaded.insert(group.clone(), m);
}
}
Ok(meta_loaded)
}
fn resolve_state_rows<S, SKF>(
&mut self,
store: &mut S,
buckets: &RollingBuckets<G, C, Accumulator::Contribution>,
meta_loaded: &MetaLoaded<G, C>,
state_key: &SKF,
) -> Result<StateRows<G>>
where
S: WindowStore,
SKF: Fn(&G) -> EncodedKey,
{
let mut state_rows: StateRows<G> = HashMap::new();
let mut resolve_order: Vec<G> = Vec::new();
let mut state_lookup_keys: Vec<EncodedKey> = Vec::new();
let mut seen: BTreeSet<G> = BTreeSet::new();
for (group, coord) in buckets.keys() {
let initial_high_water = meta_loaded.get(group).and_then(|m| m.high_water);
if initial_high_water.is_none_or(|hw| *coord >= hw) && seen.insert(group.clone()) {
resolve_order.push(group.clone());
state_lookup_keys.push(state_key(group));
}
}
let resolved_rows = store.get_or_create_row_numbers(&state_lookup_keys)?;
reifydb_assertions! {
let resolved = resolved_rows.len();
let requested = state_lookup_keys.len();
assert!(
resolved == requested,
"get_or_create_row_numbers returned {resolved} rows for {requested} group keys; \
the zip below pairs resolve_order with resolved_rows by position, so a length \
mismatch would silently leave some groups without a state_rows entry and route \
them through the per-bucket get_or_create_row_number fallback, diverging behaviour"
);
}
let state_keys: Vec<RowNumber> = resolved_rows.iter().map(|(rn, _)| *rn).collect();
for (group, (state_row_number, _)) in resolve_order.into_iter().zip(resolved_rows) {
state_rows.insert(group, state_row_number);
}
self.groups.warm(store, &state_keys)?;
Ok(state_rows)
}
fn apply_events_into_buffers<S, SKF>(
&mut self,
store: &mut S,
buckets: RollingBuckets<G, C, Accumulator::Contribution>,
meta_loaded: &mut MetaLoaded<G, C>,
state_rows: &StateRows<G>,
state_key: &SKF,
capacity: usize,
) -> Result<BTreeMap<G, GroupSlot<C, Accumulator, SK, Output>>>
where
S: WindowStore,
SKF: Fn(&G) -> EncodedKey,
{
let mut group_slots: BTreeMap<G, GroupSlot<C, Accumulator, SK, Output>> = BTreeMap::new();
for ((group, coord), events) in buckets {
let meta = meta_loaded.entry(group.clone()).or_default();
let slot = match group_slots.get_mut(&group) {
Some(s) => s,
None => {
let state_row_number = match state_rows.get(&group) {
Some(&rn) => rn,
None => {
let key = state_key(&group);
let (rn, _is_new) = store.get_or_create_row_number(&key)?;
rn
}
};
let GroupState {
buffer,
last_emit: prior_emit,
} = self.groups.get(store, &state_row_number)?.unwrap_or_default();
group_slots.insert(
group.clone(),
GroupSlot {
state_row_number,
buffer,
prior_emit,
buffer_changed: false,
},
);
group_slots.get_mut(&group).expect("just inserted")
}
};
let late = matches!(meta.high_water, Some(hw) if coord < hw)
&& matches!(self.late_policy, LatePolicy::Drop)
&& !slot.buffer.contains_key(&coord);
let mut accumulator = slot.buffer.remove(&coord).unwrap_or_default();
let mut touched = false;
for event in events {
match event {
AccumulatorEvent::Add(c) => {
if late {
continue;
}
accumulator.add(&c);
touched = true;
}
AccumulatorEvent::Remove(c) => {
if accumulator.is_empty() {
continue;
}
accumulator.remove(&c);
touched = true;
}
}
}
if !accumulator.is_empty() {
slot.buffer.insert(coord, accumulator);
}
if !touched {
continue;
}
while slot.buffer.len() > capacity {
slot.buffer.pop_first();
}
slot.buffer_changed = true;
let next_high_water = match meta.high_water {
Some(hw) if hw > coord => hw,
_ => coord,
};
reifydb_assertions! {
assert!(
next_high_water >= coord,
"high_water regressed below the window coord it just admitted, so the next batch would \
treat an already-processed window as late and silently drop its events (coord={coord:?}, \
prev_high_water={prev:?}, next_high_water={next_high_water:?})",
prev = meta.high_water
);
if let Some(prev) = meta.high_water {
assert!(
next_high_water >= prev,
"high_water moved backwards across an admit, breaking the monotonic late-event \
cutoff that buried-window dropping relies on (coord={coord:?}, prev_high_water={prev:?}, \
next_high_water={next_high_water:?})"
);
}
}
meta.high_water = Some(next_high_water);
}
Ok(group_slots)
}
fn diff_emits<S, RKF, CB>(
&mut self,
store: &mut S,
group_slots: BTreeMap<G, GroupSlot<C, Accumulator, SK, Output>>,
row_key: &RKF,
combine: &CB,
) -> Result<Vec<MultiEmit<Output>>>
where
S: WindowStore,
RKF: Fn(&G, &SK) -> EncodedKey,
CB: Fn(&G, &MultiRollingBuffer<C, Accumulator>) -> MultiRollingEmit<SK, Output>,
{
let mut emits: Vec<MultiEmit<Output>> = Vec::new();
for (group, slot) in group_slots {
if !slot.buffer_changed {
continue;
}
let new_emit = combine(&group, &slot.buffer);
for (sk, new_out) in &new_emit {
let key = row_key(&group, sk);
let (rn, _is_new_alloc) = store.get_or_create_row_number(&key)?;
match slot.prior_emit.get(sk) {
Some(prior_out) => {
if prior_out != new_out {
emits.push(MultiEmit::Update {
row_number: rn,
prior: prior_out.clone(),
value: new_out.clone(),
});
}
}
None => {
emits.push(MultiEmit::Insert {
row_number: rn,
value: new_out.clone(),
});
}
}
}
for (sk, prior_out) in &slot.prior_emit {
if !new_emit.contains_key(sk) {
let key = row_key(&group, sk);
let (rn, _is_new_alloc) = store.get_or_create_row_number(&key)?;
emits.push(MultiEmit::Remove {
row_number: rn,
value: prior_out.clone(),
});
}
}
let combined = GroupState {
buffer: slot.buffer,
last_emit: new_emit,
};
self.groups.put(store, &slot.state_row_number, combined)?;
}
Ok(emits)
}
fn persist_meta<S: WindowStore>(&mut self, store: &mut S, meta_loaded: MetaLoaded<G, C>) -> Result<()> {
for (group, meta) in meta_loaded {
self.meta.set(store, &meta_key_for(&group), &meta)?;
}
Ok(())
}
}