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

akar_processor/processor/
mod.rs
