helix_db/helix_engine/graph_core/
traversal_iter.rs

1use std::{cmp::Ordering, sync::Arc};
2
3use heed3::{RoTxn, RwTxn};
4
5use super::ops::tr_val::TraversalVal;
6use crate::{
7    helix_engine::{storage_core::storage_core::HelixGraphStorage, types::GraphError},
8    protocol::value::Value,
9    utils::filterable::Filterable,
10};
11use itertools::Itertools;
12
13pub struct RoTraversalIterator<'a, I> {
14    pub inner: I,
15    pub storage: Arc<HelixGraphStorage>,
16    pub txn: &'a RoTxn<'a>,
17}
18
19// implementing iterator for TraversalIterator
20impl<'a, I> Iterator for RoTraversalIterator<'a, I>
21where
22    I: Iterator<Item = Result<TraversalVal, GraphError>>,
23{
24    type Item = Result<TraversalVal, GraphError>;
25
26    fn next(&mut self) -> Option<Self::Item> {
27        self.inner.next()
28    }
29}
30
31impl<'a, I: Iterator<Item = Result<TraversalVal, GraphError>>> RoTraversalIterator<'a, I> {
32    pub fn take_and_collect_to<B: FromIterator<TraversalVal>>(self, n: usize) -> B {
33        self.inner
34            .filter_map(|item| item.ok())
35            .take(n)
36            .collect::<B>()
37    }
38
39    pub fn collect_to<B: FromIterator<TraversalVal>>(self) -> B {
40        self.inner.filter_map(|item| {
41            item.ok()
42        }).collect::<B>()
43    }
44
45    pub fn collect_dedup<B: FromIterator<TraversalVal>>(self) -> B {
46        self.inner
47            .filter_map(|item| item.ok())
48            .unique()
49            .collect::<B>()
50    }
51
52    pub fn collect_to_obj(self) -> TraversalVal {
53        match self.inner.filter_map(|item| item.ok()).next() {
54            Some(val) => val,
55            None => TraversalVal::Empty,
56        }
57    }
58
59    pub fn count_to_val(self) -> Value {
60        Value::from(self.inner.count())
61    }
62
63    pub fn order_by_asc(self, property: &str) -> Result<Vec<TraversalVal>, GraphError> {
64        let result = self
65            .inner
66            .filter_map(|item| item.ok())
67            .sorted_by(|a, b| match (a, b) {
68                (TraversalVal::Node(a), TraversalVal::Node(b)) => {
69                    match (a.check_property(property), b.check_property(property)) {
70                        (Ok(val_a), Ok(val_b)) => val_a.cmp(val_b),
71                        (Ok(_), Err(_)) => Ordering::Less,
72                        (Err(_), Ok(_)) => Ordering::Greater,
73                        (Err(_), Err(_)) => Ordering::Equal,
74                    }
75                }
76                (TraversalVal::Edge(a), TraversalVal::Edge(b)) => {
77                    match (a.check_property(property), b.check_property(property)) {
78                        (Ok(val_a), Ok(val_b)) => val_a.cmp(val_b),
79                        (Ok(_), Err(_)) => Ordering::Less,
80                        (Err(_), Ok(_)) => Ordering::Greater,
81                        (Err(_), Err(_)) => Ordering::Equal,
82                    }
83                }
84                (TraversalVal::Vector(a), TraversalVal::Vector(b)) => {
85                    match (a.check_property(property), b.check_property(property)) {
86                        (Ok(val_a), Ok(val_b)) => val_a.cmp(val_b),
87                        (Ok(_), Err(_)) => Ordering::Less,
88                        (Err(_), Ok(_)) => Ordering::Greater,
89                        (Err(_), Err(_)) => Ordering::Equal,
90                    }
91                }
92                (TraversalVal::Count(val_a), TraversalVal::Count(val_b)) => val_a.cmp(val_b),
93                (TraversalVal::Value(val_a), TraversalVal::Value(val_b)) => val_a.cmp(val_b),
94                _ => Ordering::Equal,
95            })
96            .collect::<Vec<_>>();
97
98        Ok(result)
99    }
100
101    pub fn order_by_desc(self, property: &str) -> Result<Vec<TraversalVal>, GraphError> {
102        let result = self
103            .inner
104            .filter_map(|item| item.ok())
105            .sorted_by(|a, b| match (a, b) {
106                (TraversalVal::Node(a), TraversalVal::Node(b)) => {
107                    match (a.check_property(property), b.check_property(property)) {
108                        (Ok(val_a), Ok(val_b)) => val_b.cmp(val_a),
109                        (Ok(_), Err(_)) => Ordering::Greater,
110                        (Err(_), Ok(_)) => Ordering::Less,
111                        (Err(_), Err(_)) => Ordering::Equal,
112                    }
113                }
114                (TraversalVal::Edge(a), TraversalVal::Edge(b)) => {
115                    match (a.check_property(property), b.check_property(property)) {
116                        (Ok(val_a), Ok(val_b)) => val_b.cmp(val_a),
117                        (Ok(_), Err(_)) => Ordering::Greater,
118                        (Err(_), Ok(_)) => Ordering::Less,
119                        (Err(_), Err(_)) => Ordering::Equal,
120                    }
121                }
122                (TraversalVal::Vector(a), TraversalVal::Vector(b)) => {
123                    match (a.check_property(property), b.check_property(property)) {
124                        (Ok(val_a), Ok(val_b)) => val_b.cmp(val_a),
125                        (Ok(_), Err(_)) => Ordering::Greater,
126                        (Err(_), Ok(_)) => Ordering::Less,
127                        (Err(_), Err(_)) => Ordering::Equal,
128                    }
129                }
130                (TraversalVal::Count(val_a), TraversalVal::Count(val_b)) => val_b.cmp(val_a),
131                (TraversalVal::Value(val_a), TraversalVal::Value(val_b)) => val_b.cmp(val_a),
132                _ => Ordering::Equal,
133            })
134            .collect::<Vec<_>>();
135
136        Ok(result)
137    }
138
139    pub fn map_value_or(
140        mut self,
141        default: bool,
142        f: impl Fn(&Value) -> bool,
143    ) -> Result<bool, GraphError> {
144        let val = match &self.inner.next() {
145            Some(Ok(TraversalVal::Value(val))) => {
146                println!("value : {:?}", val);
147                Ok(f(val))
148            }
149            Some(Ok(_)) => Err(GraphError::ConversionError(
150                "Expected value, got something else".to_string(),
151            )),
152            Some(Err(err)) => Err(GraphError::from(err.to_string())),
153            None => Ok(default),
154        };
155        println!("result: {:?}", val);
156        val
157    }
158}
159
160pub struct RwTraversalIterator<'scope, 'env, I> {
161    pub inner: I,
162    pub storage: Arc<HelixGraphStorage>,
163    pub txn: &'scope mut RwTxn<'env>,
164}
165
166// implementing iterator for TraversalIterator
167impl<'scope, 'env, I> Iterator for RwTraversalIterator<'scope, 'env, I>
168where
169    I: Iterator<Item = Result<TraversalVal, GraphError>>,
170{
171    type Item = Result<TraversalVal, GraphError>;
172
173    fn next(&mut self) -> Option<Self::Item> {
174        self.inner.next()
175    }
176}
177impl<'scope, 'env, I: Iterator<Item = Result<TraversalVal, GraphError>>>
178    RwTraversalIterator<'scope, 'env, I>
179{
180    pub fn new(storage: Arc<HelixGraphStorage>, txn: &'scope mut RwTxn<'env>, inner: I) -> Self {
181        Self {
182            inner,
183            storage,
184            txn,
185        }
186    }
187
188    pub fn collect_to<B: FromIterator<TraversalVal>>(self) -> B
189    where
190        I: Iterator<Item = Result<TraversalVal, GraphError>>,
191    {
192        self.inner.filter_map(|item| item.ok()).collect::<B>()
193    }
194
195    pub fn collect_to_val(self) -> TraversalVal
196    where
197        I: Iterator<Item = Result<TraversalVal, GraphError>>,
198    {
199        match self
200            .inner
201            .filter_map(|item| item.ok())
202            .collect::<Vec<_>>()
203            .first()
204        {
205            Some(val) => val.clone(), // TODO: Remove clone
206            None => TraversalVal::Empty,
207        }
208    }
209    pub fn collect_to_obj(self) -> TraversalVal {
210        match self.inner.filter_map(|item| item.ok()).next() {
211            Some(val) => val,
212            None => TraversalVal::Empty,
213        }
214    }
215}
216// pub trait TraversalIteratorMut<'a> {
217//     type Inner: Iterator<Item = Result<TraversalVal, GraphError>>;
218
219//     fn next<'b>(
220//         &mut self,
221//         storage: Arc<HelixGraphStorage>,
222//         txn: &'b mut RwTxn<'a>,
223//     ) -> Option<Result<TraversalVal, GraphError>>;
224
225// }
helix_db/helix_engine/graph_core/traversal_iter.rs

helix_db/helix_engine/graph_core/
traversal_iter.rs
