use std::{cmp::Reverse, sync::Arc};
use arrow_array::RecordBatch;
use reifydb_core::{
common::TimeSource,
error::diagnostic::{internal::internal, query::no_column_block},
interface::{
catalog::{
column_block::{ColumnBlock, ColumnBlockSource},
object::ObjectId,
series::Series,
},
resolved::ResolvedSeries,
},
key::{any::TaggedKey, partition::PartitionKey},
value::column::{builder::ColumnBuilder, headers::ColumnHeaders},
};
use reifydb_store_column::{block::Schema, store::ColumnStore};
use reifydb_transaction::{multi::RangeScope, transaction::Transaction};
use reifydb_value::{
error::Error,
value::{partition::Partition, system_columns::SystemColumn},
};
use crate::{
Result,
vm::volcano::{
query::{QueryContext, QueryNode},
scan::{
column_block_sequence::BlockSequenceReader, column_decode::BlockDecoder,
column_predicate::series_scan_predicate, column_prune::prune_series_blocks, empty_scan,
scan_headers, series::empty_series_columns, source_system_columns,
},
},
};
enum ScanState {
Unopened,
Reading {
reader: Box<BlockSequenceReader>,
emitted: bool,
},
Empty,
Done,
}
pub struct ColumnSeriesScanNode {
series: ResolvedSeries,
key_range_start: Option<u64>,
key_range_end: Option<u64>,
variant_tag: Option<u8>,
partition: Option<Partition>,
context: Arc<QueryContext>,
headers: ColumnHeaders,
system_columns: Vec<SystemColumn>,
state: ScanState,
}
impl ColumnSeriesScanNode {
pub fn new(
series: ResolvedSeries,
key_range_start: Option<u64>,
key_range_end: Option<u64>,
variant_tag: Option<u8>,
partition: Option<Partition>,
context: Arc<QueryContext>,
) -> Self {
let def = series.def();
let mut columns = vec![def.key.column()];
if def.tag.is_some() {
columns.push("tag");
}
for col in def.data_columns() {
columns.push(col.name.as_str());
}
let system_columns =
source_system_columns(!def.partition_by.is_empty(), def.time != TimeSource::None, false);
let headers = scan_headers(columns.into_iter(), &system_columns);
Self {
series,
key_range_start,
key_range_end,
variant_tag,
partition,
context,
headers,
system_columns,
state: ScanState::Unopened,
}
}
fn series_reads_as_empty(&self, rx: &mut Transaction<'_>, def: &Series) -> Result<bool> {
let catalog = &self.context.services.catalog;
let partitions = match self.partition {
Some(partition) => vec![partition],
None if def.partition_by.is_empty() => vec![Partition::default()],
None => self.registered_partitions(rx, def)?,
};
for partition in partitions {
let Some(metadata) = catalog.find_series_metadata(rx, def.id, partition)? else {
continue;
};
if metadata.row_count > 0 {
return Ok(false);
}
}
Ok(true)
}
fn registered_partitions(&self, rx: &mut Transaction<'_>, def: &Series) -> Result<Vec<Partition>> {
let mut partitions = Vec::new();
let stream = rx.range(PartitionKey::full_scan(ObjectId::Series(def.id)), RangeScope::All, 1024)?;
for entry in stream {
if let TaggedKey::Partition(key) = entry?.key {
partitions.push(key.partition);
}
}
Ok(partitions)
}
fn open(&self, rx: &mut Transaction<'_>) -> Result<ScanState> {
let services = &self.context.services;
let name = self.series.fully_qualified_name();
let def = self.series.def();
let blocks = match self.partition {
Some(partition) => {
services.catalog.list_column_blocks_for_series_partition(rx, def.id, partition)?
}
None => services.catalog.list_column_blocks_for_series(rx, def.id)?,
};
let blocks_were_empty = blocks.is_empty();
let predicate = series_scan_predicate(def, self.key_range_start, self.key_range_end, self.variant_tag);
let mut pruned =
prune_series_blocks(blocks, self.key_range_start, self.key_range_end, predicate.as_ref());
if blocks_were_empty && self.series_reads_as_empty(rx, def)? {
return Ok(ScanState::Empty);
}
if pruned.is_empty() && !blocks_were_empty {
return Ok(ScanState::Empty);
}
if pruned.is_empty() {
return Err(Error(Box::new(no_column_block(
self.series.identifier().clone(),
"series",
&name,
))));
}
pruned.sort_by_key(|b| Reverse(bucket_order(b)));
let store = services.ioc.try_resolve::<Arc<ColumnStore>>().ok_or_else(|| {
Error(Box::new(internal(format!(
"column store is not registered, cannot read series {}",
name
))))
})?;
Ok(ScanState::Reading {
reader: Box::new(
BlockSequenceReader::new(Arc::clone(&store), pruned, self.context.batch_size as usize)
.with_predicate(predicate)
.with_decoder(Some(BlockDecoder::new(
services.block_encodings.clone(),
store.session().clone(),
))),
),
emitted: false,
})
}
}
fn bucket_order(block: &ColumnBlock) -> (u64, Option<Partition>) {
match block.source {
ColumnBlockSource::Series {
start,
partition,
..
} => (start, partition),
ColumnBlockSource::Table {
..
} => (0, None),
}
}
fn empty_columns(schema: &Schema) -> Result<RecordBatch> {
let columns = schema
.iter()
.filter(|(name, _, _)| SystemColumn::from_name(name).is_none())
.map(|(name, ty, _)| ColumnBuilder::with_capacity(ty.clone(), 0).finish(name))
.collect();
let system: Vec<SystemColumn> =
schema.iter().filter_map(|(name, _, _)| SystemColumn::from_name(name)).collect();
empty_scan(columns, &system)
}
impl QueryNode for ColumnSeriesScanNode {
fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
Ok(())
}
fn next<'a>(&mut self, rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
if matches!(self.state, ScanState::Unopened) {
self.state = self.open(rx)?;
}
if matches!(self.state, ScanState::Empty) {
self.state = ScanState::Done;
return Ok(Some(empty_scan(empty_series_columns(self.series.def()), &self.system_columns)?));
}
let ScanState::Reading {
reader,
emitted,
} = &mut self.state
else {
return Ok(None);
};
match reader.next()? {
Some(batch) => {
*emitted = true;
Ok(Some(batch))
}
None => {
let empty = (!*emitted)
.then(|| reader.schema().map(empty_columns))
.flatten()
.transpose()?;
self.state = ScanState::Done;
Ok(empty)
}
}
}
fn headers(&self) -> Option<ColumnHeaders> {
Some(self.headers.clone())
}
}