use std::sync::Arc;
use reifydb_codec::row::{bytes::EncodedBytes, queue::EncodedQueueRow, shape::RowShape};
use reifydb_core::{
interface::{catalog::dictionary::Dictionary, resolved::ResolvedQueue, store::MultiVersionRow},
internal_error,
key::{any::TaggedKey, bound::TaggedKeyBoundRange, row::RowKeyRange},
value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns, headers::ColumnHeaders},
};
use reifydb_transaction::{multi::RangeScope, transaction::Transaction};
use reifydb_value::{
fragment::Fragment,
value::{row_number::RowNumber, system_columns::SystemColumns, value_type::ValueType},
};
use tracing::instrument;
use super::super::decode_dictionary_columns;
use crate::{
Result,
vm::volcano::query::{QueryContext, QueryNode},
};
type DrainedBatch = (Vec<EncodedBytes>, Vec<RowNumber>, Option<TaggedKey>, bool);
pub struct QueueScan {
queue: ResolvedQueue,
headers: ColumnHeaders,
shape: Option<RowShape>,
storage_types: Vec<ValueType>,
dictionaries: Vec<Option<Dictionary>>,
last_key: Option<TaggedKey>,
exhausted: bool,
context: Option<Arc<QueryContext>>,
}
impl QueueScan {
pub fn new(queue: ResolvedQueue, context: Arc<QueryContext>, rx: &mut Transaction<'_>) -> Result<Self> {
let mut storage_types = Vec::with_capacity(queue.columns().len());
let mut dictionaries = Vec::with_capacity(queue.columns().len());
for col in queue.columns() {
if let Some(dict_id) = col.dictionary_id
&& let Some(dict) = context.services.catalog.find_dictionary(rx, dict_id)?
{
storage_types.push(ValueType::DictionaryId);
dictionaries.push(Some(dict));
continue;
}
storage_types.push(col.constraint.get_type());
dictionaries.push(None);
}
let headers = ColumnHeaders {
columns: queue.columns().iter().map(|col| Fragment::internal(&col.name)).collect(),
};
Ok(Self {
queue,
headers,
shape: None,
storage_types,
dictionaries,
last_key: None,
exhausted: false,
context: Some(context),
})
}
fn get_or_load_shape(&mut self, rx: &mut Transaction, first: &EncodedBytes) -> Result<RowShape> {
if let Some(shape) = &self.shape {
return Ok(shape.clone());
}
let fingerprint = EncodedQueueRow::view(first).fingerprint();
let stored_ctx = self.context.as_ref().expect("QueueScan context not set");
let shape = stored_ctx.services.catalog.get_or_load_row_shape(fingerprint, rx)?.ok_or_else(|| {
internal_error!(
"RowShape with fingerprint {:?} not found for queue {}",
fingerprint,
self.queue.def().name
)
})?;
self.shape = Some(shape.clone());
Ok(shape)
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::queue::range_open")]
fn open_range<'rx, 'tx>(
rx: &'rx mut Transaction<'tx>,
range: TaggedKeyBoundRange,
batch_size: u64,
) -> Result<Box<dyn Iterator<Item = Result<MultiVersionRow<TaggedKey>>> + Send + 'rx>> {
rx.range_rev(range, RangeScope::All, batch_size as usize)
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::queue::drain")]
fn drain_batch(
stream: &mut dyn Iterator<Item = Result<MultiVersionRow<TaggedKey>>>,
batch_size: u64,
) -> Result<DrainedBatch> {
let mut batch: Vec<EncodedBytes> = Vec::new();
let mut row_numbers: Vec<RowNumber> = Vec::new();
let mut new_last_key = None;
let mut drained = false;
for _ in 0..batch_size {
match stream.next() {
Some(Ok(multi)) => {
if let TaggedKey::Row(key) = &multi.key {
row_numbers.push(key.row);
batch.push(multi.bytes);
new_last_key = Some(multi.key);
}
}
Some(Err(e)) => return Err(e),
None => {
drained = true;
break;
}
}
}
Ok((batch, row_numbers, new_last_key, drained))
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::queue::column_alloc")]
fn storage_columns(&self, shape: &RowShape, declared: usize) -> Result<Vec<ColumnWithName>> {
let mut storage_columns: Vec<ColumnWithName> = self
.queue
.columns()
.iter()
.enumerate()
.map(|(idx, col)| ColumnWithName {
name: Fragment::internal(&col.name),
data: ColumnBuffer::with_capacity(self.storage_types[idx].clone(), 0),
})
.collect();
for index in declared..shape.field_count() {
let field = shape.get_field(index).ok_or_else(|| {
internal_error!("queue {} shape lost field {}", self.queue.def().name, index)
})?;
storage_columns.push(ColumnWithName {
name: Fragment::internal(field.name.clone()),
data: ColumnBuffer::with_capacity(field.constraint.get_type(), 0),
});
}
Ok(storage_columns)
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::queue::append_rows")]
fn append_batch(
shape: &RowShape,
columns: &mut Columns,
bytes_vec: Vec<EncodedBytes>,
row_numbers: Vec<RowNumber>,
) -> Result<()> {
columns.append_rows(shape, bytes_vec.into_iter(), row_numbers)?;
Ok(())
}
fn enqueue_order_range(&self) -> TaggedKeyBoundRange {
RowKeyRange::scan_range(self.queue.def().id.into(), None).resume_before(self.last_key.as_ref())
}
fn empty_declared_columns(&self) -> Columns {
Columns::new(
self.queue
.columns()
.iter()
.map(|col| ColumnWithName {
name: Fragment::internal(&col.name),
data: ColumnBuffer::none_typed(col.constraint.get_type(), 0),
})
.collect(),
)
}
}
impl QueryNode for QueueScan {
#[instrument(level = "trace", skip_all, name = "volcano::scan::queue::initialize")]
fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
Ok(())
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::queue::next")]
fn next<'a>(&mut self, rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<Columns>> {
if self.exhausted {
return Ok(None);
}
let batch_size = self.context.as_ref().expect("QueueScan context not set").batch_size;
let range = self.enqueue_order_range();
let (batch_rows, row_numbers, new_last_key, drained) = {
let mut stream = Self::open_range(rx, range, batch_size)?;
Self::drain_batch(&mut stream, batch_size)?
};
if drained {
self.exhausted = true;
}
if batch_rows.is_empty() {
self.exhausted = true;
if self.last_key.is_none() {
return Ok(Some(self.empty_declared_columns()));
}
return Ok(None);
}
self.last_key = new_last_key;
let shape = self.get_or_load_shape(rx, &batch_rows[0])?;
let declared = self.queue.columns().len();
let storage_columns = self.storage_columns(&shape, declared)?;
let mut columns = Columns::with_system(storage_columns, SystemColumns::default());
Self::append_batch(&shape, &mut columns, batch_rows, row_numbers)?;
decode_dictionary_columns(&mut columns, &self.dictionaries, rx)?;
columns.columns.truncate(declared);
columns.names.truncate(declared);
Ok(Some(columns))
}
fn headers(&self) -> Option<ColumnHeaders> {
Some(self.headers.clone())
}
}