radixdb_executor/access/
scan.rs1use 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}