reifydb_core/window/engine/

tumbling.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

use std::{
	collections::{BTreeMap, BTreeSet, HashMap},
	fmt::Debug,
	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},
	util::encoding::keycode::encode_u64,
	window::{
		accumulator::WindowAccumulator,
		engine::{
			AccumulatorEvent, EmitKind, GroupMeta, LatePolicy, MetaKey, WindowResult, expiry_due_range,
			expiry_key, meta_key_for,
		},
		span::{Slot, WindowSpan},
		state::StateCache,
		store::WindowStore,
	},
};

pub type TumblingBuckets<G, C, Contribution> = BTreeMap<(G, WindowSpan<C>), Vec<AccumulatorEvent<Contribution>>>;

type MetaLoaded<G, C> = HashMap<G, GroupMeta<C>>;
type SlotResolved = Vec<Option<(RowNumber, bool)>>;

pub struct ExpiredWindow<G, C, Output> {
	pub row_number: RowNumber,
	pub group: G,
	pub window_start: C,
	pub value: Option<Output>,
}

#[derive(Serialize, Deserialize)]
#[serde(bound(serialize = "G: Serialize, C: Serialize", deserialize = "G: DeserializeOwned, C: DeserializeOwned"))]
struct TumblingIndexEntry<G, C> {
	group: G,
	window_start: C,
	row_number: u64,
}

pub fn reindex_window<S, G, C>(
	store: &mut S,
	group: &G,
	window_start: C,
	row_number: RowNumber,
	prior: Option<u64>,
	new: Option<u64>,
) -> Result<()>
where
	S: WindowStore,
	G: Clone + Serialize,
	C: Slot + Serialize,
	for<'a> &'a G: IntoEncodedKey,
{
	if prior == new {
		return Ok(());
	}
	let suffix = encode_u64(window_start.order_key());
	if let Some(old) = prior {
		store.internal_drop(&expiry_key(old, group, &suffix))?;
	}
	if let Some(new) = new {
		store.internal_set(
			&expiry_key(new, group, &suffix),
			&TumblingIndexEntry {
				group: group.clone(),
				window_start,
				row_number: row_number.0,
			},
		)?;
	}
	Ok(())
}

pub struct TumblingEngine<G, C, Accumulator> {
	accumulators: StateCache<RowNumber, Accumulator>,
	meta: StateCache<MetaKey, GroupMeta<C>>,
	late_policy: LatePolicy,
	_pd: PhantomData<G>,
}

impl<G, C, Accumulator> Default for TumblingEngine<G, C, Accumulator>
where
	G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
	C: Slot + Hash + Serialize + DeserializeOwned,
	Accumulator: WindowAccumulator,
	for<'a> &'a G: IntoEncodedKey,
{
	fn default() -> Self {
		Self::new()
	}
}

