Compare commits

..

4 Commits

Author SHA1 Message Date
Paul Masurel
8828b6d310 Support for columnar 2022-12-21 12:21:30 +09:00
Paul Masurel
2b89bf9050 Support for NotNaN in fast fields 2022-12-21 12:20:48 +09:00
Paul Masurel
3580198447 Minor refactoring 2022-12-21 12:18:33 +09:00
Paul Masurel
d96a716d20 Refactoring to prepare for the addition of dynamic fast field 2022-12-21 12:16:00 +09:00
64 changed files with 1033 additions and 1949 deletions

2
.gitignore vendored
View File

@@ -13,5 +13,3 @@ benchmark
.idea .idea
trace.dat trace.dat
cargo-timing* cargo-timing*
columnar/columnar-cli/*.json
**/perf.data*

View File

@@ -36,6 +36,7 @@ fs2 = { version = "0.4.3", optional = true }
levenshtein_automata = "0.2.1" levenshtein_automata = "0.2.1"
uuid = { version = "1.0.0", features = ["v4", "serde"] } uuid = { version = "1.0.0", features = ["v4", "serde"] }
crossbeam-channel = "0.5.4" crossbeam-channel = "0.5.4"
stable_deref_trait = "1.2.0"
rust-stemmers = "1.2.0" rust-stemmers = "1.2.0"
downcast-rs = "1.2.0" downcast-rs = "1.2.0"
bitpacking = { version = "0.8.4", default-features = false, features = ["bitpacker4x"] } bitpacking = { version = "0.8.4", default-features = false, features = ["bitpacker4x"] }
@@ -61,6 +62,7 @@ tantivy-query-grammar = { version= "0.19.0", path="./query-grammar" }
tantivy-bitpacker = { version= "0.3", path="./bitpacker" } tantivy-bitpacker = { version= "0.3", path="./bitpacker" }
common = { version= "0.5", path = "./common/", package = "tantivy-common" } common = { version= "0.5", path = "./common/", package = "tantivy-common" }
fastfield_codecs = { version= "0.3", path="./fastfield_codecs", default-features = false } fastfield_codecs = { version= "0.3", path="./fastfield_codecs", default-features = false }
ownedbytes = { version= "0.5", path="./ownedbytes" }
[target.'cfg(windows)'.dependencies] [target.'cfg(windows)'.dependencies]
winapi = "0.3.9" winapi = "0.3.9"

View File

@@ -91,15 +91,17 @@ impl BitUnpacker {
return 0u64; return 0u64;
} }
let addr_in_bits = idx * self.num_bits as u32; let addr_in_bits = idx * self.num_bits as u32;
let addr = (addr_in_bits >> 3) as usize; let addr = addr_in_bits >> 3;
let bit_shift = addr_in_bits & 7; let bit_shift = addr_in_bits & 7;
debug_assert!( debug_assert!(
addr + 8 <= data.len(), addr + 8 <= data.len() as u32,
"The fast field field should have been padded with 7 bytes." "The fast field field should have been padded with 7 bytes."
); );
let bytes: [u8; 8] = (&data[addr..addr + 8]).try_into().unwrap(); let bytes: [u8; 8] = (&data[(addr as usize)..(addr as usize) + 8])
.try_into()
.unwrap();
let val_unshifted_unmasked: u64 = u64::from_le_bytes(bytes); let val_unshifted_unmasked: u64 = u64::from_le_bytes(bytes);
let val_shifted = val_unshifted_unmasked >> bit_shift; let val_shifted: u64 = val_unshifted_unmasked >> bit_shift;
val_shifted & self.mask val_shifted & self.mask
} }
} }

View File

@@ -1,8 +1,6 @@
mod bitpacker; mod bitpacker;
mod blocked_bitpacker; mod blocked_bitpacker;
use std::cmp::Ordering;
pub use crate::bitpacker::{BitPacker, BitUnpacker}; pub use crate::bitpacker::{BitPacker, BitUnpacker};
pub use crate::blocked_bitpacker::BlockedBitpacker; pub use crate::blocked_bitpacker::BlockedBitpacker;
@@ -39,104 +37,44 @@ pub fn compute_num_bits(n: u64) -> u8 {
} }
} }
/// Computes the (min, max) of an iterator of `PartialOrd` values.
///
/// For values implementing `Ord` (in a way consistent to their `PartialOrd` impl),
/// this function behaves as expected.
///
/// For values with partial ordering, the behavior is non-trivial and may
/// depends on the order of the values.
/// For floats however, it simply returns the same results as if NaN were
/// skipped.
pub fn minmax<I, T>(mut vals: I) -> Option<(T, T)> pub fn minmax<I, T>(mut vals: I) -> Option<(T, T)>
where where
I: Iterator<Item = T>, I: Iterator<Item = T>,
T: Copy + PartialOrd, T: Copy + Ord,
{ {
let first_el = vals.find(|val| { if let Some(first_el) = vals.next() {
// We use this to make sure we skip all NaN values when return Some(vals.fold((first_el, first_el), |(min_val, max_val), el| {
// working with a float type. (min_val.min(el), max_val.max(el))
val.partial_cmp(val) == Some(Ordering::Equal) }));
})?;
let mut min_so_far: T = first_el;
let mut max_so_far: T = first_el;
for val in vals {
if val.partial_cmp(&min_so_far) == Some(Ordering::Less) {
min_so_far = val;
}
if val.partial_cmp(&max_so_far) == Some(Ordering::Greater) {
max_so_far = val;
}
} }
Some((min_so_far, max_so_far)) None
} }
#[cfg(test)] #[test]
mod tests { fn test_compute_num_bits() {
use super::*; assert_eq!(compute_num_bits(1), 1u8);
assert_eq!(compute_num_bits(0), 0u8);
#[test] assert_eq!(compute_num_bits(2), 2u8);
fn test_compute_num_bits() { assert_eq!(compute_num_bits(3), 2u8);
assert_eq!(compute_num_bits(1), 1u8); assert_eq!(compute_num_bits(4), 3u8);
assert_eq!(compute_num_bits(0), 0u8); assert_eq!(compute_num_bits(255), 8u8);
assert_eq!(compute_num_bits(2), 2u8); assert_eq!(compute_num_bits(256), 9u8);
assert_eq!(compute_num_bits(3), 2u8); assert_eq!(compute_num_bits(5_000_000_000), 33u8);
assert_eq!(compute_num_bits(4), 3u8); }
assert_eq!(compute_num_bits(255), 8u8);
assert_eq!(compute_num_bits(256), 9u8); #[test]
assert_eq!(compute_num_bits(5_000_000_000), 33u8); fn test_minmax_empty() {
} let vals: Vec<u32> = vec![];
assert_eq!(minmax(vals.into_iter()), None);
#[test] }
fn test_minmax_empty() {
let vals: Vec<u32> = vec![]; #[test]
assert_eq!(minmax(vals.into_iter()), None); fn test_minmax_one() {
} assert_eq!(minmax(vec![1].into_iter()), Some((1, 1)));
}
#[test]
fn test_minmax_one() { #[test]
assert_eq!(minmax(vec![1].into_iter()), Some((1, 1))); fn test_minmax_two() {
} assert_eq!(minmax(vec![1, 2].into_iter()), Some((1, 2)));
assert_eq!(minmax(vec![2, 1].into_iter()), Some((1, 2)));
#[test]
fn test_minmax_two() {
assert_eq!(minmax(vec![1, 2].into_iter()), Some((1, 2)));
assert_eq!(minmax(vec![2, 1].into_iter()), Some((1, 2)));
}
#[test]
fn test_minmax_nan() {
assert_eq!(
minmax(vec![f64::NAN, 1f64, 2f64].into_iter()),
Some((1f64, 2f64))
);
assert_eq!(
minmax(vec![2f64, f64::NAN, 1f64].into_iter()),
Some((1f64, 2f64))
);
assert_eq!(
minmax(vec![2f64, 1f64, f64::NAN].into_iter()),
Some((1f64, 2f64))
);
}
#[test]
fn test_minmax_inf() {
assert_eq!(
minmax(vec![f64::INFINITY, 1f64, 2f64].into_iter()),
Some((1f64, f64::INFINITY))
);
assert_eq!(
minmax(vec![-f64::INFINITY, 1f64, 2f64].into_iter()),
Some((-f64::INFINITY, 2f64))
);
assert_eq!(
minmax(vec![2f64, f64::INFINITY, 1f64].into_iter()),
Some((1f64, f64::INFINITY))
);
assert_eq!(
minmax(vec![2f64, 1f64, -f64::INFINITY].into_iter()),
Some((-f64::INFINITY, 2f64))
);
}
} }

View File

@@ -2,18 +2,25 @@
name = "tantivy-columnar" name = "tantivy-columnar"
version = "0.1.0" version = "0.1.0"
edition = "2021" edition = "2021"
license = "MIT"
[dependencies] [dependencies]
stacker = { path = "../stacker", package="tantivy-stacker"} stacker = { path = "../stacker", package="tantivy-stacker"}
serde_json = "1" serde_json = "1"
thiserror = "1" thiserror = "1"
fnv = "1" fnv = "1"
tantivy-fst = "0.4.0"
sstable = { path = "../sstable", package = "tantivy-sstable" } sstable = { path = "../sstable", package = "tantivy-sstable" }
zstd = "0.12"
common = { path = "../common", package = "tantivy-common" } common = { path = "../common", package = "tantivy-common" }
fastfield_codecs = { path = "../fastfield_codecs"} fastfield_codecs = { path = "../fastfield_codecs"}
ordered-float = "3.4"
itertools = "0.10" itertools = "0.10"
[features]
# default = ["quickwit"]
# quickwit = ["common/quickwit"]
[dev-dependencies] [dev-dependencies]
proptest = "1" proptest = "1"

View File

@@ -2,6 +2,7 @@
This crate describes columnar format used in tantivy. This crate describes columnar format used in tantivy.
## Goals ## Goals
This format is special in the following way. This format is special in the following way.
@@ -14,55 +15,14 @@ we need to be able to load columns rapidly.
and different cardinality `(required, optional, multivalued)`. and different cardinality `(required, optional, multivalued)`.
- columns, once loaded, offer cheap random access. - columns, once loaded, offer cheap random access.
# Coercion rules # Format
Users can create a columnar by appending rows to a writer. A quickwit/tantivy style sstable associated
Nothing prevents a user from recording values with different to a same `column_key`.
In that case, `tantivy-columnar`'s behavior is as follows:
- Values that corresponds to different JsonValue type are mapped to different columns. For instance, String values are treated independently from Number or boolean values. `tantivy-columnar` will simply emit several columns associated to a given column_name.
- Only one column for a given json value type is emitted. If number values with different number types are recorded (e.g. u64, i64, f64), `tantivy-columnar` will pick the first type that can represents the set of appended value, with the following prioriy order (`i64`, `u64`, `f64`). `i64` is picked over `u64` as it is likely to yield less change of types. Most use cases strictly requiring `u64` show the restriction on 50% of the values (e.g. a 64-bit hash). On the other hand, a lot of use cases can show rare negative value.
# Columnar format
Because this columnar format tries to avoid some coercion.
There can be several columns (with different type) associated to a single `column_name`.
Each column is associated to `column_key`.
The format of that key is:
`[column_name][ZERO_BYTE][column_type_header: u8]`
```
COLUMNAR:=
[COLUMNAR_DATA]
[COLUMNAR_INDEX]
[COLUMNAR_FOOTER];
# Columns are sorted by their column key.
COLUMNAR_DATA:=
[COLUMN]+;
COLUMN:=
COMPRESSED_COLUMN | NON_COMPRESSED_COLUMN;
# COLUMN_DATA is compressed when it exceeds a threshold of 100KB.
COMPRESSED_COLUMN := [b'1'][zstd(COLUMN_DATA)]
NON_COMPRESSED_COLUMN:= [b'0'][COLUMN_DATA]
COLUMNAR_INDEX := [RANGE_SSTABLE_BYTES]
COLUMNAR_FOOTER := [RANGE_SSTABLE_BYTES_LEN: 8 bytes little endian]
```
The columnar file starts by the actual column data, concatenated one after the other,
sorted by column key.
A quickwit/tantivy style sstable associates
`(column names, column_cardinality, column_type) to range of bytes. `(column names, column_cardinality, column_type) to range of bytes.
The format of the key is:
`[column_name][ZERO_BYTE][column_type_header: u8]`
Column name may not contain the zero byte. Column name may not contain the zero byte.
Listing all columns associated to `column_name` can therefore Listing all columns associated to `column_name` can therefore

View File

@@ -1,17 +0,0 @@
[package]
name = "tantivy-columnar-cli"
version = "0.1.0"
edition = "2021"
license = "MIT"
[dependencies]
columnar = {path="../", package="tantivy-columnar"}
serde_json = "1"
serde_json_borrow = {git="https://github.com/PSeitz/serde_json_borrow/"}
serde = "1"
[workspace]
members = []
[profile.release]
debug = true

View File

@@ -1,126 +0,0 @@
use columnar::ColumnarWriter;
use columnar::NumericalValue;
use serde_json_borrow;
use std::fs::File;
use std::io;
use std::io::BufRead;
use std::io::BufReader;
use std::time::Instant;
#[derive(Default)]
struct JsonStack {
path: String,
stack: Vec<usize>,
}
impl JsonStack {
fn push(&mut self, seg: &str) {
let len = self.path.len();
self.stack.push(len);
self.path.push('.');
self.path.push_str(seg);
}
fn pop(&mut self) {
if let Some(len) = self.stack.pop() {
self.path.truncate(len);
}
}
fn path(&self) -> &str {
&self.path[1..]
}
}
fn append_json_to_columnar(
doc: u32,
json_value: &serde_json_borrow::Value,
columnar: &mut ColumnarWriter,
stack: &mut JsonStack,
) -> usize {
let mut count = 0;
match json_value {
serde_json_borrow::Value::Null => {}
serde_json_borrow::Value::Bool(val) => {
columnar.record_numerical(
doc,
stack.path(),
NumericalValue::from(if *val { 1u64 } else { 0u64 }),
);
count += 1;
}
serde_json_borrow::Value::Number(num) => {
let numerical_value: NumericalValue = if let Some(num_i64) = num.as_i64() {
num_i64.into()
} else if let Some(num_u64) = num.as_u64() {
num_u64.into()
} else if let Some(num_f64) = num.as_f64() {
num_f64.into()
} else {
panic!();
};
count += 1;
columnar.record_numerical(
doc,
stack.path(),
numerical_value,
);
}
serde_json_borrow::Value::Str(msg) => {
columnar.record_str(
doc,
stack.path(),
msg.as_bytes(),
);
count += 1;
},
serde_json_borrow::Value::Array(vals) => {
for val in vals {
count += append_json_to_columnar(doc, val, columnar, stack);
}
},
serde_json_borrow::Value::Object(json_map) => {
for (child_key, child_val) in json_map {
stack.push(child_key);
count += append_json_to_columnar(doc, child_val, columnar, stack);
stack.pop();
}
},
}
count
}
fn main() -> io::Result<()> {
let file = File::open("gh_small.json")?;
let mut reader = BufReader::new(file);
let mut line = String::with_capacity(100);
let mut columnar = columnar::ColumnarWriter::default();
let mut doc = 0;
let start = Instant::now();
let mut stack = JsonStack::default();
let mut total_count = 0;
loop {
line.clear();
let len = reader.read_line(&mut line)?;
if len == 0 {
break;
}
let Ok(json_value) = serde_json::from_str::<serde_json_borrow::Value>(&line) else { continue; };
total_count += append_json_to_columnar(doc, &json_value, &mut columnar, &mut stack);
doc += 1;
}
println!("value count {total_count}");
println!("record {:?}", start.elapsed());
let mut buffer = Vec::new();
columnar.serialize(doc, &mut buffer)?;
println!("num docs: {doc}, {:?}", start.elapsed());
println!("buffer len {} MB", buffer.len() / 1_000_000);
let columnar = columnar::ColumnarReader::open(buffer)?;
for (column_name, typ, offsets, num_bytes) in columnar.list_columns()? {
if num_bytes>1_000_000 {
println!("{column_name} {typ:?} {offsets:?} {}", num_bytes / 1_000_000);
}
}
println!("{} columns", columnar.num_columns());
Ok(())
}

View File

@@ -1,26 +1,20 @@
use crate::utils::{place_bits, select_bits};
use crate::value::NumericalType; use crate::value::NumericalType;
/// Enum describing the number of values that can exist per document
/// (or per row if you will).
#[derive(Clone, Copy, Hash, Default, Debug, PartialEq, Eq, PartialOrd, Ord)] #[derive(Clone, Copy, Hash, Default, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[repr(u8)] #[repr(u8)]
pub enum Cardinality { pub enum Cardinality {
/// All documents contain exactly one value.
#[default] #[default]
Required = 0, Required = 0,
/// All documents contain at most one value.
Optional = 1, Optional = 1,
/// All documents may contain any number of values.
Multivalued = 2, Multivalued = 2,
} }
impl Cardinality { impl Cardinality {
pub(crate) fn to_code(self) -> u8 { pub fn to_code(self) -> u8 {
self as u8 self as u8
} }
pub(crate) fn try_from_code(code: u8) -> Option<Cardinality> { pub fn try_from_code(code: u8) -> Option<Cardinality> {
match code { match code {
0 => Some(Cardinality::Required), 0 => Some(Cardinality::Required),
1 => Some(Cardinality::Optional), 1 => Some(Cardinality::Optional),
@@ -34,80 +28,25 @@ impl Cardinality {
pub enum ColumnType { pub enum ColumnType {
Bytes, Bytes,
Numerical(NumericalType), Numerical(NumericalType),
Bool,
} }
impl ColumnType { impl ColumnType {
/// Encoded over 6 bits. pub fn to_code(self) -> u8 {
pub(crate) fn to_code(self) -> u8 {
let high_type;
let low_code: u8;
match self { match self {
ColumnType::Bytes => { ColumnType::Bytes => 0u8,
high_type = GeneralType::Str; ColumnType::Numerical(numerical_type) => 1u8 | (numerical_type.to_code() << 1),
low_code = 0u8;
}
ColumnType::Numerical(numerical_type) => {
high_type = GeneralType::Numerical;
low_code = numerical_type.to_code();
}
ColumnType::Bool => {
high_type = GeneralType::Bool;
low_code = 0u8;
}
} }
place_bits::<3, 6>(high_type.to_code()) | place_bits::<0, 3>(low_code)
} }
pub(crate) fn try_from_code(code: u8) -> Option<ColumnType> { pub fn try_from_code(code: u8) -> Option<ColumnType> {
if select_bits::<6, 8>(code) != 0u8 { if code == 0u8 {
return Some(ColumnType::Bytes);
}
if code & 1u8 == 0u8 {
return None; return None;
} }
let high_code = select_bits::<3, 6>(code); let numerical_type = NumericalType::try_from_code(code >> 1)?;
let low_code = select_bits::<0, 3>(code); Some(ColumnType::Numerical(numerical_type))
let high_type = GeneralType::try_from_code(high_code)?;
match high_type {
GeneralType::Bool => {
if low_code != 0u8 {
return None;
}
Some(ColumnType::Bool)
}
GeneralType::Str => {
if low_code != 0u8 {
return None;
}
Some(ColumnType::Bytes)
}
GeneralType::Numerical => {
let numerical_type = NumericalType::try_from_code(low_code)?;
Some(ColumnType::Numerical(numerical_type))
}
}
}
}
/// This corresponds to the JsonType.
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
#[repr(u8)]
pub(crate) enum GeneralType {
Bool = 0u8,
Str = 1u8,
Numerical = 2u8,
}
impl GeneralType {
pub fn to_code(self) -> u8 {
self as u8
}
pub fn try_from_code(code: u8) -> Option<Self> {
match code {
0u8 => Some(Self::Bool),
1u8 => Some(Self::Str),
2u8 => Some(Self::Numerical),
_ => None,
}
} }
} }
@@ -123,6 +62,34 @@ pub struct ColumnTypeAndCardinality {
pub typ: ColumnType, pub typ: ColumnType,
} }
#[inline]
const fn compute_mask(num_bits: u8) -> u8 {
if num_bits == 8 {
u8::MAX
} else {
(1u8 << num_bits) - 1
}
}
#[inline]
fn select_bits<const START: u8, const END: u8>(code: u8) -> u8 {
assert!(START <= END);
assert!(END <= 8);
let num_bits: u8 = END - START;
let mask: u8 = compute_mask(num_bits);
(code >> START) & mask
}
#[inline]
fn place_bits<const START: u8, const END: u8>(code: u8) -> u8 {
assert!(START <= END);
assert!(END <= 8);
let num_bits: u8 = END - START;
let mask: u8 = compute_mask(num_bits);
assert!(code <= mask);
code << START
}
impl ColumnTypeAndCardinality { impl ColumnTypeAndCardinality {
pub fn to_code(self) -> u8 { pub fn to_code(self) -> u8 {
place_bits::<6, 8>(self.cardinality.to_code()) | place_bits::<0, 6>(self.typ.to_code()) place_bits::<6, 8>(self.cardinality.to_code()) | place_bits::<0, 6>(self.typ.to_code())
@@ -156,8 +123,7 @@ mod tests {
} }
assert_eq!( assert_eq!(
column_type_header_set.len(), column_type_header_set.len(),
3 /* cardinality */ * 3 /* cardinality */ * (1 + 3) // column_types
(1 + 1 + 3) // column_types (str, bool, numerical x 3)
); );
} }
@@ -170,7 +136,7 @@ mod tests {
assert!(column_type_set.insert(column_type)); assert!(column_type_set.insert(column_type));
} }
} }
assert_eq!(column_type_set.len(), 2 + 3); assert_eq!(column_type_set.len(), 1 + 3);
} }
#[test] #[test]

View File

@@ -1,27 +1,16 @@
use std::io; use std::io;
use fnv::FnvHashMap; use fnv::FnvHashMap;
use sstable::SSTable;
pub(crate) struct IdMapping { fn fst_err_into_io_err(fst_err: tantivy_fst::Error) -> io::Error {
unordered_to_ord: Vec<OrderedId>, match fst_err {
} tantivy_fst::Error::Fst(fst_err) => {
io::Error::new(io::ErrorKind::Other, format!("FST Error: {:?}", fst_err))
impl IdMapping { }
pub fn to_ord(&self, unordered: UnorderedId) -> OrderedId { tantivy_fst::Error::Io(io_err) => io_err,
self.unordered_to_ord[unordered.0 as usize]
} }
} }
/// When we add values, we cannot know their ordered id yet.
/// For this reason, we temporarily assign them a `UnorderedId`
/// that will be mapped to an `OrderedId` upon serialization.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub struct UnorderedId(pub u32);
#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
pub struct OrderedId(pub u32);
/// `DictionaryBuilder` for dictionary encoding. /// `DictionaryBuilder` for dictionary encoding.
/// ///
/// It stores the different terms encounterred and assigns them a temporary value /// It stores the different terms encounterred and assigns them a temporary value
@@ -30,10 +19,20 @@ pub struct OrderedId(pub u32);
/// Upon serialization, we will sort the ids and hence build a `UnorderedId -> Term ordinal` /// Upon serialization, we will sort the ids and hence build a `UnorderedId -> Term ordinal`
/// mapping. /// mapping.
#[derive(Default)] #[derive(Default)]
pub(crate) struct DictionaryBuilder { pub struct DictionaryBuilder {
dict: FnvHashMap<Vec<u8>, UnorderedId>, dict: FnvHashMap<Vec<u8>, UnorderedId>,
} }
pub struct IdMapping {
unordered_to_ord: Vec<OrderedId>,
}
impl IdMapping {
pub fn to_ord(&self, unordered: UnorderedId) -> OrderedId {
self.unordered_to_ord[unordered.0 as usize]
}
}
impl DictionaryBuilder { impl DictionaryBuilder {
/// Get or allocate an unordered id. /// Get or allocate an unordered id.
/// (This ID is simply an auto-incremented id.) /// (This ID is simply an auto-incremented id.)
@@ -49,36 +48,31 @@ impl DictionaryBuilder {
/// Serialize the dictionary into an fst, and returns the /// Serialize the dictionary into an fst, and returns the
/// `UnorderedId -> TermOrdinal` map. /// `UnorderedId -> TermOrdinal` map.
pub fn serialize<'a, W: io::Write + 'a>(&self, wrt: &mut W) -> io::Result<IdMapping> { pub fn serialize<'a, W: io::Write + 'a>(&self, wrt: &mut W) -> io::Result<IdMapping> {
let mut terms: Vec<(&[u8], UnorderedId)> = serialize_inner(&self.dict, wrt).map_err(fst_err_into_io_err)
self.dict.iter().map(|(k, v)| (k.as_slice(), *v)).collect();
terms.sort_unstable_by_key(|(key, _)| *key);
// TODO Remove the allocation.
let mut unordered_to_ord: Vec<OrderedId> = vec![OrderedId(0u32); terms.len()];
let mut sstable_builder = sstable::VoidSSTable::writer(wrt);
for (ord, (key, unordered_id)) in terms.into_iter().enumerate() {
let ordered_id = OrderedId(ord as u32);
sstable_builder.insert(key, &())?;
unordered_to_ord[unordered_id.0 as usize] = ordered_id;
}
sstable_builder.finish()?;
Ok(IdMapping { unordered_to_ord })
} }
} }
#[cfg(test)] /// Helper function just there for error conversion.
mod tests { fn serialize_inner<'a, W: io::Write + 'a>(
use super::*; dict: &FnvHashMap<Vec<u8>, UnorderedId>,
wrt: &mut W,
#[test] ) -> tantivy_fst::Result<IdMapping> {
fn test_dictionary_builder() { let mut terms: Vec<(&[u8], UnorderedId)> =
let mut dictionary_builder = DictionaryBuilder::default(); dict.iter().map(|(k, v)| (k.as_slice(), *v)).collect();
let hello_uid = dictionary_builder.get_or_allocate_id(b"hello"); terms.sort_unstable_by_key(|(key, _)| *key);
let happy_uid = dictionary_builder.get_or_allocate_id(b"happy"); let mut unordered_to_ord: Vec<OrderedId> = vec![OrderedId(0u32); terms.len()];
let tax_uid = dictionary_builder.get_or_allocate_id(b"tax"); let mut fst_builder = tantivy_fst::MapBuilder::new(wrt)?;
let mut buffer = Vec::new(); for (ord, (key, unordered_id)) in terms.into_iter().enumerate() {
let id_mapping = dictionary_builder.serialize(&mut buffer).unwrap(); let ordered_id = OrderedId(ord as u32);
assert_eq!(id_mapping.to_ord(hello_uid), OrderedId(1)); fst_builder.insert(key, ord as u64)?;
assert_eq!(id_mapping.to_ord(happy_uid), OrderedId(0)); unordered_to_ord[unordered_id.0 as usize] = ordered_id;
assert_eq!(id_mapping.to_ord(tax_uid), OrderedId(2));
} }
fst_builder.finish()?;
Ok(IdMapping { unordered_to_ord })
} }
#[derive(Clone, Copy, Debug)]
pub struct UnorderedId(pub u32);
#[derive(Clone, Copy)]
pub struct OrderedId(pub u32);

View File

@@ -1,13 +1,32 @@
// Copyright (C) 2022 Quickwit, Inc.
//
// Quickwit is offered under the AGPL v3.0 and as commercial software.
// For commercial licensing, contact us at hello@quickwit.io.
//
// AGPL:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
mod column_type_header; mod column_type_header;
mod dictionary; mod dictionary;
mod reader; mod reader;
pub(crate) mod utils; mod serializer;
mod value; mod value;
mod writer; mod writer;
pub use column_type_header::Cardinality; pub use column_type_header::Cardinality;
pub use reader::ColumnarReader; pub use reader::ColumnarReader;
pub use value::{NumericalType, NumericalValue}; pub use serializer::ColumnarSerializer;
pub use writer::ColumnarWriter; pub use writer::ColumnarWriter;
pub type DocId = u32; pub type DocId = u32;
@@ -18,57 +37,21 @@ mod tests {
use common::file_slice::FileSlice; use common::file_slice::FileSlice;
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality}; use crate::column_type_header::ColumnTypeAndCardinality;
use crate::reader::ColumnarReader; use crate::reader::ColumnarReader;
use crate::serializer::ColumnarSerializer;
use crate::value::NumericalValue; use crate::value::NumericalValue;
use crate::{Cardinality, ColumnarWriter}; use crate::ColumnarWriter;
#[test] #[test]
fn test_dataframe_writer_bytes() { fn test_dataframe_writer() {
let mut dataframe_writer = ColumnarWriter::default(); let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_str(1u32, "my_string", b"hello"); dataframe_writer.record_numerical(1u32, b"srical.value", NumericalValue::U64(1u64));
dataframe_writer.record_str(3u32, "my_string", b"helloeee"); dataframe_writer.record_numerical(2u32, b"srical.value", NumericalValue::U64(2u64));
dataframe_writer.record_numerical(4u32, b"srical.value", NumericalValue::I64(2i64));
let mut buffer: Vec<u8> = Vec::new(); let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap(); let serializer = ColumnarSerializer::new(&mut buffer);
let columnar_fileslice = FileSlice::from(buffer); dataframe_writer.serialize(5, serializer).unwrap();
let columnar = ColumnarReader::open(columnar_fileslice).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
columnar.read_columns("my_string").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(cols[0].1, 0..159);
}
#[test]
fn test_dataframe_writer_bool() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_bool(1u32, "bool.value", false);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap();
let columnar_fileslice = FileSlice::from(buffer);
let columnar = ColumnarReader::open(columnar_fileslice).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
columnar.read_columns("bool.value").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(
cols[0].0,
ColumnTypeAndCardinality {
cardinality: Cardinality::Optional,
typ: ColumnType::Bool
}
);
assert_eq!(cols[0].1, 0..22);
}
#[test]
fn test_dataframe_writer_numerical() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_numerical(1u32, "srical.value", NumericalValue::U64(12u64));
dataframe_writer.record_numerical(2u32, "srical.value", NumericalValue::U64(13u64));
dataframe_writer.record_numerical(4u32, "srical.value", NumericalValue::U64(15u64));
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap();
let columnar_fileslice = FileSlice::from(buffer); let columnar_fileslice = FileSlice::from(buffer);
let columnar = ColumnarReader::open(columnar_fileslice).unwrap(); let columnar = ColumnarReader::open(columnar_fileslice).unwrap();
assert_eq!(columnar.num_columns(), 1); assert_eq!(columnar.num_columns(), 1);
@@ -81,6 +64,6 @@ mod tests {
// - vals 8 //< due to padding? could have been 1byte?. // - vals 8 //< due to padding? could have been 1byte?.
// - null footer 6 bytes // - null footer 6 bytes
// - version footer 3 bytes // Should be file-wide // - version footer 3 bytes // Should be file-wide
assert_eq!(cols[0].1, 0..32); assert_eq!(cols[0].1, 0..31);
} }
} }

