use std::{ops::Bound, sync::Arc};
use postcard::from_bytes;
use reifydb_codec::key::encoded::{EncodedKey, EncodedKeyRange};
use reifydb_core::{
interface::{catalog::dictionary::Dictionary, resolved::ResolvedDictionary, store::SingleVersionRange},
internal_error,
key::{EncodableKey, dictionary::DictionaryEntryIndexKey},
value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns, 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<EncodedKey>,
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: EncodedKeyRange,
batch_size: u64,
dict_def: &Dictionary,
) -> Result<(Vec<DictionaryEntryId>, Vec<Value>, Option<EncodedKey>)> {
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, batch_size)?;
for entry in batch.items {
new_last_key = Some(entry.key.clone());
if let Some(key) = DictionaryEntryIndexKey::decode(&entry.key) {
let entry_id = DictionaryEntryId::from_u128(key.id, dict_def.id_type.clone())?;
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<ColumnWithName> {
vec![
ColumnWithName {
name: Fragment::internal("id"),
data: ColumnBuffer::none_typed(dict_def.id_type.clone(), 0),
},
ColumnWithName {
name: Fragment::internal("value"),
data: ColumnBuffer::none_typed(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<Columns>> {
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(Columns::new(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<Columns>> {
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 full_scan = DictionaryEntryIndexKey::full_scan(dict_def.id);
let range = match &self.last_key {
None => full_scan,
Some(last) => EncodedKeyRange::new(Bound::Excluded(last.clone()), full_scan.end),
};
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(Columns::new(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<ColumnWithName> {
let data = match id_type {
ValueType::Uint1 => {
let vals: Vec<u8> = ids.iter().map(|id| id.to_u128() as u8).collect();
ColumnBuffer::uint1(vals)
}
ValueType::Uint2 => {
let vals: Vec<u16> = ids.iter().map(|id| id.to_u128() as u16).collect();
ColumnBuffer::uint2(vals)
}
ValueType::Uint4 => {
let vals: Vec<u32> = ids.iter().map(|id| id.to_u128() as u32).collect();
ColumnBuffer::uint4(vals)
}
ValueType::Uint8 => {
let vals: Vec<u64> = ids.iter().map(|id| id.to_u128() as u64).collect();
ColumnBuffer::uint8(vals)
}
ValueType::Uint16 => {
let vals: Vec<u128> = ids.iter().map(|id| id.to_u128()).collect();
ColumnBuffer::uint16(vals)
}
_ => return Err(internal_error!("Invalid dictionary id_type: {:?}", id_type)),
};
Ok(ColumnWithName {
name: Fragment::internal("id"),
data,
})
}
fn build_value_column(values: &[Value], value_type: ValueType) -> Result<ColumnWithName> {
let data = match value_type {
ValueType::Utf8 => {
let vals: Vec<String> = values
.iter()
.map(|v| match v {
Value::Utf8(s) => s.clone(),
_ => format!("{:?}", v),
})
.collect();
ColumnBuffer::utf8(vals)
}
ValueType::Int1 => {
let vals: Vec<i8> = values
.iter()
.map(|v| match v {
Value::Int1(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::int1(vals)
}
ValueType::Int2 => {
let vals: Vec<i16> = values
.iter()
.map(|v| match v {
Value::Int2(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::int2(vals)
}
ValueType::Int4 => {
let vals: Vec<i32> = values
.iter()
.map(|v| match v {
Value::Int4(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::int4(vals)
}
ValueType::Int8 => {
let vals: Vec<i64> = values
.iter()
.map(|v| match v {
Value::Int8(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::int8(vals)
}
ValueType::Uint1 => {
let vals: Vec<u8> = values
.iter()
.map(|v| match v {
Value::Uint1(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::uint1(vals)
}
ValueType::Uint2 => {
let vals: Vec<u16> = values
.iter()
.map(|v| match v {
Value::Uint2(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::uint2(vals)
}
ValueType::Uint4 => {
let vals: Vec<u32> = values
.iter()
.map(|v| match v {
Value::Uint4(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::uint4(vals)
}
ValueType::Uint8 => {
let vals: Vec<u64> = values
.iter()
.map(|v| match v {
Value::Uint8(n) => *n,
_ => 0,
})
.collect();
ColumnBuffer::uint8(vals)
}
_ => {
let vals: Vec<String> = values.iter().map(|v| format!("{:?}", v)).collect();
ColumnBuffer::utf8(vals)
}
};
Ok(ColumnWithName {
name: Fragment::internal("value"),
data,
})
}