1pub mod chunk_helpers;
10pub mod graph_source;
11pub mod join_helpers;
12pub mod mapper;
13pub mod plan_serializer;
14pub mod projection_helper;
15#[cfg(test)]
16mod tests;
17#[cfg(test)]
18mod tests_only;
19pub mod union_helpers;
20
21pub use chunk_helpers::*;
22pub use graph_source::*;
23pub use join_helpers::*;
24pub use mapper::*;
25pub use plan_serializer::*;
26pub use projection_helper::*;
27pub use union_helpers::*;
28
29use crate::physical_operator::*;
30use akar_common::error::ProcessorError;
31use akar_common::types::{PhysicalTypeID, Value};
32use akar_common::vector::{DataChunk, ValueVector};
33use akar_function::registry::{FunctionRegistry, TableFunction};
34use akar_planner::logical_operator::LogicalOperator;
35use akar_storage::table::TableCatalog;
36use akar_transaction::UndoRecord;
37use std::sync::{Arc, Mutex};
38
39pub type SequenceFn = Arc<dyn Fn(&str, bool) -> Result<Value, ProcessorError> + Send + Sync>;
40pub type SubqueryFn = Arc<dyn Fn(&akar_parser::ast::Query) -> Result<Vec<DataChunk>, ProcessorError> + Send + Sync>;
41
42#[derive(Debug, Clone)]
46pub enum SchemaDdlOp {
47 CreateSequence {
48 name: String,
49 if_not_exists: bool,
50 start_value: i64,
51 increment: i64,
52 min_value: i64,
53 max_value: i64,
54 cycle: bool,
55 },
56 DropSequence {
57 name: String,
58 if_exists: bool,
59 },
60 ExportDatabase {
61 file_path: String,
62 file_type: String,
63 schema_only: bool,
64 },
65 ImportDatabase {
66 file_path: String,
67 query: String,
68 index_query: String,
69 },
70}
71pub type SchemaDdlFn = Arc<dyn Fn(SchemaDdlOp) -> Result<String, ProcessorError> + Send + Sync>;
72
73pub trait StandaloneCallHandler: Send + Sync {
74 fn execute_call(
75 &self,
76 name: &str,
77 args: &[akar_parser::ast::Expression],
78 ) -> Result<Vec<akar_common::vector::DataChunk>, ProcessorError>;
79}
80
81pub trait StandaloneCallFn: Send + Sync {
82 fn execute(
83 &self,
84 args: &[akar_parser::ast::Expression],
85 ) -> Result<Vec<Vec<akar_common::types::Value>>, ProcessorError>;
86 fn aliases(&self) -> Vec<&'static str>;
87}
88
89#[derive(Default)]
90pub struct StandaloneCallRegistry {
91 handlers: std::collections::HashMap<String, std::sync::Arc<dyn StandaloneCallFn>>,
92}
93
94impl StandaloneCallRegistry {
95 pub fn new() -> Self {
96 Self::default()
97 }
98
99 pub fn register(&mut self, handler: std::sync::Arc<dyn StandaloneCallFn>) {
100 for alias in handler.aliases() {
101 self.handlers.insert(alias.to_lowercase(), handler.clone());
102 }
103 }
104
105 pub fn get(&self, name: &str) -> Option<std::sync::Arc<dyn StandaloneCallFn>> {
106 self.handlers.get(&name.to_lowercase()).cloned()
107 }
108}
109
110pub struct QueryProcessor {
112 function_registry: Option<Arc<Mutex<FunctionRegistry>>>,
113 table_catalog: Option<Arc<TableCatalog>>,
114 vfs: Option<Arc<akar_common::file_system::VirtualFileSystemRegistry>>,
115 standalone_call_handler: Option<Arc<dyn StandaloneCallHandler>>,
116 sequence_fn: Option<SequenceFn>,
119 subquery_fn: Option<SubqueryFn>,
121 schema_ddl_fn: Option<SchemaDdlFn>,
123 snapshot_ts: Option<u64>,
125 commit_history: Vec<(u64, u64)>,
127 written_rows: Mutex<Vec<(u64, u64)>>,
131 txn_id: Option<u64>,
134 undo_records: Arc<Mutex<Vec<UndoRecord>>>,
138}
139
140impl QueryProcessor {
141 pub fn new() -> Self {
142 Self {
143 function_registry: None,
144 table_catalog: None,
145 vfs: None,
146 standalone_call_handler: None,
147 sequence_fn: None,
148 subquery_fn: None,
149 schema_ddl_fn: None,
150 snapshot_ts: None,
151 commit_history: Vec::new(),
152 written_rows: Mutex::new(Vec::new()),
153 txn_id: None,
154 undo_records: Arc::new(Mutex::new(Vec::new())),
155 }
156 }
157
158 pub fn with_registry(registry: Arc<Mutex<FunctionRegistry>>) -> Self {
160 Self {
161 function_registry: Some(registry),
162 table_catalog: None,
163 vfs: None,
164 standalone_call_handler: None,
165 sequence_fn: None,
166 subquery_fn: None,
167 schema_ddl_fn: None,
168 snapshot_ts: None,
169 commit_history: Vec::new(),
170 written_rows: Mutex::new(Vec::new()),
171 txn_id: None,
172 undo_records: Arc::new(Mutex::new(Vec::new())),
173 }
174 }
175
176 pub fn with_catalog(
178 registry: Arc<Mutex<FunctionRegistry>>,
179 table_catalog: Arc<TableCatalog>,
180 vfs: Arc<akar_common::file_system::VirtualFileSystemRegistry>,
181 ) -> Self {
182 Self {
183 function_registry: Some(registry),
184 table_catalog: Some(table_catalog),
185 vfs: Some(vfs),
186 standalone_call_handler: None,
187 sequence_fn: None,
188 subquery_fn: None,
189 schema_ddl_fn: None,
190 snapshot_ts: None,
191 commit_history: Vec::new(),
192 written_rows: Mutex::new(Vec::new()),
193 txn_id: None,
194 undo_records: Arc::new(Mutex::new(Vec::new())),
195 }
196 }
197
198 pub fn with_standalone_call_handler(mut self, handler: Arc<dyn StandaloneCallHandler>) -> Self {
201 self.standalone_call_handler = Some(handler);
202 self
203 }
204
205 pub fn with_sequence_fn(mut self, f: SequenceFn) -> Self {
206 self.sequence_fn = Some(f);
207 self
208 }
209
210 pub fn with_subquery_fn(mut self, f: SubqueryFn) -> Self {
212 self.subquery_fn = Some(f);
213 self
214 }
215
216 pub fn with_schema_ddl_fn(mut self, f: SchemaDdlFn) -> Self {
218 self.schema_ddl_fn = Some(f);
219 self
220 }
221
222 pub fn with_snapshot(mut self, snapshot_ts: Option<u64>, commit_history: Vec<(u64, u64)>) -> Self {
224 self.snapshot_ts = snapshot_ts;
225 self.commit_history = commit_history;
226 self
227 }
228
229 pub fn with_txn_id(mut self, txn_id: Option<u64>) -> Self {
235 self.txn_id = txn_id;
236 self
237 }
238
239 pub fn record_insert_undo(&self, table_id: u64, row_id: u64) {
241 if let Ok(mut u) = self.undo_records.lock() {
242 u.push(UndoRecord::insert(table_id, row_id));
243 }
244 }
245
246 pub fn record_update_undo(&self, table_id: u64, row_id: u64, column: u32, old_data: Vec<u8>) {
248 if let Ok(mut u) = self.undo_records.lock() {
249 u.push(UndoRecord::update(table_id, row_id, column, old_data));
250 }
251 }
252
253 pub fn record_delete_undo(&self, table_id: u64, row_id: u64, old_data: Vec<u8>) {
255 if let Ok(mut u) = self.undo_records.lock() {
256 u.push(UndoRecord::delete(table_id, row_id, old_data));
257 }
258 }
259
260 pub fn take_undo_records(&self) -> Vec<UndoRecord> {
262 self.undo_records
263 .lock()
264 .map(|mut u| std::mem::take(&mut *u))
265 .unwrap_or_default()
266 }
267
268 pub fn undo_sink(&self) -> Arc<Mutex<Vec<UndoRecord>>> {
271 Arc::clone(&self.undo_records)
272 }
273
274 pub fn execute(&self, operators: &[LogicalOperator]) -> Result<Vec<DataChunk>, ProcessorError> {
276 self.execute_internal(operators)
277 }
278
279 pub fn execute_internal(&self, operators: &[LogicalOperator]) -> Result<Vec<DataChunk>, ProcessorError> {
280 if operators.is_empty() {
281 return Ok(vec![DataChunk {
282 fields: vec![],
283 field_types: vec![],
284 size: 0,
285 field_names: vec![],
286 sel_vector: None,
287 }]);
288 }
289
290 let mut intermediate_result: Option<Vec<DataChunk>> = None;
291
292 for (i, op) in operators.iter().enumerate() {
293 let current = intermediate_result.take().unwrap_or_else(|| {
294 let mut dummy = DataChunk::new(vec![], vec![]);
295 dummy.size = 1;
296 vec![dummy]
297 });
298 let next_op = operators.get(i + 1);
299
300 let mut ctx = mapper::ExecutionContext {
301 processor: self,
302 function_registry: self.function_registry.clone(),
303 table_catalog: self.table_catalog.clone(),
304 vfs: self.vfs.clone(),
305 standalone_call_handler: self.standalone_call_handler.clone(),
306 sequence_fn: self.sequence_fn.clone(),
307 subquery_fn: self.subquery_fn.clone(),
308 schema_ddl_fn: self.schema_ddl_fn.clone(),
309 snapshot_ts: self.snapshot_ts,
310 commit_history: self.commit_history.clone(),
311 written_rows: Vec::new(),
312 txn_id: self.txn_id,
313 };
314
315 let result = mapper::PlanMapper::map_and_execute(op, next_op, current, &mut ctx)?;
316
317 if !ctx.written_rows.is_empty() {
319 if let Ok(mut writes) = self.written_rows.lock() {
320 writes.append(&mut ctx.written_rows);
321 }
322 }
323
324 if let LogicalOperator::ScanRel(_) = op {
325 match &mut intermediate_result {
327 Some(existing) => existing.extend(result),
328 None => intermediate_result = Some(result),
329 }
330 } else {
331 intermediate_result = Some(result);
332 }
333 }
334
335 Ok(intermediate_result.unwrap_or_default())
336 }
337
338 pub fn take_written_rows(&self) -> Vec<(u64, u64)> {
341 self.written_rows
342 .lock()
343 .map(|mut w| std::mem::take(&mut *w))
344 .unwrap_or_default()
345 }
346
347 fn execute_table_function(
350 &self,
351 tf: &akar_planner::logical_operator::LogicalTableFunctionCall,
352 ) -> Result<Vec<DataChunk>, ProcessorError> {
353 let func_name = &tf.function_name;
354 let args: Vec<Value> = Vec::new(); if let Some(ref registry) = self.function_registry {
358 let reg = registry.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
359 if let Some(tbl_fn) = reg.get_table(func_name) {
360 match tbl_fn {
361 TableFunction::CustomTable { execute, .. } => {
362 let mut chunk = DataChunk::new(Vec::new(), Vec::new());
363 (execute)(&args, &mut chunk)?;
364 Ok(vec![chunk])
365 }
366 TableFunction::CustomTableWithGraph { execute, .. } => {
367 let mut chunk = DataChunk::new(Vec::new(), Vec::new());
368 let graph = CatalogGraphSource::new(self.table_catalog.as_ref());
369 (execute)(&args, Some(&graph), &mut chunk)?;
370 Ok(vec![chunk])
371 }
372 TableFunction::ScanCsv { .. }
373 | TableFunction::ScanParquet { .. }
374 | TableFunction::ScanJson { .. }
375 | TableFunction::ListTables
376 | TableFunction::ShowColumns { .. }
377 | TableFunction::CurrentSetting { .. } => Err(format!(
378 "Table function '{}' cannot be executed dynamically (no callback)",
379 func_name
380 )
381 .into()),
382 TableFunction::Custom { name } if name == "vector_similarity_scan" => {
383 drop(reg);
387 self.execute_vector_similarity_scan(tf)
388 }
389 TableFunction::Custom { name } => {
390 Err(format!("Custom table function '{}' has no registered handler", name).into())
391 }
392 }
393 } else {
394 Err(format!("Table function '{}' not found", func_name).into())
395 }
396 } else {
397 Err(format!(
398 "Cannot execute table function '{}': no function registry available",
399 func_name
400 )
401 .into())
402 }
403 }
404
405 fn execute_vector_similarity_scan(
410 &self,
411 tf: &akar_planner::logical_operator::LogicalTableFunctionCall,
412 ) -> Result<Vec<DataChunk>, ProcessorError> {
413 if tf.args.len() < 4 {
415 return Err(
416 "vector_similarity_scan requires 4 arguments: table_name, column_name, query_vector, top_k".into(),
417 );
418 }
419
420 fn eval_expr_to_value(expr: &akar_parser::ast::Expression) -> Option<Value> {
423 match expr {
424 akar_parser::ast::Expression::Constant(c) => match c {
425 akar_parser::ast::Constant::String(s) => Some(Value::String(s.clone())),
426 akar_parser::ast::Constant::Integer(i) => Some(Value::Int64(*i)),
427 akar_parser::ast::Constant::Float(f) => Some(Value::Double(*f)),
428 akar_parser::ast::Constant::Bool(b) => Some(Value::Bool(*b)),
429 akar_parser::ast::Constant::Null => Some(Value::Null),
430 },
431 akar_parser::ast::Expression::List(items) => {
432 let vals: Vec<Value> = items.iter().filter_map(eval_expr_to_value).collect();
433 Some(Value::List(vals))
434 }
435 _ => None, }
437 }
438
439 let table_name = match eval_expr_to_value(&tf.args[0]) {
440 Some(Value::String(s)) => s,
441 _ => return Err("First argument to vector_similarity_scan must be a table name string".into()),
442 };
443
444 let _column_name = match eval_expr_to_value(&tf.args[1]) {
445 Some(Value::String(s)) => s,
446 _ => return Err("Second argument to vector_similarity_scan must be a column name string".into()),
447 };
448
449 let query_vector = match eval_expr_to_value(&tf.args[2]) {
450 Some(Value::List(items)) => {
451 let mut vec = Vec::with_capacity(items.len());
452 for item in &items {
453 match item {
454 Value::Double(d) => vec.push(*d),
455 Value::Int64(i) => vec.push(*i as f64),
456 Value::Int32(i) => vec.push(*i as f64),
457 Value::Float(f) => vec.push(*f as f64),
458 _ => return Err("query_vector must be a list of numbers".into()),
459 }
460 }
461 vec
462 }
463 _ => return Err("Third argument to vector_similarity_scan must be a list of numbers".into()),
464 };
465
466 let top_k = match eval_expr_to_value(&tf.args[3]) {
467 Some(Value::Int64(k)) if k > 0 => k as u64,
468 _ => return Err("Fourth argument to vector_similarity_scan must be a positive integer".into()),
469 };
470
471 let tc = self
473 .table_catalog
474 .clone()
475 .ok_or_else(|| "No table catalog available for vector_similarity_scan".to_string())?;
476
477 let index_name = {
479 let mut found = None;
480 for entry in tc.all_vector_indexes() {
481 if entry.table_name == table_name {
482 found = Some(entry.name.clone());
483 break;
484 }
485 }
486 found.ok_or_else(|| format!("No vector index found on table '{}'", table_name))?
487 };
488
489 let scan = PhysicalVectorSimilarityScan {
491 index_name,
492 index_id: 0,
493 query_vector,
494 top_k,
495 table_name,
496 table_catalog: Some(tc),
497 };
498 scan.execute(vec![])
499 }
500
501 pub fn evaluate_expression(
503 _expr: &akar_parser::ast::Expression,
504 _chunk: &DataChunk,
505 ) -> Result<ValueVector, ProcessorError> {
506 let size = _chunk.size;
508 let mut v = ValueVector::new(PhysicalTypeID::Int64, size);
509 for i in 0..size {
510 v.set_i64(i, 0);
511 }
512 v.resize(size);
513 Ok(v)
514 }
515}
516
517impl Default for QueryProcessor {
518 fn default() -> Self {
519 Self::new()
520 }
521}