Skip to main content

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_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/// DDL operations that require the schema-level Catalog (from Akar-catalog).
43/// These are dispatched via callback because the processor layer doesn't
44/// directly own the schema catalog.
45#[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
110/// The query processor executes a physical plan and produces result chunks.
111pub 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    /// Callback for sequence operations (nextval/currval).
117    /// Takes (sequence_name, is_nextval) and returns the resulting value.
118    sequence_fn: Option<SequenceFn>,
119    /// Callback for executing subqueries.
120    subquery_fn: Option<SubqueryFn>,
121    /// Callback for schema-level DDL operations (CREATE/DROP SEQUENCE, EXPORT/IMPORT DATABASE).
122    schema_ddl_fn: Option<SchemaDdlFn>,
123    /// MVCC snapshot timestamp for read isolation.
124    snapshot_ts: Option<u64>,
125    /// Commit history for MVCC visibility checks.
126    commit_history: Vec<(u64, u64)>,
127    /// Row-level write set accumulated during execution.
128    /// Populated by the mapper after each write operation (SET, DELETE, INSERT).
129    /// Read by the connection layer after execution for OCC conflict detection.
130    written_rows: Mutex<Vec<(u64, u64)>>,
131    /// Active transaction id threaded into write operators so inserts/deletes
132    /// are recorded in `VersionInfo` for MVCC snapshot isolation (P52.18).
133    txn_id: Option<u64>,
134    /// Undo records captured by write operators during execution. Drained into
135    /// the active transaction by the connection layer so a rollback (including
136    /// an OCC conflict loser) can revert the in-place table writes (P52.18).
137    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    /// Create a processor with access to the function registry.
159    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    /// Create a processor with function registry, table catalog access, and VFS.
177    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    /// Set the sequence operation callback (for nextval/currval).
199
200    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    /// Set the subquery operation callback.
211    pub fn with_subquery_fn(mut self, f: SubqueryFn) -> Self {
212        self.subquery_fn = Some(f);
213        self
214    }
215
216    /// Set the schema DDL callback (for CREATE/DROP SEQUENCE, EXPORT/IMPORT DATABASE).
217    pub fn with_schema_ddl_fn(mut self, f: SchemaDdlFn) -> Self {
218        self.schema_ddl_fn = Some(f);
219        self
220    }
221
222    /// Set MVCC snapshot parameters for read isolation.
223    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    /// Set the active transaction id for write operators.
230    ///
231    /// When `Some(txn_id)`, insert/delete write operators use the MVCC-aware
232    /// `*_with_txn` storage variants so uncommitted rows are invisible to other
233    /// snapshots, and record undo entries for rollback (P52.18).
234    pub fn with_txn_id(mut self, txn_id: Option<u64>) -> Self {
235        self.txn_id = txn_id;
236        self
237    }
238
239    /// Record an insert undo (rollback deletes the row).
240    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    /// Record an update undo (rollback restores the cell).
247    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    /// Record a delete undo (rollback restores the row/edge).
254    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    /// Take the accumulated undo records (drained by the connection layer).
261    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    /// Shared undo sink handed to write operators so they can record undo
269    /// entries during execution (P52.18).
270    pub fn undo_sink(&self) -> Arc<Mutex<Vec<UndoRecord>>> {
271        Arc::clone(&self.undo_records)
272    }
273
274    /// Execute a sequence of logical operators by mapping them to physical operators.
275    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            // Accumulate written rows from this operator into the processor's write set
318            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                // Accumulate: extend rather than replace for ScanRel
326                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    /// Take the accumulated write set from the processor.
339    /// Returns all (table_id, row_id) pairs written during the last execution.
340    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    /// Execute a table function call by looking up the function in the registry
348    /// and dispatching to the appropriate handler.
349    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(); // args would be evaluated from expressions
355
356        // Look up the function in the registry
357        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                        // Evaluate args: [table_name, column_name, query_vector, top_k]
384                        // For CALL statement, args are parsed as expressions. We need to evaluate them.
385                        // For now, parse from the function args
386                        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    /// Execute a `vector_similarity_scan` table function call.
406    ///
407    /// Expects CALL vector_similarity_scan(table_name, column_name, query_vector, top_k)
408    /// and dispatches to PhysicalVectorSimilarityScan with the processor's TableCatalog.
409    fn execute_vector_similarity_scan(
410        &self,
411        tf: &akar_planner::logical_operator::LogicalTableFunctionCall,
412    ) -> Result<Vec<DataChunk>, ProcessorError> {
413        // Evaluate arguments from expressions (they should be constants or simple vars)
414        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        // For CALL statements, args arrive as Expression AST nodes.
421        // Evaluate them to Values. The simplest approach: evaluate constants inline.
422        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, // Non-constant expression — skip
436            }
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        // Find the vector index on this table
472        let tc = self
473            .table_catalog
474            .clone()
475            .ok_or_else(|| "No table catalog available for vector_similarity_scan".to_string())?;
476
477        // Look for a vector index matching this table name
478        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        // Dispatch to PhysicalVectorSimilarityScan
490        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    /// Execute a single expression against a DataChunk and return a ValueVector of results.
502    pub fn evaluate_expression(
503        _expr: &akar_parser::ast::Expression,
504        _chunk: &DataChunk,
505    ) -> Result<ValueVector, ProcessorError> {
506        // Placeholder: return a dummy Int64 vector
507        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}