akar_processor/physical/write_ops/
vectorsimilarityscan.rs1use 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
8pub 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 let vi = if self.index_name.is_empty() {
37 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 let results = vi.hnsw().search(&self.query_vector, self.top_k as usize);
61 drop(vi); if results.is_empty() {
64 return Ok(vec![DataChunk::new(vec![], vec![])]);
65 }
66
67 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 let mut output_columns: Vec<Vec<Value>> = vec![Vec::with_capacity(num_results); num_cols + 2];
79
80 for (dist, row_id) in &results {
81 output_columns[num_cols].push(Value::Double(*dist));
83 output_columns[num_cols + 1].push(Value::Int64(*row_id as i64));
84
85 for (col_idx, out_col) in output_columns.iter_mut().enumerate().take(num_cols) {
87 match node_table.get_value(*row_id, col_idx) {
88 Some(val) => out_col.push(val.clone()),
89 None => out_col.push(Value::Null),
90 }
91 }
92 }
93
94 use akar_common::arrow_vector::{ArrowVector, arrow_array_from_values};
100 use akar_common::types::{PhysicalTypeID, physical_type_from_logical};
101
102 let mut fields = Vec::with_capacity(num_cols + 2);
103 let mut field_types = Vec::with_capacity(num_cols + 2);
104
105 for col_idx in 0..num_cols {
108 let phys_type = physical_type_from_logical(node_table.columns[col_idx].logical_type);
109 fields.push(ArrowVector::new(
110 arrow_array_from_values(&output_columns[col_idx]),
111 phys_type,
112 ));
113 field_types.push(phys_type);
114 }
115
116 fields.push(ArrowVector::new(
118 arrow_array_from_values(&output_columns[num_cols]),
119 PhysicalTypeID::Double,
120 ));
121 field_types.push(PhysicalTypeID::Double);
122
123 fields.push(ArrowVector::new(
126 arrow_array_from_values(&output_columns[num_cols + 1]),
127 PhysicalTypeID::Int64,
128 ));
129 field_types.push(PhysicalTypeID::Int64);
130
131 let arrow_fields = fields.iter().map(|v| v.array.clone()).collect::<Vec<_>>();
132 let arrow_field_types = field_types;
133
134 let mut field_names: Vec<String> = node_table.columns.iter().map(|c| c.name.clone()).collect();
135 field_names.push("distance".to_string());
136 field_names.push("_id".to_string());
137
138 Ok(vec![DataChunk {
139 fields: arrow_fields,
140 field_types: arrow_field_types,
141 size: num_results,
142 field_names,
143 sel_vector: None,
144 }])
145 }
146}