1use super::*;
8use reddb_rql::ast::GraphCommandOrderBy;
9use std::cmp::Ordering;
10
11impl RedDBRuntime {
12 pub fn execute_graph_command(
14 &self,
15 raw_query: &str,
16 cmd: &GraphCommand,
17 ) -> RedDBResult<RuntimeQueryResult> {
18 match cmd {
19 GraphCommand::Neighborhood {
20 source,
21 depth,
22 direction,
23 edge_labels,
24 } => {
25 let dir = parse_direction(direction)?;
26 let res = self.graph_neighborhood(
27 source,
28 dir,
29 *depth as usize,
30 edge_labels.clone(),
31 None,
32 )?;
33 let mut result = UnifiedResult::with_columns(vec![
34 "node_id".into(),
35 "label".into(),
36 "node_type".into(),
37 "depth".into(),
38 ]);
39 for visit in &res.nodes {
40 let mut record = UnifiedRecord::new();
41 record.set("node_id", Value::text(visit.node.id.clone()));
42 record.set("label", Value::text(visit.node.label.clone()));
43 record.set("node_type", Value::text(visit.node.node_type.clone()));
44 record.set("depth", Value::Integer(visit.depth as i64));
45 result.push(record);
46 }
47 Ok(RuntimeQueryResult {
48 query: raw_query.to_string(),
49 mode: QueryMode::Sql,
50 statement: "graph_neighborhood",
51 engine: "runtime-graph",
52 result,
53 affected_rows: 0,
54 statement_type: "select",
55 bookmark: None,
56 notice: None,
57 })
58 }
59 GraphCommand::ShortestPath {
60 source,
61 target,
62 algorithm,
63 direction,
64 limit,
65 order_by,
66 } => {
67 let dir = parse_direction(direction)?;
68 let alg = parse_path_algorithm(algorithm)?;
69 let res = self.graph_shortest_path(source, target, dir, alg, None, None)?;
70 let mut result = UnifiedResult::with_columns(vec![
71 "source".into(),
72 "target".into(),
73 "nodes_visited".into(),
74 "negative_cycle_detected".into(),
75 "path_found".into(),
76 "hop_count".into(),
77 "total_weight".into(),
78 ]);
79 let mut record = UnifiedRecord::new();
80 record.set("source", Value::text(res.source));
81 record.set("target", Value::text(res.target));
82 record.set("nodes_visited", Value::Integer(res.nodes_visited as i64));
83 record.set(
84 "negative_cycle_detected",
85 match res.negative_cycle_detected {
86 Some(value) => Value::Boolean(value),
87 None => Value::Null,
88 },
89 );
90 if let Some(ref path) = res.path {
91 record.set("path_found", Value::Boolean(true));
92 record.set("hop_count", Value::Integer(path.hop_count as i64));
93 record.set("total_weight", Value::Float(path.total_weight));
94 } else {
95 record.set("path_found", Value::Boolean(false));
96 record.set("hop_count", Value::Null);
97 record.set("total_weight", Value::Null);
98 }
99 result.push(record);
100 apply_graph_order_and_limit(
101 &mut result,
102 "graph_shortest_path",
103 order_by.as_ref(),
104 limit.map(|n| n as usize),
105 )?;
106 Ok(RuntimeQueryResult {
107 query: raw_query.to_string(),
108 mode: QueryMode::Sql,
109 statement: "graph_shortest_path",
110 engine: "runtime-graph",
111 result,
112 affected_rows: 0,
113 statement_type: "select",
114 bookmark: None,
115 notice: None,
116 })
117 }
118 GraphCommand::Properties { source } => {
119 if let Some(node_ref) = source {
120 let entity = resolve_graph_node_properties(self, node_ref)?;
124 let node_id = entity.logical_id().raw().to_string();
125 let EntityKind::GraphNode(node) = entity.kind else {
126 unreachable!("property resolver returns only graph nodes");
127 };
128 let props = match entity.data {
129 EntityData::Node(data) => data.properties,
130 _ => HashMap::new(),
131 };
132
133 let mut prop_keys: Vec<&String> = props.keys().collect();
136 prop_keys.sort();
137 let mut columns: Vec<String> = Vec::with_capacity(3 + prop_keys.len());
138 columns.push("node_id".into());
139 columns.push("label".into());
140 columns.push("node_type".into());
141 for k in &prop_keys {
142 columns.push((*k).clone());
143 }
144 let mut result = UnifiedResult::with_columns(columns);
145 let mut record = UnifiedRecord::new();
146 record.set("node_id", Value::text(node_id));
147 record.set("label", Value::text(node.label));
148 record.set("node_type", Value::text(node.node_type));
149 for k in &prop_keys {
150 if let Some(v) = props.get(*k) {
151 record.set(k.as_str(), v.clone());
152 }
153 }
154 result.push(record);
155 return Ok(RuntimeQueryResult {
156 query: raw_query.to_string(),
157 mode: QueryMode::Sql,
158 statement: "graph_properties",
159 engine: "runtime-graph",
160 result,
161 affected_rows: 0,
162 statement_type: "select",
163 bookmark: None,
164 notice: None,
165 });
166 }
167 let res = self.graph_properties(None)?;
168 let mut result = UnifiedResult::with_columns(vec![
169 "node_count".into(),
170 "edge_count".into(),
171 "is_connected".into(),
172 "is_complete".into(),
173 "is_cyclic".into(),
174 "density".into(),
175 ]);
176 let mut record = UnifiedRecord::new();
177 record.set("node_count", Value::Integer(res.node_count as i64));
178 record.set("edge_count", Value::Integer(res.edge_count as i64));
179 record.set("is_connected", Value::Boolean(res.is_connected));
180 record.set("is_complete", Value::Boolean(res.is_complete));
181 record.set("is_cyclic", Value::Boolean(res.is_cyclic));
182 record.set("density", Value::Float(res.density));
183 result.push(record);
184 Ok(RuntimeQueryResult {
185 query: raw_query.to_string(),
186 mode: QueryMode::Sql,
187 statement: "graph_properties",
188 engine: "runtime-graph",
189 result,
190 affected_rows: 0,
191 statement_type: "select",
192 bookmark: None,
193 notice: None,
194 })
195 }
196 GraphCommand::Traverse {
197 source,
198 strategy,
199 depth,
200 direction,
201 edge_labels,
202 } => {
203 let dir = parse_direction(direction)?;
204 let strat = parse_traversal_strategy(strategy)?;
205 let res = self.graph_traverse(
206 source,
207 dir,
208 *depth as usize,
209 strat,
210 edge_labels.clone(),
211 None,
212 )?;
213 let mut result = UnifiedResult::with_columns(vec![
214 "node_id".into(),
215 "label".into(),
216 "node_type".into(),
217 "depth".into(),
218 ]);
219 for visit in &res.visits {
220 let mut record = UnifiedRecord::new();
221 record.set("node_id", Value::text(visit.node.id.clone()));
222 record.set("label", Value::text(visit.node.label.clone()));
223 record.set("node_type", Value::text(visit.node.node_type.clone()));
224 record.set("depth", Value::Integer(visit.depth as i64));
225 result.push(record);
226 }
227 Ok(RuntimeQueryResult {
228 query: raw_query.to_string(),
229 mode: QueryMode::Sql,
230 statement: "graph_traverse",
231 engine: "runtime-graph",
232 result,
233 affected_rows: 0,
234 statement_type: "select",
235 bookmark: None,
236 notice: None,
237 })
238 }
239 GraphCommand::Centrality {
240 algorithm,
241 limit,
242 order_by,
243 } => {
244 let alg = parse_centrality_algorithm(algorithm)?;
245 let limit_usize = limit.map(|n| n as usize);
248 let order_needs_full_set = order_by
249 .as_ref()
250 .map(|order| order.ascending)
251 .unwrap_or(false);
252 let top_k = if order_needs_full_set {
253 usize::MAX
254 } else {
255 limit_usize.unwrap_or(100).max(1)
256 };
257 let res = self.graph_centrality(alg, top_k, false, None, None, None, None)?;
258 let mut result = UnifiedResult::with_columns(vec![
259 "node_id".into(),
260 "label".into(),
261 "score".into(),
262 ]);
263 for score in &res.scores {
264 let mut record = UnifiedRecord::new();
265 record.set("node_id", Value::text(score.node.id.clone()));
266 record.set("label", Value::text(score.node.label.clone()));
267 record.set("score", Value::Float(score.score));
268 result.push(record);
269 }
270 for ds in &res.degree_scores {
271 let mut record = UnifiedRecord::new();
272 record.set("node_id", Value::text(ds.node.id.clone()));
273 record.set("label", Value::text(ds.node.label.clone()));
274 record.set("score", Value::Float(ds.total_degree as f64));
275 result.push(record);
276 }
277 apply_graph_order_and_limit(
278 &mut result,
279 "graph_centrality",
280 order_by.as_ref(),
281 Some(limit_usize.unwrap_or(100)),
282 )?;
283 Ok(RuntimeQueryResult {
284 query: raw_query.to_string(),
285 mode: QueryMode::Sql,
286 statement: "graph_centrality",
287 engine: "runtime-graph",
288 result,
289 affected_rows: 0,
290 statement_type: "select",
291 bookmark: None,
292 notice: None,
293 })
294 }
295 GraphCommand::Community {
296 algorithm,
297 max_iterations,
298 limit,
299 order_by,
300 return_assignments,
301 } => {
302 let alg = parse_community_algorithm(algorithm)?;
303 let res =
304 self.graph_communities(alg, 1, Some(*max_iterations as usize), None, None)?;
305 let result = if *return_assignments {
306 let mut result =
312 UnifiedResult::with_columns(vec!["node_id".into(), "community_id".into()]);
313 let row_cap = limit.map(|n| n as usize);
314 'outer: for community in &res.communities {
315 let mut nodes = community.nodes.clone();
316 nodes.sort();
317 for node_id in nodes {
318 if row_cap.is_some_and(|cap| result.records.len() >= cap) {
319 break 'outer;
320 }
321 let mut record = UnifiedRecord::new();
322 record.set("node_id", Value::text(node_id));
323 record.set("community_id", Value::text(community.id.clone()));
324 result.push(record);
325 }
326 }
327 result
328 } else {
329 let mut result =
330 UnifiedResult::with_columns(vec!["community_id".into(), "size".into()]);
331 for community in &res.communities {
332 let mut record = UnifiedRecord::new();
333 record.set("community_id", Value::text(community.id.clone()));
334 record.set("size", Value::Integer(community.size as i64));
335 result.push(record);
336 }
337 apply_graph_order_and_limit(
338 &mut result,
339 "graph_community",
340 order_by.as_ref(),
341 limit.map(|n| n as usize),
342 )?;
343 result
344 };
345 Ok(RuntimeQueryResult {
346 query: raw_query.to_string(),
347 mode: QueryMode::Sql,
348 statement: "graph_community",
349 engine: "runtime-graph",
350 result,
351 affected_rows: 0,
352 statement_type: "select",
353 bookmark: None,
354 notice: None,
355 })
356 }
357 GraphCommand::Components {
358 mode,
359 limit,
360 order_by,
361 } => {
362 let m = parse_components_mode(mode)?;
363 let res = self.graph_components(m, 1, None)?;
364 let mut result =
365 UnifiedResult::with_columns(vec!["component_id".into(), "size".into()]);
366 for component in &res.components {
367 let mut record = UnifiedRecord::new();
368 record.set("component_id", Value::text(component.id.clone()));
369 record.set("size", Value::Integer(component.size as i64));
370 result.push(record);
371 }
372 apply_graph_order_and_limit(
373 &mut result,
374 "graph_components",
375 order_by.as_ref(),
376 limit.map(|n| n as usize),
377 )?;
378 Ok(RuntimeQueryResult {
379 query: raw_query.to_string(),
380 mode: QueryMode::Sql,
381 statement: "graph_components",
382 engine: "runtime-graph",
383 result,
384 affected_rows: 0,
385 statement_type: "select",
386 bookmark: None,
387 notice: None,
388 })
389 }
390 GraphCommand::Cycles { max_length } => {
391 let res = self.graph_cycles(*max_length as usize, 100, None)?;
392 let mut result =
393 UnifiedResult::with_columns(vec!["cycle_index".into(), "length".into()]);
394 for (i, cycle) in res.cycles.iter().enumerate() {
395 let mut record = UnifiedRecord::new();
396 record.set("cycle_index", Value::Integer(i as i64));
397 record.set("length", Value::Integer(cycle.nodes.len() as i64));
398 result.push(record);
399 }
400 Ok(RuntimeQueryResult {
401 query: raw_query.to_string(),
402 mode: QueryMode::Sql,
403 statement: "graph_cycles",
404 engine: "runtime-graph",
405 result,
406 affected_rows: 0,
407 statement_type: "select",
408 bookmark: None,
409 notice: None,
410 })
411 }
412 GraphCommand::Clustering => {
413 let res = self.graph_clustering(100, true, None)?;
414 let mut result = UnifiedResult::with_columns(vec![
415 "node_id".into(),
416 "label".into(),
417 "score".into(),
418 ]);
419 let mut global_record = UnifiedRecord::new();
421 global_record.set("node_id", Value::text("__global__"));
422 global_record.set("label", Value::text("global_clustering"));
423 global_record.set("score", Value::Float(res.global));
424 result.push(global_record);
425 for score in &res.local {
426 let mut record = UnifiedRecord::new();
427 record.set("node_id", Value::text(score.node.id.clone()));
428 record.set("label", Value::text(score.node.label.clone()));
429 record.set("score", Value::Float(score.score));
430 result.push(record);
431 }
432 Ok(RuntimeQueryResult {
433 query: raw_query.to_string(),
434 mode: QueryMode::Sql,
435 statement: "graph_clustering",
436 engine: "runtime-graph",
437 result,
438 affected_rows: 0,
439 statement_type: "select",
440 bookmark: None,
441 notice: None,
442 })
443 }
444 GraphCommand::TopologicalSort => {
445 let res = self.graph_topological_sort(None)?;
446 let mut result = UnifiedResult::with_columns(vec![
447 "order".into(),
448 "node_id".into(),
449 "label".into(),
450 ]);
451 for (i, node) in res.ordered_nodes.iter().enumerate() {
452 let mut record = UnifiedRecord::new();
453 record.set("order", Value::Integer(i as i64));
454 record.set("node_id", Value::text(node.id.clone()));
455 record.set("label", Value::text(node.label.clone()));
456 result.push(record);
457 }
458 Ok(RuntimeQueryResult {
459 query: raw_query.to_string(),
460 mode: QueryMode::Sql,
461 statement: "graph_topological_sort",
462 engine: "runtime-graph",
463 result,
464 affected_rows: 0,
465 statement_type: "select",
466 bookmark: None,
467 notice: None,
468 })
469 }
470 }
471 }
472
473 pub fn execute_search_command(
475 &self,
476 raw_query: &str,
477 cmd: &SearchCommand,
478 ) -> RedDBResult<RuntimeQueryResult> {
479 match cmd {
480 SearchCommand::Similar {
481 vector,
482 text,
483 provider,
484 collection,
485 limit,
486 min_score,
487 vector_param,
488 limit_param,
489 min_score_param,
490 text_param,
491 } => {
492 if vector_param.is_some() {
493 return Err(RedDBError::Query(
494 "SEARCH SIMILAR $N vector parameter was not bound before execution"
495 .to_string(),
496 ));
497 }
498 if limit_param.is_some() {
499 return Err(RedDBError::Query(
500 "SEARCH SIMILAR LIMIT $N parameter was not bound before execution"
501 .to_string(),
502 ));
503 }
504 if min_score_param.is_some() {
505 return Err(RedDBError::Query(
506 "SEARCH SIMILAR MIN_SCORE $N parameter was not bound before execution"
507 .to_string(),
508 ));
509 }
510 if text_param.is_some() {
511 return Err(RedDBError::Query(
512 "SEARCH SIMILAR TEXT $N parameter was not bound before execution"
513 .to_string(),
514 ));
515 }
516 let search_vector = if let Some(query_text) = text {
518 use crate::application::ports::RuntimeEntityPort;
521 let kv_getter = |key: &str| -> RedDBResult<Option<String>> {
522 match self.get_kv("red_config", key)? {
523 Some((Value::Text(s), _)) => Ok(Some(s.to_string())),
524 _ => Ok(None),
525 }
526 };
527 let provider = match provider.as_deref() {
528 Some(p) => {
529 let provider = crate::ai::parse_provider(p)?;
530 crate::ai::ensure_provider_supports_embeddings(&provider)?;
531 provider
532 }
533 None => crate::ai::resolve_embeddings_provider(&kv_getter)?,
534 };
535 crate::runtime::ai::provider_gate::enforce(self, &provider)?;
539 let api_key = crate::ai::resolve_api_key_from_runtime(&provider, None, self)?;
540 let model = crate::ai::resolve_embeddings_model(&provider, &kv_getter);
541 let transport = crate::runtime::ai::transport::AiTransport::from_runtime(self);
542 let request = crate::ai::OpenAiEmbeddingRequest {
543 api_key,
544 model,
545 inputs: vec![query_text.clone()],
546 dimensions: None,
547 api_base: provider.resolve_api_base(),
548 };
549 let response = crate::runtime::ai::block_on_ai(async move {
550 crate::ai::openai_embeddings_async(&transport, request).await
551 })
552 .and_then(|result| result)?;
553 response.embeddings.into_iter().next().ok_or_else(|| {
554 RedDBError::Query("embedding API returned no vectors".to_string())
555 })?
556 } else {
557 vector.clone()
558 };
559 let scope = self.ai_scope();
563 let results =
564 if super::statement_frame::ReadFrame::visible_collections(&scope).is_some() {
565 crate::runtime::authorized_search::AuthorizedSearch::execute_similar(
566 self,
567 &scope,
568 collection,
569 &search_vector,
570 *limit,
571 *min_score,
572 )?
573 } else {
574 self.search_similar(collection, &search_vector, *limit, *min_score)?
576 };
577 let mut result =
578 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
579 for sr in &results {
580 let mut record = UnifiedRecord::new();
581 record.set("entity_id", Value::UnsignedInteger(sr.entity_id.raw()));
582 record.set("score", Value::Float(sr.score as f64));
583 result.push(record);
584 }
585 Ok(RuntimeQueryResult {
586 query: raw_query.to_string(),
587 mode: QueryMode::Sql,
588 statement: "search_similar",
589 engine: "runtime-search",
590 result,
591 affected_rows: 0,
592 statement_type: "select",
593 bookmark: None,
594 notice: None,
595 })
596 }
597 SearchCommand::Text {
598 query,
599 collection,
600 limit,
601 fuzzy,
602 limit_param,
603 } => {
604 if limit_param.is_some() {
605 return Err(RedDBError::Query(
606 "SEARCH TEXT LIMIT $N parameter was not bound before execution".to_string(),
607 ));
608 }
609 let collections = collection.as_ref().map(|c| vec![c.clone()]);
610 let scope = self.ai_scope();
612 let res =
613 if super::statement_frame::ReadFrame::visible_collections(&scope).is_some() {
614 crate::runtime::authorized_search::AuthorizedSearch::execute_text(
615 self,
616 &scope,
617 query.clone(),
618 collections,
619 None,
620 None,
621 None,
622 Some(*limit),
623 *fuzzy,
624 )?
625 } else {
626 self.search_text(
627 query.clone(),
628 collections,
629 None,
630 None,
631 None,
632 Some(*limit),
633 *fuzzy,
634 )?
635 };
636 let mut result =
637 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
638 for item in &res.matches {
639 let mut record = UnifiedRecord::new();
640 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
641 record.set("score", Value::Float(item.score as f64));
642 result.push(record);
643 }
644 Ok(RuntimeQueryResult {
645 query: raw_query.to_string(),
646 mode: QueryMode::Sql,
647 statement: "search_text",
648 engine: "runtime-search",
649 result,
650 affected_rows: 0,
651 statement_type: "select",
652 bookmark: None,
653 notice: None,
654 })
655 }
656 SearchCommand::Hybrid {
657 vector,
658 query,
659 collection,
660 limit,
661 limit_param,
662 } => {
663 if limit_param.is_some() {
664 return Err(RedDBError::Query(
665 "SEARCH HYBRID LIMIT $N parameter was not bound before execution"
666 .to_string(),
667 ));
668 }
669 let res = self.search_hybrid(
670 vector.clone(),
671 query.clone(),
672 Some(*limit),
673 Some(vec![collection.clone()]),
674 None,
675 None,
676 None,
677 Vec::new(),
678 None,
679 None,
680 Some(*limit),
681 )?;
682 let mut result =
683 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
684 for item in &res.matches {
685 let mut record = UnifiedRecord::new();
686 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
687 record.set("score", Value::Float(item.score as f64));
688 result.push(record);
689 }
690 Ok(RuntimeQueryResult {
691 query: raw_query.to_string(),
692 mode: QueryMode::Sql,
693 statement: "search_hybrid",
694 engine: "runtime-search",
695 result,
696 affected_rows: 0,
697 statement_type: "select",
698 bookmark: None,
699 notice: None,
700 })
701 }
702 SearchCommand::Multimodal {
703 query,
704 collection,
705 limit,
706 limit_param,
707 } => {
708 if limit_param.is_some() {
709 return Err(RedDBError::Query(
710 "SEARCH MULTIMODAL LIMIT $N parameter was not bound before execution"
711 .to_string(),
712 ));
713 }
714 let collections = collection.as_ref().map(|c| vec![c.clone()]);
715 let res =
716 self.search_multimodal(query.clone(), collections, None, None, Some(*limit))?;
717 let mut result =
718 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
719 for item in &res.matches {
720 let mut record = UnifiedRecord::new();
721 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
722 record.set("score", Value::Float(item.score as f64));
723 result.push(record);
724 }
725 Ok(RuntimeQueryResult {
726 query: raw_query.to_string(),
727 mode: QueryMode::Sql,
728 statement: "search_multimodal",
729 engine: "runtime-search",
730 result,
731 affected_rows: 0,
732 statement_type: "select",
733 bookmark: None,
734 notice: None,
735 })
736 }
737 SearchCommand::Index {
738 index,
739 value,
740 collection,
741 limit,
742 exact,
743 limit_param,
744 } => {
745 if limit_param.is_some() {
746 return Err(RedDBError::Query(
747 "SEARCH INDEX LIMIT $N parameter was not bound before execution"
748 .to_string(),
749 ));
750 }
751 let collections = collection.as_ref().map(|c| vec![c.clone()]);
752 let res = self.search_index(
753 index.clone(),
754 value.clone(),
755 *exact,
756 collections,
757 None,
758 None,
759 Some(*limit),
760 )?;
761 let mut result =
762 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
763 for item in &res.matches {
764 let mut record = UnifiedRecord::new();
765 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
766 record.set("score", Value::Float(item.score as f64));
767 result.push(record);
768 }
769 Ok(RuntimeQueryResult {
770 query: raw_query.to_string(),
771 mode: QueryMode::Sql,
772 statement: "search_index",
773 engine: "runtime-search",
774 result,
775 affected_rows: 0,
776 statement_type: "select",
777 bookmark: None,
778 notice: None,
779 })
780 }
781 SearchCommand::Context {
782 query,
783 field,
784 collection,
785 limit,
786 depth,
787 limit_param,
788 } => {
789 if limit_param.is_some() {
790 return Err(RedDBError::Query(
791 "SEARCH CONTEXT LIMIT $N parameter was not bound before execution"
792 .to_string(),
793 ));
794 }
795 use crate::application::SearchContextInput;
796 let input = SearchContextInput {
800 query: query.clone(),
801 field: field.clone(),
802 vector: None,
803 collections: collection.as_ref().map(|c| vec![c.clone()]),
804 graph_depth: Some(*depth),
805 graph_max_edges: None,
806 max_cross_refs: None,
807 follow_cross_refs: None,
808 expand_graph: None,
809 global_scan: None,
810 reindex: None,
811 limit: Some(*limit),
812 min_score: None,
813 };
814 let scope = self.ai_scope();
815 let res =
816 if super::statement_frame::ReadFrame::visible_collections(&scope).is_some() {
817 crate::runtime::authorized_search::AuthorizedSearch::execute_context(
818 self, &scope, input,
819 )?
820 } else {
821 self.search_context(input)?
822 };
823 let mut result = UnifiedResult::with_columns(vec![
824 "entity_id".into(),
825 "collection".into(),
826 "score".into(),
827 "discovery".into(),
828 "kind".into(),
829 ]);
830 let all_entities = res
831 .tables
832 .iter()
833 .map(|e| (e, "table"))
834 .chain(res.graph.nodes.iter().map(|e| (e, "graph_node")))
835 .chain(res.graph.edges.iter().map(|e| (e, "graph_edge")))
836 .chain(res.vectors.iter().map(|e| (e, "vector")))
837 .chain(res.documents.iter().map(|e| (e, "document")))
838 .chain(res.key_values.iter().map(|e| (e, "kv")));
839 for (entity, kind) in all_entities {
840 let mut record = UnifiedRecord::new();
841 record.set("entity_id", Value::UnsignedInteger(entity.entity.id.raw()));
842 record.set("collection", Value::text(entity.collection.clone()));
843 record.set("score", Value::Float(entity.score as f64));
844 record.set("discovery", Value::text(format!("{:?}", entity.discovery)));
845 record.set("kind", Value::text(kind.to_string()));
846 result.push(record);
847 }
848 Ok(RuntimeQueryResult {
849 query: raw_query.to_string(),
850 mode: QueryMode::Sql,
851 statement: "search_context",
852 engine: "runtime-context",
853 result,
854 affected_rows: 0,
855 statement_type: "select",
856 bookmark: None,
857 notice: None,
858 })
859 }
860 SearchCommand::SpatialRadius {
861 center_lat,
862 center_lon,
863 radius_km,
864 collection,
865 column,
866 limit,
867 limit_param,
868 } => {
869 if limit_param.is_some() {
870 return Err(RedDBError::Query(
871 "SEARCH SPATIAL RADIUS LIMIT $N parameter was not bound before execution"
872 .to_string(),
873 ));
874 }
875 use crate::geo::haversine_km;
876 let h3_candidates = self.h3_radius_candidate_ids(
883 collection,
884 column,
885 *center_lat,
886 *center_lon,
887 *radius_km,
888 );
889 let store = self.inner.db.store();
890 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
891
892 let mut hits: Vec<(u64, f64)> = Vec::new();
893 let mut geo_values_seen = 0usize;
894 for entity in &entities {
895 if let Some((lat, lon)) =
896 self.extract_spatial_column(entity, collection, column)
897 {
898 geo_values_seen += 1;
899 let dist = haversine_km(*center_lat, *center_lon, lat, lon);
900 if dist <= *radius_km {
901 hits.push((entity.id.raw(), dist));
902 }
903 }
904 }
905 partial_top_k(&mut hits, *limit, |a, b| {
906 a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
907 });
908
909 let mut result =
910 UnifiedResult::with_columns(vec!["entity_id".into(), "distance_km".into()]);
911 for (id, dist) in &hits {
912 let mut record = UnifiedRecord::new();
913 record.set("entity_id", Value::UnsignedInteger(*id));
914 record.set("distance_km", Value::Float(*dist));
915 result.push(record);
916 }
917 let notice = self.spatial_zero_geo_notice(
918 collection,
919 column,
920 &entities,
921 geo_values_seen,
922 result.records.is_empty(),
923 );
924 Ok(RuntimeQueryResult {
925 query: raw_query.to_string(),
926 mode: QueryMode::Sql,
927 statement: "search_spatial_radius",
928 engine: "runtime-spatial",
929 result,
930 affected_rows: 0,
931 statement_type: "select",
932 bookmark: None,
933 notice,
934 })
935 }
936 SearchCommand::SpatialBbox {
937 min_lat,
938 min_lon,
939 max_lat,
940 max_lon,
941 collection,
942 column,
943 limit,
944 limit_param,
945 } => {
946 if limit_param.is_some() {
947 return Err(RedDBError::Query(
948 "SEARCH SPATIAL BBOX LIMIT $N parameter was not bound before execution"
949 .to_string(),
950 ));
951 }
952 let h3_candidates = self.h3_bbox_candidate_ids(
957 collection, column, *min_lat, *min_lon, *max_lat, *max_lon,
958 );
959 let store = self.inner.db.store();
960 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
961
962 let mut result = UnifiedResult::with_columns(vec!["entity_id".into()]);
963 let mut count = 0;
964 let mut geo_values_seen = 0usize;
965 for entity in &entities {
966 if count >= *limit {
967 break;
968 }
969 if let Some((lat, lon)) =
970 self.extract_spatial_column(entity, collection, column)
971 {
972 geo_values_seen += 1;
973 if lat >= *min_lat && lat <= *max_lat && lon >= *min_lon && lon <= *max_lon
974 {
975 let mut record = UnifiedRecord::new();
976 record.set("entity_id", Value::UnsignedInteger(entity.id.raw()));
977 result.push(record);
978 count += 1;
979 }
980 }
981 }
982 let notice = self.spatial_zero_geo_notice(
983 collection,
984 column,
985 &entities,
986 geo_values_seen,
987 result.records.is_empty(),
988 );
989 Ok(RuntimeQueryResult {
990 query: raw_query.to_string(),
991 mode: QueryMode::Sql,
992 statement: "search_spatial_bbox",
993 engine: "runtime-spatial",
994 result,
995 affected_rows: 0,
996 statement_type: "select",
997 bookmark: None,
998 notice,
999 })
1000 }
1001 SearchCommand::SpatialWithinPolygon {
1002 vertices,
1003 collection,
1004 column,
1005 limit,
1006 limit_param,
1007 } => {
1008 if limit_param.is_some() {
1009 return Err(RedDBError::Query(
1010 "SEARCH SPATIAL WITHIN POLYGON LIMIT $N parameter was not bound before execution"
1011 .to_string(),
1012 ));
1013 }
1014 let h3_candidates = self.h3_polygon_candidate_ids(collection, column, vertices);
1015 let store = self.inner.db.store();
1016 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
1017
1018 let mut result = UnifiedResult::with_columns(vec!["entity_id".into()]);
1019 let mut count = 0;
1020 for entity in &entities {
1021 if count >= *limit {
1022 break;
1023 }
1024 if let Some((lat, lon)) =
1025 self.extract_spatial_column(entity, collection, column)
1026 {
1027 if crate::geo::point_in_polygon_even_odd(lat, lon, vertices) {
1028 let mut record = UnifiedRecord::new();
1029 record.set("entity_id", Value::UnsignedInteger(entity.id.raw()));
1030 result.push(record);
1031 count += 1;
1032 }
1033 }
1034 }
1035 Ok(RuntimeQueryResult {
1036 query: raw_query.to_string(),
1037 mode: QueryMode::Sql,
1038 statement: "search_spatial_within_polygon",
1039 engine: "runtime-spatial",
1040 result,
1041 affected_rows: 0,
1042 statement_type: "select",
1043 bookmark: None,
1044 notice: None,
1045 })
1046 }
1047 SearchCommand::SpatialNearest {
1048 lat,
1049 lon,
1050 k,
1051 collection,
1052 column,
1053 k_param,
1054 } => {
1055 if k_param.is_some() {
1056 return Err(RedDBError::Query(
1057 "SEARCH SPATIAL NEAREST K $N parameter was not bound before execution"
1058 .to_string(),
1059 ));
1060 }
1061 use crate::geo::haversine_km;
1062 let h3_candidates =
1069 self.h3_nearest_candidate_ids(collection, column, *lat, *lon, *k);
1070 let store = self.inner.db.store();
1071 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
1072
1073 let mut hits: Vec<(u64, f64)> = Vec::new();
1074 let mut geo_values_seen = 0usize;
1075 for entity in &entities {
1076 if let Some((elat, elon)) =
1077 self.extract_spatial_column(entity, collection, column)
1078 {
1079 geo_values_seen += 1;
1080 let dist = haversine_km(*lat, *lon, elat, elon);
1081 hits.push((entity.id.raw(), dist));
1082 }
1083 }
1084 partial_top_k(&mut hits, *k, |a, b| {
1085 a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
1086 });
1087
1088 let mut result =
1089 UnifiedResult::with_columns(vec!["entity_id".into(), "distance_km".into()]);
1090 for (id, dist) in &hits {
1091 let mut record = UnifiedRecord::new();
1092 record.set("entity_id", Value::UnsignedInteger(*id));
1093 record.set("distance_km", Value::Float(*dist));
1094 result.push(record);
1095 }
1096 let notice = self.spatial_zero_geo_notice(
1097 collection,
1098 column,
1099 &entities,
1100 geo_values_seen,
1101 result.records.is_empty(),
1102 );
1103 Ok(RuntimeQueryResult {
1104 query: raw_query.to_string(),
1105 mode: QueryMode::Sql,
1106 statement: "search_spatial_nearest",
1107 engine: "runtime-spatial",
1108 result,
1109 affected_rows: 0,
1110 statement_type: "select",
1111 bookmark: None,
1112 notice,
1113 })
1114 }
1115 }
1116 }
1117}
1118
1119fn partial_top_k<T: Clone>(items: &mut Vec<T>, k: usize, cmp: impl Fn(&T, &T) -> Ordering) {
1127 let n = items.len();
1128 if k == 0 {
1129 items.clear();
1130 return;
1131 }
1132 if n <= k.saturating_mul(2) {
1133 items.sort_by(|a, b| cmp(a, b));
1134 items.truncate(k);
1135 return;
1136 }
1137 let mut idxs: Vec<usize> = (0..n).collect();
1138 idxs.select_nth_unstable_by(k - 1, |&a, &b| {
1139 cmp(&items[a], &items[b]).then_with(|| a.cmp(&b))
1140 });
1141 idxs.truncate(k);
1142 idxs.sort_by(|&a, &b| cmp(&items[a], &items[b]).then_with(|| a.cmp(&b)));
1143 let orig = std::mem::take(items);
1144 *items = idxs.into_iter().map(|i| orig[i].clone()).collect();
1145}
1146
1147fn apply_graph_order_and_limit(
1148 result: &mut UnifiedResult,
1149 statement: &str,
1150 order_by: Option<&GraphCommandOrderBy>,
1151 limit: Option<usize>,
1152) -> RedDBResult<()> {
1153 if let Some(order) = order_by {
1154 let column = graph_order_metric_column(statement, &order.metric)?;
1155 let columns = result.columns.clone();
1156 let cmp = |left: &UnifiedRecord, right: &UnifiedRecord| {
1157 let cmp = compare_graph_values(left.get(column), right.get(column));
1158 let cmp = if order.ascending { cmp } else { cmp.reverse() };
1159 if cmp == Ordering::Equal {
1160 compare_graph_rows(left, right, &columns)
1161 } else {
1162 cmp
1163 }
1164 };
1165 match limit {
1167 Some(limit) => partial_top_k(&mut result.records, limit, cmp),
1168 None => result.records.sort_by(|left, right| cmp(left, right)),
1169 }
1170 } else if let Some(limit) = limit {
1171 result.records.truncate(limit);
1172 }
1173 Ok(())
1174}
1175
1176fn graph_order_metric_column(statement: &str, metric: &str) -> RedDBResult<&'static str> {
1177 let metric = metric.to_ascii_lowercase();
1178 match (statement, metric.as_str()) {
1179 ("graph_centrality", "score" | "centrality_score") => Ok("score"),
1180 ("graph_community", "size" | "community_size") => Ok("size"),
1181 ("graph_components", "size" | "component_size") => Ok("size"),
1182 ("graph_shortest_path", "hop_count" | "total_weight" | "nodes_visited") => {
1183 Ok(match metric.as_str() {
1184 "total_weight" => "total_weight",
1185 "nodes_visited" => "nodes_visited",
1186 _ => "hop_count",
1187 })
1188 }
1189 _ => Err(RedDBError::Query(format!(
1190 "unsupported ORDER BY metric '{metric}' for GRAPH {}",
1191 statement.trim_start_matches("graph_")
1192 ))),
1193 }
1194}
1195
1196fn compare_graph_rows(left: &UnifiedRecord, right: &UnifiedRecord, columns: &[String]) -> Ordering {
1197 for column in columns {
1198 let cmp = compare_graph_values(left.get(column), right.get(column));
1199 if cmp != Ordering::Equal {
1200 return cmp;
1201 }
1202 }
1203 Ordering::Equal
1204}
1205
1206fn compare_graph_values(left: Option<&Value>, right: Option<&Value>) -> Ordering {
1207 match (left, right) {
1208 (None, None) => Ordering::Equal,
1209 (None, Some(_)) => Ordering::Less,
1210 (Some(_), None) => Ordering::Greater,
1211 (Some(Value::Null), Some(Value::Null)) => Ordering::Equal,
1212 (Some(Value::Null), Some(_)) => Ordering::Less,
1213 (Some(_), Some(Value::Null)) => Ordering::Greater,
1214 (Some(Value::Integer(left)), Some(Value::Integer(right))) => left.cmp(right),
1215 (Some(Value::UnsignedInteger(left)), Some(Value::UnsignedInteger(right))) => {
1216 left.cmp(right)
1217 }
1218 (Some(Value::Float(left)), Some(Value::Float(right))) => {
1219 left.partial_cmp(right).unwrap_or(Ordering::Equal)
1220 }
1221 (Some(Value::Integer(left)), Some(Value::Float(right))) => {
1222 (*left as f64).partial_cmp(right).unwrap_or(Ordering::Equal)
1223 }
1224 (Some(Value::Float(left)), Some(Value::Integer(right))) => left
1225 .partial_cmp(&(*right as f64))
1226 .unwrap_or(Ordering::Equal),
1227 (Some(Value::UnsignedInteger(left)), Some(Value::Float(right))) => {
1228 (*left as f64).partial_cmp(right).unwrap_or(Ordering::Equal)
1229 }
1230 (Some(Value::Float(left)), Some(Value::UnsignedInteger(right))) => left
1231 .partial_cmp(&(*right as f64))
1232 .unwrap_or(Ordering::Equal),
1233 (Some(Value::Integer(left)), Some(Value::UnsignedInteger(right))) => {
1234 (*left as i128).cmp(&(*right as i128))
1235 }
1236 (Some(Value::UnsignedInteger(left)), Some(Value::Integer(right))) => {
1237 (*left as i128).cmp(&(*right as i128))
1238 }
1239 (Some(Value::Timestamp(left)), Some(Value::Timestamp(right))) => left.cmp(right),
1240 (Some(Value::Text(left)), Some(Value::Text(right))) => left.cmp(right),
1241 (Some(Value::Boolean(left)), Some(Value::Boolean(right))) => left.cmp(right),
1242 (Some(left), Some(right)) => format!("{left:?}").cmp(&format!("{right:?}")),
1243 }
1244}
1245
1246fn parse_direction(s: &str) -> RedDBResult<RuntimeGraphDirection> {
1251 match s.to_lowercase().as_str() {
1252 "outgoing" | "out" => Ok(RuntimeGraphDirection::Outgoing),
1253 "incoming" | "in" => Ok(RuntimeGraphDirection::Incoming),
1254 "both" | "any" => Ok(RuntimeGraphDirection::Both),
1255 _ => Err(RedDBError::Query(format!(
1256 "unknown direction: '{s}', expected outgoing|incoming|both"
1257 ))),
1258 }
1259}
1260
1261fn parse_path_algorithm(s: &str) -> RedDBResult<RuntimeGraphPathAlgorithm> {
1262 match s.to_lowercase().as_str() {
1263 "bfs" => Ok(RuntimeGraphPathAlgorithm::Bfs),
1264 "dijkstra" => Ok(RuntimeGraphPathAlgorithm::Dijkstra),
1265 "astar" | "a*" => Ok(RuntimeGraphPathAlgorithm::AStar),
1266 "bellman_ford" | "bellmanford" => Ok(RuntimeGraphPathAlgorithm::BellmanFord),
1267 _ => Err(RedDBError::Query(format!(
1268 "unknown path algorithm: '{s}', expected bfs|dijkstra|astar|bellman_ford"
1269 ))),
1270 }
1271}
1272
1273fn parse_traversal_strategy(s: &str) -> RedDBResult<RuntimeGraphTraversalStrategy> {
1274 match s.to_lowercase().as_str() {
1275 "bfs" => Ok(RuntimeGraphTraversalStrategy::Bfs),
1276 "dfs" => Ok(RuntimeGraphTraversalStrategy::Dfs),
1277 _ => Err(RedDBError::Query(format!(
1278 "unknown traversal strategy: '{s}', expected bfs|dfs"
1279 ))),
1280 }
1281}
1282
1283fn parse_centrality_algorithm(s: &str) -> RedDBResult<RuntimeGraphCentralityAlgorithm> {
1284 match s.to_lowercase().as_str() {
1285 "degree" => Ok(RuntimeGraphCentralityAlgorithm::Degree),
1286 "closeness" => Ok(RuntimeGraphCentralityAlgorithm::Closeness),
1287 "betweenness" => Ok(RuntimeGraphCentralityAlgorithm::Betweenness),
1288 "eigenvector" => Ok(RuntimeGraphCentralityAlgorithm::Eigenvector),
1289 "pagerank" | "page_rank" => Ok(RuntimeGraphCentralityAlgorithm::PageRank),
1290 _ => Err(RedDBError::Query(format!(
1291 "unknown centrality algorithm: '{s}', expected degree|closeness|betweenness|eigenvector|pagerank"
1292 ))),
1293 }
1294}
1295
1296fn parse_community_algorithm(s: &str) -> RedDBResult<RuntimeGraphCommunityAlgorithm> {
1297 match s.to_lowercase().as_str() {
1298 "label_propagation" | "labelpropagation" => {
1299 Ok(RuntimeGraphCommunityAlgorithm::LabelPropagation)
1300 }
1301 "louvain" => Ok(RuntimeGraphCommunityAlgorithm::Louvain),
1302 _ => Err(RedDBError::Query(format!(
1303 "unknown community algorithm: '{s}', expected label_propagation|louvain"
1304 ))),
1305 }
1306}
1307
1308fn parse_components_mode(s: &str) -> RedDBResult<RuntimeGraphComponentsMode> {
1309 match s.to_lowercase().as_str() {
1310 "connected" => Ok(RuntimeGraphComponentsMode::Connected),
1311 "weak" | "weakly_connected" => Ok(RuntimeGraphComponentsMode::Weak),
1312 "strong" | "strongly_connected" => Ok(RuntimeGraphComponentsMode::Strong),
1313 _ => Err(RedDBError::Query(format!(
1314 "unknown components mode: '{s}', expected connected|weak|strong"
1315 ))),
1316 }
1317}
1318
1319fn extract_geo_from_entity(entity: &UnifiedEntity) -> Option<(f64, f64)> {
1324 match &entity.data {
1325 EntityData::Row(row) => {
1326 if let Some(ref named) = row.named {
1327 for value in named.values() {
1328 if let Some(point) = crate::geo::recognize_geo_value(value) {
1329 return Some(point);
1330 }
1331 }
1332 if let Some(point) = crate::geo::recognize_geo_fields(|key| named.get(key)) {
1333 return Some(point);
1334 }
1335 }
1336 for value in &row.columns {
1337 if let Some(point) = crate::geo::recognize_geo_value(value) {
1338 return Some(point);
1339 }
1340 }
1341 None
1342 }
1343 EntityData::Node(node) => {
1344 for value in node.properties.values() {
1345 if let Some(point) = crate::geo::recognize_geo_value(value) {
1346 return Some(point);
1347 }
1348 }
1349 crate::geo::recognize_geo_fields(|key| node.properties.get(key))
1350 }
1351 _ => None,
1352 }
1353}
1354
1355fn index_fields_from_entity(entity: &UnifiedEntity) -> Vec<(String, Value)> {
1356 match &entity.data {
1357 EntityData::Row(row) => row
1358 .iter_fields()
1359 .map(|(field, value)| (field.to_string(), value.clone()))
1360 .collect(),
1361 EntityData::Node(node) => node
1362 .properties
1363 .iter()
1364 .map(|(field, value)| (field.clone(), value.clone()))
1365 .collect(),
1366 EntityData::Edge(edge) => edge
1367 .properties
1368 .iter()
1369 .map(|(field, value)| (field.clone(), value.clone()))
1370 .collect(),
1371 EntityData::TimeSeries(ts) => ts
1372 .fields
1373 .iter()
1374 .map(|(field, value)| (field.clone(), value.clone()))
1375 .collect(),
1376 _ => Vec::new(),
1377 }
1378}
1379
1380fn scan_collection_with_candidates(
1386 store: &std::sync::Arc<crate::storage::unified::store::UnifiedStore>,
1387 collection: &str,
1388 candidates: &Option<std::collections::HashSet<u64>>,
1389) -> Vec<UnifiedEntity> {
1390 let snapshot = crate::runtime::impl_core::capture_current_snapshot();
1391 match candidates {
1392 Some(ids) => store
1393 .get_collection(collection)
1394 .map(|manager| manager.scan(snapshot.as_ref(), |e| ids.contains(&e.id.raw())))
1395 .unwrap_or_default(),
1396 None => store
1397 .get_collection(collection)
1398 .map(|manager| manager.scan(snapshot.as_ref(), |_| true))
1399 .unwrap_or_default(),
1400 }
1401}
1402
1403fn h3_cover_cells(lat: f64, lon: f64, radius_km: f64, resolution: u8) -> Vec<u64> {
1412 let cell = crate::geo::h3::lat_lng_to_cell(lat, lon, resolution);
1413 if cell == 0 {
1414 return Vec::new();
1415 }
1416 let Some(k) = crate::geo::h3::radius_cover_ring(radius_km, resolution) else {
1419 return Vec::new();
1420 };
1421 crate::geo::h3::grid_disk(cell, k)
1422}
1423
1424impl RedDBRuntime {
1425 fn extract_spatial_column(
1426 &self,
1427 entity: &UnifiedEntity,
1428 collection: &str,
1429 column: &str,
1430 ) -> Option<(f64, f64)> {
1431 let fields = index_fields_from_entity(entity);
1432 if let Some(value) = self
1433 .inner
1434 .index_store
1435 .resolve_index_field_value(&fields, column)
1436 {
1437 return crate::geo::recognize_geo_value(&value);
1438 }
1439
1440 let is_document_collection =
1441 self.inner
1442 .db
1443 .collection_contract(collection)
1444 .is_some_and(|contract| {
1445 contract.declared_model == crate::catalog::CollectionModel::Document
1446 });
1447 if is_document_collection {
1448 return None;
1449 }
1450
1451 match entity.data {
1452 EntityData::Row(_) | EntityData::Node(_) => extract_geo_from_entity(entity),
1453 _ => None,
1454 }
1455 }
1456
1457 fn spatial_zero_geo_notice(
1458 &self,
1459 collection: &str,
1460 column: &str,
1461 scanned_entities: &[UnifiedEntity],
1462 geo_values_seen: usize,
1463 result_is_empty: bool,
1464 ) -> Option<String> {
1465 if !result_is_empty || geo_values_seen > 0 {
1466 return None;
1467 }
1468 let (entity_count, geo_count) = if scanned_entities.is_empty() {
1469 let store = self.inner.db.store();
1470 let entities = scan_collection_with_candidates(&store, collection, &None);
1471 let geo_count = entities
1472 .iter()
1473 .filter(|entity| {
1474 self.extract_spatial_column(entity, collection, column)
1475 .is_some()
1476 })
1477 .count();
1478 (entities.len(), geo_count)
1479 } else {
1480 (scanned_entities.len(), geo_values_seen)
1481 };
1482 if entity_count == 0 || geo_count > 0 {
1483 return None;
1484 }
1485 Some(format!(
1486 "no entity in '{}' has an indexable geo value in column '{}' (expected {}).",
1487 collection,
1488 column,
1489 crate::geo::RECOGNIZED_GEO_SHAPES
1490 ))
1491 }
1492
1493 fn h3_index_resolution(&self, collection: &str, column: &str) -> Option<u8> {
1496 match self
1497 .inner
1498 .index_store
1499 .find_index_for_column(collection, column)?
1500 .method
1501 {
1502 super::index_store::IndexMethodKind::H3 { resolution } => Some(resolution),
1503 _ => None,
1504 }
1505 }
1506
1507 fn h3_cell_candidate_ids(
1511 &self,
1512 collection: &str,
1513 column: &str,
1514 cells: &[u64],
1515 ) -> Option<std::collections::HashSet<u64>> {
1516 if cells.is_empty() {
1517 return None;
1518 }
1519 let keys: Vec<_> = cells
1520 .iter()
1521 .filter_map(|c| reddb_types::value_to_canonical_key(&Value::UnsignedInteger(*c)))
1522 .collect();
1523 let ids = self.inner.index_store.sorted.in_lookup_limited(
1524 collection,
1525 column,
1526 &keys,
1527 usize::MAX,
1528 )?;
1529 Some(ids.into_iter().map(|id| id.raw()).collect())
1530 }
1531
1532 fn h3_radius_candidate_ids(
1536 &self,
1537 collection: &str,
1538 column: &str,
1539 center_lat: f64,
1540 center_lon: f64,
1541 radius_km: f64,
1542 ) -> Option<std::collections::HashSet<u64>> {
1543 let resolution = self.h3_index_resolution(collection, column)?;
1544 let cells = h3_cover_cells(center_lat, center_lon, radius_km, resolution);
1545 self.h3_cell_candidate_ids(collection, column, &cells)
1546 }
1547
1548 fn h3_bbox_candidate_ids(
1552 &self,
1553 collection: &str,
1554 column: &str,
1555 min_lat: f64,
1556 min_lon: f64,
1557 max_lat: f64,
1558 max_lon: f64,
1559 ) -> Option<std::collections::HashSet<u64>> {
1560 let resolution = self.h3_index_resolution(collection, column)?;
1561 let center_lat = (min_lat + max_lat) / 2.0;
1562 let center_lon = (min_lon + max_lon) / 2.0;
1563 let radius_km = [
1564 (min_lat, min_lon),
1565 (min_lat, max_lon),
1566 (max_lat, min_lon),
1567 (max_lat, max_lon),
1568 ]
1569 .into_iter()
1570 .map(|(la, lo)| crate::geo::haversine_km(center_lat, center_lon, la, lo))
1571 .fold(0.0_f64, f64::max);
1572 let cells = h3_cover_cells(center_lat, center_lon, radius_km, resolution);
1573 self.h3_cell_candidate_ids(collection, column, &cells)
1574 }
1575
1576 fn h3_polygon_candidate_ids(
1581 &self,
1582 collection: &str,
1583 column: &str,
1584 vertices: &[(f64, f64)],
1585 ) -> Option<std::collections::HashSet<u64>> {
1586 let resolution = self.h3_index_resolution(collection, column)?;
1587 let cells = crate::geo::h3::polygon_to_cover_cells(
1588 vertices,
1589 resolution,
1590 crate::geo::h3::MAX_POLYGON_COVER_CELLS,
1591 )?;
1592 self.h3_cell_candidate_ids(collection, column, &cells)
1593 }
1594
1595 fn h3_nearest_candidate_ids(
1601 &self,
1602 collection: &str,
1603 column: &str,
1604 lat: f64,
1605 lon: f64,
1606 k: usize,
1607 ) -> Option<std::collections::HashSet<u64>> {
1608 let resolution = self.h3_index_resolution(collection, column)?;
1609 if k == 0 {
1610 return None;
1611 }
1612 let center = crate::geo::h3::lat_lng_to_cell(lat, lon, resolution);
1613 if center == 0 {
1614 return None;
1615 }
1616 let edge_km = crate::geo::h3::edge_length_km(resolution).max(f64::MIN_POSITIVE);
1617 const MAX_RING: u32 = 64;
1620 let store = self.inner.db.store();
1621 for r in 0..=MAX_RING {
1622 let cells = crate::geo::h3::grid_disk(center, r);
1623 let Some(ids) = self.h3_cell_candidate_ids(collection, column, &cells) else {
1624 continue;
1625 };
1626 if ids.len() < k {
1627 continue;
1628 }
1629 let candidates = Some(ids);
1630 let mut dists: Vec<f64> =
1631 scan_collection_with_candidates(&store, collection, &candidates)
1632 .iter()
1633 .filter_map(|e| {
1634 extract_geo_from_entity(e)
1635 .map(|(elat, elon)| crate::geo::haversine_km(lat, lon, elat, elon))
1636 })
1637 .collect();
1638 if dists.len() < k {
1639 continue;
1640 }
1641 dists.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1642 let d_k = dists[k - 1];
1643 if f64::from(r) * edge_km >= d_k {
1652 return candidates;
1653 }
1654 }
1655 None
1656 }
1657}