reifydb_engine/vm/volcano/
sort.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::cmp::Ordering::Equal;
5
6use reifydb_core::{
7	error::diagnostic::query,
8	sort::{
9		SortDirection::{Asc, Desc},
10		SortKey,
11	},
12	value::column::{buffer::ColumnBuffer, columns::Columns, headers::ColumnHeaders},
13};
14use reifydb_extension::transform::{Transform, context::TransformContext};
15use reifydb_transaction::transaction::Transaction;
16use reifydb_value::{
17	error,
18	error::Error,
19	reifydb_assertions,
20	util::cowvec::CowVec,
21	value::{
22		datetime::{CREATED_AT_COLUMN_NAME, UPDATED_AT_COLUMN_NAME},
23		row_number::ROW_NUMBER_COLUMN_NAME,
24	},
25};
26use tracing::instrument;
27
28use crate::{
29	Result,
30	vm::volcano::query::{QueryContext, QueryNode},
31};
32
33pub(crate) struct SortNode {
34	input: Box<dyn QueryNode>,
35	by: Vec<SortKey>,
36	initialized: Option<()>,
37}
38
39impl SortNode {
40	pub(crate) fn new(input: Box<dyn QueryNode>, by: Vec<SortKey>) -> Self {
41		Self {
42			input,
43			by,
44			initialized: None,
45		}
46	}
47}
48
49impl QueryNode for SortNode {
50	#[instrument(level = "trace", skip_all, name = "volcano::sort::initialize")]
51	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
52		self.input.initialize(rx, ctx)?;
53		self.initialized = Some(());
54		Ok(())
55	}
56
57	#[instrument(level = "trace", skip_all, name = "volcano::sort::next")]
58	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<Columns>> {
59		reifydb_assertions! {
60			assert!(self.initialized.is_some(), "SortNode::next() called before initialize()");
61		}
62
63		let mut columns_opt: Option<Columns> = None;
64
65		while let Some(columns) = self.input.next(rx, ctx)? {
66			if let Some(existing_columns) = &mut columns_opt {
67				existing_columns.row_numbers.make_mut().extend(columns.row_numbers.iter().copied());
68				existing_columns.created_at.make_mut().extend(columns.created_at.iter().copied());
69				existing_columns.updated_at.make_mut().extend(columns.updated_at.iter().copied());
70				for (i, col) in columns.columns.iter().enumerate() {
71					existing_columns[i].extend(col.clone())?;
72				}
73			} else {
74				columns_opt = Some(columns);
75			}
76		}
77
78		let columns = match columns_opt {
79			Some(f) => f,
80			None => return Ok(None),
81		};
82
83		let transform_ctx = TransformContext {
84			routines: &ctx.services.routines,
85			runtime_context: &ctx.services.runtime_context,
86			params: &ctx.params,
87		};
88		Ok(Some(self.apply(&transform_ctx, columns)?))
89	}
90
91	fn headers(&self) -> Option<ColumnHeaders> {
92		self.input.headers()
93	}
94}
95
96impl Transform for SortNode {
97	fn apply(&self, _ctx: &TransformContext, mut columns: Columns) -> Result<Columns> {
98		let key_refs = self
99			.by
100			.iter()
101			.map(|key| {
102				let name = key.column.fragment();
103				let stripped = name.strip_prefix('#').unwrap_or(name);
104
105				if stripped == ROW_NUMBER_COLUMN_NAME && !columns.row_numbers.is_empty() {
106					let data: Vec<u64> = columns.row_numbers.iter().map(|r| r.value()).collect();
107					return Ok::<_, Error>((ColumnBuffer::uint8(data), key.direction.clone()));
108				}
109				if stripped == CREATED_AT_COLUMN_NAME && !columns.created_at.is_empty() {
110					return Ok((
111						ColumnBuffer::datetime(columns.created_at.to_vec()),
112						key.direction.clone(),
113					));
114				}
115				if stripped == UPDATED_AT_COLUMN_NAME && !columns.updated_at.is_empty() {
116					return Ok((
117						ColumnBuffer::datetime(columns.updated_at.to_vec()),
118						key.direction.clone(),
119					));
120				}
121
122				let col = columns
123					.iter()
124					.find(|c| c.name() == name)
125					.ok_or_else(|| error!(query::column_not_found(key.column.clone())))?;
126				Ok((col.data().clone(), key.direction.clone()))
127			})
128			.collect::<Result<Vec<_>>>()?;
129
130		let row_count = columns.row_count();
131		let mut indices: Vec<usize> = (0..row_count).collect();
132
133		indices.sort_unstable_by(|&l, &r| {
134			for (col, dir) in &key_refs {
135				let vl = col.get_value(l);
136				let vr = col.get_value(r);
137				let ord = vl.partial_cmp(&vr).unwrap_or(Equal);
138				let ord = match dir {
139					Asc => ord,
140					Desc => ord.reverse(),
141				};
142				if ord != Equal {
143					return ord;
144				}
145			}
146			Equal
147		});
148
149		if !columns.row_numbers.is_empty() {
150			let reordered: Vec<_> = indices.iter().map(|&i| columns.row_numbers[i]).collect();
151			columns.row_numbers = CowVec::new(reordered);
152		}
153		if !columns.created_at.is_empty() {
154			let reordered: Vec<_> = indices.iter().map(|&i| columns.created_at[i]).collect();
155			columns.created_at = CowVec::new(reordered);
156		}
157		if !columns.updated_at.is_empty() {
158			let reordered: Vec<_> = indices.iter().map(|&i| columns.updated_at[i]).collect();
159			columns.updated_at = CowVec::new(reordered);
160		}
161
162		let cols = columns.columns.make_mut();
163		for col in cols.iter_mut() {
164			col.reorder(&indices);
165		}
166
167		Ok(columns)
168	}
169}
reifydb_engine/vm/volcano/sort.rs

reifydb_engine/vm/volcano/
sort.rs