impl<G, C, Accumulator> TumblingEngine<G, C, Accumulator>
where
	G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
	C: Slot + Hash + Serialize + DeserializeOwned,
	Accumulator: WindowAccumulator,
	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 {
			accumulators: StateCache::<RowNumber, Accumulator>::new(8),
			meta: StateCache::<MetaKey, GroupMeta<C>>::new_internal(64),
			late_policy,
			_pd: PhantomData,
		}
	}

	pub fn apply<S, K, NA>(
		&mut self,
		store: &mut S,
		buckets: TumblingBuckets<G, C, Accumulator::Contribution>,
		row_key: K,
		new_accumulator: NA,
	) -> Result<Vec<WindowResult<G, C, Accumulator::Output>>>
	where
		S: WindowStore,
		K: Fn(&G, C) -> EncodedKey,
		NA: Fn() -> Accumulator,
	{
		if buckets.is_empty() {
			return Ok(Vec::new());
		}
		let mut meta_loaded = self.warm_and_load_meta(store, &buckets)?;
		let slot_resolved = self.resolve_survivor_rows(store, &buckets, &meta_loaded, &row_key)?;
		let results =
			self.apply_events(store, buckets, slot_resolved, &mut meta_loaded, &row_key, &new_accumulator)?;
		self.persist_meta(store, meta_loaded)?;
		Ok(results)
	}

	pub fn flush<S: WindowStore>(&mut self, store: &mut S) -> Result<()> {
		self.accumulators.flush(store)?;
		self.meta.flush(store)?;
		Ok(())
	}

	fn warm_and_load_meta<S: WindowStore>(
		&mut self,
		store: &mut S,
		buckets: &TumblingBuckets<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_survivor_rows<S, K>(
		&mut self,
		store: &mut S,
		buckets: &TumblingBuckets<G, C, Accumulator::Contribution>,
		meta_loaded: &MetaLoaded<G, C>,
		row_key: &K,
	) -> Result<SlotResolved>
	where
		S: WindowStore,
		K: Fn(&G, C) -> EncodedKey,
	{
		let mut survivor_keys: Vec<EncodedKey> = Vec::new();
		let mut slot_survives: Vec<bool> = Vec::with_capacity(buckets.len());
		for (group, span) in buckets.keys() {
			let initial_high_water = meta_loaded.get(group).and_then(|m| m.high_water);
			let survives = initial_high_water.is_none_or(|hw| span.start >= hw);
			slot_survives.push(survives);
			if survives {
				survivor_keys.push(row_key(group, span.start));
			}
		}
		let resolved_rows = store.get_or_create_row_numbers(&survivor_keys)?;
		reifydb_assertions! {
			let survivors = survivor_keys.len();
			let resolved = resolved_rows.len();
			assert!(
				resolved == survivors,
				"get_or_create_row_numbers must return exactly one row per survivor key; a short batch would \
				 leave a surviving slot with no resolved row, so the slot_resolved zip below pairs it with None \
				 and apply_events silently re-creates a fresh row instead of reusing the existing window \
				 state, double-counting it (survivor_keys={survivors}, resolved_rows={resolved})"
			);
		}
		let accumulator_keys: Vec<RowNumber> = resolved_rows.iter().map(|(rn, _)| *rn).collect();
		self.accumulators.warm(store, &accumulator_keys)?;
		let mut resolved_rows = resolved_rows.into_iter();
		let slot_resolved: SlotResolved = slot_survives
			.into_iter()
			.map(|survives| {
				if survives {
					resolved_rows.next()
				} else {
					None
				}
			})
			.collect();
		Ok(slot_resolved)
	}

	fn apply_events<S, K, NA>(
		&mut self,
		store: &mut S,
		buckets: TumblingBuckets<G, C, Accumulator::Contribution>,
		slot_resolved: SlotResolved,
		meta_loaded: &mut MetaLoaded<G, C>,
		row_key: &K,
		new_accumulator: &NA,
	) -> Result<Vec<WindowResult<G, C, Accumulator::Output>>>
	where
		S: WindowStore,
		K: Fn(&G, C) -> EncodedKey,
		NA: Fn() -> Accumulator,
	{
		let mut results: Vec<WindowResult<G, C, Accumulator::Output>> = Vec::new();

		for (((group, span), events), slot_pre) in buckets.into_iter().zip(slot_resolved) {
			let entry = meta_loaded.entry(group.clone()).or_default();
			match entry.high_water {
				Some(hw) if span.start < hw => {
					if matches!(self.late_policy, LatePolicy::Drop) {
						continue;
					}
				}
				Some(hw) if span.start > hw => entry.high_water = Some(span.start),
				Some(_) => {}
				None => entry.high_water = Some(span.start),
			}

			let (row_number, is_new) = match slot_pre {
				Some(resolved) => resolved,
				None => {
					let key = row_key(&group, span.start);
					store.get_or_create_row_number(&key)?
				}
			};

			let mut accumulator: Accumulator =
				self.accumulators.get(store, &row_number)?.unwrap_or_else(new_accumulator);
			let was_empty_before = accumulator.is_empty();
			let prior = if was_empty_before {
				None
			} else {
				accumulator.finalize()
			};

			for event in events {
				match event {
					AccumulatorEvent::Add(c) => accumulator.add(&c),
					AccumulatorEvent::Remove(c) => accumulator.remove(&c),
				}
			}

			let value = accumulator.finalize();
			self.accumulators.put(store, &row_number, accumulator)?;

			match value {
				Some(value) => {
					let kind = if is_new || was_empty_before {
						EmitKind::Insert
					} else {
						EmitKind::Update
					};
					results.push(WindowResult {
						row_number,
						group,
						span,
						value,
						prior,
						kind,
					});
				}
				None => {
					if let Some(p) = prior.clone() {
						results.push(WindowResult {
							row_number,
							group,
							span,
							value: p,
							prior,
							kind: EmitKind::Remove,
						});
					}
				}
			}
		}
		Ok(results)
	}

	pub fn expire<S: WindowStore>(
		&mut self,
		store: &mut S,
		threshold: u64,
	) -> Result<Vec<ExpiredWindow<G, C, Accumulator::Output>>> {
		let mut due: Vec<(EncodedKey, TumblingIndexEntry<G, C>)> = Vec::new();
		store.internal_range_visit::<TumblingIndexEntry<G, C>>(
			expiry_due_range(threshold),
			&mut |key, entry| {
				due.push((key, entry));
				Ok(())
			},
		)?;

		let mut out: Vec<ExpiredWindow<G, C, Accumulator::Output>> = Vec::new();
		for (index_key, entry) in due {
			let row_number = RowNumber(entry.row_number);
			store.internal_drop(&index_key)?;
			let value = self
				.accumulators
				.get(store, &row_number)?
				.and_then(|accumulator| accumulator.finalize());
			self.accumulators.remove(store, &row_number)?;
			out.push(ExpiredWindow {
				row_number,
				group: entry.group,
				window_start: entry.window_start,
				value,
			});
		}
		Ok(out)
	}

	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(())
	}
}

