1use std::{
5 collections::{BTreeMap, BTreeSet},
6 num::NonZeroU64,
7 sync::Arc,
8};
9
10use arrow_array::RecordBatch;
11use reifydb_core::{
12 common::CommitVersion,
13 interface::{catalog::object::ObjectId, change::Change, resolved::ResolvedSeries},
14 internal_err,
15};
16use reifydb_evaluate::stack::SymbolTable;
17#[cfg(feature = "testing")]
18use reifydb_flow::backfill::testing::{InstalledScanHooks, TestingScan};
19use reifydb_flow::backfill::{Scan, backfill_partitioned};
20use reifydb_transaction::transaction::Transaction;
21use reifydb_value::{fragment::Fragment, params::Params, value::partition::ScanPartitions};
22
23use crate::{
24 Result,
25 vm::{
26 services::Services,
27 volcano::{
28 query::{QueryContext, QueryNode, query_budget},
29 scan::{
30 ringbuffer::RingBufferScan, series::SeriesScanNode, table::TableScanNode,
31 view::ViewScanNode,
32 },
33 },
34 },
35};
36
37#[cfg(test)]
38mod tests;
39
40pub fn run<'a>(
41 services: &Arc<Services>,
42 tx: Transaction<'a>,
43 sources: &BTreeSet<ObjectId>,
44 batch_size: NonZeroU64,
45 consume: impl FnMut(&mut TransactionScan<'a>, Change) -> Result<()>,
46) -> Result<()> {
47 let sources = sources.iter().map(|&source| (source, ScanPartitions::All)).collect();
48 run_partitioned(services, tx, &sources, batch_size, consume)
49}
50
51pub fn run_partitioned<'a>(
52 services: &Arc<Services>,
53 tx: Transaction<'a>,
54 sources: &BTreeMap<ObjectId, ScanPartitions>,
55 batch_size: NonZeroU64,
56 mut consume: impl FnMut(&mut TransactionScan<'a>, Change) -> Result<()>,
57) -> Result<()> {
58 let mut scan = TransactionScan {
59 services: Arc::clone(services),
60 tx,
61 opened: None,
62 };
63 #[cfg(feature = "testing")]
64 if let Some(InstalledScanHooks(hooks)) = services.ioc.try_resolve::<InstalledScanHooks>() {
65 return backfill_partitioned(
66 &mut TestingScan::over(scan, hooks),
67 sources,
68 batch_size,
69 |testing, change| consume(testing.scan_mut(), change),
70 );
71 }
72 backfill_partitioned(&mut scan, sources, batch_size, &mut consume)
73}
74
75pub struct TransactionScan<'a> {
76 services: Arc<Services>,
77 tx: Transaction<'a>,
78 opened: Option<Opened>,
79}
80
81struct Opened {
82 node: Box<dyn QueryNode>,
83 context: QueryContext,
84 tagged_series: bool,
85}
86
87impl<'a> TransactionScan<'a> {
88 pub fn transaction(&mut self) -> &mut Transaction<'a> {
89 &mut self.tx
90 }
91
92 fn node(
93 &mut self,
94 source: ObjectId,
95 partitions: ScanPartitions,
96 context: &Arc<QueryContext>,
97 ) -> Result<(Box<dyn QueryNode>, bool)> {
98 let catalog = &self.services.catalog;
99 let tx = &mut self.tx;
100 match source {
101 ObjectId::Table(id) => {
102 let table = catalog.resolve_table(tx, id)?;
103 let node = match partitions {
104 ScanPartitions::All => {
105 TableScanNode::new(table, None, Arc::clone(context), tx)?
106 }
107 ScanPartitions::One(partition) => {
108 TableScanNode::new(table, Some(partition), Arc::clone(context), tx)?
109 }
110 ScanPartitions::Many(listed) => {
111 TableScanNode::with_partitions(table, listed, Arc::clone(context), tx)?
112 }
113 }
114 .oldest_first()
115 .keep_dictionary_ids();
116 Ok((Box::new(node), false))
117 }
118 ObjectId::View(id) => {
119 let view = catalog.resolve_view(tx, id)?;
120 let partition = match partitions {
121 ScanPartitions::One(partition) => Some(partition),
122 ScanPartitions::All | ScanPartitions::Many(_) => None,
123 };
124 let node = ViewScanNode::new(view, partition, Arc::clone(context), tx)?
125 .oldest_first()
126 .keep_dictionary_ids();
127 Ok((Box::new(node), false))
128 }
129 ObjectId::RingBuffer(id) => {
130 let ringbuffer = catalog.resolve_ringbuffer(tx, id)?;
131 let node = RingBufferScan::new(ringbuffer, Arc::clone(context), tx)?
132 .oldest_first()
133 .keep_dictionary_ids();
134 Ok((Box::new(node), false))
135 }
136 ObjectId::Series(id) => {
137 let series = catalog.get_series(tx, id)?;
138 let namespace = catalog.resolve_namespace(tx, series.namespace)?;
139 let tagged = series.tag.is_some();
140 let series =
141 ResolvedSeries::new(Fragment::internal(series.name.clone()), namespace, series);
142 let node = SeriesScanNode::new(series, None, None, None, None, Arc::clone(context))?
143 .oldest_first()
144 .keep_dictionary_ids();
145 Ok((Box::new(node), tagged))
146 }
147 ObjectId::TableVirtual(_) | ObjectId::Dictionary(_) | ObjectId::Queue(_) => {
148 internal_err!("backfill cannot scan source {:?}: it has no row scan node", source)
149 }
150 }
151 }
152}
153
154impl Scan for TransactionScan<'_> {
155 fn version(&self) -> CommitVersion {
156 self.tx.version()
157 }
158
159 fn open(&mut self, source: ObjectId, batch_size: NonZeroU64) -> Result<()> {
160 self.open_partition(source, ScanPartitions::All, batch_size)
161 }
162
163 fn open_partition(
164 &mut self,
165 source: ObjectId,
166 partitions: ScanPartitions,
167 batch_size: NonZeroU64,
168 ) -> Result<()> {
169 let context = Arc::new(QueryContext {
170 services: Arc::clone(&self.services),
171 source: None,
172 batch_size: batch_size.get(),
173 params: Params::None,
174 symbols: SymbolTable::new(),
175 identity: self.tx.identity(),
176 memory: query_budget(&self.services),
177 });
178 let (mut node, tagged_series) = self.node(source, partitions, &context)?;
179 node.initialize(&mut self.tx, &context)?;
180 self.opened = Some(Opened {
181 node,
182 context: (*context).clone(),
183 tagged_series,
184 });
185 Ok(())
186 }
187
188 fn next(&mut self) -> Result<Option<RecordBatch>> {
189 let Some(opened) = self.opened.as_mut() else {
190 return internal_err!("transaction scan pulled before any source was opened");
191 };
192 let Some(mut chunk) = opened.node.next(&mut self.tx, &mut opened.context)? else {
193 return Ok(None);
194 };
195 if opened.tagged_series {
196 chunk.remove_column(1);
197 }
198 Ok(Some(chunk))
199 }
200}