use std::{
collections::{BTreeMap, BTreeSet},
num::NonZeroU64,
sync::Arc,
};
use arrow_array::RecordBatch;
use reifydb_core::{
common::CommitVersion,
interface::{catalog::object::ObjectId, change::Change, resolved::ResolvedSeries},
internal_err,
};
use reifydb_evaluate::stack::SymbolTable;
#[cfg(feature = "testing")]
use reifydb_flow::backfill::testing::{InstalledScanHooks, TestingScan};
use reifydb_flow::backfill::{Scan, backfill_partitioned};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{fragment::Fragment, params::Params, value::partition::ScanPartitions};
use crate::{
Result,
vm::{
services::Services,
volcano::{
query::{QueryContext, QueryNode, query_budget},
scan::{
ringbuffer::RingBufferScan, series::SeriesScanNode, table::TableScanNode,
view::ViewScanNode,
},
},
},
};
#[cfg(test)]
mod tests;
pub fn run<'a>(
services: &Arc<Services>,
tx: Transaction<'a>,
sources: &BTreeSet<ObjectId>,
batch_size: NonZeroU64,
consume: impl FnMut(&mut TransactionScan<'a>, Change) -> Result<()>,
) -> Result<()> {
let sources = sources.iter().map(|&source| (source, ScanPartitions::All)).collect();
run_partitioned(services, tx, &sources, batch_size, consume)
}
pub fn run_partitioned<'a>(
services: &Arc<Services>,
tx: Transaction<'a>,
sources: &BTreeMap<ObjectId, ScanPartitions>,
batch_size: NonZeroU64,
mut consume: impl FnMut(&mut TransactionScan<'a>, Change) -> Result<()>,
) -> Result<()> {
let mut scan = TransactionScan {
services: Arc::clone(services),
tx,
opened: None,
};
#[cfg(feature = "testing")]
if let Some(InstalledScanHooks(hooks)) = services.ioc.try_resolve::<InstalledScanHooks>() {
return backfill_partitioned(
&mut TestingScan::over(scan, hooks),
sources,
batch_size,
|testing, change| consume(testing.scan_mut(), change),
);
}
backfill_partitioned(&mut scan, sources, batch_size, &mut consume)
}
pub struct TransactionScan<'a> {
services: Arc<Services>,
tx: Transaction<'a>,
opened: Option<Opened>,
}
struct Opened {
node: Box<dyn QueryNode>,
context: QueryContext,
tagged_series: bool,
}
impl<'a> TransactionScan<'a> {
pub fn transaction(&mut self) -> &mut Transaction<'a> {
&mut self.tx
}
fn node(
&mut self,
source: ObjectId,
partitions: ScanPartitions,
context: &Arc<QueryContext>,
) -> Result<(Box<dyn QueryNode>, bool)> {
let catalog = &self.services.catalog;
let tx = &mut self.tx;
match source {
ObjectId::Table(id) => {
let table = catalog.resolve_table(tx, id)?;
let node = match partitions {
ScanPartitions::All => {
TableScanNode::new(table, None, Arc::clone(context), tx)?
}
ScanPartitions::One(partition) => {
TableScanNode::new(table, Some(partition), Arc::clone(context), tx)?
}
ScanPartitions::Many(listed) => {
TableScanNode::with_partitions(table, listed, Arc::clone(context), tx)?
}
}
.oldest_first()
.keep_dictionary_ids();
Ok((Box::new(node), false))
}
ObjectId::View(id) => {
let view = catalog.resolve_view(tx, id)?;
let partition = match partitions {
ScanPartitions::One(partition) => Some(partition),
ScanPartitions::All | ScanPartitions::Many(_) => None,
};
let node = ViewScanNode::new(view, partition, Arc::clone(context), tx)?
.oldest_first()
.keep_dictionary_ids();
Ok((Box::new(node), false))
}
ObjectId::RingBuffer(id) => {
let ringbuffer = catalog.resolve_ringbuffer(tx, id)?;
let node = RingBufferScan::new(ringbuffer, Arc::clone(context), tx)?
.oldest_first()
.keep_dictionary_ids();
Ok((Box::new(node), false))
}
ObjectId::Series(id) => {
let series = catalog.get_series(tx, id)?;
let namespace = catalog.resolve_namespace(tx, series.namespace)?;
let tagged = series.tag.is_some();
let series =
ResolvedSeries::new(Fragment::internal(series.name.clone()), namespace, series);
let node = SeriesScanNode::new(series, None, None, None, None, Arc::clone(context))?
.oldest_first()
.keep_dictionary_ids();
Ok((Box::new(node), tagged))
}
ObjectId::TableVirtual(_) | ObjectId::Dictionary(_) | ObjectId::Queue(_) => {
internal_err!("backfill cannot scan source {:?}: it has no row scan node", source)
}
}
}
}
impl Scan for TransactionScan<'_> {
fn version(&self) -> CommitVersion {
self.tx.version()
}
fn open(&mut self, source: ObjectId, batch_size: NonZeroU64) -> Result<()> {
self.open_partition(source, ScanPartitions::All, batch_size)
}
fn open_partition(
&mut self,
source: ObjectId,
partitions: ScanPartitions,
batch_size: NonZeroU64,
) -> Result<()> {
let context = Arc::new(QueryContext {
services: Arc::clone(&self.services),
source: None,
batch_size: batch_size.get(),
params: Params::None,
symbols: SymbolTable::new(),
identity: self.tx.identity(),
memory: query_budget(&self.services),
});
let (mut node, tagged_series) = self.node(source, partitions, &context)?;
node.initialize(&mut self.tx, &context)?;
self.opened = Some(Opened {
node,
context: (*context).clone(),
tagged_series,
});
Ok(())
}
fn next(&mut self) -> Result<Option<RecordBatch>> {
let Some(opened) = self.opened.as_mut() else {
return internal_err!("transaction scan pulled before any source was opened");
};
let Some(mut chunk) = opened.node.next(&mut self.tx, &mut opened.context)? else {
return Ok(None);
};
if opened.tagged_series {
chunk.remove_column(1);
}
Ok(Some(chunk))
}
}