reifydb_engine/vm/volcano/scan/

table.rs

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

use std::{collections::VecDeque, ops::Bound, sync::Arc};

use arrow_array::{ArrayRef, RecordBatch, UInt64Array};
use arrow_schema::FieldRef;
use reifydb_codec::row::{
	bytes::EncodedBytes,
	shape::{RowShape, fingerprint::RowShapeFingerprint},
	table::EncodedTableRow,
};
use reifydb_core::{
	common::TimeSource,
	error::diagnostic,
	interface::{
		catalog::{dictionary::Dictionary, storage::StorageId},
		resolved::ResolvedTable,
		store::MultiVersionRow,
	},
	key::{
		any::TaggedKey,
		row::{PartitionedRowKey, RowKeyRange, StoragePartitionedRowKey, StorageRowKey},
	},
	value::{
		batch::{append_rows_fitted, batch},
		column::{builder::ColumnBuilder, headers::ColumnHeaders},
	},
};
use reifydb_transaction::{multi::RangeScope, transaction::Transaction};
use reifydb_value::{
	error, reifydb_assertions,
	value::{
		partition::Partition,
		row_number::RowNumber,
		system_columns::{SystemColumn, stamp_system_columns},
		value_type::ValueType,
	},
};
use tracing::instrument;

use super::{
	super::decode_dictionary_columns,
	empty_scan,
	merge::{MergeLayout, PartitionMerge},
	partition_array, scan_headers, source_system_columns, storage_partitioned_row, storage_row,
};
use crate::{
	Result,
	vm::volcano::query::{QueryContext, QueryNode},
};

pub struct TableScanNode {
	table: ResolvedTable,
	context: Option<Arc<QueryContext>>,
	headers: ColumnHeaders,

	storage_types: Vec<ValueType>,

	dictionaries: Vec<Option<(Dictionary, ValueType)>>,

	shape: Option<RowShape>,
	resume: Resume,
	exhausted: bool,

	partition: Option<Partition>,

	pending: Option<VecDeque<Partition>>,

	advanced: bool,

	system_columns: Vec<SystemColumn>,

	oldest_first: bool,
	keep_dictionary_ids: bool,

	merge: Option<PartitionMerge>,

	storage_batch: Option<RecordBatch>,
}

impl TableScanNode {
	pub fn new(
		table: ResolvedTable,
		partition: Option<Partition>,
		context: Arc<QueryContext>,
		rx: &mut Transaction<'_>,
	) -> Result<Self> {
		let mut storage_types = Vec::with_capacity(table.columns().len());
		let mut dictionaries = Vec::with_capacity(table.columns().len());

		for col in table.columns() {
			if let Some(dict_id) = col.dictionary_id {
				if let Some(dict) = context.services.catalog.find_dictionary(rx, dict_id)? {
					storage_types.push(ValueType::DictionaryId);
					dictionaries.push(Some((dict, col.constraint.get_type())));
				} else {
					storage_types.push(col.constraint.get_type());
					dictionaries.push(None);
				}
			} else {
				storage_types.push(col.constraint.get_type());
				dictionaries.push(None);
			}
		}

		let system_columns = source_system_columns(
			!table.def().partition_by.is_empty(),
			table.def().time != TimeSource::None,
			true,
		);
		let headers = scan_headers(table.columns().iter().map(|col| col.name.as_str()), &system_columns);

		let resume = if table.def().partition_by.is_empty() {
			Resume::Row(None)
		} else {
			Resume::Partitioned(None)
		};

		Ok(Self {
			table,
			context: Some(context),
			headers,
			storage_types,
			dictionaries,
			shape: None,
			resume,
			exhausted: false,
			partition,
			pending: None,
			advanced: false,
			system_columns,
			oldest_first: false,
			keep_dictionary_ids: false,
			merge: None,
			storage_batch: None,
		})
	}

	pub fn with_partitions(
		table: ResolvedTable,
		partitions: Vec<Partition>,
		context: Arc<QueryContext>,
		rx: &mut Transaction<'_>,
	) -> Result<Self> {
		let mut pending = VecDeque::from(partitions);
		let first = pending.pop_front();
		let mut node = Self::new(table, first, context, rx)?;
		node.pending = Some(pending);
		Ok(node)
	}