#[cfg(test)]
mod tests {
	use std::collections::BTreeMap;

	use crate::{
		encoded::key::EncodedKey,
		window::{
			engine::{
				AccumulatorEvent, WindowResult,
				test_support::{MockStore, SumAccumulator},
				tumbling::{TumblingBuckets, TumblingEngine, reindex_window},
			},
			span::WindowSpan,
		},
	};

	fn row_key(group: &u32, window_start: u64) -> EncodedKey {
		EncodedKey::builder().u32(*group).u64(window_start).build()
	}

	fn seed_window(store: &mut MockStore, window_start: u64, contribution: i64) -> WindowResult<u32, u64, i64> {
		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		let mut buckets: TumblingBuckets<u32, u64, i64> = BTreeMap::new();
		buckets.insert(
			(1u32, WindowSpan::new(window_start, window_start + 1)),
			vec![AccumulatorEvent::Add(contribution)],
		);
		let mut results = engine.apply(store, buckets, row_key, SumAccumulator::default).expect("apply");
		engine.flush(store).expect("flush");
		results.pop().expect("one window")
	}

	#[test]
	fn expire_returns_only_due_windows_and_clears_their_state() {
		let mut store = MockStore::default();
		// Two live windows; the face indexes each by its last_event_time (10 and 90).
		let w0 = seed_window(&mut store, 0, 5);
		reindex_window(&mut store, &w0.group, w0.span.start, w0.row_number, None, Some(10)).unwrap();
		let w100 = seed_window(&mut store, 100, 7);
		reindex_window(&mut store, &w100.group, w100.span.start, w100.row_number, None, Some(90)).unwrap();
		assert_eq!(store.index_entry_count(), 2, "both live windows are indexed");

		// Threshold 10: only the window whose expiry (10) is at/under the threshold is due.
		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		let expired = engine.expire(&mut store, 10).unwrap();
		engine.flush(&mut store).unwrap();
		assert_eq!(expired.len(), 1, "exactly one window is due, not the whole population");
		assert_eq!(expired[0].window_start, 0);
		assert_eq!(expired[0].value, Some(5));
		assert_eq!(store.index_entry_count(), 1, "the due window's index entry is gone, the other remains");

		// The surviving window finalizes correctly once the threshold reaches it.
		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		let later = engine.expire(&mut store, 1000).unwrap();
		assert_eq!(later.len(), 1);
		assert_eq!(later[0].window_start, 100);
		assert_eq!(later[0].value, Some(7));
		assert_eq!(store.index_entry_count(), 0);
	}

	#[test]
	fn expire_threshold_is_inclusive() {
		let mut store = MockStore::default();
		let w = seed_window(&mut store, 0, 4);
		reindex_window(&mut store, &w.group, w.span.start, w.row_number, None, Some(50)).unwrap();

		// One below the expiry: not due, and the scan leaves the index intact.
		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		assert!(engine.expire(&mut store, 49).unwrap().is_empty());
		engine.flush(&mut store).unwrap();
		assert_eq!(store.index_entry_count(), 1);

		// Exactly at the expiry: due (the face folds the strict close boundary into the threshold).
		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		assert_eq!(engine.expire(&mut store, 50).unwrap().len(), 1);
	}

	#[test]
	fn reindex_rekeys_without_leaving_a_stale_entry() {
		let mut store = MockStore::default();
		let w = seed_window(&mut store, 0, 9);
		// Index at 10, then a later event advances the window's expiry to 80.
		reindex_window(&mut store, &w.group, w.span.start, w.row_number, None, Some(10)).unwrap();
		reindex_window(&mut store, &w.group, w.span.start, w.row_number, Some(10), Some(80)).unwrap();
		assert_eq!(store.index_entry_count(), 1, "re-keying must not leave the old entry behind");

		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		assert!(engine.expire(&mut store, 10).unwrap().is_empty(), "no longer due at the old expiry");
		let mut engine = TumblingEngine::<u32, u64, SumAccumulator>::new();
		assert_eq!(engine.expire(&mut store, 80).unwrap().len(), 1, "due at the new expiry");
	}
}