uqa_execution/
spill_scan.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Volcano scan over an owned [`crate::spill::SpillBuffer`].
8//!
9//! The scan transfers ownership of the spill file to its batch iterator at
10//! `open`, so disk batches are decoded one at a time and the file is removed
11//! even when execution stops early.
12
13use crate::batch::{Batch, RowSchema};
14use crate::physical::{ExecResult, PhysicalOperator};
15use crate::spill::{SharedSpill, SharedSpillReader, SpillBuffer, SpillDrain};
16
17/// One-shot physical scan over a disk-backed spill buffer.
18pub struct SpillScan {
19    schema: RowSchema,
20    buffer: Option<SpillBuffer>,
21    reader: Option<SpillDrain>,
22}
23
24/// Repeatable scan over an immutable shared spill. Cloning the source only
25/// clones an `Arc`; `open` creates an independent file reader.
26pub struct SharedSpillScan {
27    source: SharedSpill,
28    schema: RowSchema,
29    reader: Option<SharedSpillReader>,
30}
31
32impl SharedSpillScan {
33    pub fn new(source: SharedSpill) -> Self {
34        let schema = source.row_schema().clone();
35        Self {
36            source,
37            schema,
38            reader: None,
39        }
40    }
41}
42
43impl PhysicalOperator for SharedSpillScan {
44    fn row_schema(&self) -> &RowSchema {
45        &self.schema
46    }
47
48    fn estimated_cardinality(&self) -> Option<u64> {
49        u64::try_from(self.source.rows()).ok()
50    }
51
52    fn open(&mut self) -> ExecResult<()> {
53        self.reader = Some(self.source.reader()?);
54        Ok(())
55    }
56
57    fn next(&mut self) -> ExecResult<Option<Batch>> {
58        self.reader
59            .as_mut()
60            .map_or(Ok(None), |reader| reader.next().transpose())
61    }
62
63    fn close(&mut self) -> ExecResult<()> {
64        self.reader = None;
65        Ok(())
66    }
67}
68
69impl SpillScan {
70    pub fn new(schema: impl Into<RowSchema>, buffer: SpillBuffer) -> Self {
71        let schema = schema.into();
72        Self {
73            schema,
74            buffer: Some(buffer),
75            reader: None,
76        }
77    }
78}
79
80impl PhysicalOperator for SpillScan {
81    fn row_schema(&self) -> &RowSchema {
82        &self.schema
83    }
84
85    fn open(&mut self) -> ExecResult<()> {
86        let mut buffer = self.buffer.take().ok_or_else(|| {
87            crate::physical::ExecError::Other("spill scan cannot be reopened".into())
88        })?;
89        self.reader = Some(buffer.drain()?);
90        Ok(())
91    }
92
93    fn next(&mut self) -> ExecResult<Option<Batch>> {
94        let Some(reader) = self.reader.as_mut() else {
95            return Ok(None);
96        };
97        let Some(batch) = reader.next().transpose()? else {
98            return Ok(None);
99        };
100        if batch.schema != self.schema {
101            return Err(crate::physical::ExecError::Other(format!(
102                "spill scan schema mismatch: expected {:?}, got {:?}",
103                self.schema.columns(),
104                batch.schema.columns()
105            )));
106        }
107        Ok(Some(batch))
108    }
109
110    fn close(&mut self) -> ExecResult<()> {
111        self.reader = None;
112        Ok(())
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use std::collections::BTreeMap;
119
120    use uqa_core::Value;
121
122    use super::*;
123    use crate::physical::run_to_rows;
124
125    #[test]
126    fn scan_streams_a_forced_spill_in_input_order() {
127        let schema = RowSchema::new(vec!["x".into()]);
128        let mut spill = SpillBuffer::new(1);
129        for value in 0..300_i64 {
130            spill
131                .push(Batch::new(
132                    schema.clone(),
133                    vec![BTreeMap::from([("x".into(), Value::Int(value))])],
134                ))
135                .unwrap();
136        }
137        assert!(spill.has_spilled());
138
139        let mut scan = SpillScan::new(schema.columns().to_vec(), spill);
140        let (_, rows) = run_to_rows(&mut scan).unwrap();
141        assert_eq!(rows.len(), 300);
142        for (expected, row) in rows.iter().enumerate() {
143            assert_eq!(row.get("x"), Some(&Value::Int(expected as i64)));
144        }
145    }
146
147    #[test]
148    fn shared_spill_supports_independent_repeatable_scans() {
149        let schema = RowSchema::new(vec!["x".into()]);
150        let mut spill = SpillBuffer::new(1);
151        for value in 0..2_048_i64 {
152            spill
153                .push(Batch::new(
154                    schema.clone(),
155                    vec![BTreeMap::from([("x".into(), Value::Int(value))])],
156                ))
157                .unwrap();
158        }
159        let shared = spill.into_shared(schema.columns().to_vec()).unwrap();
160        let mut first = SharedSpillScan::new(shared.clone());
161        let mut second = SharedSpillScan::new(shared);
162        let (_, first_rows) = run_to_rows(&mut first).unwrap();
163        let (_, second_rows) = run_to_rows(&mut second).unwrap();
164        assert_eq!(first_rows, second_rows);
165        assert_eq!(first_rows.len(), 2_048);
166    }
167}
uqa_execution/spill_scan.rs

uqa_execution/
spill_scan.rs