	fn scan_partition_list<'a>(
		&mut self,
		rx: &mut Transaction<'a>,
		storage: StorageId,
		scope: RangeScope,
		batch_size: u64,
		mut last: Option<StoragePartitionedRowKey>,
	) -> Result<(ScannedBatch, Option<StoragePartitionedRowKey>)> {
		let mut batch = ScannedBatch::default();
		loop {
			let Some(partition) = self.partition else {
				batch.exhausted = true;
				return Ok((batch, None));
			};
			let room = batch_size - batch.rows.len() as u64;
			let (start, end) = partitioned_bounds(Some(partition), last);
			let (scanned, new_last) = {
				let mut stream = rx.range_partitioned_row(storage, start, end, scope, room as usize)?;
				Self::drain_batch_partitioned(&mut stream, room)?
			};
			let partition_done = scanned.exhausted;
			batch.rows.extend(scanned.rows);
			batch.row_numbers.extend(scanned.row_numbers);
			batch.partitions.extend(scanned.partitions);
			batch.commit_versions.extend(scanned.commit_versions);
			if !partition_done {
				return Ok((batch, new_last));
			}
			self.partition = self.pending.as_mut().and_then(VecDeque::pop_front);
			self.advanced = true;
			last = None;
			if batch.rows.len() as u64 >= batch_size {
				batch.exhausted = self.partition.is_none();
				return Ok((batch, 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
	}

	fn get_or_load_shape<'a>(
		&mut self,
		rx: &mut Transaction<'a>,
		fingerprint: RowShapeFingerprint,
	) -> Result<RowShape> {
		if let Some(shape) = &self.shape
			&& shape.fingerprint() == fingerprint
		{
			return Ok(shape.clone());
		}

		let stored_ctx = self.context.as_ref().expect("TableScanNode context not set");
		let shape = stored_ctx.services.catalog.get_or_load_row_shape(fingerprint, rx)?.ok_or_else(|| {
			error!(diagnostic::internal::internal(format!(
				"RowShape with fingerprint {:?} not found for table {}",
				fingerprint,
				self.table.def().name
			)))
		})?;

		self.shape = Some(shape.clone());

		Ok(shape)
	}

	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::drain_partitioned")]
	fn drain_batch_partitioned(
		stream: &mut dyn Iterator<Item = Result<MultiVersionRow<StoragePartitionedRowKey>>>,
		batch_size: u64,
	) -> Result<(ScannedBatch, Option<StoragePartitionedRowKey>)> {
		let mut batch = ScannedBatch::default();
		let mut last = None;

		for _ in 0..batch_size {
			match stream.next() {
				Some(Ok(multi)) => {
					batch.rows.push(multi.bytes);
					batch.row_numbers.push(multi.key.row());
					batch.partitions.push(multi.key.partition());
					batch.commit_versions.push(multi.version.0);
					last = Some(multi.key);
				}
				Some(Err(e)) => return Err(e),
				None => {
					batch.exhausted = true;
					break;
				}
			}
		}

		Ok((batch, last))
	}

	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::drain_row")]
	fn drain_batch_row(
		stream: &mut dyn Iterator<Item = Result<MultiVersionRow<StorageRowKey>>>,
		batch_size: u64,
	) -> Result<(ScannedBatch, Option<StorageRowKey>)> {
		let mut batch = ScannedBatch::default();
		let mut last = None;

		for _ in 0..batch_size {
			match stream.next() {
				Some(Ok(multi)) => {
					batch.rows.push(multi.bytes);
					batch.row_numbers.push(multi.key.row());
					batch.commit_versions.push(multi.version.0);
					last = Some(multi.key);
				}
				Some(Err(e)) => return Err(e),
				None => {
					batch.exhausted = true;
					break;
				}
			}
		}

		Ok((batch, last))
	}

	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::column_alloc")]
	fn storage_columns(&self) -> Vec<(FieldRef, ArrayRef)> {
		self.table
			.columns()
			.iter()
			.enumerate()
			.map(|(idx, col)| {
				ColumnBuilder::with_capacity(self.storage_types[idx].clone(), 0).finish(&col.name)
			})
			.collect()
	}

	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::empty_columns")]
	fn empty_columns(&self) -> Vec<(FieldRef, ArrayRef)> {
		self.table
			.columns()
			.iter()
			.map(|col| ColumnBuilder::with_capacity(col.constraint.get_type(), 0).finish(&col.name))
			.collect()
	}

	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::append_rows")]
	fn append_batch<'a>(
		&mut self,
		rx: &mut Transaction<'a>,
		columns: RecordBatch,
		bytes_vec: Vec<EncodedBytes>,
		row_numbers: Vec<RowNumber>,
	) -> Result<RecordBatch> {
		let mut columns = columns;
		let mut rows = bytes_vec.into_iter().peekable();
		let mut numbers = row_numbers.into_iter();
		while let Some(first) = rows.next() {
			let fingerprint = EncodedTableRow::view(&first).fingerprint();
			let mut run = vec![first];
			while let Some(next) =
				rows.next_if(|row| EncodedTableRow::view(row).fingerprint() == fingerprint)
			{
				run.push(next);
			}
			let run_numbers: Vec<RowNumber> = numbers.by_ref().take(run.len()).collect();
			let shape = self.get_or_load_shape(rx, fingerprint)?;
			columns = append_rows_fitted(columns, &shape, run, run_numbers)?;
		}
		Ok(columns)
	}
}

#[derive(Clone, Copy)]
enum Resume {
	Row(Option<StorageRowKey>),
	Partitioned(Option<StoragePartitionedRowKey>),
}

#[derive(Default)]
struct ScannedBatch {
	rows: Vec<EncodedBytes>,
	row_numbers: Vec<RowNumber>,
	partitions: Vec<Partition>,
	commit_versions: Vec<u64>,
	exhausted: bool,
}

fn partitioned_bounds(
	partition: Option<Partition>,
	last: Option<StoragePartitionedRowKey>,
) -> (Bound<StoragePartitionedRowKey>, Bound<StoragePartitionedRowKey>) {
	let start = match (last, partition) {
		(Some(key), _) => Bound::Excluded(key),
		(None, Some(partition)) => {
			Bound::Included(StoragePartitionedRowKey::new(partition, RowNumber(u64::MAX)))
		}
		(None, None) => Bound::Unbounded,
	};
	let end = match partition {
		Some(partition) => Bound::Included(StoragePartitionedRowKey::new(partition, RowNumber(u64::MIN))),
		None => Bound::Unbounded,
	};
	(start, end)
}

impl QueryNode for TableScanNode {
	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::initialize")]
	fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
		Ok(())
	}

	#[instrument(level = "trace", skip_all, name = "volcano::scan::table::next")]
	fn next<'a>(&mut self, rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
		reifydb_assertions! {
			assert!(self.context.is_some(), "TableScanNode::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 scope = RangeScope::All;

		let storage: StorageId = self.table.def().id.into();

		let (scanned, next_resume, resumed) = match self.resume {
			Resume::Partitioned(last) if self.pending.is_some() => {
				let resumed = last.is_some() || self.advanced;
				let (scanned, new_last) =
					self.scan_partition_list(rx, storage, scope, batch_size, last)?;
				(scanned, Resume::Partitioned(new_last), resumed)
			}
			Resume::Partitioned(last) if self.oldest_first => {
				let partition = self.partition;
				let merge = self.merge.get_or_insert_with(|| {
					PartitionMerge::new(
						MergeLayout::Row,
						storage,
						partition.map(|partition| {
							PartitionedRowKey::partition_range(storage, partition)
						}),
					)
				});
				let rows = merge.next(rx, batch_size)?;
				let (scanned, new_last) = Self::drain_batch_partitioned(
					&mut rows.into_iter().map(storage_partitioned_row),
					batch_size,
				)?;
				(scanned, Resume::Partitioned(new_last), last.is_some())
			}
			Resume::Partitioned(last) => {
				let (start, end) = partitioned_bounds(self.partition, last);
				let (scanned, new_last) = {
					let mut stream = rx.range_partitioned_row(
						storage,
						start,
						end,
						scope,
						batch_size as usize,
					)?;
					Self::drain_batch_partitioned(&mut stream, batch_size)?
				};
				(scanned, Resume::Partitioned(new_last), last.is_some())
			}
			Resume::Row(last) if self.oldest_first => {
				let last_key = last.map(|key| TaggedKey::Row(key.with_storage(storage)));
				let range = RowKeyRange::scan_range_rev(storage, last_key.as_ref());
				let (scanned, new_last) = {
					let mut stream = rx
						.range_rev(range, scope, batch_size as usize)?
						.map(|row| row.and_then(storage_row));
					Self::drain_batch_row(&mut stream, batch_size)?
				};
				(scanned, Resume::Row(new_last), last.is_some())
			}
			Resume::Row(last) => {
				let start = match last {
					Some(key) => Bound::Excluded(key),
					None => Bound::Unbounded,
				};
				let (scanned, new_last) = {
					let mut stream = rx.range_row(
						storage,
						start,
						Bound::Unbounded,
						scope,
						batch_size as usize,
					)?;
					Self::drain_batch_row(&mut stream, batch_size)?
				};
				(scanned, Resume::Row(new_last), last.is_some())
			}
		};

		if scanned.exhausted {
			self.exhausted = true;
		}

		if scanned.rows.is_empty() {
			self.exhausted = true;
			if !resumed {
				return Ok(Some(empty_scan(self.empty_columns(), &self.system_columns)?));
			}
			return Ok(None);
		}

		self.resume = next_resume;

		let columns = match &self.storage_batch {
			Some(storage) => storage.clone(),
			None => {
				let storage = batch(self.storage_columns())?;
				self.storage_batch = Some(storage.clone());
				storage
			}
		};
		let columns = self.append_batch(rx, columns, scanned.rows, scanned.row_numbers)?;

		let mut stamps: Vec<(SystemColumn, ArrayRef)> = Vec::new();
		if !scanned.partitions.is_empty() {
			stamps.push((SystemColumn::Partitions, partition_array(&scanned.partitions)));
		}
		stamps.push((SystemColumn::CommitVersion, Arc::new(UInt64Array::from(scanned.commit_versions))));
		let columns = stamp_system_columns(columns, stamps)?;

		if self.keep_dictionary_ids {
			return Ok(Some(columns));
		}
		Ok(Some(decode_dictionary_columns(columns, &self.dictionaries, rx)?))
	}

	fn headers(&self) -> Option<ColumnHeaders> {
		Some(self.headers.clone())
	}
}