use std::sync::Arc;
use arrow_array::RecordBatch;
use reifydb_core::{
error::diagnostic::{operation, query},
interface::{
catalog::config::{ConfigKey, GetConfig},
resolved::ResolvedObject,
},
sort::SortKey,
util::budget::MemoryBudget,
value::{batch::empty_batch, column::headers::ColumnHeaders},
};
use reifydb_evaluate::{expression::context::EvalContext, stack::SymbolTable};
use reifydb_extension::transform::context::TransformContext;
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
byte_size::ByteSize,
error,
params::Params,
reifydb_assertions,
value::{column_view::ColumnView, identity::IdentityId, system_columns::check_user_columns},
};
use crate::{Result, vm::services::Services};
pub fn query_budget(services: &Services) -> Arc<MemoryBudget> {
let limit = services.catalog.get_config_uint8(ConfigKey::QueryMemoryLimit);
Arc::new(MemoryBudget::new(ByteSize::from_bytes(limit)))
}
pub(crate) fn ensure_sort_key_orderable(key: &SortKey, data: &ColumnView<'_>) -> Result<()> {
let ty = data.get_type();
if ty.is_scalar() {
Ok(())
} else {
Err(error!(operation::sort_key_not_orderable(key.column.clone(), ty)))
}
}
pub(crate) fn charge_query_memory_bytes(budget: &MemoryBudget, charged: &mut usize, total: usize) -> Result<()> {
if total > *charged {
let delta = (total - *charged) as u64;
if !budget.try_charge(ByteSize::from_bytes(delta)) {
return Err(error!(query::memory_limit_exceeded(budget.used(), budget.limit())));
}
*charged = total;
}
Ok(())
}
pub trait QueryNode: Send + Sync {
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()>;
fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>>;
fn headers(&self) -> Option<ColumnHeaders>;
}
#[derive(Clone)]
pub struct QueryContext {
pub services: Arc<Services>,
pub source: Option<ResolvedObject>,
pub batch_size: u64,
pub params: Params,
pub symbols: SymbolTable,
pub identity: IdentityId,
pub memory: Arc<MemoryBudget>,
}
impl QueryNode for Box<dyn QueryNode> {
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
(**self).initialize(rx, ctx)
}
fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
let result = (**self).next(rx, ctx)?;
if let Some(ref batch) = result {
check_user_columns(batch)?;
reifydb_assertions! {
reifydb_value::value::canonical::assert_canonical_floats(batch, "volcano next");
}
}
Ok(result)
}
fn headers(&self) -> Option<ColumnHeaders> {
(**self).headers()
}
}
pub fn eval_context_from_query<'a>(ctx: &'a QueryContext) -> EvalContext<'a> {
EvalContext {
target: None,
batch: empty_batch(),
row_count: 1,
take: None,
params: &ctx.params,
symbols: &ctx.symbols,
routines: &ctx.services.routines,
runtime_context: &ctx.services.runtime_context,
identity: ctx.identity,
}
}
pub fn eval_context_from_transform<'a>(ctx: &'a TransformContext<'a>, stored: &'a QueryContext) -> EvalContext<'a> {
EvalContext {
target: None,
batch: empty_batch(),
row_count: 1,
take: None,
params: ctx.params,
symbols: &stored.symbols,
routines: &stored.services.routines,
runtime_context: ctx.runtime_context,
identity: stored.identity,
}
}
#[cfg(test)]
mod tests {
#[cfg(reifydb_assertions)]
use std::sync::Arc;
#[cfg(reifydb_assertions)]
use arrow_array::{ArrayRef, Float64Array, RecordBatch};
#[cfg(reifydb_assertions)]
use reifydb_core::value::column::headers::ColumnHeaders;
use reifydb_core::{
util::budget::MemoryBudget,
value::{
batch::{batch, data_size},
column::factory::int4,
},
};
#[cfg(reifydb_assertions)]
use reifydb_evaluate::stack::SymbolTable;
#[cfg(reifydb_assertions)]
use reifydb_test_harness::engine::create_test_admin_transaction;
#[cfg(reifydb_assertions)]
use reifydb_transaction::transaction::Transaction;
use reifydb_value::byte_size::ByteSize;
#[cfg(reifydb_assertions)]
use reifydb_value::{params::Params, value::identity::IdentityId};
use super::charge_query_memory_bytes;
#[cfg(reifydb_assertions)]
use super::{QueryContext, QueryNode, query_budget};
#[cfg(reifydb_assertions)]
use crate::{Result, vm::services::Services};
#[test]
fn charge_query_memory_delta_charges_and_rejects_over_budget() {
let budget = MemoryBudget::new(ByteSize::from_kib(1));
let mut charged = 0usize;
let small = batch(vec![int4("c", [1i32, 2, 3, 4])]).expect("one int4 column");
let small_size = data_size(&small).expect("an int4 batch has a data size");
charge_query_memory_bytes(&budget, &mut charged, small_size).expect("small buffer fits under 1 KiB");
let after_first = budget.used().as_bytes();
assert!(after_first > 0, "charging a non-empty buffer must consume budget");
assert_eq!(charged as u64, after_first, "charged must track exactly what the budget recorded");
charge_query_memory_bytes(&budget, &mut charged, small_size)
.expect("re-charge of the same buffer is free");
assert_eq!(
budget.used().as_bytes(),
after_first,
"delta charging must not double count an unchanged buffer"
);
let big = batch(vec![int4("c", 0..4000i32)]).expect("one int4 column");
let mut big_charged = 0usize;
let big_size = data_size(&big).expect("an int4 batch has a data size");
let err = charge_query_memory_bytes(&budget, &mut big_charged, big_size).unwrap_err();
assert_eq!(err.0.code, "QUERY_006", "an over-budget charge must raise the memory-limit diagnostic");
}
#[cfg(reifydb_assertions)]
struct NegativeZeroNode;
#[cfg(reifydb_assertions)]
impl QueryNode for NegativeZeroNode {
fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
Ok(())
}
fn next<'a>(
&mut self,
_rx: &mut Transaction<'a>,
_ctx: &mut QueryContext,
) -> Result<Option<RecordBatch>> {
let column: ArrayRef = Arc::new(Float64Array::from(vec![-0.0f64]));
Ok(Some(RecordBatch::try_from_iter([("c", column)]).unwrap()))
}
fn headers(&self) -> Option<ColumnHeaders> {
None
}
}
#[test]
#[cfg(reifydb_assertions)]
#[should_panic(expected = "is not canonical")]
fn test_box_next_with_negative_zero_panics() {
let services = Services::testing();
let mut txn = create_test_admin_transaction();
let mut ctx = QueryContext {
services: services.clone(),
source: None,
batch_size: 1,
params: Params::default(),
symbols: SymbolTable::new(),
identity: IdentityId::system(),
memory: query_budget(&services),
};
let mut node: Box<dyn QueryNode> = Box::new(NegativeZeroNode);
let _ = <Box<dyn QueryNode> as QueryNode>::next(&mut node, &mut Transaction::Admin(&mut txn), &mut ctx);
}
}