View File

@@ -3,7 +3,7 @@ use std::{io, mem};
use common::file_slice::FileSlice; use common::file_slice::FileSlice;
use common::BinarySerializable; use common::BinarySerializable;
use sstable::{Dictionary, RangeSSTable}; use sstable::{Dictionary, SSTableRange};
use crate::column_type_header::ColumnTypeAndCardinality; use crate::column_type_header::ColumnTypeAndCardinality;
@@ -11,22 +11,17 @@ fn io_invalid_data(msg: String) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, msg) // format!("Invalid key found. io::Error::new(io::ErrorKind::InvalidData, msg) // format!("Invalid key found.
// {key_bytes:?}"))); // {key_bytes:?}")));
} }
/// The ColumnarReader makes it possible to access a set of columns
/// associated to field names.
pub struct ColumnarReader { pub struct ColumnarReader {
column_dictionary: Dictionary<RangeSSTable>, column_dictionary: Dictionary<SSTableRange>,
column_data: FileSlice, column_data: FileSlice,
} }
impl ColumnarReader { impl ColumnarReader {
/// Opens a new Columnar file. pub fn num_columns(&self) -> usize {
pub fn open<F>(file_slice: F) -> io::Result<ColumnarReader> self.column_dictionary.num_terms()
where FileSlice: From<F> {
Self::open_inner(file_slice.into())
} }
fn open_inner(file_slice: FileSlice) -> io::Result<ColumnarReader> { pub fn open(file_slice: FileSlice) -> io::Result<ColumnarReader> {
let (file_slice_without_sstable_len, sstable_len_bytes) = let (file_slice_without_sstable_len, sstable_len_bytes) =
file_slice.split_from_end(mem::size_of::<u64>()); file_slice.split_from_end(mem::size_of::<u64>());
let mut sstable_len_bytes = sstable_len_bytes.read_bytes()?; let mut sstable_len_bytes = sstable_len_bytes.read_bytes()?;
@@ -40,32 +35,6 @@ impl ColumnarReader {
}) })
} }
// TODO fix ugly API
pub fn list_columns(
&self,
) -> io::Result<Vec<(String, ColumnTypeAndCardinality, Range<u64>, u64)>> {
let mut stream = self.column_dictionary.stream()?;
let mut results = Vec::new();
while stream.advance() {
let key_bytes: &[u8] = stream.key();
let column_code: u8 = key_bytes.last().cloned().unwrap();
let column_type_and_cardinality = ColumnTypeAndCardinality::try_from_code(column_code)
.ok_or_else(|| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
let range = stream.value().clone();
let column_name = String::from_utf8_lossy(&key_bytes[..key_bytes.len() - 1]);
let range_len = range.end - range.start;
results.push((
column_name.to_string(),
column_type_and_cardinality,
range,
range_len,
));
}
Ok(results)
}
/// Get all columns for the given field_name.
// TODO fix ugly API
pub fn read_columns( pub fn read_columns(
&self, &self,
field_name: &str, field_name: &str,
@@ -94,9 +63,4 @@ impl ColumnarReader {
} }
Ok(results) Ok(results)
} }
/// Return the number of columns in the columnar.
pub fn num_columns(&self) -> usize {
self.column_dictionary.num_terms()
}
} }

View File

@@ -0,0 +1,39 @@
use std::io;
use std::io::Write;
use std::ops::Range;
use common::CountingWriter;
use sstable::value::RangeWriter;
use sstable::SSTableRange;
pub struct ColumnarSerializer<W: io::Write> {
wrt: CountingWriter<W>,
sstable_range: sstable::Writer<Vec<u8>, RangeWriter>,
}
impl<W: io::Write> ColumnarSerializer<W> {
pub fn new(wrt: W) -> ColumnarSerializer<W> {
let sstable_range: sstable::Writer<Vec<u8>, RangeWriter> =
sstable::Dictionary::<SSTableRange>::builder(Vec::with_capacity(100_000)).unwrap();
ColumnarSerializer {
wrt: CountingWriter::wrap(wrt),
sstable_range,
}
}
pub fn record_column_offsets(&mut self, key: &[u8], byte_range: Range<u64>) -> io::Result<()> {
self.sstable_range.insert(key, &byte_range)
}
pub fn wrt(&mut self) -> &mut CountingWriter<W> {
&mut self.wrt
}
pub fn finalize(mut self) -> io::Result<()> {
let sstable_bytes: Vec<u8> = self.sstable_range.finish()?;
let sstable_num_bytes: u64 = sstable_bytes.len() as u64;
self.wrt.write_all(&sstable_bytes)?;
self.wrt.write_all(&sstable_num_bytes.to_le_bytes()[..])?;
Ok(())
}
}

View File

@@ -1,76 +0,0 @@
const fn compute_mask(num_bits: u8) -> u8 {
if num_bits == 8 {
u8::MAX
} else {
(1u8 << num_bits) - 1
}
}
#[inline(always)]
#[must_use]
pub(crate) fn select_bits<const START: u8, const END: u8>(code: u8) -> u8 {
assert!(START <= END);
assert!(END <= 8);
let num_bits: u8 = END - START;
let mask: u8 = compute_mask(num_bits);
(code >> START) & mask
}
#[inline(always)]
#[must_use]
pub(crate) fn place_bits<const START: u8, const END: u8>(code: u8) -> u8 {
assert!(START <= END);
assert!(END <= 8);
let num_bits: u8 = END - START;
let mask: u8 = compute_mask(num_bits);
assert!(code <= mask);
code << START
}
/// Pop-front one bytes from a slice of bytes.
#[inline(always)]
pub fn pop_first_byte(bytes: &mut &[u8]) -> Option<u8> {
if bytes.is_empty() {
return None;
}
let first_byte = bytes[0];
*bytes = &bytes[1..];
Some(first_byte)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_select_bits() {
assert_eq!(255u8, select_bits::<0, 8>(255u8));
assert_eq!(0u8, select_bits::<0, 0>(255u8));
assert_eq!(8u8, select_bits::<0, 4>(8u8));
assert_eq!(4u8, select_bits::<1, 4>(8u8));
assert_eq!(0u8, select_bits::<1, 3>(8u8));
}
#[test]
fn test_place_bits() {
assert_eq!(255u8, place_bits::<0, 8>(255u8));
assert_eq!(4u8, place_bits::<2, 3>(1u8));
assert_eq!(0u8, place_bits::<2, 2>(0u8));
}
#[test]
#[should_panic]
fn test_place_bits_overflows() {
let _ = place_bits::<1, 4>(8u8);
}
#[test]
fn test_pop_first_byte() {
let mut cursor: &[u8] = &b"abcd"[..];
assert_eq!(pop_first_byte(&mut cursor), Some(b'a'));
assert_eq!(pop_first_byte(&mut cursor), Some(b'b'));
assert_eq!(pop_first_byte(&mut cursor), Some(b'c'));
assert_eq!(pop_first_byte(&mut cursor), Some(b'd'));
assert_eq!(pop_first_byte(&mut cursor), None);
}
}

View File

