use std::sync::Arc;
use arrow_array::{Array, ArrayRef, RecordBatch, UInt64Array};
use arrow_schema::{FieldRef, SchemaRef};
use reifydb_codec::row::{bytes::EncodedBytes, series::EncodedSeriesRow, shape::RowShape};
use reifydb_core::{
common::TimeSource,
interface::{
catalog::{series::Series, storage::StorageId},
resolved::ResolvedSeries,
store::MultiVersionRow,
},
key::{
any::TaggedKey,
bound::TaggedKeyBoundRange,
series::{PartitionedSeriesRowKeyRange, SeriesRowKeyRange},
},
value::{
batch::{batch, batch_with, decode_cells},
column::{
builder::ColumnBuilder,
factory::{datetime, none_typed, uint1},
headers::ColumnHeaders,
},
},
};
use reifydb_transaction::{multi::RangeScope, transaction::Transaction};
use reifydb_value::{
reifydb_assertions,
value::{
datetime::DateTime,
partition::Partition,
system_columns::{SystemColumn, system_field},
value_type::ValueType,
},
};
use tracing::instrument;
use super::{
empty_scan,
merge::{MergeLayout, PartitionMerge},
partition_array, scan_headers, source_system_columns,
};
use crate::{
Result,
vm::{
instruction::dml::shape::get_or_create_series_shape,
volcano::{
decode_dictionary_column,
query::{QueryContext, QueryNode},
},
},
};
pub struct SeriesScanNode {
series: ResolvedSeries,
key_range_start: Option<u64>,
key_range_end: Option<u64>,
variant_tag: Option<u8>,
partition: Option<Partition>,
context: Option<Arc<QueryContext>>,
headers: ColumnHeaders,
last_key: Option<TaggedKey>,
exhausted: bool,
system_columns: Vec<SystemColumn>,
oldest_first: bool,
merge: Option<PartitionMerge>,
schema: Option<SchemaRef>,
keep_dictionary_ids: bool,
}
impl SeriesScanNode {
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>,
) -> Result<Self> {
let mut columns = vec![series.def().key.column()];
if series.def().tag.is_some() {
columns.push("tag");
}
for col in series.columns() {
columns.push(col.name.as_str());
}
let system_columns = source_system_columns(
!series.def().partition_by.is_empty(),
series.def().time != TimeSource::None,
false,
);
let headers = scan_headers(columns.into_iter(), &system_columns);
Ok(Self {
series,
key_range_start,
key_range_end,
variant_tag,
partition,
context: Some(context),
headers,
last_key: None,
exhausted: false,
system_columns,
oldest_first: false,
keep_dictionary_ids: false,
merge: None,
schema: None,
})
}
pub(crate) fn oldest_first(mut self) -> Self {
self.oldest_first = true;
self
}
pub(crate) fn keep_dictionary_ids(mut self) -> Self {
self.keep_dictionary_ids = true;
self
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::series::range_open")]
fn open_range<'rx, 'tx>(
rx: &'rx mut Transaction<'tx>,
range: TaggedKeyBoundRange,
scope: RangeScope,
batch_size: u64,
oldest_first: bool,
) -> Result<Box<dyn Iterator<Item = Result<MultiVersionRow<TaggedKey>>> + Send + 'rx>> {
if oldest_first {
return rx.range_rev(range, scope, batch_size as usize);
}
rx.range(range, scope, batch_size as usize)
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::series::drain")]
fn drain_batch(
stream: &mut dyn Iterator<Item = Result<MultiVersionRow<TaggedKey>>>,
batch_size: u64,
partitioned: bool,
has_tag: bool,
) -> Result<SeriesBatch> {
let mut batch = SeriesBatch::default();
let mut count = 0;
for entry in stream {
let entry = entry?;
let decoded: Option<(u64, u64, Option<u8>, Option<Partition>)> = if partitioned {
match &entry.key {
TaggedKey::PartitionedSeriesRow(pk) => {
Some((pk.key, pk.sequence, pk.variant_tag, Some(pk.partition)))
}
_ => None,
}
} else {
match &entry.key {
TaggedKey::SeriesRow(k) => Some((k.key, k.sequence, k.variant_tag, None)),
_ => None,
}
};
if let Some((key_val, sequence, variant_tag, partition)) = decoded {
batch.key_values.push(key_val);
batch.sequences.push(sequence);
if let Some(p) = partition {
batch.partitions.push(p);
}
let row = EncodedSeriesRow::view(&entry.bytes);
batch.created_at_values.push(row.created_at());
if let Some(time) = row.time() {
batch.time_values.push(time);
}
batch.updated_at_values.push(row.updated_at());
if has_tag {
batch.tags.push(variant_tag.unwrap_or(0));
}
batch.rows.push(entry.bytes);
batch.last_key = Some(entry.key.clone());
count += 1;
if count >= batch_size as usize {
break;
}
}
}
Ok(batch)
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::series::assemble")]
fn assemble<'a>(
&mut self,
rx: &mut Transaction<'a>,
stored_ctx: &QueryContext,
read_shape: &RowShape,
scanned: SeriesBatch,
has_tag: bool,
partitioned: bool,
) -> Result<Option<RecordBatch>> {
let series = self.series.def();
let mut result_columns = Vec::new();
result_columns.push(series.key_column_data(scanned.key_values));
if has_tag {
result_columns.push(uint1("tag", scanned.tags));
}
for (col_idx, col_def) in series.data_columns().enumerate() {
let col_type = col_def.constraint.get_type();
let dictionary = match col_def.dictionary_id {
Some(dict_id) => stored_ctx.services.catalog.find_dictionary(rx, dict_id)?,
None => None,
};
match dictionary {
Some(dictionary) => {
let mut ids = ColumnBuilder::with_capacity(
ValueType::DictionaryId,
scanned.rows.len(),
);
decode_cells(&mut ids, &col_def.name, read_shape, col_idx + 1, &scanned.rows)?;
if self.keep_dictionary_ids {
ids.set_dictionary_id(dictionary.id);
result_columns.push(ids.finish(&col_def.name));
continue;
}
let (field, array) = ids.finish(&col_def.name);
result_columns.push(decode_dictionary_column(
&field,
&array,
&dictionary,
col_type,
rx,
)?);
}
None => {
let mut data = ColumnBuilder::with_capacity(col_type, scanned.rows.len());
decode_cells(&mut data, &col_def.name, read_shape, col_idx + 1, &scanned.rows)?;
result_columns.push(data.finish(&col_def.name));
}
}
}
let mut stamps: Vec<(SystemColumn, ArrayRef)> =
vec![(SystemColumn::RowNumbers, Arc::new(UInt64Array::from(scanned.sequences)))];
if partitioned {
stamps.push((SystemColumn::Partitions, partition_array(&scanned.partitions)));
}
stamps.push((
SystemColumn::CreatedAt,
datetime(SystemColumn::CreatedAt.name(), scanned.created_at_values).1,
));
stamps.push((
SystemColumn::UpdatedAt,
datetime(SystemColumn::UpdatedAt.name(), scanned.updated_at_values).1,
));
if !scanned.time_values.is_empty() {
stamps.push((SystemColumn::Time, datetime(SystemColumn::Time.name(), scanned.time_values).1));
}
let row_count = result_columns[0].1.len();
result_columns.extend(stamps
.into_iter()
.map(|(column, array)| (system_field(column, array.logical_null_count() > 0), array)));
let out = match &self.schema {
Some(schema) => batch_with(schema, result_columns, row_count)?,
None => batch(result_columns)?,
};
self.schema = Some(out.schema());
Ok(Some(out))
}
}
#[derive(Default)]
struct SeriesBatch {
key_values: Vec<u64>,
tags: Vec<u8>,
sequences: Vec<u64>,
partitions: Vec<Partition>,
created_at_values: Vec<DateTime>,
time_values: Vec<DateTime>,
updated_at_values: Vec<DateTime>,
rows: Vec<EncodedBytes>,
last_key: Option<TaggedKey>,
}
pub(crate) fn empty_series_columns(series: &Series) -> Vec<(FieldRef, ArrayRef)> {
let key_type = series
.columns
.iter()
.find(|c| c.name == series.key.column())
.map(|c| c.constraint.get_type())
.unwrap_or(ValueType::Int8);
let mut result_columns = Vec::new();
result_columns.push(none_typed(series.key.column(), key_type, 0));
if series.tag.is_some() {
result_columns.push(none_typed("tag", ValueType::Uint1, 0));
}
for col_def in series.data_columns() {
result_columns.push(none_typed(&col_def.name, col_def.constraint.get_type(), 0));
}
result_columns
}
impl QueryNode for SeriesScanNode {
#[instrument(name = "volcano::scan::series::initialize", level = "trace", skip_all)]
fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
Ok(())
}
#[instrument(name = "volcano::scan::series::next", level = "trace", skip_all)]
fn next<'a>(&mut self, rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
reifydb_assertions! {
assert!(self.context.is_some(), "SeriesScanNode::next() called before initialize()");
}
let stored_ctx = self.context.as_ref().unwrap();
if self.exhausted {
return Ok(None);
}
let batch_size = stored_ctx.batch_size;
let series = self.series.def();
let has_tag = series.tag.is_some();
if let (Some(start), Some(end)) = (self.key_range_start, self.key_range_end)
&& start >= end
{
self.exhausted = true;
return Ok(Some(empty_scan(empty_series_columns(series), &self.system_columns)?));
}
let partitioned = !series.partition_by.is_empty();
let storage = StorageId::series(series.id);
let after = if self.oldest_first {
None
} else {
self.last_key.as_ref()
};
let range = if partitioned {
match self.partition {
Some(partition) => PartitionedSeriesRowKeyRange::scan_range(
storage,
partition,
has_tag,
self.variant_tag,
self.key_range_start,
self.key_range_end,
after,
),
None => PartitionedSeriesRowKeyRange::full_scan_range(storage, after),
}
} else {
SeriesRowKeyRange::scan_range(
storage,
has_tag,
self.variant_tag,
self.key_range_start,
self.key_range_end,
after,
)
};
let range = if self.oldest_first {
range.resume_before(self.last_key.as_ref())
} else {
range
};
let read_shape = get_or_create_series_shape(&stored_ctx.services.catalog, self.series.def(), rx)?;
let stored_ctx = stored_ctx.clone();
let scope = RangeScope::All;
let batch = if self.oldest_first && partitioned {
let fixed = self.partition.map(|_| range);
let merge = self
.merge
.get_or_insert_with(|| PartitionMerge::new(MergeLayout::Series, storage, fixed));
let rows = merge.next(rx, batch_size)?;
Self::drain_batch(&mut rows.into_iter().map(Ok), batch_size, partitioned, has_tag)?
} else {
let mut stream = Self::open_range(rx, range, scope, batch_size, self.oldest_first)?;
Self::drain_batch(&mut stream, batch_size, partitioned, has_tag)?
};
if batch.key_values.is_empty() {
self.exhausted = true;
if self.last_key.is_none() {
return Ok(Some(empty_scan(
empty_series_columns(self.series.def()),
&self.system_columns,
)?));
}
return Ok(None);
}
self.last_key = batch.last_key.clone();
self.assemble(rx, &stored_ctx, &read_shape, batch, has_tag, partitioned)
}
fn headers(&self) -> Option<ColumnHeaders> {
Some(self.headers.clone())
}
}