Skip to main content

lance_csr/
shortest_path.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! `shortestPath` / `allShortestPaths` as a breadth-first search over the adjacency.
5//!
6//! [`ShortestPathNode`] takes rows that already bind both end nodes and appends, per
7//! matching path, the path's length and its node ids (`<path>__length`, `<path>__nodes`),
8//! repeating the input row once per path. Rows with no path within the length range are
9//! dropped, as MATCH requires; OPTIONAL MATCH keeps them with nulls through its left join.
10//!
11//! The adjacency is the engine's CSR when it has one (`use_csr=True`). Otherwise the node has
12//! a second input, a scan of the relationship's two endpoint columns, and builds an adjacency
13//! from it once per query, shared by all partitions.
14//!
15//! A source with one target searches from both ends when reverse adjacency is available,
16//! expanding the smaller complete frontier. Sources with several targets share one
17//! forward search. Predecessors retain every shortest path, including parallel edges.
18
19use 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/// Which way the search follows edges.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd)]
49pub enum SearchDirection {
50    Outgoing,
51    Incoming,
52    Undirected,
53}
54
55/// The search parameters shared by the logical and physical nodes.
56#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)]
57pub struct ShortestPathSpec {
58    /// Input columns holding the source and target node ids.
59    pub source_column: String,
60    pub target_column: String,
61    pub direction: SearchDirection,
62    pub min_hops: u32,
63    pub max_hops: u32,
64    /// Every shortest path, not just one.
65    pub all: bool,
66    /// Output column names.
67    pub length_column: String,
68    pub nodes_column: String,
69    /// Endpoint columns of the edge scan (the second input), when there is no CSR.
70    pub edge_source: String,
71    pub edge_target: String,
72}
73
74/// Logical node; see the module docs.
75#[derive(Debug, Clone)]
76pub struct ShortestPathNode {
77    input: LogicalPlan,
78    /// Scan of the relationship's endpoints; `None` when `adjacency` is given.
79    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/// Physical operator for [`ShortestPathNode`].
209#[derive(Debug)]
210pub struct ShortestPathExec {
211    input: Arc<dyn ExecutionPlan>,
212    /// Single-partition edge scan, when there is no CSR.
213    edges: Option<Arc<dyn ExecutionPlan>>,
214    /// The CSR, or the adjacency built from `edges` by the first batch that needs it.
215    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    /// The edge scan must be one partition: the build reads it as a single stream. Without
281    /// this, DataFusion's distribution rule drops a coalesce added at planning time and the
282    /// build sees only partition 0's edges.
283    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
359/// The adjacency for a query without a CSR: built once from the edge scan, in the
360/// direction(s) the search follows.
361async 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
387/// Breadth-first search from `source` until every node in `targets` is reached (or the
388/// frontier empties, or `max_hops`). Returns each reached node's depth and, for every node
389/// at depth d > 0, its predecessors at depth d - 1.
390async 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        // Finish the level before stopping, so every shortest path is recorded.
406        for n in &next {
407            remaining.remove(n);
408        }
409        frontier = next;
410    }
411    (depth, parents)
412}
413
414/// Finish the whole level so parallel edges and alternative shortest paths survive.
415async 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
456/// Search a single pair using complete layers from the smaller frontier. At the
457/// first intersection every meeting node lies on the same depth cut of a shortest path.
458pub(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        // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths.
500        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
523/// Shortest paths from the search source to `target`, as node ids from source to target:
524/// all of them, or the first found.
525async 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    // Targets per source, so each source is searched once for the whole batch.
571    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
630/// Plan [`ShortestPathNode`]s; installed by [`crate::csr_expand::CsrQueryPlanner`].
631pub(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        // One stream feeds the build: merge the scan's partitions.
637        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        // Two 2-hop routes 1-2-4 and 1-3-4, and a longer 1-5-6-4.
697        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        // Beyond the hop limit, nothing is reached.
705        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        // A wide first layer makes the target's side expand. Two meeting nodes,
736        // including parallel edges on both sides of the cut, yield six paths.
737        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        // All directed graphs on three nodes, including disconnected pairs and cycles.
775        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); // Root, first layer and stitched target.
831        assert_eq!(
832            search(&forward, 0, &HashSet::from([101]), 2).await.0.len(),
833            10_101
834        );
835    }
836}