1use std::cmp::Ordering;
20use std::collections::{HashMap, HashSet};
21use std::fmt;
22use std::hash::{Hash, Hasher};
23use std::sync::Arc;
24
25use arrow::array::Int64Builder;
26use arrow::array::{Array, ArrayRef, AsArray, Int64Array, ListBuilder, RecordBatch, UInt32Array};
27use arrow::compute::{cast, take};
28use arrow::datatypes::{DataType, Field, Int64Type, Schema, SchemaRef};
29use datafusion::common::{DFSchema, DFSchemaRef};
30use datafusion::execution::{SendableRecordBatchStream, TaskContext};
31use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
32use datafusion::physical_expr::EquivalenceProperties;
33use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
34use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
35use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
36use datafusion::physical_plan::{
37 DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
38};
39use futures::lock::Mutex;
40use futures::StreamExt;
41
42use crate::csr_expand::{Adjacency, Neighbors};
43use crate::error::{GraphError, Result};
44
45type FastMap<K, V> = HashMap<K, V, ahash::RandomState>;
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd)]
49pub enum SearchDirection {
50 Outgoing,
51 Incoming,
52 Undirected,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)]
57pub struct ShortestPathSpec {
58 pub source_column: String,
60 pub target_column: String,
61 pub direction: SearchDirection,
62 pub min_hops: u32,
63 pub max_hops: u32,
64 pub all: bool,
66 pub length_column: String,
68 pub nodes_column: String,
69 pub edge_source: String,
71 pub edge_target: String,
72}
73
74#[derive(Debug, Clone)]
76pub struct ShortestPathNode {
77 input: LogicalPlan,
78 edges: Option<LogicalPlan>,
80 adjacency: Option<Neighbors>,
81 reverse_adjacency: Option<Neighbors>,
82 spec: ShortestPathSpec,
83 schema: DFSchemaRef,
84}
85
86impl ShortestPathNode {
87 pub fn try_new(
88 input: LogicalPlan,
89 edges: Option<LogicalPlan>,
90 adjacency: Option<Neighbors>,
91 spec: ShortestPathSpec,
92 ) -> Result<Self> {
93 if edges.is_none() == adjacency.is_none() {
94 return Err(plan_error(
95 "shortestPath needs exactly one of a CSR or an edge scan",
96 ));
97 }
98 let added = DFSchema::try_from(Schema::new(vec![
99 Field::new(&spec.length_column, DataType::Int64, true),
100 Field::new(
101 &spec.nodes_column,
102 DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
103 true,
104 ),
105 ]))
106 .map_err(|e| plan_error(&format!("shortestPath schema: {e}")))?;
107 let schema = input
108 .schema()
109 .join(&added)
110 .map_err(|e| plan_error(&format!("shortestPath schema: {e}")))?;
111 Ok(Self {
112 input,
113 edges,
114 adjacency,
115 reverse_adjacency: None,
116 spec,
117 schema: Arc::new(schema),
118 })
119 }
120
121 pub(crate) fn with_reverse_adjacency(mut self, adjacency: Option<Neighbors>) -> Self {
122 self.reverse_adjacency = adjacency;
123 self
124 }
125}
126
127impl PartialEq for ShortestPathNode {
128 fn eq(&self, other: &Self) -> bool {
129 self.input == other.input
130 && self.edges == other.edges
131 && self.spec == other.spec
132 && match (&self.reverse_adjacency, &other.reverse_adjacency) {
133 (Some(a), Some(b)) => a.ptr_eq(b),
134 (None, None) => true,
135 _ => false,
136 }
137 && match (&self.adjacency, &other.adjacency) {
138 (Some(a), Some(b)) => a.ptr_eq(b),
139 (None, None) => true,
140 _ => false,
141 }
142 }
143}
144impl Eq for ShortestPathNode {}
145impl Hash for ShortestPathNode {
146 fn hash<H: Hasher>(&self, state: &mut H) {
147 self.input.hash(state);
148 self.spec.hash(state);
149 }
150}
151impl PartialOrd for ShortestPathNode {
152 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
153 self.spec.partial_cmp(&other.spec)
154 }
155}
156
157impl UserDefinedLogicalNodeCore for ShortestPathNode {
158 fn name(&self) -> &str {
159 "ShortestPath"
160 }
161
162 fn inputs(&self) -> Vec<&LogicalPlan> {
163 std::iter::once(&self.input)
164 .chain(self.edges.as_ref())
165 .collect()
166 }
167
168 fn schema(&self) -> &DFSchemaRef {
169 &self.schema
170 }
171
172 fn expressions(&self) -> Vec<Expr> {
173 vec![]
174 }
175
176 fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
177 let s = &self.spec;
178 write!(
179 f,
180 "ShortestPath: {} -> {}, {:?}, hops {}..{}, {}, {}",
181 s.source_column,
182 s.target_column,
183 s.direction,
184 s.min_hops,
185 s.max_hops,
186 if s.all { "all" } else { "one" },
187 if self.adjacency.is_some() {
188 "csr"
189 } else {
190 "edge scan"
191 }
192 )
193 }
194
195 fn with_exprs_and_inputs(
196 &self,
197 _exprs: Vec<Expr>,
198 mut inputs: Vec<LogicalPlan>,
199 ) -> datafusion::common::Result<Self> {
200 let edges = (inputs.len() > 1).then(|| inputs.pop().expect("two inputs"));
201 let input = inputs.pop().expect("an input");
202 Self::try_new(input, edges, self.adjacency.clone(), self.spec.clone())
203 .map(|node| node.with_reverse_adjacency(self.reverse_adjacency.clone()))
204 .map_err(|e| datafusion::common::DataFusionError::Plan(e.to_string()))
205 }
206}
207
208#[derive(Debug)]
210pub struct ShortestPathExec {
211 input: Arc<dyn ExecutionPlan>,
212 edges: Option<Arc<dyn ExecutionPlan>>,
214 adjacency: Arc<Mutex<Option<Neighbors>>>,
216 reverse_adjacency: Option<Neighbors>,
217 spec: ShortestPathSpec,
218 schema: SchemaRef,
219 properties: Arc<PlanProperties>,
220}
221
222impl ShortestPathExec {
223 fn new(
224 input: Arc<dyn ExecutionPlan>,
225 edges: Option<Arc<dyn ExecutionPlan>>,
226 adjacency: Option<Neighbors>,
227 reverse_adjacency: Option<Neighbors>,
228 spec: ShortestPathSpec,
229 schema: SchemaRef,
230 ) -> Self {
231 let properties = Arc::new(PlanProperties::new(
232 EquivalenceProperties::new(schema.clone()),
233 input.output_partitioning().clone(),
234 EmissionType::Incremental,
235 Boundedness::Bounded,
236 ));
237 Self {
238 input,
239 edges,
240 adjacency: Arc::new(Mutex::new(adjacency)),
241 reverse_adjacency,
242 spec,
243 schema,
244 properties,
245 }
246 }
247}
248
249impl DisplayAs for ShortestPathExec {
250 fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
251 let s = &self.spec;
252 write!(
253 f,
254 "ShortestPathExec: {} -> {}, {:?}, hops {}..{}, {}",
255 s.source_column,
256 s.target_column,
257 s.direction,
258 s.min_hops,
259 s.max_hops,
260 if s.all { "all" } else { "one" }
261 )
262 }
263}
264
265impl ExecutionPlan for ShortestPathExec {
266 fn name(&self) -> &str {
267 "ShortestPathExec"
268 }
269
270 fn properties(&self) -> &Arc<PlanProperties> {
271 &self.properties
272 }
273
274 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
275 std::iter::once(&self.input)
276 .chain(self.edges.as_ref())
277 .collect()
278 }
279
280 fn required_input_distribution(&self) -> Vec<datafusion::physical_plan::Distribution> {
284 use datafusion::physical_plan::Distribution;
285 std::iter::once(Distribution::UnspecifiedDistribution)
286 .chain(self.edges.as_ref().map(|_| Distribution::SinglePartition))
287 .collect()
288 }
289
290 fn with_new_children(
291 self: Arc<Self>,
292 mut children: Vec<Arc<dyn ExecutionPlan>>,
293 ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
294 let edges = (children.len() > 1).then(|| children.pop().expect("two children"));
295 let input = children.pop().expect("a child");
296 let mut exec = Self::new(
297 input,
298 edges,
299 None,
300 self.reverse_adjacency.clone(),
301 self.spec.clone(),
302 self.schema.clone(),
303 );
304 exec.adjacency = self.adjacency.clone();
305 Ok(Arc::new(exec))
306 }
307
308 fn execute(
309 &self,
310 partition: usize,
311 context: Arc<TaskContext>,
312 ) -> datafusion::common::Result<SendableRecordBatchStream> {
313 let input = self.input.execute(partition, context.clone())?;
314 let source_index = input.schema().index_of(&self.spec.source_column)?;
315 let target_index = input.schema().index_of(&self.spec.target_column)?;
316 let (cell, edges, spec, schema) = (
317 self.adjacency.clone(),
318 self.edges.clone(),
319 self.spec.clone(),
320 self.schema.clone(),
321 );
322 let out_schema = schema.clone();
323 let reverse_adjacency = self.reverse_adjacency.clone();
324 let stream = input.then(move |batch| {
325 let reverse_adjacency = reverse_adjacency.clone();
326 let (cell, edges, spec, schema, context) = (
327 cell.clone(),
328 edges.clone(),
329 spec.clone(),
330 schema.clone(),
331 context.clone(),
332 );
333 async move {
334 let adjacency = {
335 let mut slot = cell.lock().await;
336 if slot.is_none() {
337 *slot = Some(build_adjacency(edges, &spec, context).await?);
338 }
339 slot.clone().expect("just set")
340 };
341 search_batch(
342 &batch?,
343 source_index,
344 target_index,
345 &adjacency,
346 reverse_adjacency.as_ref().or_else(|| {
347 (spec.direction == SearchDirection::Undirected).then_some(&adjacency)
348 }),
349 &spec,
350 &schema,
351 )
352 .await
353 }
354 });
355 Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream)))
356 }
357}
358
359async fn build_adjacency(
362 edges: Option<Arc<dyn ExecutionPlan>>,
363 spec: &ShortestPathSpec,
364 context: Arc<TaskContext>,
365) -> datafusion::common::Result<Neighbors> {
366 let edges = edges.expect("an edge scan when there is no CSR");
367 let read = |from: String, to: String| {
368 let (edges, context) = (edges.clone(), context.clone());
369 async move {
370 let stream = edges.execute(0, context)?;
371 Adjacency::build(stream, &from, &to)
372 .await
373 .map(Arc::new)
374 .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))
375 }
376 };
377 let (src, dst) = (spec.edge_source.clone(), spec.edge_target.clone());
378 Ok(match spec.direction {
379 SearchDirection::Outgoing => Neighbors::directed(read(src, dst).await?),
380 SearchDirection::Incoming => Neighbors::directed(read(dst, src).await?),
381 SearchDirection::Undirected => {
382 Neighbors::undirected(read(src.clone(), dst.clone()).await?, read(dst, src).await?)
383 }
384 })
385}
386
387async fn search(
391 adjacency: &Neighbors,
392 source: i64,
393 targets: &HashSet<i64>,
394 max_hops: u32,
395) -> (FastMap<i64, u32>, FastMap<i64, Vec<i64>>) {
396 let mut depth: FastMap<i64, u32> = FastMap::default();
397 let mut parents: FastMap<i64, Vec<i64>> = FastMap::default();
398 depth.insert(source, 0);
399 let mut remaining: HashSet<i64> = targets.iter().copied().filter(|t| *t != source).collect();
400 let mut frontier = vec![source];
401 let mut level = 0;
402 while !frontier.is_empty() && !remaining.is_empty() && level < max_hops {
403 level += 1;
404 let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents).await;
405 for n in &next {
407 remaining.remove(n);
408 }
409 frontier = next;
410 }
411 (depth, parents)
412}
413
414async fn expand_level(
416 adjacency: &Neighbors,
417 frontier: &[i64],
418 level: u32,
419 depth: &mut FastMap<i64, u32>,
420 parents: &mut FastMap<i64, Vec<i64>>,
421) -> Vec<i64> {
422 let mut next = Vec::new();
423 let mut work = 0_usize;
424 for &node in frontier {
425 let previous_work = work;
426 let mut position = crate::csr_expand::NeighborPosition::default();
427 for candidate in adjacency.candidates_from(node, &mut position) {
428 work += 1;
429 if work.is_multiple_of(1024) {
430 crate::csr_expand::cooperative_yield().await;
431 }
432 let Some(n) = candidate else {
433 continue;
434 };
435 let n = n as i64;
436 match depth.get(&n) {
437 Some(&d) if d == level => parents.entry(n).or_default().push(node),
438 Some(_) => {}
439 None => {
440 depth.insert(n, level);
441 parents.insert(n, vec![node]);
442 next.push(n);
443 }
444 }
445 }
446 if work == previous_work {
447 work += 1;
448 if work.is_multiple_of(1024) {
449 crate::csr_expand::cooperative_yield().await;
450 }
451 }
452 }
453 next
454}
455
456pub(crate) async fn search_bidirectional(
459 forward: &Neighbors,
460 reverse: &Neighbors,
461 source: i64,
462 target: i64,
463 max_hops: u32,
464) -> (FastMap<i64, u32>, FastMap<i64, Vec<i64>>) {
465 let mut depth = FastMap::default();
466 let mut parents = FastMap::default();
467 depth.insert(source, 0);
468 if source == target {
469 return (depth, parents);
470 }
471 let mut reverse_depth = FastMap::default();
472 let mut reverse_parents: FastMap<i64, Vec<i64>> = FastMap::default();
473 reverse_depth.insert(target, 0);
474 let (mut frontier, mut reverse_frontier) = (vec![source], vec![target]);
475 let (mut level, mut reverse_level) = (0, 0);
476 while !frontier.is_empty() && !reverse_frontier.is_empty() && level + reverse_level < max_hops {
477 if frontier.len() <= reverse_frontier.len() {
478 level += 1;
479 frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents).await;
480 } else {
481 reverse_level += 1;
482 reverse_frontier = expand_level(
483 reverse,
484 &reverse_frontier,
485 reverse_level,
486 &mut reverse_depth,
487 &mut reverse_parents,
488 )
489 .await;
490 }
491 let mut pending: Vec<_> = frontier
492 .iter()
493 .copied()
494 .filter(|n| reverse_depth.get(n) == Some(&reverse_level))
495 .collect();
496 if pending.is_empty() {
497 continue;
498 }
499 let mut visited = HashSet::new();
501 let mut work = 0_usize;
502 while let Some(node) = pending.pop() {
503 work += 1;
504 if work.is_multiple_of(1024) {
505 crate::csr_expand::cooperative_yield().await;
506 }
507 if !visited.insert(node) {
508 continue;
509 }
510 if let Some(next) = reverse_parents.get(&node) {
511 for &n in next {
512 parents.entry(n).or_default().push(node);
513 pending.push(n);
514 }
515 }
516 }
517 depth.insert(target, level + reverse_level);
518 break;
519 }
520 (depth, parents)
521}
522
523async fn paths_to(target: i64, parents: &FastMap<i64, Vec<i64>>, all: bool) -> Vec<Vec<i64>> {
526 let mut out = Vec::new();
527 let mut stack = vec![vec![target]];
528 let mut work = 0_usize;
529 while let Some(path) = stack.pop() {
530 work += 1;
531 if work.is_multiple_of(1024) {
532 crate::csr_expand::cooperative_yield().await;
533 }
534 let head = *path.last().expect("non-empty");
535 match parents.get(&head) {
536 None => {
537 let mut path = path;
538 path.reverse();
539 out.push(path);
540 if !all {
541 break;
542 }
543 }
544 Some(ps) => {
545 let take = if all { ps.len() } else { 1 };
546 for &p in ps.iter().take(take) {
547 let mut next = path.clone();
548 next.push(p);
549 stack.push(next);
550 }
551 }
552 }
553 }
554 out
555}
556
557async fn search_batch(
558 batch: &RecordBatch,
559 source_index: usize,
560 target_index: usize,
561 adjacency: &Neighbors,
562 reverse_adjacency: Option<&Neighbors>,
563 spec: &ShortestPathSpec,
564 schema: &SchemaRef,
565) -> datafusion::common::Result<RecordBatch> {
566 let sources = cast(batch.column(source_index), &DataType::Int64)?;
567 let sources = sources.as_primitive::<Int64Type>();
568 let targets = cast(batch.column(target_index), &DataType::Int64)?;
569 let targets = targets.as_primitive::<Int64Type>();
570 let mut by_source: FastMap<i64, HashSet<i64>> = FastMap::default();
572 for i in 0..batch.num_rows() {
573 if sources.is_valid(i) && targets.is_valid(i) {
574 by_source
575 .entry(sources.value(i))
576 .or_default()
577 .insert(targets.value(i));
578 }
579 }
580 let mut searches = FastMap::default();
581 for (&source, targets) in &by_source {
582 let result = match reverse_adjacency.filter(|_| targets.len() == 1) {
583 Some(reverse) => {
584 search_bidirectional(
585 adjacency,
586 reverse,
587 source,
588 *targets.iter().next().unwrap(),
589 spec.max_hops,
590 )
591 .await
592 }
593 None => search(adjacency, source, targets, spec.max_hops).await,
594 };
595 searches.insert(source, result);
596 }
597 let mut rows = Vec::new();
598 let mut lengths = Vec::new();
599 let mut nodes = ListBuilder::new(Int64Builder::new());
600 for i in 0..batch.num_rows() {
601 if !(sources.is_valid(i) && targets.is_valid(i)) {
602 continue;
603 }
604 let (s, t) = (sources.value(i), targets.value(i));
605 let (depth, parents) = &searches[&s];
606 let Some(&d) = depth.get(&t) else {
607 continue;
608 };
609 if d < spec.min_hops || d > spec.max_hops {
610 continue;
611 }
612 for path in paths_to(t, parents, spec.all).await {
613 rows.push(i as u32);
614 lengths.push(d as i64);
615 nodes.values().append_slice(&path);
616 nodes.append(true);
617 }
618 }
619 let rows = UInt32Array::from(rows);
620 let mut columns: Vec<ArrayRef> = batch
621 .columns()
622 .iter()
623 .map(|c| take(c.as_ref(), &rows, None))
624 .collect::<std::result::Result<_, _>>()?;
625 columns.push(Arc::new(Int64Array::from(lengths)));
626 columns.push(Arc::new(nodes.finish()));
627 Ok(RecordBatch::try_new(schema.clone(), columns)?)
628}
629
630pub(crate) fn plan_shortest_path(
632 node: &ShortestPathNode,
633 physical_inputs: &[Arc<dyn ExecutionPlan>],
634) -> Arc<dyn ExecutionPlan> {
635 let edges = physical_inputs.get(1).map(|scan| {
636 Arc::new(CoalescePartitionsExec::new(scan.clone())) as Arc<dyn ExecutionPlan>
638 });
639 Arc::new(ShortestPathExec::new(
640 physical_inputs[0].clone(),
641 edges,
642 node.adjacency.clone(),
643 node.reverse_adjacency.clone(),
644 node.spec.clone(),
645 Arc::new(node.schema.as_arrow().clone()),
646 ))
647}
648
649fn plan_error(message: &str) -> GraphError {
650 GraphError::PlanError {
651 message: message.to_string(),
652 location: snafu::Location::new(file!(), line!(), column!()),
653 }
654}
655
656#[cfg(test)]
657mod tests {
658 use super::*;
659
660 async fn adjacencies(edges: &[(i64, i64)]) -> (Arc<Adjacency>, Arc<Adjacency>) {
661 let schema = Arc::new(Schema::new(vec![
662 Field::new("src", DataType::Int64, false),
663 Field::new("dst", DataType::Int64, false),
664 ]));
665 let batch = RecordBatch::try_new(
666 schema.clone(),
667 vec![
668 Arc::new(Int64Array::from(
669 edges.iter().map(|e| e.0).collect::<Vec<_>>(),
670 )),
671 Arc::new(Int64Array::from(
672 edges.iter().map(|e| e.1).collect::<Vec<_>>(),
673 )),
674 ],
675 )
676 .unwrap();
677 let stream = || {
678 let batch = batch.clone();
679 Box::pin(RecordBatchStreamAdapter::new(
680 schema.clone(),
681 futures::stream::once(async move { Ok(batch) }),
682 )) as SendableRecordBatchStream
683 };
684 let out = Adjacency::build(stream(), "src", "dst").await.unwrap();
685 let inc = Adjacency::build(stream(), "dst", "src").await.unwrap();
686 (Arc::new(out), Arc::new(inc))
687 }
688
689 async fn neighbors(edges: &[(i64, i64)]) -> Neighbors {
690 let (out, inc) = adjacencies(edges).await;
691 Neighbors::undirected(out, inc)
692 }
693
694 #[tokio::test]
695 async fn finds_every_shortest_path_and_only_shortest_ones() {
696 let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]).await;
698 let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10).await;
699 assert_eq!(depth[&4], 2);
700 let mut all = paths_to(4, &parents, true).await;
701 all.sort();
702 assert_eq!(all, vec![vec![1, 2, 4], vec![1, 3, 4]]);
703 assert_eq!(paths_to(4, &parents, false).await.len(), 1);
704 let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1).await;
706 assert!(!depth.contains_key(&4));
707 }
708
709 #[tokio::test]
710 async fn cancellation_without_edges_yields_while_scanning_frontier() {
711 let (out, _) = adjacencies(&[]).await;
712 let adjacency = Neighbors::directed(out);
713 let frontier = vec![99; 100_000];
714 let mut depth = FastMap::default();
715 let mut parents = FastMap::default();
716 let mut next = Box::pin(expand_level(
717 &adjacency,
718 &frontier,
719 1,
720 &mut depth,
721 &mut parents,
722 ));
723 assert!(futures::poll!(next.as_mut()).is_pending());
724 drop(next);
725 assert!(depth.is_empty() && parents.is_empty());
726 assert!(
727 expand_level(&adjacency, &frontier, 1, &mut depth, &mut parents)
728 .await
729 .is_empty()
730 );
731 }
732
733 #[tokio::test]
734 async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() {
735 let (out, inc) = adjacencies(&[
738 (1, 2),
739 (1, 2),
740 (1, 3),
741 (1, 9),
742 (2, 4),
743 (2, 4),
744 (3, 4),
745 (3, 4),
746 (4, 5),
747 (5, 5),
748 (5, 1),
749 ])
750 .await;
751 let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc));
752 let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3).await;
753 assert_eq!(depth[&5], 3);
754 let mut all = paths_to(5, &parents, true).await;
755 all.sort();
756 assert_eq!(
757 all,
758 vec![vec![1, 2, 4, 5]; 4]
759 .into_iter()
760 .chain(vec![vec![1, 3, 4, 5]; 2])
761 .collect::<Vec<_>>()
762 );
763 let one = paths_to(5, &parents, false).await;
764 assert_eq!(one.len(), 1);
765 assert!(all.contains(&one[0]));
766 assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2)
767 .await
768 .0
769 .contains_key(&5));
770 }
771
772 #[tokio::test]
773 async fn bidirectional_matches_forward_search_for_small_graphs() {
774 let edges = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)];
776 for mask in 0..(1 << edges.len()) {
777 let selected: Vec<_> = edges
778 .iter()
779 .enumerate()
780 .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge))
781 .collect();
782 let (out, inc) = adjacencies(&selected).await;
783 let directed = Neighbors::directed(out.clone());
784 let reverse = Neighbors::directed(inc.clone());
785 let undirected = Neighbors::undirected(out, inc);
786 for (forward, reverse) in [(&directed, &reverse), (&undirected, &undirected)] {
787 for source in 0..3 {
788 for target in 0..3 {
789 for max_hops in 0..=3 {
790 let (expected_depth, expected_parents) =
791 search(forward, source, &HashSet::from([target]), max_hops).await;
792 let (depth, parents) =
793 search_bidirectional(forward, reverse, source, target, max_hops)
794 .await;
795 assert_eq!(
796 depth.get(&target),
797 expected_depth.get(&target),
798 "graph {mask}, {source}->{target}, max {max_hops}"
799 );
800 if depth.contains_key(&target) {
801 let mut expected = paths_to(target, &expected_parents, true).await;
802 let mut actual = paths_to(target, &parents, true).await;
803 expected.sort();
804 actual.sort();
805 assert_eq!(
806 actual, expected,
807 "graph {mask}, {source}->{target}, max {max_hops}"
808 );
809 }
810 }
811 }
812 }
813 }
814 }
815 }
816
817 #[tokio::test]
818 async fn bidirectional_avoids_expanding_the_wide_middle_layer() {
819 let mut edges = Vec::new();
820 for branch in 1..=100 {
821 edges.push((0, branch));
822 for leaf in 0..100 {
823 edges.push((branch, 101 + (branch - 1) * 100 + leaf));
824 }
825 }
826 let (out, inc) = adjacencies(&edges).await;
827 let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc));
828 let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2).await;
829 assert_eq!(paths_to(101, &parents, true).await, vec![vec![0, 1, 101]]);
830 assert_eq!(depth.len(), 102); assert_eq!(
832 search(&forward, 0, &HashSet::from([101]), 2).await.0.len(),
833 10_101
834 );
835 }
836}