1use 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#[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
88pub 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 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 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 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 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 if self.in_memory_bytes > self.budget_bytes {
207 self.spill_pending()?;
208 spilled = true;
209 }
210 Ok(spilled)
211 }
212
213 pub fn encoded_size(batch: &Batch) -> ExecResult<usize> {
216 encoded_batch_size(batch)
217 }
218
219 pub fn rows(&self) -> usize {
221 self.rows
222 }
223
224 pub fn in_memory_rows(&self) -> usize {
226 self.in_memory_rows
227 }
228
229 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 pub fn spill_path(&self) -> Option<&Path> {
263 self.spill_file.as_ref().map(NamedTempFile::path)
264 }
265
266 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 pub fn spill_pending(&mut self) -> ExecResult<bool> {
283 if self.batches.is_empty() {
284 return Ok(false);
285 }
286
287 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 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 pub fn read_rows(&self) -> ExecResult<SpillRows<SpillReader<'_>>> {
348 self.reader().map(SpillRows::new)
349 }
350
351 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 pub fn drain_all(&mut self) -> ExecResult<Vec<Batch>> {
389 self.drain()?.collect()
390 }
391
392 pub fn drain_rows(&mut self) -> ExecResult<SpillRows<SpillDrain>> {
394 self.drain().map(SpillRows::new)
395 }
396
397 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 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#[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 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 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 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
588pub 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
649pub struct SpillDrain {
651 reader: Option<BufReader<File>>,
652 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
717pub 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
782pub 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;