Skip to main content

uqa_execution/
spill.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Disk-backed spill buffer for blocking operators (`Sort`,
8//! `HashAggregate`, `Window`).
9//!
10//! The budget is measured in the exact number of bytes each batch occupies in
11//! the spill encoding. [`SpillBuffer::push`] automatically flushes before a
12//! successful push could leave more encoded bytes in memory than the budget.
13//! Draining restores spilled batches first and then any in-memory tail,
14//! preserving input order. The temporary file is removed when the buffer (or
15//! its active drain iterator) is dropped.
16
17use std::fs::File;
18use std::io::{BufReader, Write};
19use std::path::{Path, PathBuf};
20use std::sync::Arc;
21
22use crate::batch::{Batch, OwnedPhysicalRow, PhysicalRow, RowSchema};
23use crate::physical::ExecResult;
24use tempfile::NamedTempFile;
25
26mod format;
27mod indexed;
28
29use format::{
30    append_batches, decode_batch, encoded_batch_overhead_size, encoded_batch_size,
31    encoded_physical_row_record_size, open_spill_reader, read_bounded_spill_record, spill_error,
32};
33pub use indexed::IndexedSpill;
34
35const SPILL_MAGIC: &[u8] = b"UQA-SPILL\x01\n";
36
37/// Incremental exact size of one not-yet-encoded batch. When the first row with lock origins arrives, the binary format adds an empty origin-count field to every preceding origin-free row, so accounting must update those already-buffered rows as well as the new record.
38#[derive(Clone, Copy)]
39pub(crate) struct EncodedBatchSizer {
40    physical_width: usize,
41    bytes: usize,
42    origin_free_rows: usize,
43    has_lock_origins: bool,
44}
45
46impl EncodedBatchSizer {
47    pub(crate) fn new(schema: &RowSchema) -> ExecResult<Self> {
48        Ok(Self {
49            physical_width: schema.physical_width(),
50            bytes: encoded_batch_overhead_size(schema)?,
51            origin_free_rows: 0,
52            has_lock_origins: false,
53        })
54    }
55
56    pub(crate) fn append(&mut self, row: &PhysicalRow) -> ExecResult<()> {
57        let mut additional = encoded_physical_row_record_size(row, self.physical_width)?;
58        if row.lock_origins().is_empty() {
59            self.origin_free_rows = self.origin_free_rows.checked_add(1).ok_or_else(|| {
60                spill_error("incremental spill batch origin-free row count overflow")
61            })?;
62            if self.has_lock_origins {
63                additional = additional.checked_add(8).ok_or_else(|| {
64                    spill_error("incremental spill batch lock-origin size overflow")
65                })?;
66            }
67        } else if !self.has_lock_origins {
68            let preceding_metadata = self.origin_free_rows.checked_mul(8).ok_or_else(|| {
69                spill_error("incremental spill batch lock-origin metadata overflow")
70            })?;
71            additional = additional
72                .checked_add(preceding_metadata)
73                .ok_or_else(|| spill_error("incremental spill batch lock-origin size overflow"))?;
74            self.has_lock_origins = true;
75        }
76        self.bytes = self
77            .bytes
78            .checked_add(additional)
79            .ok_or_else(|| spill_error("incremental spill batch size overflow"))?;
80        Ok(())
81    }
82
83    pub(crate) fn bytes(self) -> usize {
84        self.bytes
85    }
86}
87
88/// Append-only batch buffer with an encoded-byte memory budget.
89///
90/// The budget is exact for the serialized representation and does not claim to
91/// be the Rust allocator's resident-byte accounting. At most one incoming or
92/// decoded batch can itself be larger than the budget; successful pushes do not
93/// retain such an oversized batch in memory.
94pub struct SpillBuffer {
95    schema: Option<RowSchema>,
96    batches: Vec<Batch>,
97    rows: usize,
98    in_memory_rows: usize,
99    in_memory_bytes: usize,
100    max_in_memory_record_bytes: usize,
101    /// Encoded-byte budget. Set to `usize::MAX` to disable spilling.
102    budget_bytes: usize,
103    spill_directory: Option<PathBuf>,
104    spill_file: Option<NamedTempFile>,
105    spilled_batches: usize,
106    spilled_rows: usize,
107    spilled_bytes: usize,
108    max_spilled_record_bytes: usize,
109}
110
111impl SpillBuffer {
112    pub fn new(budget_bytes: usize) -> Self {
113        Self {
114            schema: None,
115            batches: Vec::new(),
116            rows: 0,
117            in_memory_rows: 0,
118            in_memory_bytes: 0,
119            max_in_memory_record_bytes: 0,
120            budget_bytes,
121            spill_directory: None,
122            spill_file: None,
123            spilled_batches: 0,
124            spilled_rows: 0,
125            spilled_bytes: 0,
126            max_spilled_record_bytes: 0,
127        }
128    }
129
130    /// Create a buffer whose temporary spill file will be placed in `directory`.
131    ///
132    /// File creation is deferred until the first spill. This is primarily useful
133    /// when an engine has a dedicated temporary-data volume.
134    pub fn new_in(budget_bytes: usize, directory: impl Into<PathBuf>) -> Self {
135        let mut buffer = Self::new(budget_bytes);
136        buffer.spill_directory = Some(directory.into());
137        buffer
138    }
139
140    pub fn unbounded() -> Self {
141        Self::new(usize::MAX)
142    }
143
144    /// Append a batch, spilling automatically when required by the byte budget.
145    ///
146    /// Returns `true` if this push wrote one or more batches to disk. If disk
147    /// creation, encoding, or writing fails, the new batch and all earlier
148    /// batches remain owned by the buffer and the error is returned.
149    pub fn push(&mut self, batch: Batch) -> ExecResult<bool> {
150        if let Some(schema) = self.schema.as_ref() {
151            if schema != &batch.schema {
152                return Err(spill_error(format!(
153                    "spill buffer schema mismatch: expected {:?}, got {:?}",
154                    schema.columns(),
155                    batch.schema.columns()
156                )));
157            }
158        } else {
159            self.schema = Some(batch.schema.clone());
160        }
161        let batch_rows = batch.rows.len();
162        let next_rows = self
163            .rows
164            .checked_add(batch_rows)
165            .ok_or_else(|| spill_error("spill buffer row count overflow"))?;
166        let batch_bytes = match Self::encoded_size(&batch) {
167            Ok(bytes) => bytes,
168            Err(error) => {
169                // Preserve ownership even when an exotic value or an encoded
170                // size overflow prevents budget accounting. The failed
171                // operator will abort, but no row silently disappears.
172                self.retain_batch(batch, usize::MAX);
173                return Err(error);
174            }
175        };
176        let would_exceed = self
177            .in_memory_bytes
178            .checked_add(batch_bytes)
179            .is_none_or(|bytes| bytes > self.budget_bytes);
180
181        let mut spilled = false;
182        if would_exceed && !self.batches.is_empty() {
183            if let Err(error) = self.spill_pending() {
184                self.retain_batch(batch, batch_bytes);
185                return Err(error);
186            }
187            spilled = true;
188        }
189
190        let next_in_memory_rows = self
191            .in_memory_rows
192            .checked_add(batch_rows)
193            .ok_or_else(|| spill_error("spill buffer in-memory row count overflow"))?;
194        let next_in_memory_bytes = self
195            .in_memory_bytes
196            .checked_add(batch_bytes)
197            .ok_or_else(|| spill_error("spill buffer in-memory byte count overflow"))?;
198        self.rows = next_rows;
199        self.in_memory_rows = next_in_memory_rows;
200        self.in_memory_bytes = next_in_memory_bytes;
201        self.max_in_memory_record_bytes = self.max_in_memory_record_bytes.max(batch_bytes);
202        self.batches.push(batch);
203
204        // A single encoded batch may exceed work_mem. It must pass through
205        // memory once, but a successful push never retains it there.
206        if self.in_memory_bytes > self.budget_bytes {
207            self.spill_pending()?;
208            spilled = true;
209        }
210        Ok(spilled)
211    }
212
213    /// Exact byte count used for budget accounting, including the record
214    /// length prefix written to disk.
215    pub fn encoded_size(batch: &Batch) -> ExecResult<usize> {
216        encoded_batch_size(batch)
217    }
218
219    /// Total buffered rows, including rows already written to disk.
220    pub fn rows(&self) -> usize {
221        self.rows
222    }
223
224    /// Rows currently retained in memory.
225    pub fn in_memory_rows(&self) -> usize {
226        self.in_memory_rows
227    }
228
229    /// Exact encoded bytes currently retained in memory.
230    pub fn in_memory_bytes(&self) -> usize {
231        self.in_memory_bytes
232    }
233
234    pub fn budget_bytes(&self) -> usize {
235        self.budget_bytes
236    }
237
238    pub fn over_budget(&self) -> bool {
239        self.in_memory_bytes > self.budget_bytes
240    }
241
242    pub fn has_spilled(&self) -> bool {
243        self.spill_file.is_some()
244    }
245
246    pub fn spilled_rows(&self) -> usize {
247        self.spilled_rows
248    }
249
250    pub fn spilled_batches(&self) -> usize {
251        self.spilled_batches
252    }
253
254    pub fn spilled_bytes(&self) -> usize {
255        self.spilled_bytes
256    }
257
258    /// Path of the live spill file, if one has been created.
259    ///
260    /// The path is diagnostic only and becomes invalid as soon as the buffer or
261    /// the drain iterator that owns the file is dropped.
262    pub fn spill_path(&self) -> Option<&Path> {
263        self.spill_file.as_ref().map(NamedTempFile::path)
264    }
265
266    /// Flush all pending in-memory batches when the byte budget is exceeded.
267    ///
268    /// Returns `true` when batches were written. A failed append is rolled back
269    /// to the previous file length and the pending batches remain in memory, so
270    /// callers never observe a silent partial spill.
271    pub fn spill_if_over_budget(&mut self) -> ExecResult<bool> {
272        if !self.over_budget() || self.batches.is_empty() {
273            return Ok(false);
274        }
275        self.spill_pending()
276    }
277
278    /// Force all pending in-memory batches to disk regardless of the budget.
279    ///
280    /// This is useful at a blocking-operator phase boundary. It returns `false`
281    /// when there is nothing pending.
282    pub fn spill_pending(&mut self) -> ExecResult<bool> {
283        if self.batches.is_empty() {
284            return Ok(false);
285        }
286
287        // Reject metadata overflow before writing. An error after a successful
288        // append would leave a retry able to duplicate records while the
289        // published counters disagreed with the file.
290        let next_spilled_batches = self
291            .spilled_batches
292            .checked_add(self.batches.len())
293            .ok_or_else(|| spill_error("spill batch count overflow"))?;
294        let next_spilled_rows = self
295            .spilled_rows
296            .checked_add(self.in_memory_rows)
297            .ok_or_else(|| spill_error("spill row count overflow"))?;
298        let next_spilled_bytes = self
299            .spilled_bytes
300            .checked_add(self.in_memory_bytes)
301            .ok_or_else(|| spill_error("spill byte count overflow"))?;
302        let next_max_spilled_record_bytes = self
303            .max_spilled_record_bytes
304            .max(self.max_in_memory_record_bytes);
305
306        if let Some(file) = self.spill_file.as_mut() {
307            append_batches(file.as_file_mut(), &self.batches)?;
308        } else {
309            let mut file = self.create_spill_file()?;
310            append_batches(file.as_file_mut(), &self.batches)?;
311            self.spill_file = Some(file);
312        }
313
314        self.spilled_batches = next_spilled_batches;
315        self.spilled_rows = next_spilled_rows;
316        self.spilled_bytes = next_spilled_bytes;
317        self.max_spilled_record_bytes = next_max_spilled_record_bytes;
318        self.batches.clear();
319        self.in_memory_rows = 0;
320        self.in_memory_bytes = 0;
321        self.max_in_memory_record_bytes = 0;
322        Ok(true)
323    }
324
325    /// Open a repeatable streaming reader without consuming this buffer.
326    ///
327    /// Spilled batches are decoded one at a time. The in-memory tail is cloned
328    /// one batch at a time only when the reader reaches it.
329    pub fn reader(&self) -> ExecResult<SpillReader<'_>> {
330        let reader = self
331            .spill_file
332            .as_ref()
333            .map(open_spill_reader)
334            .transpose()?;
335        let disk_finished = reader.is_none();
336        Ok(SpillReader {
337            reader,
338            memory: self.batches.iter(),
339            disk_finished,
340            failed: false,
341            max_record_bytes: self.max_spilled_record_bytes,
342            expected_schema: self.schema.clone(),
343        })
344    }
345
346    /// Open a repeatable physical-row stream without collecting all batches.
347    pub fn read_rows(&self) -> ExecResult<SpillRows<SpillReader<'_>>> {
348        self.reader().map(SpillRows::new)
349    }
350
351    /// Drain buffered batches in their original input order.
352    ///
353    /// The returned iterator owns the temporary file. Each disk read or decode
354    /// failure is returned as a [`crate::physical::ExecError`], and dropping the
355    /// iterator early still removes the temporary file.
356    pub fn drain(&mut self) -> ExecResult<SpillDrain> {
357        let reader = self
358            .spill_file
359            .as_ref()
360            .map(open_spill_reader)
361            .transpose()?;
362        let spill_file = self.spill_file.take();
363        let memory = std::mem::take(&mut self.batches).into_iter();
364        let expected_schema = self.schema.take();
365
366        self.rows = 0;
367        self.in_memory_rows = 0;
368        self.in_memory_bytes = 0;
369        self.max_in_memory_record_bytes = 0;
370        self.spilled_batches = 0;
371        self.spilled_rows = 0;
372        self.spilled_bytes = 0;
373        let max_record_bytes = std::mem::take(&mut self.max_spilled_record_bytes);
374
375        let disk_finished = reader.is_none();
376        Ok(SpillDrain {
377            reader,
378            spill_file,
379            memory,
380            disk_finished,
381            failed: false,
382            max_record_bytes,
383            expected_schema,
384        })
385    }
386
387    /// Drain and materialize every restored batch.
388    pub fn drain_all(&mut self) -> ExecResult<Vec<Batch>> {
389        self.drain()?.collect()
390    }
391
392    /// Consume the buffer as a physical-row stream without collecting batches.
393    pub fn drain_rows(&mut self) -> ExecResult<SpillRows<SpillDrain>> {
394        self.drain().map(SpillRows::new)
395    }
396
397    /// Discard all buffered data and remove any spill file.
398    pub fn clear(&mut self) {
399        self.schema = None;
400        self.batches.clear();
401        self.spill_file = None;
402        self.rows = 0;
403        self.in_memory_rows = 0;
404        self.in_memory_bytes = 0;
405        self.max_in_memory_record_bytes = 0;
406        self.spilled_batches = 0;
407        self.spilled_rows = 0;
408        self.spilled_bytes = 0;
409        self.max_spilled_record_bytes = 0;
410    }
411
412    /// Seal this buffer as an immutable, cheaply cloneable materialization.
413    /// Batches that fit within the configured byte budget remain in memory;
414    /// once spilling has started, every pending batch is flushed and readers
415    /// reopen the file independently. Both forms support repeatable scans
416    /// without collecting the complete input again.
417    pub fn into_shared(mut self, schema: impl Into<RowSchema>) -> ExecResult<SharedSpill> {
418        let schema = schema.into();
419        if let Some(actual) = self.schema.as_ref() {
420            if actual != &schema {
421                return Err(spill_error(format!(
422                    "shared spill schema mismatch: expected {:?}, got {:?}",
423                    schema.columns(),
424                    actual.columns()
425                )));
426            }
427        }
428        let rows = self.rows;
429        let storage = if self.spill_file.is_none() {
430            let batches = std::mem::take(&mut self.batches);
431            SharedSpillStorage::Memory(batches)
432        } else {
433            self.spill_pending()?;
434            SharedSpillStorage::Disk(
435                self.spill_file
436                    .take()
437                    .expect("spill file exists after flushing shared materialization"),
438            )
439        };
440        Ok(SharedSpill {
441            inner: Arc::new(SharedSpillInner {
442                storage,
443                schema,
444                rows,
445                max_record_bytes: self.max_spilled_record_bytes,
446            }),
447        })
448    }
449
450    fn create_spill_file(&self) -> ExecResult<NamedTempFile> {
451        let mut file = match &self.spill_directory {
452            Some(directory) => NamedTempFile::new_in(directory).map_err(|error| {
453                spill_error(format!(
454                    "failed to create spill file in {}: {error}",
455                    directory.display()
456                ))
457            })?,
458            None => NamedTempFile::new()
459                .map_err(|error| spill_error(format!("failed to create spill file: {error}")))?,
460        };
461        file.as_file_mut()
462            .write_all(SPILL_MAGIC)
463            .map_err(|error| spill_error(format!("failed to initialize spill file: {error}")))?;
464        file.as_file_mut()
465            .flush()
466            .map_err(|error| spill_error(format!("failed to flush spill header: {error}")))?;
467        Ok(file)
468    }
469
470    fn retain_batch(&mut self, batch: Batch, encoded_bytes: usize) {
471        self.rows = self.rows.saturating_add(batch.rows.len());
472        self.in_memory_rows = self.in_memory_rows.saturating_add(batch.rows.len());
473        self.in_memory_bytes = self.in_memory_bytes.saturating_add(encoded_bytes);
474        self.max_in_memory_record_bytes = self.max_in_memory_record_bytes.max(encoded_bytes);
475        self.batches.push(batch);
476    }
477}
478
479enum SharedSpillStorage {
480    Memory(Vec<Batch>),
481    Disk(NamedTempFile),
482}
483
484struct SharedSpillInner {
485    storage: SharedSpillStorage,
486    schema: RowSchema,
487    rows: usize,
488    max_record_bytes: usize,
489}
490
491/// Immutable repeatable row materialization bounded by the source buffer's
492/// memory budget and backed by a temporary file after that budget is exceeded.
493#[derive(Clone)]
494pub struct SharedSpill {
495    inner: Arc<SharedSpillInner>,
496}
497
498impl SharedSpill {
499    pub fn schema(&self) -> &[String] {
500        self.inner.schema.columns()
501    }
502
503    pub fn row_schema(&self) -> &RowSchema {
504        &self.inner.schema
505    }
506
507    pub fn rows(&self) -> usize {
508        self.inner.rows
509    }
510
511    /// Whether this materialization crossed its memory budget and uses disk.
512    pub fn has_spilled(&self) -> bool {
513        matches!(self.inner.storage, SharedSpillStorage::Disk(_))
514    }
515
516    pub fn reader(&self) -> ExecResult<SharedSpillReader> {
517        let source = Arc::clone(&self.inner);
518        Self::reader_from_source(source)
519    }
520
521    /// Consume this materialization into a one-shot reader.
522    ///
523    /// When the in-memory materialization has no other owners, batches move
524    /// directly into the reader instead of being deep-cloned. Shared and disk
525    /// materializations retain the independent-reader behavior of [`Self::reader`].
526    pub fn into_reader(self) -> ExecResult<SharedSpillReader> {
527        match Arc::try_unwrap(self.inner) {
528            Ok(SharedSpillInner {
529                storage: SharedSpillStorage::Memory(batches),
530                schema,
531                max_record_bytes,
532                ..
533            }) => Ok(SharedSpillReader {
534                reader: SharedSpillReaderSource::OwnedMemory(batches.into_iter()),
535                source: None,
536                failed: false,
537                max_record_bytes,
538                expected_schema: Some(schema),
539            }),
540            Ok(inner) => Self::reader_from_source(Arc::new(inner)),
541            Err(source) => Self::reader_from_source(source),
542        }
543    }
544
545    fn reader_from_source(source: Arc<SharedSpillInner>) -> ExecResult<SharedSpillReader> {
546        let reader = match &source.storage {
547            SharedSpillStorage::Memory(_) => SharedSpillReaderSource::Memory { next_batch: 0 },
548            SharedSpillStorage::Disk(file) => {
549                SharedSpillReaderSource::Disk(open_spill_reader(file)?)
550            }
551        };
552        let max_record_bytes = source.max_record_bytes;
553        let expected_schema = Some(source.schema.clone());
554        Ok(SharedSpillReader {
555            reader,
556            source: Some(source),
557            failed: false,
558            max_record_bytes,
559            expected_schema,
560        })
561    }
562
563    /// Open an independent physical-row reader without collecting the spill's
564    /// batches or row count in memory.
565    pub fn read_rows(&self) -> ExecResult<SpillRows<SharedSpillReader>> {
566        self.reader().map(SpillRows::new)
567    }
568}
569
570enum SharedSpillReaderSource {
571    Memory { next_batch: usize },
572    OwnedMemory(std::vec::IntoIter<Batch>),
573    Disk(BufReader<File>),
574}
575
576fn validate_decoded_schema(batch: Batch, expected_schema: Option<&RowSchema>) -> ExecResult<Batch> {
577    if expected_schema.is_none_or(|expected| expected == &batch.schema) {
578        return Ok(batch);
579    }
580    let expected = expected_schema.expect("schema presence checked above");
581    Err(spill_error(format!(
582        "spill batch schema mismatch: expected {:?}, got {:?}",
583        expected.columns(),
584        batch.schema.columns()
585    )))
586}
587
588/// Reader for a [`SharedSpill`]. Independent readers retain the shared source;
589/// a consuming reader may instead own unique in-memory batches directly.
590pub struct SharedSpillReader {
591    reader: SharedSpillReaderSource,
592    source: Option<Arc<SharedSpillInner>>,
593    failed: bool,
594    max_record_bytes: usize,
595    expected_schema: Option<RowSchema>,
596}
597
598impl Iterator for SharedSpillReader {
599    type Item = ExecResult<Batch>;
600
601    fn next(&mut self) -> Option<Self::Item> {
602        if self.failed {
603            return None;
604        }
605        match &mut self.reader {
606            SharedSpillReaderSource::Memory { next_batch } => {
607                let source = self
608                    .source
609                    .as_ref()
610                    .expect("shared memory reader retains its source");
611                let SharedSpillStorage::Memory(batches) = &source.storage else {
612                    unreachable!("shared materialization reader/storage mismatch")
613                };
614                let batch = batches.get(*next_batch)?.clone();
615                *next_batch += 1;
616                Some(validate_decoded_schema(
617                    batch,
618                    self.expected_schema.as_ref(),
619                ))
620            }
621            SharedSpillReaderSource::OwnedMemory(batches) => batches
622                .next()
623                .map(|batch| validate_decoded_schema(batch, self.expected_schema.as_ref())),
624            SharedSpillReaderSource::Disk(reader) => {
625                match read_bounded_spill_record(reader, self.max_record_bytes, "shared spill batch")
626                {
627                    Ok(None) => None,
628                    Ok(Some(record)) => {
629                        let decoded = decode_batch(&record).and_then(|batch| {
630                            validate_decoded_schema(batch, self.expected_schema.as_ref())
631                        });
632                        if decoded.is_err() {
633                            self.failed = true;
634                        }
635                        Some(decoded)
636                    }
637                    Err(error) => {
638                        self.failed = true;
639                        Some(Err(spill_error(format!(
640                            "failed to read shared spill batch: {error}"
641                        ))))
642                    }
643                }
644            }
645        }
646    }
647}
648
649/// Restoring iterator returned by [`SpillBuffer::drain`].
650pub struct SpillDrain {
651    reader: Option<BufReader<File>>,
652    // Keep the named file alive until disk iteration finishes or the iterator
653    // is dropped. Its Drop implementation unlinks the temporary file.
654    spill_file: Option<NamedTempFile>,
655    memory: std::vec::IntoIter<Batch>,
656    disk_finished: bool,
657    failed: bool,
658    max_record_bytes: usize,
659    expected_schema: Option<RowSchema>,
660}
661
662impl Iterator for SpillDrain {
663    type Item = ExecResult<Batch>;
664
665    fn next(&mut self) -> Option<Self::Item> {
666        if self.failed {
667            return None;
668        }
669
670        if !self.disk_finished {
671            let Some(reader) = self.reader.as_mut() else {
672                self.failed = true;
673                self.disk_finished = true;
674                return Some(Err(spill_error(
675                    "spill drain entered disk phase without a reader",
676                )));
677            };
678            match read_bounded_spill_record(reader, self.max_record_bytes, "spill batch") {
679                Ok(None) => {
680                    self.disk_finished = true;
681                    self.reader = None;
682                    self.spill_file = None;
683                }
684                Ok(Some(record)) => {
685                    let decoded = decode_batch(&record).and_then(|batch| {
686                        validate_decoded_schema(batch, self.expected_schema.as_ref())
687                    });
688                    if decoded.is_err() {
689                        self.failed = true;
690                    }
691                    return Some(decoded);
692                }
693                Err(error) => {
694                    self.failed = true;
695                    return Some(Err(spill_error(format!(
696                        "failed to read spill batch: {error}"
697                    ))));
698                }
699            }
700        }
701
702        self.memory
703            .next()
704            .map(|batch| validate_decoded_schema(batch, self.expected_schema.as_ref()))
705    }
706
707    fn size_hint(&self) -> (usize, Option<usize>) {
708        let lower = if self.disk_finished {
709            self.memory.len()
710        } else {
711            0
712        };
713        (lower, None)
714    }
715}
716
717/// Repeatable, non-consuming batch reader returned by [`SpillBuffer::reader`].
718pub struct SpillReader<'a> {
719    reader: Option<BufReader<File>>,
720    memory: std::slice::Iter<'a, Batch>,
721    disk_finished: bool,
722    failed: bool,
723    max_record_bytes: usize,
724    expected_schema: Option<RowSchema>,
725}
726
727impl Iterator for SpillReader<'_> {
728    type Item = ExecResult<Batch>;
729
730    fn next(&mut self) -> Option<Self::Item> {
731        if self.failed {
732            return None;
733        }
734
735        if !self.disk_finished {
736            let Some(reader) = self.reader.as_mut() else {
737                self.failed = true;
738                self.disk_finished = true;
739                return Some(Err(spill_error(
740                    "spill reader entered disk phase without a file reader",
741                )));
742            };
743            match read_bounded_spill_record(reader, self.max_record_bytes, "spill batch") {
744                Ok(None) => {
745                    self.disk_finished = true;
746                    self.reader = None;
747                }
748                Ok(Some(record)) => {
749                    let decoded = decode_batch(&record).and_then(|batch| {
750                        validate_decoded_schema(batch, self.expected_schema.as_ref())
751                    });
752                    if decoded.is_err() {
753                        self.failed = true;
754                    }
755                    return Some(decoded);
756                }
757                Err(error) => {
758                    self.failed = true;
759                    return Some(Err(spill_error(format!(
760                        "failed to read spill batch: {error}"
761                    ))));
762                }
763            }
764        }
765
766        self.memory
767            .next()
768            .cloned()
769            .map(|batch| validate_decoded_schema(batch, self.expected_schema.as_ref()))
770    }
771
772    fn size_hint(&self) -> (usize, Option<usize>) {
773        let lower = if self.disk_finished {
774            self.memory.len()
775        } else {
776            0
777        };
778        (lower, None)
779    }
780}
781
782/// Physical-row flattening adapter for [`SpillReader`] and [`SpillDrain`].
783pub struct SpillRows<I> {
784    batches: I,
785    current_schema: Option<RowSchema>,
786    current: std::vec::IntoIter<PhysicalRow>,
787}
788
789impl<I> SpillRows<I> {
790    fn new(batches: I) -> Self {
791        Self {
792            batches,
793            current_schema: None,
794            current: Vec::new().into_iter(),
795        }
796    }
797}
798
799impl<I> Iterator for SpillRows<I>
800where
801    I: Iterator<Item = ExecResult<Batch>>,
802{
803    type Item = ExecResult<OwnedPhysicalRow>;
804
805    fn next(&mut self) -> Option<Self::Item> {
806        loop {
807            if let Some(row) = self.current.next() {
808                let schema = self
809                    .current_schema
810                    .as_ref()
811                    .expect("spill row iterator retains the current batch schema")
812                    .clone();
813                return Some(Ok(OwnedPhysicalRow::new(schema, row)));
814            }
815            match self.batches.next()? {
816                Ok(batch) => {
817                    self.current_schema = Some(batch.schema);
818                    self.current = batch.rows.into_iter();
819                }
820                Err(error) => return Some(Err(error)),
821            }
822        }
823    }
824}
825
826#[cfg(test)]
827mod tests;