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 crate::ai::AiConnection { api_key, api_base } =
540 crate::ai::resolve_connection_from_runtime(&provider, None, self)?;
541 let model = crate::ai::resolve_embeddings_model(&provider, &kv_getter);
542 let transport = crate::runtime::ai::transport::AiTransport::from_runtime(self);
543 let request = crate::ai::OpenAiEmbeddingRequest {
544 api_key,
545 model,
546 inputs: vec![query_text.clone()],
547 dimensions: None,
548 api_base,
549 };
550 let response = crate::runtime::ai::block_on_ai(async move {
551 crate::ai::openai_embeddings_async(&transport, request).await
552 })
553 .and_then(|result| result)?;
554 response.embeddings.into_iter().next().ok_or_else(|| {
555 RedDBError::Query("embedding API returned no vectors".to_string())
556 })?
557 } else {
558 vector.clone()
559 };
560 let scope = self.ai_scope();
564 let results =
565 if super::statement_frame::ReadFrame::visible_collections(&scope).is_some() {
566 crate::runtime::authorized_search::AuthorizedSearch::execute_similar(
567 self,
568 &scope,
569 collection,
570 &search_vector,
571 *limit,
572 *min_score,
573 )?
574 } else {
575 self.search_similar(collection, &search_vector, *limit, *min_score)?
577 };
578 let mut result =
579 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
580 for sr in &results {
581 let mut record = UnifiedRecord::new();
582 record.set("entity_id", Value::UnsignedInteger(sr.entity_id.raw()));
583 record.set("score", Value::Float(sr.score as f64));
584 result.push(record);
585 }
586 Ok(RuntimeQueryResult {
587 query: raw_query.to_string(),
588 mode: QueryMode::Sql,
589 statement: "search_similar",
590 engine: "runtime-search",
591 result,
592 affected_rows: 0,
593 statement_type: "select",
594 bookmark: None,
595 notice: None,
596 })
597 }
598 SearchCommand::Text {
599 query,
600 collection,
601 limit,
602 fuzzy,
603 limit_param,
604 } => {
605 if limit_param.is_some() {
606 return Err(RedDBError::Query(
607 "SEARCH TEXT LIMIT $N parameter was not bound before execution".to_string(),
608 ));
609 }
610 let collections = collection.as_ref().map(|c| vec![c.clone()]);
611 let scope = self.ai_scope();
613 let res =
614 if super::statement_frame::ReadFrame::visible_collections(&scope).is_some() {
615 crate::runtime::authorized_search::AuthorizedSearch::execute_text(
616 self,
617 &scope,
618 query.clone(),
619 collections,
620 None,
621 None,
622 None,
623 Some(*limit),
624 *fuzzy,
625 )?
626 } else {
627 self.search_text(
628 query.clone(),
629 collections,
630 None,
631 None,
632 None,
633 Some(*limit),
634 *fuzzy,
635 )?
636 };
637 let mut result =
638 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
639 for item in &res.matches {
640 let mut record = UnifiedRecord::new();
641 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
642 record.set("score", Value::Float(item.score as f64));
643 result.push(record);
644 }
645 Ok(RuntimeQueryResult {
646 query: raw_query.to_string(),
647 mode: QueryMode::Sql,
648 statement: "search_text",
649 engine: "runtime-search",
650 result,
651 affected_rows: 0,
652 statement_type: "select",
653 bookmark: None,
654 notice: None,
655 })
656 }
657 SearchCommand::Hybrid {
658 vector,
659 query,
660 collection,
661 limit,
662 limit_param,
663 } => {
664 if limit_param.is_some() {
665 return Err(RedDBError::Query(
666 "SEARCH HYBRID LIMIT $N parameter was not bound before execution"
667 .to_string(),
668 ));
669 }
670 let res = self.search_hybrid(
671 vector.clone(),
672 query.clone(),
673 Some(*limit),
674 Some(vec![collection.clone()]),
675 None,
676 None,
677 None,
678 Vec::new(),
679 None,
680 None,
681 Some(*limit),
682 )?;
683 let mut result =
684 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
685 for item in &res.matches {
686 let mut record = UnifiedRecord::new();
687 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
688 record.set("score", Value::Float(item.score as f64));
689 result.push(record);
690 }
691 Ok(RuntimeQueryResult {
692 query: raw_query.to_string(),
693 mode: QueryMode::Sql,
694 statement: "search_hybrid",
695 engine: "runtime-search",
696 result,
697 affected_rows: 0,
698 statement_type: "select",
699 bookmark: None,
700 notice: None,
701 })
702 }
703 SearchCommand::Multimodal {
704 query,
705 collection,
706 limit,
707 limit_param,
708 } => {
709 if limit_param.is_some() {
710 return Err(RedDBError::Query(
711 "SEARCH MULTIMODAL LIMIT $N parameter was not bound before execution"
712 .to_string(),
713 ));
714 }
715 let collections = collection.as_ref().map(|c| vec![c.clone()]);
716 let res =
717 self.search_multimodal(query.clone(), collections, None, None, Some(*limit))?;
718 let mut result =
719 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
720 for item in &res.matches {
721 let mut record = UnifiedRecord::new();
722 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
723 record.set("score", Value::Float(item.score as f64));
724 result.push(record);
725 }
726 Ok(RuntimeQueryResult {
727 query: raw_query.to_string(),
728 mode: QueryMode::Sql,
729 statement: "search_multimodal",
730 engine: "runtime-search",
731 result,
732 affected_rows: 0,
733 statement_type: "select",
734 bookmark: None,
735 notice: None,
736 })
737 }
738 SearchCommand::Index {
739 index,
740 value,
741 collection,
742 limit,
743 exact,
744 limit_param,
745 } => {
746 if limit_param.is_some() {
747 return Err(RedDBError::Query(
748 "SEARCH INDEX LIMIT $N parameter was not bound before execution"
749 .to_string(),
750 ));
751 }
752 let collections = collection.as_ref().map(|c| vec![c.clone()]);
753 let res = self.search_index(
754 index.clone(),
755 value.clone(),
756 *exact,
757 collections,
758 None,
759 None,
760 Some(*limit),
761 )?;
762 let mut result =
763 UnifiedResult::with_columns(vec!["entity_id".into(), "score".into()]);
764 for item in &res.matches {
765 let mut record = UnifiedRecord::new();
766 record.set("entity_id", Value::UnsignedInteger(item.entity.id.raw()));
767 record.set("score", Value::Float(item.score as f64));
768 result.push(record);
769 }
770 Ok(RuntimeQueryResult {
771 query: raw_query.to_string(),
772 mode: QueryMode::Sql,
773 statement: "search_index",
774 engine: "runtime-search",
775 result,
776 affected_rows: 0,
777 statement_type: "select",
778 bookmark: None,
779 notice: None,
780 })
781 }
782 SearchCommand::Context {
783 query,
784 field,
785 collection,
786 limit,
787 depth,
788 limit_param,
789 } => {
790 if limit_param.is_some() {
791 return Err(RedDBError::Query(
792 "SEARCH CONTEXT LIMIT $N parameter was not bound before execution"
793 .to_string(),
794 ));
795 }
796 use crate::application::SearchContextInput;
797 let input = SearchContextInput {
801 query: query.clone(),
802 field: field.clone(),
803 vector: None,
804 collections: collection.as_ref().map(|c| vec![c.clone()]),
805 graph_depth: Some(*depth),
806 graph_max_edges: None,
807 max_cross_refs: None,
808 follow_cross_refs: None,
809 expand_graph: None,
810 global_scan: None,
811 reindex: None,
812 limit: Some(*limit),
813 min_score: None,
814 };
815 let scope = self.ai_scope();
816 let res =
817 if super::statement_frame::ReadFrame::visible_collections(&scope).is_some() {
818 crate::runtime::authorized_search::AuthorizedSearch::execute_context(
819 self, &scope, input,
820 )?
821 } else {
822 self.search_context(input)?
823 };
824 let mut result = UnifiedResult::with_columns(vec![
825 "entity_id".into(),
826 "collection".into(),
827 "score".into(),
828 "discovery".into(),
829 "kind".into(),
830 ]);
831 let all_entities = res
832 .tables
833 .iter()
834 .map(|e| (e, "table"))
835 .chain(res.graph.nodes.iter().map(|e| (e, "graph_node")))
836 .chain(res.graph.edges.iter().map(|e| (e, "graph_edge")))
837 .chain(res.vectors.iter().map(|e| (e, "vector")))
838 .chain(res.documents.iter().map(|e| (e, "document")))
839 .chain(res.key_values.iter().map(|e| (e, "kv")));
840 for (entity, kind) in all_entities {
841 let mut record = UnifiedRecord::new();
842 record.set("entity_id", Value::UnsignedInteger(entity.entity.id.raw()));
843 record.set("collection", Value::text(entity.collection.clone()));
844 record.set("score", Value::Float(entity.score as f64));
845 record.set("discovery", Value::text(format!("{:?}", entity.discovery)));
846 record.set("kind", Value::text(kind.to_string()));
847 result.push(record);
848 }
849 Ok(RuntimeQueryResult {
850 query: raw_query.to_string(),
851 mode: QueryMode::Sql,
852 statement: "search_context",
853 engine: "runtime-context",
854 result,
855 affected_rows: 0,
856 statement_type: "select",
857 bookmark: None,
858 notice: None,
859 })
860 }
861 SearchCommand::SpatialRadius {
862 center_lat,
863 center_lon,
864 radius_km,
865 collection,
866 column,
867 limit,
868 limit_param,
869 } => {
870 if limit_param.is_some() {
871 return Err(RedDBError::Query(
872 "SEARCH SPATIAL RADIUS LIMIT $N parameter was not bound before execution"
873 .to_string(),
874 ));
875 }
876 use crate::geo::haversine_km;
877 let h3_candidates = self.h3_radius_candidate_ids(
884 collection,
885 column,
886 *center_lat,
887 *center_lon,
888 *radius_km,
889 );
890 let store = self.inner.db.store();
891 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
892
893 let mut hits: Vec<(u64, f64)> = Vec::new();
894 let mut geo_values_seen = 0usize;
895 for entity in &entities {
896 if let Some((lat, lon)) =
897 self.extract_spatial_column(entity, collection, column)
898 {
899 geo_values_seen += 1;
900 let dist = haversine_km(*center_lat, *center_lon, lat, lon);
901 if dist <= *radius_km {
902 hits.push((entity.id.raw(), dist));
903 }
904 }
905 }
906 partial_top_k(&mut hits, *limit, |a, b| {
907 a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
908 });
909
910 let mut result =
911 UnifiedResult::with_columns(vec!["entity_id".into(), "distance_km".into()]);
912 for (id, dist) in &hits {
913 let mut record = UnifiedRecord::new();
914 record.set("entity_id", Value::UnsignedInteger(*id));
915 record.set("distance_km", Value::Float(*dist));
916 result.push(record);
917 }
918 let notice = self.spatial_zero_geo_notice(
919 collection,
920 column,
921 &entities,
922 geo_values_seen,
923 result.records.is_empty(),
924 );
925 Ok(RuntimeQueryResult {
926 query: raw_query.to_string(),
927 mode: QueryMode::Sql,
928 statement: "search_spatial_radius",
929 engine: "runtime-spatial",
930 result,
931 affected_rows: 0,
932 statement_type: "select",
933 bookmark: None,
934 notice,
935 })
936 }
937 SearchCommand::SpatialBbox {
938 min_lat,
939 min_lon,
940 max_lat,
941 max_lon,
942 collection,
943 column,
944 limit,
945 limit_param,
946 } => {
947 if limit_param.is_some() {
948 return Err(RedDBError::Query(
949 "SEARCH SPATIAL BBOX LIMIT $N parameter was not bound before execution"
950 .to_string(),
951 ));
952 }
953 let h3_candidates = self.h3_bbox_candidate_ids(
958 collection, column, *min_lat, *min_lon, *max_lat, *max_lon,
959 );
960 let store = self.inner.db.store();
961 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
962
963 let mut result = UnifiedResult::with_columns(vec!["entity_id".into()]);
964 let mut count = 0;
965 let mut geo_values_seen = 0usize;
966 for entity in &entities {
967 if count >= *limit {
968 break;
969 }
970 if let Some((lat, lon)) =
971 self.extract_spatial_column(entity, collection, column)
972 {
973 geo_values_seen += 1;
974 if lat >= *min_lat && lat <= *max_lat && lon >= *min_lon && lon <= *max_lon
975 {
976 let mut record = UnifiedRecord::new();
977 record.set("entity_id", Value::UnsignedInteger(entity.id.raw()));
978 result.push(record);
979 count += 1;
980 }
981 }
982 }
983 let notice = self.spatial_zero_geo_notice(
984 collection,
985 column,
986 &entities,
987 geo_values_seen,
988 result.records.is_empty(),
989 );
990 Ok(RuntimeQueryResult {
991 query: raw_query.to_string(),
992 mode: QueryMode::Sql,
993 statement: "search_spatial_bbox",
994 engine: "runtime-spatial",
995 result,
996 affected_rows: 0,
997 statement_type: "select",
998 bookmark: None,
999 notice,
1000 })
1001 }
1002 SearchCommand::SpatialWithinPolygon {
1003 vertices,
1004 collection,
1005 column,
1006 limit,
1007 limit_param,
1008 } => {
1009 if limit_param.is_some() {
1010 return Err(RedDBError::Query(
1011 "SEARCH SPATIAL WITHIN POLYGON LIMIT $N parameter was not bound before execution"
1012 .to_string(),
1013 ));
1014 }
1015 let h3_candidates = self.h3_polygon_candidate_ids(collection, column, vertices);
1016 let store = self.inner.db.store();
1017 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
1018
1019 let mut result = UnifiedResult::with_columns(vec!["entity_id".into()]);
1020 let mut count = 0;
1021 for entity in &entities {
1022 if count >= *limit {
1023 break;
1024 }
1025 if let Some((lat, lon)) =
1026 self.extract_spatial_column(entity, collection, column)
1027 {
1028 if crate::geo::point_in_polygon_even_odd(lat, lon, vertices) {
1029 let mut record = UnifiedRecord::new();
1030 record.set("entity_id", Value::UnsignedInteger(entity.id.raw()));
1031 result.push(record);
1032 count += 1;
1033 }
1034 }
1035 }
1036 Ok(RuntimeQueryResult {
1037 query: raw_query.to_string(),
1038 mode: QueryMode::Sql,
1039 statement: "search_spatial_within_polygon",
1040 engine: "runtime-spatial",
1041 result,
1042 affected_rows: 0,
1043 statement_type: "select",
1044 bookmark: None,
1045 notice: None,
1046 })
1047 }
1048 SearchCommand::SpatialNearest {
1049 lat,
1050 lon,
1051 k,
1052 collection,
1053 column,
1054 k_param,
1055 } => {
1056 if k_param.is_some() {
1057 return Err(RedDBError::Query(
1058 "SEARCH SPATIAL NEAREST K $N parameter was not bound before execution"
1059 .to_string(),
1060 ));
1061 }
1062 use crate::geo::haversine_km;
1063 let h3_candidates =
1070 self.h3_nearest_candidate_ids(collection, column, *lat, *lon, *k);
1071 let store = self.inner.db.store();
1072 let entities = scan_collection_with_candidates(&store, collection, &h3_candidates);
1073
1074 let mut hits: Vec<(u64, f64)> = Vec::new();
1075 let mut geo_values_seen = 0usize;
1076 for entity in &entities {
1077 if let Some((elat, elon)) =
1078 self.extract_spatial_column(entity, collection, column)
1079 {
1080 geo_values_seen += 1;
1081 let dist = haversine_km(*lat, *lon, elat, elon);
1082 hits.push((entity.id.raw(), dist));
1083 }
1084 }
1085 partial_top_k(&mut hits, *k, |a, b| {
1086 a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
1087 });
1088
1089 let mut result =
1090 UnifiedResult::with_columns(vec!["entity_id".into(), "distance_km".into()]);
1091 for (id, dist) in &hits {
1092 let mut record = UnifiedRecord::new();
1093 record.set("entity_id", Value::UnsignedInteger(*id));
1094 record.set("distance_km", Value::Float(*dist));
1095 result.push(record);
1096 }
1097 let notice = self.spatial_zero_geo_notice(
1098 collection,
1099 column,
1100 &entities,
1101 geo_values_seen,
1102 result.records.is_empty(),
1103 );
1104 Ok(RuntimeQueryResult {
1105 query: raw_query.to_string(),
1106 mode: QueryMode::Sql,
1107 statement: "search_spatial_nearest",
1108 engine: "runtime-spatial",
1109 result,
1110 affected_rows: 0,
1111 statement_type: "select",
1112 bookmark: None,
1113 notice,
1114 })
1115 }
1116 }
1117 }
1118}
1119
1120fn partial_top_k<T: Clone>(items: &mut Vec<T>, k: usize, cmp: impl Fn(&T, &T) -> Ordering) {
1128 let n = items.len();
1129 if k == 0 {
1130 items.clear();
1131 return;
1132 }
1133 if n <= k.saturating_mul(2) {
1134 items.sort_by(|a, b| cmp(a, b));
1135 items.truncate(k);
1136 return;
1137 }
1138 let mut idxs: Vec<usize> = (0..n).collect();
1139 idxs.select_nth_unstable_by(k - 1, |&a, &b| {
1140 cmp(&items[a], &items[b]).then_with(|| a.cmp(&b))
1141 });
1142 idxs.truncate(k);
1143 idxs.sort_by(|&a, &b| cmp(&items[a], &items[b]).then_with(|| a.cmp(&b)));
1144 let orig = std::mem::take(items);
1145 *items = idxs.into_iter().map(|i| orig[i].clone()).collect();
1146}
1147
1148fn apply_graph_order_and_limit(
1149 result: &mut UnifiedResult,
1150 statement: &str,
1151 order_by: Option<&GraphCommandOrderBy>,
1152 limit: Option<usize>,
1153) -> RedDBResult<()> {
1154 if let Some(order) = order_by {
1155 let column = graph_order_metric_column(statement, &order.metric)?;
1156 let columns = result.columns.clone();
1157 let cmp = |left: &UnifiedRecord, right: &UnifiedRecord| {
1158 let cmp = compare_graph_values(left.get(column), right.get(column));
1159 let cmp = if order.ascending { cmp } else { cmp.reverse() };
1160 if cmp == Ordering::Equal {
1161 compare_graph_rows(left, right, &columns)
1162 } else {
1163 cmp
1164 }
1165 };
1166 match limit {
1168 Some(limit) => partial_top_k(&mut result.records, limit, cmp),
1169 None => result.records.sort_by(|left, right| cmp(left, right)),
1170 }
1171 } else if let Some(limit) = limit {
1172 result.records.truncate(limit);
1173 }
1174 Ok(())
1175}
1176
1177fn graph_order_metric_column(statement: &str, metric: &str) -> RedDBResult<&'static str> {
1178 let metric = metric.to_ascii_lowercase();
1179 match (statement, metric.as_str()) {
1180 ("graph_centrality", "score" | "centrality_score") => Ok("score"),
1181 ("graph_community", "size" | "community_size") => Ok("size"),
1182 ("graph_components", "size" | "component_size") => Ok("size"),
1183 ("graph_shortest_path", "hop_count" | "total_weight" | "nodes_visited") => {
1184 Ok(match metric.as_str() {
1185 "total_weight" => "total_weight",
1186 "nodes_visited" => "nodes_visited",
1187 _ => "hop_count",
1188 })
1189 }
1190 _ => Err(RedDBError::Query(format!(
1191 "unsupported ORDER BY metric '{metric}' for GRAPH {}",
1192 statement.trim_start_matches("graph_")
1193 ))),
1194 }
1195}
1196
1197fn compare_graph_rows(left: &UnifiedRecord, right: &UnifiedRecord, columns: &[String]) -> Ordering {
1198 for column in columns {
1199 let cmp = compare_graph_values(left.get(column), right.get(column));
1200 if cmp != Ordering::Equal {
1201 return cmp;
1202 }
1203 }
1204 Ordering::Equal
1205}
1206
1207fn compare_graph_values(left: Option<&Value>, right: Option<&Value>) -> Ordering {
1208 match (left, right) {
1209 (None, None) => Ordering::Equal,
1210 (None, Some(_)) => Ordering::Less,
1211 (Some(_), None) => Ordering::Greater,
1212 (Some(Value::Null), Some(Value::Null)) => Ordering::Equal,
1213 (Some(Value::Null), Some(_)) => Ordering::Less,
1214 (Some(_), Some(Value::Null)) => Ordering::Greater,
1215 (Some(Value::Integer(left)), Some(Value::Integer(right))) => left.cmp(right),
1216 (Some(Value::UnsignedInteger(left)), Some(Value::UnsignedInteger(right))) => {
1217 left.cmp(right)
1218 }
1219 (Some(Value::Float(left)), Some(Value::Float(right))) => {
1220 left.partial_cmp(right).unwrap_or(Ordering::Equal)
1221 }
1222 (Some(Value::Integer(left)), Some(Value::Float(right))) => {
1223 (*left as f64).partial_cmp(right).unwrap_or(Ordering::Equal)
1224 }
1225 (Some(Value::Float(left)), Some(Value::Integer(right))) => left
1226 .partial_cmp(&(*right as f64))
1227 .unwrap_or(Ordering::Equal),
1228 (Some(Value::UnsignedInteger(left)), Some(Value::Float(right))) => {
1229 (*left as f64).partial_cmp(right).unwrap_or(Ordering::Equal)
1230 }
1231 (Some(Value::Float(left)), Some(Value::UnsignedInteger(right))) => left
1232 .partial_cmp(&(*right as f64))
1233 .unwrap_or(Ordering::Equal),
1234 (Some(Value::Integer(left)), Some(Value::UnsignedInteger(right))) => {
1235 (*left as i128).cmp(&(*right as i128))
1236 }
1237 (Some(Value::UnsignedInteger(left)), Some(Value::Integer(right))) => {
1238 (*left as i128).cmp(&(*right as i128))
1239 }
1240 (Some(Value::Timestamp(left)), Some(Value::Timestamp(right))) => left.cmp(right),
1241 (Some(Value::Text(left)), Some(Value::Text(right))) => left.cmp(right),
1242 (Some(Value::Boolean(left)), Some(Value::Boolean(right))) => left.cmp(right),
1243 (Some(left), Some(right)) => format!("{left:?}").cmp(&format!("{right:?}")),
1244 }
1245}
1246
1247fn parse_direction(s: &str) -> RedDBResult<RuntimeGraphDirection> {
1252 match s.to_lowercase().as_str() {
1253 "outgoing" | "out" => Ok(RuntimeGraphDirection::Outgoing),
1254 "incoming" | "in" => Ok(RuntimeGraphDirection::Incoming),
1255 "both" | "any" => Ok(RuntimeGraphDirection::Both),
1256 _ => Err(RedDBError::Query(format!(
1257 "unknown direction: '{s}', expected outgoing|incoming|both"
1258 ))),
1259 }
1260}
1261
1262fn parse_path_algorithm(s: &str) -> RedDBResult<RuntimeGraphPathAlgorithm> {
1263 match s.to_lowercase().as_str() {
1264 "bfs" => Ok(RuntimeGraphPathAlgorithm::Bfs),
1265 "dijkstra" => Ok(RuntimeGraphPathAlgorithm::Dijkstra),
1266 "astar" | "a*" => Ok(RuntimeGraphPathAlgorithm::AStar),
1267 "bellman_ford" | "bellmanford" => Ok(RuntimeGraphPathAlgorithm::BellmanFord),
1268 _ => Err(RedDBError::Query(format!(
1269 "unknown path algorithm: '{s}', expected bfs|dijkstra|astar|bellman_ford"
1270 ))),
1271 }
1272}
1273
1274fn parse_traversal_strategy(s: &str) -> RedDBResult<RuntimeGraphTraversalStrategy> {
1275 match s.to_lowercase().as_str() {
1276 "bfs" => Ok(RuntimeGraphTraversalStrategy::Bfs),
1277 "dfs" => Ok(RuntimeGraphTraversalStrategy::Dfs),
1278 _ => Err(RedDBError::Query(format!(
1279 "unknown traversal strategy: '{s}', expected bfs|dfs"
1280 ))),
1281 }
1282}
1283
1284fn parse_centrality_algorithm(s: &str) -> RedDBResult<RuntimeGraphCentralityAlgorithm> {
1285 match s.to_lowercase().as_str() {
1286 "degree" => Ok(RuntimeGraphCentralityAlgorithm::Degree),
1287 "closeness" => Ok(RuntimeGraphCentralityAlgorithm::Closeness),
1288 "betweenness" => Ok(RuntimeGraphCentralityAlgorithm::Betweenness),
1289 "eigenvector" => Ok(RuntimeGraphCentralityAlgorithm::Eigenvector),
1290 "pagerank" | "page_rank" => Ok(RuntimeGraphCentralityAlgorithm::PageRank),
1291 _ => Err(RedDBError::Query(format!(
1292 "unknown centrality algorithm: '{s}', expected degree|closeness|betweenness|eigenvector|pagerank"
1293 ))),
1294 }
1295}
1296
1297fn parse_community_algorithm(s: &str) -> RedDBResult<RuntimeGraphCommunityAlgorithm> {
1298 match s.to_lowercase().as_str() {
1299 "label_propagation" | "labelpropagation" => {
1300 Ok(RuntimeGraphCommunityAlgorithm::LabelPropagation)
1301 }
1302 "louvain" => Ok(RuntimeGraphCommunityAlgorithm::Louvain),
1303 _ => Err(RedDBError::Query(format!(
1304 "unknown community algorithm: '{s}', expected label_propagation|louvain"
1305 ))),
1306 }
1307}
1308
1309fn parse_components_mode(s: &str) -> RedDBResult<RuntimeGraphComponentsMode> {
1310 match s.to_lowercase().as_str() {
1311 "connected" => Ok(RuntimeGraphComponentsMode::Connected),
1312 "weak" | "weakly_connected" => Ok(RuntimeGraphComponentsMode::Weak),
1313 "strong" | "strongly_connected" => Ok(RuntimeGraphComponentsMode::Strong),
1314 _ => Err(RedDBError::Query(format!(
1315 "unknown components mode: '{s}', expected connected|weak|strong"
1316 ))),
1317 }
1318}
1319
1320fn extract_geo_from_entity(entity: &UnifiedEntity) -> Option<(f64, f64)> {
1325 match &entity.data {
1326 EntityData::Row(row) => {
1327 if let Some(ref named) = row.named {
1328 for value in named.values() {
1329 if let Some(point) = crate::geo::recognize_geo_value(value) {
1330 return Some(point);
1331 }
1332 }
1333 if let Some(point) = crate::geo::recognize_geo_fields(|key| named.get(key)) {
1334 return Some(point);
1335 }
1336 }
1337 for value in &row.columns {
1338 if let Some(point) = crate::geo::recognize_geo_value(value) {
1339 return Some(point);
1340 }
1341 }
1342 None
1343 }
1344 EntityData::Node(node) => {
1345 for value in node.properties.values() {
1346 if let Some(point) = crate::geo::recognize_geo_value(value) {
1347 return Some(point);
1348 }
1349 }
1350 crate::geo::recognize_geo_fields(|key| node.properties.get(key))
1351 }
1352 _ => None,
1353 }
1354}
1355
1356fn index_fields_from_entity(entity: &UnifiedEntity) -> Vec<(String, Value)> {
1357 match &entity.data {
1358 EntityData::Row(row) => row
1359 .iter_fields()
1360 .map(|(field, value)| (field.to_string(), value.clone()))
1361 .collect(),
1362 EntityData::Node(node) => node
1363 .properties
1364 .iter()
1365 .map(|(field, value)| (field.clone(), value.clone()))
1366 .collect(),
1367 EntityData::Edge(edge) => edge
1368 .properties
1369 .iter()
1370 .map(|(field, value)| (field.clone(), value.clone()))
1371 .collect(),
1372 EntityData::TimeSeries(ts) => ts
1373 .fields
1374 .iter()
1375 .map(|(field, value)| (field.clone(), value.clone()))
1376 .collect(),
1377 _ => Vec::new(),
1378 }
1379}
1380
1381fn scan_collection_with_candidates(
1387 store: &std::sync::Arc<crate::storage::unified::store::UnifiedStore>,
1388 collection: &str,
1389 candidates: &Option<std::collections::HashSet<u64>>,
1390) -> Vec<UnifiedEntity> {
1391 let snapshot = crate::runtime::impl_core::capture_current_snapshot();
1392 match candidates {
1393 Some(ids) => store
1394 .get_collection(collection)
1395 .map(|manager| manager.scan(snapshot.as_ref(), |e| ids.contains(&e.id.raw())))
1396 .unwrap_or_default(),
1397 None => store
1398 .get_collection(collection)
1399 .map(|manager| manager.scan(snapshot.as_ref(), |_| true))
1400 .unwrap_or_default(),
1401 }
1402}
1403
1404fn h3_cover_cells(lat: f64, lon: f64, radius_km: f64, resolution: u8) -> Vec<u64> {
1413 let cell = crate::geo::h3::lat_lng_to_cell(lat, lon, resolution);
1414 if cell == 0 {
1415 return Vec::new();
1416 }
1417 let Some(k) = crate::geo::h3::radius_cover_ring(radius_km, resolution) else {
1420 return Vec::new();
1421 };
1422 crate::geo::h3::grid_disk(cell, k)
1423}
1424
1425impl RedDBRuntime {
1426 fn extract_spatial_column(
1427 &self,
1428 entity: &UnifiedEntity,
1429 collection: &str,
1430 column: &str,
1431 ) -> Option<(f64, f64)> {
1432 let fields = index_fields_from_entity(entity);
1433 if let Some(value) = self
1434 .inner
1435 .index_store
1436 .resolve_index_field_value(&fields, column)
1437 {
1438 return crate::geo::recognize_geo_value(&value);
1439 }
1440
1441 let is_document_collection =
1442 self.inner
1443 .db
1444 .collection_contract(collection)
1445 .is_some_and(|contract| {
1446 contract.declared_model == crate::catalog::CollectionModel::Document
1447 });
1448 if is_document_collection {
1449 return None;
1450 }
1451
1452 match entity.data {
1453 EntityData::Row(_) | EntityData::Node(_) => extract_geo_from_entity(entity),
1454 _ => None,
1455 }
1456 }
1457
1458 fn spatial_zero_geo_notice(
1459 &self,
1460 collection: &str,
1461 column: &str,
1462 scanned_entities: &[UnifiedEntity],
1463 geo_values_seen: usize,
1464 result_is_empty: bool,
1465 ) -> Option<String> {
1466 if !result_is_empty || geo_values_seen > 0 {
1467 return None;
1468 }
1469 let (entity_count, geo_count) = if scanned_entities.is_empty() {
1470 let store = self.inner.db.store();
1471 let entities = scan_collection_with_candidates(&store, collection, &None);
1472 let geo_count = entities
1473 .iter()
1474 .filter(|entity| {
1475 self.extract_spatial_column(entity, collection, column)
1476 .is_some()
1477 })
1478 .count();
1479 (entities.len(), geo_count)
1480 } else {
1481 (scanned_entities.len(), geo_values_seen)
1482 };
1483 if entity_count == 0 || geo_count > 0 {
1484 return None;
1485 }
1486 Some(format!(
1487 "no entity in '{}' has an indexable geo value in column '{}' (expected {}).",
1488 collection,
1489 column,
1490 crate::geo::RECOGNIZED_GEO_SHAPES
1491 ))
1492 }
1493
1494 fn h3_index_resolution(&self, collection: &str, column: &str) -> Option<u8> {
1497 match self
1498 .inner
1499 .index_store
1500 .find_index_for_column(collection, column)?
1501 .method
1502 {
1503 super::index_store::IndexMethodKind::H3 { resolution } => Some(resolution),
1504 _ => None,
1505 }
1506 }
1507
1508 fn h3_cell_candidate_ids(
1512 &self,
1513 collection: &str,
1514 column: &str,
1515 cells: &[u64],
1516 ) -> Option<std::collections::HashSet<u64>> {
1517 if cells.is_empty() {
1518 return None;
1519 }
1520 let keys: Vec<_> = cells
1521 .iter()
1522 .filter_map(|c| reddb_types::value_to_canonical_key(&Value::UnsignedInteger(*c)))
1523 .collect();
1524 let ids = self.inner.index_store.sorted.in_lookup_limited(
1525 collection,
1526 column,
1527 &keys,
1528 usize::MAX,
1529 )?;
1530 Some(ids.into_iter().map(|id| id.raw()).collect())
1531 }
1532
1533 fn h3_radius_candidate_ids(
1537 &self,
1538 collection: &str,
1539 column: &str,
1540 center_lat: f64,
1541 center_lon: f64,
1542 radius_km: f64,
1543 ) -> Option<std::collections::HashSet<u64>> {
1544 let resolution = self.h3_index_resolution(collection, column)?;
1545 let cells = h3_cover_cells(center_lat, center_lon, radius_km, resolution);
1546 self.h3_cell_candidate_ids(collection, column, &cells)
1547 }
1548
1549 fn h3_bbox_candidate_ids(
1553 &self,
1554 collection: &str,
1555 column: &str,
1556 min_lat: f64,
1557 min_lon: f64,
1558 max_lat: f64,
1559 max_lon: f64,
1560 ) -> Option<std::collections::HashSet<u64>> {
1561 let resolution = self.h3_index_resolution(collection, column)?;
1562 let center_lat = (min_lat + max_lat) / 2.0;
1563 let center_lon = (min_lon + max_lon) / 2.0;
1564 let radius_km = [
1565 (min_lat, min_lon),
1566 (min_lat, max_lon),
1567 (max_lat, min_lon),
1568 (max_lat, max_lon),
1569 ]
1570 .into_iter()
1571 .map(|(la, lo)| crate::geo::haversine_km(center_lat, center_lon, la, lo))
1572 .fold(0.0_f64, f64::max);
1573 let cells = h3_cover_cells(center_lat, center_lon, radius_km, resolution);
1574 self.h3_cell_candidate_ids(collection, column, &cells)
1575 }
1576
1577 fn h3_polygon_candidate_ids(
1582 &self,
1583 collection: &str,
1584 column: &str,
1585 vertices: &[(f64, f64)],
1586 ) -> Option<std::collections::HashSet<u64>> {
1587 let resolution = self.h3_index_resolution(collection, column)?;
1588 let cells = crate::geo::h3::polygon_to_cover_cells(
1589 vertices,
1590 resolution,
1591 crate::geo::h3::MAX_POLYGON_COVER_CELLS,
1592 )?;
1593 self.h3_cell_candidate_ids(collection, column, &cells)
1594 }
1595
1596 fn h3_nearest_candidate_ids(
1602 &self,
1603 collection: &str,
1604 column: &str,
1605 lat: f64,
1606 lon: f64,
1607 k: usize,
1608 ) -> Option<std::collections::HashSet<u64>> {
1609 let resolution = self.h3_index_resolution(collection, column)?;
1610 if k == 0 {
1611 return None;
1612 }
1613 let center = crate::geo::h3::lat_lng_to_cell(lat, lon, resolution);
1614 if center == 0 {
1615 return None;
1616 }
1617 let edge_km = crate::geo::h3::edge_length_km(resolution).max(f64::MIN_POSITIVE);
1618 const MAX_RING: u32 = 64;
1621 let store = self.inner.db.store();
1622 for r in 0..=MAX_RING {
1623 let cells = crate::geo::h3::grid_disk(center, r);
1624 let Some(ids) = self.h3_cell_candidate_ids(collection, column, &cells) else {
1625 continue;
1626 };
1627 if ids.len() < k {
1628 continue;
1629 }
1630 let candidates = Some(ids);
1631 let mut dists: Vec<f64> =
1632 scan_collection_with_candidates(&store, collection, &candidates)
1633 .iter()
1634 .filter_map(|e| {
1635 extract_geo_from_entity(e)
1636 .map(|(elat, elon)| crate::geo::haversine_km(lat, lon, elat, elon))
1637 })
1638 .collect();
1639 if dists.len() < k {
1640 continue;
1641 }
1642 dists.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1643 let d_k = dists[k - 1];
1644 if f64::from(r) * edge_km >= d_k {
1653 return candidates;
1654 }
1655 }
1656 None
1657 }
1658}