Skip to main content

radixdb_executor/access/
scan.rs

1//! Scanner construction and bounded scanner draining.
2
3use radixdb_core::{Error, Result, Row, RowVec};
4use radixdb_storage::expression::Expression as StorageExpression;
5use radixdb_storage::traits::{Scanner, Table, TypedBatchFallbackReason};
6
7use crate::context::ExecutionContext;
8
9pub fn open_scan(
10    table: &dyn Table,
11    columns: &[usize],
12    predicate: Option<&dyn StorageExpression>,
13    context: &ExecutionContext,
14) -> Result<Box<dyn Scanner>> {
15    table
16        .scan(columns, predicate)
17        .map(|scanner| budget_scanner(scanner, context))
18}
19
20pub fn open_exact_projection_scan(
21    table: &dyn Table,
22    columns: &[usize],
23    predicate: Option<&dyn StorageExpression>,
24    context: &ExecutionContext,
25) -> Result<Box<dyn Scanner>> {
26    table
27        .scan_exact_projection(columns, predicate)
28        .map(|scanner| budget_scanner(scanner, context))
29}
30
31fn budget_scanner(scanner: Box<dyn Scanner>, context: &ExecutionContext) -> Box<dyn Scanner> {
32    if !context.has_public_scan_budget() {
33        return scanner;
34    }
35    Box::new(PublicBudgetScanner {
36        inner: scanner,
37        context: context.clone(),
38        error: None,
39    })
40}
41
42struct PublicBudgetScanner {
43    inner: Box<dyn Scanner>,
44    context: ExecutionContext,
45    error: Option<Error>,
46}
47
48impl Scanner for PublicBudgetScanner {
49    fn next(&mut self) -> bool {
50        if self.error.is_some() || !self.inner.next() {
51            return false;
52        }
53        if let Err(error) = self.context.claim_public_scan_rows(1) {
54            self.error = Some(error);
55            return false;
56        }
57        true
58    }
59
60    fn row(&self) -> &Row {
61        self.inner.row()
62    }
63
64    fn err(&self) -> Option<&Error> {
65        self.error.as_ref().or_else(|| self.inner.err())
66    }
67
68    fn close(&mut self) -> Result<()> {
69        self.inner.close()
70    }
71
72    fn take_row(&mut self) -> Row {
73        self.inner.take_row()
74    }
75
76    fn estimated_count(&self) -> Option<usize> {
77        self.inner.estimated_count()
78    }
79
80    fn take_row_with_id(&mut self) -> Result<(i64, Row)> {
81        self.inner.take_row_with_id()
82    }
83
84    fn current_row_id(&self) -> Result<i64> {
85        self.inner.current_row_id()
86    }
87
88    fn warmup(&mut self) {
89        self.inner.warmup();
90    }
91
92    fn typed_batch_fallback_reason(&self) -> Option<TypedBatchFallbackReason> {
93        Some(TypedBatchFallbackReason::UnsupportedResultShape)
94    }
95}
96
97pub fn collect_scanner_rows(
98    mut scanner: Box<dyn Scanner>,
99    context: &ExecutionContext,
100) -> Result<RowVec> {
101    let mut rows = RowVec::with_capacity(scanner.estimated_count().unwrap_or(64).min(1024));
102    let mut seen = 0u64;
103    while scanner.next() {
104        seen += 1;
105        if seen.is_multiple_of(100) {
106            context.check_cancelled()?;
107        }
108        rows.push(scanner.take_row_with_id()?);
109    }
110    if let Some(error) = scanner.err() {
111        return Err(error.clone());
112    }
113    scanner.close()?;
114    Ok(rows)
115}