mentat/query-projector/src/relresult.rs

184 lines
5.2 KiB
Rust
Raw Normal View History

// Copyright 2018 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
use core_traits::{Binding, TypedValue};
/// The result you get from a 'rel' query, like:
///
/// ```edn
/// [:find ?person ?name
/// :where [?person :person/name ?name]]
/// ```
///
/// There are three ways to get data out of a `RelResult`:
/// - By iterating over rows as slices. Use `result.rows()`. This is efficient and is
/// recommended in two cases:
/// 1. If you don't need to take ownership of the resulting values (e.g., you're comparing
/// or making a modified clone).
/// 2. When the data you're retrieving is cheap to clone. All scalar values are relatively
/// cheap: they're either small values or `Rc`.
/// - By direct reference to a row by index, using `result.row(i)`. This also returns
/// a reference.
/// - By consuming the results using `into_iter`. This allocates short-lived vectors,
/// but gives you ownership of the enclosed `TypedValue`s.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RelResult<T> {
pub width: usize,
pub values: Vec<T>,
}
pub type StructuredRelResult = RelResult<Binding>;
impl<T> RelResult<T> {
pub fn empty(width: usize) -> RelResult<T> {
RelResult {
2020-02-21 15:27:39 +00:00
width,
values: Vec::new(),
}
}
pub fn is_empty(&self) -> bool {
self.values.is_empty()
}
pub fn row_count(&self) -> usize {
self.values.len() / self.width
}
pub fn rows(&self) -> ::std::slice::Chunks<T> {
2020-02-21 15:27:39 +00:00
// TODO(gburd): Nightly-only API `exact_chunks`. #47115.
self.values.chunks(self.width)
}
pub fn row(&self, index: usize) -> Option<&[T]> {
let end = self.width * (index + 1);
if end > self.values.len() {
None
} else {
let start = self.width * index;
Some(&self.values[start..end])
}
}
}
#[test]
fn test_rel_result() {
let empty = StructuredRelResult::empty(3);
let unit = StructuredRelResult {
width: 1,
values: vec![TypedValue::Long(5).into()],
};
let two_by_two = StructuredRelResult {
width: 2,
values: vec![
TypedValue::Long(5).into(),
TypedValue::Boolean(true).into(),
TypedValue::Long(-2).into(),
TypedValue::Boolean(false).into(),
],
};
assert!(empty.is_empty());
assert!(!unit.is_empty());
assert!(!two_by_two.is_empty());
assert_eq!(empty.row_count(), 0);
assert_eq!(unit.row_count(), 1);
assert_eq!(two_by_two.row_count(), 2);
assert_eq!(empty.row(0), None);
assert_eq!(unit.row(1), None);
assert_eq!(two_by_two.row(2), None);
assert_eq!(
unit.row(0),
Some(vec![TypedValue::Long(5).into()].as_slice())
);
assert_eq!(
two_by_two.row(0),
Some(vec![TypedValue::Long(5).into(), TypedValue::Boolean(true).into()].as_slice())
);
assert_eq!(
two_by_two.row(1),
Some(
vec![
TypedValue::Long(-2).into(),
TypedValue::Boolean(false).into()
]
.as_slice()
)
);
let mut rr = two_by_two.rows();
assert_eq!(
rr.next(),
Some(vec![TypedValue::Long(5).into(), TypedValue::Boolean(true).into()].as_slice())
);
assert_eq!(
rr.next(),
Some(
vec![
TypedValue::Long(-2).into(),
TypedValue::Boolean(false).into()
]
.as_slice()
)
);
assert_eq!(rr.next(), None);
}
// Primarily for testing.
impl From<Vec<Vec<TypedValue>>> for RelResult<Binding> {
fn from(src: Vec<Vec<TypedValue>>) -> Self {
if src.is_empty() {
RelResult::empty(0)
} else {
let width = src.get(0).map(|r| r.len()).unwrap_or(0);
RelResult {
2020-02-21 15:27:39 +00:00
width,
values: src
.into_iter()
.flat_map(|r| r.into_iter().map(|v| v.into()))
.collect(),
}
}
}
}
pub struct SubvecIntoIterator<T> {
width: usize,
values: ::std::vec::IntoIter<T>,
}
impl<T> Iterator for SubvecIntoIterator<T> {
// TODO: this is a good opportunity to use `SmallVec` instead: most queries return a handful of columns.
type Item = Vec<T>;
fn next(&mut self) -> Option<Self::Item> {
let result: Vec<_> = (&mut self.values).take(self.width).collect();
if result.is_empty() {
None
} else {
Some(result)
}
}
}
impl<T> IntoIterator for RelResult<T> {
type Item = Vec<T>;
type IntoIter = SubvecIntoIterator<T>;
fn into_iter(self) -> Self::IntoIter {
SubvecIntoIterator {
width: self.width,
values: self.values.into_iter(),
}
}
}