@@ -1,8 +1,10 @@
use ordered_float::NotNan;
#[derive(Copy, Clone, Debug, PartialEq)] #[derive(Copy, Clone, Debug, PartialEq)]
pub enum NumericalValue { pub enum NumericalValue {
I64(i64), I64(i64),
U64(u64), U64(u64),
F64(f64), F64(NotNan<f64>),
} }
impl From<u64> for NumericalValue { impl From<u64> for NumericalValue {
@@ -17,8 +19,8 @@ impl From<i64> for NumericalValue {
} }
} }
impl From<f64> for NumericalValue { impl From<NotNan<f64>> for NumericalValue {
fn from(val: f64) -> Self { fn from(val: NotNan<f64>) -> Self {
NumericalValue::F64(val) NumericalValue::F64(val)
} }
} }
@@ -93,11 +95,11 @@ impl Coerce for u64 {
} }
} }
impl Coerce for f64 { impl Coerce for NotNan<f64> {
fn coerce(value: NumericalValue) -> Self { fn coerce(value: NumericalValue) -> Self {
match value { match value {
NumericalValue::I64(val) => val as f64, NumericalValue::I64(val) => unsafe { NotNan::new_unchecked(val as f64) },
NumericalValue::U64(val) => val as f64, NumericalValue::U64(val) => unsafe { NotNan::new_unchecked(val as f64) },
NumericalValue::F64(val) => val, NumericalValue::F64(val) => val,
} }
} }

View File

@@ -1,7 +1,12 @@
use std::fmt;
use std::num::NonZeroU8;
use ordered_float::NotNan;
use thiserror::Error;
use crate::dictionary::UnorderedId; use crate::dictionary::UnorderedId;
use crate::utils::{place_bits, pop_first_byte, select_bits};
use crate::value::NumericalValue; use crate::value::NumericalValue;
use crate::{DocId, NumericalType}; use crate::DocId;
/// When we build a columnar dataframe, we first just group /// When we build a columnar dataframe, we first just group
/// all mutations per column, and append them in append-only object. /// all mutations per column, and append them in append-only object.
@@ -13,76 +18,6 @@ pub(crate) enum ColumnOperation<T> {
Value(T), Value(T),
} }
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct ColumnOperationHeader {
typ_code: u8,
len: u8,
}
impl ColumnOperationHeader {
fn to_code(self) -> u8 {
place_bits::<0, 4>(self.len) | place_bits::<4, 8>(self.typ_code)
}
fn from_code(code: u8) -> Self {
let len = select_bits::<0, 4>(code);
let typ_code = select_bits::<4, 8>(code);
ColumnOperationHeader { typ_code, len }
}
}
const NEW_DOC_CODE: u8 = 0u8;
const NEW_VALUE_CODE: u8 = 1u8;
impl<V: SymbolValue> ColumnOperation<V> {
pub fn serialize(self) -> impl AsRef<[u8]> {
let mut minibuf = MiniBuffer::default();
let header = match self {
ColumnOperation::NewDoc(new_doc) => {
let symbol_len = new_doc.serialize(&mut minibuf.bytes[1..]);
ColumnOperationHeader {
typ_code: NEW_DOC_CODE,
len: symbol_len,
}
}
ColumnOperation::Value(val) => {
let symbol_len = val.serialize(&mut minibuf.bytes[1..]);
ColumnOperationHeader {
typ_code: NEW_VALUE_CODE,
len: symbol_len,
}
}
};
minibuf.bytes[0] = header.to_code();
minibuf.len = 1 + header.len;
minibuf
}
/// Deserialize a colummn operation.
/// Returns None if the buffer is empty.
///
/// Panics if the payload is invalid.
pub fn deserialize(bytes: &mut &[u8]) -> Option<Self> {
let header_byte = pop_first_byte(bytes)?;
let column_op_header = ColumnOperationHeader::from_code(header_byte);
let symbol_bytes: &[u8];
(symbol_bytes, *bytes) = bytes.split_at(column_op_header.len as usize);
match column_op_header.typ_code {
NEW_DOC_CODE => {
let new_doc = u32::deserialize(symbol_bytes);
Some(ColumnOperation::NewDoc(new_doc))
}
NEW_VALUE_CODE => {
let value = V::deserialize(symbol_bytes);
Some(ColumnOperation::Value(value))
}
_ => {
panic!("Unknown code {}", column_op_header.typ_code);
}
}
}
}
impl<T> From<T> for ColumnOperation<T> { impl<T> From<T> for ColumnOperation<T> {
fn from(value: T) -> Self { fn from(value: T) -> Self {
ColumnOperation::Value(value) ColumnOperation::Value(value)
@@ -90,116 +25,207 @@ impl<T> From<T> for ColumnOperation<T> {
} }
#[allow(clippy::from_over_into)] #[allow(clippy::from_over_into)]
pub(crate) trait SymbolValue: Clone + Copy { pub(crate) trait SymbolValue: Into<MiniBuffer> + Clone + Copy + fmt::Debug {
fn serialize(self, buffer: &mut [u8]) -> u8; fn deserialize(header: NonZeroU8, bytes: &mut &[u8]) -> Result<Self, ParseError>;
// Reads the header type and the given bytes.
//
// `bytes` does not contain the header byte.
// This method should advance bytes by the number of bytes that were consumed.
fn deserialize(bytes: &[u8]) -> Self;
} }
impl SymbolValue for bool { pub(crate) struct MiniBuffer {
fn serialize(self, buffer: &mut [u8]) -> u8 { pub bytes: [u8; 9],
buffer[0] = if self { 1u8 } else { 0u8 }; pub len: usize,
1u8 }
}
fn deserialize(bytes: &[u8]) -> Self { impl MiniBuffer {
bytes[0] == 1u8 pub fn as_slice(&self) -> &[u8] {
&self.bytes[..self.len]
} }
} }
#[derive(Default)] fn compute_header_byte(typ: SymbolType, len: usize) -> u8 {
struct MiniBuffer { assert!(len <= 9);
pub bytes: [u8; 10], (len << 4) as u8 | typ as u8
pub len: u8,
}
impl AsRef<[u8]> for MiniBuffer {
fn as_ref(&self) -> &[u8] {
&self.bytes[..self.len as usize]
}
} }
impl SymbolValue for NumericalValue { impl SymbolValue for NumericalValue {
fn deserialize(mut bytes: &[u8]) -> Self { fn deserialize(header_byte: NonZeroU8, bytes: &mut &[u8]) -> Result<Self, ParseError> {
let type_code = pop_first_byte(&mut bytes).unwrap(); let (typ, len) = parse_header_byte(header_byte)?;
let symbol_type = NumericalType::try_from_code(type_code).unwrap(); let value_bytes: &[u8];
let mut octet: [u8; 8] = [0u8; 8]; (value_bytes, *bytes) = bytes.split_at(len);
octet[..bytes.len()].copy_from_slice(bytes); let symbol: NumericalValue = match typ {
match symbol_type { SymbolType::U64 => {
NumericalType::U64 => { let mut octet: [u8; 8] = [0u8; 8];
octet[..value_bytes.len()].copy_from_slice(value_bytes);
let val: u64 = u64::from_le_bytes(octet); let val: u64 = u64::from_le_bytes(octet);
NumericalValue::U64(val) NumericalValue::U64(val)
} }
NumericalType::I64 => { SymbolType::I64 => {
let mut octet: [u8; 8] = [0u8; 8];
octet[..value_bytes.len()].copy_from_slice(value_bytes);
let encoded: u64 = u64::from_le_bytes(octet); let encoded: u64 = u64::from_le_bytes(octet);
let val: i64 = decode_zig_zag(encoded); let val: i64 = decode_zig_zag(encoded);
NumericalValue::I64(val) NumericalValue::I64(val)
} }
NumericalType::F64 => { SymbolType::Float => {
debug_assert_eq!(bytes.len(), 8); let octet: [u8; 8] =
let val: f64 = f64::from_le_bytes(octet); value_bytes.try_into().map_err(|_| ParseError::InvalidLen {
NumericalValue::F64(val) typ: SymbolType::Float,
len,
})?;
let val_possibly_nan = f64::from_le_bytes(octet);
let val_not_nan = NotNan::new(val_possibly_nan)
.map_err(|_| ParseError::NaN)?;
NumericalValue::F64(val_not_nan)
} }
} };
Ok(symbol)
} }
}
fn serialize(self, output: &mut [u8]) -> u8 { #[allow(clippy::from_over_into)]
impl Into<MiniBuffer> for NumericalValue {
fn into(self) -> MiniBuffer {
let mut bytes = [0u8; 9];
match self { match self {
NumericalValue::F64(val) => { NumericalValue::F64(val) => {
output[0] = NumericalType::F64 as u8; let len = 8;
output[1..9].copy_from_slice(&val.to_le_bytes()); let header_byte = compute_header_byte(SymbolType::Float, len);
9u8 bytes[0] = header_byte;
bytes[1..].copy_from_slice(&val.to_le_bytes());
MiniBuffer {
bytes,
len: len + 1,
}
} }
NumericalValue::U64(val) => { NumericalValue::U64(val) => {
let len = compute_num_bytes_for_u64(val) as u8; let len = compute_num_bytes_for_u64(val);
output[0] = NumericalType::U64 as u8; let header_byte = compute_header_byte(SymbolType::U64, len);
output[1..9].copy_from_slice(&val.to_le_bytes()); bytes[0] = header_byte;
len + 1u8 bytes[1..].copy_from_slice(&val.to_le_bytes());
MiniBuffer {
bytes,
len: len + 1,
}
} }
NumericalValue::I64(val) => { NumericalValue::I64(val) => {
let zig_zag_encoded = encode_zig_zag(val); let encoded = encode_zig_zag(val);
let len = compute_num_bytes_for_u64(zig_zag_encoded) as u8; let len = compute_num_bytes_for_u64(encoded);
output[0] = NumericalType::I64 as u8; let header_byte = compute_header_byte(SymbolType::I64, len);
output[1..9].copy_from_slice(&zig_zag_encoded.to_le_bytes()); bytes[0] = header_byte;
len + 1u8 bytes[1..].copy_from_slice(&encoded.to_le_bytes());
MiniBuffer {
bytes,
len: len + 1,
}
} }
} }
} }
} }
impl SymbolValue for u32 { #[allow(clippy::from_over_into)]
fn serialize(self, output: &mut [u8]) -> u8 { impl Into<MiniBuffer> for UnorderedId {
let len = compute_num_bytes_for_u64(self as u64); fn into(self) -> MiniBuffer {
output[0..4].copy_from_slice(&self.to_le_bytes()); let mut bytes = [0u8; 9];
len as u8 let val = self.0 as u64;
} let len = compute_num_bytes_for_u64(val) + 1;
bytes[0] = len as u8;
fn deserialize(bytes: &[u8]) -> Self { bytes[1..].copy_from_slice(&val.to_le_bytes());
let mut quartet: [u8; 4] = [0u8; 4]; MiniBuffer { bytes, len }
quartet[..bytes.len()].copy_from_slice(bytes);
u32::from_le_bytes(quartet)
} }
} }
impl SymbolValue for UnorderedId { impl SymbolValue for UnorderedId {
fn serialize(self, output: &mut [u8]) -> u8 { fn deserialize(header: NonZeroU8, bytes: &mut &[u8]) -> Result<UnorderedId, ParseError> {
self.0.serialize(output) let len = header.get() as usize;
} let symbol_bytes: &[u8];
(symbol_bytes, *bytes) = bytes.split_at(len);
fn deserialize(bytes: &[u8]) -> Self { let mut value_bytes = [0u8; 4];
UnorderedId(u32::deserialize(bytes)) value_bytes[..len - 1].copy_from_slice(&symbol_bytes[1..]);
let value = u32::from_le_bytes(value_bytes);
Ok(UnorderedId(value))
} }
} }
const HEADER_MASK: u8 = (1u8 << 4) - 1u8;
fn compute_num_bytes_for_u64(val: u64) -> usize { fn compute_num_bytes_for_u64(val: u64) -> usize {
let msb = (64u32 - val.leading_zeros()) as usize; let msb = (64u32 - val.leading_zeros()) as usize;
(msb + 7) / 8 (msb + 7) / 8
} }
fn parse_header_byte(byte: NonZeroU8) -> Result<(SymbolType, usize), ParseError> {
let len = (byte.get() as usize) >> 4;
let typ_code = byte.get() & HEADER_MASK;
let typ = SymbolType::try_from(typ_code)?;
Ok((typ, len))
}
#[derive(Error, Debug)]
pub enum ParseError {
#[error("Type byte unknown `{0}`")]
UnknownType(u8),
#[error("Invalid len for type `{len}` for type `{typ:?}`.")]
InvalidLen { typ: SymbolType, len: usize },
#[error("Missing bytes.")]
MissingBytes,
#[error("Not a number value.")]
NaN,
}
impl<V: SymbolValue> ColumnOperation<V> {
pub fn serialize(self) -> MiniBuffer {
match self {
ColumnOperation::NewDoc(doc) => {
let mut minibuf: [u8; 9] = [0u8; 9];
minibuf[0] = 0u8;
minibuf[1..5].copy_from_slice(&doc.to_le_bytes());
MiniBuffer {
bytes: minibuf,
len: 5,
}
}
ColumnOperation::Value(val) => val.into(),
}
}
pub fn deserialize(bytes: &mut &[u8]) -> Result<Self, ParseError> {
if bytes.is_empty() {
return Err(ParseError::MissingBytes);
}
let header_byte = bytes[0];
*bytes = &bytes[1..];
if let Some(header_byte) = NonZeroU8::new(header_byte) {
let value = V::deserialize(header_byte, bytes)?;
Ok(ColumnOperation::Value(value))
} else {
let doc_bytes: &[u8];
(doc_bytes, *bytes) = bytes.split_at(4);
let doc: u32 =
u32::from_le_bytes(doc_bytes.try_into().map_err(|_| ParseError::MissingBytes)?);
Ok(ColumnOperation::NewDoc(doc))
}
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum SymbolType {
U64 = 1u8,
I64 = 2u8,
Float = 3u8,
}
impl TryFrom<u8> for SymbolType {
type Error = ParseError;
fn try_from(byte: u8) -> Result<Self, ParseError> {
match byte {
1u8 => Ok(SymbolType::U64),
2u8 => Ok(SymbolType::I64),
3u8 => Ok(SymbolType::Float),
_ => Err(ParseError::UnknownType(byte)),
}
}
}
fn encode_zig_zag(n: i64) -> u64 { fn encode_zig_zag(n: i64) -> u64 {
((n << 1) ^ (n >> 63)) as u64 ((n << 1) ^ (n >> 63)) as u64
} }
@@ -210,7 +236,7 @@ fn decode_zig_zag(n: u64) -> i64 {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::{SymbolType, *};
#[track_caller] #[track_caller]
fn test_zig_zag_aux(val: i64) { fn test_zig_zag_aux(val: i64) {
@@ -242,27 +268,31 @@ mod tests {
} }
} }
#[track_caller]
fn ser_deser_header_byte_aux(symbol_type: SymbolType, len: usize) {
let header_byte = compute_header_byte(symbol_type, len);
let (serdeser_numerical_type, serdeser_len) =
parse_header_byte(NonZeroU8::new(header_byte).unwrap()).unwrap();
assert_eq!(symbol_type, serdeser_numerical_type);
assert_eq!(len, serdeser_len);
}
#[test] #[test]
fn test_header_byte_serialization() { fn test_header_byte_serialization() {
for len in 0..=15 { for len in 1..9 {
for typ_code in 0..=15 { ser_deser_header_byte_aux(SymbolType::Float, len);
let header = ColumnOperationHeader { typ_code, len }; ser_deser_header_byte_aux(SymbolType::I64, len);
let header_code = header.to_code(); ser_deser_header_byte_aux(SymbolType::U64, len);
let serdeser_header = ColumnOperationHeader::from_code(header_code);
assert_eq!(header, serdeser_header);
}
} }
} }
#[track_caller] #[track_caller]
fn ser_deser_symbol(column_op: ColumnOperation<NumericalValue>) { fn ser_deser_symbol(symbol: ColumnOperation<NumericalValue>) {
let buf = column_op.serialize(); let buf = symbol.serialize();
let mut buffer = buf.as_ref().to_vec(); let mut bytes = &buf.bytes[..];
buffer.extend_from_slice(b"234234");
let mut bytes = &buffer[..];
let serdeser_symbol = ColumnOperation::deserialize(&mut bytes).unwrap(); let serdeser_symbol = ColumnOperation::deserialize(&mut bytes).unwrap();
assert_eq!(bytes.len() + buf.as_ref().len() as usize, buffer.len()); assert_eq!(bytes.len() + buf.len, buf.bytes.len());
assert_eq!(column_op, serdeser_symbol); assert_eq!(symbol, serdeser_symbol);
} }
#[test] #[test]
@@ -288,24 +318,4 @@ mod tests {
ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(u64::MIN))); ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(u64::MIN)));
ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(u64::MAX))); ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(u64::MAX)));
} }
fn test_column_operation_unordered_aux(val: u32, expected_len: usize) {
let column_op = ColumnOperation::Value(UnorderedId(val));
let minibuf = column_op.serialize();
assert_eq!(minibuf.as_ref().len() as usize, expected_len);
let mut buf = minibuf.as_ref().to_vec();
buf.extend_from_slice(&[2, 2, 2, 2, 2, 2]);
let mut cursor = &buf[..];
let column_op_serdeser: ColumnOperation<UnorderedId> =
ColumnOperation::deserialize(&mut cursor).unwrap();
assert_eq!(column_op_serdeser, ColumnOperation::Value(UnorderedId(val)));
assert_eq!(cursor.len() + expected_len, buf.len());
}
#[test]
fn test_column_operation_unordered() {
test_column_operation_unordered_aux(300u32, 3);
test_column_operation_unordered_aux(1u32, 2);
test_column_operation_unordered_aux(0u32, 1);
}
} }

View File

@@ -1,270 +0,0 @@
use std::cmp::Ordering;
use stacker::{ExpUnrolledLinkedList, MemoryArena};
use crate::dictionary::{DictionaryBuilder, UnorderedId};
use crate::writer::column_operation::{ColumnOperation, SymbolValue};
use crate::{Cardinality, DocId, NumericalType, NumericalValue};
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[repr(u8)]
enum DocumentStep {
SameDoc = 0,
NextDoc = 1,
SkippedDoc = 2,
}
#[inline(always)]
fn delta_with_last_doc(last_doc_opt: Option<u32>, doc: u32) -> DocumentStep {
let expected_next_doc = last_doc_opt.map(|last_doc| last_doc + 1).unwrap_or(0u32);
match doc.cmp(&expected_next_doc) {
Ordering::Less => DocumentStep::SameDoc,
Ordering::Equal => DocumentStep::NextDoc,
Ordering::Greater => DocumentStep::SkippedDoc,
}
}
#[derive(Copy, Clone, Default)]
pub struct ColumnWriter {
// Detected cardinality of the column so far.
cardinality: Cardinality,
// Last document inserted.
// None if no doc has been added yet.
last_doc_opt: Option<u32>,
// Buffer containing the serialized values.
values: ExpUnrolledLinkedList,
}
impl ColumnWriter {
/// Returns an iterator over the Symbol that have been recorded
/// for the given column.
pub(crate) fn operation_iterator<'a, V: SymbolValue>(
&self,
arena: &MemoryArena,
buffer: &'a mut Vec<u8>,
) -> impl Iterator<Item = ColumnOperation<V>> + 'a {
buffer.clear();
self.values.read_to_end(arena, buffer);
let mut cursor: &[u8] = &buffer[..];
std::iter::from_fn(move || ColumnOperation::deserialize(&mut cursor))
}
/// Records a change of the document being recorded.
///
/// This function will also update the cardinality of the column
/// if necessary.
pub(crate) fn record<S: SymbolValue>(&mut self, doc: DocId, value: S, arena: &mut MemoryArena) {
// Difference between `doc` and the last doc.
match delta_with_last_doc(self.last_doc_opt, doc) {
DocumentStep::SameDoc => {
// This is the last encounterred document.
self.cardinality = Cardinality::Multivalued;
}
DocumentStep::NextDoc => {
self.last_doc_opt = Some(doc);
self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena);
}
DocumentStep::SkippedDoc => {
self.cardinality = self.cardinality.max(Cardinality::Optional);
self.last_doc_opt = Some(doc);
self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena);
}
}
self.write_symbol(ColumnOperation::Value(value), arena);
}
// Get the cardinality.
// The overall number of docs in the column is necessary to
// deal with the case where the all docs contain 1 value, except some documents
// at the end of the column.
pub fn get_cardinality(&self, num_docs: DocId) -> Cardinality {
match delta_with_last_doc(self.last_doc_opt, num_docs) {
DocumentStep::SameDoc | DocumentStep::NextDoc => self.cardinality,
DocumentStep::SkippedDoc => self.cardinality.max(Cardinality::Optional),
}
}
/// Appends a new symbol to the `ColumnWriter`.
fn write_symbol<V: SymbolValue>(
&mut self,
column_operation: ColumnOperation<V>,
arena: &mut MemoryArena,
) {
self.values
.writer(arena)
.extend_from_slice(column_operation.serialize().as_ref());
}
}
#[derive(Clone, Copy, Default)]
pub(crate) struct NumericalColumnWriter {
compatible_numerical_types: CompatibleNumericalTypes,
column_writer: ColumnWriter,
}
/// State used to store what types are still acceptable
/// after having seen a set of numerical values.
#[derive(Clone, Copy)]
pub(crate) struct CompatibleNumericalTypes {
all_values_within_i64_range: bool,
all_values_within_u64_range: bool,
// f64 is always acceptable.
}
impl Default for CompatibleNumericalTypes {
fn default() -> CompatibleNumericalTypes {
CompatibleNumericalTypes {
all_values_within_i64_range: true,
all_values_within_u64_range: true,
}
}
}
impl CompatibleNumericalTypes {
fn accept_value(&mut self, numerical_value: NumericalValue) {
match numerical_value {
NumericalValue::I64(val_i64) => {
let value_within_u64_range = val_i64 >= 0i64;
self.all_values_within_u64_range &= value_within_u64_range;
}
NumericalValue::U64(val_u64) => {
let value_within_i64_range = val_u64 < i64::MAX as u64;
self.all_values_within_i64_range &= value_within_i64_range;
}
NumericalValue::F64(_) => {
self.all_values_within_i64_range = false;
self.all_values_within_u64_range = false;
}
}
}
pub fn to_numerical_type(self) -> NumericalType {
if self.all_values_within_i64_range {
NumericalType::I64
} else if self.all_values_within_u64_range {
NumericalType::U64
} else {
NumericalType::F64
}
}
}
impl NumericalColumnWriter {
pub fn column_type_and_cardinality(&self, num_docs: DocId) -> (NumericalType, Cardinality) {
let numerical_type = self.compatible_numerical_types.to_numerical_type();
let cardinality = self.column_writer.get_cardinality(num_docs);
(numerical_type, cardinality)
}
pub fn record_numerical_value(
&mut self,
doc: DocId,
value: NumericalValue,
arena: &mut MemoryArena,
) {
self.compatible_numerical_types.accept_value(value);
self.column_writer.record(doc, value, arena);
}
pub fn operation_iterator<'a>(
self,
arena: &MemoryArena,
buffer: &'a mut Vec<u8>,
) -> impl Iterator<Item = ColumnOperation<NumericalValue>> + 'a {
self.column_writer.operation_iterator(arena, buffer)
}
}
#[derive(Copy, Clone, Default)]
pub struct StrColumnWriter {
pub(crate) dictionary_id: u32,
pub(crate) column_writer: ColumnWriter,
}
impl StrColumnWriter {
pub fn with_dictionary_id(dictionary_id: u32) -> StrColumnWriter {
StrColumnWriter {
dictionary_id,
column_writer: Default::default(),
}
}
pub(crate) fn record_bytes(
&mut self,
doc: DocId,
bytes: &[u8],
dictionaries: &mut [DictionaryBuilder],
arena: &mut MemoryArena,
) {
let unordered_id = dictionaries[self.dictionary_id as usize].get_or_allocate_id(bytes);
self.column_writer.record(doc, unordered_id, arena);
}
pub(crate) fn operation_iterator<'a>(
&self,
arena: &MemoryArena,
byte_buffer: &'a mut Vec<u8>,
) -> impl Iterator<Item = ColumnOperation<UnorderedId>> + 'a {
self.column_writer.operation_iterator(arena, byte_buffer)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_delta_with_last_doc() {
assert_eq!(delta_with_last_doc(None, 0u32), DocumentStep::NextDoc);
assert_eq!(delta_with_last_doc(None, 1u32), DocumentStep::SkippedDoc);
assert_eq!(delta_with_last_doc(None, 2u32), DocumentStep::SkippedDoc);
assert_eq!(delta_with_last_doc(Some(0u32), 0u32), DocumentStep::SameDoc);
assert_eq!(delta_with_last_doc(Some(1u32), 1u32), DocumentStep::SameDoc);
assert_eq!(delta_with_last_doc(Some(1u32), 2u32), DocumentStep::NextDoc);
assert_eq!(
delta_with_last_doc(Some(1u32), 3u32),
DocumentStep::SkippedDoc
);
assert_eq!(
delta_with_last_doc(Some(1u32), 4u32),
DocumentStep::SkippedDoc
);
}
#[track_caller]
fn test_column_writer_coercion_iter_aux(
values: impl Iterator<Item = NumericalValue>,
expected_numerical_type: NumericalType,
) {
let mut compatible_numerical_types = CompatibleNumericalTypes::default();
for value in values {
compatible_numerical_types.accept_value(value);
}
assert_eq!(
compatible_numerical_types.to_numerical_type(),
expected_numerical_type
);
}
#[track_caller]
fn test_column_writer_coercion_aux(
values: &[NumericalValue],
expected_numerical_type: NumericalType,
) {
test_column_writer_coercion_iter_aux(values.iter().copied(), expected_numerical_type);
test_column_writer_coercion_iter_aux(values.iter().rev().copied(), expected_numerical_type);
}
#[test]
fn test_column_writer_coercion() {
test_column_writer_coercion_aux(&[], NumericalType::I64);
test_column_writer_coercion_aux(&[1i64.into()], NumericalType::I64);
test_column_writer_coercion_aux(&[1u64.into()], NumericalType::I64);
// We don't detect exact integer at the moment. We could!
test_column_writer_coercion_aux(&[1f64.into()], NumericalType::F64);
test_column_writer_coercion_aux(&[u64::MAX.into()], NumericalType::U64);
test_column_writer_coercion_aux(&[(i64::MAX as u64).into()], NumericalType::U64);
test_column_writer_coercion_aux(&[(1u64 << 63).into()], NumericalType::U64);
test_column_writer_coercion_aux(&[1i64.into(), 1u64.into()], NumericalType::I64);
test_column_writer_coercion_aux(&[u64::MAX.into(), (-1i64).into()], NumericalType::F64);
}
}

View File

