use std::sync::Arc;
use arrow_array::{ArrayRef, RecordBatch};
use arrow_schema::FieldRef;
use postcard::from_bytes;
use reifydb_core::{
interface::{catalog::dictionary::Dictionary, resolved::ResolvedDictionary, store::SingleVersionRange},
internal_error,
key::{any::TaggedKey, bound::TaggedKeyBoundRange, catalog::DictionaryEntryIndexKey},
value::{
batch::batch,
column::{
builder::ColumnBuilder,
factory::{none_typed, uint1, uint2, uint4, uint8, uint16},
headers::ColumnHeaders,
},
},
};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
fragment::Fragment,
reifydb_assertions,
value::{Value, dictionary::DictionaryEntryId, value_type::ValueType},
};
use tracing::instrument;
use crate::{
Result,
vm::volcano::query::{QueryContext, QueryNode},
};
pub struct DictionaryScanNode {
dictionary: ResolvedDictionary,
context: Option<Arc<QueryContext>>,
headers: ColumnHeaders,
last_key: Option<TaggedKey>,
exhausted: bool,
}
impl DictionaryScanNode {
pub fn new(dictionary: ResolvedDictionary, context: Arc<QueryContext>) -> Result<Self> {
let headers = ColumnHeaders {
columns: vec![Fragment::internal("id"), Fragment::internal("value")],
};
Ok(Self {
dictionary,
context: Some(context),
headers,
last_key: None,
exhausted: false,
})
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::dictionary::drain")]
fn drain_batch<'a>(
rx: &mut Transaction<'a>,
range: TaggedKeyBoundRange,
batch_size: u64,
dict_def: &Dictionary,
) -> Result<(Vec<DictionaryEntryId>, Vec<Value>, Option<TaggedKey>)> {
let mut ids: Vec<DictionaryEntryId> = Vec::new();
let mut values: Vec<Value> = Vec::new();
let mut new_last_key = None;
let single = rx
.single()
.ok_or_else(|| internal_error!("single-version store is not available for dictionary scans"))?;
let store = single.read_store();
let batch = SingleVersionRange::range_batch(&store, range.encode(), batch_size)?;
for entry in batch.items {
let Some(key) = DictionaryEntryIndexKey::decode(&entry.key) else {
panic!(
"dictionary {} holds an entry index key that does not decode: {:?}",
dict_def.id, entry.key
);
};
let entry_id = DictionaryEntryId::from_u128(key.id, dict_def.id_type.clone())?;
new_last_key = Some(TaggedKey::from(key));
let value: Value = from_bytes(&entry.bytes)
.map_err(|e| internal_error!("Failed to deserialize dictionary value: {}", e))?;
ids.push(entry_id);
values.push(value);
}
Ok((ids, values, new_last_key))
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::dictionary::empty_columns")]
fn empty_columns(dict_def: &Dictionary) -> Vec<(FieldRef, ArrayRef)> {
vec![none_typed("id", dict_def.id_type.clone(), 0), none_typed("value", dict_def.value_type.clone(), 0)]
}
#[instrument(level = "trace", skip_all, name = "volcano::scan::dictionary::assemble")]
fn assemble(ids: &[DictionaryEntryId], values: &[Value], dict_def: &Dictionary) -> Result<Option<RecordBatch>> {
let id_column = build_id_column(ids, dict_def.id_type.clone())?;
let value_column = build_value_column(values, dict_def.value_type.clone());
Ok(Some(batch(vec![id_column, value_column])?))
}
}
impl QueryNode for DictionaryScanNode {
#[instrument(name = "volcano::scan::dictionary::initialize", level = "trace", skip_all)]
fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
Ok(())
}
#[instrument(name = "volcano::scan::dictionary::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(), "DictionaryScan::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 dict_def = self.dictionary.def();
let range = DictionaryEntryIndexKey::full_scan(dict_def.id).resume_after(self.last_key.as_ref());
let (ids, values, new_last_key) = Self::drain_batch(rx, range, batch_size, dict_def)?;
if ids.is_empty() {
self.exhausted = true;
if self.last_key.is_none() {
return Ok(Some(batch(Self::empty_columns(dict_def))?));
}
return Ok(None);
}
self.last_key = new_last_key;
Self::assemble(&ids, &values, dict_def)
}
fn headers(&self) -> Option<ColumnHeaders> {
Some(self.headers.clone())
}
}
fn build_id_column(ids: &[DictionaryEntryId], id_type: ValueType) -> Result<(FieldRef, ArrayRef)> {
let data = match id_type {
ValueType::Uint1 => {
let vals: Vec<u8> = ids.iter().map(|id| id.to_u128() as u8).collect();
uint1("id", vals)
}
ValueType::Uint2 => {
let vals: Vec<u16> = ids.iter().map(|id| id.to_u128() as u16).collect();
uint2("id", vals)
}
ValueType::Uint4 => {
let vals: Vec<u32> = ids.iter().map(|id| id.to_u128() as u32).collect();
uint4("id", vals)
}
ValueType::Uint8 => {
let vals: Vec<u64> = ids.iter().map(|id| id.to_u128() as u64).collect();
uint8("id", vals)
}
ValueType::Uint16 => {
let vals: Vec<u128> = ids.iter().map(|id| id.to_u128()).collect();
uint16("id", vals)
}
_ => return Err(internal_error!("Invalid dictionary id_type: {:?}", id_type)),
};
Ok(data)
}
fn build_value_column(values: &[Value], value_type: ValueType) -> (FieldRef, ArrayRef) {
let mut data = ColumnBuilder::with_capacity(value_type, values.len());
for value in values {
data.push_value(value.clone());
}
data.finish("value")
}