yykv-index 0.0.1

Indexing service for YYKV using Tantivy for full-text search
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use std::collections::{BTreeMap, BTreeSet};
use uuid::Uuid;
use yykv_types::DsResult;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EdgeDirection {
    Outgoing,
    Incoming,
    Both,
}

pub struct MemoryAdjacencyIndex {
    // tenant_id -> (edge_type -> (from_id -> set of to_ids))
    forward: BTreeMap<Uuid, BTreeMap<String, BTreeMap<Uuid, BTreeSet<Uuid>>>>,
    backward: BTreeMap<Uuid, BTreeMap<String, BTreeMap<Uuid, BTreeSet<Uuid>>>>,
}

impl Default for MemoryAdjacencyIndex {
    fn default() -> Self {
        Self::new()
    }
}

impl MemoryAdjacencyIndex {
    pub fn new() -> Self {
        Self {
            forward: BTreeMap::new(),
            backward: BTreeMap::new(),
        }
    }

    pub fn add_edge(
        &mut self,
        tenant_id: Uuid,
        from: Uuid,
        to: Uuid,
        edge_type: &str,
    ) -> DsResult<()> {
        self.forward
            .entry(tenant_id)
            .or_default()
            .entry(edge_type.to_string())
            .or_default()
            .entry(from)
            .or_default()
            .insert(to);
        self.backward
            .entry(tenant_id)
            .or_default()
            .entry(edge_type.to_string())
            .or_default()
            .entry(to)
            .or_default()
            .insert(from);
        Ok(())
    }

    pub fn neighbors(
        &self,
        tenant_id: Uuid,
        from: Uuid,
        edge_type: &str,
        direction: EdgeDirection,
    ) -> DsResult<Vec<Uuid>> {
        let mut results = BTreeSet::new();

        if matches!(direction, EdgeDirection::Outgoing | EdgeDirection::Both)
            && let Some(tos) = self
                .forward
                .get(&tenant_id)
                .and_then(|m| m.get(edge_type))
                .and_then(|m| m.get(&from))
        {
            results.extend(tos);
        }

        if matches!(direction, EdgeDirection::Incoming | EdgeDirection::Both)
            && let Some(froms) = self
                .backward
                .get(&tenant_id)
                .and_then(|m| m.get(edge_type))
                .and_then(|m| m.get(&from))
        {
            results.extend(froms);
        }

        Ok(results.into_iter().collect())
    }

    pub fn delete_node(&mut self, id: Uuid, tenant_id: Uuid) -> DsResult<()> {
        // Remove all edges from/to this node in all edge types for this tenant
        if let Some(forward_map) = self.forward.get_mut(&tenant_id) {
            for edge_map in forward_map.values_mut() {
                edge_map.remove(&id);
                for tos in edge_map.values_mut() {
                    tos.remove(&id);
                }
            }
        }
        if let Some(backward_map) = self.backward.get_mut(&tenant_id) {
            for edge_map in backward_map.values_mut() {
                edge_map.remove(&id);
                for froms in edge_map.values_mut() {
                    froms.remove(&id);
                }
            }
        }
        Ok(())
    }
}