@@ -1,42 +1,196 @@
mod column_operation; mod column_operation;
mod column_writers;
mod serializer;
mod value_index; mod value_index;
use std::io::{self, Write}; use std::io::{self, Write};
use column_operation::ColumnOperation; use column_operation::ColumnOperation;
use common::CountingWriter;
use fastfield_codecs::serialize::ValueIndexInfo; use fastfield_codecs::serialize::ValueIndexInfo;
use fastfield_codecs::{Column, MonotonicallyMappableToU64, VecColumn}; use fastfield_codecs::{Column, MonotonicallyMappableToU64, VecColumn};
use serializer::ColumnarSerializer; use ordered_float::NotNan;
use stacker::{Addr, ArenaHashMap, MemoryArena}; use stacker::{Addr, ArenaHashMap, ExpUnrolledLinkedList, MemoryArena};
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality, GeneralType}; use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality};
use crate::dictionary::{DictionaryBuilder, IdMapping, UnorderedId}; use crate::dictionary::{DictionaryBuilder, IdMapping, UnorderedId};
use crate::value::{Coerce, NumericalType, NumericalValue}; use crate::value::{Coerce, NumericalType, NumericalValue};
use crate::writer::column_writers::{ColumnWriter, NumericalColumnWriter, StrColumnWriter}; use crate::writer::column_operation::SymbolValue;
use crate::writer::value_index::{IndexBuilder, SpareIndexBuilders}; use crate::writer::value_index::{IndexBuilder, SpareIndexBuilders};
use crate::{Cardinality, DocId}; use crate::{Cardinality, ColumnarSerializer, DocId};
/// Threshold above which a column data will be compressed #[derive(Copy, Clone, Default)]
/// using ZSTD. struct ColumnWriter {
const COLUMN_COMPRESSION_THRESHOLD: usize = 100_000; // Detected cardinality of the column so far.
cardinality: Cardinality,
// Last document inserted.
// None if no doc has been added yet.
last_doc_opt: Option<u32>,
// Buffer containing the serialized values.
values: ExpUnrolledLinkedList,
}
/// This is a set of buffers that are only here #[derive(Clone, Copy, Default)]
/// to limit the amount of allocation. pub struct NumericalColumnWriter {
#[derive(Default)] compatible_numerical_types: CompatibleNumericalTypes,
struct SpareBuffers { column_writer: ColumnWriter,
value_index_builders: SpareIndexBuilders, }
i64_values: Vec<i64>,
u64_values: Vec<u64>, #[derive(Clone, Copy)]
f64_values: Vec<f64>, struct CompatibleNumericalTypes {
bool_values: Vec<bool>, all_values_within_i64_range: bool,
column_buffer: Vec<u8>, all_values_within_u64_range: bool,
}
impl Default for CompatibleNumericalTypes {
fn default() -> CompatibleNumericalTypes {
CompatibleNumericalTypes {
all_values_within_i64_range: true,
all_values_within_u64_range: true,
}
}
}
impl CompatibleNumericalTypes {
pub fn accept_value(&mut self, numerical_value: NumericalValue) {
match numerical_value {
NumericalValue::I64(val_i64) => {
let value_within_u64_range = val_i64 >= 0i64;
self.all_values_within_u64_range &= value_within_u64_range;
}
NumericalValue::U64(val_u64) => {
let value_within_i64_range = val_u64 < i64::MAX as u64;
self.all_values_within_i64_range &= value_within_i64_range;
}
NumericalValue::F64(_) => {
self.all_values_within_i64_range = false;
self.all_values_within_u64_range = false;
}
}
}
pub fn to_numerical_type(self) -> NumericalType {
if self.all_values_within_i64_range {
NumericalType::I64
} else if self.all_values_within_u64_range {
NumericalType::U64
} else {
NumericalType::F64
}
}
}
impl NumericalColumnWriter {
pub fn record_numerical_value(
&mut self,
doc: DocId,
value: NumericalValue,
arena: &mut MemoryArena,
) {
self.compatible_numerical_types.accept_value(value);
self.column_writer.record(doc, value, arena);
}
}
impl ColumnWriter {
fn symbol_iterator<'a, V: SymbolValue>(
&self,
arena: &MemoryArena,
buffer: &'a mut Vec<u8>,
) -> impl Iterator<Item = ColumnOperation<V>> + 'a {
buffer.clear();
self.values.read_to_end(arena, buffer);
let mut cursor: &[u8] = &buffer[..];
std::iter::from_fn(move || {
if cursor.is_empty() {
return None;
}
let symbol = ColumnOperation::deserialize(&mut cursor)
.expect("Failed to deserialize symbol from in-memory. This should never happen.");
Some(symbol)
})
}
fn delta_with_last_doc(&self, doc: DocId) -> u32 {
self.last_doc_opt
.map(|last_doc| doc - last_doc)
.unwrap_or(doc + 1u32)
}
/// Records a change of the document being recorded.
///
/// This function will also update the cardinality of the column
/// if necessary.
fn record(&mut self, doc: DocId, value: NumericalValue, arena: &mut MemoryArena) {
// Difference between `doc` and the last doc.
match self.delta_with_last_doc(doc) {
0 => {
// This is the last encounterred document.
self.cardinality = Cardinality::Multivalued;
}
1 => {
self.last_doc_opt = Some(doc);
self.write_symbol::<NumericalValue>(ColumnOperation::NewDoc(doc), arena);
}
_ => {
self.cardinality = self.cardinality.max(Cardinality::Optional);
self.last_doc_opt = Some(doc);
self.write_symbol::<NumericalValue>(ColumnOperation::NewDoc(doc), arena);
}
}
self.write_symbol(ColumnOperation::Value(value), arena);
}
// Get the cardinality.
// The overall number of docs in the column is necessary to
// deal with the case where the all docs contain 1 value, except some documents
// at the end of the column.
fn get_cardinality(&self, num_docs: DocId) -> Cardinality {
if self.delta_with_last_doc(num_docs) > 1 {
self.cardinality.max(Cardinality::Optional)
} else {
self.cardinality
}
}
fn write_symbol<V: SymbolValue>(
&mut self,
symbol: ColumnOperation<V>,
arena: &mut MemoryArena,
) {
self.values
.writer(arena)
.extend_from_slice(symbol.serialize().as_slice());
}
}
#[derive(Copy, Clone, Default)]
pub struct BytesColumnWriter {
dictionary_id: u32,
column_writer: ColumnWriter,
}
impl BytesColumnWriter {
pub fn with_dictionary_id(dictionary_id: u32) -> BytesColumnWriter {
BytesColumnWriter {
dictionary_id,
column_writer: Default::default(),
}
}
pub fn record_bytes(
&mut self,
doc: DocId,
bytes: &[u8],
dictionaries: &mut [DictionaryBuilder],
arena: &mut MemoryArena,
) {
let unordered_id = dictionaries[self.dictionary_id as usize].get_or_allocate_id(bytes);
let numerical_value = NumericalValue::U64(unordered_id.0 as u64);
self.column_writer.record(doc, numerical_value, arena);
}
} }
pub struct ColumnarWriter { pub struct ColumnarWriter {
numerical_field_hash_map: ArenaHashMap, numerical_field_hash_map: ArenaHashMap,
bool_field_hash_map: ArenaHashMap,
bytes_field_hash_map: ArenaHashMap, bytes_field_hash_map: ArenaHashMap,
arena: MemoryArena, arena: MemoryArena,
// Dictionaries used to store dictionary-encoded values. // Dictionaries used to store dictionary-encoded values.
@@ -44,11 +198,19 @@ pub struct ColumnarWriter {
buffers: SpareBuffers, buffers: SpareBuffers,
} }
#[derive(Default)]
struct SpareBuffers {
byte_buffer: Vec<u8>,
value_index_builders: SpareIndexBuilders,
i64_values: Vec<i64>,
u64_values: Vec<u64>,
f64_values: Vec<ordered_float::NotNan<f64>>,
}
impl Default for ColumnarWriter { impl Default for ColumnarWriter {
fn default() -> Self { fn default() -> Self {
ColumnarWriter { ColumnarWriter {
numerical_field_hash_map: ArenaHashMap::new(10_000), numerical_field_hash_map: ArenaHashMap::new(10_000),
bool_field_hash_map: ArenaHashMap::new(10_000),
bytes_field_hash_map: ArenaHashMap::new(10_000), bytes_field_hash_map: ArenaHashMap::new(10_000),
dictionaries: Vec::new(), dictionaries: Vec::new(),
arena: MemoryArena::default(), arena: MemoryArena::default(),
@@ -57,304 +219,235 @@ impl Default for ColumnarWriter {
} }
} }
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
enum BytesOrNumerical {
Bytes,
Numerical,
}
impl ColumnarWriter { impl ColumnarWriter {
pub fn record_numerical( pub fn record_numerical(&mut self, doc: DocId, key: &[u8], numerical_value: NumericalValue) {
&mut self,
doc: DocId,
column_name: &str,
numerical_value: NumericalValue,
) {
assert!(
!column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte"
);
let (hash_map, arena) = (&mut self.numerical_field_hash_map, &mut self.arena); let (hash_map, arena) = (&mut self.numerical_field_hash_map, &mut self.arena);
hash_map.mutate_or_create( hash_map.mutate_or_create(key, |column_opt: Option<NumericalColumnWriter>| {
column_name.as_bytes(), let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
|column_opt: Option<NumericalColumnWriter>| { column.record_numerical_value(doc, numerical_value, arena);
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default(); column
column.record_numerical_value(doc, numerical_value, arena); });
column
},
);
} }
pub fn record_bool(&mut self, doc: DocId, column_name: &str, val: bool) { pub fn record_bytes(&mut self, doc: DocId, key: &[u8], value: &[u8]) {
assert!(
!column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte"
);
let (hash_map, arena) = (&mut self.bool_field_hash_map, &mut self.arena);
hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<ColumnWriter>| {
let mut column: ColumnWriter = column_opt.unwrap_or_default();
column.record(doc, val, arena);
column
},
);
}
pub fn record_str(&mut self, doc: DocId, column_name: &str, value: &[u8]) {
assert!(
!column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte"
);
let (hash_map, arena, dictionaries) = ( let (hash_map, arena, dictionaries) = (
&mut self.bytes_field_hash_map, &mut self.bytes_field_hash_map,
&mut self.arena, &mut self.arena,
&mut self.dictionaries, &mut self.dictionaries,
); );
hash_map.mutate_or_create( hash_map.mutate_or_create(key, |column_opt: Option<BytesColumnWriter>| {
column_name.as_bytes(), let mut column: BytesColumnWriter = column_opt.unwrap_or_else(|| {
|column_opt: Option<StrColumnWriter>| { let dictionary_id = dictionaries.len() as u32;
let mut column: StrColumnWriter = column_opt.unwrap_or_else(|| { dictionaries.push(DictionaryBuilder::default());
let dictionary_id = dictionaries.len() as u32; BytesColumnWriter::with_dictionary_id(dictionary_id)
dictionaries.push(DictionaryBuilder::default()); });
StrColumnWriter::with_dictionary_id(dictionary_id) column.record_bytes(doc, value, dictionaries, arena);
}); column
column.record_bytes(doc, value, dictionaries, arena); });
column
},
);
} }
pub fn serialize(&mut self, num_docs: DocId, wrt: &mut dyn io::Write) -> io::Result<()> { pub fn serialize<W: io::Write>(
let mut serializer = ColumnarSerializer::new(wrt); &mut self,
let mut field_columns: Vec<(&[u8], GeneralType, Addr)> = self num_docs: DocId,
mut serializer: ColumnarSerializer<W>,
) -> io::Result<()> {
let mut field_columns: Vec<(&[u8], BytesOrNumerical, Addr)> = self
.numerical_field_hash_map .numerical_field_hash_map
.iter() .iter()
.map(|(term, addr, _)| (term, GeneralType::Numerical, addr)) .map(|(term, addr, _)| (term, BytesOrNumerical::Numerical, addr))
.collect(); .collect();
field_columns.extend( field_columns.extend(
self.bytes_field_hash_map self.bytes_field_hash_map
.iter() .iter()
.map(|(term, addr, _)| (term, GeneralType::Str, addr)), .map(|(term, addr, _)| (term, BytesOrNumerical::Bytes, addr)),
); );
field_columns.extend( let mut key_buffer = Vec::new();
self.bool_field_hash_map field_columns.sort_unstable_by_key(|(key, col_type, _)| (*key, *col_type));
.iter()
.map(|(term, addr, _)| (term, GeneralType::Bool, addr)),
);
field_columns.sort_unstable_by_key(|(column_name, col_type, _)| (*column_name, *col_type));
let (arena, buffers, dictionaries) = (&self.arena, &mut self.buffers, &self.dictionaries); let (arena, buffers, dictionaries) = (&self.arena, &mut self.buffers, &self.dictionaries);
let mut symbol_byte_buffer: Vec<u8> = Vec::new(); for (key, bytes_or_numerical, addr) in field_columns {
for (column_name, bytes_or_numerical, addr) in field_columns { let wrt = serializer.wrt();
match bytes_or_numerical { let start_offset = wrt.written_bytes();
GeneralType::Bool => { let column_type_and_cardinality: ColumnTypeAndCardinality =
let column_writer: ColumnWriter = self.bool_field_hash_map.read(addr); match bytes_or_numerical {
let cardinality = column_writer.get_cardinality(num_docs); BytesOrNumerical::Bytes => {
let column_type_and_cardinality = ColumnTypeAndCardinality { let BytesColumnWriter { dictionary_id, column_writer } =
cardinality, self.bytes_field_hash_map.read(addr);
typ: ColumnType::Bool,
};
let column_serializer =
serializer.serialize_column(column_name, column_type_and_cardinality);
serialize_bool_column(
cardinality,
num_docs,
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
column_serializer,
)?;
}
GeneralType::Str => {
let str_column_writer: StrColumnWriter = self.bytes_field_hash_map.read(addr);
let dictionary_builder = let dictionary_builder =
&dictionaries[str_column_writer.dictionary_id as usize]; &dictionaries[dictionary_id as usize];
let cardinality = str_column_writer.column_writer.get_cardinality(num_docs);
let column_type_and_cardinality = ColumnTypeAndCardinality {
cardinality,
typ: ColumnType::Bytes,
};
let column_serializer =
serializer.serialize_column(column_name, column_type_and_cardinality);
serialize_bytes_column( serialize_bytes_column(
cardinality, &column_writer,
num_docs, num_docs,
dictionary_builder, dictionary_builder,
str_column_writer.operation_iterator(arena, &mut symbol_byte_buffer), arena,
buffers, buffers,
column_serializer, wrt,
)?; )?;
ColumnTypeAndCardinality {
cardinality: column_writer.get_cardinality(num_docs),
typ: ColumnType::Bytes,
}
} }
GeneralType::Numerical => { BytesOrNumerical::Numerical => {
let numerical_column_writer: NumericalColumnWriter = let NumericalColumnWriter { compatible_numerical_types, column_writer } =
self.numerical_field_hash_map.read(addr); self.numerical_field_hash_map.read(addr);
let (numerical_type, cardinality) = let cardinality = column_writer.get_cardinality(num_docs);
numerical_column_writer.column_type_and_cardinality(num_docs); let numerical_type = compatible_numerical_types.to_numerical_type();
let column_type_and_cardinality = ColumnTypeAndCardinality {
cardinality,
typ: ColumnType::Numerical(numerical_type),
};
let column_serializer =
serializer.serialize_column(column_name, column_type_and_cardinality);
serialize_numerical_column( serialize_numerical_column(
cardinality, cardinality,
num_docs,
numerical_type, numerical_type,
numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer), &column_writer,
num_docs,
arena,
buffers, buffers,
column_serializer, wrt,
)?; )?;
ColumnTypeAndCardinality {
cardinality,
typ: ColumnType::Numerical(numerical_type),
}
} }
}; };
let end_offset = wrt.written_bytes();
let key_with_type = prepare_key(key, column_type_and_cardinality, &mut key_buffer);
serializer.record_column_offsets(key_with_type, start_offset..end_offset)?;
} }
serializer.finalize()?; serializer.finalize()?;
Ok(()) Ok(())
} }
} }
fn compress_and_write_column<W: io::Write>(column_bytes: &[u8], wrt: &mut W) -> io::Result<()> { /// Returns a key consisting of the concatenation of the key and the column_type_and_cardinality
if column_bytes.len() >= COLUMN_COMPRESSION_THRESHOLD { /// code.
wrt.write_all(&[1])?; fn prepare_key<'a>(
let mut encoder = zstd::Encoder::new(wrt, 3)?; key: &[u8],
encoder.write_all(column_bytes)?; column_type_cardinality: ColumnTypeAndCardinality,
encoder.finish()?; buffer: &'a mut Vec<u8>,
} else { ) -> &'a [u8] {
wrt.write_all(&[0])?; buffer.clear();
wrt.write_all(column_bytes)?; buffer.extend_from_slice(key);
} buffer.push(0u8);
Ok(()) buffer.push(column_type_cardinality.to_code());
&buffer[..]
} }
fn serialize_bytes_column<W: io::Write>( fn serialize_bytes_column<W: io::Write>(
cardinality: Cardinality, column_writer: &ColumnWriter,
num_docs: DocId, num_docs: DocId,
dictionary_builder: &DictionaryBuilder, dictionary_builder: &DictionaryBuilder,
operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>, arena: &MemoryArena,
buffers: &mut SpareBuffers, buffers: &mut SpareBuffers,
mut wrt: W, wrt: &mut CountingWriter<W>,
) -> io::Result<()> { ) -> io::Result<()> {
let start_offset = wrt.written_bytes();
let id_mapping: IdMapping = dictionary_builder.serialize(wrt)?;
let dictionary_num_bytes: u32 = (wrt.written_bytes() - start_offset) as u32;
let cardinality = column_writer.get_cardinality(num_docs);
let SpareBuffers { let SpareBuffers {
byte_buffer,
value_index_builders, value_index_builders,
u64_values, u64_values,
column_buffer,
.. ..
} = buffers; } = buffers;
column_buffer.clear(); let symbol_iterator = column_writer
let id_mapping: IdMapping = dictionary_builder.serialize(column_buffer)?; .symbol_iterator(arena, byte_buffer)
let dictionary_num_bytes: u32 = column_buffer.len() as u32; .map(|symbol: ColumnOperation<UnorderedId>| {
let operation_iterator = operation_it.map(|symbol: ColumnOperation<UnorderedId>| { // We map unordered ids to ordered ids.
// We map unordered ids to ordered ids. match symbol {
match symbol { ColumnOperation::Value(unordered_id) => {
ColumnOperation::Value(unordered_id) => { let ordered_id = id_mapping.to_ord(unordered_id);
let ordered_id = id_mapping.to_ord(unordered_id); ColumnOperation::Value(ordered_id.0 as u64)
ColumnOperation::Value(ordered_id.0 as u64) }
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
} }
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc), });
}
});
serialize_column( serialize_column(
operation_iterator, symbol_iterator,
cardinality, cardinality,
num_docs, num_docs,
value_index_builders, value_index_builders,
u64_values, u64_values,
column_buffer, wrt,
)?; )?;
column_buffer.write_all(&dictionary_num_bytes.to_le_bytes()[..])?; wrt.write_all(&dictionary_num_bytes.to_le_bytes()[..])?;
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(()) Ok(())
} }
fn serialize_numerical_column<W: io::Write>( fn serialize_numerical_column<W: io::Write>(
cardinality: Cardinality, cardinality: Cardinality,
num_docs: DocId,
numerical_type: NumericalType, numerical_type: NumericalType,
op_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>, column_writer: &ColumnWriter,
num_docs: DocId,
arena: &MemoryArena,
buffers: &mut SpareBuffers, buffers: &mut SpareBuffers,
mut wrt: W, wrt: &mut W,
) -> io::Result<()> { ) -> io::Result<()> {
let SpareBuffers { let SpareBuffers {
byte_buffer,
value_index_builders, value_index_builders,
u64_values, u64_values,
i64_values, i64_values,
f64_values, f64_values,
column_buffer,
..
} = buffers; } = buffers;
column_buffer.clear(); let symbol_iterator = column_writer.symbol_iterator(arena, byte_buffer);
match numerical_type { match numerical_type {
NumericalType::I64 => { NumericalType::I64 => {
serialize_column( serialize_column(
coerce_numerical_symbol::<i64>(op_iterator), coerce_numerical_symbol::<i64>(symbol_iterator),
cardinality, cardinality,
num_docs, num_docs,
value_index_builders, value_index_builders,
i64_values, i64_values,
column_buffer, wrt,
)?; )?;
} }
NumericalType::U64 => { NumericalType::U64 => {
serialize_column( serialize_column(
coerce_numerical_symbol::<u64>(op_iterator), coerce_numerical_symbol::<u64>(symbol_iterator),
cardinality, cardinality,
num_docs, num_docs,
value_index_builders, value_index_builders,
u64_values, u64_values,
column_buffer, wrt,
)?; )?;
} }
NumericalType::F64 => { NumericalType::F64 => {
serialize_column( serialize_column(
coerce_numerical_symbol::<f64>(op_iterator), coerce_numerical_symbol::<NotNan<f64>>(symbol_iterator),
cardinality, cardinality,
num_docs, num_docs,
value_index_builders, value_index_builders,
f64_values, f64_values,
column_buffer, wrt,
)?; )?;
} }
}; };
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(())
}
fn serialize_bool_column<W: io::Write>(
cardinality: Cardinality,
num_docs: DocId,
column_operations_it: impl Iterator<Item = ColumnOperation<bool>>,
buffers: &mut SpareBuffers,
mut wrt: W,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
bool_values,
column_buffer,
..
} = buffers;
column_buffer.clear();
serialize_column(
column_operations_it,
cardinality,
num_docs,
value_index_builders,
bool_values,
column_buffer,
)?;
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(()) Ok(())
} }
fn serialize_column< fn serialize_column<
T: Copy + Default + std::fmt::Debug + Send + Sync + MonotonicallyMappableToU64 + PartialOrd, T: Copy + Ord + Default + Send + Sync + MonotonicallyMappableToU64,
W: io::Write,
>( >(
op_iterator: impl Iterator<Item = ColumnOperation<T>>, symbol_iterator: impl Iterator<Item = ColumnOperation<T>>,
cardinality: Cardinality, cardinality: Cardinality,
num_docs: DocId, num_docs: DocId,
value_index_builders: &mut SpareIndexBuilders, value_index_builders: &mut SpareIndexBuilders,
values: &mut Vec<T>, values: &mut Vec<T>,
wrt: &mut Vec<u8>, wrt: &mut W,
) -> io::Result<()> ) -> io::Result<()>
where where
for<'a> VecColumn<'a, T>: Column<T>, for<'a> VecColumn<'a, T>: Column<T>,
{ {
values.clear();
match cardinality { match cardinality {
Cardinality::Required => { Cardinality::Required => {
consume_operation_iterator( consume_symbol_iterator(
op_iterator, symbol_iterator,
value_index_builders.borrow_required_index_builder(), value_index_builders.borrow_required_index_builder(),
values, values,
); );
@@ -366,7 +459,7 @@ where
} }
Cardinality::Optional => { Cardinality::Optional => {
let optional_index_builder = value_index_builders.borrow_optional_index_builder(); let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values); consume_symbol_iterator(symbol_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_docs); let optional_index = optional_index_builder.finish(num_docs);
fastfield_codecs::serialize::serialize_new( fastfield_codecs::serialize::serialize_new(
ValueIndexInfo::SingleValue(Box::new(optional_index)), ValueIndexInfo::SingleValue(Box::new(optional_index)),
@@ -377,7 +470,7 @@ where
} }
Cardinality::Multivalued => { Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder(); let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values); consume_symbol_iterator(symbol_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_docs); let multivalued_index = multivalued_index_builder.finish(num_docs);
fastfield_codecs::serialize::serialize_new( fastfield_codecs::serialize::serialize_new(
ValueIndexInfo::MultiValue(Box::new(multivalued_index)), ValueIndexInfo::MultiValue(Box::new(multivalued_index)),
@@ -391,10 +484,10 @@ where
} }
fn coerce_numerical_symbol<T>( fn coerce_numerical_symbol<T>(
operation_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>, symbol_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
) -> impl Iterator<Item = ColumnOperation<T>> ) -> impl Iterator<Item = ColumnOperation<T>>
where T: Coerce { where T: Coerce {
operation_iterator.map(|symbol| match symbol { symbol_iterator.map(|symbol| match symbol {
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc), ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
ColumnOperation::Value(numerical_value) => { ColumnOperation::Value(numerical_value) => {
ColumnOperation::Value(Coerce::coerce(numerical_value)) ColumnOperation::Value(Coerce::coerce(numerical_value))
@@ -402,12 +495,12 @@ where T: Coerce {
}) })
} }
fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>( fn consume_symbol_iterator<T, TIndexBuilder: IndexBuilder>(
operation_iterator: impl Iterator<Item = ColumnOperation<T>>, symbol_iterator: impl Iterator<Item = ColumnOperation<T>>,
index_builder: &mut TIndexBuilder, index_builder: &mut TIndexBuilder,
values: &mut Vec<T>, values: &mut Vec<T>,
) { ) {
for symbol in operation_iterator { for symbol in symbol_iterator {
match symbol { match symbol {
ColumnOperation::NewDoc(doc) => { ColumnOperation::NewDoc(doc) => {
index_builder.record_doc(doc); index_builder.record_doc(doc);
@@ -422,24 +515,42 @@ fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use column_operation::ColumnOperation;
use ordered_float::NotNan;
use stacker::MemoryArena; use stacker::MemoryArena;
use super::*; use super::prepare_key;
use crate::value::NumericalValue; use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality};
use crate::value::{NumericalType, NumericalValue};
use crate::writer::column_operation::ColumnOperation;
use crate::writer::CompatibleNumericalTypes;
use crate::Cardinality; use crate::Cardinality;
#[test]
fn test_prepare_key_bytes() {
let mut buffer: Vec<u8> = b"somegarbage".to_vec();
let column_type_and_cardinality = ColumnTypeAndCardinality {
typ: ColumnType::Bytes,
cardinality: Cardinality::Optional,
};
let prepared_key = prepare_key(b"root\0child", column_type_and_cardinality, &mut buffer);
assert_eq!(prepared_key.len(), 12);
assert_eq!(&prepared_key[..10], b"root\0child");
assert_eq!(prepared_key[10], 0u8);
assert_eq!(prepared_key[11], column_type_and_cardinality.to_code());
}
#[test] #[test]
fn test_column_writer_required_simple() { fn test_column_writer_required_simple() {
let mut arena = MemoryArena::default(); let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default(); let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(14i64), &mut arena); column_writer.record(0u32, 14i64.into(), &mut arena);
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena); column_writer.record(1u32, 15i64.into(), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena); column_writer.record(2u32, (-16i64).into(), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Required); assert_eq!(column_writer.get_cardinality(3), Cardinality::Required);
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer) .symbol_iterator(&mut arena, &mut buffer)
.collect(); .collect();
assert_eq!(symbols.len(), 6); assert_eq!(symbols.len(), 6);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32))); assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
@@ -463,12 +574,12 @@ mod tests {
fn test_column_writer_optional_cardinality_missing_first() { fn test_column_writer_optional_cardinality_missing_first() {
let mut arena = MemoryArena::default(); let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default(); let mut column_writer = super::ColumnWriter::default();
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena); column_writer.record(1u32, 15i64.into(), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena); column_writer.record(2u32, (-16i64).into(), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Optional); assert_eq!(column_writer.get_cardinality(3), Cardinality::Optional);
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer) .symbol_iterator(&mut arena, &mut buffer)
.collect(); .collect();
assert_eq!(symbols.len(), 4); assert_eq!(symbols.len(), 4);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(1u32))); assert!(matches!(symbols[0], ColumnOperation::NewDoc(1u32)));
@@ -487,11 +598,11 @@ mod tests {
fn test_column_writer_optional_cardinality_missing_last() { fn test_column_writer_optional_cardinality_missing_last() {
let mut arena = MemoryArena::default(); let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default(); let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(15i64), &mut arena); column_writer.record(0u32, 15i64.into(), &mut arena);
assert_eq!(column_writer.get_cardinality(2), Cardinality::Optional); assert_eq!(column_writer.get_cardinality(2), Cardinality::Optional);
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer) .symbol_iterator(&mut arena, &mut buffer)
.collect(); .collect();
assert_eq!(symbols.len(), 2); assert_eq!(symbols.len(), 2);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32))); assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
@@ -505,12 +616,12 @@ mod tests {
fn test_column_writer_multivalued() { fn test_column_writer_multivalued() {
let mut arena = MemoryArena::default(); let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default(); let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(16i64), &mut arena); column_writer.record(0u32, 16i64.into(), &mut arena);
column_writer.record(0u32, NumericalValue::from(17i64), &mut arena); column_writer.record(0u32, 17i64.into(), &mut arena);
assert_eq!(column_writer.get_cardinality(1), Cardinality::Multivalued); assert_eq!(column_writer.get_cardinality(1), Cardinality::Multivalued);
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer) .symbol_iterator(&mut arena, &mut buffer)
.collect(); .collect();
assert_eq!(symbols.len(), 3); assert_eq!(symbols.len(), 3);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32))); assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
@@ -523,4 +634,42 @@ mod tests {
ColumnOperation::Value(NumericalValue::I64(17i64)) ColumnOperation::Value(NumericalValue::I64(17i64))
)); ));
} }
#[track_caller]
fn test_column_writer_coercion_iter_aux(
values: impl Iterator<Item = NumericalValue>,
expected_numerical_type: NumericalType,
) {
let mut compatible_numerical_types = CompatibleNumericalTypes::default();
for value in values {
compatible_numerical_types.accept_value(value);
}
assert_eq!(
compatible_numerical_types.to_numerical_type(),
expected_numerical_type
);
}
#[track_caller]
fn test_column_writer_coercion_aux(
values: &[NumericalValue],
expected_numerical_type: NumericalType,
) {
test_column_writer_coercion_iter_aux(values.iter().copied(), expected_numerical_type);
test_column_writer_coercion_iter_aux(values.iter().rev().copied(), expected_numerical_type);
}
#[test]
fn test_column_writer_coercion() {
test_column_writer_coercion_aux(&[], NumericalType::I64);
test_column_writer_coercion_aux(&[1i64.into()], NumericalType::I64);
test_column_writer_coercion_aux(&[1u64.into()], NumericalType::I64);
// We don't detect exact integer at the moment. We could!
test_column_writer_coercion_aux(&[NotNan::new(1f64).unwrap().into()], NumericalType::F64);
test_column_writer_coercion_aux(&[u64::MAX.into()], NumericalType::U64);
test_column_writer_coercion_aux(&[(i64::MAX as u64).into()], NumericalType::U64);
test_column_writer_coercion_aux(&[(1u64 << 63).into()], NumericalType::U64);
test_column_writer_coercion_aux(&[1i64.into(), 1u64.into()], NumericalType::I64);
test_column_writer_coercion_aux(&[u64::MAX.into(), (-1i64).into()], NumericalType::F64);
}
} }

