reifydb_engine/vm/volcano/scan/

dictionary.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

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,
	})
}