reifydb_evaluate/expression/
lookup.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer};
5use reifydb_rql::expression::ColumnExpression;
6use reifydb_value::value::{
7	Value,
8	blob::Blob,
9	date::Date,
10	datetime::DateTime,
11	decimal::Decimal,
12	dictionary::DictionaryEntryId,
13	duration::Duration,
14	identity::IdentityId,
15	int::Int,
16	time::Time,
17	uint::Uint,
18	uuid::{Uuid4, Uuid7},
19	value_type::ValueType,
20};
21
22use crate::{Result, expression::context::EvalContext, stack::Variable};
23
24macro_rules! extract_typed_column {
25	($col:expr, $take:expr, $variant:ident($x:ident) => $transform:expr, $default:expr, $constructor:ident) => {{
26		let mut data = Vec::new();
27		let mut bitvec = Vec::new();
28		let mut count = 0;
29		for v in $col.data().iter() {
30			if count >= $take {
31				break;
32			}
33			match v {
34				Value::$variant($x) => {
35					data.push($transform);
36					bitvec.push(true);
37				}
38				_ => {
39					data.push($default);
40					bitvec.push(false);
41				}
42			}
43			count += 1;
44		}
45		Ok($col.with_new_data(ColumnBuffer::$constructor(data, bitvec)))
46	}};
47}
48
49pub(crate) fn column_lookup(ctx: &EvalContext, column: &ColumnExpression) -> Result<ColumnWithName> {
50	let name = column.0.name.text();
51
52	if let Some(data) = ctx.columns.system_column(name) {
53		return Ok(ColumnWithName::new(name.to_string(), data));
54	}
55
56	if let Some(col) = ctx.columns.iter().find(|c| c.name() == name) {
57		let owned = ColumnWithName::new(col.name().clone(), col.data().clone());
58		return extract_column_data(&owned, ctx);
59	}
60
61	if let Some(Variable::Columns {
62		columns: scalar_cols,
63	}) = ctx.symbols.get(name)
64		&& scalar_cols.is_scalar()
65		&& let Some(col) = scalar_cols.columns.first()
66	{
67		let owned = ColumnWithName::new(scalar_cols.name_at(0).clone(), col.clone());
68		return extract_column_data(&owned, ctx);
69	}
70
71	Ok(ColumnWithName::new(name.to_string(), ColumnBuffer::none_typed(ValueType::Boolean, ctx.row_count)))
72}
73
74fn extract_column_data(col: &ColumnWithName, ctx: &EvalContext) -> Result<ColumnWithName> {
75	let take = ctx.take.unwrap_or(usize::MAX);
76
77	if take >= col.data().len() {
78		return Ok(col.clone());
79	}
80
81	let col_type = col.data().get_type();
82	let effective_type = match col_type {
83		ValueType::Option(inner) => *inner,
84		other => other,
85	};
86
87	extract_column_data_by_type(col, take, effective_type)
88}
89
90fn extract_column_data_by_type(col: &ColumnWithName, take: usize, col_type: ValueType) -> Result<ColumnWithName> {
91	match col_type {
92		ValueType::Boolean => extract_typed_column!(col, take, Boolean(b) => b, false, bool_with_bitvec),
93		ValueType::Float4 => {
94			extract_typed_column!(col, take, Float4(v) => v.value(), 0.0f32, float4_with_bitvec)
95		}
96		ValueType::Float8 => {
97			extract_typed_column!(col, take, Float8(v) => v.value(), 0.0f64, float8_with_bitvec)
98		}
99		ValueType::Int1 => extract_typed_column!(col, take, Int1(n) => n, 0, int1_with_bitvec),
100		ValueType::Int2 => extract_typed_column!(col, take, Int2(n) => n, 0, int2_with_bitvec),
101		ValueType::Int4 => extract_typed_column!(col, take, Int4(n) => n, 0, int4_with_bitvec),
102		ValueType::Int8 => extract_typed_column!(col, take, Int8(n) => n, 0, int8_with_bitvec),
103		ValueType::Int16 => extract_typed_column!(col, take, Int16(n) => n, 0, int16_with_bitvec),
104		ValueType::Utf8 => {
105			extract_typed_column!(col, take, Utf8(s) => s.clone(), "".to_string(), utf8_with_bitvec)
106		}
107		ValueType::Uint1 => extract_typed_column!(col, take, Uint1(n) => n, 0, uint1_with_bitvec),
108		ValueType::Uint2 => extract_typed_column!(col, take, Uint2(n) => n, 0, uint2_with_bitvec),
109		ValueType::Uint4 => extract_typed_column!(col, take, Uint4(n) => n, 0, uint4_with_bitvec),
110		ValueType::Uint8 => extract_typed_column!(col, take, Uint8(n) => n, 0, uint8_with_bitvec),
111		ValueType::Uint16 => extract_typed_column!(col, take, Uint16(n) => n, 0, uint16_with_bitvec),
112		ValueType::Date => extract_typed_column!(col, take, Date(d) => d, Date::default(), date_with_bitvec),
113		ValueType::DateTime => {
114			extract_typed_column!(col, take, DateTime(dt) => dt, DateTime::default(), datetime_with_bitvec)
115		}
116		ValueType::Time => extract_typed_column!(col, take, Time(t) => t, Time::default(), time_with_bitvec),
117		ValueType::Duration => {
118			extract_typed_column!(col, take, Duration(i) => i, Duration::default(), duration_with_bitvec)
119		}
120		ValueType::IdentityId => {
121			extract_typed_column!(col, take, IdentityId(i) => i, IdentityId::default(), identity_id_with_bitvec)
122		}
123		ValueType::Uuid4 => {
124			extract_typed_column!(col, take, Uuid4(i) => i, Uuid4::default(), uuid4_with_bitvec)
125		}
126		ValueType::Uuid7 => {
127			extract_typed_column!(col, take, Uuid7(i) => i, Uuid7::default(), uuid7_with_bitvec)
128		}
129		ValueType::DictionaryId => {
130			extract_typed_column!(col, take, DictionaryId(i) => i, DictionaryEntryId::default(), dictionary_id_with_bitvec)
131		}
132		ValueType::Blob => {
133			extract_typed_column!(col, take, Blob(b) => b.clone(), Blob::new(vec![]), blob_with_bitvec)
134		}
135		ValueType::Int => extract_typed_column!(col, take, Int(b) => b.clone(), Int::zero(), int_with_bitvec),
136		ValueType::Uint => {
137			extract_typed_column!(col, take, Uint(b) => b.clone(), Uint::zero(), uint_with_bitvec)
138		}
139		ValueType::Any => {
140			extract_typed_column!(col, take, Any(boxed) => *boxed.clone(), Value::none(), any_with_bitvec)
141		}
142		ValueType::Decimal => {
143			extract_typed_column!(col, take, Decimal(b) => b.clone(), Decimal::from_i64(0), decimal_with_bitvec)
144		}
145		ValueType::Option(inner) => extract_column_data_by_type(col, take, *inner),
146		ValueType::List(_) => {
147			extract_typed_column!(col, take, Any(boxed) => *boxed.clone(), Value::none(), any_with_bitvec)
148		}
149		ValueType::Record(_) => {
150			extract_typed_column!(col, take, Any(boxed) => *boxed.clone(), Value::none(), any_with_bitvec)
151		}
152		ValueType::Tuple(_) => {
153			extract_typed_column!(col, take, Any(boxed) => *boxed.clone(), Value::none(), any_with_bitvec)
154		}
155	}
156}
157
158#[cfg(test)]
159pub mod tests {
160	use reifydb_core::{
161		interface::identifier::{ColumnIdentifier, ColumnObject},
162		value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns},
163	};
164	use reifydb_routine_abi::registry::Routines;
165	use reifydb_rql::expression::ColumnExpression;
166	use reifydb_runtime::context::{RuntimeContext, clock::Clock};
167	use reifydb_value::{fragment::Fragment, params::Params, value::identity::IdentityId};
168
169	use super::column_lookup;
170	use crate::{expression::context::EvalContext, stack::SymbolTable};
171
172	#[test]
173	fn test_column_not_found_returns_correct_row_count() {
174		let columns = Columns::new(vec![ColumnWithName::new(
175			"existing_col".to_string(),
176			ColumnBuffer::int4([1, 2, 3, 4, 5]),
177		)]);
178
179		let runtime_ctx = RuntimeContext::with_clock(Clock::Real);
180		let routines = Routines::empty();
181		let base = EvalContext {
182			params: &Params::None,
183			symbols: &SymbolTable::new(),
184			routines: &routines,
185			runtime_context: &runtime_ctx,
186			identity: IdentityId::root(),
187			is_aggregate_context: false,
188			columns: Columns::empty(),
189			row_count: 1,
190			target: None,
191			take: None,
192		};
193		let ctx = base.with_eval(columns, 5);
194
195		// A missing column resolves to all-none, and its length must match the other columns
196		// or every downstream zip would misalign.
197		let result = column_lookup(
198			&ctx,
199			&ColumnExpression(ColumnIdentifier {
200				object: ColumnObject::Alias(Fragment::internal("nonexistent_col")),
201				name: Fragment::internal("nonexistent_col"),
202			}),
203		)
204		.unwrap();
205
206		assert_eq!(
207			result.data().len(),
208			5,
209			"Column not found should return column with ctx.row_count rows, not 0"
210		);
211	}
212}
reifydb_evaluate/expression/lookup.rs

reifydb_evaluate/expression/
lookup.rs