View File

@@ -1,116 +0,0 @@
use std::io;
use std::io::Write;
use common::CountingWriter;
use sstable::value::RangeValueWriter;
use sstable::RangeSSTable;
use crate::column_type_header::ColumnTypeAndCardinality;
pub struct ColumnarSerializer<W: io::Write> {
wrt: CountingWriter<W>,
sstable_range: sstable::Writer<Vec<u8>, RangeValueWriter>,
prepare_key_buffer: Vec<u8>,
}
/// Returns a key consisting of the concatenation of the key and the column_type_and_cardinality
/// code.
fn prepare_key<'a>(
key: &[u8],
column_type_cardinality: ColumnTypeAndCardinality,
buffer: &'a mut Vec<u8>,
) {
buffer.clear();
buffer.extend_from_slice(key);
buffer.push(0u8);
buffer.push(column_type_cardinality.to_code());
}
impl<W: io::Write> ColumnarSerializer<W> {
pub(crate) fn new(wrt: W) -> ColumnarSerializer<W> {
let sstable_range: sstable::Writer<Vec<u8>, RangeValueWriter> =
sstable::Dictionary::<RangeSSTable>::builder(Vec::with_capacity(100_000)).unwrap();
ColumnarSerializer {
wrt: CountingWriter::wrap(wrt),
sstable_range,
prepare_key_buffer: Vec::new(),
}
}
pub fn serialize_column<'a>(
&'a mut self,
column_name: &[u8],
column_type_cardinality: ColumnTypeAndCardinality,
) -> impl io::Write + 'a {
let start_offset = self.wrt.written_bytes();
prepare_key(
column_name,
column_type_cardinality,
&mut self.prepare_key_buffer,
);
ColumnSerializer {
columnar_serializer: self,
start_offset,
}
}
pub(crate) fn finalize(mut self) -> io::Result<()> {
let sstable_bytes: Vec<u8> = self.sstable_range.finish()?;
let sstable_num_bytes: u64 = sstable_bytes.len() as u64;
self.wrt.write_all(&sstable_bytes)?;
self.wrt.write_all(&sstable_num_bytes.to_le_bytes()[..])?;
Ok(())
}
}
struct ColumnSerializer<'a, W: io::Write> {
columnar_serializer: &'a mut ColumnarSerializer<W>,
start_offset: u64,
}
impl<'a, W: io::Write> Drop for ColumnSerializer<'a, W> {
fn drop(&mut self) {
let end_offset: u64 = self.columnar_serializer.wrt.written_bytes();
let byte_range = self.start_offset..end_offset;
self.columnar_serializer.sstable_range.insert_cannot_fail(
&self.columnar_serializer.prepare_key_buffer[..],
&byte_range,
);
self.columnar_serializer.prepare_key_buffer.clear();
}
}
impl<'a, W: io::Write> io::Write for ColumnSerializer<'a, W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.columnar_serializer.wrt.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.columnar_serializer.wrt.flush()
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
self.columnar_serializer.wrt.write_all(buf)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::column_type_header::ColumnType;
use crate::Cardinality;
#[test]
fn test_prepare_key_bytes() {
let mut buffer: Vec<u8> = b"somegarbage".to_vec();
let column_type_and_cardinality = ColumnTypeAndCardinality {
typ: ColumnType::Bytes,
cardinality: Cardinality::Optional,
};
prepare_key(b"root\0child", column_type_and_cardinality, &mut buffer);
assert_eq!(buffer.len(), 12);
assert_eq!(&buffer[..10], b"root\0child");
assert_eq!(buffer[10], 0u8);
assert_eq!(buffer[11], column_type_and_cardinality.to_code());
}
}

View File

