Skip to main content

akar_processor/physical/write_ops/
ddl_fts.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
3use akar_common::types::{PhysicalTypeID, Value};
4use akar_common::vector::{DataChunk, ValueVector};
5use akar_storage::table::TableCatalog;
6use std::sync::Arc;
7
8// ==================== DDL & FTS ====================
9
10/// Physical COUNT on rel table — optimized via CSR metadata (Ladybug).
11/// Instead of scanning all edges, directly reads the edge count from the RelTable.
12pub struct PhysicalCountRelTable {
13    pub table_name: String,
14    pub table_id: u64,
15    pub table_catalog: Option<Arc<TableCatalog>>,
16}
17
18impl PhysicalOperatorExec for PhysicalCountRelTable {
19    fn operator_type(&self) -> &str {
20        "count_rel_table"
21    }
22
23    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
24        let tc = self
25            .table_catalog
26            .as_ref()
27            .ok_or_else(|| "No table catalog for CountRelTable".to_string())?;
28
29        let count = if let Some(table) = tc.get_rel_table(self.table_id) {
30            table.num_rows as i64
31        } else {
32            0
33        };
34
35        let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
36        v.resize(1);
37        v.set_i64(0, count);
38        let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&v).array;
39        Ok(vec![DataChunk::new(vec![arr], vec![PhysicalTypeID::Int64])])
40    }
41}
42
43/// Physical operator for `CREATE FTS INDEX` — builds 3 macro tables:
44/// 1. `fts_{idx}_docs`: node table (doc_id INT64, text STRING)
45/// 2. `fts_{idx}_terms`: node table (term_id INT64, term STRING, doc_freq INT64)
46/// 3. `fts_{idx}_appears_in`: rel table (FROM terms TO docs, term_freq INT64)
47pub struct PhysicalCreateFtsIndex {
48    pub index_name: String,
49    pub table_name: String,
50    pub column_name: String,
51    pub docs_table: String,
52    pub terms_table: String,
53    pub posting_table: String,
54    pub table_catalog: Arc<TableCatalog>,
55}
56
57impl PhysicalOperatorExec for PhysicalCreateFtsIndex {
58    fn operator_type(&self) -> &str {
59        "create_fts_index"
60    }
61
62    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
63        // Locate source table
64        let source_table = match self.table_catalog.get_node_table_by_name(&self.table_name) {
65            Some(t) => t,
66            None => return Err(format!("Table '{}' not found", self.table_name).into()),
67        };
68        let col_idx = source_table
69            .columns
70            .iter()
71            .position(|c| c.name == self.column_name)
72            .ok_or_else(|| format!("Column '{}' not found in '{}'", self.column_name, self.table_name))?;
73
74        // Ensure macro tables exist; create if needed
75        if self.table_catalog.get_node_table_by_name(&self.docs_table).is_none() {
76            let docs_cols = vec![
77                akar_storage::table::ColumnDefinition {
78                    name: "doc_id".into(),
79                    logical_type: akar_common::types::LogicalTypeID::Int64,
80                    is_primary_key: true,
81                    compression: akar_common::enums::CompressionType::Uncompressed,
82                },
83                akar_storage::table::ColumnDefinition {
84                    name: "text".into(),
85                    logical_type: akar_common::types::LogicalTypeID::String,
86                    is_primary_key: false,
87                    compression: akar_common::enums::CompressionType::Uncompressed,
88                },
89            ];
90            self.table_catalog.create_node_table(self.docs_table.clone(), docs_cols);
91        }
92        if self.table_catalog.get_node_table_by_name(&self.terms_table).is_none() {
93            let terms_cols = vec![
94                akar_storage::table::ColumnDefinition {
95                    name: "term_id".into(),
96                    logical_type: akar_common::types::LogicalTypeID::Int64,
97                    is_primary_key: true,
98                    compression: akar_common::enums::CompressionType::Uncompressed,
99                },
100                akar_storage::table::ColumnDefinition {
101                    name: "term".into(),
102                    logical_type: akar_common::types::LogicalTypeID::String,
103                    is_primary_key: false,
104                    compression: akar_common::enums::CompressionType::Uncompressed,
105                },
106                akar_storage::table::ColumnDefinition {
107                    name: "doc_freq".into(),
108                    logical_type: akar_common::types::LogicalTypeID::Int64,
109                    is_primary_key: false,
110                    compression: akar_common::enums::CompressionType::Uncompressed,
111                },
112            ];
113            self.table_catalog
114                .create_node_table(self.terms_table.clone(), terms_cols);
115        }
116
117        // Collect docs data
118        let source_data = source_table.to_column_major_data();
119        let num_rows = source_table.num_rows as usize;
120
121        // term -> (term_id, doc_freq)
122        let mut term_map: std::collections::HashMap<String, (i64, i64)> = std::collections::HashMap::new();
123        // (doc_id, text) rows
124        let mut doc_rows: Vec<Vec<Value>> = Vec::new();
125        // posting: (term_id, doc_id, term_freq)
126        let mut postings: Vec<(i64, i64, i64)> = Vec::new();
127
128        for row_idx in 0..num_rows {
129            let text = if let Some(col_data) = source_data.get(col_idx) {
130                if let Some(Value::String(s)) = col_data.get(row_idx) {
131                    s.clone()
132                } else {
133                    String::new()
134                }
135            } else {
136                String::new()
137            };
138
139            let doc_id = row_idx as i64;
140            doc_rows.push(vec![Value::Int64(doc_id), Value::String(text.clone())]);
141
142            // Tokenize using Akar-fts utilities
143            let tokens = akar_fts::tokenize(&text);
144            let mut freq_map: std::collections::HashMap<String, i64> = std::collections::HashMap::new();
145            for token in tokens {
146                let stemmed = akar_fts::stem_word(&token);
147                if !akar_fts::STOP_WORDS.contains(&stemmed.as_str()) {
148                    *freq_map.entry(stemmed).or_insert(0) += 1;
149                }
150            }
151
152            for (term, freq) in freq_map {
153                let next_id = term_map.len() as i64;
154                let (term_id, doc_freq) = term_map.entry(term).or_insert((next_id, 0));
155                *doc_freq += 1;
156                postings.push((*term_id, doc_id, freq));
157            }
158        }
159
160        // Insert docs
161        {
162            let mut docs_table = self.table_catalog.get_node_table_by_name_mut(&self.docs_table).unwrap();
163            for row in doc_rows {
164                docs_table.insert_row(row)?;
165            }
166        }
167
168        // Insert terms
169        if self.table_catalog.get_node_table_by_name(&self.terms_table).is_some() {
170            let mut terms_table = self
171                .table_catalog
172                .get_node_table_by_name_mut(&self.terms_table)
173                .unwrap();
174            let mut term_list: Vec<(String, i64, i64)> =
175                term_map.into_iter().map(|(t, (id, df))| (t, id, df)).collect();
176            term_list.sort_by_key(|(_, id, _)| *id);
177            for (term, term_id, doc_freq) in term_list {
178                terms_table.insert_row(vec![Value::Int64(term_id), Value::String(term), Value::Int64(doc_freq)])?;
179            }
180        }
181
182        // Create and populate posting (appears_in) table
183        let docs_table_id = self
184            .table_catalog
185            .get_node_table_by_name(&self.docs_table)
186            .unwrap()
187            .table_id;
188        let terms_table_id = self
189            .table_catalog
190            .get_node_table_by_name(&self.terms_table)
191            .unwrap()
192            .table_id;
193
194        if self.table_catalog.get_rel_table_by_name(&self.posting_table).is_none() {
195            let posting_cols = vec![akar_storage::table::ColumnDefinition {
196                name: "term_freq".into(),
197                logical_type: akar_common::types::LogicalTypeID::Int64,
198                is_primary_key: false,
199                compression: akar_common::enums::CompressionType::Uncompressed,
200            }];
201            // FROM terms TO docs
202            self.table_catalog.create_rel_table(
203                self.posting_table.clone(),
204                terms_table_id,
205                docs_table_id,
206                posting_cols,
207            );
208        }
209
210        {
211            let mut posting_table = self
212                .table_catalog
213                .get_rel_table_by_name_mut(&self.posting_table)
214                .unwrap();
215            for (term_id, doc_id, freq) in postings {
216                posting_table.insert_rel(term_id as u64, doc_id as u64, vec![Value::Int64(freq)])?;
217            }
218        }
219
220        let mut result_vec = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::String, 1);
221        result_vec.resize(1);
222        result_vec
223            .set_value(
224                0,
225                &Value::String(format!("FTS index '{}' built successfully.", self.index_name)),
226            )
227            .unwrap();
228        let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&result_vec).array;
229        let mut result = DataChunk::new(vec![arr], vec![akar_common::types::PhysicalTypeID::String]);
230        result.size = 1;
231        result.field_names = vec!["result".to_string()];
232        Ok(vec![result])
233    }
234}
235
236/// Physical operator for `USING FTS INDEX` scan — queries the 3 macro tables
237/// and returns ranked (node_id, score) pairs using BM25 scoring.
238#[derive(Debug, Clone)]
239pub struct PhysicalFtsScan {
240    pub index_name: String,
241    pub query_string: String,
242    pub docs_table: String,
243    pub terms_table: String,
244    pub posting_table: String,
245    pub table_catalog: Arc<TableCatalog>,
246}
247
248impl PhysicalOperatorExec for PhysicalFtsScan {
249    fn operator_type(&self) -> &str {
250        "fts_scan"
251    }
252
253    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
254        // Tokenize query
255        let query_tokens: Vec<String> = akar_fts::tokenize(&self.query_string)
256            .into_iter()
257            .map(|t| akar_fts::stem_word(&t))
258            .filter(|t| !akar_fts::STOP_WORDS.contains(&t.as_str()))
259            .collect();
260
261        // Lookup terms table for matching terms
262        let terms_table = match self.table_catalog.get_node_table_by_name(&self.terms_table) {
263            Some(t) => t,
264            None => {
265                return Err(format!(
266                    "FTS terms table '{}' not found. Has the index been created?",
267                    self.terms_table
268                )
269                .into());
270            }
271        };
272
273        // Get total doc count from docs table
274        let num_docs = self
275            .table_catalog
276            .get_node_table_by_name(&self.docs_table)
277            .map(|t| t.num_rows as f64)
278            .unwrap_or(1.0);
279
280        // Build map: term -> (term_id, doc_freq)
281        let terms_data = terms_table.to_column_major_data();
282        let num_terms = terms_table.num_rows as usize;
283        let mut matching_terms: Vec<(i64, i64)> = Vec::new(); // (term_id, doc_freq)
284
285        for row_idx in 0..num_terms {
286            let term_val = terms_data.get(1).and_then(|d| d.get(row_idx));
287            let term_str = if let Some(Value::String(s)) = term_val {
288                s.clone()
289            } else {
290                continue;
291            };
292            if query_tokens.contains(&term_str) {
293                let term_id = if let Some(Value::Int64(id)) = terms_data.first().and_then(|d| d.get(row_idx)) {
294                    *id
295                } else {
296                    continue;
297                };
298                let doc_freq = if let Some(Value::Int64(df)) = terms_data.get(2).and_then(|d| d.get(row_idx)) {
299                    *df
300                } else {
301                    0
302                };
303                matching_terms.push((term_id, doc_freq));
304            }
305        }
306
307        // Accumulate per-doc BM25 scores from posting table
308        let mut doc_scores: std::collections::HashMap<i64, f64> = std::collections::HashMap::new();
309
310        if let Some(posting_table) = self.table_catalog.get_rel_table_by_name(&self.posting_table) {
311            for &(term_id, doc_freq) in &matching_terms {
312                let idf = ((num_docs - doc_freq as f64 + 0.5) / (doc_freq as f64 + 0.5) + 1.0).ln();
313                // Scan posting table for this term using get_outgoing_edges(term_id)
314                let posting_rels = posting_table.get_outgoing_edges(term_id as u64);
315                for (doc_id, rel_vals) in posting_rels {
316                    let tf = if let Some(Value::Int64(freq)) = rel_vals.first() {
317                        *freq as f64
318                    } else {
319                        1.0
320                    };
321                    // BM25: k1=1.5, b=0.75 (simplified, no avg doc len)
322                    let k1 = 1.5_f64;
323                    let score = idf * (tf * (k1 + 1.0)) / (tf + k1);
324                    *doc_scores.entry(doc_id as i64).or_insert(0.0) += score;
325                }
326            }
327        }
328
329        // Sort by score descending
330        let mut ranked: Vec<(i64, f64)> = doc_scores.into_iter().collect();
331        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
332
333        // Return (doc_id, score) data chunks
334        let n = ranked.len();
335        let mut id_vec = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, n);
336        let mut score_vec = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Double, n);
337        id_vec.resize(n);
338        score_vec.resize(n);
339        for (i, (doc_id, score)) in ranked.into_iter().enumerate() {
340            id_vec.set_i64(i, doc_id);
341            score_vec.set_double(i, score);
342        }
343        let arr1 = akar_common::arrow_vector::ArrowVector::from_legacy(&id_vec).array;
344        let arr2 = akar_common::arrow_vector::ArrowVector::from_legacy(&score_vec).array;
345        let mut chunk = DataChunk::new(
346            vec![arr1, arr2],
347            vec![
348                akar_common::types::PhysicalTypeID::Int64,
349                akar_common::types::PhysicalTypeID::Double,
350            ],
351        );
352        chunk.size = n;
353        chunk.field_names = vec!["doc_id".to_string(), "score".to_string()];
354        Ok(vec![chunk])
355    }
356}