akar_processor/processor/
mod.rs

1//! Query processor — maps logical operators to physical operators and executes them.
2//!
3//! Pipeline execution model:
4//! 1. Scan operators produce raw DataChunks
5//! 2. Filter removes non-matching rows
6//! 3. Projection selects/transforms columns
7//! 4. Limit/OrderBy/Aggregate are applied last
8
9pub 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_storage::wal::{WALRecord, WalSink};
37use akar_transaction::UndoRecord;
38use std::sync::{Arc, Mutex};
39
40pub type SequenceFn = Arc<dyn Fn(&str, bool) -> Result<Value, ProcessorError> + Send + Sync>;
41pub type SubqueryFn = Arc<dyn Fn(&akar_parser::ast::Query) -> Result<Vec<DataChunk>, ProcessorError> + Send + Sync>;
42
43/// DDL operations that require the schema-level Catalog (from Akar-catalog).
44/// These are dispatched via callback because the processor layer doesn't
45/// directly own the schema catalog.
46#[derive(Debug, Clone)]
47pub enum SchemaDdlOp {
48    CreateSequence {
49        name: String,
50        if_not_exists: bool,
51        start_value: i64,
52        increment: i64,
53        min_value: i64,
54        max_value: i64,
55        cycle: bool,
56    },
57    DropSequence {
58        name: String,
59        if_exists: bool,
60    },
61    ExportDatabase {
62        file_path: String,
63        file_type: String,
64        schema_only: bool,
65    },
66    ImportDatabase {
67        file_path: String,
68        query: String,
69        index_query: String,
70    },
71}
72pub type SchemaDdlFn = Arc<dyn Fn(SchemaDdlOp) -> Result<String, ProcessorError> + Send + Sync>;
73
74pub trait StandaloneCallHandler: Send + Sync {
75    fn execute_call(
76        &self,
77        name: &str,
78        args: &[akar_parser::ast::Expression],
79    ) -> Result<Vec<akar_common::vector::DataChunk>, ProcessorError>;
80}
81
82pub trait StandaloneCallFn: Send + Sync {
83    fn execute(
84        &self,
85        args: &[akar_parser::ast::Expression],
86    ) -> Result<Vec<Vec<akar_common::types::Value>>, ProcessorError>;
87    fn aliases(&self) -> Vec<&'static str>;
88}
89
90#[derive(Default)]
91pub struct StandaloneCallRegistry {
92    handlers: std::collections::HashMap<String, std::sync::Arc<dyn StandaloneCallFn>>,
93}
94
95impl StandaloneCallRegistry {
96    pub fn new() -> Self {
97        Self::default()
98    }
99
100    pub fn register(&mut self, handler: std::sync::Arc<dyn StandaloneCallFn>) {
101        for alias in handler.aliases() {
102            self.handlers.insert(alias.to_lowercase(), handler.clone());
103        }
104    }
105
106    pub fn get(&self, name: &str) -> Option<std::sync::Arc<dyn StandaloneCallFn>> {
107        self.handlers.get(&name.to_lowercase()).cloned()
108    }
109}
110
111/// The query processor executes a physical plan and produces result chunks.
112pub struct QueryProcessor {
113    function_registry: Option<Arc<Mutex<FunctionRegistry>>>,
114    table_catalog: Option<Arc<TableCatalog>>,
115    vfs: Option<Arc<akar_common::file_system::VirtualFileSystemRegistry>>,
116    standalone_call_handler: Option<Arc<dyn StandaloneCallHandler>>,
117    /// Callback for sequence operations (nextval/currval).
118    /// Takes (sequence_name, is_nextval) and returns the resulting value.
119    sequence_fn: Option<SequenceFn>,
120    /// Callback for executing subqueries.
121    subquery_fn: Option<SubqueryFn>,
122    /// Callback for schema-level DDL operations (CREATE/DROP SEQUENCE, EXPORT/IMPORT DATABASE).
123    schema_ddl_fn: Option<SchemaDdlFn>,
124    /// MVCC snapshot timestamp for read isolation.
125    snapshot_ts: Option<u64>,
126    /// Commit history for MVCC visibility checks.
127    commit_history: Vec<(u64, u64)>,
128    /// Row-level write set accumulated during execution.
129    /// Populated by the mapper after each write operation (SET, DELETE, INSERT).
130    /// Read by the connection layer after execution for OCC conflict detection.
131    written_rows: Mutex<Vec<(u64, u64)>>,
132    /// Active transaction id threaded into write operators so inserts/deletes
133    /// are recorded in `VersionInfo` for MVCC snapshot isolation (P52.18).
134    txn_id: Option<u64>,
135    /// Undo records captured by write operators during execution. Drained into
136    /// the active transaction by the connection layer so a rollback (including
137    /// an OCC conflict loser) can revert the in-place table writes (P52.18).
138    undo_records: Arc<Mutex<Vec<UndoRecord>>>,
139    /// Typed WAL records captured by SQL write operators during execution.
140    /// Drained into the transaction's `LocalWAL` by the connection layer and
141    /// bulk-copied into the global WAL at commit, making WAL replay
142    /// self-sufficient without per-commit mirror persistence (P60.2).
143    wal_records: WalSink,
144}
145
146impl QueryProcessor {
147    pub fn new() -> Self {
148        Self {
149            function_registry: None,
150            table_catalog: None,
151            vfs: None,
152            standalone_call_handler: None,
153            sequence_fn: None,
154            subquery_fn: None,
155            schema_ddl_fn: None,
156            snapshot_ts: None,
157            commit_history: Vec::new(),
158            written_rows: Mutex::new(Vec::new()),
159            txn_id: None,
160            undo_records: Arc::new(Mutex::new(Vec::new())),
161            wal_records: Arc::new(Mutex::new(Vec::new())),
162        }
163    }
164
165    /// Create a processor with access to the function registry.
166    pub fn with_registry(registry: Arc<Mutex<FunctionRegistry>>) -> Self {
167        Self {
168            function_registry: Some(registry),
169            table_catalog: None,
170            vfs: None,
171            standalone_call_handler: None,
172            sequence_fn: None,
173            subquery_fn: None,
174            schema_ddl_fn: None,
175            snapshot_ts: None,
176            commit_history: Vec::new(),
177            written_rows: Mutex::new(Vec::new()),
178            txn_id: None,
179            undo_records: Arc::new(Mutex::new(Vec::new())),
180            wal_records: Arc::new(Mutex::new(Vec::new())),
181        }
182    }
183
184    /// Create a processor with function registry, table catalog access, and VFS.
185    pub fn with_catalog(
186        registry: Arc<Mutex<FunctionRegistry>>,
187        table_catalog: Arc<TableCatalog>,
188        vfs: Arc<akar_common::file_system::VirtualFileSystemRegistry>,
189    ) -> Self {
190        Self {
191            function_registry: Some(registry),
192            table_catalog: Some(table_catalog),
193            vfs: Some(vfs),
194            standalone_call_handler: None,
195            sequence_fn: None,
196            subquery_fn: None,
197            schema_ddl_fn: None,
198            snapshot_ts: None,
199            commit_history: Vec::new(),
200            written_rows: Mutex::new(Vec::new()),
201            txn_id: None,
202            undo_records: Arc::new(Mutex::new(Vec::new())),
203            wal_records: Arc::new(Mutex::new(Vec::new())),
204        }
205    }
206
207    /// Set the sequence operation callback (for nextval/currval).
208
209    pub fn with_standalone_call_handler(mut self, handler: Arc<dyn StandaloneCallHandler>) -> Self {
210        self.standalone_call_handler = Some(handler);
211        self
212    }
213
214    pub fn with_sequence_fn(mut self, f: SequenceFn) -> Self {
215        self.sequence_fn = Some(f);
216        self
217    }
218
219    /// Set the subquery operation callback.
220    pub fn with_subquery_fn(mut self, f: SubqueryFn) -> Self {
221        self.subquery_fn = Some(f);
222        self
223    }
224
225    /// Set the schema DDL callback (for CREATE/DROP SEQUENCE, EXPORT/IMPORT DATABASE).
226    pub fn with_schema_ddl_fn(mut self, f: SchemaDdlFn) -> Self {
227        self.schema_ddl_fn = Some(f);
228        self
229    }
230
231    /// Set MVCC snapshot parameters for read isolation.
232    pub fn with_snapshot(mut self, snapshot_ts: Option<u64>, commit_history: Vec<(u64, u64)>) -> Self {
233        self.snapshot_ts = snapshot_ts;
234        self.commit_history = commit_history;
235        self
236    }
237
238    /// Set the active transaction id for write operators.
239    ///
240    /// When `Some(txn_id)`, insert/delete write operators use the MVCC-aware
241    /// `*_with_txn` storage variants so uncommitted rows are invisible to other
242    /// snapshots, and record undo entries for rollback (P52.18).
243    pub fn with_txn_id(mut self, txn_id: Option<u64>) -> Self {
244        self.txn_id = txn_id;
245        self
246    }
247
248    /// Record an insert undo (rollback deletes the row).
249    pub fn record_insert_undo(&self, table_id: u64, row_id: u64) {
250        if let Ok(mut u) = self.undo_records.lock() {
251            u.push(UndoRecord::insert(table_id, row_id));
252        }
253    }
254
255    /// Record an update undo (rollback restores the cell).
256    pub fn record_update_undo(&self, table_id: u64, row_id: u64, column: u32, old_data: Vec<u8>) {
257        if let Ok(mut u) = self.undo_records.lock() {
258            u.push(UndoRecord::update(table_id, row_id, column, old_data));
259        }
260    }
261
262    /// Record a delete undo (rollback restores the row/edge).
263    pub fn record_delete_undo(&self, table_id: u64, row_id: u64, old_data: Vec<u8>) {
264        if let Ok(mut u) = self.undo_records.lock() {
265            u.push(UndoRecord::delete(table_id, row_id, old_data));
266        }
267    }
268
269    /// Take the accumulated undo records (drained by the connection layer).
270    pub fn take_undo_records(&self) -> Vec<UndoRecord> {
271        self.undo_records
272            .lock()
273            .map(|mut u| std::mem::take(&mut *u))
274            .unwrap_or_default()
275    }
276
277    /// Shared undo sink handed to write operators so they can record undo
278    /// entries during execution (P52.18).
279    pub fn undo_sink(&self) -> Arc<Mutex<Vec<UndoRecord>>> {
280        Arc::clone(&self.undo_records)
281    }
282
283    /// Shared WAL sink handed to write operators so they can log typed
284    /// Insert/Delete/Update records during execution (P60.2).
285    pub fn wal_sink(&self) -> WalSink {
286        Arc::clone(&self.wal_records)
287    }
288
289    /// Take the accumulated typed WAL records (drained by the connection
290    /// layer after execution; empty when nothing was logged).
291    pub fn take_wal_records(&self) -> Vec<WALRecord> {
292        self.wal_records
293            .lock()
294            .map(|mut w| std::mem::take(&mut *w))
295            .unwrap_or_default()
296    }
297
298    /// Execute a sequence of logical operators by mapping them to physical operators.
299    pub fn execute(&self, operators: &[LogicalOperator]) -> Result<Vec<DataChunk>, ProcessorError> {
300        self.execute_internal(operators)
301    }
302
303    pub fn execute_internal(&self, operators: &[LogicalOperator]) -> Result<Vec<DataChunk>, ProcessorError> {
304        if operators.is_empty() {
305            return Ok(vec![DataChunk {
306                fields: vec![],
307                field_types: vec![],
308                size: 0,
309                field_names: vec![],
310                sel_vector: None,
311            }]);
312        }
313
314        let mut intermediate_result: Option<Vec<DataChunk>> = None;
315
316        for (i, op) in operators.iter().enumerate() {
317            let current = intermediate_result.take().unwrap_or_else(|| {
318                let mut dummy = DataChunk::new(vec![], vec![]);
319                dummy.size = 1;
320                vec![dummy]
321            });
322            let next_op = operators.get(i + 1);
323
324            let mut ctx = mapper::ExecutionContext {
325                processor: self,
326                function_registry: self.function_registry.clone(),
327                table_catalog: self.table_catalog.clone(),
328                vfs: self.vfs.clone(),
329                standalone_call_handler: self.standalone_call_handler.clone(),
330                sequence_fn: self.sequence_fn.clone(),
331                subquery_fn: self.subquery_fn.clone(),
332                schema_ddl_fn: self.schema_ddl_fn.clone(),
333                snapshot_ts: self.snapshot_ts,
334                commit_history: self.commit_history.clone(),
335                written_rows: Vec::new(),
336                txn_id: self.txn_id,
337            };
338
339            let result = mapper::PlanMapper::map_and_execute(op, next_op, current, &mut ctx)?;
340
341            // Accumulate written rows from this operator into the processor's write set
342            if !ctx.written_rows.is_empty() {
343                if let Ok(mut writes) = self.written_rows.lock() {
344                    writes.append(&mut ctx.written_rows);
345                }
346            }
347
348            if let LogicalOperator::ScanRel(_) = op {
349                // Accumulate: extend rather than replace for ScanRel
350                match &mut intermediate_result {
351                    Some(existing) => existing.extend(result),
352                    None => intermediate_result = Some(result),
353                }
354            } else {
355                intermediate_result = Some(result);
356            }
357        }
358
359        Ok(intermediate_result.unwrap_or_default())
360    }
361
362    /// Take the accumulated write set from the processor.
363    /// Returns all (table_id, row_id) pairs written during the last execution.
364    pub fn take_written_rows(&self) -> Vec<(u64, u64)> {
365        self.written_rows
366            .lock()
367            .map(|mut w| std::mem::take(&mut *w))
368            .unwrap_or_default()
369    }
370
371    /// Execute a table function call by looking up the function in the registry
372    /// and dispatching to the appropriate handler.
373    fn execute_table_function(
374        &self,
375        tf: &akar_planner::logical_operator::LogicalTableFunctionCall,
376    ) -> Result<Vec<DataChunk>, ProcessorError> {
377        let func_name = &tf.function_name;
378        let args: Vec<Value> = Vec::new(); // args would be evaluated from expressions
379
380        // Look up the function in the registry
381        if let Some(ref registry) = self.function_registry {
382            let reg = registry.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
383            if let Some(tbl_fn) = reg.get_table(func_name) {
384                match tbl_fn {
385                    TableFunction::CustomTable { execute, .. } => {
386                        let mut chunk = DataChunk::new(Vec::new(), Vec::new());
387                        (execute)(&args, &mut chunk)?;
388                        Ok(vec![chunk])
389                    }
390                    TableFunction::CustomTableWithGraph { execute, .. } => {
391                        let mut chunk = DataChunk::new(Vec::new(), Vec::new());
392                        let graph = CatalogGraphSource::new(self.table_catalog.as_ref());
393                        (execute)(&args, Some(&graph), &mut chunk)?;
394                        Ok(vec![chunk])
395                    }
396                    TableFunction::ScanCsv { .. }
397                    | TableFunction::ScanParquet { .. }
398                    | TableFunction::ScanJson { .. }
399                    | TableFunction::ListTables
400                    | TableFunction::ShowColumns { .. }
401                    | TableFunction::CurrentSetting { .. } => Err(format!(
402                        "Table function '{}' cannot be executed dynamically (no callback)",
403                        func_name
404                    )
405                    .into()),
406                    TableFunction::Custom { name } if name == "vector_similarity_scan" => {
407                        // Evaluate args: [table_name, column_name, query_vector, top_k]
408                        // For CALL statement, args are parsed as expressions. We need to evaluate them.
409                        // For now, parse from the function args
410                        drop(reg);
411                        self.execute_vector_similarity_scan(tf)
412                    }
413                    TableFunction::Custom { name } => {
414                        Err(format!("Custom table function '{}' has no registered handler", name).into())
415                    }
416                }
417            } else {
418                Err(format!("Table function '{}' not found", func_name).into())
419            }
420        } else {
421            Err(format!(
422                "Cannot execute table function '{}': no function registry available",
423                func_name
424            )
425            .into())
426        }
427    }
428
429    /// Execute a `vector_similarity_scan` table function call.
430    ///
431    /// Expects CALL vector_similarity_scan(table_name, column_name, query_vector, top_k)
432    /// and dispatches to PhysicalVectorSimilarityScan with the processor's TableCatalog.
433    fn execute_vector_similarity_scan(
434        &self,
435        tf: &akar_planner::logical_operator::LogicalTableFunctionCall,
436    ) -> Result<Vec<DataChunk>, ProcessorError> {
437        // Evaluate arguments from expressions (they should be constants or simple vars)
438        if tf.args.len() < 4 {
439            return Err(
440                "vector_similarity_scan requires 4 arguments: table_name, column_name, query_vector, top_k".into(),
441            );
442        }
443
444        // For CALL statements, args arrive as Expression AST nodes.
445        // Evaluate them to Values. The simplest approach: evaluate constants inline.
446        fn eval_expr_to_value(expr: &akar_parser::ast::Expression) -> Option<Value> {
447            match expr {
448                akar_parser::ast::Expression::Constant(c) => match c {
449                    akar_parser::ast::Constant::String(s) => Some(Value::String(s.clone())),
450                    akar_parser::ast::Constant::Integer(i) => Some(Value::Int64(*i)),
451                    akar_parser::ast::Constant::Float(f) => Some(Value::Double(*f)),
452                    akar_parser::ast::Constant::Bool(b) => Some(Value::Bool(*b)),
453                    akar_parser::ast::Constant::Null => Some(Value::Null),
454                },
455                akar_parser::ast::Expression::List(items) => {
456                    let vals: Vec<Value> = items.iter().filter_map(eval_expr_to_value).collect();
457                    Some(Value::List(vals))
458                }
459                _ => None, // Non-constant expression — skip
460            }
461        }
462
463        let table_name = match eval_expr_to_value(&tf.args[0]) {
464            Some(Value::String(s)) => s,
465            _ => return Err("First argument to vector_similarity_scan must be a table name string".into()),
466        };
467
468        let _column_name = match eval_expr_to_value(&tf.args[1]) {
469            Some(Value::String(s)) => s,
470            _ => return Err("Second argument to vector_similarity_scan must be a column name string".into()),
471        };
472
473        let query_vector = match eval_expr_to_value(&tf.args[2]) {
474            Some(Value::List(items)) => {
475                let mut vec = Vec::with_capacity(items.len());
476                for item in &items {
477                    match item {
478                        Value::Double(d) => vec.push(*d),
479                        Value::Int64(i) => vec.push(*i as f64),
480                        Value::Int32(i) => vec.push(*i as f64),
481                        Value::Float(f) => vec.push(*f as f64),
482                        _ => return Err("query_vector must be a list of numbers".into()),
483                    }
484                }
485                vec
486            }
487            _ => return Err("Third argument to vector_similarity_scan must be a list of numbers".into()),
488        };
489
490        let top_k = match eval_expr_to_value(&tf.args[3]) {
491            Some(Value::Int64(k)) if k > 0 => k as u64,
492            _ => return Err("Fourth argument to vector_similarity_scan must be a positive integer".into()),
493        };
494
495        // Find the vector index on this table
496        let tc = self
497            .table_catalog
498            .clone()
499            .ok_or_else(|| "No table catalog available for vector_similarity_scan".to_string())?;
500
501        // Look for a vector index matching this table name
502        let index_name = {
503            let mut found = None;
504            for entry in tc.all_vector_indexes() {
505                if entry.table_name == table_name {
506                    found = Some(entry.name.clone());
507                    break;
508                }
509            }
510            found.ok_or_else(|| format!("No vector index found on table '{}'", table_name))?
511        };
512
513        // Dispatch to PhysicalVectorSimilarityScan
514        let scan = PhysicalVectorSimilarityScan {
515            index_name,
516            index_id: 0,
517            query_vector,
518            top_k,
519            table_name,
520            table_catalog: Some(tc),
521        };
522        scan.execute(vec![])
523    }
524
525    /// Execute a single expression against a DataChunk and return a ValueVector of results.
526    pub fn evaluate_expression(
527        _expr: &akar_parser::ast::Expression,
528        _chunk: &DataChunk,
529    ) -> Result<ValueVector, ProcessorError> {
530        // Placeholder: return a dummy Int64 vector
531        let size = _chunk.size;
532        let mut v = ValueVector::new(PhysicalTypeID::Int64, size);
533        for i in 0..size {
534            v.set_i64(i, 0);
535        }
536        v.resize(size);
537        Ok(v)
538    }
539}
540
541impl Default for QueryProcessor {
542    fn default() -> Self {
543        Self::new()
544    }
545}
akar_processor/processor/mod.rs

akar_processor/processor/
mod.rs