@@ -24,11 +24,8 @@ pub trait FileHandle: 'static + Send + Sync + HasLen + fmt::Debug {
fn read_bytes(&self, range: Range<usize>) -> io::Result<OwnedBytes>; fn read_bytes(&self, range: Range<usize>) -> io::Result<OwnedBytes>;
#[doc(hidden)] #[doc(hidden)]
async fn read_bytes_async(&self, _byte_range: Range<usize>) -> io::Result<OwnedBytes> { async fn read_bytes_async(&self, byte_range: Range<usize>) -> io::Result<OwnedBytes> {
Err(io::Error::new( self.read_bytes(byte_range)
io::ErrorKind::Unsupported,
"Async read is not supported.",
))
} }
} }
@@ -39,6 +36,7 @@ impl FileHandle for &'static [u8] {
Ok(OwnedBytes::new(bytes)) Ok(OwnedBytes::new(bytes))
} }
#[cfg(feature = "quickwit")]
async fn read_bytes_async(&self, byte_range: Range<usize>) -> io::Result<OwnedBytes> { async fn read_bytes_async(&self, byte_range: Range<usize>) -> io::Result<OwnedBytes> {
Ok(self.read_bytes(byte_range)?) Ok(self.read_bytes(byte_range)?)
} }
@@ -67,16 +65,7 @@ impl fmt::Debug for FileSlice {
} }
} }
/// Takes a range, a `RangeBounds` object, and returns #[inline]
/// a `Range` that corresponds to the relative application of the
/// `RangeBounds` object to the original `Range`.
///
/// For instance, combine_ranges(`[2..11)`, `[5..7]`) returns `[7..10]`
/// as it reads, what is the sub-range that starts at the 5 element of
/// `[2..11)` and ends at the 9th element included.
///
/// This function panics, if the result would suggest something outside
/// of the bounds of the original range.
fn combine_ranges<R: RangeBounds<usize>>(orig_range: Range<usize>, rel_range: R) -> Range<usize> { fn combine_ranges<R: RangeBounds<usize>>(orig_range: Range<usize>, rel_range: R) -> Range<usize> {
let start: usize = orig_range.start let start: usize = orig_range.start
+ match rel_range.start_bound().cloned() { + match rel_range.start_bound().cloned() {
@@ -224,6 +213,7 @@ impl FileHandle for FileSlice {
self.read_bytes_slice(range) self.read_bytes_slice(range)
} }
#[cfg(feature = "quickwit")]
async fn read_bytes_async(&self, byte_range: Range<usize>) -> io::Result<OwnedBytes> { async fn read_bytes_async(&self, byte_range: Range<usize>) -> io::Result<OwnedBytes> {
self.read_bytes_slice_async(byte_range).await self.read_bytes_slice_async(byte_range).await
} }
@@ -241,15 +231,16 @@ impl FileHandle for OwnedBytes {
Ok(self.slice(range)) Ok(self.slice(range))
} }
#[cfg(feature = "quickwit")]
async fn read_bytes_async(&self, range: Range<usize>) -> io::Result<OwnedBytes> { async fn read_bytes_async(&self, range: Range<usize>) -> io::Result<OwnedBytes> {
self.read_bytes(range) let bytes = self.read_bytes(range)?;
Ok(bytes)
} }
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::io; use std::io;
use std::ops::Bound;
use std::sync::Arc; use std::sync::Arc;
use super::{FileHandle, FileSlice}; use super::{FileHandle, FileSlice};
@@ -332,11 +323,6 @@ mod tests {
assert_eq!(combine_ranges(1..3, 1..), 2..3); assert_eq!(combine_ranges(1..3, 1..), 2..3);
assert_eq!(combine_ranges(1..4, ..2), 1..3); assert_eq!(combine_ranges(1..4, ..2), 1..3);
assert_eq!(combine_ranges(3..10, 2..5), 5..8); assert_eq!(combine_ranges(3..10, 2..5), 5..8);
assert_eq!(combine_ranges(2..11, 5..=7), 7..10);
assert_eq!(
combine_ranges(2..11, (Bound::Excluded(5), Bound::Unbounded)),
8..11
);
} }
#[test] #[test]

View File

@@ -10,7 +10,7 @@ mod serialize;
mod vint; mod vint;
mod writer; mod writer;
pub use bitset::*; pub use bitset::*;
pub use ownedbytes::{OwnedBytes, StableDeref}; pub use ownedbytes::OwnedBytes;
pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize}; pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize};
pub use vint::{ pub use vint::{
deserialize_vint_u128, read_u32_vint, read_u32_vint_no_advance, serialize_vint_u128, deserialize_vint_u128, read_u32_vint, read_u32_vint_no_advance, serialize_vint_u128,

View File

@@ -14,12 +14,14 @@ repository = "https://github.com/quickwit-oss/tantivy"
[dependencies] [dependencies]
common = { version = "0.5", path = "../common/", package = "tantivy-common" } common = { version = "0.5", path = "../common/", package = "tantivy-common" }
tantivy-bitpacker = { version= "0.3", path = "../bitpacker/" } tantivy-bitpacker = { version= "0.3", path = "../bitpacker/" }
ownedbytes = { version = "0.5", path = "../ownedbytes" }
prettytable-rs = {version="0.9.0", optional= true} prettytable-rs = {version="0.9.0", optional= true}
rand = {version="0.8.3", optional= true} rand = {version="0.8.3", optional= true}
fastdivide = "0.4" fastdivide = "0.4"
log = "0.4" log = "0.4"
itertools = { version = "0.10.3" } itertools = { version = "0.10.3" }
measure_time = { version="0.8.2", optional=true} measure_time = { version="0.8.2", optional=true}
ordered-float = "3.4"
[dev-dependencies] [dev-dependencies]
more-asserts = "0.3.0" more-asserts = "0.3.0"

View File

@@ -7,8 +7,8 @@ mod tests {
use std::iter; use std::iter;
use std::sync::Arc; use std::sync::Arc;
use common::OwnedBytes;
use fastfield_codecs::*; use fastfield_codecs::*;
use ownedbytes::OwnedBytes;
use rand::prelude::*; use rand::prelude::*;
use test::Bencher; use test::Bencher;

View File

@@ -1,6 +1,6 @@
use std::io::{self, Write}; use std::io::{self, Write};
use common::OwnedBytes; use ownedbytes::OwnedBytes;
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker}; use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
use crate::serialize::NormalizedHeader; use crate::serialize::NormalizedHeader;

View File

@@ -1,7 +1,8 @@
use std::sync::Arc; use std::sync::Arc;
use std::{io, iter}; use std::{io, iter};
use common::{BinarySerializable, CountingWriter, DeserializeFrom, OwnedBytes}; use common::{BinarySerializable, CountingWriter, DeserializeFrom};
use ownedbytes::OwnedBytes;
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker}; use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
use crate::line::Line; use crate::line::Line;
@@ -46,7 +47,7 @@ impl FastFieldCodec for BlockwiseLinearCodec {
type Reader = BlockwiseLinearReader; type Reader = BlockwiseLinearReader;
fn open_from_bytes( fn open_from_bytes(
bytes: common::OwnedBytes, bytes: ownedbytes::OwnedBytes,
normalized_header: NormalizedHeader, normalized_header: NormalizedHeader,
) -> io::Result<Self::Reader> { ) -> io::Result<Self::Reader> {
let footer_len: u32 = (&bytes[bytes.len() - 4..]).deserialize()?; let footer_len: u32 = (&bytes[bytes.len() - 4..]).deserialize()?;
@@ -170,18 +171,15 @@ impl Column for BlockwiseLinearReader {
interpoled_val.wrapping_add(bitpacked_diff) interpoled_val.wrapping_add(bitpacked_diff)
} }
#[inline(always)]
fn min_value(&self) -> u64 { fn min_value(&self) -> u64 {
// The BlockwiseLinearReader assumes a normalized vector. // The BlockwiseLinearReader assumes a normalized vector.
0u64 0u64
} }
#[inline(always)]
fn max_value(&self) -> u64 { fn max_value(&self) -> u64 {
self.normalized_header.max_value self.normalized_header.max_value
} }
#[inline(always)]
fn num_vals(&self) -> u32 { fn num_vals(&self) -> u32 {
self.normalized_header.num_vals self.normalized_header.num_vals
} }

View File

@@ -135,7 +135,7 @@ impl<'a, T: Copy + PartialOrd + Send + Sync> Column<T> for VecColumn<'a, T> {
} }
} }
impl<'a, T: Copy + PartialOrd + Default, V> From<&'a V> for VecColumn<'a, T> impl<'a, T: Copy + Ord + Default, V> From<&'a V> for VecColumn<'a, T>
where V: AsRef<[T]> + ?Sized where V: AsRef<[T]> + ?Sized
{ {
fn from(values: &'a V) -> Self { fn from(values: &'a V) -> Self {

View File

@@ -17,7 +17,8 @@ use std::{
ops::{Range, RangeInclusive}, ops::{Range, RangeInclusive},
}; };
use common::{BinarySerializable, CountingWriter, OwnedBytes, VInt, VIntU128}; use common::{BinarySerializable, CountingWriter, VInt, VIntU128};
use ownedbytes::OwnedBytes;
use tantivy_bitpacker::{self, BitPacker, BitUnpacker}; use tantivy_bitpacker::{self, BitPacker, BitUnpacker};
use crate::compact_space::build_compact_space::get_compact_space; use crate::compact_space::build_compact_space::get_compact_space;

View File

@@ -1,6 +1,7 @@
use std::io; use std::io;
use common::{BinarySerializable, OwnedBytes}; use common::BinarySerializable;
use ownedbytes::OwnedBytes;
const MAGIC_NUMBER: u16 = 4335u16; const MAGIC_NUMBER: u16 = 4335u16;
const FASTFIELD_FORMAT_VERSION: u8 = 1; const FASTFIELD_FORMAT_VERSION: u8 = 1;

View File

@@ -45,7 +45,7 @@ mod tests {
use std::io; use std::io;
use std::num::NonZeroU64; use std::num::NonZeroU64;
use common::OwnedBytes; use ownedbytes::OwnedBytes;
use crate::gcd::{compute_gcd, find_gcd}; use crate::gcd::{compute_gcd, find_gcd};
use crate::{FastFieldCodecType, VecColumn}; use crate::{FastFieldCodecType, VecColumn};

View File

@@ -18,7 +18,7 @@ use std::io;
use std::io::Write; use std::io::Write;
use std::sync::Arc; use std::sync::Arc;
use common::{BinarySerializable, OwnedBytes}; use common::BinarySerializable;
use compact_space::CompactSpaceDecompressor; use compact_space::CompactSpaceDecompressor;
use format_version::read_format_version; use format_version::read_format_version;
use monotonic_mapping::{ use monotonic_mapping::{
@@ -26,6 +26,7 @@ use monotonic_mapping::{
StrictlyMonotonicMappingToInternalBaseval, StrictlyMonotonicMappingToInternalGCDBaseval, StrictlyMonotonicMappingToInternalBaseval, StrictlyMonotonicMappingToInternalGCDBaseval,
}; };
use null_index_footer::read_null_index_footer; use null_index_footer::read_null_index_footer;
use ownedbytes::OwnedBytes;
use serialize::{Header, U128Header}; use serialize::{Header, U128Header};
mod bitpacked; mod bitpacked;
@@ -44,6 +45,8 @@ mod column;
mod gcd; mod gcd;
pub mod serialize; pub mod serialize;
pub use ordered_float;
use self::bitpacked::BitpackedCodec; use self::bitpacked::BitpackedCodec;
use self::blockwise_linear::BlockwiseLinearCodec; use self::blockwise_linear::BlockwiseLinearCodec;
pub use self::column::{monotonic_map_column, Column, IterColumn, VecColumn}; pub use self::column::{monotonic_map_column, Column, IterColumn, VecColumn};
@@ -435,7 +438,7 @@ mod tests {
mod bench { mod bench {
use std::sync::Arc; use std::sync::Arc;
use common::OwnedBytes; use ownedbytes::OwnedBytes;
use rand::rngs::StdRng; use rand::rngs::StdRng;
use rand::{Rng, SeedableRng}; use rand::{Rng, SeedableRng};
use test::{self, Bencher}; use test::{self, Bencher};

View File

@@ -1,6 +1,7 @@
use std::io::{self, Write}; use std::io::{self, Write};
use common::{BinarySerializable, OwnedBytes}; use common::BinarySerializable;
use ownedbytes::OwnedBytes;
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker}; use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
use crate::line::Line; use crate::line::Line;
@@ -24,13 +25,13 @@ impl Column for LinearReader {
interpoled_val.wrapping_add(bitpacked_diff) interpoled_val.wrapping_add(bitpacked_diff)
} }
#[inline(always)] #[inline]
fn min_value(&self) -> u64 { fn min_value(&self) -> u64 {
// The LinearReader assumes a normalized vector. // The LinearReader assumes a normalized vector.
0u64 0u64
} }
#[inline(always)] #[inline]
fn max_value(&self) -> u64 { fn max_value(&self) -> u64 {
self.header.max_value self.header.max_value
} }

View File

@@ -6,10 +6,10 @@ use std::io::BufRead;
use std::net::{IpAddr, Ipv6Addr}; use std::net::{IpAddr, Ipv6Addr};
use std::str::FromStr; use std::str::FromStr;
use common::OwnedBytes;
use fastfield_codecs::{open_u128, serialize_u128, Column, FastFieldCodecType, VecColumn}; use fastfield_codecs::{open_u128, serialize_u128, Column, FastFieldCodecType, VecColumn};
use itertools::Itertools; use itertools::Itertools;
use measure_time::print_time; use measure_time::print_time;
use ownedbytes::OwnedBytes;
use prettytable::{Cell, Row, Table}; use prettytable::{Cell, Row, Table};
fn print_set_stats(ip_addrs: &[u128]) { fn print_set_stats(ip_addrs: &[u128]) {

View File

@@ -1,6 +1,7 @@
use std::marker::PhantomData; use std::marker::PhantomData;
use fastdivide::DividerU64; use fastdivide::DividerU64;
use ordered_float::NotNan;
use crate::MonotonicallyMappableToU128; use crate::MonotonicallyMappableToU128;
@@ -56,12 +57,10 @@ impl<T> From<T> for StrictlyMonotonicMappingInverter<T> {
impl<From, To, T> StrictlyMonotonicFn<To, From> for StrictlyMonotonicMappingInverter<T> impl<From, To, T> StrictlyMonotonicFn<To, From> for StrictlyMonotonicMappingInverter<T>
where T: StrictlyMonotonicFn<From, To> where T: StrictlyMonotonicFn<From, To>
{ {
#[inline(always)]
fn mapping(&self, val: To) -> From { fn mapping(&self, val: To) -> From {
self.orig_mapping.inverse(val) self.orig_mapping.inverse(val)
} }
#[inline(always)]
fn inverse(&self, val: From) -> To { fn inverse(&self, val: From) -> To {
self.orig_mapping.mapping(val) self.orig_mapping.mapping(val)
} }
@@ -84,12 +83,10 @@ impl<External: MonotonicallyMappableToU128, T: MonotonicallyMappableToU128>
StrictlyMonotonicFn<External, u128> for StrictlyMonotonicMappingToInternal<T> StrictlyMonotonicFn<External, u128> for StrictlyMonotonicMappingToInternal<T>
where T: MonotonicallyMappableToU128 where T: MonotonicallyMappableToU128
{ {
#[inline(always)]
fn mapping(&self, inp: External) -> u128 { fn mapping(&self, inp: External) -> u128 {
External::to_u128(inp) External::to_u128(inp)
} }
#[inline(always)]
fn inverse(&self, out: u128) -> External { fn inverse(&self, out: u128) -> External {
External::from_u128(out) External::from_u128(out)
} }
@@ -99,12 +96,10 @@ impl<External: MonotonicallyMappableToU64, T: MonotonicallyMappableToU64>
StrictlyMonotonicFn<External, u64> for StrictlyMonotonicMappingToInternal<T> StrictlyMonotonicFn<External, u64> for StrictlyMonotonicMappingToInternal<T>
where T: MonotonicallyMappableToU64 where T: MonotonicallyMappableToU64
{ {
#[inline(always)]
fn mapping(&self, inp: External) -> u64 { fn mapping(&self, inp: External) -> u64 {
External::to_u64(inp) External::to_u64(inp)
} }
#[inline(always)]
fn inverse(&self, out: u64) -> External { fn inverse(&self, out: u64) -> External {
External::from_u64(out) External::from_u64(out)
} }
@@ -132,13 +127,11 @@ impl StrictlyMonotonicMappingToInternalGCDBaseval {
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64> impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
for StrictlyMonotonicMappingToInternalGCDBaseval for StrictlyMonotonicMappingToInternalGCDBaseval
{ {
#[inline(always)]
fn mapping(&self, inp: External) -> u64 { fn mapping(&self, inp: External) -> u64 {
self.gcd_divider self.gcd_divider
.divide(External::to_u64(inp) - self.min_value) .divide(External::to_u64(inp) - self.min_value)
} }
#[inline(always)]
fn inverse(&self, out: u64) -> External { fn inverse(&self, out: u64) -> External {
External::from_u64(self.min_value + out * self.gcd) External::from_u64(self.min_value + out * self.gcd)
} }
@@ -149,7 +142,6 @@ pub(crate) struct StrictlyMonotonicMappingToInternalBaseval {
min_value: u64, min_value: u64,
} }
impl StrictlyMonotonicMappingToInternalBaseval { impl StrictlyMonotonicMappingToInternalBaseval {
#[inline(always)]
pub(crate) fn new(min_value: u64) -> Self { pub(crate) fn new(min_value: u64) -> Self {
Self { min_value } Self { min_value }
} }
@@ -158,24 +150,20 @@ impl StrictlyMonotonicMappingToInternalBaseval {
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64> impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
for StrictlyMonotonicMappingToInternalBaseval for StrictlyMonotonicMappingToInternalBaseval
{ {
#[inline(always)]
fn mapping(&self, val: External) -> u64 { fn mapping(&self, val: External) -> u64 {
External::to_u64(val) - self.min_value External::to_u64(val) - self.min_value
} }
#[inline(always)]
fn inverse(&self, val: u64) -> External { fn inverse(&self, val: u64) -> External {
External::from_u64(self.min_value + val) External::from_u64(self.min_value + val)
} }
} }
impl MonotonicallyMappableToU64 for u64 { impl MonotonicallyMappableToU64 for u64 {
#[inline(always)]
fn to_u64(self) -> u64 { fn to_u64(self) -> u64 {
self self
} }
#[inline(always)]
fn from_u64(val: u64) -> Self { fn from_u64(val: u64) -> Self {
val val
} }
@@ -208,22 +196,51 @@ impl MonotonicallyMappableToU64 for bool {
// TODO remove me. // TODO remove me.
// Tantivy should refuse NaN values and work with NotNaN internally. // Tantivy should refuse NaN values and work with NotNaN internally.
impl MonotonicallyMappableToU64 for f64 { impl MonotonicallyMappableToU64 for f64 {
#[inline(always)]
fn to_u64(self) -> u64 { fn to_u64(self) -> u64 {
common::f64_to_u64(self) common::f64_to_u64(self)
} }
#[inline(always)]
fn from_u64(val: u64) -> Self { fn from_u64(val: u64) -> Self {
common::u64_to_f64(val) common::u64_to_f64(val)
} }
} }
impl MonotonicallyMappableToU64 for ordered_float::NotNan<f64> {
fn to_u64(self) -> u64 {
common::f64_to_u64(self.into_inner())
}
fn from_u64(val: u64) -> Self {
NotNan::new(common::u64_to_f64(val)).expect("Invalid NotNaN f64 value.")
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn test_from_u64_pos_inf() {
let inf_as_u64 = common::f64_to_u64(f64::INFINITY);
let inf_back_to_f64 = NotNan::from_u64(inf_as_u64);
assert_eq!(inf_back_to_f64, NotNan::new(f64::INFINITY).unwrap());
}
#[test]
fn test_from_u64_neg_inf() {
let inf_as_u64 = common::f64_to_u64(-f64::INFINITY);
let inf_back_to_f64 = NotNan::from_u64(inf_as_u64);
assert_eq!(inf_back_to_f64, NotNan::new(-f64::INFINITY).unwrap());
}
#[test]
#[should_panic(expected = "Invalid NotNaN")]
fn test_from_u64_nan_panics() {
let nan_as_u64 = common::f64_to_u64(f64::NAN);
NotNan::from_u64(nan_as_u64);
}
#[test] #[test]
fn strictly_monotonic_test() { fn strictly_monotonic_test() {
// identity mapping // identity mapping

View File

@@ -1,8 +1,9 @@
use std::convert::TryInto; use std::convert::TryInto;
use std::io::{self, Write}; use std::io::{self, Write};
use common::{BinarySerializable, OwnedBytes}; use common::BinarySerializable;
use itertools::Itertools; use itertools::Itertools;
use ownedbytes::OwnedBytes;
use super::{get_bit_at, set_bit_at}; use super::{get_bit_at, set_bit_at};

View File

@@ -1,6 +1,7 @@
use std::io::{self, Write}; use std::io::{self, Write};
use common::{BitSet, OwnedBytes}; use common::BitSet;
use ownedbytes::OwnedBytes;
use super::{serialize_dense_codec, DenseCodec}; use super::{serialize_dense_codec, DenseCodec};
@@ -58,7 +59,6 @@ enum SparseCodecBlockVariant {
impl SparseCodecBlockVariant { impl SparseCodecBlockVariant {
/// The number of non-null values that preceeded that block. /// The number of non-null values that preceeded that block.
#[inline]
fn offset(&self) -> u32 { fn offset(&self) -> u32 {
match self { match self {
SparseCodecBlockVariant::Empty { offset } => *offset, SparseCodecBlockVariant::Empty { offset } => *offset,

View File

@@ -1,7 +1,8 @@
use std::io::{self, Write}; use std::io::{self, Write};
use std::ops::Range; use std::ops::Range;
use common::{BinarySerializable, CountingWriter, OwnedBytes, VInt}; use common::{BinarySerializable, CountingWriter, VInt};
use ownedbytes::OwnedBytes;
#[derive(Debug, Clone, Copy, Eq, PartialEq)] #[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) enum FastFieldCardinality { pub(crate) enum FastFieldCardinality {

View File

@@ -1,28 +1,10 @@
// Copyright (C) 2022 Quickwit, Inc.
//
// Quickwit is offered under the AGPL v3.0 and as commercial software.
// For commercial licensing, contact us at hello@quickwit.io.
//
// AGPL:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::io; use std::io;
use std::num::NonZeroU64; use std::num::NonZeroU64;
use std::sync::Arc; use std::sync::Arc;
use common::{BinarySerializable, OwnedBytes, VInt}; use common::{BinarySerializable, VInt};
use log::warn; use log::warn;
use ownedbytes::OwnedBytes;
use crate::bitpacked::BitpackedCodec; use crate::bitpacked::BitpackedCodec;
use crate::blockwise_linear::BlockwiseLinearCodec; use crate::blockwise_linear::BlockwiseLinearCodec;
@@ -202,7 +184,6 @@ pub enum ValueIndexInfo<'a> {
SingleValue(Box<dyn SingleValueIndexInfo + 'a>), SingleValue(Box<dyn SingleValueIndexInfo + 'a>),
} }
// TODO Remove me
impl Default for ValueIndexInfo<'static> { impl Default for ValueIndexInfo<'static> {
fn default() -> Self { fn default() -> Self {
struct Dummy {} struct Dummy {}

View File

@@ -6,10 +6,10 @@ use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock, Weak}; use std::sync::{Arc, RwLock, Weak};
use std::{fmt, result}; use std::{fmt, result};
use common::StableDeref;
use fs2::FileExt; use fs2::FileExt;
use memmap2::Mmap; use memmap2::Mmap;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use stable_deref_trait::StableDeref;
use tempfile::TempDir; use tempfile::TempDir;
use crate::core::META_FILEPATH; use crate::core::META_FILEPATH;

View File

@@ -20,7 +20,8 @@ use std::io::BufWriter;
use std::path::PathBuf; use std::path::PathBuf;
pub use common::file_slice::{FileHandle, FileSlice}; pub use common::file_slice::{FileHandle, FileSlice};
pub use common::{AntiCallToken, OwnedBytes, TerminatingWrite}; pub use common::{AntiCallToken, TerminatingWrite};
pub use ownedbytes::OwnedBytes;
pub(crate) use self::composite_file::{CompositeFile, CompositeWrite}; pub(crate) use self::composite_file::{CompositeFile, CompositeWrite};
pub use self::directory::{Directory, DirectoryClone, DirectoryLock}; pub use self::directory::{Directory, DirectoryClone, DirectoryLock};

View File

@@ -1,7 +1,8 @@
use std::io; use std::io;
use std::io::Write; use std::io::Write;
use common::{intersect_bitsets, BitSet, OwnedBytes, ReadOnlyBitSet}; use common::{intersect_bitsets, BitSet, ReadOnlyBitSet};
use ownedbytes::OwnedBytes;
use crate::space_usage::ByteCount; use crate::space_usage::ByteCount;
use crate::DocId; use crate::DocId;

View File

@@ -80,7 +80,6 @@ impl MultiValueIndex {
/// ///
/// TODO: Instead of a linear scan we can employ a exponential search into binary search to /// TODO: Instead of a linear scan we can employ a exponential search into binary search to
/// match a docid to its value position. /// match a docid to its value position.
#[allow(clippy::bool_to_int_with_if)]
pub(crate) fn positions_to_docids(&self, doc_id_range: Range<u32>, positions: &mut Vec<u32>) { pub(crate) fn positions_to_docids(&self, doc_id_range: Range<u32>, positions: &mut Vec<u32>) {
if positions.is_empty() { if positions.is_empty() {
return; return;

View File

@@ -20,7 +20,7 @@ mod atomic_impl {
} }
pub fn fetch_add(&self, val: u64, order: Ordering) -> u64 { pub fn fetch_add(&self, val: u64, order: Ordering) -> u64 {
self.0.fetch_add(val, order) self.0.fetch_add(val, order) as u64
} }
pub fn revert(&self, val: u64, order: Ordering) -> u64 { pub fn revert(&self, val: u64, order: Ordering) -> u64 {

View File

@@ -213,21 +213,21 @@ impl<'a> FieldSerializer<'a> {
fail_point!("FieldSerializer::close_term", |msg: Option<String>| { fail_point!("FieldSerializer::close_term", |msg: Option<String>| {
Err(io::Error::new(io::ErrorKind::Other, format!("{:?}", msg))) Err(io::Error::new(io::ErrorKind::Other, format!("{:?}", msg)))
}); });
if self.term_open { if !self.term_open {
self.postings_serializer return Ok(());
.close_term(self.current_term_info.doc_freq)?;
self.current_term_info.postings_range.end =
self.postings_serializer.written_bytes() as usize;
if let Some(positions_serializer) = self.positions_serializer_opt.as_mut() {
positions_serializer.close_term()?;
self.current_term_info.positions_range.end =
positions_serializer.written_bytes() as usize;
}
self.term_dictionary_builder
.insert_value(&self.current_term_info)?;
self.term_open = false;
} }
self.postings_serializer
.close_term(self.current_term_info.doc_freq)?;
self.current_term_info.postings_range.end =
self.postings_serializer.written_bytes() as usize;
if let Some(positions_serializer) = self.positions_serializer_opt.as_mut() {
positions_serializer.close_term()?;
self.current_term_info.positions_range.end =
positions_serializer.written_bytes() as usize;
}
self.term_dictionary_builder
.insert_value(&self.current_term_info)?;
self.term_open = false;
Ok(()) Ok(())
} }

View File

@@ -3,7 +3,7 @@ use tantivy_query_grammar::Occur;
use crate::query::{BooleanWeight, DisjunctionMaxCombiner, EnableScoring, Query, Weight}; use crate::query::{BooleanWeight, DisjunctionMaxCombiner, EnableScoring, Query, Weight};
use crate::{Score, Term}; use crate::{Score, Term};
/// The disjunction max query returns documents matching one or more wrapped queries, /// The disjunction max query кeturns documents matching one or more wrapped queries,
/// called query clauses or clauses. /// called query clauses or clauses.
/// ///
/// If a returned document matches multiple query clauses, /// If a returned document matches multiple query clauses,

View File

@@ -280,7 +280,9 @@ impl DocSet for IpRangeDocSet {
#[inline] #[inline]
fn doc(&self) -> DocId { fn doc(&self) -> DocId {
self.loaded_docs.current().unwrap_or(TERMINATED) self.loaded_docs
.current()
.unwrap_or(TERMINATED)
} }
/// Advances the `DocSet` forward until reaching the target, or going to the /// Advances the `DocSet` forward until reaching the target, or going to the

View File

@@ -4,8 +4,9 @@ use std::ops::{AddAssign, Range};
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use common::{BinarySerializable, HasLen, OwnedBytes}; use common::{BinarySerializable, HasLen};
use lru::LruCache; use lru::LruCache;
use ownedbytes::OwnedBytes;
use super::footer::DocStoreFooter; use super::footer::DocStoreFooter;
use super::index::SkipIndex; use super::index::SkipIndex;

View File

@@ -1,6 +1,7 @@
use std::io; use std::io;
mod merger; mod merger;
mod termdict;
use std::iter::ExactSizeIterator; use std::iter::ExactSizeIterator;
@@ -12,39 +13,26 @@ use tantivy_fst::automaton::AlwaysMatch;
pub use self::merger::TermMerger; pub use self::merger::TermMerger;
use crate::postings::TermInfo; use crate::postings::TermInfo;
/// The term dictionary contains all of the terms in
/// `tantivy index` in a sorted manner.
///
/// The `Fst` crate is used to associate terms to their
/// respective `TermOrdinal`. The `TermInfoStore` then makes it
/// possible to fetch the associated `TermInfo`.
pub type TermDictionary = sstable::Dictionary<TermSSTable>; pub type TermDictionary = sstable::Dictionary<TermSSTable>;
pub type TermDictionaryBuilder<W> = sstable::Writer<W, TermInfoWriter>;
/// Builder for the new term dictionary.
pub type TermDictionaryBuilder<W> = sstable::Writer<W, TermInfoValueWriter>;
/// `TermStreamer` acts as a cursor over a range of terms of a segment.
/// Terms are guaranteed to be sorted.
pub type TermStreamer<'a, A = AlwaysMatch> = sstable::Streamer<'a, TermSSTable, A>; pub type TermStreamer<'a, A = AlwaysMatch> = sstable::Streamer<'a, TermSSTable, A>;
/// SSTable used to store TermInfo objects.
pub struct TermSSTable; pub struct TermSSTable;
impl SSTable for TermSSTable { impl SSTable for TermSSTable {
type Value = TermInfo; type Value = TermInfo;
type ValueReader = TermInfoValueReader; type ValueReader = TermInfoReader;
type ValueWriter = TermInfoValueWriter; type ValueWriter = TermInfoWriter;
} }
#[derive(Default)] #[derive(Default)]
pub struct TermInfoValueReader { pub struct TermInfoReader {
term_infos: Vec<TermInfo>, term_infos: Vec<TermInfo>,
} }
impl ValueReader for TermInfoValueReader { impl ValueReader for TermInfoReader {
type Value = TermInfo; type Value = TermInfo;
#[inline(always)]
fn value(&self, idx: usize) -> &TermInfo { fn value(&self, idx: usize) -> &TermInfo {
&self.term_infos[idx] &self.term_infos[idx]
} }
@@ -76,18 +64,18 @@ impl ValueReader for TermInfoValueReader {
} }
#[derive(Default)] #[derive(Default)]
pub struct TermInfoValueWriter { pub struct TermInfoWriter {
term_infos: Vec<TermInfo>, term_infos: Vec<TermInfo>,
} }
impl ValueWriter for TermInfoValueWriter { impl ValueWriter for TermInfoWriter {
type Value = TermInfo; type Value = TermInfo;
fn write(&mut self, term_info: &TermInfo) { fn write(&mut self, term_info: &TermInfo) {
self.term_infos.push(term_info.clone()); self.term_infos.push(term_info.clone());
} }
fn serialize_block(&self, buffer: &mut Vec<u8>) { fn serialize_block(&mut self, buffer: &mut Vec<u8>) {
VInt(self.term_infos.len() as u64).serialize_into_vec(buffer); VInt(self.term_infos.len() as u64).serialize_into_vec(buffer);
if self.term_infos.is_empty() { if self.term_infos.is_empty() {
return; return;
@@ -99,9 +87,6 @@ impl ValueWriter for TermInfoValueWriter {
VInt(term_info.postings_range.len() as u64).serialize_into_vec(buffer); VInt(term_info.postings_range.len() as u64).serialize_into_vec(buffer);
VInt(term_info.positions_range.len() as u64).serialize_into_vec(buffer); VInt(term_info.positions_range.len() as u64).serialize_into_vec(buffer);
} }
}
fn clear(&mut self) {
self.term_infos.clear(); self.term_infos.clear();
} }
} }
@@ -111,11 +96,11 @@ mod tests {
use sstable::value::{ValueReader, ValueWriter}; use sstable::value::{ValueReader, ValueWriter};
use crate::postings::TermInfo; use crate::postings::TermInfo;
use crate::termdict::sstable_termdict::TermInfoValueReader; use crate::termdict::sstable_termdict::TermInfoReader;
#[test] #[test]
fn test_block_terminfos() { fn test_block_terminfos() {
let mut term_info_writer = super::TermInfoValueWriter::default(); let mut term_info_writer = super::TermInfoWriter::default();
term_info_writer.write(&TermInfo { term_info_writer.write(&TermInfo {
doc_freq: 120u32, doc_freq: 120u32,
postings_range: 17..45, postings_range: 17..45,
@@ -133,7 +118,8 @@ mod tests {
}); });
let mut buffer = Vec::new(); let mut buffer = Vec::new();
term_info_writer.serialize_block(&mut buffer); term_info_writer.serialize_block(&mut buffer);
let mut term_info_reader = TermInfoValueReader::default(); // let mut block_reader = make_block_reader(&buffer[..]);
let mut term_info_reader = TermInfoReader::default();
let num_bytes: usize = term_info_reader.load(&buffer[..]).unwrap(); let num_bytes: usize = term_info_reader.load(&buffer[..]).unwrap();
assert_eq!( assert_eq!(
term_info_reader.value(0), term_info_reader.value(0),

View File

@@ -0,0 +1,11 @@
use sstable::SSTable;
use crate::postings::TermInfo;
use crate::termdict::sstable_termdict::{TermInfoReader, TermInfoWriter};
pub struct TermInfoSSTable;
impl SSTable for TermInfoSSTable {
type Value = TermInfo;
type ValueReader = TermInfoReader;
type ValueWriter = TermInfoWriter;
}

View File

@@ -2,7 +2,6 @@
name = "tantivy-sstable" name = "tantivy-sstable"
version = "0.1.0" version = "0.1.0"
edition = "2021" edition = "2021"
license = "MIT"
[dependencies] [dependencies]
common = {path="../common", package="tantivy-common"} common = {path="../common", package="tantivy-common"}
@@ -12,10 +11,3 @@ tantivy-fst = "0.4"
[dev-dependencies] [dev-dependencies]
proptest = "1" proptest = "1"
criterion = "0.4"
names = "0.14"
rand = "0.8"
[[bench]]
name = "stream_bench"
harness = false

View File

@@ -1,28 +0,0 @@
# SSTable
The `tantivy-sstable` crate is yet another sstable crate.
It has been designed to be used in `quickwit`:
- as an alternative to the default tantivy fst dictionary.
- as a way to store the column index for dynamic fast fields.
The benefit compared to the fst crate is locality.
Searching a key in the fst crate requires downloading the entire dictionary.
Once the sstable index is downloaded, running a `get` in the sstable
crate only requires a single fetch.
Right now, the block index and the default block size have been thought
for quickwit, and the performance of a get is very bad.
# Sorted strings?
SSTable stands for Sorted String Table.
Strings have to be insert in sorted order.
That sorted order is used in different ways:
- it makes gets and streaming ranges of keys
possible.
- it allows incremental encoding of the keys
- the front compression is leveraged to optimize
the intersection with an automaton

View File

@@ -1,87 +0,0 @@
use std::collections::BTreeSet;
use std::io;
use common::file_slice::FileSlice;
use criterion::{criterion_group, criterion_main, Criterion};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use tantivy_sstable::{self, Dictionary, MonotonicU64SSTable};
const CHARSET: &'static [u8] = b"abcdefghij";
fn generate_key(rng: &mut impl Rng) -> String {
let len = rng.gen_range(3..12);
std::iter::from_fn(|| {
let idx = rng.gen_range(0..CHARSET.len());
Some(CHARSET[idx] as char)
})
.take(len)
.collect()
}
fn prepare_sstable() -> io::Result<Dictionary<MonotonicU64SSTable>> {
let mut rng = StdRng::from_seed([3u8; 32]);
let mut els = BTreeSet::new();
while els.len() < 100_000 {
els.insert(generate_key(&mut rng));
}
let mut dictionary_builder = Dictionary::<MonotonicU64SSTable>::builder(Vec::new())?;
for (ord, word) in els.iter().enumerate() {
dictionary_builder.insert(word, &(ord as u64))?;
}
let buffer = dictionary_builder.finish()?;
let dictionary = Dictionary::open(FileSlice::from(buffer))?;
Ok(dictionary)
}
fn stream_bench(
dictionary: &Dictionary<MonotonicU64SSTable>,
lower: &[u8],
upper: &[u8],
do_scan: bool,
) -> usize {
let mut stream = dictionary
.range()
.ge(lower)
.lt(upper)
.into_stream()
.unwrap();
if !do_scan {
return 0;
}
let mut count = 0;
while stream.advance() {
count += 1;
}
count
}
pub fn criterion_benchmark(c: &mut Criterion) {
let dict = prepare_sstable().unwrap();
c.bench_function("short_scan_init", |b| {
b.iter(|| stream_bench(&dict, b"fa", b"fana", false))
});
c.bench_function("short_scan_init_and_scan", |b| {
b.iter(|| {
assert_eq!(stream_bench(&dict, b"fa", b"faz", true), 971);
})
});
c.bench_function("full_scan_init_and_scan_full_with_bound", |b| {
b.iter(|| {
assert_eq!(stream_bench(&dict, b"", b"z", true), 100_000);
})
});
c.bench_function("full_scan_init_and_scan_full_no_bounds", |b| {
b.iter(|| {
let mut stream = dict.stream().unwrap();
let mut count = 0;
while stream.advance() {
count += 1;
}
count
})
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);

View File

@@ -1,5 +1,4 @@
use std::io; use std::io;
use std::ops::Range;
pub struct BlockReader<'a> { pub struct BlockReader<'a> {
buffer: Vec<u8>, buffer: Vec<u8>,
@@ -30,8 +29,8 @@ impl<'a> BlockReader<'a> {
} }
#[inline(always)] #[inline(always)]
pub fn buffer_from_to(&self, range: Range<usize>) -> &[u8] { pub fn buffer_from_to(&self, start: usize, end: usize) -> &[u8] {
&self.buffer[range] &self.buffer[start..end]
} }
pub fn read_block(&mut self) -> io::Result<bool> { pub fn read_block(&mut self) -> io::Result<bool> {
@@ -52,17 +51,14 @@ impl<'a> BlockReader<'a> {
Ok(true) Ok(true)
} }
#[inline(always)]
pub fn offset(&self) -> usize { pub fn offset(&self) -> usize {
self.offset self.offset
} }
#[inline(always)]
pub fn advance(&mut self, num_bytes: usize) { pub fn advance(&mut self, num_bytes: usize) {
self.offset += num_bytes; self.offset += num_bytes;
} }
#[inline(always)]
pub fn buffer(&self) -> &[u8] { pub fn buffer(&self) -> &[u8] {
&self.buffer[self.offset..] &self.buffer[self.offset..]
} }

View File

@@ -16,8 +16,6 @@ where W: io::Write
block: Vec<u8>, block: Vec<u8>,
write: CountingWriter<BufWriter<W>>, write: CountingWriter<BufWriter<W>>,
value_writer: TValueWriter, value_writer: TValueWriter,
// Only here to avoid allocations.
stateless_buffer: Vec<u8>,
} }
impl<W, TValueWriter> DeltaWriter<W, TValueWriter> impl<W, TValueWriter> DeltaWriter<W, TValueWriter>
@@ -30,7 +28,6 @@ where
block: Vec::with_capacity(BLOCK_LEN * 2), block: Vec::with_capacity(BLOCK_LEN * 2),
write: CountingWriter::wrap(BufWriter::new(wrt)), write: CountingWriter::wrap(BufWriter::new(wrt)),
value_writer: TValueWriter::default(), value_writer: TValueWriter::default(),
stateless_buffer: Vec::new(),
} }
} }
} }
@@ -45,16 +42,15 @@ where
return Ok(None); return Ok(None);
} }
let start_offset = self.write.written_bytes() as usize; let start_offset = self.write.written_bytes() as usize;
let buffer: &mut Vec<u8> = &mut self.stateless_buffer; // TODO avoid buffer allocation
self.value_writer.serialize_block(buffer); let mut buffer = Vec::new();
self.value_writer.clear(); self.value_writer.serialize_block(&mut buffer);
let block_len = buffer.len() + self.block.len(); let block_len = buffer.len() + self.block.len();
self.write.write_all(&(block_len as u32).to_le_bytes())?; self.write.write_all(&(block_len as u32).to_le_bytes())?;
self.write.write_all(&buffer[..])?; self.write.write_all(&buffer[..])?;
self.write.write_all(&self.block[..])?; self.write.write_all(&self.block[..])?;
let end_offset = self.write.written_bytes() as usize; let end_offset = self.write.written_bytes() as usize;
self.block.clear(); self.block.clear();
buffer.clear();
Ok(Some(start_offset..end_offset)) Ok(Some(start_offset..end_offset))
} }
@@ -95,7 +91,8 @@ where
pub struct DeltaReader<'a, TValueReader> { pub struct DeltaReader<'a, TValueReader> {
common_prefix_len: usize, common_prefix_len: usize,
suffix_range: Range<usize>, suffix_start: usize,
suffix_end: usize,
value_reader: TValueReader, value_reader: TValueReader,
block_reader: BlockReader<'a>, block_reader: BlockReader<'a>,
idx: usize, idx: usize,
@@ -108,7 +105,8 @@ where TValueReader: value::ValueReader
DeltaReader { DeltaReader {
idx: 0, idx: 0,
common_prefix_len: 0, common_prefix_len: 0,
suffix_range: 0..0, suffix_start: 0,
suffix_end: 0,
value_reader: TValueReader::default(), value_reader: TValueReader::default(),
block_reader: BlockReader::new(Box::new(reader)), block_reader: BlockReader::new(Box::new(reader)),
} }
@@ -146,14 +144,15 @@ where TValueReader: value::ValueReader
} }
fn read_delta_key(&mut self) -> bool { fn read_delta_key(&mut self) -> bool {
let Some((keep, add)) = self.read_keep_add() else { if let Some((keep, add)) = self.read_keep_add() {
return false; self.common_prefix_len = keep;
}; self.suffix_start = self.block_reader.offset();
self.common_prefix_len = keep; self.suffix_end = self.suffix_start + add;
let suffix_start = self.block_reader.offset(); self.block_reader.advance(add);
self.suffix_range = suffix_start..(suffix_start + add); true
self.block_reader.advance(add); } else {
true false
}
} }
pub fn advance(&mut self) -> io::Result<bool> { pub fn advance(&mut self) -> io::Result<bool> {
@@ -173,17 +172,15 @@ where TValueReader: value::ValueReader
Ok(true) Ok(true)
} }
#[inline(always)]
pub fn common_prefix_len(&self) -> usize { pub fn common_prefix_len(&self) -> usize {
self.common_prefix_len self.common_prefix_len
} }
#[inline(always)]
pub fn suffix(&self) -> &[u8] { pub fn suffix(&self) -> &[u8] {
self.block_reader.buffer_from_to(self.suffix_range.clone()) self.block_reader
.buffer_from_to(self.suffix_start, self.suffix_end)
} }
#[inline(always)]
pub fn value(&self) -> &TValueReader::Value { pub fn value(&self) -> &TValueReader::Value {
self.value_reader.value(self.idx) self.value_reader.value(self.idx)
} }
@@ -192,11 +189,11 @@ where TValueReader: value::ValueReader
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::DeltaReader; use super::DeltaReader;
use crate::value::U64MonotonicValueReader; use crate::value::U64MonotonicReader;
#[test] #[test]
fn test_empty() { fn test_empty() {
let mut delta_reader: DeltaReader<U64MonotonicValueReader> = DeltaReader::empty(); let mut delta_reader: DeltaReader<U64MonotonicReader> = DeltaReader::empty();
assert!(!delta_reader.advance().unwrap()); assert!(!delta_reader.advance().unwrap());
} }
} }

View File

@@ -11,25 +11,12 @@ use tantivy_fst::Automaton;
use crate::streamer::{Streamer, StreamerBuilder}; use crate::streamer::{Streamer, StreamerBuilder};
use crate::{BlockAddr, DeltaReader, Reader, SSTable, SSTableIndex, TermOrdinal}; use crate::{BlockAddr, DeltaReader, Reader, SSTable, SSTableIndex, TermOrdinal};
/// An SSTable is a sorted map that associates sorted `&[u8]` keys /// The term dictionary contains all of the terms in
/// to any kind of typed values. /// `tantivy index` in a sorted manner.
/// ///
/// The SSTable is organized in blocks. /// The `Fst` crate is used to associate terms to their
/// In each block, keys and values are encoded separately. /// respective `TermOrdinal`. The `TermInfoStore` then makes it
/// /// possible to fetch the associated `TermInfo`.
/// The keys are encoded using incremental encoding.
/// The values on the other hand, are encoded according to a value-specific
/// codec defined in the TSSTable generic argument.
///
/// Finally, an index is joined to the Dictionary to make it possible,
/// given a key to identify which block contains this key.
///
/// The codec was designed in such a way that the sstable
/// reader is not aware of block, and yet can read any sequence of blocks,
/// as long as the slice of bytes it is given starts and stops at
/// block boundary.
///
/// (See also README.md)
pub struct Dictionary<TSSTable: SSTable> { pub struct Dictionary<TSSTable: SSTable> {
pub sstable_slice: FileSlice, pub sstable_slice: FileSlice,
pub sstable_index: SSTableIndex, pub sstable_index: SSTableIndex,
@@ -75,23 +62,6 @@ impl<TSSTable: SSTable> Dictionary<TSSTable> {
Ok(TSSTable::delta_reader(data)) Ok(TSSTable::delta_reader(data))
} }
/// This function returns a file slice covering a set of sstable blocks
/// that include the key range passed in arguments.
///
/// It works by identifying
/// - `first_block`: the block containing the start boudary key
/// - `last_block`: the block containing the end boundary key.
///
/// And then returning the range that spans over all blocks between.
/// and including first_block and last_block, aka:
/// `[first_block.start_offset .. last_block.end_offset)`
///
/// Technically this function does not provide the tightest fit, as
/// for simplification, it treats the start bound of the `key_range`
/// as if it was inclusive, even if it is exclusive.
/// On the rare edge case where a user asks for `(start_key, end_key]`
/// and `start_key` happens to be the last key of a block, we return a
/// slice that is the first block was not necessary.
fn file_slice_for_range(&self, key_range: impl RangeBounds<[u8]>) -> FileSlice { fn file_slice_for_range(&self, key_range: impl RangeBounds<[u8]>) -> FileSlice {
let start_bound: Bound<usize> = match key_range.start_bound() { let start_bound: Bound<usize> = match key_range.start_bound() {
Bound::Included(key) | Bound::Excluded(key) => { Bound::Included(key) | Bound::Excluded(key) => {

View File

@@ -20,15 +20,13 @@ mod block_reader;
pub use self::block_reader::BlockReader; pub use self::block_reader::BlockReader;
pub use self::delta::{DeltaReader, DeltaWriter}; pub use self::delta::{DeltaReader, DeltaWriter};
pub use self::merge::VoidMerge; pub use self::merge::VoidMerge;
use self::value::{U64MonotonicValueReader, U64MonotonicValueWriter, ValueReader, ValueWriter}; use self::value::{U64MonotonicReader, U64MonotonicWriter, ValueReader, ValueWriter};
use crate::value::{RangeValueReader, RangeValueWriter}; use crate::value::{RangeReader, RangeWriter};
pub type TermOrdinal = u64; pub type TermOrdinal = u64;
const DEFAULT_KEY_CAPACITY: usize = 50; const DEFAULT_KEY_CAPACITY: usize = 50;
/// Given two byte string returns the length of
/// the longest common prefix.
fn common_prefix_len(left: &[u8], right: &[u8]) -> usize { fn common_prefix_len(left: &[u8], right: &[u8]) -> usize {
left.iter() left.iter()
.cloned() .cloned()
@@ -40,8 +38,6 @@ fn common_prefix_len(left: &[u8], right: &[u8]) -> usize {
#[derive(Debug, Copy, Clone)] #[derive(Debug, Copy, Clone)]
pub struct SSTableDataCorruption; pub struct SSTableDataCorruption;
/// SSTable makes it possible to read and write
/// sstables with typed values.
pub trait SSTable: Sized { pub trait SSTable: Sized {
type Value: Clone; type Value: Clone;
type ValueReader: ValueReader<Value = Self::Value>; type ValueReader: ValueReader<Value = Self::Value>;
@@ -51,8 +47,14 @@ pub trait SSTable: Sized {
DeltaWriter::new(write) DeltaWriter::new(write)
} }
fn writer<W: io::Write>(wrt: W) -> Writer<W, Self::ValueWriter> { fn writer<W: io::Write>(write: W) -> Writer<W, Self::ValueWriter> {
Writer::new(wrt) Writer {
previous_key: Vec::with_capacity(DEFAULT_KEY_CAPACITY),
num_terms: 0u64,
index_builder: SSTableIndexBuilder::default(),
delta_writer: Self::delta_writer(write),
first_ordinal_of_the_block: 0u64,
}
} }
fn delta_reader<'a, R: io::Read + 'a>(reader: R) -> DeltaReader<'a, Self::ValueReader> { fn delta_reader<'a, R: io::Read + 'a>(reader: R) -> DeltaReader<'a, Self::ValueReader> {
@@ -87,47 +89,32 @@ pub struct VoidSSTable;
impl SSTable for VoidSSTable { impl SSTable for VoidSSTable {
type Value = (); type Value = ();
type ValueReader = value::VoidValueReader; type ValueReader = value::VoidReader;
type ValueWriter = value::VoidValueWriter; type ValueWriter = value::VoidWriter;
} }
/// SSTable associated keys to u64
/// sorted in order.
///
/// In other words, two keys `k1` and `k2`
/// such that `k1` <= `k2`, are required to observe
/// `range_sstable[k1] <= range_sstable[k2]`.
#[allow(dead_code)] #[allow(dead_code)]
pub struct MonotonicU64SSTable; pub struct SSTableMonotonicU64;
impl SSTable for MonotonicU64SSTable { impl SSTable for SSTableMonotonicU64 {
type Value = u64; type Value = u64;
type ValueReader = U64MonotonicValueReader; type ValueReader = U64MonotonicReader;
type ValueWriter = U64MonotonicValueWriter; type ValueWriter = U64MonotonicWriter;
} }
/// SSTable associating keys to ranges. /// Retpresent
/// The range are required to partition the pub struct SSTableRange;
/// space.
///
/// In other words, two consecutive keys `k1` and `k2`
/// are required to observe
/// `range_sstable[k1].end == range_sstable[k2].start`.
///
/// The first range is not required to start at `0`.
pub struct RangeSSTable;
impl SSTable for RangeSSTable { impl SSTable for SSTableRange {
type Value = Range<u64>; type Value = Range<u64>;
type ValueReader = RangeValueReader; type ValueReader = RangeReader;
type ValueWriter = RangeValueWriter; type ValueWriter = RangeWriter;
} }
/// SSTable reader.
pub struct Reader<'a, TValueReader> { pub struct Reader<'a, TValueReader> {
key: Vec<u8>, key: Vec<u8>,
delta_reader: DeltaReader<'a, TValueReader>, delta_reader: DeltaReader<'a, TValueReader>,
@@ -148,19 +135,16 @@ where TValueReader: ValueReader
Ok(true) Ok(true)
} }
#[inline(always)]
pub fn key(&self) -> &[u8] { pub fn key(&self) -> &[u8] {
&self.key &self.key
} }
#[inline(always)]
pub fn value(&self) -> &TValueReader::Value { pub fn value(&self) -> &TValueReader::Value {
self.delta_reader.value() self.delta_reader.value()
} }
} }
impl<'a, TValueReader> AsRef<[u8]> for Reader<'a, TValueReader> { impl<'a, TValueReader> AsRef<[u8]> for Reader<'a, TValueReader> {
#[inline(always)]
fn as_ref(&self) -> &[u8] { fn as_ref(&self) -> &[u8] {
&self.key &self.key
} }
@@ -181,54 +165,21 @@ where
W: io::Write, W: io::Write,
TValueWriter: value::ValueWriter, TValueWriter: value::ValueWriter,
{ {
/// Use `Self::new`. This method only exists to match its
/// equivalent in fst.
/// TODO remove this function. (See Issue #1727)
#[doc(hidden)]
pub fn create(wrt: W) -> io::Result<Self> { pub fn create(wrt: W) -> io::Result<Self> {
Ok(Self::new(wrt)) Ok(Writer {
}
/// Creates a new `TermDictionaryBuilder`.
pub fn new(wrt: W) -> Self {
Writer {
previous_key: Vec::with_capacity(DEFAULT_KEY_CAPACITY), previous_key: Vec::with_capacity(DEFAULT_KEY_CAPACITY),
num_terms: 0u64, num_terms: 0u64,
index_builder: SSTableIndexBuilder::default(), index_builder: SSTableIndexBuilder::default(),
delta_writer: DeltaWriter::new(wrt), delta_writer: DeltaWriter::new(wrt),
first_ordinal_of_the_block: 0u64, first_ordinal_of_the_block: 0u64,
} })
} }
/// Returns the last inserted key.
/// If no key has been inserted yet, or the block was just
/// flushed, this function returns "".
#[inline(always)] #[inline(always)]
pub(crate) fn last_inserted_key(&self) -> &[u8] { pub(crate) fn current_key(&self) -> &[u8] {
&self.previous_key[..] &self.previous_key[..]
} }
/// Inserts a `(key, value)` pair in the term dictionary.
/// Keys have to be inserted in order.
///
/// # Panics
///
/// Will panics if keys are inserted in an invalid order.
#[inline]
pub fn insert<K: AsRef<[u8]>>(
&mut self,
key: K,
value: &TValueWriter::Value,
) -> io::Result<()> {
self.insert_key(key.as_ref())?;
self.insert_value(value)?;
Ok(())
}
/// # Warning
///
/// Horribly dangerous internal API. See `.insert(...)`.
#[doc(hidden)]
#[inline] #[inline]
pub fn insert_key(&mut self, key: &[u8]) -> io::Result<()> { pub fn insert_key(&mut self, key: &[u8]) -> io::Result<()> {
// If this is the first key in the block, we use it to // If this is the first key in the block, we use it to
@@ -253,10 +204,17 @@ where
Ok(()) Ok(())
} }
/// # Warning #[inline]
/// pub fn insert<K: AsRef<[u8]>>(
/// Horribly dangerous internal API. See `.insert(...)`. &mut self,
#[doc(hidden)] key: K,
value: &TValueWriter::Value,
) -> io::Result<()> {
self.insert_key(key.as_ref())?;
self.insert_value(value)?;
Ok(())
}
#[inline] #[inline]
pub fn insert_value(&mut self, value: &TValueWriter::Value) -> io::Result<()> { pub fn insert_value(&mut self, value: &TValueWriter::Value) -> io::Result<()> {
self.delta_writer.write_value(value); self.delta_writer.write_value(value);
@@ -298,23 +256,11 @@ where
Ok(wrt.into_inner()?) Ok(wrt.into_inner()?)
} }
} }
impl<TValueWriter> Writer<Vec<u8>, TValueWriter>
where TValueWriter: value::ValueWriter
{
#[inline]
pub fn insert_cannot_fail<K: AsRef<[u8]>>(&mut self, key: K, value: &TValueWriter::Value) {
self.insert(key, value)
.expect("SSTable over a Vec should never fail");
}
}
#[cfg(test)] #[cfg(test)]
mod test { mod test {
use std::io; use std::io;
use std::ops::Bound;
use super::{common_prefix_len, MonotonicU64SSTable, SSTable, VoidMerge, VoidSSTable}; use super::{common_prefix_len, SSTable, SSTableMonotonicU64, VoidMerge, VoidSSTable};
fn aux_test_common_prefix_len(left: &str, right: &str, expect_len: usize) { fn aux_test_common_prefix_len(left: &str, right: &str, expect_len: usize) {
assert_eq!( assert_eq!(
@@ -423,11 +369,8 @@ mod test {
let mut buffer = Vec::new(); let mut buffer = Vec::new();
{ {
let mut writer = VoidSSTable::writer(&mut buffer); let mut writer = VoidSSTable::writer(&mut buffer);
assert_eq!(writer.last_inserted_key(), b"");
writer.insert(b"abcd", &()).unwrap(); writer.insert(b"abcd", &()).unwrap();
assert_eq!(writer.last_inserted_key(), b"abcd");
writer.insert(b"abe", &()).unwrap(); writer.insert(b"abe", &()).unwrap();
assert_eq!(writer.last_inserted_key(), b"abe");
writer.finish().unwrap(); writer.finish().unwrap();
} }
let mut output = Vec::new(); let mut output = Vec::new();
@@ -438,12 +381,12 @@ mod test {
#[test] #[test]
fn test_sstable_u64() -> io::Result<()> { fn test_sstable_u64() -> io::Result<()> {
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let mut writer = MonotonicU64SSTable::writer(&mut buffer); let mut writer = SSTableMonotonicU64::writer(&mut buffer);
writer.insert(b"abcd", &1u64)?; writer.insert(b"abcd", &1u64)?;
writer.insert(b"abe", &4u64)?; writer.insert(b"abe", &4u64)?;
writer.insert(b"gogo", &4324234234234234u64)?; writer.insert(b"gogo", &4324234234234234u64)?;
writer.finish()?; writer.finish()?;
let mut reader = MonotonicU64SSTable::reader(&buffer[..]); let mut reader = SSTableMonotonicU64::reader(&buffer[..]);
assert!(reader.advance()?); assert!(reader.advance()?);
assert_eq!(reader.key(), b"abcd"); assert_eq!(reader.key(), b"abcd");
assert_eq!(reader.value(), &1u64); assert_eq!(reader.value(), &1u64);
@@ -459,72 +402,7 @@ mod test {
#[test] #[test]
fn test_sstable_empty() { fn test_sstable_empty() {
let mut sstable_range_empty = crate::RangeSSTable::create_empty_reader(); let mut sstable_range_empty = crate::SSTableRange::create_empty_reader();
assert!(!sstable_range_empty.advance().unwrap()); assert!(!sstable_range_empty.advance().unwrap());
} }
use common::file_slice::FileSlice;
use proptest::prelude::*;
use crate::Dictionary;
fn bound_strategy() -> impl Strategy<Value = Bound<String>> {
prop_oneof![
Just(Bound::<String>::Unbounded),
"[a-d]*".prop_map(|key| Bound::Included(key)),
"[a-d]*".prop_map(|key| Bound::Excluded(key)),
]
}
fn extract_key(bound: Bound<&String>) -> Option<&str> {
match bound.as_ref() {
Bound::Included(key) => Some(key.as_str()),
Bound::Excluded(key) => Some(key.as_str()),
Bound::Unbounded => None,
}
}
fn bounds_strategy() -> impl Strategy<Value = (Bound<String>, Bound<String>)> {
(bound_strategy(), bound_strategy()).prop_filter(
"Lower bound <= Upper bound",
|(left, right)| match (extract_key(left.as_ref()), extract_key(right.as_ref())) {
(None, _) => true,
(_, None) => true,
(left, right) => left <= right,
},
)
}
proptest! {
#[test]
fn test_prop_test_ranges(words in prop::collection::btree_set("[a-d]*", 1..100),
(lower_bound, upper_bound) in bounds_strategy(),
) {
// TODO tweak block size.
let mut builder = Dictionary::<VoidSSTable>::builder(Vec::new()).unwrap();
for word in &words {
builder.insert(word.as_bytes(), &()).unwrap();
}
let buffer: Vec<u8> = builder.finish().unwrap();
let dictionary: Dictionary<VoidSSTable> = Dictionary::open(FileSlice::from(buffer)).unwrap();
let mut range_builder = dictionary.range();
range_builder = match lower_bound.as_ref() {
Bound::Included(key) => range_builder.ge(key.as_bytes()),
Bound::Excluded(key) => range_builder.gt(key.as_bytes()),
Bound::Unbounded => range_builder,
};
range_builder = match upper_bound.as_ref() {
Bound::Included(key) => range_builder.le(key.as_bytes()),
Bound::Excluded(key) => range_builder.lt(key.as_bytes()),
Bound::Unbounded => range_builder,
};
let mut stream = range_builder.into_stream().unwrap();
let mut btree_set_range = words.range((lower_bound, upper_bound));
while stream.advance() {
let val = btree_set_range.next().unwrap();
assert_eq!(val.as_bytes(), stream.key());
}
assert!(btree_set_range.next().is_none());
}
}
} }

View File

@@ -53,7 +53,7 @@ pub fn merge_sstable<SST: SSTable, W: io::Write, M: ValueMerger<SST::Value>>(
} }
for _ in 0..len - 1 { for _ in 0..len - 1 {
if let Some(mut head) = heap.peek_mut() { if let Some(mut head) = heap.peek_mut() {
if head.0.key() == writer.last_inserted_key() { if head.0.key() == writer.current_key() {
value_merger.add(head.0.value()); value_merger.add(head.0.value());
if !head.0.advance()? { if !head.0.advance()? {
PeekMut::pop(head); PeekMut::pop(head);

View File

@@ -71,7 +71,7 @@ mod tests {
use std::collections::{BTreeMap, BTreeSet}; use std::collections::{BTreeMap, BTreeSet};
use std::str; use std::str;
use super::super::{MonotonicU64SSTable, SSTable, VoidSSTable}; use super::super::{SSTable, SSTableMonotonicU64, VoidSSTable};
use super::{U64Merge, VoidMerge}; use super::{U64Merge, VoidMerge};
fn write_sstable(keys: &[&'static str]) -> Vec<u8> { fn write_sstable(keys: &[&'static str]) -> Vec<u8> {
@@ -89,7 +89,7 @@ mod tests {
fn write_sstable_u64(keys: &[(&'static str, u64)]) -> Vec<u8> { fn write_sstable_u64(keys: &[(&'static str, u64)]) -> Vec<u8> {
let mut buffer: Vec<u8> = vec![]; let mut buffer: Vec<u8> = vec![];
{ {
let mut sstable_writer = MonotonicU64SSTable::writer(&mut buffer); let mut sstable_writer = SSTableMonotonicU64::writer(&mut buffer);
for (key, val) in keys { for (key, val) in keys {
assert!(sstable_writer.insert(key.as_bytes(), val).is_ok()); assert!(sstable_writer.insert(key.as_bytes(), val).is_ok());
} }
@@ -132,8 +132,8 @@ mod tests {
} }
} }
let mut w = Vec::new(); let mut w = Vec::new();
assert!(MonotonicU64SSTable::merge(sstables_ref, &mut w, U64Merge).is_ok()); assert!(SSTableMonotonicU64::merge(sstables_ref, &mut w, U64Merge).is_ok());
let mut reader = MonotonicU64SSTable::reader(&w[..]); let mut reader = SSTableMonotonicU64::reader(&w[..]);
for (k, v) in merged { for (k, v) in merged {
assert!(reader.advance().unwrap()); assert!(reader.advance().unwrap());
assert_eq!(reader.key(), k.as_bytes()); assert_eq!(reader.key(), k.as_bytes());

View File

@@ -16,9 +16,16 @@ impl SSTableIndex {
} }
pub fn search_block(&self, key: &[u8]) -> Option<BlockAddr> { pub fn search_block(&self, key: &[u8]) -> Option<BlockAddr> {
self.search_block_from(key).next()
}
pub fn search_block_from<'key, 'slf: 'key>(
&'slf self,
key: &'key [u8],
) -> impl Iterator<Item = BlockAddr> + Clone + 'key {
self.blocks self.blocks
.iter() .iter()
.find(|block| &block.last_key_or_greater[..] >= key) .skip_while(|block| &block.last_key_or_greater[..] < key)
.map(|block| block.block_addr.clone()) .map(|block| block.block_addr.clone())
} }
} }

View File

@@ -211,10 +211,10 @@ mod tests {
use common::OwnedBytes; use common::OwnedBytes;
use crate::{Dictionary, MonotonicU64SSTable}; use crate::{Dictionary, SSTableMonotonicU64};
fn create_test_dictionary() -> io::Result<Dictionary<MonotonicU64SSTable>> { fn create_test_dictionary() -> io::Result<Dictionary<SSTableMonotonicU64>> {
let mut dict_builder = Dictionary::<MonotonicU64SSTable>::builder(Vec::new())?; let mut dict_builder = Dictionary::<SSTableMonotonicU64>::builder(Vec::new())?;
dict_builder.insert(b"abaisance", &0)?; dict_builder.insert(b"abaisance", &0)?;
dict_builder.insert(b"abalation", &1)?; dict_builder.insert(b"abalation", &1)?;
dict_builder.insert(b"abalienate", &2)?; dict_builder.insert(b"abalienate", &2)?;

View File

@@ -20,30 +20,24 @@ pub trait ValueReader: Default {
fn load(&mut self, data: &[u8]) -> io::Result<usize>; fn load(&mut self, data: &[u8]) -> io::Result<usize>;
} }
/// `ValueWriter` is a trait to make it possible to write blocks
/// of value.
pub trait ValueWriter: Default { pub trait ValueWriter: Default {
/// Type of the value being written. /// Type of the value being written.
type Value; type Value;
/// Records a new value. /// Records a new value.
/// This method usually just accumulates data in a `Vec`, /// This method usually just accumulates data in a `Vec`,
/// only to be serialized on the call to `ValueWriter::serialize_block`. /// only to be serialized on the call to `ValueWriter::write_block`.
fn write(&mut self, val: &Self::Value); fn write(&mut self, val: &Self::Value);
/// Serializes the accumulated values into the output buffer. /// Serializes the accumulated values into the output buffer.
fn serialize_block(&self, output: &mut Vec<u8>); fn serialize_block(&mut self, output: &mut Vec<u8>);
/// Clears the `ValueWriter`. After a call to clear, the `ValueWriter`
/// should behave like a fresh `ValueWriter::default()`.
fn clear(&mut self);
} }
pub use range::{RangeValueReader, RangeValueWriter}; pub use range::{RangeReader, RangeWriter};
pub use u64_monotonic::{U64MonotonicValueReader, U64MonotonicValueWriter}; pub use u64_monotonic::{U64MonotonicReader, U64MonotonicWriter};
pub use void::{VoidValueReader, VoidValueWriter}; pub use void::{VoidReader, VoidWriter};
fn deserialize_vint_u64(data: &mut &[u8]) -> u64 { fn deserialize_u64(data: &mut &[u8]) -> u64 {
let (num_bytes, val) = super::vint::deserialize_read(data); let (num_bytes, val) = super::vint::deserialize_read(data);
*data = &data[num_bytes..]; *data = &data[num_bytes..];
val val
@@ -69,7 +63,6 @@ pub(crate) mod tests {
writer.write(value); writer.write(value);
} }
writer.serialize_block(&mut buffer); writer.serialize_block(&mut buffer);
writer.clear();
} }
let data_len = buffer.len(); let data_len = buffer.len();
buffer.extend_from_slice(&b"extradata"[..]); buffer.extend_from_slice(&b"extradata"[..]);

View File

@@ -1,18 +1,16 @@
use std::io; use std::io;
use std::ops::Range; use std::ops::Range;
use crate::value::{deserialize_vint_u64, ValueReader, ValueWriter}; use crate::value::{deserialize_u64, ValueReader, ValueWriter};
/// See module comment.
#[derive(Default)] #[derive(Default)]
pub struct RangeValueReader { pub struct RangeReader {
vals: Vec<Range<u64>>, vals: Vec<Range<u64>>,
} }
impl ValueReader for RangeValueReader { impl ValueReader for RangeReader {
type Value = Range<u64>; type Value = Range<u64>;
#[inline(always)]
fn value(&self, idx: usize) -> &Range<u64> { fn value(&self, idx: usize) -> &Range<u64> {
&self.vals[idx] &self.vals[idx]
} }
@@ -20,11 +18,11 @@ impl ValueReader for RangeValueReader {
fn load(&mut self, mut data: &[u8]) -> io::Result<usize> { fn load(&mut self, mut data: &[u8]) -> io::Result<usize> {
self.vals.clear(); self.vals.clear();
let original_num_bytes = data.len(); let original_num_bytes = data.len();
let len = deserialize_vint_u64(&mut data) as usize; let len = deserialize_u64(&mut data) as usize;
if len != 0 { if len != 0 {
let mut prev_val = deserialize_vint_u64(&mut data); let mut prev_val = deserialize_u64(&mut data);
for _ in 1..len { for _ in 1..len {
let next_val = prev_val + deserialize_vint_u64(&mut data); let next_val = prev_val + deserialize_u64(&mut data);
self.vals.push(prev_val..next_val); self.vals.push(prev_val..next_val);
prev_val = next_val; prev_val = next_val;
} }
@@ -33,23 +31,12 @@ impl ValueReader for RangeValueReader {
} }
} }
/// Range writer. The range are required to partition the
/// space.
///
/// In other words, two consecutive keys `k1` and `k2`
/// are required to observe
/// `range_sstable[k1].end == range_sstable[k2].start`.
///
/// The writer will panic if the inserted value do not follow
/// this property.
///
/// The first range is not required to start at `0`.
#[derive(Default)] #[derive(Default)]
pub struct RangeValueWriter { pub struct RangeWriter {
vals: Vec<u64>, vals: Vec<u64>,
} }
impl ValueWriter for RangeValueWriter { impl ValueWriter for RangeWriter {
type Value = Range<u64>; type Value = Range<u64>;
fn write(&mut self, val: &Range<u64>) { fn write(&mut self, val: &Range<u64>) {
@@ -62,7 +49,7 @@ impl ValueWriter for RangeValueWriter {
} }
} }
fn serialize_block(&self, writer: &mut Vec<u8>) { fn serialize_block(&mut self, writer: &mut Vec<u8>) {
let mut prev_val = 0u64; let mut prev_val = 0u64;
crate::vint::serialize_into_vec(self.vals.len() as u64, writer); crate::vint::serialize_into_vec(self.vals.len() as u64, writer);
for &val in &self.vals { for &val in &self.vals {
@@ -70,9 +57,6 @@ impl ValueWriter for RangeValueWriter {
crate::vint::serialize_into_vec(delta, writer); crate::vint::serialize_into_vec(delta, writer);
prev_val = val; prev_val = val;
} }
}
fn clear(&mut self) {
self.vals.clear(); self.vals.clear();
} }
} }
@@ -83,19 +67,17 @@ mod tests {
#[test] #[test]
fn test_range_reader_writer() { fn test_range_reader_writer() {
crate::value::tests::test_value_reader_writer::<_, RangeValueReader, RangeValueWriter>(&[]); crate::value::tests::test_value_reader_writer::<_, RangeReader, RangeWriter>(&[]);
crate::value::tests::test_value_reader_writer::<_, RangeValueReader, RangeValueWriter>(&[ crate::value::tests::test_value_reader_writer::<_, RangeReader, RangeWriter>(&[0..3]);
0..3, crate::value::tests::test_value_reader_writer::<_, RangeReader, RangeWriter>(&[
]);
crate::value::tests::test_value_reader_writer::<_, RangeValueReader, RangeValueWriter>(&[
0..3, 0..3,
3..10, 3..10,
]); ]);
crate::value::tests::test_value_reader_writer::<_, RangeValueReader, RangeValueWriter>(&[ crate::value::tests::test_value_reader_writer::<_, RangeReader, RangeWriter>(&[
0..0, 0..0,
0..10, 0..10,
]); ]);
crate::value::tests::test_value_reader_writer::<_, RangeValueReader, RangeValueWriter>(&[ crate::value::tests::test_value_reader_writer::<_, RangeReader, RangeWriter>(&[
100..110, 100..110,
110..121, 110..121,
121..1250, 121..1250,
@@ -105,7 +87,7 @@ mod tests {
#[test] #[test]
#[should_panic] #[should_panic]
fn test_range_reader_writer_panics() { fn test_range_reader_writer_panics() {
crate::value::tests::test_value_reader_writer::<_, RangeValueReader, RangeValueWriter>(&[ crate::value::tests::test_value_reader_writer::<_, RangeReader, RangeWriter>(&[
1..3, 1..3,
4..10, 4..10,
]); ]);

View File

@@ -1,28 +1,27 @@
use std::io; use std::io;
use crate::value::{deserialize_vint_u64, ValueReader, ValueWriter}; use crate::value::{deserialize_u64, ValueReader, ValueWriter};
use crate::vint; use crate::vint;
#[derive(Default)] #[derive(Default)]
pub struct U64MonotonicValueReader { pub struct U64MonotonicReader {
vals: Vec<u64>, vals: Vec<u64>,
} }
impl ValueReader for U64MonotonicValueReader { impl ValueReader for U64MonotonicReader {
type Value = u64; type Value = u64;
#[inline(always)]
fn value(&self, idx: usize) -> &Self::Value { fn value(&self, idx: usize) -> &Self::Value {
&self.vals[idx] &self.vals[idx]
} }
fn load(&mut self, mut data: &[u8]) -> io::Result<usize> { fn load(&mut self, mut data: &[u8]) -> io::Result<usize> {
let original_num_bytes = data.len(); let original_num_bytes = data.len();
let num_vals = deserialize_vint_u64(&mut data) as usize; let num_vals = deserialize_u64(&mut data) as usize;
self.vals.clear(); self.vals.clear();
let mut prev_val = 0u64; let mut prev_val = 0u64;
for _ in 0..num_vals { for _ in 0..num_vals {
let delta = deserialize_vint_u64(&mut data); let delta = deserialize_u64(&mut data);
let val = prev_val + delta; let val = prev_val + delta;
self.vals.push(val); self.vals.push(val);
prev_val = val; prev_val = val;
@@ -32,18 +31,18 @@ impl ValueReader for U64MonotonicValueReader {
} }
#[derive(Default)] #[derive(Default)]
pub struct U64MonotonicValueWriter { pub struct U64MonotonicWriter {
vals: Vec<u64>, vals: Vec<u64>,
} }
impl ValueWriter for U64MonotonicValueWriter { impl ValueWriter for U64MonotonicWriter {
type Value = u64; type Value = u64;
fn write(&mut self, val: &Self::Value) { fn write(&mut self, val: &Self::Value) {
self.vals.push(*val); self.vals.push(*val);
} }
fn serialize_block(&self, output: &mut Vec<u8>) { fn serialize_block(&mut self, output: &mut Vec<u8>) {
let mut prev_val = 0u64; let mut prev_val = 0u64;
vint::serialize_into_vec(self.vals.len() as u64, output); vint::serialize_into_vec(self.vals.len() as u64, output);
for &val in &self.vals { for &val in &self.vals {
@@ -51,9 +50,6 @@ impl ValueWriter for U64MonotonicValueWriter {
vint::serialize_into_vec(delta, output); vint::serialize_into_vec(delta, output);
prev_val = val; prev_val = val;
} }
}
fn clear(&mut self) {
self.vals.clear(); self.vals.clear();
} }
} }
@@ -64,20 +60,14 @@ mod tests {
#[test] #[test]
fn test_u64_monotonic_reader_writer() { fn test_u64_monotonic_reader_writer() {
crate::value::tests::test_value_reader_writer::< crate::value::tests::test_value_reader_writer::<_, U64MonotonicReader, U64MonotonicWriter>(
_, &[],
U64MonotonicValueReader, );
U64MonotonicValueWriter, crate::value::tests::test_value_reader_writer::<_, U64MonotonicReader, U64MonotonicWriter>(
>(&[]); &[5],
crate::value::tests::test_value_reader_writer::< );
_, crate::value::tests::test_value_reader_writer::<_, U64MonotonicReader, U64MonotonicWriter>(
U64MonotonicValueReader, &[1u64, 30u64],
U64MonotonicValueWriter, );
>(&[5]);
crate::value::tests::test_value_reader_writer::<
_,
U64MonotonicValueReader,
U64MonotonicValueWriter,
>(&[1u64, 30u64]);
} }
} }

View File

@@ -3,12 +3,11 @@ use std::io;
use crate::value::{ValueReader, ValueWriter}; use crate::value::{ValueReader, ValueWriter};
#[derive(Default)] #[derive(Default)]
pub struct VoidValueReader; pub struct VoidReader;
impl ValueReader for VoidValueReader { impl ValueReader for VoidReader {
type Value = (); type Value = ();
#[inline(always)]
fn value(&self, _idx: usize) -> &() { fn value(&self, _idx: usize) -> &() {
&() &()
} }
@@ -19,16 +18,14 @@ impl ValueReader for VoidValueReader {
} }
#[derive(Default)] #[derive(Default)]
pub struct VoidValueWriter; pub struct VoidWriter;
impl ValueWriter for VoidValueWriter { impl ValueWriter for VoidWriter {
type Value = (); type Value = ();
fn write(&mut self, _val: &()) {} fn write(&mut self, _val: &()) {}
fn serialize_block(&self, _output: &mut Vec<u8>) {} fn serialize_block(&mut self, _output: &mut Vec<u8>) {}
fn clear(&mut self) {}
} }
#[cfg(test)] #[cfg(test)]
@@ -37,12 +34,8 @@ mod tests {
#[test] #[test]
fn test_range_reader_writer() { fn test_range_reader_writer() {
crate::value::tests::test_value_reader_writer::<_, VoidValueReader, VoidValueWriter>(&[]); crate::value::tests::test_value_reader_writer::<_, VoidReader, VoidWriter>(&[]);
crate::value::tests::test_value_reader_writer::<_, VoidValueReader, VoidValueWriter>(&[()]); crate::value::tests::test_value_reader_writer::<_, VoidReader, VoidWriter>(&[()]);
crate::value::tests::test_value_reader_writer::<_, VoidValueReader, VoidValueWriter>(&[ crate::value::tests::test_value_reader_writer::<_, VoidReader, VoidWriter>(&[(), (), ()]);
(),
(),
(),
]);
} }
} }

View File

@@ -2,7 +2,6 @@
name = "tantivy-stacker" name = "tantivy-stacker"
version = "0.1.0" version = "0.1.0"
edition = "2021" edition = "2021"
license = "MIT"
[dependencies] [dependencies]
murmurhash32 = "0.2" murmurhash32 = "0.2"