akar_processor/physical/write_ops/
vectorsimilarityscan.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
3use akar_common::types::Value;
4use akar_common::vector::DataChunk;
5use akar_storage::table::TableCatalog;
6use std::sync::Arc;
7
8// ==================== VectorSimilarityScan ====================
9
10/// Physical operator for vector similarity search using an HNSW index.
11///
12/// Searches the `VectorIndexTable` for the top-K nearest neighbours and
13/// looks up the corresponding rows from the `NodeTable` to produce
14/// output columns including a `distance` column.
15pub struct PhysicalVectorSimilarityScan {
16    pub index_name: String,
17    pub index_id: u64,
18    pub query_vector: Vec<f64>,
19    pub top_k: u64,
20    pub table_name: String,
21    pub table_catalog: Option<Arc<TableCatalog>>,
22}
23
24impl PhysicalOperatorExec for PhysicalVectorSimilarityScan {
25    fn operator_type(&self) -> &str {
26        "vector_similarity_scan"
27    }
28
29    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
30        let tc = self
31            .table_catalog
32            .clone()
33            .ok_or_else(|| "No table catalog available for VectorSimilarityScan".to_string())?;
34
35        // Resolve the vector index — by name if given, or find first index on the table
36        let vi = if self.index_name.is_empty() {
37            // Find the first vector index on this table
38            // Scan all vector indexes to find one matching this table
39            let index_name = {
40                let mut found_name = String::new();
41                for entry in tc.all_vector_indexes() {
42                    if entry.table_name == self.table_name {
43                        found_name = entry.name.clone();
44                        break;
45                    }
46                }
47                if found_name.is_empty() {
48                    return Err(format!("No vector index found on table '{}'", self.table_name).into());
49                }
50                found_name
51            };
52            tc.get_vector_index_by_name(&index_name)
53                .ok_or_else(|| format!("Vector index '{}' not found", index_name))?
54        } else {
55            tc.get_vector_index_by_name(&self.index_name)
56                .ok_or_else(|| format!("Vector index '{}' not found", self.index_name))?
57        };
58
59        // Search the HNSW index for top-K nearest neighbours
60        let results = vi.hnsw().search(&self.query_vector, self.top_k as usize);
61        drop(vi); // Release the DashMap reference
62
63        if results.is_empty() {
64            return Ok(vec![DataChunk::new(vec![], vec![])]);
65        }
66
67        // Look up rows from the node table
68        let node_table = tc
69            .get_node_table_by_name(&self.table_name)
70            .ok_or_else(|| format!("Node table '{}' not found", self.table_name))?;
71
72        let num_cols = node_table.columns.len();
73        let num_results = results.len();
74
75        // Build output columns: all table columns + distance column
76        let mut output_columns: Vec<Vec<Value>> = vec![Vec::with_capacity(num_results); num_cols + 1];
77
78        for (dist, row_id) in &results {
79            // Add distance as the last column
80            output_columns[num_cols].push(Value::Double(*dist));
81
82            // Look up each column value from the node table
83            for (col_idx, out_col) in output_columns.iter_mut().enumerate().take(num_cols) {
84                match node_table.get_value(*row_id, col_idx) {
85                    Some(val) => out_col.push(val.clone()),
86                    None => out_col.push(Value::Null),
87                }
88            }
89        }
90
91        drop(node_table);
92
93        // Convert column-major Vec<Vec<Value>> to DataChunks
94        use akar_common::types::PhysicalTypeID;
95        use akar_common::vector::{DataChunk, ValueVector};
96
97        let mut fields = Vec::with_capacity(num_cols + 1);
98
99        // Add table columns
100        for (_col_idx, col_data) in output_columns.iter().enumerate().take(num_cols) {
101            let mut v = ValueVector::new(PhysicalTypeID::Double, num_results);
102            v.resize(num_results);
103            for (i, val) in col_data.iter().enumerate() {
104                match val {
105                    Value::Double(d) => v.set_double(i, *d),
106                    Value::Int64(x) => {
107                        let buf = &mut v.data_mut()[i * 8..(i + 1) * 8];
108                        buf.copy_from_slice(&x.to_le_bytes());
109                        v.set_null(i, false);
110                    }
111                    Value::String(_) => {
112                        v.set_value(i, val)?;
113                    }
114                    Value::Null => {
115                        v.set_null(i, true);
116                    }
117                    _ => {
118                        v.set_null(i, true);
119                    }
120                }
121            }
122            fields.push(v);
123        }
124
125        // Add distance column
126        let dist_data = &output_columns[num_cols];
127        let mut dist_v = ValueVector::new(PhysicalTypeID::Double, num_results);
128        dist_v.resize(num_results);
129        for (i, val) in dist_data.iter().enumerate() {
130            if let Value::Double(d) = val {
131                dist_v.set_double(i, *d);
132            } else {
133                dist_v.set_null(i, true);
134            }
135        }
136        fields.push(dist_v);
137
138        let arrow_fields = fields
139            .iter()
140            .map(|v| akar_common::arrow_vector::ArrowVector::from_legacy(v).array)
141            .collect::<Vec<_>>();
142        let arrow_field_types = fields.iter().map(|v| v.physical_type()).collect::<Vec<_>>();
143
144        Ok(vec![DataChunk {
145            fields: arrow_fields,
146            field_types: arrow_field_types,
147            size: num_results,
148            field_names: vec![],
149            sel_vector: None,
150        }])
151    }
152}
akar_processor/physical/write_ops/vectorsimilarityscan.rs

akar_processor/physical/write_ops/
vectorsimilarityscan.rs
