1use std::collections::HashMap;
40use std::fs::{File, OpenOptions};
41use std::hash::Hash;
42use std::io::{BufReader, BufWriter, Read, Write};
43use std::path::{Path, PathBuf};
44
45pub trait Mergeable {
57 fn merge(&mut self, other: Self);
60}
61
62#[derive(Debug)]
64pub enum SpillError {
65 Io(std::io::Error),
67 BatchTooLarge { size: usize, limit: usize },
71 Codec(String),
74}
75
76impl From<std::io::Error> for SpillError {
77 fn from(e: std::io::Error) -> Self {
78 Self::Io(e)
79 }
80}
81
82impl std::fmt::Display for SpillError {
83 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84 match self {
85 Self::Io(e) => write!(f, "spill i/o: {e}"),
86 Self::BatchTooLarge { size, limit } => {
87 write!(f, "spill batch {size} bytes exceeds limit {limit}")
88 }
89 Self::Codec(msg) => write!(f, "spill codec: {msg}"),
90 }
91 }
92}
93
94impl std::error::Error for SpillError {}
95
96struct SpillBatch {
99 path: PathBuf,
100 bytes: u64,
101 entries: usize,
102}
103
104impl Drop for SpillBatch {
105 fn drop(&mut self) {
106 let _ = std::fs::remove_file(&self.path);
107 }
108}
109
110pub(crate) fn read_spill_header<R: Read, const N: usize>(
113 reader: &mut R,
114) -> std::io::Result<Option<[u8; N]>> {
115 assert!(N > 0, "invariant: spill headers have a positive size");
116 let mut bytes = [0; N];
117 let read = loop {
118 match reader.read(&mut bytes) {
119 Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
120 result => break result?,
121 }
122 };
123 if read == 0 {
124 return Ok(None);
125 }
126 reader.read_exact(&mut bytes[read..])?;
127 Ok(Some(bytes))
128}
129
130pub const DEFAULT_SPILL_DECODE_BYTES: usize = 64 * 1024 * 1024;
133
134pub struct SpillDecodeBudget {
138 limit_bytes: usize,
139 remaining_bytes: usize,
140}
141
142impl SpillDecodeBudget {
143 pub fn new(limit_bytes: usize) -> Self {
144 Self {
145 limit_bytes,
146 remaining_bytes: limit_bytes,
147 }
148 }
149 pub fn remaining_bytes(&self) -> usize {
150 self.remaining_bytes
151 }
152 pub fn charge_bytes(&mut self, bytes: usize) -> std::io::Result<()> {
153 let Some(remaining) = self.remaining_bytes.checked_sub(bytes) else {
154 return Err(std::io::Error::new(
155 std::io::ErrorKind::OutOfMemory,
156 format!(
157 "spill decode allocation limit {} exceeded: needs {bytes}, remaining {}",
158 self.limit_bytes, self.remaining_bytes
159 ),
160 ));
161 };
162 self.remaining_bytes = remaining;
163 Ok(())
164 }
165 pub fn charge_elements<T>(&mut self, count: usize) -> std::io::Result<()> {
166 let bytes = count.checked_mul(std::mem::size_of::<T>()).ok_or_else(|| {
167 std::io::Error::new(
168 std::io::ErrorKind::OutOfMemory,
169 "spill decode allocation limit: element count overflow",
170 )
171 })?;
172 self.charge_bytes(bytes)
173 }
174}
175
176pub trait SpillCodec: Sized {
184 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError>;
187 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError>;
191
192 fn decode_with_budget<R: Read>(
195 _r: &mut R,
196 _budget: &mut SpillDecodeBudget,
197 ) -> Result<Option<Self>, SpillError> {
198 Err(SpillError::Codec(
199 "bounded spill decoding is not implemented for this codec".into(),
200 ))
201 }
202}
203
204impl SpillCodec for u64 {
209 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
210 w.write_all(&self.to_le_bytes())?;
211 Ok(8)
212 }
213 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
214 Ok(read_spill_header(r)?.map(u64::from_le_bytes))
215 }
216 fn decode_with_budget<R: Read>(
217 r: &mut R,
218 _budget: &mut SpillDecodeBudget,
219 ) -> Result<Option<Self>, SpillError> {
220 Self::decode(r)
221 }
222}
223
224impl SpillCodec for i64 {
225 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
226 w.write_all(&self.to_le_bytes())?;
227 Ok(8)
228 }
229 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
230 Ok(read_spill_header(r)?.map(i64::from_le_bytes))
231 }
232 fn decode_with_budget<R: Read>(
233 r: &mut R,
234 _budget: &mut SpillDecodeBudget,
235 ) -> Result<Option<Self>, SpillError> {
236 Self::decode(r)
237 }
238}
239
240impl SpillCodec for f64 {
241 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
242 w.write_all(&self.to_le_bytes())?;
243 Ok(8)
244 }
245 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
246 Ok(read_spill_header(r)?.map(f64::from_le_bytes))
247 }
248 fn decode_with_budget<R: Read>(
249 r: &mut R,
250 _budget: &mut SpillDecodeBudget,
251 ) -> Result<Option<Self>, SpillError> {
252 Self::decode(r)
253 }
254}
255
256impl SpillCodec for String {
257 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
258 let bytes = self.as_bytes();
259 let len = bytes.len() as u32;
260 w.write_all(&len.to_le_bytes())?;
261 w.write_all(bytes)?;
262 Ok(4 + bytes.len())
263 }
264 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
265 Self::decode_with_budget(r, &mut SpillDecodeBudget::new(DEFAULT_SPILL_DECODE_BYTES))
266 }
267 fn decode_with_budget<R: Read>(
268 r: &mut R,
269 budget: &mut SpillDecodeBudget,
270 ) -> Result<Option<Self>, SpillError> {
271 let Some(lenbuf) = read_spill_header(r)? else {
272 return Ok(None);
273 };
274 let len = u32::from_le_bytes(lenbuf) as usize;
275 budget.charge_bytes(len)?;
276 let mut buf = vec![0u8; len];
277 r.read_exact(&mut buf)?;
278 String::from_utf8(buf)
279 .map(Some)
280 .map_err(|e| SpillError::Codec(format!("invalid utf-8: {e}")))
281 }
282}
283
284pub struct SpilledHashAgg<K, S>
291where
292 K: Hash + Eq + Clone + SpillCodec,
293 S: Clone + Mergeable + SpillCodec,
294{
295 table: HashMap<K, S>,
297 avg_entry_bytes: usize,
301 mem_limit_bytes: usize,
304 spill_dir: PathBuf,
307 spilled_batches: Vec<SpillBatch>,
309 next_seq: u64,
311 pub total_spilled_bytes: u64,
313 pub spill_count: u64,
315}
316
317impl<K, S> SpilledHashAgg<K, S>
318where
319 K: Hash + Eq + Clone + SpillCodec,
320 S: Clone + Mergeable + SpillCodec,
321{
322 pub fn new(
327 spill_dir: impl AsRef<Path>,
328 mem_limit_bytes: usize,
329 avg_entry_bytes: usize,
330 ) -> Self {
331 Self {
332 table: HashMap::new(),
333 avg_entry_bytes,
334 mem_limit_bytes,
335 spill_dir: spill_dir.as_ref().to_path_buf(),
336 spilled_batches: Vec::new(),
337 next_seq: 0,
338 total_spilled_bytes: 0,
339 spill_count: 0,
340 }
341 }
342
343 pub fn accumulate(&mut self, key: K, increment: S) -> Result<(), SpillError> {
348 match self.table.get_mut(&key) {
349 Some(existing) => existing.merge(increment),
350 None => {
351 self.table.insert(key, increment);
352 if self.should_spill() {
353 self.spill_partition()?;
354 }
355 }
356 }
357 Ok(())
358 }
359
360 fn should_spill(&self) -> bool {
365 if self.mem_limit_bytes == 0 {
366 return false;
367 }
368 let estimated = self.table.len().saturating_mul(self.avg_entry_bytes);
369 estimated > self.mem_limit_bytes
370 }
371
372 pub fn spill_partition(&mut self) -> Result<(), SpillError> {
377 if self.table.is_empty() {
378 return Ok(());
379 }
380 let path = self.spill_dir.join(format!("spill_{}.bin", self.next_seq));
381 self.next_seq += 1;
382 let file = OpenOptions::new()
383 .write(true)
384 .create_new(true)
385 .open(&path)?;
386 let mut batch = SpillBatch {
387 path,
388 bytes: 0,
389 entries: self.table.len(),
390 };
391 let mut writer = BufWriter::new(file);
392 let mut bytes_written = 0usize;
393 for (k, s) in &self.table {
396 bytes_written = bytes_written
397 .checked_add(k.encode(&mut writer)?)
398 .ok_or_else(|| SpillError::Codec("spill byte count overflow".into()))?;
399 bytes_written = bytes_written
400 .checked_add(s.encode(&mut writer)?)
401 .ok_or_else(|| SpillError::Codec("spill byte count overflow".into()))?;
402 }
403 writer.flush()?;
404 batch.bytes = u64::try_from(bytes_written)
405 .map_err(|_| SpillError::Codec("spill byte count overflow".into()))?;
406 let total_bytes = self
407 .total_spilled_bytes
408 .checked_add(batch.bytes)
409 .ok_or_else(|| SpillError::Codec("spill byte count overflow".into()))?;
410 self.table.clear();
411 self.total_spilled_bytes = total_bytes;
412 self.spill_count += 1;
413 self.spilled_batches.push(batch);
414 Ok(())
415 }
416
417 pub fn drain(mut self) -> Result<HashMap<K, S>, SpillError> {
427 let mut merged = std::mem::take(&mut self.table);
430 for batch in self.spilled_batches {
431 let file = File::open(&batch.path)?;
432 let metadata = file.metadata()?;
433 let file_bytes = usize::try_from(metadata.len()).unwrap_or(usize::MAX);
434 if self.mem_limit_bytes > 0 && file_bytes > self.mem_limit_bytes {
435 return Err(SpillError::BatchTooLarge {
436 size: file_bytes,
437 limit: self.mem_limit_bytes,
438 });
439 }
440 if metadata.len() != batch.bytes {
441 return Err(SpillError::Codec(
442 "spill batch length changed after publication".into(),
443 ));
444 }
445 let mut reader = BufReader::new(file);
446 let mut budget = SpillDecodeBudget::new(if self.mem_limit_bytes == 0 {
447 DEFAULT_SPILL_DECODE_BYTES
448 } else {
449 self.mem_limit_bytes
450 });
451 for _ in 0..batch.entries {
452 let key = K::decode_with_budget(&mut reader, &mut budget)?.ok_or_else(|| {
453 SpillError::Codec("spill batch ended before its published entry count".into())
454 })?;
455 let state = match S::decode_with_budget(&mut reader, &mut budget)? {
456 Some(s) => s,
457 None => {
458 return Err(SpillError::Codec(
459 "spill batch ended mid-entry: state missing".to_string(),
460 ))
461 }
462 };
463 match merged.get_mut(&key) {
464 Some(existing) => existing.merge(state),
465 None => {
466 merged.insert(key, state);
467 }
468 }
469 }
470 if read_spill_header::<_, 1>(&mut reader)?.is_some() {
471 return Err(SpillError::Codec("spill batch has trailing entries".into()));
472 }
473 }
474 Ok(merged)
475 }
476
477 pub fn spilled_batch_count(&self) -> usize {
480 self.spilled_batches.len()
481 }
482
483 pub fn in_memory_groups(&self) -> usize {
485 self.table.len()
486 }
487}
488
489#[derive(Debug, Clone, Copy)]
499pub struct SumState<T>(pub T);
500
501impl Mergeable for SumState<i64> {
502 fn merge(&mut self, other: Self) {
503 self.0 = self.0.saturating_add(other.0);
504 }
505}
506impl SpillCodec for SumState<i64> {
507 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
508 self.0.encode(w)
509 }
510 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
511 Ok(i64::decode(r)?.map(SumState))
512 }
513 fn decode_with_budget<R: Read>(
514 r: &mut R,
515 _budget: &mut SpillDecodeBudget,
516 ) -> Result<Option<Self>, SpillError> {
517 Self::decode(r)
518 }
519}
520
521impl Mergeable for SumState<f64> {
522 fn merge(&mut self, other: Self) {
523 self.0 += other.0;
524 }
525}
526impl SpillCodec for SumState<f64> {
527 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
528 self.0.encode(w)
529 }
530 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
531 Ok(f64::decode(r)?.map(SumState))
532 }
533 fn decode_with_budget<R: Read>(
534 r: &mut R,
535 _budget: &mut SpillDecodeBudget,
536 ) -> Result<Option<Self>, SpillError> {
537 Self::decode(r)
538 }
539}
540
541#[derive(Debug, Clone, Copy)]
543pub struct CountState(pub u64);
544
545impl Mergeable for CountState {
546 fn merge(&mut self, other: Self) {
547 self.0 = self.0.saturating_add(other.0);
548 }
549}
550impl SpillCodec for CountState {
551 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
552 self.0.encode(w)
553 }
554 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
555 Ok(u64::decode(r)?.map(CountState))
556 }
557 fn decode_with_budget<R: Read>(
558 r: &mut R,
559 _budget: &mut SpillDecodeBudget,
560 ) -> Result<Option<Self>, SpillError> {
561 Self::decode(r)
562 }
563}
564
565#[derive(Debug, Clone, Copy)]
568pub struct MinState<T>(pub T);
569#[derive(Debug, Clone, Copy)]
570pub struct MaxState<T>(pub T);
571
572impl Mergeable for MinState<i64> {
573 fn merge(&mut self, other: Self) {
574 if other.0 < self.0 {
575 self.0 = other.0;
576 }
577 }
578}
579impl SpillCodec for MinState<i64> {
580 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
581 self.0.encode(w)
582 }
583 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
584 Ok(i64::decode(r)?.map(MinState))
585 }
586 fn decode_with_budget<R: Read>(
587 r: &mut R,
588 _budget: &mut SpillDecodeBudget,
589 ) -> Result<Option<Self>, SpillError> {
590 Self::decode(r)
591 }
592}
593
594impl Mergeable for MaxState<i64> {
595 fn merge(&mut self, other: Self) {
596 if other.0 > self.0 {
597 self.0 = other.0;
598 }
599 }
600}
601impl SpillCodec for MaxState<i64> {
602 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
603 self.0.encode(w)
604 }
605 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
606 Ok(i64::decode(r)?.map(MaxState))
607 }
608 fn decode_with_budget<R: Read>(
609 r: &mut R,
610 _budget: &mut SpillDecodeBudget,
611 ) -> Result<Option<Self>, SpillError> {
612 Self::decode(r)
613 }
614}
615
616#[derive(Debug, Clone, Copy)]
618pub struct AvgState {
619 pub sum: f64,
620 pub count: u64,
621}
622
623impl Mergeable for AvgState {
624 fn merge(&mut self, other: Self) {
625 self.sum += other.sum;
626 self.count += other.count;
627 }
628}
629impl SpillCodec for AvgState {
630 fn encode<W: Write>(&self, w: &mut W) -> Result<usize, SpillError> {
631 let a = self.sum.encode(w)?;
632 let b = self.count.encode(w)?;
633 Ok(a + b)
634 }
635 fn decode<R: Read>(r: &mut R) -> Result<Option<Self>, SpillError> {
636 let sum = match f64::decode(r)? {
637 Some(v) => v,
638 None => return Ok(None),
639 };
640 let count = match u64::decode(r)? {
641 Some(v) => v,
642 None => {
643 return Err(SpillError::Codec(
644 "AvgState ended after sum: count missing".to_string(),
645 ))
646 }
647 };
648 Ok(Some(AvgState { sum, count }))
649 }
650 fn decode_with_budget<R: Read>(
651 r: &mut R,
652 _budget: &mut SpillDecodeBudget,
653 ) -> Result<Option<Self>, SpillError> {
654 Self::decode(r)
655 }
656}
657
658impl AvgState {
659 pub fn finalize(self) -> Option<f64> {
662 if self.count == 0 {
663 None
664 } else {
665 Some(self.sum / self.count as f64)
666 }
667 }
668}
669
670pub fn spilled_hash_agg_default<K, S>() -> std::io::Result<SpilledHashAgg<K, S>>
685where
686 K: std::hash::Hash + Eq + Clone + SpillCodec,
687 S: Clone + Mergeable + SpillCodec,
688{
689 use std::sync::atomic::{AtomicU64, Ordering};
690 static SEQ: AtomicU64 = AtomicU64::new(0);
691 let seq = SEQ.fetch_add(1, Ordering::Relaxed);
692 let pid = std::process::id();
693 let dir = std::env::temp_dir().join(format!("reddb-spill-{pid}-{seq}"));
694 std::fs::create_dir_all(&dir)?;
695 Ok(SpilledHashAgg::new(
696 dir,
697 64 * 1024 * 1024, 128, ))
700}