akar_processor/physical/write_ops/
ddl_fts.rs1use 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
8pub 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
43pub 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 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 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 let source_data = source_table.to_column_major_data();
119 let num_rows = source_table.num_rows as usize;
120
121 let mut term_map: std::collections::HashMap<String, (i64, i64)> = std::collections::HashMap::new();
123 let mut doc_rows: Vec<Vec<Value>> = Vec::new();
125 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 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 {
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 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 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 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#[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 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 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 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 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(); 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 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 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 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 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 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}