use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::hash::Hash;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
pub trait Mergeable {
fn merge(&mut self, other: Self);
}
#[derive(Debug)]
pub enum SpillError {
Io(std::io::Error),
BatchTooLarge { size: usize, limit: usize },
Codec(String),
}
impl From<std::io::Error> for SpillError {
fn from(e: std::io::Error) -> Self {
Self::Io(e)
}
}
impl std::fmt::Display for SpillError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(e) => write!(f, "spill i/o: {e}"),
Self::BatchTooLarge { size, limit } => {
write!(f, "spill batch {size} bytes exceeds limit {limit}")
}
Self::Codec(msg) => write!(f, "spill codec: {msg}"),
}
}
}
impl std::error::Error for SpillError {}
struct SpillBatch {
path: PathBuf,
bytes: u64,
entries: usize,
}
impl Drop for SpillBatch {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.path);
}
}
pub(crate) fn read_spill_header<R: Read, const N: usize>(
reader: &mut R,
) -> std::io::Result<Option<[u8; N]>> {
assert!(N > 0, "invariant: spill headers have a positive size");
let mut bytes = [0; N];
let read = loop {
match reader.read(&mut bytes) {
Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
result => break result?,
}
};
if read == 0 {
return Ok(None);
}
reader.read_exact(&mut bytes[read..])?;
Ok(Some(bytes))
}
pub const DEFAULT_SPILL_DECODE_BYTES: usize = 64 * 1024 * 1024;
pub struct SpillDecodeBudget {
limit_bytes: usize,
remaining_bytes: usize,
}
impl SpillDecodeBudget {
pub fn new(limit_bytes: usize) -> Self {
Self {
limit_bytes,
remaining_bytes: limit_bytes,
}
}
pub fn remaining_bytes(&self) -> usize {
self.remaining_bytes
}
pub fn charge_bytes(&mut self, bytes: usize) -> std::io::Result<()> {
let Some(remaining) = self.remaining_bytes.checked_sub(bytes) else {
return Err(std::io::Error::new(
std::io::ErrorKind::OutOfMemory,
format!(
"spill decode allocation limit {} exceeded: needs {bytes}, remaining {}",
self.limit_bytes, self.remaining_bytes
),
));
};
self.remaining_bytes = remaining;
Ok(())
}
pub fn charge_elements<T>(&mut self, count: usize) -> std::io::Result<()> {
let bytes = count.checked_mul(std::mem::size_of::<T>()).ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::OutOfMemory,
"spill decode allocation limit: element count overflow",
)
})?;
self.charge_bytes(bytes)
}
}
pub trait SpillCodec: Sized {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError>;
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError>;
fn decode_with_budget<R: Read>(
_r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Err(SpillError::Codec(
"bounded spill decoding is not implemented for this codec".into(),
))
}
}
impl SpillCodec for u64 {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
w.write_all(&self.to_le_bytes())?;
Ok(8)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(read_spill_header(r)?.map(u64::from_le_bytes))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
impl SpillCodec for i64 {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
w.write_all(&self.to_le_bytes())?;
Ok(8)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(read_spill_header(r)?.map(i64::from_le_bytes))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
impl SpillCodec for f64 {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
w.write_all(&self.to_le_bytes())?;
Ok(8)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(read_spill_header(r)?.map(f64::from_le_bytes))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
impl SpillCodec for String {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
let bytes = self.as_bytes();
let len = bytes.len() as u32;
w.write_all(&len.to_le_bytes())?;
w.write_all(bytes)?;
Ok(4 + bytes.len())
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Self::decode_with_budget(r, &mut SpillDecodeBudget::new(DEFAULT_SPILL_DECODE_BYTES))
}
fn decode_with_budget<R: Read>(
r: &mut R,
budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
let Some(lenbuf) = read_spill_header(r)? else {
return Ok(None);
};
let len = u32::from_le_bytes(lenbuf) as usize;
budget.charge_bytes(len)?;
let mut buf = vec![0u8; len];
r.read_exact(&mut buf)?;
String::from_utf8(buf)
.map(Some)
.map_err(|e| SpillError::Codec(format!("invalid utf-8: {e}")))
}
}
pub struct SpilledHashAgg<K, S>
where
K: Hash + Eq + Clone + SpillCodec,
S: Clone + Mergeable + SpillCodec,
{
table: HashMap<K, S>,
avg_entry_bytes: usize,
mem_limit_bytes: usize,
spill_dir: PathBuf,
spilled_batches: Vec<SpillBatch>,
next_seq: u64,
pub total_spilled_bytes: u64,
pub spill_count: u64,
}
impl<K, S> SpilledHashAgg<K, S>
where
K: Hash + Eq + Clone + SpillCodec,
S: Clone + Mergeable + SpillCodec,
{
pub fn new(
spill_dir: impl AsRef<Path>,
mem_limit_bytes: usize,
avg_entry_bytes: usize,
) -> Self {
Self {
table: HashMap::new(),
avg_entry_bytes,
mem_limit_bytes,
spill_dir: spill_dir.as_ref().to_path_buf(),
spilled_batches: Vec::new(),
next_seq: 0,
total_spilled_bytes: 0,
spill_count: 0,
}
}
pub fn accumulate(&mut self, key: K, increment: S) -> Result<(), SpillError> {
match self.table.get_mut(&key) {
Some(existing) => existing.merge(increment),
None => {
self.table.insert(key, increment);
if self.should_spill() {
self.spill_partition()?;
}
}
}
Ok(())
}
fn should_spill(&self) -> bool {
if self.mem_limit_bytes == 0 {
return false;
}
let estimated = self.table.len().saturating_mul(self.avg_entry_bytes);
estimated > self.mem_limit_bytes
}
pub fn spill_partition(&mut self) -> Result<(), SpillError> {
if self.table.is_empty() {
return Ok(());
}
let path = self.spill_dir.join(format!("spill_{}.bin", self.next_seq));
self.next_seq += 1;
let file = OpenOptions::new()
.write(true)
.create_new(true)
.open(&path)?;
let mut batch = SpillBatch {
path,
bytes: 0,
entries: self.table.len(),
};
let mut writer = BufWriter::new(file);
let mut bytes_written = 0usize;
for (k, s) in &self.table {
bytes_written = bytes_written
.checked_add(k.encode(&mut writer)?)
.ok_or_else(|| SpillError::Codec("spill byte count overflow".into()))?;
bytes_written = bytes_written
.checked_add(s.encode(&mut writer)?)
.ok_or_else(|| SpillError::Codec("spill byte count overflow".into()))?;
}
writer.flush()?;
batch.bytes = u64::try_from(bytes_written)
.map_err(|_| SpillError::Codec("spill byte count overflow".into()))?;
let total_bytes = self
.total_spilled_bytes
.checked_add(batch.bytes)
.ok_or_else(|| SpillError::Codec("spill byte count overflow".into()))?;
self.table.clear();
self.total_spilled_bytes = total_bytes;
self.spill_count += 1;
self.spilled_batches.push(batch);
Ok(())
}
pub fn drain(mut self) -> Result<HashMap<K, S>, SpillError> {
let mut merged = std::mem::take(&mut self.table);
for batch in self.spilled_batches {
let file = File::open(&batch.path)?;
let metadata = file.metadata()?;
let file_bytes = usize::try_from(metadata.len()).unwrap_or(usize::MAX);
if self.mem_limit_bytes > 0 && file_bytes > self.mem_limit_bytes {
return Err(SpillError::BatchTooLarge {
size: file_bytes,
limit: self.mem_limit_bytes,
});
}
if metadata.len() != batch.bytes {
return Err(SpillError::Codec(
"spill batch length changed after publication".into(),
));
}
let mut reader = BufReader::new(file);
let mut budget = SpillDecodeBudget::new(if self.mem_limit_bytes == 0 {
DEFAULT_SPILL_DECODE_BYTES
} else {
self.mem_limit_bytes
});
for _ in 0..batch.entries {
let key = K::decode_with_budget(&mut reader, &mut budget)?.ok_or_else(|| {
SpillError::Codec("spill batch ended before its published entry count".into())
})?;
let state = match S::decode_with_budget(&mut reader, &mut budget)? {
Some(s) => s,
None => {
return Err(SpillError::Codec(
"spill batch ended mid-entry: state missing".to_string(),
))
}
};
match merged.get_mut(&key) {
Some(existing) => existing.merge(state),
None => {
merged.insert(key, state);
}
}
}
if read_spill_header::<_, 1>(&mut reader)?.is_some() {
return Err(SpillError::Codec("spill batch has trailing entries".into()));
}
}
Ok(merged)
}
pub fn spilled_batch_count(&self) -> usize {
self.spilled_batches.len()
}
pub fn in_memory_groups(&self) -> usize {
self.table.len()
}
}
#[derive(Debug, Clone, Copy)]
pub struct SumState<T>(pub T);
impl Mergeable for SumState<i64> {
fn merge(&mut self, other: Self) {
self.0 = self.0.saturating_add(other.0);
}
}
impl SpillCodec for SumState<i64> {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
self.0.encode(w)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(i64::decode(r)?.map(SumState))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
impl Mergeable for SumState<f64> {
fn merge(&mut self, other: Self) {
self.0 += other.0;
}
}
impl SpillCodec for SumState<f64> {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
self.0.encode(w)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(f64::decode(r)?.map(SumState))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
#[derive(Debug, Clone, Copy)]
pub struct CountState(pub u64);
impl Mergeable for CountState {
fn merge(&mut self, other: Self) {
self.0 = self.0.saturating_add(other.0);
}
}
impl SpillCodec for CountState {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
self.0.encode(w)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(u64::decode(r)?.map(CountState))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
#[derive(Debug, Clone, Copy)]
pub struct MinState<T>(pub T);
#[derive(Debug, Clone, Copy)]
pub struct MaxState<T>(pub T);
impl Mergeable for MinState<i64> {
fn merge(&mut self, other: Self) {
if other.0 < self.0 {
self.0 = other.0;
}
}
}
impl SpillCodec for MinState<i64> {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
self.0.encode(w)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(i64::decode(r)?.map(MinState))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
impl Mergeable for MaxState<i64> {
fn merge(&mut self, other: Self) {
if other.0 > self.0 {
self.0 = other.0;
}
}
}
impl SpillCodec for MaxState<i64> {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
self.0.encode(w)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
Ok(i64::decode(r)?.map(MaxState))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
#[derive(Debug, Clone, Copy)]
pub struct AvgState {
pub sum: f64,
pub count: u64,
}
impl Mergeable for AvgState {
fn merge(&mut self, other: Self) {
self.sum += other.sum;
self.count += other.count;
}
}
impl SpillCodec for AvgState {
fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
let a = self.sum.encode(w)?;
let b = self.count.encode(w)?;
Ok(a + b)
}
fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
let sum = match f64::decode(r)? {
Some(v) => v,
None => return Ok(None),
};
let count = match u64::decode(r)? {
Some(v) => v,
None => {
return Err(SpillError::Codec(
"AvgState ended after sum: count missing".to_string(),
))
}
};
Ok(Some(AvgState { sum, count }))
}
fn decode_with_budget<R: Read>(
r: &mut R,
_budget: &mut SpillDecodeBudget,
) -> Result<Option<Self>, SpillError> {
Self::decode(r)
}
}
impl AvgState {
pub fn finalize(self) -> Option<f64> {
if self.count == 0 {
None
} else {
Some(self.sum / self.count as f64)
}
}
}
pub fn spilled_hash_agg_default<K, S>() -> std::io::Result<SpilledHashAgg<K, S>>
where
K: std::hash::Hash + Eq + Clone + SpillCodec,
S: Clone + Mergeable + SpillCodec,
{
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let seq = SEQ.fetch_add(1, Ordering::Relaxed);
let pid = std::process::id();
let dir = std::env::temp_dir().join(format!("reddb-spill-{pid}-{seq}"));
std::fs::create_dir_all(&dir)?;
Ok(SpilledHashAgg::new(
dir,
64 * 1024 * 1024, 128, ))
}