Skip to main content

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 + internal `_id`
76        // (physical row offset, the identity column other operators resolve).
77        // Convention matches `PhysicalArtIndexRangeScan` (copyfrom.rs).
78        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            // Add distance (second-to-last) and `_id` (last) columns.
82            output_columns[num_cols].push(Value::Double(*dist));
83            output_columns[num_cols + 1].push(Value::Int64(*row_id as i64));
84
85            // Look up each column value from the node table
86            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        // Convert column-major Vec<Vec<Value>> into an Arrow-native DataChunk.
95        // Every column is built via arrow_array_from_values, which encodes
96        // `Value::List` into a proper ListArray. The previous ValueVector path
97        // collapsed FLOAT[] (List) columns to NULL when read back downstream,
98        // breaking the cosine threshold filter / ORDER BY re-evaluation (P71.4).
99        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        // Add table columns — typed per column from the node table schema
106        // (P52.40: forcing every column to Double corrupted Int64/String data).
107        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        // Add the distance column (Double).
117        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        // Add internal `_id` (physical row offset) column, found by name by
124        // DELETE/SET/INSERT/extend machinery (row_id_column_index).
125        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}