mentat_query_projector/
relresult.rs

1// Copyright 2018 Mozilla
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
4// this file except in compliance with the License. You may obtain a copy of the
5// License at http://www.apache.org/licenses/LICENSE-2.0
6// Unless required by applicable law or agreed to in writing, software distributed
7// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
8// CONDITIONS OF ANY KIND, either express or implied. See the License for the
9// specific language governing permissions and limitations under the License.
10
11use core_traits::{Binding, TypedValue};
12
13/// The result you get from a 'rel' query, like:
14///
15/// ```edn
16/// [:find ?person ?name
17///  :where [?person :person/name ?name]]
18/// ```
19///
20/// There are three ways to get data out of a `RelResult`:
21/// - By iterating over rows as slices. Use `result.rows()`. This is efficient and is
22///   recommended in two cases:
23///   1. If you don't need to take ownership of the resulting values (e.g., you're comparing
24///      or making a modified clone).
25///   2. When the data you're retrieving is cheap to clone. All scalar values are relatively
26///      cheap: they're either small values or `Rc`.
27/// - By direct reference to a row by index, using `result.row(i)`. This also returns
28///   a reference.
29/// - By consuming the results using `into_iter`. This allocates short-lived vectors,
30///   but gives you ownership of the enclosed `TypedValue`s.
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct RelResult<T> {
33    pub width: usize,
34    pub values: Vec<T>,
35}
36
37pub type StructuredRelResult = RelResult<Binding>;
38
39impl<T> RelResult<T> {
40    pub fn empty(width: usize) -> RelResult<T> {
41        RelResult {
42            width,
43            values: Vec::new(),
44        }
45    }
46
47    pub fn is_empty(&self) -> bool {
48        self.values.is_empty()
49    }
50
51    pub fn row_count(&self) -> usize {
52        self.values.len() / self.width
53    }
54
55    pub fn rows(&self) -> ::std::slice::Chunks<'_, T> {
56        // TODO(gburd): Nightly-only API `exact_chunks`. #47115.
57        self.values.chunks(self.width)
58    }
59
60    pub fn row(&self, index: usize) -> Option<&[T]> {
61        let end = self.width * (index + 1);
62        if end > self.values.len() {
63            None
64        } else {
65            let start = self.width * index;
66            Some(&self.values[start..end])
67        }
68    }
69}
70
71#[test]
72fn test_rel_result() {
73    let empty = StructuredRelResult::empty(3);
74    let unit = StructuredRelResult {
75        width: 1,
76        values: vec![TypedValue::Long(5).into()],
77    };
78    let two_by_two = StructuredRelResult {
79        width: 2,
80        values: vec![
81            TypedValue::Long(5).into(),
82            TypedValue::Boolean(true).into(),
83            TypedValue::Long(-2).into(),
84            TypedValue::Boolean(false).into(),
85        ],
86    };
87
88    assert!(empty.is_empty());
89    assert!(!unit.is_empty());
90    assert!(!two_by_two.is_empty());
91
92    assert_eq!(empty.row_count(), 0);
93    assert_eq!(unit.row_count(), 1);
94    assert_eq!(two_by_two.row_count(), 2);
95
96    assert_eq!(empty.row(0), None);
97    assert_eq!(unit.row(1), None);
98    assert_eq!(two_by_two.row(2), None);
99
100    assert_eq!(
101        unit.row(0),
102        Some(vec![TypedValue::Long(5).into()].as_slice())
103    );
104    assert_eq!(
105        two_by_two.row(0),
106        Some(vec![TypedValue::Long(5).into(), TypedValue::Boolean(true).into()].as_slice())
107    );
108    assert_eq!(
109        two_by_two.row(1),
110        Some(
111            vec![
112                TypedValue::Long(-2).into(),
113                TypedValue::Boolean(false).into()
114            ]
115            .as_slice()
116        )
117    );
118
119    let mut rr = two_by_two.rows();
120    assert_eq!(
121        rr.next(),
122        Some(vec![TypedValue::Long(5).into(), TypedValue::Boolean(true).into()].as_slice())
123    );
124    assert_eq!(
125        rr.next(),
126        Some(
127            vec![
128                TypedValue::Long(-2).into(),
129                TypedValue::Boolean(false).into()
130            ]
131            .as_slice()
132        )
133    );
134    assert_eq!(rr.next(), None);
135}
136
137// Primarily for testing.
138impl From<Vec<Vec<TypedValue>>> for RelResult<Binding> {
139    fn from(src: Vec<Vec<TypedValue>>) -> Self {
140        if src.is_empty() {
141            RelResult::empty(0)
142        } else {
143            let width = src.get(0).map(|r| r.len()).unwrap_or(0);
144            RelResult {
145                width,
146                values: src
147                    .into_iter()
148                    .flat_map(|r| r.into_iter().map(|v| v.into()))
149                    .collect(),
150            }
151        }
152    }
153}
154
155pub struct SubvecIntoIterator<T> {
156    width: usize,
157    values: ::std::vec::IntoIter<T>,
158}
159
160impl<T> Iterator for SubvecIntoIterator<T> {
161    // TODO: this is a good opportunity to use `SmallVec` instead: most queries return a handful of columns.
162    type Item = Vec<T>;
163    fn next(&mut self) -> Option<Self::Item> {
164        let result: Vec<_> = (&mut self.values).take(self.width).collect();
165        if result.is_empty() {
166            None
167        } else {
168            Some(result)
169        }
170    }
171}
172
173impl<T> IntoIterator for RelResult<T> {
174    type Item = Vec<T>;
175    type IntoIter = SubvecIntoIterator<T>;
176
177    fn into_iter(self) -> Self::IntoIter {
178        SubvecIntoIterator {
179            width: self.width,
180            values: self.values.into_iter(),
181        }
182    }
183}
mentat_query_projector/relresult.rs

mentat_query_projector/
relresult.rs
