Skip to main content

akar_processor/physical/write_ops/
set.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
3use crate::physical::write_ops::delete::ast_constant_to_value;
4use akar_common::types::PhysicalTypeID;
5use akar_common::types::Value;
6use akar_common::vector::{DataChunk, ValueVector};
7use akar_function::registry::FunctionRegistry;
8use akar_function::scalar::evaluate_scalar;
9use akar_parser::ast::Expression;
10use akar_storage::table::TableCatalog;
11use std::sync::Arc;
12
13// ==================== Set ====================
14
15/// Physical operator for SET — updates a property on matched rows.
16pub struct PhysicalSet {
17    pub table_name: String,
18    pub table_id: u64,
19    pub column_name: String,
20    pub column_idx: usize,
21    pub value: akar_parser::ast::Expression,
22    pub is_node: bool,
23    pub table_catalog: Arc<TableCatalog>,
24}
25
26impl PhysicalOperatorExec for PhysicalSet {
27    fn operator_type(&self) -> &str {
28        "set"
29    }
30
31    fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
32        // Collect row indices from input chunks (first column has row index)
33        let mut rows_to_update: Vec<(u64, akar_common::types::Value)> = Vec::new();
34
35        for chunk in &input {
36            for row in 0..chunk.size {
37                if !chunk.fields.is_empty()
38                    && let Some(akar_common::types::Value::Int64(val)) = chunk.get_value(0, row)
39                {
40                    // Evaluate the SET value expression against the current row
41                    let set_val = evaluate_expression_for_row(&self.value, chunk, row);
42                    rows_to_update.push((val as u64, set_val));
43                }
44            }
45        }
46
47        if rows_to_update.is_empty() {
48            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
49            v.resize(1);
50            v.set_i64(0, 0);
51            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
52            return Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])]);
53        }
54
55        // Apply updates to the table
56        let mut updated = 0u64;
57        if self.is_node {
58            if let Some(mut table) = self.table_catalog.get_node_table_by_name_mut(&self.table_name) {
59                for (row_idx, val) in &rows_to_update {
60                    if table.update_cell(*row_idx, self.column_idx, val.clone()).is_ok() {
61                        updated += 1;
62                    }
63                }
64            } else {
65                return Err(format!("Node table '{}' not found for SET", self.table_name).into());
66            }
67        } else {
68            if let Some(mut table) = self.table_catalog.get_rel_table_by_name_mut(&self.table_name) {
69                for (edge_idx, val) in &rows_to_update {
70                    if table
71                        .update_cell(*edge_idx as usize, self.column_idx, val.clone())
72                        .is_ok()
73                    {
74                        updated += 1;
75                    }
76                }
77            } else {
78                return Err(format!("Rel table '{}' not found for SET", self.table_name).into());
79            }
80        }
81
82        tracing::info!("SET: updated {updated} rows in '{}'", self.table_name);
83
84        let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
85        v.resize(1);
86        v.set_i64(0, updated as i64);
87        let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
88        Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])])
89    }
90}
91
92/// Simple expression evaluator for SET value expressions against a DataChunk row.
93pub fn evaluate_expression_for_row(
94    expr: &akar_parser::ast::Expression,
95    chunk: &DataChunk,
96    row: usize,
97) -> akar_common::types::Value {
98    match expr {
99        akar_parser::ast::Expression::Constant(c) => match c {
100            akar_parser::ast::Constant::Null => akar_common::types::Value::Null,
101            akar_parser::ast::Constant::Bool(b) => akar_common::types::Value::Bool(*b),
102            akar_parser::ast::Constant::Integer(i) => akar_common::types::Value::Int64(*i),
103            akar_parser::ast::Constant::Float(f) => akar_common::types::Value::Double(*f),
104            akar_parser::ast::Constant::String(s) => akar_common::types::Value::String(s.clone()),
105        },
106        _ => {
107            // Fallback: try to get value from chunk fields
108            if chunk.fields.len() > 1 {
109                chunk.get_value(1, row).unwrap_or(akar_common::types::Value::Null)
110            } else {
111                akar_common::types::Value::Null
112            }
113        }
114    }
115}
116
117/// Evaluate a constant-only expression (literal or function call over
118/// literals) into a `Value`, using the function registry. Used by the
119/// CREATE DML write path to support expressions like `DATE('2024-01-15')`.
120/// Returns `Value::Null` for expressions that reference variables or
121/// otherwise cannot be folded without a row context.
122pub fn evaluate_constant_expr(expr: &Expression, registry: &FunctionRegistry) -> Value {
123    match expr {
124        Expression::Constant(c) => ast_constant_to_value(c),
125        Expression::FunctionCall(name, args) => {
126            let arg_values: Vec<Value> = args.iter().map(|a| evaluate_constant_expr(a, registry)).collect();
127            if arg_values.iter().any(|v| matches!(v, Value::Null)) {
128                return Value::Null;
129            }
130            let func = match registry.get_scalar(name).cloned() {
131                Some(f) => f,
132                None => return Value::Null,
133            };
134            evaluate_scalar(&func, &arg_values).unwrap_or(Value::Null)
135        }
136        _ => Value::Null,
137    }
138}