Compare commits

..

2 Commits

Author SHA1 Message Date
Paul Masurel
1a72844048 Added simple columnar CLI program 2022-12-23 22:25:45 +09:00
Paul Masurel
d91df6cc7e Added support for dynamic fast field.
See README for more information.
2022-12-23 22:24:40 +09:00
173 changed files with 3280 additions and 11804 deletions

2
.gitignore vendored
View File

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

View File

@@ -1,6 +1,6 @@
[package] [package]
name = "tantivy" name = "tantivy"
version = "0.19.1-quickwit" version = "0.19.0"
authors = ["Paul Masurel <paul.masurel@gmail.com>"] authors = ["Paul Masurel <paul.masurel@gmail.com>"]
license = "MIT" license = "MIT"
categories = ["database-implementations", "data-structures"] categories = ["database-implementations", "data-structures"]
@@ -15,7 +15,7 @@ rust-version = "1.62"
[dependencies] [dependencies]
oneshot = "0.1.5" oneshot = "0.1.5"
base64 = "0.21.0" base64 = "0.20.0"
byteorder = "1.4.3" byteorder = "1.4.3"
crc32fast = "1.3.2" crc32fast = "1.3.2"
once_cell = "1.10.0" once_cell = "1.10.0"
@@ -23,7 +23,7 @@ regex = { version = "1.5.5", default-features = false, features = ["std", "unico
aho-corasick = "0.7" aho-corasick = "0.7"
tantivy-fst = "0.4.0" tantivy-fst = "0.4.0"
memmap2 = { version = "0.5.3", optional = true } memmap2 = { version = "0.5.3", optional = true }
lz4_flex = { version = "0.10", default-features = false, features = ["checked-decode"], optional = true } lz4_flex = { version = "0.9.2", default-features = false, features = ["checked-decode"], optional = true }
brotli = { version = "3.3.4", optional = true } brotli = { version = "3.3.4", optional = true }
zstd = { version = "0.12", optional = true, default-features = false } zstd = { version = "0.12", optional = true, default-features = false }
snap = { version = "1.0.5", optional = true } snap = { version = "1.0.5", optional = true }
@@ -48,21 +48,19 @@ murmurhash32 = "0.2.0"
time = { version = "0.3.10", features = ["serde-well-known"] } time = { version = "0.3.10", features = ["serde-well-known"] }
smallvec = "1.8.0" smallvec = "1.8.0"
rayon = "1.5.2" rayon = "1.5.2"
lru = "0.9.0" lru = "0.7.5"
fastdivide = "0.4.0" fastdivide = "0.4.0"
itertools = "0.10.3" itertools = "0.10.3"
measure_time = "0.8.2" measure_time = "0.8.2"
async-trait = "0.1.53" async-trait = "0.1.53"
arc-swap = "1.5.0" arc-swap = "1.5.0"
#columnar = { version="0.1", path="./columnar", package ="tantivy-columnar" }
sstable = { version="0.1", path="./sstable", package ="tantivy-sstable", optional = true } sstable = { version="0.1", path="./sstable", package ="tantivy-sstable", optional = true }
stacker = { version="0.1", path="./stacker", package ="tantivy-stacker" } stacker = { version="0.1", path="./stacker", package ="tantivy-stacker" }
tantivy-query-grammar = { version= "0.19.0", path="./query-grammar" } 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 }
tokenizer-api = { version="0.1", path="./tokenizer-api", package="tantivy-tokenizer-api" }
[target.'cfg(windows)'.dependencies] [target.'cfg(windows)'.dependencies]
winapi = "0.3.9" winapi = "0.3.9"
@@ -108,7 +106,7 @@ unstable = [] # useful for benches.
quickwit = ["sstable"] quickwit = ["sstable"]
[workspace] [workspace]
members = ["query-grammar", "bitpacker", "common", "fastfield_codecs", "ownedbytes", "stacker", "sstable", "tokenizer-api", "columnar"] members = ["query-grammar", "bitpacker", "common", "fastfield_codecs", "ownedbytes", "stacker", "sstable", "columnar"]
# Following the "fail" crate best practises, we isolate # Following the "fail" crate best practises, we isolate
# tests that define specific behavior in fail check points # tests that define specific behavior in fail check points

View File

@@ -29,7 +29,7 @@ Your mileage WILL vary depending on the nature of queries and their load.
# Features # Features
- Full-text search - Full-text search
- Configurable tokenizer (stemming available for 17 Latin languages) with third party support for Chinese ([tantivy-jieba](https://crates.io/crates/tantivy-jieba) and [cang-jie](https://crates.io/crates/cang-jie)), Japanese ([lindera](https://github.com/lindera-morphology/lindera-tantivy), [Vaporetto](https://crates.io/crates/vaporetto_tantivy), and [tantivy-tokenizer-tiny-segmenter](https://crates.io/crates/tantivy-tokenizer-tiny-segmenter)) and Korean ([lindera](https://github.com/lindera-morphology/lindera-tantivy) + [lindera-ko-dic-builder](https://github.com/lindera-morphology/lindera-ko-dic-builder)) - Configurable tokenizer (stemming available for 17 Latin languages with third party support for Chinese ([tantivy-jieba](https://crates.io/crates/tantivy-jieba) and [cang-jie](https://crates.io/crates/cang-jie)), Japanese ([lindera](https://github.com/lindera-morphology/lindera-tantivy), [Vaporetto](https://crates.io/crates/vaporetto_tantivy), and [tantivy-tokenizer-tiny-segmenter](https://crates.io/crates/tantivy-tokenizer-tiny-segmenter)) and Korean ([lindera](https://github.com/lindera-morphology/lindera-tantivy) + [lindera-ko-dic-builder](https://github.com/lindera-morphology/lindera-ko-dic-builder))
- Fast (check out the :racehorse: :sparkles: [benchmark](https://tantivy-search.github.io/bench/) :sparkles: :racehorse:) - Fast (check out the :racehorse: :sparkles: [benchmark](https://tantivy-search.github.io/bench/) :sparkles: :racehorse:)
- Tiny startup time (<10ms), perfect for command-line tools - Tiny startup time (<10ms), perfect for command-line tools
- BM25 scoring (the same as Lucene) - BM25 scoring (the same as Lucene)
@@ -41,13 +41,13 @@ Your mileage WILL vary depending on the nature of queries and their load.
- SIMD integer compression when the platform/CPU includes the SSE2 instruction set - SIMD integer compression when the platform/CPU includes the SSE2 instruction set
- Single valued and multivalued u64, i64, and f64 fast fields (equivalent of doc values in Lucene) - Single valued and multivalued u64, i64, and f64 fast fields (equivalent of doc values in Lucene)
- `&[u8]` fast fields - `&[u8]` fast fields
- Text, i64, u64, f64, dates, ip, bool, and hierarchical facet fields - Text, i64, u64, f64, dates, and hierarchical facet fields
- Compressed document store (LZ4, Zstd, None, Brotli, Snap) - LZ4 compressed document store
- Range queries - Range queries
- Faceted search - Faceted search
- Configurable indexing (optional term frequency and position indexing) - Configurable indexing (optional term frequency and position indexing)
- JSON Field - JSON Field
- Aggregation Collector: histogram, range buckets, average, and stats metrics - Aggregation Collector: range buckets, average, and stats metrics
- LogMergePolicy with deletes - LogMergePolicy with deletes
- Searcher Warmer API - Searcher Warmer API
- Cheesy logo with a horse - Cheesy logo with a horse
@@ -80,11 +80,10 @@ There are many ways to support this project.
# Contributing code # Contributing code
We use the GitHub Pull Request workflow: reference a GitHub ticket and/or include a comprehensive commit message when opening a PR. We use the GitHub Pull Request workflow: reference a GitHub ticket and/or include a comprehensive commit message when opening a PR.
Feel free to update CHANGELOG.md with your contribution.
## Tokenizer ## Minimum supported Rust version
When implementing a tokenizer for tantivy depend on the `tantivy-tokenizer-api` crate. Tantivy currently requires at least Rust 1.62 or later to compile.
## Clone and build locally ## Clone and build locally
@@ -92,9 +91,41 @@ Tantivy compiles on stable Rust.
To check out and run tests, you can simply run: To check out and run tests, you can simply run:
```bash ```bash
git clone https://github.com/quickwit-oss/tantivy.git git clone https://github.com/quickwit-oss/tantivy.git
cd tantivy cd tantivy
cargo test cargo build
```
## Run tests
Some tests will not run with just `cargo test` because of `fail-rs`.
To run the tests exhaustively, run `./run-tests.sh`.
## Debug
You might find it useful to step through the programme with a debugger.
### A failing test
Make sure you haven't run `cargo clean` after the most recent `cargo test` or `cargo build` to guarantee that the `target/` directory exists. Use this bash script to find the name of the most recent debug build of Tantivy and run it under `rust-gdb`:
```bash
find target/debug/ -maxdepth 1 -executable -type f -name "tantivy*" -printf '%TY-%Tm-%Td %TT %p\n' | sort -r | cut -d " " -f 3 | xargs -I RECENT_DBG_TANTIVY rust-gdb RECENT_DBG_TANTIVY
```
Now that you are in `rust-gdb`, you can set breakpoints on lines and methods that match your source code and run the debug executable with flags that you normally pass to `cargo test` like this:
```bash
$gdb run --test-threads 1 --test $NAME_OF_TEST
```
### An example
By default, `rustc` compiles everything in the `examples/` directory in debug mode. This makes it easy for you to make examples to reproduce bugs:
```bash
rust-gdb target/debug/examples/$EXAMPLE_NAME
$ gdb run
``` ```
# Companies Using Tantivy # Companies Using Tantivy

View File

@@ -34,7 +34,7 @@ pub fn hdfs_index_benchmark(c: &mut Criterion) {
let index = Index::create_in_ram(schema.clone()); let index = Index::create_in_ram(schema.clone());
let index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap(); let index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap();
for _ in 0..NUM_REPEATS { for _ in 0..NUM_REPEATS {
for doc_json in HDFS_LOGS.trim().split('\n') { for doc_json in HDFS_LOGS.trim().split("\n") {
let doc = schema.parse_document(doc_json).unwrap(); let doc = schema.parse_document(doc_json).unwrap();
index_writer.add_document(doc).unwrap(); index_writer.add_document(doc).unwrap();
} }
@@ -46,7 +46,7 @@ pub fn hdfs_index_benchmark(c: &mut Criterion) {
let index = Index::create_in_ram(schema.clone()); let index = Index::create_in_ram(schema.clone());
let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap(); let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap();
for _ in 0..NUM_REPEATS { for _ in 0..NUM_REPEATS {
for doc_json in HDFS_LOGS.trim().split('\n') { for doc_json in HDFS_LOGS.trim().split("\n") {
let doc = schema.parse_document(doc_json).unwrap(); let doc = schema.parse_document(doc_json).unwrap();
index_writer.add_document(doc).unwrap(); index_writer.add_document(doc).unwrap();
} }
@@ -59,7 +59,7 @@ pub fn hdfs_index_benchmark(c: &mut Criterion) {
let index = Index::create_in_ram(schema_with_store.clone()); let index = Index::create_in_ram(schema_with_store.clone());
let index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap(); let index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap();
for _ in 0..NUM_REPEATS { for _ in 0..NUM_REPEATS {
for doc_json in HDFS_LOGS.trim().split('\n') { for doc_json in HDFS_LOGS.trim().split("\n") {
let doc = schema.parse_document(doc_json).unwrap(); let doc = schema.parse_document(doc_json).unwrap();
index_writer.add_document(doc).unwrap(); index_writer.add_document(doc).unwrap();
} }
@@ -71,7 +71,7 @@ pub fn hdfs_index_benchmark(c: &mut Criterion) {
let index = Index::create_in_ram(schema_with_store.clone()); let index = Index::create_in_ram(schema_with_store.clone());
let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap(); let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap();
for _ in 0..NUM_REPEATS { for _ in 0..NUM_REPEATS {
for doc_json in HDFS_LOGS.trim().split('\n') { for doc_json in HDFS_LOGS.trim().split("\n") {
let doc = schema.parse_document(doc_json).unwrap(); let doc = schema.parse_document(doc_json).unwrap();
index_writer.add_document(doc).unwrap(); index_writer.add_document(doc).unwrap();
} }
@@ -85,7 +85,7 @@ pub fn hdfs_index_benchmark(c: &mut Criterion) {
let json_field = dynamic_schema.get_field("json").unwrap(); let json_field = dynamic_schema.get_field("json").unwrap();
let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap(); let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap();
for _ in 0..NUM_REPEATS { for _ in 0..NUM_REPEATS {
for doc_json in HDFS_LOGS.trim().split('\n') { for doc_json in HDFS_LOGS.trim().split("\n") {
let json_val: serde_json::Map<String, serde_json::Value> = let json_val: serde_json::Map<String, serde_json::Value> =
serde_json::from_str(doc_json).unwrap(); serde_json::from_str(doc_json).unwrap();
let doc = tantivy::doc!(json_field=>json_val); let doc = tantivy::doc!(json_field=>json_val);
@@ -101,7 +101,7 @@ pub fn hdfs_index_benchmark(c: &mut Criterion) {
let json_field = dynamic_schema.get_field("json").unwrap(); let json_field = dynamic_schema.get_field("json").unwrap();
let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap(); let mut index_writer = index.writer_with_num_threads(1, 100_000_000).unwrap();
for _ in 0..NUM_REPEATS { for _ in 0..NUM_REPEATS {
for doc_json in HDFS_LOGS.trim().split('\n') { for doc_json in HDFS_LOGS.trim().split("\n") {
let json_val: serde_json::Map<String, serde_json::Value> = let json_val: serde_json::Map<String, serde_json::Value> =
serde_json::from_str(doc_json).unwrap(); serde_json::from_str(doc_json).unwrap();
let doc = tantivy::doc!(json_field=>json_val); let doc = tantivy::doc!(json_field=>json_val);

View File

@@ -15,7 +15,3 @@ homepage = "https://github.com/quickwit-oss/tantivy"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies] [dependencies]
[dev-dependencies]
rand = "0.8"
proptest = "1"

View File

@@ -4,39 +4,9 @@ extern crate test;
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use rand::seq::IteratorRandom; use tantivy_bitpacker::BlockedBitpacker;
use rand::thread_rng;
use tantivy_bitpacker::{BitPacker, BitUnpacker, BlockedBitpacker};
use test::Bencher; use test::Bencher;
#[inline(never)]
fn create_bitpacked_data(bit_width: u8, num_els: u32) -> Vec<u8> {
let mut bitpacker = BitPacker::new();
let mut buffer = Vec::new();
for _ in 0..num_els {
// the values do not matter.
bitpacker.write(0u64, bit_width, &mut buffer).unwrap();
bitpacker.flush(&mut buffer).unwrap();
}
buffer
}
#[bench]
fn bench_bitpacking_read(b: &mut Bencher) {
let bit_width = 3;
let num_els = 1_000_000u32;
let bit_unpacker = BitUnpacker::new(bit_width);
let data = create_bitpacked_data(bit_width, num_els);
let idxs: Vec<u32> = (0..num_els).choose_multiple(&mut thread_rng(), 100_000);
b.iter(|| {
let mut out = 0u64;
for &idx in &idxs {
out = out.wrapping_add(bit_unpacker.get(idx, &data[..]));
}
out
});
}
#[bench] #[bench]
fn bench_blockedbitp_read(b: &mut Bencher) { fn bench_blockedbitp_read(b: &mut Bencher) {
let mut blocked_bitpacker = BlockedBitpacker::new(); let mut blocked_bitpacker = BlockedBitpacker::new();
@@ -44,9 +14,9 @@ mod tests {
blocked_bitpacker.add(val * val); blocked_bitpacker.add(val * val);
} }
b.iter(|| { b.iter(|| {
let mut out = 0u64; let mut out = 0;
for val in 0..=21500 { for val in 0..=21500 {
out = out.wrapping_add(blocked_bitpacker.get(val)); out = blocked_bitpacker.get(val);
} }
out out
}); });

View File

@@ -56,31 +56,27 @@ impl BitPacker {
pub fn close<TWrite: io::Write>(&mut self, output: &mut TWrite) -> io::Result<()> { pub fn close<TWrite: io::Write>(&mut self, output: &mut TWrite) -> io::Result<()> {
self.flush(output)?; self.flush(output)?;
// Padding the write file to simplify reads.
output.write_all(&[0u8; 7])?;
Ok(()) Ok(())
} }
} }
#[derive(Clone, Debug, Default, Copy)] #[derive(Clone, Debug, Default)]
pub struct BitUnpacker { pub struct BitUnpacker {
num_bits: u32, num_bits: u64,
mask: u64, mask: u64,
} }
impl BitUnpacker { impl BitUnpacker {
/// Creates a bit unpacker, that assumes the same bitwidth for all values.
///
/// The bitunpacker works by doing an unaligned read of 8 bytes.
/// For this reason, values of `num_bits` between
/// [57..63] are forbidden.
pub fn new(num_bits: u8) -> BitUnpacker { pub fn new(num_bits: u8) -> BitUnpacker {
assert!(num_bits <= 7 * 8 || num_bits == 64);
let mask: u64 = if num_bits == 64 { let mask: u64 = if num_bits == 64 {
!0u64 !0u64
} else { } else {
(1u64 << num_bits) - 1u64 (1u64 << num_bits) - 1u64
}; };
BitUnpacker { BitUnpacker {
num_bits: u32::from(num_bits), num_bits: u64::from(num_bits),
mask, mask,
} }
} }
@@ -91,40 +87,28 @@ impl BitUnpacker {
#[inline] #[inline]
pub fn get(&self, idx: u32, data: &[u8]) -> u64 { pub fn get(&self, idx: u32, data: &[u8]) -> u64 {
let addr_in_bits = idx * self.num_bits; if self.num_bits == 0 {
let addr = (addr_in_bits >> 3) as usize; return 0u64;
if addr + 8 > data.len() {
if self.num_bits == 0 {
return 0;
}
let bit_shift = addr_in_bits & 7;
return self.get_slow_path(addr, bit_shift, data);
} }
let addr_in_bits = idx * self.num_bits as u32;
let addr = (addr_in_bits >> 3) as usize;
let bit_shift = addr_in_bits & 7; let bit_shift = addr_in_bits & 7;
debug_assert!(
addr + 8 <= data.len(),
"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..addr + 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 = val_unshifted_unmasked >> bit_shift;
val_shifted & self.mask val_shifted & self.mask
} }
#[inline(never)]
fn get_slow_path(&self, addr: usize, bit_shift: u32, data: &[u8]) -> u64 {
let mut bytes: [u8; 8] = [0u8; 8];
let available_bytes = data.len() - addr;
// This function is meant to only be called if we did not have 8 bytes to load.
debug_assert!(available_bytes < 8);
bytes[..available_bytes].copy_from_slice(&data[addr..]);
let val_unshifted_unmasked: u64 = u64::from_le_bytes(bytes);
let val_shifted = val_unshifted_unmasked >> bit_shift;
val_shifted & self.mask
}
} }
#[cfg(test)] #[cfg(test)]
mod test { mod test {
use super::{BitPacker, BitUnpacker}; use super::{BitPacker, BitUnpacker};
fn create_bitpacker(len: usize, num_bits: u8) -> (BitUnpacker, Vec<u64>, Vec<u8>) { fn create_fastfield_bitpacker(len: usize, num_bits: u8) -> (BitUnpacker, Vec<u64>, Vec<u8>) {
let mut data = Vec::new(); let mut data = Vec::new();
let mut bitpacker = BitPacker::new(); let mut bitpacker = BitPacker::new();
let max_val: u64 = (1u64 << num_bits as u64) - 1u64; let max_val: u64 = (1u64 << num_bits as u64) - 1u64;
@@ -135,13 +119,13 @@ mod test {
bitpacker.write(val, num_bits, &mut data).unwrap(); bitpacker.write(val, num_bits, &mut data).unwrap();
} }
bitpacker.close(&mut data).unwrap(); bitpacker.close(&mut data).unwrap();
assert_eq!(data.len(), ((num_bits as usize) * len + 7) / 8); assert_eq!(data.len(), ((num_bits as usize) * len + 7) / 8 + 7);
let bitunpacker = BitUnpacker::new(num_bits); let bitunpacker = BitUnpacker::new(num_bits);
(bitunpacker, vals, data) (bitunpacker, vals, data)
} }
fn test_bitpacker_util(len: usize, num_bits: u8) { fn test_bitpacker_util(len: usize, num_bits: u8) {
let (bitunpacker, vals, data) = create_bitpacker(len, num_bits); let (bitunpacker, vals, data) = create_fastfield_bitpacker(len, num_bits);
for (i, val) in vals.iter().enumerate() { for (i, val) in vals.iter().enumerate() {
assert_eq!(bitunpacker.get(i as u32, &data), *val); assert_eq!(bitunpacker.get(i as u32, &data), *val);
} }
@@ -155,49 +139,4 @@ mod test {
test_bitpacker_util(6, 14); test_bitpacker_util(6, 14);
test_bitpacker_util(1000, 14); test_bitpacker_util(1000, 14);
} }
use proptest::prelude::*;
fn num_bits_strategy() -> impl Strategy<Value = u8> {
prop_oneof!(Just(0), Just(1), 2u8..56u8, Just(56), Just(64),)
}
fn vals_strategy() -> impl Strategy<Value = (u8, Vec<u64>)> {
(num_bits_strategy(), 0usize..100usize).prop_flat_map(|(num_bits, len)| {
let max_val = if num_bits == 64 {
u64::MAX
} else {
(1u64 << num_bits as u32) - 1
};
let vals = proptest::collection::vec(0..=max_val, len);
vals.prop_map(move |vals| (num_bits, vals))
})
}
fn test_bitpacker_aux(num_bits: u8, vals: &[u64]) {
let mut buffer: Vec<u8> = Vec::new();
let mut bitpacker = BitPacker::new();
for &val in vals {
bitpacker.write(val, num_bits, &mut buffer).unwrap();
}
bitpacker.flush(&mut buffer).unwrap();
assert_eq!(buffer.len(), (vals.len() * num_bits as usize + 7) / 8);
let bitunpacker = BitUnpacker::new(num_bits);
let max_val = if num_bits == 64 {
u64::MAX
} else {
(1u64 << num_bits) - 1
};
for (i, val) in vals.iter().copied().enumerate() {
assert!(val <= max_val);
assert_eq!(bitunpacker.get(i as u32, &buffer), val);
}
}
proptest::proptest! {
#[test]
fn test_bitpacker_proptest((num_bits, vals) in vals_strategy()) {
test_bitpacker_aux(num_bits, &vals);
}
}
} }

View File

@@ -5,23 +5,15 @@ edition = "2021"
license = "MIT" license = "MIT"
[dependencies] [dependencies]
itertools = "0.10.5"
log = "0.4.17"
fnv = "1.0.7"
fastdivide = "0.4.0"
rand = { version = "0.8.5", optional = true }
measure_time = { version = "0.8.2", optional = true }
prettytable-rs = { version = "0.10.0", optional = true }
stacker = { path = "../stacker", package="tantivy-stacker"} stacker = { path = "../stacker", package="tantivy-stacker"}
serde_json = "1"
thiserror = "1"
fnv = "1"
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" }
tantivy-bitpacker = { version= "0.3", path = "../bitpacker/" } fastfield_codecs = { path = "../fastfield_codecs"}
itertools = "0.10"
[dev-dependencies] [dev-dependencies]
proptest = "1.0.0" proptest = "1"
more-asserts = "0.3.1"
rand = "0.8.5"
[features]
unstable = []

View File

@@ -6,45 +6,52 @@ This crate describes columnar format used in tantivy.
This format is special in the following way. This format is special in the following way.
- it needs to be compact - it needs to be compact
- accessing a specific column does not require to load the entire columnar. It can be done in 2 to 3 random access. - it does not required to be loaded in memory.
- it is designed to fit well with quickwit's strange constraint:
we need to be able to load columns rapidly.
- columns of several types can be associated with the same column name. - columns of several types can be associated with the same column name.
- it needs to support columns with different types `(str, u64, i64, f64)` - it needs to support columns with different types `(str, u64, i64, f64)`
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.
- it is designed to allow range queries.
# Coercion rules # Coercion rules
Users can create a columnar by inserting rows to a `ColumnarWriter`, Users can create a columnar by appending rows to a writer.
and serializing it into a `Write` object. Nothing prevents a user from recording values with different to a same `column_key`.
Nothing prevents a user from recording values with different type to the same `column_name`.
In that case, `tantivy-columnar`'s behavior is as follows: In that case, `tantivy-columnar`'s behavior is as follows:
- JsonValues are grouped into 3 types (String, Number, bool). - 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.
Values that corresponds to different groups are mapped to different columns. For instance, String values are treated independently - 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.
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 # Columnar format
This columnar format may have more than one column (with different types) associated to the same `column_name` (see [Coercion rules](#coercion-rules) above). Because this columnar format tries to avoid some coercion.
The `(column_name, columne_type)` couple however uniquely identifies a column. There can be several columns (with different type) associated to a single `column_name`.
That couple is serialized as a column `column_key`. The format of that key is:
Each column is associated to `column_key`.
The format of that key is:
`[column_name][ZERO_BYTE][column_type_header: u8]` `[column_name][ZERO_BYTE][column_type_header: u8]`
``` ```
COLUMNAR:= COLUMNAR:=
[COLUMNAR_DATA] [COLUMNAR_DATA]
[COLUMNAR_KEY_TO_DATA_INDEX] [COLUMNAR_INDEX]
[COLUMNAR_FOOTER]; [COLUMNAR_FOOTER];
# Columns are sorted by their column key. # Columns are sorted by their column key.
COLUMNAR_DATA:= COLUMNAR_DATA:=
[COLUMN_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] COLUMNAR_FOOTER := [RANGE_SSTABLE_BYTES_LEN: 8 bytes little endian]
@@ -53,10 +60,10 @@ COLUMNAR_FOOTER := [RANGE_SSTABLE_BYTES_LEN: 8 bytes little endian]
The columnar file starts by the actual column data, concatenated one after the other, The columnar file starts by the actual column data, concatenated one after the other,
sorted by column key. sorted by column key.
A sstable associates A quickwit/tantivy style sstable associates
`(column name, column_cardinality, column_type) to range of bytes. `(column names, column_cardinality, column_type) to range of bytes.
Column name may not contain the zero byte `\0`. 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
be done by listing all keys prefixed by be done by listing all keys prefixed by
@@ -64,46 +71,3 @@ be done by listing all keys prefixed by
The associated range of bytes refer to a range of bytes The associated range of bytes refer to a range of bytes
This crate exposes a columnar format for tantivy.
This format is described in README.md
The crate introduces the following concepts.
`Columnar` is an equivalent of a dataframe.
It maps `column_key` to `Column`.
A `Column<T>` asssociates a `RowId` (u32) to any
number of values.
This is made possible by wrapping a `ColumnIndex` and a `ColumnValue` object.
The `ColumnValue<T>` represents a mapping that associates each `RowId` to
exactly one single value.
The `ColumnIndex` then maps each RowId to a set of `RowId` in the
`ColumnValue`.
For optimization, and compression purposes, the `ColumnIndex` has three
possible representation, each for different cardinalities.
- Full
All RowId have exactly one value. The ColumnIndex is the trivial mapping.
- Optional
All RowIds can have at most one value. The ColumnIndex is the trivial mapping `ColumnRowId -> Option<ColumnValueRowId>`.
- Multivalued
All RowIds can have any number of values.
The column index is mapping values to a range.
All these objects are implemented an unit tested independently
in their own module:
- columnar
- column_index
- column_values
- column

View File

@@ -0,0 +1,17 @@
[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

@@ -0,0 +1,126 @@
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,46 +0,0 @@
# zero to one
* merges
* full still needs a num_values
* replug u128
* add dictionary encoded stuff
* fix multivalued
* find a way to make columnar work with strict types
* plug to tantivy
- indexing
- aggregations
- merge
# Perf and Size
* re-add ZSTD compression for dictionaries
no systematic monotonic mapping
consider removing multilinear
f32?
adhoc solution for bool?
add metrics helper for aggregate. sum(row_id)
review inline absence/presence
improv perf of select using PDEP
compare with roaring bitmap/elias fano etc etc.
SIMD range? (see blog post)
Add alignment?
Consider another codec to bridge the gap between few and 5k elements
# Cleanup and rationalization
in benchmark, unify percent vs ratio, f32 vs f64.
investigate if should have better errors? io::Error is overused at the moment.
rename rank/select in unit tests
Review the public API via cargo doc
go through TODOs
remove all doc_id occurences -> row_id
use the rank & select naming in unit tests branch.
multi-linear -> blockwise
linear codec -> simply a multiplication for the index column
rename columnar to something more explicit, like column_dictionary or columnar_table
# Other
fix enhance column-cli
# Santa claus
autodetect datetime ipaddr, plug customizable tokenizer.

View File

@@ -1,78 +0,0 @@
use std::io;
use std::ops::Deref;
use std::sync::Arc;
use sstable::{Dictionary, VoidSSTable};
use crate::column::Column;
use crate::RowId;
/// Dictionary encoded column.
///
/// The column simply gives access to a regular u64-column that, in
/// which the values are term-ordinals.
///
/// These ordinals are ids uniquely identify the bytes that are stored in
/// the column. These ordinals are small, and sorted in the same order
/// as the term_ord_column.
#[derive(Clone)]
pub struct BytesColumn {
pub(crate) dictionary: Arc<Dictionary<VoidSSTable>>,
pub(crate) term_ord_column: Column<u64>,
}
impl BytesColumn {
/// Fills the given `output` buffer with the term associated to the ordinal `ord`.
///
/// Returns `false` if the term does not exist (e.g. `term_ord` is greater or equal to the
/// overll number of terms).
pub fn ord_to_bytes(&self, ord: u64, output: &mut Vec<u8>) -> io::Result<bool> {
self.dictionary.ord_to_term(ord, output)
}
/// Returns the number of rows in the column.
pub fn num_rows(&self) -> RowId {
self.term_ord_column.num_rows()
}
/// Returns the column of ordinals
pub fn ords(&self) -> &Column<u64> {
&self.term_ord_column
}
}
#[derive(Clone)]
pub struct StrColumn(BytesColumn);
impl From<BytesColumn> for StrColumn {
fn from(bytes_col: BytesColumn) -> Self {
StrColumn(bytes_col)
}
}
impl StrColumn {
/// Fills the buffer
pub fn ord_to_str(&self, term_ord: u64, output: &mut String) -> io::Result<bool> {
unsafe {
let buf = output.as_mut_vec();
self.0.dictionary.ord_to_term(term_ord, buf)?;
// TODO consider remove checks if it hurts performance.
if std::str::from_utf8(buf.as_slice()).is_err() {
buf.clear();
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"Not valid utf-8",
));
}
}
Ok(true)
}
}
impl Deref for StrColumn {
type Target = BytesColumn;
fn deref(&self) -> &Self::Target {
&self.0
}
}

View File

@@ -1,109 +0,0 @@
mod dictionary_encoded;
mod serialize;
use std::ops::Deref;
use std::sync::Arc;
use common::BinarySerializable;
pub use dictionary_encoded::{BytesColumn, StrColumn};
pub use serialize::{
open_column_bytes, open_column_u128, open_column_u64, serialize_column_mappable_to_u128,
serialize_column_mappable_to_u64,
};
use crate::column_index::ColumnIndex;
use crate::column_values::ColumnValues;
use crate::{Cardinality, RowId};
#[derive(Clone)]
pub struct Column<T> {
pub idx: ColumnIndex<'static>,
pub values: Arc<dyn ColumnValues<T>>,
}
impl<T: PartialOrd> Column<T> {
pub fn num_rows(&self) -> RowId {
match &self.idx {
ColumnIndex::Full => self.values.num_vals() as u32,
ColumnIndex::Optional(optional_index) => optional_index.num_rows(),
ColumnIndex::Multivalued(col_index) => {
// The multivalued index contains all value start row_id,
// and one extra value at the end with the overall number of rows.
col_index.num_vals() - 1
}
}
}
pub fn min_value(&self) -> T {
self.values.min_value()
}
pub fn max_value(&self) -> T {
self.values.max_value()
}
}
impl<T: PartialOrd + Copy + Send + Sync + 'static> Column<T> {
pub fn first(&self, row_id: RowId) -> Option<T> {
self.values(row_id).next()
}
pub fn values(&self, row_id: RowId) -> impl Iterator<Item = T> + '_ {
self.value_row_ids(row_id)
.map(|value_row_id: RowId| self.values.get_val(value_row_id))
}
pub fn first_or_default_col(self, default_value: T) -> Arc<dyn ColumnValues<T>> {
Arc::new(FirstValueWithDefault {
column: self,
default_value,
})
}
}
impl<T> Deref for Column<T> {
type Target = ColumnIndex<'static>;
fn deref(&self) -> &Self::Target {
&self.idx
}
}
impl BinarySerializable for Cardinality {
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
self.to_code().serialize(writer)
}
fn deserialize<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
let cardinality_code = u8::deserialize(reader)?;
let cardinality = Cardinality::try_from_code(cardinality_code)?;
Ok(cardinality)
}
}
// TODO simplify or optimize
struct FirstValueWithDefault<T: Copy> {
column: Column<T>,
default_value: T,
}
impl<T: PartialOrd + Send + Sync + Copy + 'static> ColumnValues<T> for FirstValueWithDefault<T> {
fn get_val(&self, idx: u32) -> T {
self.column.first(idx).unwrap_or(self.default_value)
}
fn min_value(&self) -> T {
self.column.values.min_value()
}
fn max_value(&self) -> T {
self.column.values.max_value()
}
fn num_vals(&self) -> u32 {
match &self.column.idx {
ColumnIndex::Full => self.column.values.num_vals(),
ColumnIndex::Optional(optional_idx) => optional_idx.num_rows(),
ColumnIndex::Multivalued(_) => todo!(),
}
}
}

View File

@@ -1,101 +0,0 @@
use std::io;
use std::io::Write;
use std::sync::Arc;
use common::OwnedBytes;
use sstable::Dictionary;
use crate::column::{BytesColumn, Column};
use crate::column_index::{serialize_column_index, SerializableColumnIndex};
use crate::column_values::serialize::serialize_column_values_u128;
use crate::column_values::{
serialize_column_values, ColumnValues, FastFieldCodecType, MonotonicallyMappableToU128,
MonotonicallyMappableToU64,
};
pub fn serialize_column_mappable_to_u128<
F: Fn() -> I,
I: Iterator<Item = T>,
T: MonotonicallyMappableToU128,
>(
column_index: SerializableColumnIndex<'_>,
column_values: F,
num_vals: u32,
output: &mut impl Write,
) -> io::Result<()> {
let column_index_num_bytes = serialize_column_index(column_index, output)?;
serialize_column_values_u128(
|| column_values().map(|val| val.to_u128()),
num_vals,
output,
)?;
output.write_all(&column_index_num_bytes.to_le_bytes())?;
Ok(())
}
pub fn serialize_column_mappable_to_u64<T: MonotonicallyMappableToU64>(
column_index: SerializableColumnIndex<'_>,
column_values: &impl ColumnValues<T>,
output: &mut impl Write,
) -> io::Result<()> {
let column_index_num_bytes = serialize_column_index(column_index, output)?;
serialize_column_values(
column_values,
&[
FastFieldCodecType::Bitpacked,
FastFieldCodecType::BlockwiseLinear,
],
output,
)?;
output.write_all(&column_index_num_bytes.to_le_bytes())?;
Ok(())
}
pub fn open_column_u64<T: MonotonicallyMappableToU64>(bytes: OwnedBytes) -> io::Result<Column<T>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data)?;
let column_values = crate::column_values::open_u64_mapped(column_values_data)?;
Ok(Column {
idx: column_index,
values: column_values,
})
}
pub fn open_column_u128<T: MonotonicallyMappableToU128>(
bytes: OwnedBytes,
) -> io::Result<Column<T>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data)?;
let column_values = crate::column_values::open_u128_mapped(column_values_data)?;
Ok(Column {
idx: column_index,
values: column_values,
})
}
pub fn open_column_bytes<T: From<BytesColumn>>(data: OwnedBytes) -> io::Result<T> {
let (body, dictionary_len_bytes) = data.rsplit(4);
let dictionary_len = u32::from_le_bytes(dictionary_len_bytes.as_slice().try_into().unwrap());
let (dictionary_bytes, column_bytes) = body.split(dictionary_len as usize);
let dictionary = Arc::new(Dictionary::from_bytes(dictionary_bytes)?);
let term_ord_column = crate::column::open_column_u64::<u64>(column_bytes)?;
let bytes_column = BytesColumn {
dictionary,
term_ord_column,
};
Ok(bytes_column.into())
}

View File

@@ -1,54 +0,0 @@
mod multivalued_index;
mod optional_index;
mod serialize;
use std::ops::Range;
use std::sync::Arc;
pub use optional_index::{OptionalIndex, SerializableOptionalIndex, Set};
pub use serialize::{open_column_index, serialize_column_index, SerializableColumnIndex};
use crate::column_values::ColumnValues;
use crate::{Cardinality, RowId};
#[derive(Clone)]
pub enum ColumnIndex<'a> {
Full,
Optional(OptionalIndex),
// TODO Remove the static by fixing the codec if possible.
/// The column values enclosed contains for all row_id,
/// the value start_index.
///
/// In addition, at index num_rows, an extra value is added
/// containing the overal number of values.
Multivalued(Arc<dyn ColumnValues<RowId> + 'a>),
}
impl<'a> ColumnIndex<'a> {
pub fn get_cardinality(&self) -> Cardinality {
match self {
ColumnIndex::Full => Cardinality::Full,
ColumnIndex::Optional(_) => Cardinality::Optional,
ColumnIndex::Multivalued(_) => Cardinality::Multivalued,
}
}
pub fn value_row_ids(&self, row_id: RowId) -> Range<RowId> {
match self {
ColumnIndex::Full => row_id..row_id + 1,
ColumnIndex::Optional(optional_index) => {
if let Some(val) = optional_index.rank_if_exists(row_id) {
val..val + 1
} else {
0..0
}
}
ColumnIndex::Multivalued(multivalued_index) => {
let multivalued_index_ref = &**multivalued_index;
let start: u32 = multivalued_index_ref.get_val(row_id);
let end: u32 = multivalued_index_ref.get_val(row_id + 1);
start..end
}
}
}
}

View File

@@ -1,29 +0,0 @@
use std::io;
use std::io::Write;
use std::sync::Arc;
use common::OwnedBytes;
use crate::column_values::{ColumnValues, FastFieldCodecType};
use crate::RowId;
#[derive(Clone)]
pub struct MultivaluedIndex(Arc<dyn ColumnValues<RowId>>);
pub fn serialize_multivalued_index(
multivalued_index: &dyn ColumnValues<RowId>,
output: &mut impl Write,
) -> io::Result<()> {
crate::column_values::serialize_column_values(
&*multivalued_index,
&[FastFieldCodecType::Bitpacked, FastFieldCodecType::Linear],
output,
)?;
Ok(())
}
pub fn open_multivalued_index(bytes: OwnedBytes) -> io::Result<Arc<dyn ColumnValues<RowId>>> {
let start_index_column: Arc<dyn ColumnValues<RowId>> =
crate::column_values::open_u64_mapped(bytes)?;
Ok(start_index_column)
}

View File

@@ -1,453 +0,0 @@
use std::io::{self, Write};
use std::ops::Range;
use std::sync::Arc;
mod set;
mod set_block;
use common::{BinarySerializable, GroupByIteratorExtended, OwnedBytes, VInt};
pub use set::{Set, SetCodec};
use set_block::{
DenseBlock, DenseBlockCodec, SparseBlock, SparseBlockCodec, DENSE_BLOCK_NUM_BYTES,
};
use crate::{InvalidData, RowId};
/// The threshold for for number of elements after which we switch to dense block encoding.
///
/// We simply pick the value that minimize the size of the blocks.
const DENSE_BLOCK_THRESHOLD: u32 =
set_block::DENSE_BLOCK_NUM_BYTES / std::mem::size_of::<u16>() as u32; //< 5_120
const ELEMENTS_PER_BLOCK: u32 = u16::MAX as u32 + 1;
const BLOCK_SIZE: RowId = 1 << 16;
#[derive(Copy, Clone, Debug)]
struct BlockMeta {
non_null_rows_before_block: u32,
start_byte_offset: u32,
block_variant: BlockVariant,
}
#[derive(Clone, Copy, Debug)]
enum BlockVariant {
Dense,
Sparse { num_vals: u16 },
}
impl BlockVariant {
pub fn empty() -> Self {
Self::Sparse { num_vals: 0 }
}
pub fn num_bytes_in_block(&self) -> u32 {
match *self {
BlockVariant::Dense => set_block::DENSE_BLOCK_NUM_BYTES,
BlockVariant::Sparse { num_vals } => num_vals as u32 * 2,
}
}
}
/// This codec is inspired by roaring bitmaps.
/// In the dense blocks, however, in order to accelerate `select`
/// we interleave an offset over two bytes. (more on this lower)
///
/// The lower 16 bits of doc ids are stored as u16 while the upper 16 bits are given by the block
/// id. Each block contains 1<<16 docids.
///
/// # Serialized Data Layout
/// The data starts with the block data. Each block is either dense or sparse encoded, depending on
/// the number of values in the block. A block is sparse when it contains less than
/// DENSE_BLOCK_THRESHOLD (6144) values.
/// [Sparse data block | dense data block, .. #repeat*; Desc: Either a sparse or dense encoded
/// block]
/// ### Sparse block data
/// [u16 LE, .. #repeat*; Desc: Positions with values in a block]
/// ### Dense block data
/// [Dense codec for the whole block; Desc: Similar to a bitvec(0..ELEMENTS_PER_BLOCK) + Metadata
/// for faster lookups. See dense.rs]
///
/// The data is followed by block metadata, to know which area of the raw block data belongs to
/// which block. Only metadata for blocks with elements is recorded to
/// keep the overhead low for scenarios with many very sparse columns. The block metadata consists
/// of the block index and the number of values in the block. Since we don't store empty blocks
/// num_vals is incremented by 1, e.g. 0 means 1 value.
///
/// The last u16 is storing the number of metadata blocks.
/// [u16 LE, .. #repeat*; Desc: Positions with values in a block][(u16 LE, u16 LE), .. #repeat*;
/// Desc: (Block Id u16, Num Elements u16)][u16 LE; Desc: num blocks with values u16]
///
/// # Opening
/// When opening the data layout, the data is expanded to `Vec<SparseCodecBlockVariant>`, where the
/// index is the block index. For each block `byte_start` and `offset` is computed.
#[derive(Clone)]
pub struct OptionalIndex {
num_rows: RowId,
num_non_null_rows: RowId,
block_data: OwnedBytes,
block_metas: Arc<[BlockMeta]>,
}
impl OptionalIndex {
pub fn num_rows(&self) -> RowId {
self.num_rows
}
pub fn num_non_nulls(&self) -> RowId {
self.num_non_null_rows
}
}
/// Splits a value address into lower and upper 16bits.
/// The lower 16 bits are the value in the block
/// The upper 16 bits are the block index
#[derive(Copy, Debug, Clone)]
struct RowAddr {
block_id: u16,
in_block_row_id: u16,
}
#[inline(always)]
fn row_addr_from_row_id(row_id: RowId) -> RowAddr {
RowAddr {
block_id: (row_id / BLOCK_SIZE) as u16,
in_block_row_id: (row_id % BLOCK_SIZE) as u16,
}
}
impl Set<RowId> for OptionalIndex {
// Check if value at position is not null.
#[inline]
fn contains(&self, row_id: RowId) -> bool {
let RowAddr {
block_id,
in_block_row_id,
} = row_addr_from_row_id(row_id);
let block_meta = self.block_metas[block_id as usize];
match self.block(block_meta) {
Block::Dense(dense_block) => dense_block.contains(in_block_row_id),
Block::Sparse(sparse_block) => sparse_block.contains(in_block_row_id),
}
}
#[inline]
fn rank_if_exists(&self, row_id: RowId) -> Option<RowId> {
let RowAddr {
block_id,
in_block_row_id,
} = row_addr_from_row_id(row_id);
let block_meta = self.block_metas[block_id as usize];
let block = self.block(block_meta);
let block_offset_row_id = match block {
Block::Dense(dense_block) => dense_block.rank_if_exists(in_block_row_id),
Block::Sparse(sparse_block) => sparse_block.rank_if_exists(in_block_row_id),
}? as u32;
Some(block_meta.non_null_rows_before_block + block_offset_row_id)
}
#[inline]
fn select(&self, rank: RowId) -> RowId {
let block_pos = self.find_block(rank, 0);
let block_doc_idx_start = block_pos * ELEMENTS_PER_BLOCK;
let block_meta = self.block_metas[block_pos as usize];
let block: Block<'_> = self.block(block_meta);
let index_in_block = (rank - block_meta.non_null_rows_before_block) as u16;
let in_block_rank = match block {
Block::Dense(dense_block) => dense_block.select(index_in_block),
Block::Sparse(sparse_block) => sparse_block.select(index_in_block),
};
block_doc_idx_start + in_block_rank as u32
}
fn select_batch(&self, ranks: &[u32], output_idxs: &mut [u32]) {
let mut block_pos = 0u32;
let mut start = 0;
let group_by_it = ranks.iter().copied().group_by(move |codec_idx| {
block_pos = self.find_block(*codec_idx, block_pos);
block_pos
});
for (block_pos, block_iter) in group_by_it {
let block_doc_idx_start = block_pos * ELEMENTS_PER_BLOCK;
let block_meta = self.block_metas[block_pos as usize];
let block: Block<'_> = self.block(block_meta);
let offset = block_meta.non_null_rows_before_block;
let indexes_in_block_iter =
block_iter.map(move |codec_idx| (codec_idx - offset) as u16);
match block {
Block::Dense(dense_block) => {
for in_offset in dense_block.select_iter(indexes_in_block_iter) {
output_idxs[start] = in_offset as u32 + block_doc_idx_start;
start += 1;
}
}
Block::Sparse(sparse_block) => {
for in_offset in sparse_block.select_iter(indexes_in_block_iter) {
output_idxs[start] = in_offset as u32 + block_doc_idx_start;
start += 1;
}
}
};
}
}
}
impl OptionalIndex {
#[inline]
fn block<'a>(&'a self, block_meta: BlockMeta) -> Block<'a> {
let BlockMeta {
start_byte_offset,
block_variant,
..
} = block_meta;
let start_byte_offset = start_byte_offset as usize;
let bytes = self.block_data.as_slice();
match block_variant {
BlockVariant::Dense => Block::Dense(DenseBlockCodec::open(
&bytes[start_byte_offset..start_byte_offset + DENSE_BLOCK_NUM_BYTES as usize],
)),
BlockVariant::Sparse { num_vals } => {
let end_byte_offset = start_byte_offset + num_vals as usize * 2;
let sparse_bytes = &bytes[start_byte_offset..end_byte_offset];
Block::Sparse(SparseBlockCodec::open(sparse_bytes))
}
}
}
#[inline]
fn find_block(&self, dense_idx: u32, start_block_pos: u32) -> u32 {
for block_pos in start_block_pos..self.block_metas.len() as u32 {
let offset = self.block_metas[block_pos as usize].non_null_rows_before_block;
if offset > dense_idx {
return block_pos - 1;
}
}
self.block_metas.len() as u32 - 1u32
}
// TODO Add a good API for the codec_idx to original_idx translation.
// The Iterator API is a probably a bad idea
}
#[derive(Copy, Clone)]
enum Block<'a> {
Dense(DenseBlock<'a>),
Sparse(SparseBlock<'a>),
}
#[derive(Debug, Copy, Clone)]
enum OptionalIndexCodec {
Dense = 0,
Sparse = 1,
}
impl OptionalIndexCodec {
fn to_code(self) -> u8 {
self as u8
}
fn try_from_code(code: u8) -> Result<Self, InvalidData> {
match code {
0 => Ok(Self::Dense),
1 => Ok(Self::Sparse),
_ => Err(InvalidData),
}
}
}
impl BinarySerializable for OptionalIndexCodec {
fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
writer.write_all(&[self.to_code()])
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let optional_codec_code = u8::deserialize(reader)?;
let optional_codec = Self::try_from_code(optional_codec_code)?;
Ok(optional_codec)
}
}
fn serialize_optional_index_block(block_els: &[u16], out: &mut impl io::Write) -> io::Result<()> {
let is_sparse = is_sparse(block_els.len() as u32);
if is_sparse {
SparseBlockCodec::serialize(block_els.iter().copied(), out)?;
} else {
DenseBlockCodec::serialize(block_els.iter().copied(), out)?;
}
Ok(())
}
pub fn serialize_optional_index<'a, W: io::Write>(
serializable_optional_index: &dyn SerializableOptionalIndex<'a>,
output: &mut W,
) -> io::Result<()> {
VInt(serializable_optional_index.num_rows() as u64).serialize(output)?;
let mut rows_it = serializable_optional_index.non_null_rows();
let mut block_metadata: Vec<SerializedBlockMeta> = Vec::new();
let mut current_block = Vec::new();
// This if-statement for the first element ensures that
// `block_metadata` is not empty in the loop below.
let Some(idx) = rows_it.next() else {
output.write_all(&0u16.to_le_bytes())?;
return Ok(());
};
let row_addr = row_addr_from_row_id(idx);
let mut current_block_id = row_addr.block_id;
current_block.push(row_addr.in_block_row_id);
for idx in rows_it {
let value_addr = row_addr_from_row_id(idx);
if current_block_id != value_addr.block_id {
serialize_optional_index_block(&current_block[..], output)?;
block_metadata.push(SerializedBlockMeta {
block_id: current_block_id,
num_non_null_rows: current_block.len() as u32,
});
current_block.clear();
current_block_id = value_addr.block_id;
}
current_block.push(value_addr.in_block_row_id);
}
// handle last block
serialize_optional_index_block(&current_block[..], output)?;
block_metadata.push(SerializedBlockMeta {
block_id: current_block_id,
num_non_null_rows: current_block.len() as u32,
});
for block in &block_metadata {
output.write_all(&block.to_bytes())?;
}
output.write_all((block_metadata.len() as u16).to_le_bytes().as_ref())?;
Ok(())
}
const SERIALIZED_BLOCK_META_NUM_BYTES: usize = 4;
#[derive(Clone, Copy, Debug)]
struct SerializedBlockMeta {
block_id: u16,
num_non_null_rows: u32, //< takes values in 1..=u16::MAX
}
// TODO unit tests
impl SerializedBlockMeta {
#[inline]
fn from_bytes(bytes: [u8; SERIALIZED_BLOCK_META_NUM_BYTES]) -> SerializedBlockMeta {
let block_id = u16::from_le_bytes(bytes[0..2].try_into().unwrap());
let num_non_null_rows: u32 =
u16::from_le_bytes(bytes[2..4].try_into().unwrap()) as u32 + 1u32;
SerializedBlockMeta {
block_id,
num_non_null_rows,
}
}
#[inline]
fn to_bytes(&self) -> [u8; SERIALIZED_BLOCK_META_NUM_BYTES] {
assert!(self.num_non_null_rows > 0);
let mut bytes = [0u8; SERIALIZED_BLOCK_META_NUM_BYTES];
bytes[0..2].copy_from_slice(&self.block_id.to_le_bytes());
// We don't store empty blocks, therefore we can subtract 1.
// This way we will be able to use u16 when the number of elements is 1 << 16 or u16::MAX+1
bytes[2..4].copy_from_slice(&((self.num_non_null_rows - 1u32) as u16).to_le_bytes());
bytes
}
}
#[inline]
fn is_sparse(num_rows_in_block: u32) -> bool {
num_rows_in_block < DENSE_BLOCK_THRESHOLD as u32
}
fn deserialize_optional_index_block_metadatas(
data: &[u8],
num_rows: u32,
) -> (Box<[BlockMeta]>, u32) {
let num_blocks = data.len() / SERIALIZED_BLOCK_META_NUM_BYTES;
let mut block_metas = Vec::with_capacity(num_blocks as usize + 1);
let mut start_byte_offset = 0;
let mut non_null_rows_before_block = 0;
for block_meta_bytes in data.chunks_exact(SERIALIZED_BLOCK_META_NUM_BYTES) {
let block_meta_bytes: [u8; SERIALIZED_BLOCK_META_NUM_BYTES] =
block_meta_bytes.try_into().unwrap();
let SerializedBlockMeta {
block_id,
num_non_null_rows,
} = SerializedBlockMeta::from_bytes(block_meta_bytes);
block_metas.resize(
block_id as usize,
BlockMeta {
non_null_rows_before_block,
start_byte_offset,
block_variant: BlockVariant::empty(),
},
);
let block_variant = if is_sparse(num_non_null_rows) {
BlockVariant::Sparse {
num_vals: num_non_null_rows as u16,
}
} else {
BlockVariant::Dense
};
block_metas.push(BlockMeta {
non_null_rows_before_block,
start_byte_offset,
block_variant,
});
start_byte_offset += block_variant.num_bytes_in_block();
non_null_rows_before_block += num_non_null_rows as u32;
}
block_metas.resize(
((num_rows + BLOCK_SIZE - 1) / BLOCK_SIZE) as usize,
BlockMeta {
non_null_rows_before_block,
start_byte_offset,
block_variant: BlockVariant::empty(),
},
);
(block_metas.into_boxed_slice(), non_null_rows_before_block)
}
pub fn open_optional_index(bytes: OwnedBytes) -> io::Result<OptionalIndex> {
let (mut bytes, num_non_empty_blocks_bytes) = bytes.rsplit(2);
let num_non_empty_block_bytes =
u16::from_le_bytes(num_non_empty_blocks_bytes.as_slice().try_into().unwrap());
let num_rows = VInt::deserialize_u64(&mut bytes)? as u32;
let block_metas_num_bytes =
num_non_empty_block_bytes as usize * SERIALIZED_BLOCK_META_NUM_BYTES;
let (block_data, block_metas) = bytes.rsplit(block_metas_num_bytes);
let (block_metas, num_non_null_rows) =
deserialize_optional_index_block_metadatas(block_metas.as_slice(), num_rows).into();
let optional_index = OptionalIndex {
num_rows,
num_non_null_rows,
block_data,
block_metas: block_metas.into(),
};
Ok(optional_index)
}
pub trait SerializableOptionalIndex<'a> {
fn num_rows(&self) -> RowId;
fn non_null_rows(&self) -> Box<dyn Iterator<Item = RowId> + 'a>;
}
impl SerializableOptionalIndex<'static> for Range<u32> {
fn num_rows(&self) -> RowId {
self.end
}
fn non_null_rows(&self) -> Box<dyn Iterator<Item = RowId> + 'static> {
Box::new(self.clone())
}
}
#[cfg(test)]
mod tests;

View File

@@ -1,38 +0,0 @@
use std::io;
/// A codec makes it possible to serialize a set of
/// elements, and open the resulting Set representation.
pub trait SetCodec {
type Item: Copy + TryFrom<usize> + Eq + std::hash::Hash + std::fmt::Debug;
type Reader<'a>: Set<Self::Item>;
/// Serializes a set of unique sorted u16 elements.
///
/// May panic if the elements are not sorted.
fn serialize(els: impl Iterator<Item = Self::Item>, wrt: impl io::Write) -> io::Result<()>;
fn open<'a>(data: &'a [u8]) -> Self::Reader<'a>;
}
pub trait Set<T> {
/// Returns true if the elements is contained in the Set
fn contains(&self, el: T) -> bool;
/// If the set contains `el` returns its position in the sortd set of elements.
/// If the set does not contain the element, it returns `None`.
fn rank_if_exists(&self, el: T) -> Option<T>;
/// Return the rank-th value stored in this bitmap.
///
/// # Panics
///
/// May panic if rank is greater than the number of elements in the Set.
fn select(&self, rank: T) -> T;
/// Batch version of select.
/// `ranks` is assumed to be sorted.
///
/// # Panics
///
/// May panic if rank is greater than the number of elements in the Set.
fn select_batch(&self, ranks: &[T], outputs: &mut [T]);
}

View File

@@ -1,271 +0,0 @@
use std::convert::TryInto;
use std::io::{self, Write};
use common::BinarySerializable;
use crate::column_index::optional_index::{Set, SetCodec, ELEMENTS_PER_BLOCK};
#[inline(always)]
fn get_bit_at(input: u64, n: u16) -> bool {
input & (1 << n) != 0
}
#[inline]
fn set_bit_at(input: &mut u64, n: u16) {
*input |= 1 << n;
}
/// For the `DenseCodec`, `data` which contains the encoded blocks.
/// Each block consists of [u8; 12]. The first 8 bytes is a bitvec for 64 elements.
/// The last 4 bytes are the offset, the number of set bits so far.
///
/// When translating the original index to a dense index, the correct block can be computed
/// directly `orig_idx/64`. Inside the block the position is `orig_idx%64`.
///
/// When translating a dense index to the original index, we can use the offset to find the correct
/// block. Direct computation is not possible, but we can employ a linear or binary search.
const ELEMENTS_PER_MINI_BLOCK: u16 = 64;
const MINI_BLOCK_BITVEC_NUM_BYTES: usize = 8;
const MINI_BLOCK_OFFSET_NUM_BYTES: usize = 2;
pub const MINI_BLOCK_NUM_BYTES: usize = MINI_BLOCK_BITVEC_NUM_BYTES + MINI_BLOCK_OFFSET_NUM_BYTES;
/// Number of bytes in a dense block.
pub const DENSE_BLOCK_NUM_BYTES: u32 =
(ELEMENTS_PER_BLOCK as u32 / ELEMENTS_PER_MINI_BLOCK as u32) * MINI_BLOCK_NUM_BYTES as u32;
pub struct DenseBlockCodec;
impl SetCodec for DenseBlockCodec {
type Item = u16;
type Reader<'a> = DenseBlock<'a>;
fn serialize(els: impl Iterator<Item = u16>, wrt: impl io::Write) -> io::Result<()> {
serialize_dense_codec(els, wrt)
}
#[inline]
fn open<'a>(data: &'a [u8]) -> Self::Reader<'a> {
assert_eq!(data.len(), DENSE_BLOCK_NUM_BYTES as usize);
DenseBlock(data)
}
}
/// Interpreting the bitvec as a set of integer within 0..=63
/// and given an element, returns the number of elements in the
/// set lesser than the element.
///
/// # Panics
///
/// May panic or return a wrong result if el <= 64.
#[inline(always)]
fn rank_u64(bitvec: u64, el: u16) -> u16 {
debug_assert!(el < 64);
let mask = (1u64 << el) - 1;
let masked_bitvec = bitvec & mask;
masked_bitvec.count_ones() as u16
}
#[inline(always)]
fn select_u64(mut bitvec: u64, rank: u16) -> u16 {
for _ in 0..rank {
bitvec &= bitvec - 1;
}
bitvec.trailing_zeros() as u16
}
// TODO test the following solution on Intel... on Ryzen Zen <3 it is a catastrophy.
// #[target_feature(enable = "bmi2")]
// unsafe fn select_bitvec_unsafe(bitvec: u64, rank: u16) -> u16 {
// let pdep = _pdep_u64(1u64 << rank, bitvec);
// pdep.trailing_zeros() as u16
// }
#[derive(Clone, Copy, Debug)]
struct DenseMiniBlock {
bitvec: u64,
rank: u16,
}
impl DenseMiniBlock {
fn from_bytes(data: [u8; MINI_BLOCK_NUM_BYTES]) -> Self {
let bitvec = u64::from_le_bytes(data[..MINI_BLOCK_BITVEC_NUM_BYTES].try_into().unwrap());
let rank = u16::from_le_bytes(data[MINI_BLOCK_BITVEC_NUM_BYTES..].try_into().unwrap());
Self { bitvec, rank }
}
fn to_bytes(&self) -> [u8; MINI_BLOCK_NUM_BYTES] {
let mut bytes = [0u8; MINI_BLOCK_NUM_BYTES];
bytes[..MINI_BLOCK_BITVEC_NUM_BYTES].copy_from_slice(&self.bitvec.to_le_bytes());
bytes[MINI_BLOCK_BITVEC_NUM_BYTES..].copy_from_slice(&self.rank.to_le_bytes());
bytes
}
}
#[derive(Copy, Clone)]
pub struct DenseBlock<'a>(&'a [u8]);
impl<'a> Set<u16> for DenseBlock<'a> {
#[inline(always)]
fn contains(&self, el: u16) -> bool {
let mini_block_id = el / ELEMENTS_PER_MINI_BLOCK;
let bitvec = self.mini_block(mini_block_id).bitvec;
let pos_in_bitvec = el % ELEMENTS_PER_MINI_BLOCK;
get_bit_at(bitvec, pos_in_bitvec)
}
#[inline(always)]
fn rank_if_exists(&self, el: u16) -> Option<u16> {
let block_pos = el / ELEMENTS_PER_MINI_BLOCK;
let index_block = self.mini_block(block_pos);
let pos_in_block_bit_vec = el % ELEMENTS_PER_MINI_BLOCK;
let ones_in_block = rank_u64(index_block.bitvec, pos_in_block_bit_vec);
let rank = index_block.rank + ones_in_block;
if get_bit_at(index_block.bitvec, pos_in_block_bit_vec) {
Some(rank)
} else {
None
}
}
#[inline(always)]
fn select(&self, rank: u16) -> u16 {
let block_id = self.find_miniblock_containing_rank(rank, 0).unwrap();
let index_block = self.mini_block(block_id);
let in_block_rank = rank - index_block.rank;
block_id * ELEMENTS_PER_MINI_BLOCK + select_u64(index_block.bitvec, in_block_rank)
}
fn select_batch(&self, ranks: &[u16], outputs: &mut [u16]) {
let orig_ids = self.select_iter(ranks.iter().copied());
for (output, original_id) in outputs.iter_mut().zip(orig_ids) {
*output = original_id;
}
}
}
impl<'a> DenseBlock<'a> {
/// Iterator verison of select.
///
/// # Panics
/// Panics if one of the rank is higher than the number of elements in the set.
pub fn select_iter<'b>(
&self,
rank_it: impl Iterator<Item = u16> + 'b,
) -> impl Iterator<Item = u16> + 'b
where
Self: 'b,
{
let mut block_id = 0u16;
let me = *self;
rank_it.map(move |rank| {
block_id = me.find_miniblock_containing_rank(rank, block_id).unwrap();
let index_block = me.mini_block(block_id);
let in_block_rank = rank - index_block.rank;
block_id * ELEMENTS_PER_MINI_BLOCK + select_u64(index_block.bitvec, in_block_rank)
})
}
}
impl<'a> DenseBlock<'a> {
#[inline]
fn mini_block(&self, mini_block_id: u16) -> DenseMiniBlock {
let data_start_pos = mini_block_id as usize * MINI_BLOCK_NUM_BYTES;
DenseMiniBlock::from_bytes(
self.0[data_start_pos..data_start_pos + MINI_BLOCK_NUM_BYTES]
.try_into()
.unwrap(),
)
}
#[inline]
fn iter_miniblocks(
&self,
from_block_id: u16,
) -> impl Iterator<Item = (u16, DenseMiniBlock)> + '_ {
self.0
.chunks_exact(MINI_BLOCK_NUM_BYTES)
.enumerate()
.skip(from_block_id as usize)
.map(|(block_id, bytes)| {
let mini_block = DenseMiniBlock::from_bytes(bytes.try_into().unwrap());
(block_id as u16, mini_block)
})
}
/// Finds the block position containing the dense_idx.
///
/// # Correctness
/// dense_idx needs to be smaller than the number of values in the index
///
/// The last offset number is equal to the number of values in the index.
#[inline]
fn find_miniblock_containing_rank(&self, rank: u16, from_block_id: u16) -> Option<u16> {
self.iter_miniblocks(from_block_id)
.take_while(|(_, block)| block.rank <= rank)
.map(|(block_id, _)| block_id)
.last()
}
}
/// Iterator over all values, true if set, otherwise false
pub fn serialize_dense_codec(
els: impl Iterator<Item = u16>,
mut output: impl Write,
) -> io::Result<()> {
let mut non_null_rows_before: u16 = 0u16;
let mut block = 0u64;
let mut current_block_id = 0u16;
for el in els {
let block_id = el / ELEMENTS_PER_MINI_BLOCK;
let in_offset = el % ELEMENTS_PER_MINI_BLOCK;
while block_id > current_block_id {
let dense_mini_block = DenseMiniBlock {
bitvec: block,
rank: non_null_rows_before as u16,
};
output.write_all(&dense_mini_block.to_bytes())?;
non_null_rows_before += block.count_ones() as u16;
block = 0u64;
current_block_id += 1u16;
}
set_bit_at(&mut block, in_offset);
}
while current_block_id <= u16::MAX / ELEMENTS_PER_MINI_BLOCK {
block.serialize(&mut output)?;
non_null_rows_before.serialize(&mut output)?;
// This will overflow to 0 exactly if all bits are set.
// This is however not problem as we won't use this last value.
non_null_rows_before = non_null_rows_before.wrapping_add(block.count_ones() as u16);
block = 0u64;
current_block_id += 1u16;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_select_bitvec() {
assert_eq!(select_u64(1u64, 0), 0);
assert_eq!(select_u64(2u64, 0), 1);
assert_eq!(select_u64(4u64, 0), 2);
assert_eq!(select_u64(8u64, 0), 3);
assert_eq!(select_u64(1 | 8u64, 0), 0);
assert_eq!(select_u64(1 | 8u64, 1), 3);
}
#[test]
fn test_count_ones() {
for i in 0..=63 {
assert_eq!(rank_u64(u64::MAX, i), i);
}
}
#[test]
fn test_dense() {
assert_eq!(DENSE_BLOCK_NUM_BYTES, 10_240);
}
}

View File

@@ -1,8 +0,0 @@
mod dense;
mod sparse;
pub use dense::{DenseBlock, DenseBlockCodec, DENSE_BLOCK_NUM_BYTES};
pub use sparse::{SparseBlock, SparseBlockCodec};
#[cfg(test)]
mod tests;

View File

@@ -1,112 +0,0 @@
use crate::column_index::optional_index::{Set, SetCodec};
pub struct SparseBlockCodec;
impl SetCodec for SparseBlockCodec {
type Item = u16;
type Reader<'a> = SparseBlock<'a>;
fn serialize(
els: impl Iterator<Item = u16>,
mut wrt: impl std::io::Write,
) -> std::io::Result<()> {
for el in els {
wrt.write_all(&el.to_le_bytes())?;
}
Ok(())
}
fn open<'a>(data: &'a [u8]) -> Self::Reader<'a> {
SparseBlock(data)
}
}
#[derive(Copy, Clone)]
pub struct SparseBlock<'a>(&'a [u8]);
impl<'a> Set<u16> for SparseBlock<'a> {
#[inline(always)]
fn contains(&self, el: u16) -> bool {
self.binary_search(el).is_ok()
}
#[inline(always)]
fn rank_if_exists(&self, el: u16) -> Option<u16> {
self.binary_search(el).ok()
}
#[inline(always)]
fn select(&self, rank: u16) -> u16 {
let offset = rank as usize * 2;
u16::from_le_bytes(self.0[offset..offset + 2].try_into().unwrap())
}
fn select_batch(&self, ranks: &[u16], outputs: &mut [u16]) {
let orig_ids = self.select_iter(ranks.iter().copied());
for (output, original_id) in outputs.iter_mut().zip(orig_ids) {
*output = original_id;
}
}
}
#[inline(always)]
fn get_u16(data: &[u8], byte_position: usize) -> u16 {
let bytes: [u8; 2] = data[byte_position..byte_position + 2].try_into().unwrap();
u16::from_le_bytes(bytes)
}
impl<'a> SparseBlock<'a> {
#[inline(always)]
fn value_at_idx(&self, data: &[u8], idx: u16) -> u16 {
let start_offset: usize = idx as usize * 2;
get_u16(data, start_offset)
}
#[inline]
fn num_vals(&self) -> u16 {
(self.0.len() / 2) as u16
}
#[inline]
#[allow(clippy::comparison_chain)]
// Looks for the element in the block. Returns the positions if found.
fn binary_search(&self, target: u16) -> Result<u16, u16> {
let data = &self.0;
let mut size = self.num_vals();
let mut left = 0;
let mut right = size;
// TODO try different implem.
// e.g. exponential search into binary search
while left < right {
let mid = left + size / 2;
// TODO do boundary check only once, and then use an
// unsafe `value_at_idx`
let mid_val = self.value_at_idx(data, mid);
if target > mid_val {
left = mid + 1;
} else if target < mid_val {
right = mid;
} else {
return Ok(mid);
}
size = right - left;
}
Err(left)
}
pub fn select_iter<'b>(
&self,
iter: impl Iterator<Item = u16> + 'b,
) -> impl Iterator<Item = u16> + 'b
where
Self: 'b,
{
iter.map(|codec_id| {
let offset = codec_id as usize * 2;
u16::from_le_bytes(self.0[offset..offset + 2].try_into().unwrap())
})
}
}

View File

@@ -1,112 +0,0 @@
use std::collections::HashMap;
use crate::column_index::optional_index::set_block::{
DenseBlockCodec, SparseBlockCodec, DENSE_BLOCK_NUM_BYTES,
};
use crate::column_index::optional_index::{Set, SetCodec};
fn test_set_helper<C: SetCodec<Item = u16>>(vals: &[u16]) -> usize {
let mut buffer = Vec::new();
C::serialize(vals.iter().copied(), &mut buffer).unwrap();
let tested_set = C::open(buffer.as_slice());
let hash_set: HashMap<C::Item, C::Item> = vals
.iter()
.copied()
.enumerate()
.map(|(ord, val)| (val, C::Item::try_from(ord).ok().unwrap()))
.collect();
for val in 0u16..=u16::MAX {
assert_eq!(tested_set.contains(val), hash_set.contains_key(&val));
assert_eq!(tested_set.rank_if_exists(val), hash_set.get(&val).copied());
}
for rank in 0..vals.len() {
assert_eq!(tested_set.select(rank as u16), vals[rank]);
}
buffer.len()
}
#[test]
fn test_dense_block_set_u16_empty() {
let buffer_len = test_set_helper::<DenseBlockCodec>(&[]);
assert_eq!(buffer_len, DENSE_BLOCK_NUM_BYTES as usize);
}
#[test]
fn test_dense_block_set_u16_max() {
let buffer_len = test_set_helper::<DenseBlockCodec>(&[u16::MAX]);
assert_eq!(buffer_len, DENSE_BLOCK_NUM_BYTES as usize);
}
#[test]
fn test_sparse_block_set_u16_empty() {
let buffer_len = test_set_helper::<SparseBlockCodec>(&[]);
assert_eq!(buffer_len, 0);
}
#[test]
fn test_sparse_block_set_u16_max() {
let buffer_len = test_set_helper::<SparseBlockCodec>(&[u16::MAX]);
assert_eq!(buffer_len, 2);
}
use proptest::prelude::*;
proptest! {
#![proptest_config(ProptestConfig::with_cases(1))]
#[test]
fn test_prop_test_dense(els in proptest::collection::btree_set(0..=u16::MAX, 0..=u16::MAX as usize)) {
let vals: Vec<u16> = els.into_iter().collect();
let buffer_len = test_set_helper::<DenseBlockCodec>(&vals);
assert_eq!(buffer_len, DENSE_BLOCK_NUM_BYTES as usize);
}
#[test]
fn test_prop_test_sparse(els in proptest::collection::btree_set(0..=u16::MAX, 0..=u16::MAX as usize)) {
let vals: Vec<u16> = els.into_iter().collect();
let buffer_len = test_set_helper::<SparseBlockCodec>(&vals);
assert_eq!(buffer_len, vals.len() * 2);
}
}
#[test]
fn test_simple_translate_codec_codec_idx_to_original_idx_dense() {
let mut buffer = Vec::new();
DenseBlockCodec::serialize([1, 3, 17, 32, 30_000, 30_001].iter().copied(), &mut buffer)
.unwrap();
let tested_set = DenseBlockCodec::open(buffer.as_slice());
assert!(tested_set.contains(1));
assert_eq!(
&tested_set
.select_iter([0, 1, 2, 5].iter().copied())
.collect::<Vec<u16>>(),
&[1, 3, 17, 30_001]
);
}
#[test]
fn test_simple_translate_codec_idx_to_original_idx_sparse() {
let mut buffer = Vec::new();
SparseBlockCodec::serialize([1, 3, 17].iter().copied(), &mut buffer).unwrap();
let tested_set = SparseBlockCodec::open(buffer.as_slice());
assert!(tested_set.contains(1));
assert_eq!(
&tested_set
.select_iter([0, 1, 2].iter().copied())
.collect::<Vec<u16>>(),
&[1, 3, 17]
);
}
#[test]
fn test_simple_translate_codec_idx_to_original_idx_dense() {
let mut buffer = Vec::new();
DenseBlockCodec::serialize(0u16..150u16, &mut buffer).unwrap();
let tested_set = DenseBlockCodec::open(buffer.as_slice());
assert!(tested_set.contains(1));
let rg = 0u16..150u16;
let els: Vec<u16> = rg.clone().collect();
assert_eq!(
&tested_set.select_iter(rg.clone()).collect::<Vec<u16>>(),
&els
);
}

View File

@@ -1,327 +0,0 @@
use proptest::prelude::{any, prop, *};
use proptest::strategy::Strategy;
use proptest::{prop_oneof, proptest};
use super::*;
#[test]
fn test_dense_block_threshold() {
assert_eq!(super::DENSE_BLOCK_THRESHOLD, 5_120);
}
fn random_bitvec() -> BoxedStrategy<Vec<bool>> {
prop_oneof![
1 => prop::collection::vec(proptest::bool::weighted(1.0), 0..100),
1 => prop::collection::vec(proptest::bool::weighted(0.00), 0..(ELEMENTS_PER_BLOCK as usize * 3)), // empty blocks
1 => prop::collection::vec(proptest::bool::weighted(1.00), 0..(ELEMENTS_PER_BLOCK as usize + 10)), // full block
1 => prop::collection::vec(proptest::bool::weighted(0.01), 0..100),
1 => prop::collection::vec(proptest::bool::weighted(0.01), 0..u16::MAX as usize),
8 => vec![any::<bool>()],
]
.boxed()
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(50))]
#[test]
fn test_with_random_bitvecs(bitvec1 in random_bitvec(), bitvec2 in random_bitvec(), bitvec3 in random_bitvec()) {
let mut bitvec = Vec::new();
bitvec.extend_from_slice(&bitvec1);
bitvec.extend_from_slice(&bitvec2);
bitvec.extend_from_slice(&bitvec3);
test_null_index(&bitvec[..]);
}
}
#[test]
fn test_with_random_sets_simple() {
let vals = 10..BLOCK_SIZE * 2;
let mut out: Vec<u8> = Vec::new();
serialize_optional_index(&vals.clone(), &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
let ranks: Vec<u32> = (65_472u32..65_473u32).collect();
let els: Vec<u32> = ranks.iter().copied().map(|rank| rank + 10).collect();
let mut output = vec![0u32; ranks.len()];
null_index.select_batch(&ranks[..], &mut output[..]);
assert_eq!(&output, &els);
}
#[test]
fn test_optional_index_trailing_empty_blocks() {
test_null_index(&[false]);
}
#[test]
fn test_optional_index_one_block_false() {
let mut iter = vec![false; ELEMENTS_PER_BLOCK as usize];
iter.push(true);
test_null_index(&iter[..]);
}
#[test]
fn test_optional_index_one_block_true() {
let mut iter = vec![true; ELEMENTS_PER_BLOCK as usize];
iter.push(true);
test_null_index(&iter[..]);
}
impl<'a> SerializableOptionalIndex<'a> for &'a [bool] {
fn num_rows(&self) -> RowId {
self.len() as u32
}
fn non_null_rows(&self) -> Box<dyn Iterator<Item = RowId> + 'a> {
Box::new(
self.iter()
.cloned()
.enumerate()
.filter(|(_pos, val)| *val)
.map(|(pos, _val)| pos as u32),
)
}
}
fn test_null_index(data: &[bool]) {
let mut out: Vec<u8> = Vec::new();
serialize_optional_index(&data, &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
let orig_idx_with_value: Vec<u32> = data
.iter()
.enumerate()
.filter(|(_pos, val)| **val)
.map(|(pos, _val)| pos as u32)
.collect();
let ids: Vec<u32> = (0..orig_idx_with_value.len() as u32).collect();
let mut output = vec![0u32; ids.len()];
null_index.select_batch(&ids[..], &mut output);
// assert_eq!(&output[0..100], &orig_idx_with_value[0..100]);
assert_eq!(output, orig_idx_with_value);
let step_size = (orig_idx_with_value.len() / 100).max(1);
for (dense_idx, orig_idx) in orig_idx_with_value.iter().enumerate().step_by(step_size) {
assert_eq!(null_index.rank_if_exists(*orig_idx), Some(dense_idx as u32));
}
// 100 samples
let step_size = (data.len() / 100).max(1);
for (pos, value) in data.iter().enumerate().step_by(step_size) {
assert_eq!(null_index.contains(pos as u32), *value);
}
}
#[test]
fn test_optional_index_test_translation() {
let mut out = vec![];
let iter = &[true, false, true, false];
serialize_optional_index(&&iter[..], &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
let mut output = vec![0u32; 2];
null_index.select_batch(&[0, 1], &mut output);
assert_eq!(output, &[0, 2]);
}
#[test]
fn test_optional_index_translate() {
let mut out = vec![];
let iter = &[true, false, true, false];
serialize_optional_index(&&iter[..], &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
assert_eq!(null_index.rank_if_exists(0), Some(0));
assert_eq!(null_index.rank_if_exists(2), Some(1));
}
#[test]
fn test_optional_index_small() {
let mut out = vec![];
let iter = &[true, false, true, false];
serialize_optional_index(&&iter[..], &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
assert!(null_index.contains(0));
assert!(!null_index.contains(1));
assert!(null_index.contains(2));
assert!(!null_index.contains(3));
}
#[test]
fn test_optional_index_large() {
let mut docs = vec![];
docs.extend((0..ELEMENTS_PER_BLOCK).map(|_idx| false));
docs.extend((0..=1).map(|_idx| true));
let mut out = vec![];
serialize_optional_index(&&docs[..], &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
assert!(!null_index.contains(0));
assert!(!null_index.contains(100));
assert!(!null_index.contains(ELEMENTS_PER_BLOCK - 1));
assert!(null_index.contains(ELEMENTS_PER_BLOCK));
assert!(null_index.contains(ELEMENTS_PER_BLOCK + 1));
}
#[cfg(all(test, feature = "unstable"))]
mod bench {
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use test::Bencher;
use super::*;
const TOTAL_NUM_VALUES: u32 = 1_000_000;
fn gen_bools(fill_ratio: f64) -> OptionalIndex {
let mut out = Vec::new();
let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
let vals: Vec<bool> = (0..TOTAL_NUM_VALUES)
.map(|_| rng.gen_bool(fill_ratio))
.collect();
serialize_optional_index(&&vals[..], &mut out).unwrap();
let codec = open_optional_index(OwnedBytes::new(out)).unwrap();
codec
}
fn random_range_iterator(
start: u32,
end: u32,
avg_step_size: u32,
avg_deviation: u32,
) -> impl Iterator<Item = u32> {
let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
let mut current = start;
std::iter::from_fn(move || {
current += rng.gen_range(avg_step_size - avg_deviation..=avg_step_size + avg_deviation);
if current >= end {
None
} else {
Some(current)
}
})
}
fn n_percent_step_iterator(percent: f32, num_values: u32) -> impl Iterator<Item = u32> {
let ratio = percent as f32 / 100.0;
let step_size = (1f32 / ratio) as u32;
let deviation = step_size - 1;
random_range_iterator(0, num_values, step_size, deviation)
}
fn walk_over_data(codec: &OptionalIndex, avg_step_size: u32) -> Option<u32> {
walk_over_data_from_positions(
codec,
random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, 0),
)
}
fn walk_over_data_from_positions(
codec: &OptionalIndex,
positions: impl Iterator<Item = u32>,
) -> Option<u32> {
let mut dense_idx: Option<u32> = None;
for idx in positions {
dense_idx = dense_idx.or(codec.rank_if_exists(idx));
}
dense_idx
}
#[bench]
fn bench_translate_orig_to_codec_1percent_filled_10percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_translate_orig_to_codec_5percent_filled_10percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.05f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_translate_orig_to_codec_5percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.05f64);
bench.iter(|| walk_over_data(&codec, 1000));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_1percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_10percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_90percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_10percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_translate_orig_to_codec_50percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.5f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_translate_orig_to_codec_90percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) {
bench_translate_codec_to_orig_util(0.01f64, 0.005f32, bench);
}
#[bench]
fn bench_translate_codec_to_orig_10percent_filled_0comma005percent_hit(bench: &mut Bencher) {
bench_translate_codec_to_orig_util(0.1f64, 0.005f32, bench);
}
#[bench]
fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
bench_translate_codec_to_orig_util(0.01f64, 10f32, bench);
}
#[bench]
fn bench_translate_codec_to_orig_1percent_filled_full_scan(bench: &mut Bencher) {
bench_translate_codec_to_orig_util(0.01f64, 100f32, bench);
}
fn bench_translate_codec_to_orig_util(
percent_filled: f64,
percent_hit: f32,
bench: &mut Bencher,
) {
let codec = gen_bools(percent_filled);
let num_non_nulls = codec.num_non_nulls();
let idxs: Vec<u32> = if percent_hit == 100.0f32 {
(0..num_non_nulls).collect()
} else {
n_percent_step_iterator(percent_hit, num_non_nulls).collect()
};
let mut output = vec![0u32; idxs.len()];
bench.iter(|| {
codec.select_batch(&idxs[..], &mut output);
});
}
#[bench]
fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) {
bench_translate_codec_to_orig_util(0.9f64, 0.005, bench);
}
#[bench]
fn bench_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
bench_translate_codec_to_orig_util(0.9f64, 100.0f32, bench);
}
}

View File

@@ -1,73 +0,0 @@
use std::io;
use std::io::Write;
use common::{CountingWriter, OwnedBytes};
use crate::column_index::multivalued_index::serialize_multivalued_index;
use crate::column_index::optional_index::serialize_optional_index;
use crate::column_index::{ColumnIndex, SerializableOptionalIndex};
use crate::column_values::ColumnValues;
use crate::{Cardinality, RowId};
pub enum SerializableColumnIndex<'a> {
Full,
Optional(Box<dyn SerializableOptionalIndex<'a> + 'a>),
// TODO remove the Arc<dyn> apart from serialization this is not
// dynamic at all.
Multivalued(Box<dyn ColumnValues<RowId> + 'a>),
}
impl<'a> SerializableColumnIndex<'a> {
pub fn get_cardinality(&self) -> Cardinality {
match self {
SerializableColumnIndex::Full => Cardinality::Full,
SerializableColumnIndex::Optional(_) => Cardinality::Optional,
SerializableColumnIndex::Multivalued(_) => Cardinality::Multivalued,
}
}
}
pub fn serialize_column_index(
column_index: SerializableColumnIndex,
output: &mut impl Write,
) -> io::Result<u32> {
let mut output = CountingWriter::wrap(output);
let cardinality = column_index.get_cardinality().to_code();
output.write_all(&[cardinality])?;
match column_index {
SerializableColumnIndex::Full => {}
SerializableColumnIndex::Optional(optional_index) => {
serialize_optional_index(&*optional_index, &mut output)?
}
SerializableColumnIndex::Multivalued(multivalued_index) => {
serialize_multivalued_index(&*multivalued_index, &mut output)?
}
}
let column_index_num_bytes = output.written_bytes() as u32;
Ok(column_index_num_bytes)
}
pub fn open_column_index(mut bytes: OwnedBytes) -> io::Result<ColumnIndex<'static>> {
if bytes.is_empty() {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"Failed to deserialize column index. Empty buffer.",
));
}
let cardinality_code = bytes[0];
let cardinality = Cardinality::try_from_code(cardinality_code)?;
bytes.advance(1);
match cardinality {
Cardinality::Full => Ok(ColumnIndex::Full),
Cardinality::Optional => {
let optional_index = super::optional_index::open_optional_index(bytes)?;
Ok(ColumnIndex::Optional(optional_index))
}
Cardinality::Multivalued => {
let multivalued_index = super::multivalued_index::open_multivalued_index(bytes)?;
Ok(ColumnIndex::Multivalued(multivalued_index))
}
}
}
// TODO unit tests

View File

@@ -0,0 +1,188 @@
use crate::utils::{place_bits, select_bits};
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)]
#[repr(u8)]
pub enum Cardinality {
/// All documents contain exactly one value.
#[default]
Required = 0,
/// All documents contain at most one value.
Optional = 1,
/// All documents may contain any number of values.
Multivalued = 2,
}
impl Cardinality {
pub(crate) fn to_code(self) -> u8 {
self as u8
}
pub(crate) fn try_from_code(code: u8) -> Option<Cardinality> {
match code {
0 => Some(Cardinality::Required),
1 => Some(Cardinality::Optional),
2 => Some(Cardinality::Multivalued),
_ => None,
}
}
}
#[derive(Hash, Eq, PartialEq, Debug, Clone, Copy)]
pub enum ColumnType {
Bytes,
Numerical(NumericalType),
Bool,
}
impl ColumnType {
/// Encoded over 6 bits.
pub(crate) fn to_code(self) -> u8 {
let high_type;
let low_code: u8;
match self {
ColumnType::Bytes => {
high_type = GeneralType::Str;
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> {
if select_bits::<6, 8>(code) != 0u8 {
return None;
}
let high_code = select_bits::<3, 6>(code);
let low_code = select_bits::<0, 3>(code);
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,
}
}
}
/// Represents the type and cardinality of a column.
/// This is encoded over one-byte and added to a column key in the
/// columnar sstable.
///
/// Cardinality is encoded as the first two highest two bits.
/// The low 6 bits encode the column type.
#[derive(Eq, Hash, PartialEq, Debug, Copy, Clone)]
pub struct ColumnTypeAndCardinality {
pub cardinality: Cardinality,
pub typ: ColumnType,
}
impl ColumnTypeAndCardinality {
pub fn to_code(self) -> u8 {
place_bits::<6, 8>(self.cardinality.to_code()) | place_bits::<0, 6>(self.typ.to_code())
}
pub fn try_from_code(code: u8) -> Option<ColumnTypeAndCardinality> {
let typ_code = select_bits::<0, 6>(code);
let cardinality_code = select_bits::<6, 8>(code);
let cardinality = Cardinality::try_from_code(cardinality_code)?;
let typ = ColumnType::try_from_code(typ_code)?;
assert_eq!(typ.to_code(), typ_code);
Some(ColumnTypeAndCardinality { cardinality, typ })
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::ColumnTypeAndCardinality;
use crate::column_type_header::{Cardinality, ColumnType};
#[test]
fn test_column_type_header_to_code() {
let mut column_type_header_set: HashSet<ColumnTypeAndCardinality> = HashSet::new();
for code in u8::MIN..=u8::MAX {
if let Some(column_type_header) = ColumnTypeAndCardinality::try_from_code(code) {
assert_eq!(column_type_header.to_code(), code);
assert!(column_type_header_set.insert(column_type_header));
}
}
assert_eq!(
column_type_header_set.len(),
3 /* cardinality */ *
(1 + 1 + 3) // column_types (str, bool, numerical x 3)
);
}
#[test]
fn test_column_type_to_code() {
let mut column_type_set: HashSet<ColumnType> = HashSet::new();
for code in u8::MIN..=u8::MAX {
if let Some(column_type) = ColumnType::try_from_code(code) {
assert_eq!(column_type.to_code(), code);
assert!(column_type_set.insert(column_type));
}
}
assert_eq!(column_type_set.len(), 2 + 3);
}
#[test]
fn test_cardinality_to_code() {
let mut num_cardinality = 0;
for code in u8::MIN..=u8::MAX {
let cardinality_opt = Cardinality::try_from_code(code);
if let Some(cardinality) = cardinality_opt {
assert_eq!(cardinality.to_code(), code);
num_cardinality += 1;
}
}
assert_eq!(num_cardinality, 3);
}
}

View File

@@ -1,115 +0,0 @@
use std::io::{self, Write};
use common::OwnedBytes;
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
use super::serialize::NormalizedHeader;
use super::{ColumnValues, FastFieldCodec, FastFieldCodecType};
/// Depending on the field type, a different
/// fast field is required.
#[derive(Clone)]
pub struct BitpackedReader {
data: OwnedBytes,
bit_unpacker: BitUnpacker,
normalized_header: NormalizedHeader,
}
impl ColumnValues for BitpackedReader {
#[inline]
fn get_val(&self, doc: u32) -> u64 {
self.bit_unpacker.get(doc, &self.data)
}
#[inline]
fn min_value(&self) -> u64 {
// The BitpackedReader assumes a normalized vector.
0
}
#[inline]
fn max_value(&self) -> u64 {
self.normalized_header.max_value
}
#[inline]
fn num_vals(&self) -> u32 {
self.normalized_header.num_vals
}
}
pub struct BitpackedCodec;
impl FastFieldCodec for BitpackedCodec {
/// The CODEC_TYPE is an enum value used for serialization.
const CODEC_TYPE: FastFieldCodecType = FastFieldCodecType::Bitpacked;
type Reader = BitpackedReader;
/// Opens a fast field given a file.
fn open_from_bytes(
data: OwnedBytes,
normalized_header: NormalizedHeader,
) -> io::Result<Self::Reader> {
let num_bits = compute_num_bits(normalized_header.max_value);
let bit_unpacker = BitUnpacker::new(num_bits);
Ok(BitpackedReader {
data,
bit_unpacker,
normalized_header,
})
}
/// Serializes data with the BitpackedFastFieldSerializer.
///
/// The bitpacker assumes that the column has been normalized.
/// i.e. It has already been shifted by its minimum value, so that its
/// current minimum value is 0.
///
/// Ideally, we made a shift upstream on the column so that `col.min_value() == 0`.
fn serialize(column: &dyn ColumnValues, write: &mut impl Write) -> io::Result<()> {
assert_eq!(column.min_value(), 0u64);
let num_bits = compute_num_bits(column.max_value());
let mut bit_packer = BitPacker::new();
for val in column.iter() {
bit_packer.write(val, num_bits, write)?;
}
bit_packer.close(write)?;
Ok(())
}
fn estimate(column: &dyn ColumnValues) -> Option<f32> {
let num_bits = compute_num_bits(column.max_value());
let num_bits_uncompressed = 64;
Some(num_bits as f32 / num_bits_uncompressed as f32)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::column_values::tests::create_and_validate;
fn create_and_validate_bitpacked_codec(data: &[u64], name: &str) {
create_and_validate::<BitpackedCodec>(data, name);
}
#[test]
fn test_with_codec_data_sets() {
let data_sets = crate::column_values::tests::get_codec_test_datasets();
for (mut data, name) in data_sets {
create_and_validate_bitpacked_codec(&data, name);
data.reverse();
create_and_validate::<BitpackedCodec>(&data, name);
}
}
#[test]
fn bitpacked_fast_field_rand() {
for _ in 0..500 {
let mut data = (0..1 + rand::random::<u8>() as usize)
.map(|_| rand::random::<i64>() as u64 / 2)
.collect::<Vec<_>>();
create_and_validate_bitpacked_codec(&data, "rand");
data.reverse();
create_and_validate::<BitpackedCodec>(&data, "rand");
}
}
}

View File

@@ -1,188 +0,0 @@
use std::sync::Arc;
use std::{io, iter};
use common::{BinarySerializable, CountingWriter, DeserializeFrom, OwnedBytes};
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
use crate::column_values::line::Line;
use crate::column_values::serialize::NormalizedHeader;
use crate::column_values::{ColumnValues, FastFieldCodec, FastFieldCodecType, VecColumn};
const CHUNK_SIZE: usize = 512;
#[derive(Debug, Default)]
struct Block {
line: Line,
bit_unpacker: BitUnpacker,
data_start_offset: usize,
}
impl BinarySerializable for Block {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
self.line.serialize(writer)?;
self.bit_unpacker.bit_width().serialize(writer)?;
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let line = Line::deserialize(reader)?;
let bit_width = u8::deserialize(reader)?;
Ok(Block {
line,
bit_unpacker: BitUnpacker::new(bit_width),
data_start_offset: 0,
})
}
}
fn compute_num_blocks(num_vals: u32) -> usize {
(num_vals as usize + CHUNK_SIZE - 1) / CHUNK_SIZE
}
pub struct BlockwiseLinearCodec;
impl FastFieldCodec for BlockwiseLinearCodec {
const CODEC_TYPE: FastFieldCodecType = FastFieldCodecType::BlockwiseLinear;
type Reader = BlockwiseLinearReader;
fn open_from_bytes(
bytes: common::OwnedBytes,
normalized_header: NormalizedHeader,
) -> io::Result<Self::Reader> {
let footer_len: u32 = (&bytes[bytes.len() - 4..]).deserialize()?;
let footer_offset = bytes.len() - 4 - footer_len as usize;
let (data, mut footer) = bytes.split(footer_offset);
let num_blocks = compute_num_blocks(normalized_header.num_vals);
let mut blocks: Vec<Block> = iter::repeat_with(|| Block::deserialize(&mut footer))
.take(num_blocks)
.collect::<io::Result<_>>()?;
let mut start_offset = 0;
for block in &mut blocks {
block.data_start_offset = start_offset;
start_offset += (block.bit_unpacker.bit_width() as usize) * CHUNK_SIZE / 8;
}
Ok(BlockwiseLinearReader {
blocks: Arc::new(blocks),
data,
normalized_header,
})
}
// Estimate first_chunk and extrapolate
fn estimate(column: &dyn ColumnValues) -> Option<f32> {
if column.num_vals() < 10 * CHUNK_SIZE as u32 {
return None;
}
let mut first_chunk: Vec<u64> = column.iter().take(CHUNK_SIZE).collect();
let line = Line::train(&VecColumn::from(&first_chunk));
for (i, buffer_val) in first_chunk.iter_mut().enumerate() {
let interpolated_val = line.eval(i as u32);
*buffer_val = buffer_val.wrapping_sub(interpolated_val);
}
let estimated_bit_width = first_chunk
.iter()
.map(|el| ((el + 1) as f32 * 3.0) as u64)
.map(compute_num_bits)
.max()
.unwrap();
let metadata_per_block = {
let mut out = vec![];
Block::default().serialize(&mut out).unwrap();
out.len()
};
let num_bits = estimated_bit_width as u64 * column.num_vals() as u64
// function metadata per block
+ metadata_per_block as u64 * (column.num_vals() as u64 / CHUNK_SIZE as u64);
let num_bits_uncompressed = 64 * column.num_vals();
Some(num_bits as f32 / num_bits_uncompressed as f32)
}
fn serialize(column: &dyn ColumnValues, wrt: &mut impl io::Write) -> io::Result<()> {
// The BitpackedReader assumes a normalized vector.
assert_eq!(column.min_value(), 0);
let mut buffer = Vec::with_capacity(CHUNK_SIZE);
let num_vals = column.num_vals();
let num_blocks = compute_num_blocks(num_vals);
let mut blocks = Vec::with_capacity(num_blocks);
let mut vals = column.iter();
let mut bit_packer = BitPacker::new();
for _ in 0..num_blocks {
buffer.clear();
buffer.extend((&mut vals).take(CHUNK_SIZE));
let line = Line::train(&VecColumn::from(&buffer));
assert!(!buffer.is_empty());
for (i, buffer_val) in buffer.iter_mut().enumerate() {
let interpolated_val = line.eval(i as u32);
*buffer_val = buffer_val.wrapping_sub(interpolated_val);
}
let bit_width = buffer.iter().copied().map(compute_num_bits).max().unwrap();
for &buffer_val in &buffer {
bit_packer.write(buffer_val, bit_width, wrt)?;
}
blocks.push(Block {
line,
bit_unpacker: BitUnpacker::new(bit_width),
data_start_offset: 0,
});
}
bit_packer.close(wrt)?;
assert_eq!(blocks.len(), compute_num_blocks(num_vals));
let mut counting_wrt = CountingWriter::wrap(wrt);
for block in &blocks {
block.serialize(&mut counting_wrt)?;
}
let footer_len = counting_wrt.written_bytes();
(footer_len as u32).serialize(&mut counting_wrt)?;
Ok(())
}
}
#[derive(Clone)]
pub struct BlockwiseLinearReader {
blocks: Arc<Vec<Block>>,
normalized_header: NormalizedHeader,
data: OwnedBytes,
}
impl ColumnValues for BlockwiseLinearReader {
#[inline(always)]
fn get_val(&self, idx: u32) -> u64 {
let block_id = (idx / CHUNK_SIZE as u32) as usize;
let idx_within_block = idx % (CHUNK_SIZE as u32);
let block = &self.blocks[block_id];
let interpoled_val: u64 = block.line.eval(idx_within_block);
let block_bytes = &self.data[block.data_start_offset..];
let bitpacked_diff = block.bit_unpacker.get(idx_within_block, block_bytes);
interpoled_val.wrapping_add(bitpacked_diff)
}
#[inline(always)]
fn min_value(&self) -> u64 {
// The BlockwiseLinearReader assumes a normalized vector.
0u64
}
#[inline(always)]
fn max_value(&self) -> u64 {
self.normalized_header.max_value
}
#[inline(always)]
fn num_vals(&self) -> u32 {
self.normalized_header.num_vals
}
}

View File

@@ -1,376 +0,0 @@
use std::marker::PhantomData;
use std::ops::{Range, RangeInclusive};
use tantivy_bitpacker::minmax;
use crate::column_values::monotonic_mapping::StrictlyMonotonicFn;
/// `ColumnValues` provides access to a dense field column.
///
/// `Column` are just a wrapper over `ColumnValues` and a `ColumnIndex`.
pub trait ColumnValues<T: PartialOrd = u64>: Send + Sync {
/// Return the value associated with the given idx.
///
/// This accessor should return as fast as possible.
///
/// # Panics
///
/// May panic if `idx` is greater than the column length.
fn get_val(&self, idx: u32) -> T;
/// Fills an output buffer with the fast field values
/// associated with the `DocId` going from
/// `start` to `start + output.len()`.
///
/// # Panics
///
/// Must panic if `start + output.len()` is greater than
/// the segment's `maxdoc`.
#[inline]
fn get_range(&self, start: u64, output: &mut [T]) {
for (out, idx) in output.iter_mut().zip(start..) {
*out = self.get_val(idx as u32);
}
}
/// Get the positions of values which are in the provided value range.
///
/// Note that position == docid for single value fast fields
#[inline]
fn get_docids_for_value_range(
&self,
value_range: RangeInclusive<T>,
doc_id_range: Range<u32>,
positions: &mut Vec<u32>,
) {
let doc_id_range = doc_id_range.start..doc_id_range.end.min(self.num_vals());
for idx in doc_id_range.start..doc_id_range.end {
let val = self.get_val(idx);
if value_range.contains(&val) {
positions.push(idx);
}
}
}
/// Returns the minimum value for this fast field.
///
/// This min_value may not be exact.
/// For instance, the min value does not take in account of possible
/// deleted document. All values are however guaranteed to be higher than
/// `.min_value()`.
fn min_value(&self) -> T;
/// Returns the maximum value for this fast field.
///
/// This max_value may not be exact.
/// For instance, the max value does not take in account of possible
/// deleted document. All values are however guaranteed to be higher than
/// `.max_value()`.
fn max_value(&self) -> T;
/// The number of values in the column.
fn num_vals(&self) -> u32;
/// Returns a iterator over the data
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = T> + 'a> {
Box::new((0..self.num_vals()).map(|idx| self.get_val(idx)))
}
}
impl<T: Copy + PartialOrd> ColumnValues<T> for std::sync::Arc<dyn ColumnValues<T>> {
fn get_val(&self, idx: u32) -> T {
self.as_ref().get_val(idx)
}
fn min_value(&self) -> T {
self.as_ref().min_value()
}
fn max_value(&self) -> T {
self.as_ref().max_value()
}
fn num_vals(&self) -> u32 {
self.as_ref().num_vals()
}
fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = T> + 'b> {
self.as_ref().iter()
}
fn get_range(&self, start: u64, output: &mut [T]) {
self.as_ref().get_range(start, output)
}
}
impl<'a, C: ColumnValues<T> + ?Sized, T: Copy + PartialOrd> ColumnValues<T> for &'a C {
fn get_val(&self, idx: u32) -> T {
(*self).get_val(idx)
}
fn min_value(&self) -> T {
(*self).min_value()
}
fn max_value(&self) -> T {
(*self).max_value()
}
fn num_vals(&self) -> u32 {
(*self).num_vals()
}
fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = T> + 'b> {
(*self).iter()
}
fn get_range(&self, start: u64, output: &mut [T]) {
(*self).get_range(start, output)
}
}
/// VecColumn provides `Column` over a slice.
pub struct VecColumn<'a, T = u64> {
pub(crate) values: &'a [T],
pub(crate) min_value: T,
pub(crate) max_value: T,
}
impl<'a, T: Copy + PartialOrd + Send + Sync> ColumnValues<T> for VecColumn<'a, T> {
fn get_val(&self, position: u32) -> T {
self.values[position as usize]
}
fn iter(&self) -> Box<dyn Iterator<Item = T> + '_> {
Box::new(self.values.iter().copied())
}
fn min_value(&self) -> T {
self.min_value
}
fn max_value(&self) -> T {
self.max_value
}
fn num_vals(&self) -> u32 {
self.values.len() as u32
}
fn get_range(&self, start: u64, output: &mut [T]) {
output.copy_from_slice(&self.values[start as usize..][..output.len()])
}
}
impl<'a, T: Copy + PartialOrd + Default, V> From<&'a V> for VecColumn<'a, T>
where V: AsRef<[T]> + ?Sized
{
fn from(values: &'a V) -> Self {
let values = values.as_ref();
let (min_value, max_value) = minmax(values.iter().copied()).unwrap_or_default();
Self {
values,
min_value,
max_value,
}
}
}
struct MonotonicMappingColumn<C, T, Input> {
from_column: C,
monotonic_mapping: T,
_phantom: PhantomData<Input>,
}
/// Creates a view of a column transformed by a strictly monotonic mapping. See
/// [`StrictlyMonotonicFn`].
///
/// E.g. apply a gcd monotonic_mapping([100, 200, 300]) == [1, 2, 3]
/// monotonic_mapping.mapping() is expected to be injective, and we should always have
/// monotonic_mapping.inverse(monotonic_mapping.mapping(el)) == el
///
/// The inverse of the mapping is required for:
/// `fn get_positions_for_value_range(&self, range: RangeInclusive<T>) -> Vec<u64> `
/// The user provides the original value range and we need to monotonic map them in the same way the
/// serialization does before calling the underlying column.
///
/// Note that when opening a codec, the monotonic_mapping should be the inverse of the mapping
/// during serialization. And therefore the monotonic_mapping_inv when opening is the same as
/// monotonic_mapping during serialization.
pub fn monotonic_map_column<C, T, Input, Output>(
from_column: C,
monotonic_mapping: T,
) -> impl ColumnValues<Output>
where
C: ColumnValues<Input>,
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
Input: PartialOrd + Send + Sync + Clone,
Output: PartialOrd + Send + Sync + Clone,
{
MonotonicMappingColumn {
from_column,
monotonic_mapping,
_phantom: PhantomData,
}
}
impl<C, T, Input, Output> ColumnValues<Output> for MonotonicMappingColumn<C, T, Input>
where
C: ColumnValues<Input>,
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
Input: PartialOrd + Send + Sync + Clone,
Output: PartialOrd + Send + Sync + Clone,
{
#[inline]
fn get_val(&self, idx: u32) -> Output {
let from_val = self.from_column.get_val(idx);
self.monotonic_mapping.mapping(from_val)
}
fn min_value(&self) -> Output {
let from_min_value = self.from_column.min_value();
self.monotonic_mapping.mapping(from_min_value)
}
fn max_value(&self) -> Output {
let from_max_value = self.from_column.max_value();
self.monotonic_mapping.mapping(from_max_value)
}
fn num_vals(&self) -> u32 {
self.from_column.num_vals()
}
fn iter(&self) -> Box<dyn Iterator<Item = Output> + '_> {
Box::new(
self.from_column
.iter()
.map(|el| self.monotonic_mapping.mapping(el)),
)
}
fn get_docids_for_value_range(
&self,
range: RangeInclusive<Output>,
doc_id_range: Range<u32>,
positions: &mut Vec<u32>,
) {
self.from_column.get_docids_for_value_range(
self.monotonic_mapping.inverse(range.start().clone())
..=self.monotonic_mapping.inverse(range.end().clone()),
doc_id_range,
positions,
)
}
// We voluntarily do not implement get_range as it yields a regression,
// and we do not have any specialized implementation anyway.
}
/// Wraps an iterator into a `Column`.
pub struct IterColumn<T>(T);
impl<T> From<T> for IterColumn<T>
where T: Iterator + Clone + ExactSizeIterator
{
fn from(iter: T) -> Self {
IterColumn(iter)
}
}
impl<T> ColumnValues<T::Item> for IterColumn<T>
where
T: Iterator + Clone + ExactSizeIterator + Send + Sync,
T::Item: PartialOrd,
{
fn get_val(&self, idx: u32) -> T::Item {
self.0.clone().nth(idx as usize).unwrap()
}
fn min_value(&self) -> T::Item {
self.0.clone().next().unwrap()
}
fn max_value(&self) -> T::Item {
self.0.clone().last().unwrap()
}
fn num_vals(&self) -> u32 {
self.0.len() as u32
}
fn iter(&self) -> Box<dyn Iterator<Item = T::Item> + '_> {
Box::new(self.0.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::column_values::monotonic_mapping::{
StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternalBaseval,
StrictlyMonotonicMappingToInternalGCDBaseval,
};
#[test]
fn test_monotonic_mapping() {
let vals = &[3u64, 5u64][..];
let col = VecColumn::from(vals);
let mapped = monotonic_map_column(col, StrictlyMonotonicMappingToInternalBaseval::new(2));
assert_eq!(mapped.min_value(), 1u64);
assert_eq!(mapped.max_value(), 3u64);
assert_eq!(mapped.num_vals(), 2);
assert_eq!(mapped.num_vals(), 2);
assert_eq!(mapped.get_val(0), 1);
assert_eq!(mapped.get_val(1), 3);
}
#[test]
fn test_range_as_col() {
let col = IterColumn::from(10..100);
assert_eq!(col.num_vals(), 90);
assert_eq!(col.max_value(), 99);
}
#[test]
fn test_monotonic_mapping_iter() {
let vals: Vec<u64> = (10..110u64).map(|el| el * 10).collect();
let col = VecColumn::from(&vals);
let mapped = monotonic_map_column(
col,
StrictlyMonotonicMappingInverter::from(
StrictlyMonotonicMappingToInternalGCDBaseval::new(10, 100),
),
);
let val_i64s: Vec<u64> = mapped.iter().collect();
for i in 0..100 {
assert_eq!(val_i64s[i as usize], mapped.get_val(i));
}
}
#[test]
fn test_monotonic_mapping_get_range() {
let vals: Vec<u64> = (0..100u64).map(|el| el * 10).collect();
let col = VecColumn::from(&vals);
let mapped = monotonic_map_column(
col,
StrictlyMonotonicMappingInverter::from(
StrictlyMonotonicMappingToInternalGCDBaseval::new(10, 0),
),
);
assert_eq!(mapped.min_value(), 0u64);
assert_eq!(mapped.max_value(), 9900u64);
assert_eq!(mapped.num_vals(), 100);
let val_u64s: Vec<u64> = mapped.iter().collect();
assert_eq!(val_u64s.len(), 100);
for i in 0..100 {
assert_eq!(val_u64s[i as usize], mapped.get_val(i));
assert_eq!(val_u64s[i as usize], vals[i as usize] * 10);
}
let mut buf = [0u64; 20];
mapped.get_range(7, &mut buf[..]);
assert_eq!(&val_u64s[7..][..20], &buf);
}
}

View File

@@ -1,19 +0,0 @@
// 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/>.
//

View File

@@ -1,43 +0,0 @@
use std::ops::RangeInclusive;
/// The range of a blank in value space.
///
/// A blank is an unoccupied space in the data.
/// Use try_into() to construct.
/// A range has to have at least length of 3. Invalid ranges will be rejected.
///
/// Ordered by range length.
#[derive(Debug, Eq, PartialEq, Clone)]
pub(crate) struct BlankRange {
blank_range: RangeInclusive<u128>,
}
impl TryFrom<RangeInclusive<u128>> for BlankRange {
type Error = &'static str;
fn try_from(range: RangeInclusive<u128>) -> Result<Self, Self::Error> {
let blank_size = range.end().saturating_sub(*range.start());
if blank_size < 2 {
Err("invalid range")
} else {
Ok(BlankRange { blank_range: range })
}
}
}
impl BlankRange {
pub(crate) fn blank_size(&self) -> u128 {
self.blank_range.end() - self.blank_range.start() + 1
}
pub(crate) fn blank_range(&self) -> RangeInclusive<u128> {
self.blank_range.clone()
}
}
impl Ord for BlankRange {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.blank_size().cmp(&other.blank_size())
}
}
impl PartialOrd for BlankRange {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.blank_size().cmp(&other.blank_size()))
}
}

View File

@@ -1,231 +0,0 @@
use std::collections::{BTreeSet, BinaryHeap};
use std::iter;
use std::ops::RangeInclusive;
use itertools::Itertools;
use super::blank_range::BlankRange;
use super::{CompactSpace, RangeMapping};
/// Put the blanks for the sorted values into a binary heap
fn get_blanks(values_sorted: &BTreeSet<u128>) -> BinaryHeap<BlankRange> {
let mut blanks: BinaryHeap<BlankRange> = BinaryHeap::new();
for (first, second) in values_sorted.iter().tuple_windows() {
// Correctness Overflow: the values are deduped and sorted (BTreeSet property), that means
// there's always space between two values.
let blank_range = first + 1..=second - 1;
let blank_range: Result<BlankRange, _> = blank_range.try_into();
if let Ok(blank_range) = blank_range {
blanks.push(blank_range);
}
}
blanks
}
struct BlankCollector {
blanks: Vec<BlankRange>,
staged_blanks_sum: u128,
}
impl BlankCollector {
fn new() -> Self {
Self {
blanks: vec![],
staged_blanks_sum: 0,
}
}
fn stage_blank(&mut self, blank: BlankRange) {
self.staged_blanks_sum += blank.blank_size();
self.blanks.push(blank);
}
fn drain(&mut self) -> impl Iterator<Item = BlankRange> + '_ {
self.staged_blanks_sum = 0;
self.blanks.drain(..)
}
fn staged_blanks_sum(&self) -> u128 {
self.staged_blanks_sum
}
fn num_staged_blanks(&self) -> usize {
self.blanks.len()
}
}
fn num_bits(val: u128) -> u8 {
(128u32 - val.leading_zeros()) as u8
}
/// Will collect blanks and add them to compact space if more bits are saved than cost from
/// metadata.
pub fn get_compact_space(
values_deduped_sorted: &BTreeSet<u128>,
total_num_values: u32,
cost_per_blank: usize,
) -> CompactSpace {
let mut compact_space_builder = CompactSpaceBuilder::new();
if values_deduped_sorted.is_empty() {
return compact_space_builder.finish();
}
let mut blanks: BinaryHeap<BlankRange> = get_blanks(values_deduped_sorted);
// Replace after stabilization of https://github.com/rust-lang/rust/issues/62924
// We start by space that's limited to min_value..=max_value
let min_value = *values_deduped_sorted.iter().next().unwrap_or(&0);
let max_value = *values_deduped_sorted.iter().last().unwrap_or(&0);
// +1 for null, in case min and max covers the whole space, we are off by one.
let mut amplitude_compact_space = (max_value - min_value).saturating_add(1);
if min_value != 0 {
compact_space_builder.add_blanks(iter::once(0..=min_value - 1));
}
if max_value != u128::MAX {
compact_space_builder.add_blanks(iter::once(max_value + 1..=u128::MAX));
}
let mut amplitude_bits: u8 = num_bits(amplitude_compact_space);
let mut blank_collector = BlankCollector::new();
// We will stage blanks until they reduce the compact space by at least 1 bit and then flush
// them if the metadata cost is lower than the total number of saved bits.
// Binary heap to process the gaps by their size
while let Some(blank_range) = blanks.pop() {
blank_collector.stage_blank(blank_range);
let staged_spaces_sum: u128 = blank_collector.staged_blanks_sum();
let amplitude_new_compact_space = amplitude_compact_space - staged_spaces_sum;
let amplitude_new_bits = num_bits(amplitude_new_compact_space);
if amplitude_bits == amplitude_new_bits {
continue;
}
let saved_bits = (amplitude_bits - amplitude_new_bits) as usize * total_num_values as usize;
// TODO: Maybe calculate exact cost of blanks and run this more expensive computation only,
// when amplitude_new_bits changes
let cost = blank_collector.num_staged_blanks() * cost_per_blank;
if cost >= saved_bits {
// Continue here, since although we walk over the blanks by size,
// we can potentially save a lot at the last bits, which are smaller blanks
//
// E.g. if the first range reduces the compact space by 1000 from 2000 to 1000, which
// saves 11-10=1 bit and the next range reduces the compact space by 950 to
// 50, which saves 10-6=4 bit
continue;
}
amplitude_compact_space = amplitude_new_compact_space;
amplitude_bits = amplitude_new_bits;
compact_space_builder.add_blanks(blank_collector.drain().map(|blank| blank.blank_range()));
}
// special case, when we don't collected any blanks because:
// * the data is empty (early exit)
// * the algorithm did decide it's not worth the cost, which can be the case for single values
//
// We drain one collected blank unconditionally, so the empty case is reserved for empty
// data, and therefore empty compact_space means the data is empty and no data is covered
// (conversely to all data) and we can assign null to it.
if compact_space_builder.is_empty() {
compact_space_builder.add_blanks(
blank_collector
.drain()
.map(|blank| blank.blank_range())
.take(1),
);
}
let compact_space = compact_space_builder.finish();
if max_value - min_value != u128::MAX {
debug_assert_eq!(
compact_space.amplitude_compact_space(),
amplitude_compact_space
);
}
compact_space
}
#[derive(Debug, Clone, Eq, PartialEq)]
struct CompactSpaceBuilder {
blanks: Vec<RangeInclusive<u128>>,
}
impl CompactSpaceBuilder {
/// Creates a new compact space builder which will initially cover the whole space.
fn new() -> Self {
Self { blanks: Vec::new() }
}
/// Assumes that repeated add_blank calls don't overlap and are not adjacent,
/// e.g. [3..=5, 5..=10] is not allowed
///
/// Both of those assumptions are true when blanks are produced from sorted values.
fn add_blanks(&mut self, blank: impl Iterator<Item = RangeInclusive<u128>>) {
self.blanks.extend(blank);
}
fn is_empty(&self) -> bool {
self.blanks.is_empty()
}
/// Convert blanks to covered space and assign null value
fn finish(mut self) -> CompactSpace {
// sort by start. ranges are not allowed to overlap
self.blanks.sort_unstable_by_key(|blank| *blank.start());
let mut covered_space = Vec::with_capacity(self.blanks.len());
// begining of the blanks
if let Some(first_blank_start) = self.blanks.first().map(RangeInclusive::start) {
if *first_blank_start != 0 {
covered_space.push(0..=first_blank_start - 1);
}
}
// Between the blanks
let between_blanks = self.blanks.iter().tuple_windows().map(|(left, right)| {
assert!(
left.end() < right.start(),
"overlapping or adjacent ranges detected"
);
*left.end() + 1..=*right.start() - 1
});
covered_space.extend(between_blanks);
// end of the blanks
if let Some(last_blank_end) = self.blanks.last().map(RangeInclusive::end) {
if *last_blank_end != u128::MAX {
covered_space.push(last_blank_end + 1..=u128::MAX);
}
}
if covered_space.is_empty() {
covered_space.push(0..=0); // empty data case
};
let mut compact_start: u64 = 1; // 0 is reserved for `null`
let mut ranges_mapping: Vec<RangeMapping> = Vec::with_capacity(covered_space.len());
for cov in covered_space {
let range_mapping = super::RangeMapping {
value_range: cov,
compact_start,
};
let covered_range_len = range_mapping.range_length();
ranges_mapping.push(range_mapping);
compact_start += covered_range_len;
}
// println!("num ranges {}", ranges_mapping.len());
CompactSpace { ranges_mapping }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_binary_heap_pop_order() {
let mut blanks: BinaryHeap<BlankRange> = BinaryHeap::new();
blanks.push((0..=10).try_into().unwrap());
blanks.push((100..=200).try_into().unwrap());
blanks.push((100..=110).try_into().unwrap());
assert_eq!(blanks.pop().unwrap().blank_size(), 101);
assert_eq!(blanks.pop().unwrap().blank_size(), 11);
}
}

View File

@@ -1,813 +0,0 @@
/// This codec takes a large number space (u128) and reduces it to a compact number space.
///
/// It will find spaces in the number range. For example:
///
/// 100, 101, 102, 103, 104, 50000, 50001
/// could be mapped to
/// 100..104 -> 0..4
/// 50000..50001 -> 5..6
///
/// Compact space 0..=6 requires much less bits than 100..=50001
///
/// The codec is created to compress ip addresses, but may be employed in other use cases.
use std::{
cmp::Ordering,
collections::BTreeSet,
io::{self, Write},
ops::{Range, RangeInclusive},
};
use common::{BinarySerializable, CountingWriter, OwnedBytes, VInt, VIntU128};
use tantivy_bitpacker::{self, BitPacker, BitUnpacker};
use crate::column_values::compact_space::build_compact_space::get_compact_space;
use crate::column_values::ColumnValues;
mod blank_range;
mod build_compact_space;
/// The cost per blank is quite hard actually, since blanks are delta encoded, the actual cost of
/// blanks depends on the number of blanks.
///
/// The number is taken by looking at a real dataset. It is optimized for larger datasets.
const COST_PER_BLANK_IN_BITS: usize = 36;
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct CompactSpace {
ranges_mapping: Vec<RangeMapping>,
}
/// Maps the range from the original space to compact_start + range.len()
#[derive(Debug, Clone, Eq, PartialEq)]
struct RangeMapping {
value_range: RangeInclusive<u128>,
compact_start: u64,
}
impl RangeMapping {
fn range_length(&self) -> u64 {
(self.value_range.end() - self.value_range.start()) as u64 + 1
}
// The last value of the compact space in this range
fn compact_end(&self) -> u64 {
self.compact_start + self.range_length() - 1
}
}
impl BinarySerializable for CompactSpace {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
VInt(self.ranges_mapping.len() as u64).serialize(writer)?;
let mut prev_value = 0;
for value_range in self
.ranges_mapping
.iter()
.map(|range_mapping| &range_mapping.value_range)
{
let blank_delta_start = value_range.start() - prev_value;
VIntU128(blank_delta_start).serialize(writer)?;
prev_value = *value_range.start();
let blank_delta_end = value_range.end() - prev_value;
VIntU128(blank_delta_end).serialize(writer)?;
prev_value = *value_range.end();
}
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let num_ranges = VInt::deserialize(reader)?.0;
let mut ranges_mapping: Vec<RangeMapping> = vec![];
let mut value = 0u128;
let mut compact_start = 1u64; // 0 is reserved for `null`
for _ in 0..num_ranges {
let blank_delta_start = VIntU128::deserialize(reader)?.0;
value += blank_delta_start;
let blank_start = value;
let blank_delta_end = VIntU128::deserialize(reader)?.0;
value += blank_delta_end;
let blank_end = value;
let range_mapping = RangeMapping {
value_range: blank_start..=blank_end,
compact_start,
};
let range_length = range_mapping.range_length();
ranges_mapping.push(range_mapping);
compact_start += range_length;
}
Ok(Self { ranges_mapping })
}
}
impl CompactSpace {
/// Amplitude is the value range of the compact space including the sentinel value used to
/// identify null values. The compact space is 0..=amplitude .
///
/// It's only used to verify we don't exceed u64 number space, which would indicate a bug.
fn amplitude_compact_space(&self) -> u128 {
self.ranges_mapping
.last()
.map(|last_range| last_range.compact_end() as u128)
.unwrap_or(1) // compact space starts at 1, 0 == null
}
fn get_range_mapping(&self, pos: usize) -> &RangeMapping {
&self.ranges_mapping[pos]
}
/// Returns either Ok(the value in the compact space) or if it is outside the compact space the
/// Err(position where it would be inserted)
fn u128_to_compact(&self, value: u128) -> Result<u64, usize> {
self.ranges_mapping
.binary_search_by(|probe| {
let value_range = &probe.value_range;
if value < *value_range.start() {
Ordering::Greater
} else if value > *value_range.end() {
Ordering::Less
} else {
Ordering::Equal
}
})
.map(|pos| {
let range_mapping = &self.ranges_mapping[pos];
let pos_in_range = (value - range_mapping.value_range.start()) as u64;
range_mapping.compact_start + pos_in_range
})
}
/// Unpacks a value from compact space u64 to u128 space
fn compact_to_u128(&self, compact: u64) -> u128 {
let pos = self
.ranges_mapping
.binary_search_by_key(&compact, |range_mapping| range_mapping.compact_start)
// Correctness: Overflow. The first range starts at compact space 0, the error from
// binary search can never be 0
.map_or_else(|e| e - 1, |v| v);
let range_mapping = &self.ranges_mapping[pos];
let diff = compact - range_mapping.compact_start;
range_mapping.value_range.start() + diff as u128
}
}
pub struct CompactSpaceCompressor {
params: IPCodecParams,
}
#[derive(Debug, Clone)]
pub struct IPCodecParams {
compact_space: CompactSpace,
bit_unpacker: BitUnpacker,
min_value: u128,
max_value: u128,
num_vals: u32,
num_bits: u8,
}
impl CompactSpaceCompressor {
/// Taking the vals as Vec may cost a lot of memory. It is used to sort the vals.
pub fn train_from(iter: impl Iterator<Item = u128>, num_vals: u32) -> Self {
let mut values_sorted = BTreeSet::new();
values_sorted.extend(iter);
let total_num_values = num_vals;
let compact_space =
get_compact_space(&values_sorted, total_num_values, COST_PER_BLANK_IN_BITS);
let amplitude_compact_space = compact_space.amplitude_compact_space();
assert!(
amplitude_compact_space <= u64::MAX as u128,
"case unsupported."
);
let num_bits = tantivy_bitpacker::compute_num_bits(amplitude_compact_space as u64);
let min_value = *values_sorted.iter().next().unwrap_or(&0);
let max_value = *values_sorted.iter().last().unwrap_or(&0);
assert_eq!(
compact_space
.u128_to_compact(max_value)
.expect("could not convert max value to compact space"),
amplitude_compact_space as u64
);
CompactSpaceCompressor {
params: IPCodecParams {
compact_space,
bit_unpacker: BitUnpacker::new(num_bits),
min_value,
max_value,
num_vals: total_num_values,
num_bits,
},
}
}
fn write_footer(self, writer: &mut impl Write) -> io::Result<()> {
let writer = &mut CountingWriter::wrap(writer);
self.params.serialize(writer)?;
let footer_len = writer.written_bytes() as u32;
footer_len.serialize(writer)?;
Ok(())
}
pub fn compress_into(
self,
vals: impl Iterator<Item = u128>,
write: &mut impl Write,
) -> io::Result<()> {
let mut bitpacker = BitPacker::default();
for val in vals {
let compact = self
.params
.compact_space
.u128_to_compact(val)
.map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidData,
"Could not convert value to compact_space. This is a bug.",
)
})?;
bitpacker.write(compact, self.params.num_bits, write)?;
}
bitpacker.close(write)?;
self.write_footer(write)?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct CompactSpaceDecompressor {
data: OwnedBytes,
params: IPCodecParams,
}
impl BinarySerializable for IPCodecParams {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
// header flags for future optional dictionary encoding
let footer_flags = 0u64;
footer_flags.serialize(writer)?;
VIntU128(self.min_value).serialize(writer)?;
VIntU128(self.max_value).serialize(writer)?;
VIntU128(self.num_vals as u128).serialize(writer)?;
self.num_bits.serialize(writer)?;
self.compact_space.serialize(writer)?;
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let _header_flags = u64::deserialize(reader)?;
let min_value = VIntU128::deserialize(reader)?.0;
let max_value = VIntU128::deserialize(reader)?.0;
let num_vals = VIntU128::deserialize(reader)?.0 as u32;
let num_bits = u8::deserialize(reader)?;
let compact_space = CompactSpace::deserialize(reader)?;
Ok(Self {
compact_space,
bit_unpacker: BitUnpacker::new(num_bits),
min_value,
max_value,
num_vals,
num_bits,
})
}
}
impl ColumnValues<u128> for CompactSpaceDecompressor {
#[inline]
fn get_val(&self, doc: u32) -> u128 {
self.get(doc)
}
fn min_value(&self) -> u128 {
self.min_value()
}
fn max_value(&self) -> u128 {
self.max_value()
}
fn num_vals(&self) -> u32 {
self.params.num_vals
}
#[inline]
fn iter(&self) -> Box<dyn Iterator<Item = u128> + '_> {
Box::new(self.iter())
}
#[inline]
fn get_docids_for_value_range(
&self,
value_range: RangeInclusive<u128>,
positions_range: Range<u32>,
positions: &mut Vec<u32>,
) {
self.get_positions_for_value_range(value_range, positions_range, positions)
}
}
impl CompactSpaceDecompressor {
pub fn open(data: OwnedBytes) -> io::Result<CompactSpaceDecompressor> {
let (data_slice, footer_len_bytes) = data.split_at(data.len() - 4);
let footer_len = u32::deserialize(&mut &footer_len_bytes[..])?;
let data_footer = &data_slice[data_slice.len() - footer_len as usize..];
let params = IPCodecParams::deserialize(&mut &data_footer[..])?;
let decompressor = CompactSpaceDecompressor { data, params };
Ok(decompressor)
}
/// Converting to compact space for the decompressor is more complex, since we may get values
/// which are outside the compact space. e.g. if we map
/// 1000 => 5
/// 2000 => 6
///
/// and we want a mapping for 1005, there is no equivalent compact space. We instead return an
/// error with the index of the next range.
fn u128_to_compact(&self, value: u128) -> Result<u64, usize> {
self.params.compact_space.u128_to_compact(value)
}
fn compact_to_u128(&self, compact: u64) -> u128 {
self.params.compact_space.compact_to_u128(compact)
}
/// Comparing on compact space: Random dataset 0,24 (50% random hit) - 1.05 GElements/s
/// Comparing on compact space: Real dataset 1.08 GElements/s
///
/// Comparing on original space: Real dataset .06 GElements/s (not completely optimized)
#[inline]
pub fn get_positions_for_value_range(
&self,
value_range: RangeInclusive<u128>,
position_range: Range<u32>,
positions: &mut Vec<u32>,
) {
if value_range.start() > value_range.end() {
return;
}
let position_range = position_range.start..position_range.end.min(self.num_vals());
let from_value = *value_range.start();
let to_value = *value_range.end();
assert!(to_value >= from_value);
let compact_from = self.u128_to_compact(from_value);
let compact_to = self.u128_to_compact(to_value);
// Quick return, if both ranges fall into the same non-mapped space, the range can't cover
// any values, so we can early exit
match (compact_to, compact_from) {
(Err(pos1), Err(pos2)) if pos1 == pos2 => return,
_ => {}
}
let compact_from = compact_from.unwrap_or_else(|pos| {
// Correctness: Out of bounds, if this value is Err(last_index + 1), we early exit,
// since the to_value also mapps into the same non-mapped space
let range_mapping = self.params.compact_space.get_range_mapping(pos);
range_mapping.compact_start
});
// If there is no compact space, we go to the closest upperbound compact space
let compact_to = compact_to.unwrap_or_else(|pos| {
// Correctness: Overflow, if this value is Err(0), we early exit,
// since the from_value also mapps into the same non-mapped space
// Get end of previous range
let pos = pos - 1;
let range_mapping = self.params.compact_space.get_range_mapping(pos);
range_mapping.compact_end()
});
let range = compact_from..=compact_to;
let scan_num_docs = position_range.end - position_range.start;
let step_size = 4;
let cutoff = position_range.start + scan_num_docs - scan_num_docs % step_size;
let mut push_if_in_range = |idx, val| {
if range.contains(&val) {
positions.push(idx);
}
};
let get_val = |idx| self.params.bit_unpacker.get(idx, &self.data);
// unrolled loop
for idx in (position_range.start..cutoff).step_by(step_size as usize) {
let idx1 = idx;
let idx2 = idx + 1;
let idx3 = idx + 2;
let idx4 = idx + 3;
let val1 = get_val(idx1);
let val2 = get_val(idx2);
let val3 = get_val(idx3);
let val4 = get_val(idx4);
push_if_in_range(idx1, val1);
push_if_in_range(idx2, val2);
push_if_in_range(idx3, val3);
push_if_in_range(idx4, val4);
}
// handle rest
for idx in cutoff..position_range.end {
push_if_in_range(idx, get_val(idx));
}
}
#[inline]
fn iter_compact(&self) -> impl Iterator<Item = u64> + '_ {
(0..self.params.num_vals).map(move |idx| self.params.bit_unpacker.get(idx, &self.data))
}
#[inline]
fn iter(&self) -> impl Iterator<Item = u128> + '_ {
// TODO: Performance. It would be better to iterate on the ranges and check existence via
// the bit_unpacker.
self.iter_compact()
.map(|compact| self.compact_to_u128(compact))
}
#[inline]
pub fn get(&self, idx: u32) -> u128 {
let compact = self.params.bit_unpacker.get(idx, &self.data);
self.compact_to_u128(compact)
}
pub fn min_value(&self) -> u128 {
self.params.min_value
}
pub fn max_value(&self) -> u128 {
self.params.max_value
}
}
// TODO reenable what can be reenabled.
// #[cfg(test)]
// mod tests {
//
// use super::*;
// use crate::column::format_version::read_format_version;
// use crate::column::column_footer::read_null_index_footer;
// use crate::column::serialize::U128Header;
// use crate::column::{open_u128, serialize_u128};
//
// #[test]
// fn compact_space_test() {
// let ips = &[
// 2u128, 4u128, 1000, 1001, 1002, 1003, 1004, 1005, 1008, 1010, 1012, 1260,
// ]
// .into_iter()
// .collect();
// let compact_space = get_compact_space(ips, ips.len() as u32, 11);
// let amplitude = compact_space.amplitude_compact_space();
// assert_eq!(amplitude, 17);
// assert_eq!(1, compact_space.u128_to_compact(2).unwrap());
// assert_eq!(2, compact_space.u128_to_compact(3).unwrap());
// assert_eq!(compact_space.u128_to_compact(100).unwrap_err(), 1);
//
// for (num1, num2) in (0..3).tuple_windows() {
// assert_eq!(
// compact_space.get_range_mapping(num1).compact_end() + 1,
// compact_space.get_range_mapping(num2).compact_start
// );
// }
//
// let mut output: Vec<u8> = Vec::new();
// compact_space.serialize(&mut output).unwrap();
//
// assert_eq!(
// compact_space,
// CompactSpace::deserialize(&mut &output[..]).unwrap()
// );
//
// for ip in ips {
// let compact = compact_space.u128_to_compact(*ip).unwrap();
// assert_eq!(compact_space.compact_to_u128(compact), *ip);
// }
// }
//
// #[test]
// fn compact_space_amplitude_test() {
// let ips = &[100000u128, 1000000].into_iter().collect();
// let compact_space = get_compact_space(ips, ips.len() as u32, 1);
// let amplitude = compact_space.amplitude_compact_space();
// assert_eq!(amplitude, 2);
// }
//
// fn test_all(mut data: OwnedBytes, expected: &[u128]) {
// let _header = U128Header::deserialize(&mut data);
// let decompressor = CompactSpaceDecompressor::open(data).unwrap();
// for (idx, expected_val) in expected.iter().cloned().enumerate() {
// let val = decompressor.get(idx as u32);
// assert_eq!(val, expected_val);
//
// let test_range = |range: RangeInclusive<u128>| {
// let expected_positions = expected
// .iter()
// .positions(|val| range.contains(val))
// .map(|pos| pos as u32)
// .collect::<Vec<_>>();
// let mut positions = Vec::new();
// decompressor.get_positions_for_value_range(
// range,
// 0..decompressor.num_vals(),
// &mut positions,
// );
// assert_eq!(positions, expected_positions);
// };
//
// test_range(expected_val.saturating_sub(1)..=expected_val);
// test_range(expected_val..=expected_val);
// test_range(expected_val..=expected_val.saturating_add(1));
// test_range(expected_val.saturating_sub(1)..=expected_val.saturating_add(1));
// }
// }
//
// fn test_aux_vals(u128_vals: &[u128]) -> OwnedBytes {
// let mut out = Vec::new();
// serialize_u128(
// || u128_vals.iter().cloned(),
// u128_vals.len() as u32,
// &mut out,
// )
// .unwrap();
//
// let data = OwnedBytes::new(out);
// let (data, _format_version) = read_format_version(data).unwrap();
// let (data, _null_index_footer) = read_null_index_footer(data).unwrap();
// test_all(data.clone(), u128_vals);
//
// data
// }
//
// #[test]
// fn test_range_1() {
// let vals = &[
// 1u128,
// 100u128,
// 3u128,
// 99999u128,
// 100000u128,
// 100001u128,
// 4_000_211_221u128,
// 4_000_211_222u128,
// 333u128,
// ];
// let mut data = test_aux_vals(vals);
//
// let _header = U128Header::deserialize(&mut data);
// let decomp = CompactSpaceDecompressor::open(data).unwrap();
// let complete_range = 0..vals.len() as u32;
// for (pos, val) in vals.iter().enumerate() {
// let val = *val;
// let pos = pos as u32;
// let mut positions = Vec::new();
// decomp.get_positions_for_value_range(val..=val, pos..pos + 1, &mut positions);
// assert_eq!(positions, vec![pos]);
// }
//
// handle docid range out of bounds
// let positions: Vec<u32> = get_positions_for_value_range_helper(&decomp, 0..=1, 1..u32::MAX);
// assert!(positions.is_empty());
//
// let positions =
// get_positions_for_value_range_helper(&decomp, 0..=1, complete_range.clone());
// assert_eq!(positions, vec![0]);
// let positions =
// get_positions_for_value_range_helper(&decomp, 0..=2, complete_range.clone());
// assert_eq!(positions, vec![0]);
// let positions =
// get_positions_for_value_range_helper(&decomp, 0..=3, complete_range.clone());
// assert_eq!(positions, vec![0, 2]);
// assert_eq!(
// get_positions_for_value_range_helper(
// &decomp,
// 99999u128..=99999u128,
// complete_range.clone()
// ),
// vec![3]
// );
// assert_eq!(
// get_positions_for_value_range_helper(
// &decomp,
// 99999u128..=100000u128,
// complete_range.clone()
// ),
// vec![3, 4]
// );
// assert_eq!(
// get_positions_for_value_range_helper(
// &decomp,
// 99998u128..=100000u128,
// complete_range.clone()
// ),
// vec![3, 4]
// );
// assert_eq!(
// &get_positions_for_value_range_helper(
// &decomp,
// 99998u128..=99999u128,
// complete_range.clone()
// ),
// &[3]
// );
// assert!(get_positions_for_value_range_helper(
// &decomp,
// 99998u128..=99998u128,
// complete_range.clone()
// )
// .is_empty());
// assert_eq!(
// &get_positions_for_value_range_helper(
// &decomp,
// 333u128..=333u128,
// complete_range.clone()
// ),
// &[8]
// );
// assert_eq!(
// &get_positions_for_value_range_helper(
// &decomp,
// 332u128..=333u128,
// complete_range.clone()
// ),
// &[8]
// );
// assert_eq!(
// &get_positions_for_value_range_helper(
// &decomp,
// 332u128..=334u128,
// complete_range.clone()
// ),
// &[8]
// );
// assert_eq!(
// &get_positions_for_value_range_helper(
// &decomp,
// 333u128..=334u128,
// complete_range.clone()
// ),
// &[8]
// );
//
// assert_eq!(
// &get_positions_for_value_range_helper(
// &decomp,
// 4_000_211_221u128..=5_000_000_000u128,
// complete_range
// ),
// &[6, 7]
// );
// }
//
// #[test]
// fn test_empty() {
// let vals = &[];
// let data = test_aux_vals(vals);
// let _decomp = CompactSpaceDecompressor::open(data).unwrap();
// }
//
// #[test]
// fn test_range_2() {
// let vals = &[
// 100u128,
// 99999u128,
// 100000u128,
// 100001u128,
// 4_000_211_221u128,
// 4_000_211_222u128,
// 333u128,
// ];
// let mut data = test_aux_vals(vals);
// let _header = U128Header::deserialize(&mut data);
// let decomp = CompactSpaceDecompressor::open(data).unwrap();
// let complete_range = 0..vals.len() as u32;
// assert!(
// &get_positions_for_value_range_helper(&decomp, 0..=5, complete_range.clone())
// .is_empty(),
// );
// assert_eq!(
// &get_positions_for_value_range_helper(&decomp, 0..=100, complete_range.clone()),
// &[0]
// );
// assert_eq!(
// &get_positions_for_value_range_helper(&decomp, 0..=105, complete_range),
// &[0]
// );
// }
//
// fn get_positions_for_value_range_helper<C: Column<T> + ?Sized, T: PartialOrd>(
// column: &C,
// value_range: RangeInclusive<T>,
// doc_id_range: Range<u32>,
// ) -> Vec<u32> {
// let mut positions = Vec::new();
// column.get_docids_for_value_range(value_range, doc_id_range, &mut positions);
// positions
// }
//
// #[test]
// fn test_range_3() {
// let vals = &[
// 200u128,
// 201,
// 202,
// 203,
// 204,
// 204,
// 206,
// 207,
// 208,
// 209,
// 210,
// 1_000_000,
// 5_000_000_000,
// ];
// let mut out = Vec::new();
// serialize_u128(|| vals.iter().cloned(), vals.len() as u32, &mut out).unwrap();
// let decomp = open_u128::<u128>(OwnedBytes::new(out)).unwrap();
// let complete_range = 0..vals.len() as u32;
//
// assert_eq!(
// get_positions_for_value_range_helper(&*decomp, 199..=200, complete_range.clone()),
// vec![0]
// );
//
// assert_eq!(
// get_positions_for_value_range_helper(&*decomp, 199..=201, complete_range.clone()),
// vec![0, 1]
// );
//
// assert_eq!(
// get_positions_for_value_range_helper(&*decomp, 200..=200, complete_range.clone()),
// vec![0]
// );
//
// assert_eq!(
// get_positions_for_value_range_helper(&*decomp, 1_000_000..=1_000_000, complete_range),
// vec![11]
// );
// }
//
// #[test]
// fn test_bug1() {
// let vals = &[9223372036854775806];
// let _data = test_aux_vals(vals);
// }
//
// #[test]
// fn test_bug2() {
// let vals = &[340282366920938463463374607431768211455u128];
// let _data = test_aux_vals(vals);
// }
//
// #[test]
// fn test_bug3() {
// let vals = &[340282366920938463463374607431768211454];
// let _data = test_aux_vals(vals);
// }
//
// #[test]
// fn test_bug4() {
// let vals = &[340282366920938463463374607431768211455, 0];
// let _data = test_aux_vals(vals);
// }
//
// #[test]
// fn test_first_large_gaps() {
// let vals = &[1_000_000_000u128; 100];
// let _data = test_aux_vals(vals);
// }
// use itertools::Itertools;
// use proptest::prelude::*;
//
// fn num_strategy() -> impl Strategy<Value = u128> {
// prop_oneof![
// 1 => prop::num::u128::ANY.prop_map(|num| u128::MAX - (num % 10) ),
// 1 => prop::num::u128::ANY.prop_map(|num| i64::MAX as u128 + 5 - (num % 10) ),
// 1 => prop::num::u128::ANY.prop_map(|num| i128::MAX as u128 + 5 - (num % 10) ),
// 1 => prop::num::u128::ANY.prop_map(|num| num % 10 ),
// 20 => prop::num::u128::ANY,
// ]
// }
//
// proptest! {
// #![proptest_config(ProptestConfig::with_cases(10))]
//
// #[test]
// fn compress_decompress_random(vals in proptest::collection::vec(num_strategy()
// , 1..1000)) {
// let _data = test_aux_vals(&vals);
// }
// }
// }
//

View File

@@ -1,75 +0,0 @@
use std::num::NonZeroU64;
use fastdivide::DividerU64;
/// Compute the gcd of two non null numbers.
///
/// It is recommended, but not required, to feed values such that `large >= small`.
fn compute_gcd(mut large: NonZeroU64, mut small: NonZeroU64) -> NonZeroU64 {
loop {
let rem: u64 = large.get() % small;
if let Some(new_small) = NonZeroU64::new(rem) {
(large, small) = (small, new_small);
} else {
return small;
}
}
}
// Find GCD for iterator of numbers
pub fn find_gcd(numbers: impl Iterator<Item = u64>) -> Option<NonZeroU64> {
let mut numbers = numbers.flat_map(NonZeroU64::new);
let mut gcd: NonZeroU64 = numbers.next()?;
if gcd.get() == 1 {
return Some(gcd);
}
let mut gcd_divider = DividerU64::divide_by(gcd.get());
for val in numbers {
let remainder = val.get() - (gcd_divider.divide(val.get())) * gcd.get();
if remainder == 0 {
continue;
}
gcd = compute_gcd(val, gcd);
if gcd.get() == 1 {
return Some(gcd);
}
gcd_divider = DividerU64::divide_by(gcd.get());
}
Some(gcd)
}
#[cfg(test)]
mod tests {
use std::num::NonZeroU64;
use crate::column_values::gcd::{compute_gcd, find_gcd};
#[test]
fn test_compute_gcd() {
let test_compute_gcd_aux = |large, small, expected| {
let large = NonZeroU64::new(large).unwrap();
let small = NonZeroU64::new(small).unwrap();
let expected = NonZeroU64::new(expected).unwrap();
assert_eq!(compute_gcd(small, large), expected);
assert_eq!(compute_gcd(large, small), expected);
};
test_compute_gcd_aux(1, 4, 1);
test_compute_gcd_aux(2, 4, 2);
test_compute_gcd_aux(10, 25, 5);
test_compute_gcd_aux(25, 25, 25);
}
#[test]
fn find_gcd_test() {
assert_eq!(find_gcd([0].into_iter()), None);
assert_eq!(find_gcd([0, 10].into_iter()), NonZeroU64::new(10));
assert_eq!(find_gcd([10, 0].into_iter()), NonZeroU64::new(10));
assert_eq!(find_gcd([].into_iter()), None);
assert_eq!(find_gcd([15, 30, 5, 10].into_iter()), NonZeroU64::new(5));
assert_eq!(find_gcd([15, 16, 10].into_iter()), NonZeroU64::new(1));
assert_eq!(find_gcd([0, 5, 5, 5].into_iter()), NonZeroU64::new(5));
assert_eq!(find_gcd([0, 0].into_iter()), None);
}
}

View File

@@ -1,222 +0,0 @@
use std::io;
use std::num::NonZeroU32;
use common::{BinarySerializable, VInt};
use crate::column_values::ColumnValues;
const MID_POINT: u64 = (1u64 << 32) - 1u64;
/// `Line` describes a line function `y: ax + b` using integer
/// arithmetics.
///
/// The slope is in fact a decimal split into a 32 bit integer value,
/// and a 32-bit decimal value.
///
/// The multiplication then becomes.
/// `y = m * x >> 32 + b`
#[derive(Debug, Clone, Copy, Default)]
pub struct Line {
slope: u64,
intercept: u64,
}
/// Compute the line slope.
///
/// This function has the nice property of being
/// invariant by translation.
/// `
/// compute_slope(y0, y1)
/// = compute_slope(y0 + X % 2^64, y1 + X % 2^64)
/// `
fn compute_slope(y0: u64, y1: u64, num_vals: NonZeroU32) -> u64 {
let dy = y1.wrapping_sub(y0);
let sign = dy <= (1 << 63);
let abs_dy = if sign {
y1.wrapping_sub(y0)
} else {
y0.wrapping_sub(y1)
};
if abs_dy >= 1 << 32 {
// This is outside of realm we handle.
// Let's just bail.
return 0u64;
}
let abs_slope = (abs_dy << 32) / num_vals.get() as u64;
if sign {
abs_slope
} else {
// The complement does indeed create the
// opposite decreasing slope...
//
// Intuitively (without the bitshifts and % u64::MAX)
// ```
// (x + shift)*(u64::MAX - abs_slope)
// - (x * (u64::MAX - abs_slope))
// = - shift * abs_slope
// ```
u64::MAX - abs_slope
}
}
impl Line {
#[inline(always)]
pub fn eval(&self, x: u32) -> u64 {
let linear_part = ((x as u64).wrapping_mul(self.slope) >> 32) as i32 as u64;
self.intercept.wrapping_add(linear_part)
}
// Same as train, but the intercept is only estimated from provided sample positions
pub fn estimate(sample_positions_and_values: &[(u64, u64)]) -> Self {
let first_val = sample_positions_and_values[0].1;
let last_val = sample_positions_and_values[sample_positions_and_values.len() - 1].1;
let num_vals = sample_positions_and_values[sample_positions_and_values.len() - 1].0 + 1;
Self::train_from(
first_val,
last_val,
num_vals as u32,
sample_positions_and_values.iter().cloned(),
)
}
// Intercept is only computed from provided positions
fn train_from(
first_val: u64,
last_val: u64,
num_vals: u32,
positions_and_values: impl Iterator<Item = (u64, u64)>,
) -> Self {
// TODO replace with let else
let idx_last_val = if let Some(idx_last_val) = NonZeroU32::new(num_vals - 1) {
idx_last_val
} else {
return Line::default();
};
let y0 = first_val;
let y1 = last_val;
// We first independently pick our slope.
let slope = compute_slope(y0, y1, idx_last_val);
// We picked our slope. Note that it does not have to be perfect.
// Now we need to compute the best intercept.
//
// Intuitively, the best intercept is such that line passes through one of the
// `(i, ys[])`.
//
// The best intercept therefore has the form
// `y[i] - line.eval(i)` (using wrapping arithmetics).
// In other words, the best intercept is one of the `y - Line::eval(ys[i])`
// and our task is just to pick the one that minimizes our error.
//
// Without sorting our values, this is a difficult problem.
// We however rely on the following trick...
//
// We only focus on the case where the interpolation is half decent.
// If the line interpolation is doing its job on a dataset suited for it,
// we can hope that the maximum error won't be larger than `u64::MAX / 2`.
//
// In other words, even without the intercept the values `y - Line::eval(ys[i])` will all be
// within an interval that takes less than half of the modulo space of `u64`.
//
// Our task is therefore to identify this interval.
// Here we simply translate all of our values by `y0 - 2^63` and pick the min.
let mut line = Line {
slope,
intercept: 0,
};
let heuristic_shift = y0.wrapping_sub(MID_POINT);
line.intercept = positions_and_values
.map(|(pos, y)| y.wrapping_sub(line.eval(pos as u32)))
.min_by_key(|&val| val.wrapping_sub(heuristic_shift))
.unwrap_or(0u64); //< Never happens.
line
}
/// Returns a line that attemps to approximate a function
/// f: i in 0..[ys.num_vals()) -> ys[i].
///
/// - The approximation is always lower than the actual value.
/// Or more rigorously, formally `f(i).wrapping_sub(ys[i])` is small
/// for any i in [0..ys.len()).
/// - It computes without panicking for any value of it.
///
/// This function is only invariable by translation if all of the
/// `ys` are packaged into half of the space. (See heuristic below)
pub fn train(ys: &dyn ColumnValues) -> Self {
let first_val = ys.iter().next().unwrap();
let last_val = ys.iter().nth(ys.num_vals() as usize - 1).unwrap();
Self::train_from(
first_val,
last_val,
ys.num_vals(),
ys.iter().enumerate().map(|(pos, val)| (pos as u64, val)),
)
}
}
impl BinarySerializable for Line {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
VInt(self.slope).serialize(writer)?;
VInt(self.intercept).serialize(writer)?;
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let slope = VInt::deserialize(reader)?.0;
let intercept = VInt::deserialize(reader)?.0;
Ok(Line { slope, intercept })
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::column_values::VecColumn;
/// Test training a line and ensuring that the maximum difference between
/// the data points and the line is `expected`.
///
/// This function operates translation over the data for better coverage.
#[track_caller]
fn test_line_interpol_with_translation(ys: &[u64], expected: Option<u64>) {
let mut translations = vec![0, 100, u64::MAX / 2, u64::MAX, u64::MAX - 1];
translations.extend_from_slice(ys);
for translation in translations {
let translated_ys: Vec<u64> = ys
.iter()
.copied()
.map(|y| y.wrapping_add(translation))
.collect();
let largest_err = test_eval_max_err(&translated_ys);
assert_eq!(largest_err, expected);
}
}
fn test_eval_max_err(ys: &[u64]) -> Option<u64> {
let line = Line::train(&VecColumn::from(&ys));
ys.iter()
.enumerate()
.map(|(x, y)| y.wrapping_sub(line.eval(x as u32)))
.max()
}
#[test]
fn test_train() {
test_line_interpol_with_translation(&[11, 11, 11, 12, 12, 13], Some(1));
test_line_interpol_with_translation(&[13, 12, 12, 11, 11, 11], Some(1));
test_line_interpol_with_translation(&[13, 13, 12, 11, 11, 11], Some(1));
test_line_interpol_with_translation(&[13, 13, 12, 11, 11, 11], Some(1));
test_line_interpol_with_translation(&[u64::MAX - 1, 0, 0, 1], Some(1));
test_line_interpol_with_translation(&[u64::MAX - 1, u64::MAX, 0, 1], Some(0));
test_line_interpol_with_translation(&[0, 1, 2, 3, 5], Some(0));
test_line_interpol_with_translation(&[1, 2, 3, 4], Some(0));
let data: Vec<u64> = (0..255).collect();
test_line_interpol_with_translation(&data, Some(0));
let data: Vec<u64> = (0..255).map(|el| el * 2).collect();
test_line_interpol_with_translation(&data, Some(0));
}
}

View File

@@ -1,230 +0,0 @@
use std::io::{self, Write};
use common::{BinarySerializable, OwnedBytes};
use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
use super::line::Line;
use super::serialize::NormalizedHeader;
use super::{ColumnValues, FastFieldCodec, FastFieldCodecType};
/// Depending on the field type, a different
/// fast field is required.
#[derive(Clone)]
pub struct LinearReader {
data: OwnedBytes,
linear_params: LinearParams,
header: NormalizedHeader,
}
impl ColumnValues for LinearReader {
#[inline]
fn get_val(&self, doc: u32) -> u64 {
let interpoled_val: u64 = self.linear_params.line.eval(doc);
let bitpacked_diff = self.linear_params.bit_unpacker.get(doc, &self.data);
interpoled_val.wrapping_add(bitpacked_diff)
}
#[inline(always)]
fn min_value(&self) -> u64 {
// The LinearReader assumes a normalized vector.
0u64
}
#[inline(always)]
fn max_value(&self) -> u64 {
self.header.max_value
}
#[inline]
fn num_vals(&self) -> u32 {
self.header.num_vals
}
}
/// Fastfield serializer, which tries to guess values by linear interpolation
/// and stores the difference bitpacked.
pub struct LinearCodec;
#[derive(Debug, Clone)]
struct LinearParams {
line: Line,
bit_unpacker: BitUnpacker,
}
impl BinarySerializable for LinearParams {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
self.line.serialize(writer)?;
self.bit_unpacker.bit_width().serialize(writer)?;
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let line = Line::deserialize(reader)?;
let bit_width = u8::deserialize(reader)?;
Ok(Self {
line,
bit_unpacker: BitUnpacker::new(bit_width),
})
}
}
impl FastFieldCodec for LinearCodec {
const CODEC_TYPE: FastFieldCodecType = FastFieldCodecType::Linear;
type Reader = LinearReader;
/// Opens a fast field given a file.
fn open_from_bytes(mut data: OwnedBytes, header: NormalizedHeader) -> io::Result<Self::Reader> {
let linear_params = LinearParams::deserialize(&mut data)?;
Ok(LinearReader {
data,
linear_params,
header,
})
}
/// Creates a new fast field serializer.
fn serialize(column: &dyn ColumnValues, write: &mut impl Write) -> io::Result<()> {
assert_eq!(column.min_value(), 0);
let line = Line::train(column);
let max_offset_from_line = column
.iter()
.enumerate()
.map(|(pos, actual_value)| {
let calculated_value = line.eval(pos as u32);
actual_value.wrapping_sub(calculated_value)
})
.max()
.unwrap();
let num_bits = compute_num_bits(max_offset_from_line);
let linear_params = LinearParams {
line,
bit_unpacker: BitUnpacker::new(num_bits),
};
linear_params.serialize(write)?;
let mut bit_packer = BitPacker::new();
for (pos, actual_value) in column.iter().enumerate() {
let calculated_value = line.eval(pos as u32);
let offset = actual_value.wrapping_sub(calculated_value);
bit_packer.write(offset, num_bits, write)?;
}
bit_packer.close(write)?;
Ok(())
}
/// estimation for linear interpolation is hard because, you don't know
/// where the local maxima for the deviation of the calculated value are and
/// the offset to shift all values to >=0 is also unknown.
#[allow(clippy::question_mark)]
fn estimate(column: &dyn ColumnValues) -> Option<f32> {
if column.num_vals() < 3 {
return None; // disable compressor for this case
}
let limit_num_vals = column.num_vals().min(100_000);
let num_samples = 100;
let step_size = (limit_num_vals / num_samples).max(1); // 20 samples
let mut sample_positions_and_values: Vec<_> = Vec::new();
for (pos, val) in column.iter().enumerate().step_by(step_size as usize) {
sample_positions_and_values.push((pos as u64, val));
}
let line = Line::estimate(&sample_positions_and_values);
let estimated_bit_width = sample_positions_and_values
.into_iter()
.map(|(pos, actual_value)| {
let interpolated_val = line.eval(pos as u32);
actual_value.wrapping_sub(interpolated_val)
})
.map(|diff| ((diff as f32 * 1.5) * 2.0) as u64)
.map(compute_num_bits)
.max()
.unwrap_or(0);
// Extrapolate to whole column
let num_bits = (estimated_bit_width as u64 * column.num_vals() as u64) + 64;
let num_bits_uncompressed = 64 * column.num_vals();
Some(num_bits as f32 / num_bits_uncompressed as f32)
}
}
#[cfg(test)]
mod tests {
use rand::RngCore;
use super::*;
use crate::column_values::tests;
fn create_and_validate(data: &[u64], name: &str) -> Option<(f32, f32)> {
tests::create_and_validate::<LinearCodec>(data, name)
}
#[test]
fn test_compression() {
let data = (10..=6_000_u64).collect::<Vec<_>>();
let (estimate, actual_compression) =
create_and_validate(&data, "simple monotonically large").unwrap();
assert_le!(actual_compression, 0.001);
assert_le!(estimate, 0.02);
}
#[test]
fn test_with_codec_datasets() {
let data_sets = tests::get_codec_test_datasets();
for (mut data, name) in data_sets {
create_and_validate(&data, name);
data.reverse();
create_and_validate(&data, name);
}
}
#[test]
fn linear_interpol_fast_field_test_large_amplitude() {
let data = vec![
i64::MAX as u64 / 2,
i64::MAX as u64 / 3,
i64::MAX as u64 / 2,
];
create_and_validate(&data, "large amplitude");
}
#[test]
fn overflow_error_test() {
let data = vec![1572656989877777, 1170935903116329, 720575940379279, 0];
create_and_validate(&data, "overflow test");
}
#[test]
fn linear_interpol_fast_concave_data() {
let data = vec![0, 1, 2, 5, 8, 10, 20, 50];
create_and_validate(&data, "concave data");
}
#[test]
fn linear_interpol_fast_convex_data() {
let data = vec![0, 40, 60, 70, 75, 77];
create_and_validate(&data, "convex data");
}
#[test]
fn linear_interpol_fast_field_test_simple() {
let data = (10..=20_u64).collect::<Vec<_>>();
create_and_validate(&data, "simple monotonically");
}
#[test]
fn linear_interpol_fast_field_rand() {
let mut rng = rand::thread_rng();
for _ in 0..50 {
let mut data = (0..10_000).map(|_| rng.next_u64()).collect::<Vec<_>>();
create_and_validate(&data, "random");
data.reverse();
create_and_validate(&data, "random");
}
}
}

View File

@@ -1,222 +0,0 @@
#[macro_use]
extern crate prettytable;
use std::collections::HashSet;
use std::env;
use std::io::BufRead;
use std::net::{IpAddr, Ipv6Addr};
use std::str::FromStr;
use common::OwnedBytes;
use fastfield_codecs::{open_u128, serialize_u128, Column, FastFieldCodecType, VecColumn};
use itertools::Itertools;
use measure_time::print_time;
use prettytable::{Cell, Row, Table};
fn print_set_stats(ip_addrs: &[u128]) {
println!("NumIps\t{}", ip_addrs.len());
let ip_addr_set: HashSet<u128> = ip_addrs.iter().cloned().collect();
println!("NumUniqueIps\t{}", ip_addr_set.len());
let ratio_unique = ip_addr_set.len() as f64 / ip_addrs.len() as f64;
println!("RatioUniqueOverTotal\t{ratio_unique:.4}");
// histogram
let mut ip_addrs = ip_addrs.to_vec();
ip_addrs.sort();
let mut cnts: Vec<usize> = ip_addrs
.into_iter()
.dedup_with_count()
.map(|(cnt, _)| cnt)
.collect();
cnts.sort();
let top_256_cnt: usize = cnts.iter().rev().take(256).sum();
let top_128_cnt: usize = cnts.iter().rev().take(128).sum();
let top_64_cnt: usize = cnts.iter().rev().take(64).sum();
let top_8_cnt: usize = cnts.iter().rev().take(8).sum();
let total: usize = cnts.iter().sum();
println!("{}", total);
println!("{}", top_256_cnt);
println!("{}", top_128_cnt);
println!("Percentage Top8 {:02}", top_8_cnt as f32 / total as f32);
println!("Percentage Top64 {:02}", top_64_cnt as f32 / total as f32);
println!("Percentage Top128 {:02}", top_128_cnt as f32 / total as f32);
println!("Percentage Top256 {:02}", top_256_cnt as f32 / total as f32);
let mut cnts: Vec<(usize, usize)> = cnts.into_iter().dedup_with_count().collect();
cnts.sort_by(|a, b| {
if a.1 == b.1 {
a.0.cmp(&b.0)
} else {
b.1.cmp(&a.1)
}
});
}
fn ip_dataset() -> Vec<u128> {
let mut ip_addr_v4 = 0;
let stdin = std::io::stdin();
let ip_addrs: Vec<u128> = stdin
.lock()
.lines()
.flat_map(|line| {
let line = line.unwrap();
let line = line.trim();
let ip_addr = IpAddr::from_str(line.trim()).ok()?;
if ip_addr.is_ipv4() {
ip_addr_v4 += 1;
}
let ip_addr_v6: Ipv6Addr = match ip_addr {
IpAddr::V4(v4) => v4.to_ipv6_mapped(),
IpAddr::V6(v6) => v6,
};
Some(ip_addr_v6)
})
.map(|ip_v6| u128::from_be_bytes(ip_v6.octets()))
.collect();
println!("IpAddrsAny\t{}", ip_addrs.len());
println!("IpAddrsV4\t{}", ip_addr_v4);
ip_addrs
}
fn bench_ip() {
let dataset = ip_dataset();
print_set_stats(&dataset);
// Chunks
{
let mut data = vec![];
for dataset in dataset.chunks(500_000) {
serialize_u128(|| dataset.iter().cloned(), dataset.len() as u32, &mut data).unwrap();
}
let compression = data.len() as f64 / (dataset.len() * 16) as f64;
println!("Compression 50_000 chunks {:.4}", compression);
println!(
"Num Bits per elem {:.2}",
(data.len() * 8) as f32 / dataset.len() as f32
);
}
let mut data = vec![];
{
print_time!("creation");
serialize_u128(|| dataset.iter().cloned(), dataset.len() as u32, &mut data).unwrap();
}
let compression = data.len() as f64 / (dataset.len() * 16) as f64;
println!("Compression {:.2}", compression);
println!(
"Num Bits per elem {:.2}",
(data.len() * 8) as f32 / dataset.len() as f32
);
let decompressor = open_u128::<u128>(OwnedBytes::new(data)).unwrap();
// Sample some ranges
let mut doc_values = Vec::new();
for value in dataset.iter().take(1110).skip(1100).cloned() {
doc_values.clear();
print_time!("get range");
decompressor.get_docids_for_value_range(
value..=value,
0..decompressor.num_vals(),
&mut doc_values,
);
println!("{:?}", doc_values.len());
}
}
fn main() {
if env::args().nth(1).unwrap() == "bench_ip" {
bench_ip();
return;
}
let mut table = Table::new();
// Add a row per time
table.add_row(row!["", "Compression Ratio", "Compression Estimation"]);
for (data, data_set_name) in get_codec_test_data_sets() {
let results: Vec<(f32, f32, FastFieldCodecType)> = [
serialize_with_codec(&data, FastFieldCodecType::Bitpacked),
serialize_with_codec(&data, FastFieldCodecType::Linear),
serialize_with_codec(&data, FastFieldCodecType::BlockwiseLinear),
]
.into_iter()
.flatten()
.collect();
let best_compression_ratio_codec = results
.iter()
.min_by(|&res1, &res2| res1.partial_cmp(res2).unwrap())
.cloned()
.unwrap();
table.add_row(Row::new(vec![Cell::new(data_set_name).style_spec("Bbb")]));
for (est, comp, codec_type) in results {
let est_cell = est.to_string();
let ratio_cell = comp.to_string();
let style = if comp == best_compression_ratio_codec.1 {
"Fb"
} else {
""
};
table.add_row(Row::new(vec![
Cell::new(&format!("{codec_type:?}")).style_spec("bFg"),
Cell::new(&ratio_cell).style_spec(style),
Cell::new(&est_cell).style_spec(""),
]));
}
}
table.printstd();
}
pub fn get_codec_test_data_sets() -> Vec<(Vec<u64>, &'static str)> {
let mut data_and_names = vec![];
let data = (1000..=200_000_u64).collect::<Vec<_>>();
data_and_names.push((data, "Autoincrement"));
let mut current_cumulative = 0;
let data = (1..=200_000_u64)
.map(|num| {
let num = (num as f32 + num as f32).log10() as u64;
current_cumulative += num;
current_cumulative
})
.collect::<Vec<_>>();
// let data = (1..=200000_u64).map(|num| num + num).collect::<Vec<_>>();
data_and_names.push((data, "Monotonically increasing concave"));
let mut current_cumulative = 0;
let data = (1..=200_000_u64)
.map(|num| {
let num = (200_000.0 - num as f32).log10() as u64;
current_cumulative += num;
current_cumulative
})
.collect::<Vec<_>>();
data_and_names.push((data, "Monotonically increasing convex"));
let data = (1000..=200_000_u64)
.map(|num| num + rand::random::<u8>() as u64)
.collect::<Vec<_>>();
data_and_names.push((data, "Almost monotonically increasing"));
data_and_names
}
pub fn serialize_with_codec(
data: &[u64],
codec_type: FastFieldCodecType,
) -> Option<(f32, f32, FastFieldCodecType)> {
let col = VecColumn::from(data);
let estimation = fastfield_codecs::estimate(&col, codec_type)?;
let mut out = Vec::new();
fastfield_codecs::serialize(&col, &mut out, &[codec_type]).ok()?;
let actual_compression = out.len() as f32 / (col.num_vals() * 8) as f32;
Some((estimation, actual_compression, codec_type))
}

View File

@@ -1,326 +0,0 @@
#![warn(missing_docs)]
#![cfg_attr(all(feature = "unstable", test), feature(test))]
//! # `fastfield_codecs`
//!
//! - Columnar storage of data for tantivy [`Column`].
//! - Encode data in different codecs.
//! - Monotonically map values to u64/u128
#[cfg(test)]
mod tests;
use std::io;
use std::io::Write;
use std::sync::Arc;
use common::{BinarySerializable, OwnedBytes};
use compact_space::CompactSpaceDecompressor;
use monotonic_mapping::{
StrictlyMonotonicMappingInverter, StrictlyMonotonicMappingToInternal,
StrictlyMonotonicMappingToInternalBaseval, StrictlyMonotonicMappingToInternalGCDBaseval,
};
use serialize::{Header, U128Header};
mod bitpacked;
mod blockwise_linear;
mod compact_space;
mod line;
mod linear;
pub(crate) mod monotonic_mapping;
pub(crate) mod monotonic_mapping_u128;
mod column;
mod column_with_cardinality;
mod gcd;
pub mod serialize;
pub use self::column::{monotonic_map_column, ColumnValues, IterColumn, VecColumn};
pub use self::monotonic_mapping::{MonotonicallyMappableToU64, StrictlyMonotonicFn};
pub use self::monotonic_mapping_u128::MonotonicallyMappableToU128;
#[cfg(test)]
pub use self::serialize::tests::serialize_and_load;
pub use self::serialize::{serialize_column_values, NormalizedHeader};
use crate::column_values::bitpacked::BitpackedCodec;
use crate::column_values::blockwise_linear::BlockwiseLinearCodec;
use crate::column_values::linear::LinearCodec;
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
#[repr(u8)]
/// Available codecs to use to encode the u64 (via [`MonotonicallyMappableToU64`]) converted data.
pub enum FastFieldCodecType {
/// Bitpack all values in the value range. The number of bits is defined by the amplitude
/// `column.max_value() - column.min_value()`
Bitpacked = 1,
/// Linear interpolation puts a line between the first and last value and then bitpacks the
/// values by the offset from the line. The number of bits is defined by the max deviation from
/// the line.
Linear = 2,
/// Same as [`FastFieldCodecType::Linear`], but encodes in blocks of 512 elements.
BlockwiseLinear = 3,
}
impl BinarySerializable for FastFieldCodecType {
fn serialize<W: Write>(&self, wrt: &mut W) -> io::Result<()> {
self.to_code().serialize(wrt)
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let code = u8::deserialize(reader)?;
let codec_type: Self = Self::from_code(code)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Unknown code `{code}.`"))?;
Ok(codec_type)
}
}
impl FastFieldCodecType {
pub(crate) fn to_code(self) -> u8 {
self as u8
}
pub(crate) fn from_code(code: u8) -> Option<Self> {
match code {
1 => Some(Self::Bitpacked),
2 => Some(Self::Linear),
3 => Some(Self::BlockwiseLinear),
_ => None,
}
}
}
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
#[repr(u8)]
/// Available codecs to use to encode the u128 (via [`MonotonicallyMappableToU128`]) converted data.
pub enum U128FastFieldCodecType {
/// This codec takes a large number space (u128) and reduces it to a compact number space, by
/// removing the holes.
CompactSpace = 1,
}
impl BinarySerializable for U128FastFieldCodecType {
fn serialize<W: Write>(&self, wrt: &mut W) -> io::Result<()> {
self.to_code().serialize(wrt)
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let code = u8::deserialize(reader)?;
let codec_type: Self = Self::from_code(code)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Unknown code `{code}.`"))?;
Ok(codec_type)
}
}
impl U128FastFieldCodecType {
pub(crate) fn to_code(self) -> u8 {
self as u8
}
pub(crate) fn from_code(code: u8) -> Option<Self> {
match code {
1 => Some(Self::CompactSpace),
_ => None,
}
}
}
/// Returns the correct codec reader wrapped in the `Arc` for the data.
pub fn open_u128_mapped<T: MonotonicallyMappableToU128>(
mut bytes: OwnedBytes,
) -> io::Result<Arc<dyn ColumnValues<T>>> {
let header = U128Header::deserialize(&mut bytes)?;
assert_eq!(header.codec_type, U128FastFieldCodecType::CompactSpace);
let reader = CompactSpaceDecompressor::open(bytes)?;
let inverted: StrictlyMonotonicMappingInverter<StrictlyMonotonicMappingToInternal<T>> =
StrictlyMonotonicMappingToInternal::<T>::new().into();
Ok(Arc::new(monotonic_map_column(reader, inverted)))
}
/// Returns the correct codec reader wrapped in the `Arc` for the data.
pub fn open_u64_mapped<T: MonotonicallyMappableToU64>(
mut bytes: OwnedBytes,
) -> io::Result<Arc<dyn ColumnValues<T>>> {
let header = Header::deserialize(&mut bytes)?;
match header.codec_type {
FastFieldCodecType::Bitpacked => open_specific_codec::<BitpackedCodec, _>(bytes, &header),
FastFieldCodecType::Linear => open_specific_codec::<LinearCodec, _>(bytes, &header),
FastFieldCodecType::BlockwiseLinear => {
open_specific_codec::<BlockwiseLinearCodec, _>(bytes, &header)
}
}
}
fn open_specific_codec<C: FastFieldCodec, Item: MonotonicallyMappableToU64>(
bytes: OwnedBytes,
header: &Header,
) -> io::Result<Arc<dyn ColumnValues<Item>>> {
let normalized_header = header.normalized();
let reader = C::open_from_bytes(bytes, normalized_header)?;
let min_value = header.min_value;
if let Some(gcd) = header.gcd {
let mapping = StrictlyMonotonicMappingInverter::from(
StrictlyMonotonicMappingToInternalGCDBaseval::new(gcd.get(), min_value),
);
Ok(Arc::new(monotonic_map_column(reader, mapping)))
} else {
let mapping = StrictlyMonotonicMappingInverter::from(
StrictlyMonotonicMappingToInternalBaseval::new(min_value),
);
Ok(Arc::new(monotonic_map_column(reader, mapping)))
}
}
/// The FastFieldSerializerEstimate trait is required on all variants
/// of fast field compressions, to decide which one to choose.
pub(crate) trait FastFieldCodec: 'static {
/// A codex needs to provide a unique name and id, which is
/// used for debugging and de/serialization.
const CODEC_TYPE: FastFieldCodecType;
type Reader: ColumnValues<u64> + 'static;
/// Reads the metadata and returns the CodecReader
fn open_from_bytes(bytes: OwnedBytes, header: NormalizedHeader) -> io::Result<Self::Reader>;
/// Serializes the data using the serializer into write.
///
/// The column iterator should be preferred over using column `get_val` method for
/// performance reasons.
fn serialize(column: &dyn ColumnValues, write: &mut impl Write) -> io::Result<()>;
/// Returns an estimate of the compression ratio.
/// If the codec is not applicable, returns `None`.
///
/// The baseline is uncompressed 64bit data.
///
/// It could make sense to also return a value representing
/// computational complexity.
fn estimate(column: &dyn ColumnValues) -> Option<f32>;
}
#[cfg(all(test, feature = "unstable"))]
mod bench {
use std::sync::Arc;
use common::OwnedBytes;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use test::{self, Bencher};
use super::*;
fn get_data() -> Vec<u64> {
let mut rng = StdRng::seed_from_u64(2u64);
let mut data: Vec<_> = (100..55000_u64)
.map(|num| num + rng.gen::<u8>() as u64)
.collect();
data.push(99_000);
data.insert(1000, 2000);
data.insert(2000, 100);
data.insert(3000, 4100);
data.insert(4000, 100);
data.insert(5000, 800);
data
}
#[inline(never)]
fn value_iter() -> impl Iterator<Item = u64> {
0..20_000
}
fn get_reader_for_bench<Codec: FastFieldCodec>(data: &[u64]) -> Codec::Reader {
let mut bytes = Vec::new();
let min_value = *data.iter().min().unwrap();
let data = data.iter().map(|el| *el - min_value).collect::<Vec<_>>();
let col = VecColumn::from(&data);
let normalized_header = NormalizedHeader {
num_vals: col.num_vals(),
max_value: col.max_value(),
};
Codec::serialize(&VecColumn::from(&data), &mut bytes).unwrap();
Codec::open_from_bytes(OwnedBytes::new(bytes), normalized_header).unwrap()
}
fn bench_get<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
let col = get_reader_for_bench::<Codec>(data);
b.iter(|| {
let mut sum = 0u64;
for pos in value_iter() {
let val = col.get_val(pos as u32);
sum = sum.wrapping_add(val);
}
sum
});
}
#[inline(never)]
fn bench_get_dynamic_helper(b: &mut Bencher, col: Arc<dyn ColumnValues>) {
b.iter(|| {
let mut sum = 0u64;
for pos in value_iter() {
let val = col.get_val(pos as u32);
sum = sum.wrapping_add(val);
}
sum
});
}
fn bench_get_dynamic<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
let col = Arc::new(get_reader_for_bench::<Codec>(data));
bench_get_dynamic_helper(b, col);
}
fn bench_create<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
let min_value = *data.iter().min().unwrap();
let data = data.iter().map(|el| *el - min_value).collect::<Vec<_>>();
let mut bytes = Vec::new();
b.iter(|| {
bytes.clear();
Codec::serialize(&VecColumn::from(&data), &mut bytes).unwrap();
});
}
#[bench]
fn bench_fastfield_bitpack_create(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_create::<BitpackedCodec>(b, &data);
}
#[bench]
fn bench_fastfield_linearinterpol_create(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_create::<LinearCodec>(b, &data);
}
#[bench]
fn bench_fastfield_multilinearinterpol_create(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_create::<BlockwiseLinearCodec>(b, &data);
}
#[bench]
fn bench_fastfield_bitpack_get(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_get::<BitpackedCodec>(b, &data);
}
#[bench]
fn bench_fastfield_bitpack_get_dynamic(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_get_dynamic::<BitpackedCodec>(b, &data);
}
#[bench]
fn bench_fastfield_linearinterpol_get(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_get::<LinearCodec>(b, &data);
}
#[bench]
fn bench_fastfield_linearinterpol_get_dynamic(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_get_dynamic::<LinearCodec>(b, &data);
}
#[bench]
fn bench_fastfield_multilinearinterpol_get(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_get::<BlockwiseLinearCodec>(b, &data);
}
#[bench]
fn bench_fastfield_multilinearinterpol_get_dynamic(b: &mut Bencher) {
let data: Vec<_> = get_data();
bench_get_dynamic::<BlockwiseLinearCodec>(b, &data);
}
}

View File

@@ -1,278 +0,0 @@
use std::marker::PhantomData;
use fastdivide::DividerU64;
use super::MonotonicallyMappableToU128;
use crate::RowId;
/// Monotonic maps a value to u64 value space.
/// Monotonic mapping enables `PartialOrd` on u64 space without conversion to original space.
pub trait MonotonicallyMappableToU64: 'static + PartialOrd + Copy + Send + Sync {
/// Converts a value to u64.
///
/// Internally all fast field values are encoded as u64.
fn to_u64(self) -> u64;
/// Converts a value from u64
///
/// Internally all fast field values are encoded as u64.
/// **Note: To be used for converting encoded Term, Posting values.**
fn from_u64(val: u64) -> Self;
}
/// Values need to be strictly monotonic mapped to a `Internal` value (u64 or u128) that can be
/// used in fast field codecs.
///
/// The monotonic mapping is required so that `PartialOrd` can be used on `Internal` without
/// converting to `External`.
///
/// All strictly monotonic functions are invertible because they are guaranteed to have a one-to-one
/// mapping from their range to their domain. The `inverse` method is required when opening a codec,
/// so a value can be converted back to its original domain (e.g. ip address or f64) from its
/// internal representation.
pub trait StrictlyMonotonicFn<External, Internal> {
/// Strictly monotonically maps the value from External to Internal.
fn mapping(&self, inp: External) -> Internal;
/// Inverse of `mapping`. Maps the value from Internal to External.
fn inverse(&self, out: Internal) -> External;
}
/// Inverts a strictly monotonic mapping from `StrictlyMonotonicFn<A, B>` to
/// `StrictlyMonotonicFn<B, A>`.
///
/// # Warning
///
/// This type comes with a footgun. A type being strictly monotonic does not impose that the inverse
/// mapping is strictly monotonic over the entire space External. e.g. a -> a * 2. Use at your own
/// risks.
pub(crate) struct StrictlyMonotonicMappingInverter<T> {
orig_mapping: T,
}
impl<T> From<T> for StrictlyMonotonicMappingInverter<T> {
fn from(orig_mapping: T) -> Self {
Self { orig_mapping }
}
}
impl<From, To, T> StrictlyMonotonicFn<To, From> for StrictlyMonotonicMappingInverter<T>
where T: StrictlyMonotonicFn<From, To>
{
#[inline(always)]
fn mapping(&self, val: To) -> From {
self.orig_mapping.inverse(val)
}
#[inline(always)]
fn inverse(&self, val: From) -> To {
self.orig_mapping.mapping(val)
}
}
/// Applies the strictly monotonic mapping from `T` without any additional changes.
pub(crate) struct StrictlyMonotonicMappingToInternal<T> {
_phantom: PhantomData<T>,
}
impl<T> StrictlyMonotonicMappingToInternal<T> {
pub(crate) fn new() -> StrictlyMonotonicMappingToInternal<T> {
Self {
_phantom: PhantomData,
}
}
}
impl<External: MonotonicallyMappableToU128, T: MonotonicallyMappableToU128>
StrictlyMonotonicFn<External, u128> for StrictlyMonotonicMappingToInternal<T>
where T: MonotonicallyMappableToU128
{
#[inline(always)]
fn mapping(&self, inp: External) -> u128 {
External::to_u128(inp)
}
#[inline(always)]
fn inverse(&self, out: u128) -> External {
External::from_u128(out)
}
}
impl<External: MonotonicallyMappableToU64, T: MonotonicallyMappableToU64>
StrictlyMonotonicFn<External, u64> for StrictlyMonotonicMappingToInternal<T>
where T: MonotonicallyMappableToU64
{
#[inline(always)]
fn mapping(&self, inp: External) -> u64 {
External::to_u64(inp)
}
#[inline(always)]
fn inverse(&self, out: u64) -> External {
External::from_u64(out)
}
}
/// Mapping dividing by gcd and a base value.
///
/// The function is assumed to be only called on values divided by passed
/// gcd value. (It is necessary for the function to be monotonic.)
pub(crate) struct StrictlyMonotonicMappingToInternalGCDBaseval {
gcd_divider: DividerU64,
gcd: u64,
min_value: u64,
}
impl StrictlyMonotonicMappingToInternalGCDBaseval {
pub(crate) fn new(gcd: u64, min_value: u64) -> Self {
let gcd_divider = DividerU64::divide_by(gcd);
Self {
gcd_divider,
gcd,
min_value,
}
}
}
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
for StrictlyMonotonicMappingToInternalGCDBaseval
{
#[inline(always)]
fn mapping(&self, inp: External) -> u64 {
self.gcd_divider
.divide(External::to_u64(inp) - self.min_value)
}
#[inline(always)]
fn inverse(&self, out: u64) -> External {
External::from_u64(self.min_value + out * self.gcd)
}
}
/// Strictly monotonic mapping with a base value.
pub(crate) struct StrictlyMonotonicMappingToInternalBaseval {
min_value: u64,
}
impl StrictlyMonotonicMappingToInternalBaseval {
#[inline(always)]
pub(crate) fn new(min_value: u64) -> Self {
Self { min_value }
}
}
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
for StrictlyMonotonicMappingToInternalBaseval
{
#[inline(always)]
fn mapping(&self, val: External) -> u64 {
External::to_u64(val) - self.min_value
}
#[inline(always)]
fn inverse(&self, val: u64) -> External {
External::from_u64(self.min_value + val)
}
}
impl MonotonicallyMappableToU64 for u64 {
#[inline(always)]
fn to_u64(self) -> u64 {
self
}
#[inline(always)]
fn from_u64(val: u64) -> Self {
val
}
}
impl MonotonicallyMappableToU64 for i64 {
#[inline(always)]
fn to_u64(self) -> u64 {
common::i64_to_u64(self)
}
#[inline(always)]
fn from_u64(val: u64) -> Self {
common::u64_to_i64(val)
}
}
impl MonotonicallyMappableToU64 for crate::DateTime {
#[inline(always)]
fn to_u64(self) -> u64 {
common::i64_to_u64(self.timestamp_micros)
}
#[inline(always)]
fn from_u64(val: u64) -> Self {
crate::DateTime {
timestamp_micros: common::u64_to_i64(val),
}
}
}
impl MonotonicallyMappableToU64 for bool {
#[inline(always)]
fn to_u64(self) -> u64 {
u64::from(self)
}
#[inline(always)]
fn from_u64(val: u64) -> Self {
val > 0
}
}
impl MonotonicallyMappableToU64 for RowId {
#[inline(always)]
fn to_u64(self) -> u64 {
u64::from(self)
}
#[inline(always)]
fn from_u64(val: u64) -> RowId {
val as RowId
}
}
// TODO remove me.
// Tantivy should refuse NaN values and work with NotNaN internally.
impl MonotonicallyMappableToU64 for f64 {
#[inline(always)]
fn to_u64(self) -> u64 {
common::f64_to_u64(self)
}
#[inline(always)]
fn from_u64(val: u64) -> Self {
common::u64_to_f64(val)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn strictly_monotonic_test() {
// identity mapping
test_round_trip(&StrictlyMonotonicMappingToInternal::<u64>::new(), 100u64);
// round trip to i64
test_round_trip(&StrictlyMonotonicMappingToInternal::<i64>::new(), 100u64);
// TODO
// identity mapping
// test_round_trip(&StrictlyMonotonicMappingToInternal::<u128>::new(), 100u128);
// base value to i64 round trip
let mapping = StrictlyMonotonicMappingToInternalBaseval::new(100);
test_round_trip::<_, _, u64>(&mapping, 100i64);
// base value and gcd to u64 round trip
let mapping = StrictlyMonotonicMappingToInternalGCDBaseval::new(10, 100);
test_round_trip::<_, _, u64>(&mapping, 100u64);
}
fn test_round_trip<T: StrictlyMonotonicFn<K, L>, K: std::fmt::Debug + Eq + Copy, L>(
mapping: &T,
test_val: K,
) {
assert_eq!(mapping.inverse(mapping.mapping(test_val)), test_val);
}
}

View File

@@ -1,40 +0,0 @@
use std::net::Ipv6Addr;
/// Montonic maps a value to u128 value space
/// Monotonic mapping enables `PartialOrd` on u128 space without conversion to original space.
pub trait MonotonicallyMappableToU128: 'static + PartialOrd + Copy + Send + Sync {
/// Converts a value to u128.
///
/// Internally all fast field values are encoded as u64.
fn to_u128(self) -> u128;
/// Converts a value from u128
///
/// Internally all fast field values are encoded as u64.
/// **Note: To be used for converting encoded Term, Posting values.**
fn from_u128(val: u128) -> Self;
}
impl MonotonicallyMappableToU128 for u128 {
fn to_u128(self) -> u128 {
self
}
fn from_u128(val: u128) -> Self {
val
}
}
impl MonotonicallyMappableToU128 for Ipv6Addr {
fn to_u128(self) -> u128 {
ip_to_u128(self)
}
fn from_u128(val: u128) -> Self {
Ipv6Addr::from(val.to_be_bytes())
}
}
fn ip_to_u128(ip_addr: Ipv6Addr) -> u128 {
u128::from_be_bytes(ip_addr.octets())
}

View File

@@ -1,320 +0,0 @@
// 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::num::NonZeroU64;
use common::{BinarySerializable, VInt};
use log::warn;
use super::bitpacked::BitpackedCodec;
use super::blockwise_linear::BlockwiseLinearCodec;
use super::linear::LinearCodec;
use super::monotonic_mapping::{
StrictlyMonotonicFn, StrictlyMonotonicMappingToInternal,
StrictlyMonotonicMappingToInternalGCDBaseval,
};
use super::{
monotonic_map_column, ColumnValues, FastFieldCodec, FastFieldCodecType,
MonotonicallyMappableToU64, U128FastFieldCodecType,
};
use crate::column_values::compact_space::CompactSpaceCompressor;
/// The normalized header gives some parameters after applying the following
/// normalization of the vector:
/// `val -> (val - min_value) / gcd`
///
/// By design, after normalization, `min_value = 0` and `gcd = 1`.
#[derive(Debug, Copy, Clone)]
pub struct NormalizedHeader {
/// The number of values in the underlying column.
pub num_vals: u32,
/// The max value of the underlying column.
pub max_value: u64,
}
#[derive(Debug, Copy, Clone)]
pub(crate) struct Header {
pub num_vals: u32,
pub min_value: u64,
pub max_value: u64,
pub gcd: Option<NonZeroU64>,
pub codec_type: FastFieldCodecType,
}
impl Header {
pub fn normalized(self) -> NormalizedHeader {
let gcd = self.gcd.map(|gcd| gcd.get()).unwrap_or(1);
let gcd_min_val_mapping =
StrictlyMonotonicMappingToInternalGCDBaseval::new(gcd, self.min_value);
let max_value = gcd_min_val_mapping.mapping(self.max_value);
NormalizedHeader {
num_vals: self.num_vals,
max_value,
}
}
pub(crate) fn normalize_column<C: ColumnValues>(&self, from_column: C) -> impl ColumnValues {
normalize_column(from_column, self.min_value, self.gcd)
}
pub fn compute_header(
column: impl ColumnValues<u64>,
codecs: &[FastFieldCodecType],
) -> Option<Header> {
let num_vals = column.num_vals();
let min_value = column.min_value();
let max_value = column.max_value();
let gcd = super::gcd::find_gcd(column.iter().map(|val| val - min_value))
.filter(|gcd| gcd.get() > 1u64);
let normalized_column = normalize_column(column, min_value, gcd);
let codec_type = detect_codec(normalized_column, codecs)?;
Some(Header {
num_vals,
min_value,
max_value,
gcd,
codec_type,
})
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub(crate) struct U128Header {
pub num_vals: u32,
pub codec_type: U128FastFieldCodecType,
}
impl BinarySerializable for U128Header {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
VInt(self.num_vals as u64).serialize(writer)?;
self.codec_type.serialize(writer)?;
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let num_vals = VInt::deserialize(reader)?.0 as u32;
let codec_type = U128FastFieldCodecType::deserialize(reader)?;
Ok(U128Header {
num_vals,
codec_type,
})
}
}
fn normalize_column<C: ColumnValues>(
from_column: C,
min_value: u64,
gcd: Option<NonZeroU64>,
) -> impl ColumnValues {
let gcd = gcd.map(|gcd| gcd.get()).unwrap_or(1);
let mapping = StrictlyMonotonicMappingToInternalGCDBaseval::new(gcd, min_value);
monotonic_map_column(from_column, mapping)
}
impl BinarySerializable for Header {
fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
VInt(self.num_vals as u64).serialize(writer)?;
VInt(self.min_value).serialize(writer)?;
VInt(self.max_value - self.min_value).serialize(writer)?;
if let Some(gcd) = self.gcd {
VInt(gcd.get()).serialize(writer)?;
} else {
VInt(0u64).serialize(writer)?;
}
self.codec_type.serialize(writer)?;
Ok(())
}
fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let num_vals = VInt::deserialize(reader)?.0 as u32;
let min_value = VInt::deserialize(reader)?.0;
let amplitude = VInt::deserialize(reader)?.0;
let max_value = min_value + amplitude;
let gcd_u64 = VInt::deserialize(reader)?.0;
let codec_type = FastFieldCodecType::deserialize(reader)?;
Ok(Header {
num_vals,
min_value,
max_value,
gcd: NonZeroU64::new(gcd_u64),
codec_type,
})
}
}
/// Serializes u128 values with the compact space codec.
pub fn serialize_column_values_u128<F: Fn() -> I, I: Iterator<Item = u128>>(
iter_gen: F,
num_vals: u32,
output: &mut impl io::Write,
) -> io::Result<()> {
let header = U128Header {
num_vals,
codec_type: U128FastFieldCodecType::CompactSpace,
};
header.serialize(output)?;
let compressor = CompactSpaceCompressor::train_from(iter_gen(), num_vals);
compressor.compress_into(iter_gen(), output)?;
Ok(())
}
/// Serializes the column with the codec with the best estimate on the data.
pub fn serialize_column_values<T: MonotonicallyMappableToU64>(
typed_column: impl ColumnValues<T>,
codecs: &[FastFieldCodecType],
output: &mut impl io::Write,
) -> io::Result<()> {
let column = monotonic_map_column(typed_column, StrictlyMonotonicMappingToInternal::<T>::new());
let header = Header::compute_header(&column, codecs).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"Data cannot be serialized with this list of codec. {:?}",
codecs
),
)
})?;
header.serialize(output)?;
let normalized_column = header.normalize_column(column);
assert_eq!(normalized_column.min_value(), 0u64);
serialize_given_codec(normalized_column, header.codec_type, output)?;
Ok(())
}
fn detect_codec(
column: impl ColumnValues<u64>,
codecs: &[FastFieldCodecType],
) -> Option<FastFieldCodecType> {
let mut estimations = Vec::new();
for &codec in codecs {
let estimation_opt = match codec {
FastFieldCodecType::Bitpacked => BitpackedCodec::estimate(&column),
FastFieldCodecType::Linear => LinearCodec::estimate(&column),
FastFieldCodecType::BlockwiseLinear => BlockwiseLinearCodec::estimate(&column),
};
if let Some(estimation) = estimation_opt {
estimations.push((estimation, codec));
}
}
if let Some(broken_estimation) = estimations.iter().find(|estimation| estimation.0.is_nan()) {
warn!(
"broken estimation for fast field codec {:?}",
broken_estimation.1
);
}
// removing nan values for codecs with broken calculations, and max values which disables
// codecs
estimations.retain(|estimation| !estimation.0.is_nan() && estimation.0 != f32::MAX);
estimations.sort_by(|(score_left, _), (score_right, _)| score_left.total_cmp(score_right));
Some(estimations.first()?.1)
}
pub(crate) fn serialize_given_codec(
column: impl ColumnValues<u64>,
codec_type: FastFieldCodecType,
output: &mut impl io::Write,
) -> io::Result<()> {
match codec_type {
FastFieldCodecType::Bitpacked => {
BitpackedCodec::serialize(&column, output)?;
}
FastFieldCodecType::Linear => {
LinearCodec::serialize(&column, output)?;
}
FastFieldCodecType::BlockwiseLinear => {
BlockwiseLinearCodec::serialize(&column, output)?;
}
}
Ok(())
}
#[cfg(test)]
pub mod tests {
use std::sync::Arc;
use common::OwnedBytes;
use super::*;
use crate::column_values::{open_u64_mapped, VecColumn};
const ALL_CODEC_TYPES: [FastFieldCodecType; 3] = [
FastFieldCodecType::Bitpacked,
FastFieldCodecType::Linear,
FastFieldCodecType::BlockwiseLinear,
];
/// Helper function to serialize a column (autodetect from all codecs) and then open it
pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default>(
column: &[T],
) -> Arc<dyn ColumnValues<T>> {
let mut buffer = Vec::new();
serialize_column_values(&VecColumn::from(&column), &ALL_CODEC_TYPES, &mut buffer).unwrap();
open_u64_mapped(OwnedBytes::new(buffer)).unwrap()
}
#[test]
fn test_serialize_deserialize_u128_header() {
let original = U128Header {
num_vals: 11,
codec_type: U128FastFieldCodecType::CompactSpace,
};
let mut out = Vec::new();
original.serialize(&mut out).unwrap();
let restored = U128Header::deserialize(&mut &out[..]).unwrap();
assert_eq!(restored, original);
}
#[test]
fn test_serialize_deserialize() {
let original = [1u64, 5u64, 10u64];
let restored: Vec<u64> = serialize_and_load(&original[..]).iter().collect();
assert_eq!(&restored, &original[..]);
}
#[test]
fn test_fastfield_bool_size_bitwidth_1() {
let mut buffer = Vec::new();
let col = VecColumn::from(&[false, true][..]);
serialize_column_values(&col, &ALL_CODEC_TYPES, &mut buffer).unwrap();
// TODO put the header as a footer so that it serves as a padding.
// 5 bytes of header, 1 byte of value, 7 bytes of padding.
assert_eq!(buffer.len(), 5 + 1);
}
#[test]
fn test_fastfield_bool_bit_size_bitwidth_0() {
let mut buffer = Vec::new();
let col = VecColumn::from(&[true][..]);
serialize_column_values(&col, &ALL_CODEC_TYPES, &mut buffer).unwrap();
// 5 bytes of header, 0 bytes of value, 7 bytes of padding.
assert_eq!(buffer.len(), 5);
}
#[test]
fn test_fastfield_gcd() {
let mut buffer = Vec::new();
let vals: Vec<u64> = (0..80).map(|val| (val % 7) * 1_000u64).collect();
let col = VecColumn::from(&vals[..]);
serialize_column_values(&col, &[FastFieldCodecType::Bitpacked], &mut buffer).unwrap();
// Values are stored over 3 bits.
assert_eq!(buffer.len(), 7 + (3 * 80 / 8));
}
}

View File

@@ -1,309 +0,0 @@
use proptest::prelude::*;
use proptest::strategy::Strategy;
use proptest::{prop_oneof, proptest};
use super::bitpacked::BitpackedCodec;
use super::blockwise_linear::BlockwiseLinearCodec;
use super::linear::LinearCodec;
use super::serialize::Header;
pub(crate) fn create_and_validate<Codec: FastFieldCodec>(
data: &[u64],
name: &str,
) -> Option<(f32, f32)> {
let col = &VecColumn::from(data);
let header = Header::compute_header(col, &[Codec::CODEC_TYPE])?;
let normalized_col = header.normalize_column(col);
let estimation = Codec::estimate(&normalized_col)?;
let mut out = Vec::new();
let col = VecColumn::from(data);
serialize_column_values(&col, &[Codec::CODEC_TYPE], &mut out).unwrap();
let actual_compression = out.len() as f32 / (data.len() as f32 * 8.0);
let reader = super::open_u64_mapped::<u64>(OwnedBytes::new(out)).unwrap();
assert_eq!(reader.num_vals(), data.len() as u32);
for (doc, orig_val) in data.iter().copied().enumerate() {
let val = reader.get_val(doc as u32);
assert_eq!(
val, orig_val,
"val `{val}` does not match orig_val {orig_val:?}, in data set {name}, data `{data:?}`",
);
}
if !data.is_empty() {
let test_rand_idx = rand::thread_rng().gen_range(0..=data.len() - 1);
let expected_positions: Vec<u32> = data
.iter()
.enumerate()
.filter(|(_, el)| **el == data[test_rand_idx])
.map(|(pos, _)| pos as u32)
.collect();
let mut positions = Vec::new();
reader.get_docids_for_value_range(
data[test_rand_idx]..=data[test_rand_idx],
0..data.len() as u32,
&mut positions,
);
assert_eq!(expected_positions, positions);
}
Some((estimation, actual_compression))
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
#[test]
fn test_proptest_small_bitpacked(data in proptest::collection::vec(num_strategy(), 1..10)) {
create_and_validate::<BitpackedCodec>(&data, "proptest bitpacked");
}
#[test]
fn test_proptest_small_linear(data in proptest::collection::vec(num_strategy(), 1..10)) {
create_and_validate::<LinearCodec>(&data, "proptest linearinterpol");
}
#[test]
fn test_proptest_small_blockwise_linear(data in proptest::collection::vec(num_strategy(), 1..10)) {
create_and_validate::<BlockwiseLinearCodec>(&data, "proptest multilinearinterpol");
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(10))]
#[test]
fn test_proptest_large_bitpacked(data in proptest::collection::vec(num_strategy(), 1..6000)) {
create_and_validate::<BitpackedCodec>(&data, "proptest bitpacked");
}
#[test]
fn test_proptest_large_linear(data in proptest::collection::vec(num_strategy(), 1..6000)) {
create_and_validate::<LinearCodec>(&data, "proptest linearinterpol");
}
#[test]
fn test_proptest_large_blockwise_linear(data in proptest::collection::vec(num_strategy(), 1..6000)) {
create_and_validate::<BlockwiseLinearCodec>(&data, "proptest multilinearinterpol");
}
}
fn num_strategy() -> impl Strategy<Value = u64> {
prop_oneof![
1 => prop::num::u64::ANY.prop_map(|num| u64::MAX - (num % 10) ),
1 => prop::num::u64::ANY.prop_map(|num| num % 10 ),
20 => prop::num::u64::ANY,
]
}
pub fn get_codec_test_datasets() -> Vec<(Vec<u64>, &'static str)> {
let mut data_and_names = vec![];
let data = (10..=10_000_u64).collect::<Vec<_>>();
data_and_names.push((data, "simple monotonically increasing"));
data_and_names.push((
vec![5, 6, 7, 8, 9, 10, 99, 100],
"offset in linear interpol",
));
data_and_names.push((vec![5, 50, 3, 13, 1, 1000, 35], "rand small"));
data_and_names.push((vec![10], "single value"));
data_and_names.push((
vec![1572656989877777, 1170935903116329, 720575940379279, 0],
"overflow error",
));
data_and_names
}
fn test_codec<C: FastFieldCodec>() {
let codec_name = format!("{:?}", C::CODEC_TYPE);
for (data, dataset_name) in get_codec_test_datasets() {
let estimate_actual_opt: Option<(f32, f32)> =
tests::create_and_validate::<C>(&data, dataset_name);
let result = if let Some((estimate, actual)) = estimate_actual_opt {
format!("Estimate `{estimate}` Actual `{actual}`")
} else {
"Disabled".to_string()
};
println!("Codec {codec_name}, DataSet {dataset_name}, {result}");
}
}
#[test]
fn test_codec_bitpacking() {
test_codec::<BitpackedCodec>();
}
#[test]
fn test_codec_interpolation() {
test_codec::<LinearCodec>();
}
#[test]
fn test_codec_multi_interpolation() {
test_codec::<BlockwiseLinearCodec>();
}
use super::*;
#[test]
fn estimation_good_interpolation_case() {
let data = (10..=20000_u64).collect::<Vec<_>>();
let data: VecColumn = data.as_slice().into();
let linear_interpol_estimation = LinearCodec::estimate(&data).unwrap();
assert_le!(linear_interpol_estimation, 0.01);
let multi_linear_interpol_estimation = BlockwiseLinearCodec::estimate(&data).unwrap();
assert_le!(multi_linear_interpol_estimation, 0.2);
assert_lt!(linear_interpol_estimation, multi_linear_interpol_estimation);
let bitpacked_estimation = BitpackedCodec::estimate(&data).unwrap();
assert_lt!(linear_interpol_estimation, bitpacked_estimation);
}
#[test]
fn estimation_test_bad_interpolation_case() {
let data: &[u64] = &[200, 10, 10, 10, 10, 1000, 20];
let data: VecColumn = data.into();
let linear_interpol_estimation = LinearCodec::estimate(&data).unwrap();
assert_le!(linear_interpol_estimation, 0.34);
let bitpacked_estimation = BitpackedCodec::estimate(&data).unwrap();
assert_lt!(bitpacked_estimation, linear_interpol_estimation);
}
#[test]
fn estimation_prefer_bitpacked() {
let data = VecColumn::from(&[10, 10, 10, 10]);
let linear_interpol_estimation = LinearCodec::estimate(&data).unwrap();
let bitpacked_estimation = BitpackedCodec::estimate(&data).unwrap();
assert_lt!(bitpacked_estimation, linear_interpol_estimation);
}
#[test]
fn estimation_test_bad_interpolation_case_monotonically_increasing() {
let mut data: Vec<u64> = (201..=20000_u64).collect();
data.push(1_000_000);
let data: VecColumn = data.as_slice().into();
// in this case the linear interpolation can't in fact not be worse than bitpacking,
// but the estimator adds some threshold, which leads to estimated worse behavior
let linear_interpol_estimation = LinearCodec::estimate(&data).unwrap();
assert_le!(linear_interpol_estimation, 0.35);
let bitpacked_estimation = BitpackedCodec::estimate(&data).unwrap();
assert_le!(bitpacked_estimation, 0.32);
assert_le!(bitpacked_estimation, linear_interpol_estimation);
}
#[test]
fn test_fast_field_codec_type_to_code() {
let mut count_codec = 0;
for code in 0..=255 {
if let Some(codec_type) = FastFieldCodecType::from_code(code) {
assert_eq!(codec_type.to_code(), code);
count_codec += 1;
}
}
assert_eq!(count_codec, 3);
}
fn test_fastfield_gcd_i64_with_codec(
codec_type: FastFieldCodecType,
num_vals: usize,
) -> io::Result<()> {
let mut vals: Vec<i64> = (-4..=(num_vals as i64) - 5).map(|val| val * 1000).collect();
let mut buffer: Vec<u8> = Vec::new();
crate::column_values::serialize_column_values(
&VecColumn::from(&vals),
&[codec_type],
&mut buffer,
)?;
let buffer = OwnedBytes::new(buffer);
let column = crate::column_values::open_u64_mapped::<i64>(buffer.clone())?;
assert_eq!(column.get_val(0), -4000i64);
assert_eq!(column.get_val(1), -3000i64);
assert_eq!(column.get_val(2), -2000i64);
assert_eq!(column.max_value(), (num_vals as i64 - 5) * 1000);
assert_eq!(column.min_value(), -4000i64);
// Can't apply gcd
let mut buffer_without_gcd = Vec::new();
vals.pop();
vals.push(1001i64);
crate::column_values::serialize_column_values(
&VecColumn::from(&vals),
&[codec_type],
&mut buffer_without_gcd,
)?;
let buffer_without_gcd = OwnedBytes::new(buffer_without_gcd);
assert!(buffer_without_gcd.len() > buffer.len());
Ok(())
}
#[test]
fn test_fastfield_gcd_i64() -> io::Result<()> {
for &codec_type in &[
FastFieldCodecType::Bitpacked,
FastFieldCodecType::BlockwiseLinear,
FastFieldCodecType::Linear,
] {
test_fastfield_gcd_i64_with_codec(codec_type, 5500)?;
}
Ok(())
}
fn test_fastfield_gcd_u64_with_codec(
codec_type: FastFieldCodecType,
num_vals: usize,
) -> io::Result<()> {
let mut vals: Vec<u64> = (1..=num_vals).map(|i| i as u64 * 1000u64).collect();
let mut buffer: Vec<u8> = Vec::new();
crate::column_values::serialize_column_values(
&VecColumn::from(&vals),
&[codec_type],
&mut buffer,
)?;
let buffer = OwnedBytes::new(buffer);
let column = crate::column_values::open_u64_mapped::<u64>(buffer.clone())?;
assert_eq!(column.get_val(0), 1000u64);
assert_eq!(column.get_val(1), 2000u64);
assert_eq!(column.get_val(2), 3000u64);
assert_eq!(column.max_value(), num_vals as u64 * 1000);
assert_eq!(column.min_value(), 1000u64);
// Can't apply gcd
let mut buffer_without_gcd = Vec::new();
vals.pop();
vals.push(1001u64);
crate::column_values::serialize_column_values(
&VecColumn::from(&vals),
&[codec_type],
&mut buffer_without_gcd,
)?;
let buffer_without_gcd = OwnedBytes::new(buffer_without_gcd);
assert!(buffer_without_gcd.len() > buffer.len());
Ok(())
}
#[test]
fn test_fastfield_gcd_u64() -> io::Result<()> {
for &codec_type in &[
FastFieldCodecType::Bitpacked,
FastFieldCodecType::BlockwiseLinear,
FastFieldCodecType::Linear,
] {
test_fastfield_gcd_u64_with_codec(codec_type, 5500)?;
}
Ok(())
}
#[test]
pub fn test_fastfield2() {
let test_fastfield = crate::column_values::serialize_and_load(&[100u64, 200u64, 300u64]);
assert_eq!(test_fastfield.get_val(0), 100);
assert_eq!(test_fastfield.get_val(1), 200);
assert_eq!(test_fastfield.get_val(2), 300);
}

View File

@@ -1,160 +0,0 @@
use std::net::Ipv6Addr;
use crate::value::NumericalType;
use crate::InvalidData;
/// The column type represents the column type.
/// Any changes need to be propagated to `COLUMN_TYPES`.
#[derive(Hash, Eq, PartialEq, Debug, Clone, Copy, Ord, PartialOrd)]
#[repr(u8)]
pub enum ColumnType {
I64 = 0u8,
U64 = 1u8,
F64 = 2u8,
Bytes = 3u8,
Str = 4u8,
Bool = 5u8,
IpAddr = 6u8,
DateTime = 7u8,
}
// The order needs to match _exactly_ the order in the enum
const COLUMN_TYPES: [ColumnType; 8] = [
ColumnType::I64,
ColumnType::U64,
ColumnType::F64,
ColumnType::Bytes,
ColumnType::Str,
ColumnType::Bool,
ColumnType::IpAddr,
ColumnType::DateTime,
];
impl ColumnType {
pub fn to_code(self) -> u8 {
self as u8
}
pub(crate) fn try_from_code(code: u8) -> Result<ColumnType, InvalidData> {
COLUMN_TYPES.get(code as usize).copied().ok_or(InvalidData)
}
}
impl From<NumericalType> for ColumnType {
fn from(numerical_type: NumericalType) -> Self {
match numerical_type {
NumericalType::I64 => ColumnType::I64,
NumericalType::U64 => ColumnType::U64,
NumericalType::F64 => ColumnType::F64,
}
}
}
impl ColumnType {
pub fn numerical_type(&self) -> Option<NumericalType> {
match self {
ColumnType::I64 => Some(NumericalType::I64),
ColumnType::U64 => Some(NumericalType::U64),
ColumnType::F64 => Some(NumericalType::F64),
ColumnType::Bytes
| ColumnType::Str
| ColumnType::Bool
| ColumnType::IpAddr
| ColumnType::DateTime => None,
}
}
}
// TODO remove if possible
pub trait HasAssociatedColumnType: 'static + Send + Sync + Copy + PartialOrd {
fn column_type() -> ColumnType;
fn default_value() -> Self;
}
impl HasAssociatedColumnType for u64 {
fn column_type() -> ColumnType {
ColumnType::U64
}
fn default_value() -> Self {
0u64
}
}
impl HasAssociatedColumnType for i64 {
fn column_type() -> ColumnType {
ColumnType::I64
}
fn default_value() -> Self {
0i64
}
}
impl HasAssociatedColumnType for f64 {
fn column_type() -> ColumnType {
ColumnType::F64
}
fn default_value() -> Self {
Default::default()
}
}
impl HasAssociatedColumnType for bool {
fn column_type() -> ColumnType {
ColumnType::Bool
}
fn default_value() -> Self {
Default::default()
}
}
impl HasAssociatedColumnType for crate::DateTime {
fn column_type() -> ColumnType {
ColumnType::DateTime
}
fn default_value() -> Self {
Default::default()
}
}
impl HasAssociatedColumnType for Ipv6Addr {
fn column_type() -> ColumnType {
ColumnType::IpAddr
}
fn default_value() -> Self {
Ipv6Addr::from([0u8; 16])
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Cardinality;
#[test]
fn test_column_type_to_code() {
for (code, expected_column_type) in super::COLUMN_TYPES.iter().copied().enumerate() {
if let Ok(column_type) = ColumnType::try_from_code(code as u8) {
assert_eq!(column_type, expected_column_type);
}
}
for code in COLUMN_TYPES.len() as u8..=u8::MAX {
assert!(ColumnType::try_from_code(code as u8).is_err());
}
}
#[test]
fn test_cardinality_to_code() {
let mut num_cardinality = 0;
for code in u8::MIN..=u8::MAX {
if let Ok(cardinality) = Cardinality::try_from_code(code) {
assert_eq!(cardinality.to_code(), code);
num_cardinality += 1;
}
}
assert_eq!(num_cardinality, 3);
}
}

View File

@@ -1,73 +0,0 @@
use crate::InvalidData;
pub const VERSION_FOOTER_NUM_BYTES: usize = MAGIC_BYTES.len() + std::mem::size_of::<u32>();
/// We end the file by these 4 bytes just to somewhat identify that
/// this is indeed a columnar file.
const MAGIC_BYTES: [u8; 4] = [2, 113, 119, 066];
pub fn footer() -> [u8; VERSION_FOOTER_NUM_BYTES] {
let mut footer_bytes = [0u8; VERSION_FOOTER_NUM_BYTES];
footer_bytes[0..4].copy_from_slice(&Version::V1.to_bytes());
footer_bytes[4..8].copy_from_slice(&MAGIC_BYTES[..]);
footer_bytes
}
pub fn parse_footer(footer_bytes: [u8; VERSION_FOOTER_NUM_BYTES]) -> Result<Version, InvalidData> {
if footer_bytes[4..8] != MAGIC_BYTES {
return Err(InvalidData);
}
Version::try_from_bytes(footer_bytes[0..4].try_into().unwrap())
}
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
#[repr(u32)]
pub enum Version {
V1 = 1u32,
}
impl Version {
fn to_bytes(&self) -> [u8; 4] {
(*self as u32).to_le_bytes()
}
fn try_from_bytes(bytes: [u8; 4]) -> Result<Version, InvalidData> {
let code = u32::from_le_bytes(bytes);
match code {
1u32 => Ok(Version::V1),
_ => Err(InvalidData),
}
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::*;
#[test]
fn test_footer_dserialization() {
let parsed_version: Version = parse_footer(footer()).unwrap();
assert_eq!(Version::V1, parsed_version);
}
#[test]
fn test_version_serialization() {
let version_to_tests: Vec<u32> = [0, 1 << 8, 1 << 16, 1 << 24]
.iter()
.copied()
.flat_map(|offset| (0..255).map(move |el| el + offset))
.collect();
let mut valid_versions: HashSet<u32> = HashSet::default();
for &i in &version_to_tests {
let version_res = Version::try_from_bytes(i.to_le_bytes());
if let Ok(version) = version_res {
assert_eq!(version, Version::V1);
assert_eq!(version.to_bytes(), i.to_le_bytes());
valid_versions.insert(i);
}
}
assert_eq!(valid_versions.len(), 1);
}
}

View File

@@ -1,208 +0,0 @@
use std::collections::HashMap;
use std::io;
use crate::columnar::ColumnarReader;
use crate::dynamic_column::DynamicColumn;
use crate::ColumnType;
pub enum MergeDocOrder {
/// Columnar tables are simply stacked one above the other.
/// If the i-th columnar_readers has n_rows_i rows, then
/// in the resulting columnar,
/// rows [r0..n_row_0) contains the row of columnar_readers[0], in ordder
/// rows [n_row_0..n_row_0 + n_row_1 contains the row of columnar_readers[1], in order.
/// ..
Stack,
/// Some more complex mapping, that can interleaves rows from the different readers and
/// possibly drop rows.
Complex(()),
}
pub fn merge_columnar(
_columnar_readers: &[ColumnarReader],
mapping: MergeDocOrder,
_output: &mut impl io::Write,
) -> io::Result<()> {
match mapping {
MergeDocOrder::Stack => {
// implement me :)
todo!();
}
MergeDocOrder::Complex(_) => {
// for later
todo!();
}
}
}
/// Column types are grouped into different categories.
/// After merge, all columns belonging to the same category are coerced to
/// the same column type.
///
/// In practise, today, only Numerical colummns are coerced into one type today.
///
/// See also [README.md].
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
#[repr(u8)]
enum ColumnTypeCategory {
Bool,
Str,
Numerical,
DateTime,
Bytes,
IpAddr,
}
impl From<ColumnType> for ColumnTypeCategory {
fn from(column_type: ColumnType) -> Self {
match column_type {
ColumnType::I64 => ColumnTypeCategory::Numerical,
ColumnType::U64 => ColumnTypeCategory::Numerical,
ColumnType::F64 => ColumnTypeCategory::Numerical,
ColumnType::Bytes => ColumnTypeCategory::Bytes,
ColumnType::Str => ColumnTypeCategory::Str,
ColumnType::Bool => ColumnTypeCategory::Bool,
ColumnType::IpAddr => ColumnTypeCategory::IpAddr,
ColumnType::DateTime => ColumnTypeCategory::DateTime,
}
}
}
fn collect_columns(
columnar_readers: &[&ColumnarReader],
) -> io::Result<HashMap<String, HashMap<ColumnTypeCategory, Vec<DynamicColumn>>>> {
// Each column name may have multiple types of column associated.
// For merging we are interested in the same column type category since they can be merged.
let mut field_name_to_group: HashMap<String, HashMap<ColumnTypeCategory, Vec<DynamicColumn>>> =
HashMap::new();
for columnar_reader in columnar_readers {
let column_name_and_handle = columnar_reader.list_columns()?;
for (column_name, handle) in column_name_and_handle {
let column_type_to_handles = field_name_to_group
.entry(column_name.to_string())
.or_default();
let columns = column_type_to_handles
.entry(handle.column_type().into())
.or_default();
columns.push(handle.open()?);
}
}
normalize_columns(&mut field_name_to_group);
Ok(field_name_to_group)
}
/// Coerce numerical type columns to the same type
/// TODO rename to `coerce_columns`
fn normalize_columns(map: &mut HashMap<String, HashMap<ColumnTypeCategory, Vec<DynamicColumn>>>) {
for (_field_name, type_category_to_columns) in map.iter_mut() {
for (type_category, columns) in type_category_to_columns {
if type_category == &ColumnTypeCategory::Numerical {
let casted_columns = cast_to_common_numerical_column(&columns);
*columns = casted_columns;
}
}
}
}
/// Receives a list of columns of numerical types (u64, i64, f64)
///
/// Returns a list of `DynamicColumn` which are all of the same numerical type
fn cast_to_common_numerical_column(columns: &[DynamicColumn]) -> Vec<DynamicColumn> {
assert!(columns
.iter()
.all(|column| column.column_type().numerical_type().is_some()));
let coerce_to_i64: Vec<_> = columns
.iter()
.map(|column| column.clone().coerce_to_i64())
.collect();
if coerce_to_i64.iter().all(|column| column.is_some()) {
return coerce_to_i64
.into_iter()
.map(|column| column.unwrap())
.collect();
}
let coerce_to_u64: Vec<_> = columns
.iter()
.map(|column| column.clone().coerce_to_u64())
.collect();
if coerce_to_u64.iter().all(|column| column.is_some()) {
return coerce_to_u64
.into_iter()
.map(|column| column.unwrap())
.collect();
}
columns
.iter()
.map(|column| {
column
.clone()
.coerce_to_f64()
.expect("couldn't cast column to f64")
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ColumnarWriter;
#[test]
fn test_column_coercion() {
// i64 type
let columnar1 = {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_numerical(1u32, "numbers", 1i64);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(2, &mut buffer).unwrap();
ColumnarReader::open(buffer).unwrap()
};
// u64 type
let columnar2 = {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_numerical(1u32, "numbers", u64::MAX - 100);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(2, &mut buffer).unwrap();
ColumnarReader::open(buffer).unwrap()
};
// f64 type
let columnar3 = {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_numerical(1u32, "numbers", 30.5);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(2, &mut buffer).unwrap();
ColumnarReader::open(buffer).unwrap()
};
let column_map = collect_columns(&[&columnar1, &columnar2, &columnar3]).unwrap();
assert_eq!(column_map.len(), 1);
let cat_to_columns = column_map.get("numbers").unwrap();
assert_eq!(cat_to_columns.len(), 1);
let numerical = cat_to_columns.get(&ColumnTypeCategory::Numerical).unwrap();
assert!(numerical.iter().all(|column| column.is_f64()));
let column_map = collect_columns(&[&columnar1, &columnar1]).unwrap();
assert_eq!(column_map.len(), 1);
let cat_to_columns = column_map.get("numbers").unwrap();
assert_eq!(cat_to_columns.len(), 1);
let numerical = cat_to_columns.get(&ColumnTypeCategory::Numerical).unwrap();
assert!(numerical.iter().all(|column| column.is_i64()));
let column_map = collect_columns(&[&columnar2, &columnar2]).unwrap();
assert_eq!(column_map.len(), 1);
let cat_to_columns = column_map.get("numbers").unwrap();
assert_eq!(cat_to_columns.len(), 1);
let numerical = cat_to_columns.get(&ColumnTypeCategory::Numerical).unwrap();
assert!(numerical.iter().all(|column| column.is_u64()));
}
}

View File

@@ -1,10 +0,0 @@
mod column_type;
mod format_version;
mod merge;
mod reader;
mod writer;
pub use column_type::{ColumnType, HasAssociatedColumnType};
pub use merge::{merge_columnar, MergeDocOrder};
pub use reader::ColumnarReader;
pub use writer::ColumnarWriter;

View File

@@ -1,122 +0,0 @@
use std::{io, mem};
use common::file_slice::FileSlice;
use common::BinarySerializable;
use sstable::{Dictionary, RangeSSTable};
use crate::columnar::{format_version, ColumnType};
use crate::dynamic_column::DynamicColumnHandle;
fn io_invalid_data(msg: String) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, msg)
}
/// The ColumnarReader makes it possible to access a set of columns
/// associated to field names.
pub struct ColumnarReader {
column_dictionary: Dictionary<RangeSSTable>,
column_data: FileSlice,
}
impl ColumnarReader {
/// Opens a new Columnar file.
pub fn open<F>(file_slice: F) -> io::Result<ColumnarReader>
where FileSlice: From<F> {
Self::open_inner(file_slice.into())
}
fn open_inner(file_slice: FileSlice) -> io::Result<ColumnarReader> {
let (file_slice_without_sstable_len, footer_slice) = file_slice
.split_from_end(mem::size_of::<u64>() + format_version::VERSION_FOOTER_NUM_BYTES);
let footer_bytes = footer_slice.read_bytes()?;
let (mut sstable_len_bytes, version_footer_bytes) =
footer_bytes.rsplit(format_version::VERSION_FOOTER_NUM_BYTES);
let version_footer_bytes: [u8; format_version::VERSION_FOOTER_NUM_BYTES] =
version_footer_bytes.as_slice().try_into().unwrap();
let _version = format_version::parse_footer(version_footer_bytes)?;
let sstable_len = u64::deserialize(&mut sstable_len_bytes)?;
let (column_data, sstable) =
file_slice_without_sstable_len.split_from_end(sstable_len as usize);
let column_dictionary = Dictionary::open(sstable)?;
Ok(ColumnarReader {
column_dictionary,
column_data,
})
}
// TODO Add unit tests
pub fn list_columns(&self) -> io::Result<Vec<(String, DynamicColumnHandle)>> {
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: ColumnType = ColumnType::try_from_code(column_code)
.map_err(|_| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
let range = stream.value().clone();
let column_name =
// The last two bytes are respectively the 0u8 separator and the column_type.
String::from_utf8_lossy(&key_bytes[..key_bytes.len() - 2]).to_string();
let file_slice = self
.column_data
.slice(range.start as usize..range.end as usize);
let column_handle = DynamicColumnHandle {
file_slice,
column_type,
};
results.push((column_name, column_handle));
}
Ok(results)
}
/// Get all columns for the given column name.
///
/// There can be more than one column associated to a given column name, provided they have
/// different types.
// TODO fix ugly API
pub fn read_columns(&self, column_name: &str) -> io::Result<Vec<DynamicColumnHandle>> {
// Each column is a associated to a given `column_key`,
// that starts by `column_name\0column_header`.
//
// Listing the columns associated to the given column name is therefore equivalent to
// listing `column_key` with the prefix `column_name\0`.
//
// This is in turn equivalent to searching for the range
// `[column_name,\0`..column_name\1)`.
// TODO can we get some more generic `prefix(..)` logic in the dictioanry.
let mut start_key = column_name.to_string();
start_key.push('\0');
let mut end_key = column_name.to_string();
end_key.push(1u8 as char);
let mut stream = self
.column_dictionary
.range()
.ge(start_key.as_bytes())
.lt(end_key.as_bytes())
.into_stream()?;
let mut results = Vec::new();
while stream.advance() {
let key_bytes: &[u8] = stream.key();
assert!(key_bytes.starts_with(start_key.as_bytes()));
let column_code: u8 = key_bytes.last().cloned().unwrap();
let column_type = ColumnType::try_from_code(column_code)
.map_err(|_| io_invalid_data(format!("Unknown column code `{column_code}`")))?;
let range = stream.value().clone();
let file_slice = self
.column_data
.slice(range.start as usize..range.end as usize);
let dynamic_column_handle = DynamicColumnHandle {
file_slice,
column_type,
};
results.push(dynamic_column_handle);
}
Ok(results)
}
/// Return the number of columns in the columnar.
pub fn num_columns(&self) -> usize {
self.column_dictionary.num_terms()
}
}

View File

@@ -1,753 +0,0 @@
mod column_operation;
mod column_writers;
mod serializer;
mod value_index;
use std::io;
use std::net::Ipv6Addr;
use column_operation::ColumnOperation;
use common::CountingWriter;
use serializer::ColumnarSerializer;
use stacker::{Addr, ArenaHashMap, MemoryArena};
use crate::column_index::SerializableColumnIndex;
use crate::column_values::{
ColumnValues, MonotonicallyMappableToU128, MonotonicallyMappableToU64, VecColumn,
};
use crate::columnar::column_type::ColumnType;
use crate::columnar::writer::column_writers::{
ColumnWriter, NumericalColumnWriter, StrOrBytesColumnWriter,
};
use crate::columnar::writer::value_index::{IndexBuilder, PreallocatedIndexBuilders};
use crate::dictionary::{DictionaryBuilder, TermIdMapping, UnorderedId};
use crate::value::{Coerce, NumericalType, NumericalValue};
use crate::{Cardinality, RowId};
/// This is a set of buffers that are used to temporarily write the values into before passing them
/// to the fast field codecs.
#[derive(Default)]
struct SpareBuffers {
value_index_builders: PreallocatedIndexBuilders,
i64_values: Vec<i64>,
u64_values: Vec<u64>,
f64_values: Vec<f64>,
bool_values: Vec<bool>,
ip_addr_values: Vec<Ipv6Addr>,
}
/// Makes it possible to create a new columnar.
///
/// ```rust
/// use tantivy_columnar::ColumnarWriter;
///
/// let mut columnar_writer = ColumnarWriter::default();
/// columnar_writer.record_str(0u32 /* doc id */, "product_name", "Red backpack");
/// columnar_writer.record_numerical(0u32 /* doc id */, "price", 10u64);
/// columnar_writer.record_str(1u32 /* doc id */, "product_name", "Apple");
/// columnar_writer.record_numerical(0u32 /* doc id */, "price", 10.5f64); //< uh oh we ended up mixing integer and floats.
/// let mut wrt: Vec<u8> = Vec::new();
/// columnar_writer.serialize(2u32, &mut wrt).unwrap();
/// ```
pub struct ColumnarWriter {
numerical_field_hash_map: ArenaHashMap,
datetime_field_hash_map: ArenaHashMap,
bool_field_hash_map: ArenaHashMap,
ip_addr_field_hash_map: ArenaHashMap,
bytes_field_hash_map: ArenaHashMap,
str_field_hash_map: ArenaHashMap,
arena: MemoryArena,
// Dictionaries used to store dictionary-encoded values.
dictionaries: Vec<DictionaryBuilder>,
buffers: SpareBuffers,
}
impl Default for ColumnarWriter {
fn default() -> Self {
ColumnarWriter {
numerical_field_hash_map: ArenaHashMap::new(10_000),
bool_field_hash_map: ArenaHashMap::new(10_000),
ip_addr_field_hash_map: ArenaHashMap::new(10_000),
bytes_field_hash_map: ArenaHashMap::new(10_000),
str_field_hash_map: ArenaHashMap::new(10_000),
datetime_field_hash_map: ArenaHashMap::new(10_000),
dictionaries: Vec::new(),
arena: MemoryArena::default(),
buffers: SpareBuffers::default(),
}
}
}
#[inline]
fn mutate_or_create_column<V, TMutator>(
arena_hash_map: &mut ArenaHashMap,
column_name: &str,
updater: TMutator,
) where
V: Copy + 'static,
TMutator: FnMut(Option<V>) -> V,
{
assert!(
!column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte"
);
arena_hash_map.mutate_or_create(column_name.as_bytes(), updater);
}
impl ColumnarWriter {
pub fn mem_usage(&self) -> usize {
// TODO add dictionary builders.
self.arena.mem_usage()
+ self.numerical_field_hash_map.mem_usage()
+ self.bool_field_hash_map.mem_usage()
+ self.bytes_field_hash_map.mem_usage()
+ self.str_field_hash_map.mem_usage()
+ self.ip_addr_field_hash_map.mem_usage()
+ self.datetime_field_hash_map.mem_usage()
}
pub fn record_column_type(&mut self, column_name: &str, column_type: ColumnType) {
match column_type {
ColumnType::Str | ColumnType::Bytes => {
let (hash_map, dictionaries) = (
if column_type == ColumnType::Str {
&mut self.str_field_hash_map
} else {
&mut self.bytes_field_hash_map
},
&mut self.dictionaries,
);
mutate_or_create_column(
hash_map,
column_name,
|column_opt: Option<StrOrBytesColumnWriter>| {
if let Some(column_writer) = column_opt {
column_writer
} else {
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::with_dictionary_id(dictionary_id)
}
},
);
}
ColumnType::Bool => {
mutate_or_create_column(
&mut self.bool_field_hash_map,
column_name,
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
);
}
ColumnType::DateTime => {
mutate_or_create_column(
&mut self.datetime_field_hash_map,
column_name,
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
);
}
ColumnType::I64 | ColumnType::F64 | ColumnType::U64 => {
let numerical_type = column_type.numerical_type().unwrap();
mutate_or_create_column(
&mut self.numerical_field_hash_map,
column_name,
|column_opt: Option<NumericalColumnWriter>| {
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
column.force_numerical_type(numerical_type);
column
},
);
}
ColumnType::IpAddr => mutate_or_create_column(
&mut self.ip_addr_field_hash_map,
column_name,
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
),
}
}
pub fn force_numerical_type(&mut self, column_name: &str, numerical_type: NumericalType) {
mutate_or_create_column(
&mut self.numerical_field_hash_map,
column_name,
|column_opt: Option<NumericalColumnWriter>| {
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
column.force_numerical_type(numerical_type);
column
},
);
}
pub fn record_numerical<T: Into<NumericalValue> + Copy>(
&mut self,
doc: RowId,
column_name: &str,
numerical_value: T,
) {
let (hash_map, arena) = (&mut self.numerical_field_hash_map, &mut self.arena);
mutate_or_create_column(
hash_map,
column_name,
|column_opt: Option<NumericalColumnWriter>| {
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
column.record_numerical_value(doc, numerical_value.into(), arena);
column
},
);
}
pub fn record_ip_addr(&mut self, doc: RowId, column_name: &str, ip_addr: Ipv6Addr) {
assert!(
!column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte"
);
let (hash_map, arena) = (&mut self.ip_addr_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, ip_addr, arena);
column
},
);
}
pub fn record_bool(&mut self, doc: RowId, column_name: &str, val: bool) {
let (hash_map, arena) = (&mut self.bool_field_hash_map, &mut self.arena);
mutate_or_create_column(hash_map, column_name, |column_opt: Option<ColumnWriter>| {
let mut column: ColumnWriter = column_opt.unwrap_or_default();
column.record(doc, val, arena);
column
});
}
pub fn record_datetime(&mut self, doc: RowId, column_name: &str, datetime: crate::DateTime) {
let (hash_map, arena) = (&mut self.datetime_field_hash_map, &mut self.arena);
mutate_or_create_column(hash_map, column_name, |column_opt: Option<ColumnWriter>| {
let mut column: ColumnWriter = column_opt.unwrap_or_default();
column.record(doc, NumericalValue::I64(datetime.timestamp_micros), arena);
column
});
}
pub fn record_str(&mut self, doc: RowId, column_name: &str, value: &str) {
let (hash_map, arena, dictionaries) = (
&mut self.str_field_hash_map,
&mut self.arena,
&mut self.dictionaries,
);
hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<StrOrBytesColumnWriter>| {
let mut column: StrOrBytesColumnWriter = column_opt.unwrap_or_else(|| {
// Each column has its own dictionary
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::with_dictionary_id(dictionary_id)
});
column.record_bytes(doc, value.as_bytes(), dictionaries, arena);
column
},
);
}
pub fn record_bytes(&mut self, doc: RowId, column_name: &str, value: &[u8]) {
assert!(
!column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte"
);
let (hash_map, arena, dictionaries) = (
&mut self.bytes_field_hash_map,
&mut self.arena,
&mut self.dictionaries,
);
hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<StrOrBytesColumnWriter>| {
let mut column: StrOrBytesColumnWriter = column_opt.unwrap_or_else(|| {
// Each column has its own dictionary
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::with_dictionary_id(dictionary_id)
});
column.record_bytes(doc, value, dictionaries, arena);
column
},
);
}
pub fn serialize(&mut self, num_docs: RowId, wrt: &mut dyn io::Write) -> io::Result<()> {
let mut serializer = ColumnarSerializer::new(wrt);
let mut columns: Vec<(&[u8], ColumnType, Addr)> = self
.numerical_field_hash_map
.iter()
.map(|(column_name, addr, _)| {
let numerical_column_writer: NumericalColumnWriter =
self.numerical_field_hash_map.read(addr);
let column_type = numerical_column_writer.numerical_type().into();
(column_name, column_type, addr)
})
.collect();
columns.extend(
self.bytes_field_hash_map
.iter()
.map(|(term, addr, _)| (term, ColumnType::Bytes, addr)),
);
columns.extend(
self.str_field_hash_map
.iter()
.map(|(column_name, addr, _)| (column_name, ColumnType::Str, addr)),
);
columns.extend(
self.bool_field_hash_map
.iter()
.map(|(column_name, addr, _)| (column_name, ColumnType::Bool, addr)),
);
columns.extend(
self.ip_addr_field_hash_map
.iter()
.map(|(column_name, addr, _)| (column_name, ColumnType::IpAddr, addr)),
);
columns.extend(
self.datetime_field_hash_map
.iter()
.map(|(column_name, addr, _)| (column_name, ColumnType::DateTime, addr)),
);
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 mut symbol_byte_buffer: Vec<u8> = Vec::new();
for (column_name, column_type, addr) in columns {
match column_type {
ColumnType::Bool | ColumnType::DateTime => {
let column_writer: ColumnWriter = if column_type == ColumnType::Bool {
self.bool_field_hash_map.read(addr)
} else {
self.datetime_field_hash_map.read(addr)
};
let cardinality = column_writer.get_cardinality(num_docs);
let mut column_serializer =
serializer.serialize_column(column_name, ColumnType::Bool);
serialize_bool_column(
cardinality,
num_docs,
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
&mut column_serializer,
)?;
}
ColumnType::IpAddr => {
let column_writer: ColumnWriter = self.ip_addr_field_hash_map.read(addr);
let cardinality = column_writer.get_cardinality(num_docs);
let mut column_serializer =
serializer.serialize_column(column_name, ColumnType::IpAddr);
serialize_ip_addr_column(
cardinality,
num_docs,
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
&mut column_serializer,
)?;
}
ColumnType::Bytes | ColumnType::Str => {
let str_or_bytes_column_writer: StrOrBytesColumnWriter =
if column_type == ColumnType::Bytes {
self.bytes_field_hash_map.read(addr)
} else {
self.str_field_hash_map.read(addr)
};
let dictionary_builder =
&dictionaries[str_or_bytes_column_writer.dictionary_id as usize];
let cardinality = str_or_bytes_column_writer
.column_writer
.get_cardinality(num_docs);
let mut column_serializer =
serializer.serialize_column(column_name, column_type);
serialize_bytes_or_str_column(
cardinality,
num_docs,
dictionary_builder,
str_or_bytes_column_writer
.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
&mut column_serializer,
)?;
}
ColumnType::I64 | ColumnType::F64 | ColumnType::U64 => {
let numerical_column_writer: NumericalColumnWriter =
self.numerical_field_hash_map.read(addr);
let numerical_type = column_type.numerical_type().unwrap();
let cardinality = numerical_column_writer.cardinality(num_docs);
let mut column_serializer =
serializer.serialize_column(column_name, ColumnType::from(numerical_type));
serialize_numerical_column(
cardinality,
num_docs,
numerical_type,
numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
&mut column_serializer,
)?;
}
};
}
serializer.finalize()?;
Ok(())
}
}
fn serialize_bytes_or_str_column(
cardinality: Cardinality,
num_docs: RowId,
dictionary_builder: &DictionaryBuilder,
operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>,
buffers: &mut SpareBuffers,
wrt: impl io::Write,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
u64_values,
..
} = buffers;
let mut counting_writer = CountingWriter::wrap(wrt);
let term_id_mapping: TermIdMapping = dictionary_builder.serialize(&mut counting_writer)?;
let dictionary_num_bytes: u32 = counting_writer.written_bytes() as u32;
let mut wrt = counting_writer.finish();
let operation_iterator = operation_it.map(|symbol: ColumnOperation<UnorderedId>| {
// We map unordered ids to ordered ids.
match symbol {
ColumnOperation::Value(unordered_id) => {
let ordered_id = term_id_mapping.to_ord(unordered_id);
ColumnOperation::Value(ordered_id.0 as u64)
}
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
}
});
send_to_serialize_column_mappable_to_u64(
operation_iterator,
cardinality,
num_docs,
value_index_builders,
u64_values,
&mut wrt,
)?;
wrt.write_all(&dictionary_num_bytes.to_le_bytes()[..])?;
Ok(())
}
fn serialize_numerical_column(
cardinality: Cardinality,
num_docs: RowId,
numerical_type: NumericalType,
op_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
buffers: &mut SpareBuffers,
wrt: &mut impl io::Write,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
u64_values,
i64_values,
f64_values,
..
} = buffers;
match numerical_type {
NumericalType::I64 => {
send_to_serialize_column_mappable_to_u64(
coerce_numerical_symbol::<i64>(op_iterator),
cardinality,
num_docs,
value_index_builders,
i64_values,
wrt,
)?;
}
NumericalType::U64 => {
send_to_serialize_column_mappable_to_u64(
coerce_numerical_symbol::<u64>(op_iterator),
cardinality,
num_docs,
value_index_builders,
u64_values,
wrt,
)?;
}
NumericalType::F64 => {
send_to_serialize_column_mappable_to_u64(
coerce_numerical_symbol::<f64>(op_iterator),
cardinality,
num_docs,
value_index_builders,
f64_values,
wrt,
)?;
}
};
Ok(())
}
fn serialize_bool_column(
cardinality: Cardinality,
num_docs: RowId,
column_operations_it: impl Iterator<Item = ColumnOperation<bool>>,
buffers: &mut SpareBuffers,
wrt: &mut impl io::Write,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
bool_values,
..
} = buffers;
send_to_serialize_column_mappable_to_u64(
column_operations_it,
cardinality,
num_docs,
value_index_builders,
bool_values,
wrt,
)?;
Ok(())
}
fn serialize_ip_addr_column(
cardinality: Cardinality,
num_docs: RowId,
column_operations_it: impl Iterator<Item = ColumnOperation<Ipv6Addr>>,
buffers: &mut SpareBuffers,
wrt: &mut impl io::Write,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
ip_addr_values,
..
} = buffers;
send_to_serialize_column_mappable_to_u128(
column_operations_it,
cardinality,
num_docs,
value_index_builders,
ip_addr_values,
wrt,
)?;
Ok(())
}
fn send_to_serialize_column_mappable_to_u128<
T: Copy + std::fmt::Debug + Send + Sync + MonotonicallyMappableToU128 + PartialOrd,
>(
op_iterator: impl Iterator<Item = ColumnOperation<T>>,
cardinality: Cardinality,
num_docs: RowId,
value_index_builders: &mut PreallocatedIndexBuilders,
values: &mut Vec<T>,
mut wrt: impl io::Write,
) -> io::Result<()>
where
for<'a> VecColumn<'a, T>: ColumnValues<T>,
{
values.clear();
// TODO: split index and values
let serializable_column_index = match cardinality {
Cardinality::Full => {
consume_operation_iterator(
op_iterator,
value_index_builders.borrow_required_index_builder(),
values,
);
SerializableColumnIndex::Full
}
Cardinality::Optional => {
let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_docs);
SerializableColumnIndex::Optional(Box::new(optional_index))
}
Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_docs);
SerializableColumnIndex::Multivalued(Box::new(multivalued_index))
}
};
crate::column::serialize_column_mappable_to_u128(
serializable_column_index,
|| values.iter().cloned(),
values.len() as u32,
&mut wrt,
)?;
Ok(())
}
fn send_to_serialize_column_mappable_to_u64<
T: Copy + Default + std::fmt::Debug + Send + Sync + MonotonicallyMappableToU64 + PartialOrd,
>(
op_iterator: impl Iterator<Item = ColumnOperation<T>>,
cardinality: Cardinality,
num_docs: RowId,
value_index_builders: &mut PreallocatedIndexBuilders,
values: &mut Vec<T>,
mut wrt: impl io::Write,
) -> io::Result<()>
where
for<'a> VecColumn<'a, T>: ColumnValues<T>,
{
values.clear();
let serializable_column_index = match cardinality {
Cardinality::Full => {
consume_operation_iterator(
op_iterator,
value_index_builders.borrow_required_index_builder(),
values,
);
SerializableColumnIndex::Full
}
Cardinality::Optional => {
let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_docs);
SerializableColumnIndex::Optional(Box::new(optional_index))
}
Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_docs);
SerializableColumnIndex::Multivalued(Box::new(multivalued_index))
}
};
crate::column::serialize_column_mappable_to_u64(
serializable_column_index,
&VecColumn::from(&values[..]),
&mut wrt,
)?;
Ok(())
}
fn coerce_numerical_symbol<T>(
operation_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
) -> impl Iterator<Item = ColumnOperation<T>>
where T: Coerce {
operation_iterator.map(|symbol| match symbol {
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
ColumnOperation::Value(numerical_value) => {
ColumnOperation::Value(Coerce::coerce(numerical_value))
}
})
}
fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>(
operation_iterator: impl Iterator<Item = ColumnOperation<T>>,
index_builder: &mut TIndexBuilder,
values: &mut Vec<T>,
) {
for symbol in operation_iterator {
match symbol {
ColumnOperation::NewDoc(doc) => {
index_builder.record_row(doc);
}
ColumnOperation::Value(value) => {
index_builder.record_value();
values.push(value);
}
}
}
}
#[cfg(test)]
mod tests {
use stacker::MemoryArena;
use crate::columnar::writer::column_operation::ColumnOperation;
use crate::{Cardinality, NumericalValue};
#[test]
fn test_column_writer_required_simple() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(14i64), &mut arena);
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Full);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 6);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(14i64))
));
assert!(matches!(symbols[2], ColumnOperation::NewDoc(1u32)));
assert!(matches!(
symbols[3],
ColumnOperation::Value(NumericalValue::I64(15i64))
));
assert!(matches!(symbols[4], ColumnOperation::NewDoc(2u32)));
assert!(matches!(
symbols[5],
ColumnOperation::Value(NumericalValue::I64(-16i64))
));
}
#[test]
fn test_column_writer_optional_cardinality_missing_first() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Optional);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 4);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(1u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(15i64))
));
assert!(matches!(symbols[2], ColumnOperation::NewDoc(2u32)));
assert!(matches!(
symbols[3],
ColumnOperation::Value(NumericalValue::I64(-16i64))
));
}
#[test]
fn test_column_writer_optional_cardinality_missing_last() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(15i64), &mut arena);
assert_eq!(column_writer.get_cardinality(2), Cardinality::Optional);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 2);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(15i64))
));
}
#[test]
fn test_column_writer_multivalued() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(16i64), &mut arena);
column_writer.record(0u32, NumericalValue::from(17i64), &mut arena);
assert_eq!(column_writer.get_cardinality(1), Cardinality::Multivalued);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 3);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(16i64))
));
assert!(matches!(
symbols[2],
ColumnOperation::Value(NumericalValue::I64(17i64))
));
}
}

View File

@@ -3,11 +3,11 @@ use std::io;
use fnv::FnvHashMap; use fnv::FnvHashMap;
use sstable::SSTable; use sstable::SSTable;
pub(crate) struct TermIdMapping { pub(crate) struct IdMapping {
unordered_to_ord: Vec<OrderedId>, unordered_to_ord: Vec<OrderedId>,
} }
impl TermIdMapping { impl IdMapping {
pub fn to_ord(&self, unordered: UnorderedId) -> OrderedId { pub fn to_ord(&self, unordered: UnorderedId) -> OrderedId {
self.unordered_to_ord[unordered.0 as usize] self.unordered_to_ord[unordered.0 as usize]
} }
@@ -48,7 +48,7 @@ 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<TermIdMapping> { pub fn serialize<'a, W: io::Write + 'a>(&self, wrt: &mut W) -> io::Result<IdMapping> {
let mut terms: Vec<(&[u8], UnorderedId)> = let mut terms: Vec<(&[u8], UnorderedId)> =
self.dict.iter().map(|(k, v)| (k.as_slice(), *v)).collect(); self.dict.iter().map(|(k, v)| (k.as_slice(), *v)).collect();
terms.sort_unstable_by_key(|(key, _)| *key); terms.sort_unstable_by_key(|(key, _)| *key);
@@ -61,7 +61,7 @@ impl DictionaryBuilder {
unordered_to_ord[unordered_id.0 as usize] = ordered_id; unordered_to_ord[unordered_id.0 as usize] = ordered_id;
} }
sstable_builder.finish()?; sstable_builder.finish()?;
Ok(TermIdMapping { unordered_to_ord }) Ok(IdMapping { unordered_to_ord })
} }
} }

View File

@@ -1,241 +0,0 @@
use std::io;
use std::net::Ipv6Addr;
use std::sync::Arc;
use common::file_slice::FileSlice;
use common::{HasLen, OwnedBytes};
use crate::column::{BytesColumn, Column, StrColumn};
use crate::column_values::{monotonic_map_column, StrictlyMonotonicFn};
use crate::columnar::ColumnType;
use crate::{DateTime, NumericalType};
#[derive(Clone)]
pub enum DynamicColumn {
Bool(Column<bool>),
I64(Column<i64>),
U64(Column<u64>),
F64(Column<f64>),
IpAddr(Column<Ipv6Addr>),
DateTime(Column<DateTime>),
Bytes(BytesColumn),
Str(StrColumn),
}
impl DynamicColumn {
pub fn column_type(&self) -> ColumnType {
match self {
DynamicColumn::Bool(_) => ColumnType::Bool,
DynamicColumn::I64(_) => ColumnType::I64,
DynamicColumn::U64(_) => ColumnType::U64,
DynamicColumn::F64(_) => ColumnType::F64,
DynamicColumn::IpAddr(_) => ColumnType::IpAddr,
DynamicColumn::DateTime(_) => ColumnType::DateTime,
DynamicColumn::Bytes(_) => ColumnType::Bytes,
DynamicColumn::Str(_) => ColumnType::Str,
}
}
pub fn is_numerical(&self) -> bool {
self.column_type().numerical_type().is_some()
}
pub fn is_f64(&self) -> bool {
self.column_type().numerical_type() == Some(NumericalType::F64)
}
pub fn is_i64(&self) -> bool {
self.column_type().numerical_type() == Some(NumericalType::I64)
}
pub fn is_u64(&self) -> bool {
self.column_type().numerical_type() == Some(NumericalType::U64)
}
pub fn coerce_to_f64(self) -> Option<DynamicColumn> {
match self {
DynamicColumn::I64(column) => Some(DynamicColumn::F64(Column {
idx: column.idx,
values: Arc::new(monotonic_map_column(column.values, MapI64ToF64)),
})),
DynamicColumn::U64(column) => Some(DynamicColumn::F64(Column {
idx: column.idx,
values: Arc::new(monotonic_map_column(column.values, MapU64ToF64)),
})),
DynamicColumn::F64(_) => Some(self),
_ => None,
}
}
pub fn coerce_to_i64(self) -> Option<DynamicColumn> {
match self {
DynamicColumn::U64(column) => {
if column.max_value() > i64::MAX as u64 {
return None;
}
Some(DynamicColumn::I64(Column {
idx: column.idx,
values: Arc::new(monotonic_map_column(column.values, MapU64ToI64)),
}))
}
DynamicColumn::I64(_) => Some(self),
_ => None,
}
}
pub fn coerce_to_u64(self) -> Option<DynamicColumn> {
match self {
DynamicColumn::I64(column) => {
if column.min_value() < 0 {
return None;
}
Some(DynamicColumn::U64(Column {
idx: column.idx,
values: Arc::new(monotonic_map_column(column.values, MapI64ToU64)),
}))
}
DynamicColumn::U64(_) => Some(self),
_ => None,
}
}
}
struct MapI64ToF64;
impl StrictlyMonotonicFn<i64, f64> for MapI64ToF64 {
#[inline(always)]
fn mapping(&self, inp: i64) -> f64 {
inp as f64
}
#[inline(always)]
fn inverse(&self, out: f64) -> i64 {
out as i64
}
}
struct MapU64ToF64;
impl StrictlyMonotonicFn<u64, f64> for MapU64ToF64 {
#[inline(always)]
fn mapping(&self, inp: u64) -> f64 {
inp as f64
}
#[inline(always)]
fn inverse(&self, out: f64) -> u64 {
out as u64
}
}
struct MapU64ToI64;
impl StrictlyMonotonicFn<u64, i64> for MapU64ToI64 {
#[inline(always)]
fn mapping(&self, inp: u64) -> i64 {
inp as i64
}
#[inline(always)]
fn inverse(&self, out: i64) -> u64 {
out as u64
}
}
struct MapI64ToU64;
impl StrictlyMonotonicFn<i64, u64> for MapI64ToU64 {
#[inline(always)]
fn mapping(&self, inp: i64) -> u64 {
inp as u64
}
#[inline(always)]
fn inverse(&self, out: u64) -> i64 {
out as i64
}
}
macro_rules! static_dynamic_conversions {
($typ:ty, $enum_name:ident) => {
impl Into<Option<$typ>> for DynamicColumn {
fn into(self) -> Option<$typ> {
if let DynamicColumn::$enum_name(col) = self {
Some(col)
} else {
None
}
}
}
impl From<$typ> for DynamicColumn {
fn from(typed_column: $typ) -> Self {
DynamicColumn::$enum_name(typed_column)
}
}
};
}
static_dynamic_conversions!(Column<bool>, Bool);
static_dynamic_conversions!(Column<u64>, U64);
static_dynamic_conversions!(Column<i64>, I64);
static_dynamic_conversions!(Column<f64>, F64);
static_dynamic_conversions!(Column<crate::DateTime>, DateTime);
static_dynamic_conversions!(StrColumn, Str);
static_dynamic_conversions!(BytesColumn, Bytes);
static_dynamic_conversions!(Column<Ipv6Addr>, IpAddr);
#[derive(Clone)]
pub struct DynamicColumnHandle {
pub(crate) file_slice: FileSlice,
pub(crate) column_type: ColumnType,
}
impl DynamicColumnHandle {
// TODO rename load
pub fn open(&self) -> io::Result<DynamicColumn> {
let column_bytes: OwnedBytes = self.file_slice.read_bytes()?;
self.open_internal(column_bytes)
}
// TODO rename load_async
pub async fn open_async(&self) -> io::Result<DynamicColumn> {
let column_bytes: OwnedBytes = self.file_slice.read_bytes_async().await?;
self.open_internal(column_bytes)
}
/// Returns the `u64` fast field reader reader associated with `fields` of types
/// Str, u64, i64, f64, or datetime.
///
/// If not, the fastfield reader will returns the u64-value associated with the original
/// FastValue.
pub fn open_u64_lenient(&self) -> io::Result<Option<Column<u64>>> {
let column_bytes = self.file_slice.read_bytes()?;
match self.column_type {
ColumnType::Str | ColumnType::Bytes => {
let column: BytesColumn = crate::column::open_column_bytes(column_bytes)?;
Ok(Some(column.term_ord_column))
}
ColumnType::Bool => Ok(None),
ColumnType::IpAddr => Ok(None),
ColumnType::I64 | ColumnType::U64 | ColumnType::F64 | ColumnType::DateTime => {
let column = crate::column::open_column_u64::<u64>(column_bytes)?;
Ok(Some(column))
}
}
}
fn open_internal(&self, column_bytes: OwnedBytes) -> io::Result<DynamicColumn> {
let dynamic_column: DynamicColumn = match self.column_type {
ColumnType::Bytes => {
crate::column::open_column_bytes::<BytesColumn>(column_bytes)?.into()
}
ColumnType::Str => crate::column::open_column_bytes::<StrColumn>(column_bytes)?.into(),
ColumnType::I64 => crate::column::open_column_u64::<i64>(column_bytes)?.into(),
ColumnType::U64 => crate::column::open_column_u64::<u64>(column_bytes)?.into(),
ColumnType::F64 => crate::column::open_column_u64::<f64>(column_bytes)?.into(),
ColumnType::Bool => crate::column::open_column_u64::<bool>(column_bytes)?.into(),
ColumnType::IpAddr => crate::column::open_column_u128::<Ipv6Addr>(column_bytes)?.into(),
ColumnType::DateTime => {
crate::column::open_column_u64::<crate::DateTime>(column_bytes)?.into()
}
};
Ok(dynamic_column)
}
pub fn num_bytes(&self) -> usize {
self.file_slice.len()
}
pub fn column_type(&self) -> ColumnType {
self.column_type
}
}

View File

@@ -1,80 +1,86 @@
#![cfg_attr(all(feature = "unstable", test), feature(test))] mod column_type_header;
#[cfg(test)]
#[macro_use]
extern crate more_asserts;
#[cfg(all(test, feature = "unstable"))]
extern crate test;
use std::io;
mod column;
mod column_index;
mod column_values;
mod columnar;
mod dictionary; mod dictionary;
mod dynamic_column; mod reader;
pub(crate) mod utils; pub(crate) mod utils;
mod value; mod value;
mod writer;
pub use column::{BytesColumn, Column, StrColumn}; pub use column_type_header::Cardinality;
pub use column_values::ColumnValues; pub use reader::ColumnarReader;
pub use columnar::{
merge_columnar, ColumnType, ColumnarReader, ColumnarWriter, HasAssociatedColumnType,
MergeDocOrder,
};
pub use value::{NumericalType, NumericalValue}; pub use value::{NumericalType, NumericalValue};
pub use writer::ColumnarWriter;
pub use self::dynamic_column::{DynamicColumn, DynamicColumnHandle}; pub type DocId = u32;
pub type RowId = u32;
#[derive(Clone, Copy, PartialOrd, PartialEq, Default, Debug)]
pub struct DateTime {
pub timestamp_micros: i64,
}
#[derive(Copy, Clone, Debug)]
pub struct InvalidData;
impl From<InvalidData> for io::Error {
fn from(_: InvalidData) -> Self {
io::Error::new(io::ErrorKind::InvalidData, "Invalid data")
}
}
/// Enum describing the number of values that can exist per document
/// (or per row if you will).
///
/// The cardinality must fit on 2 bits.
#[derive(Clone, Copy, Hash, Default, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[repr(u8)]
pub enum Cardinality {
/// All documents contain exactly one value.
/// `Full` is the default for auto-detecting the Cardinality, since it is the most strict.
#[default]
Full = 0,
/// All documents contain at most one value.
Optional = 1,
/// All documents may contain any number of values.
Multivalued = 2,
}
impl Cardinality {
pub(crate) fn to_code(self) -> u8 {
self as u8
}
pub(crate) fn try_from_code(code: u8) -> Result<Cardinality, InvalidData> {
match code {
0 => Ok(Cardinality::Full),
1 => Ok(Cardinality::Optional),
2 => Ok(Cardinality::Multivalued),
_ => Err(InvalidData),
}
}
}
#[cfg(test)] #[cfg(test)]
mod tests; mod tests {
use std::ops::Range;
use common::file_slice::FileSlice;
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality};
use crate::reader::ColumnarReader;
use crate::value::NumericalValue;
use crate::{Cardinality, ColumnarWriter};
#[test]
fn test_dataframe_writer_bytes() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_str(1u32, "my_string", b"hello");
dataframe_writer.record_str(3u32, "my_string", b"helloeee");
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("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 = ColumnarReader::open(columnar_fileslice).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
columnar.read_columns("srical.value").unwrap();
assert_eq!(cols.len(), 1);
// Right now this 31 bytes are spent as follows
//
// - header 14 bytes
// - vals 8 //< due to padding? could have been 1byte?.
// - null footer 6 bytes
// - version footer 3 bytes // Should be file-wide
assert_eq!(cols[0].1, 0..32);
}
}

102
columnar/src/reader/mod.rs Normal file
View File

@@ -0,0 +1,102 @@
use std::ops::Range;
use std::{io, mem};
use common::file_slice::FileSlice;
use common::BinarySerializable;
use sstable::{Dictionary, RangeSSTable};
use crate::column_type_header::ColumnTypeAndCardinality;
fn io_invalid_data(msg: String) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, msg) // format!("Invalid key found.
// {key_bytes:?}")));
}
/// The ColumnarReader makes it possible to access a set of columns
/// associated to field names.
pub struct ColumnarReader {
column_dictionary: Dictionary<RangeSSTable>,
column_data: FileSlice,
}
impl ColumnarReader {
/// Opens a new Columnar file.
pub fn open<F>(file_slice: F) -> io::Result<ColumnarReader>
where FileSlice: From<F> {
Self::open_inner(file_slice.into())
}
fn open_inner(file_slice: FileSlice) -> io::Result<ColumnarReader> {
let (file_slice_without_sstable_len, sstable_len_bytes) =
file_slice.split_from_end(mem::size_of::<u64>());
let mut sstable_len_bytes = sstable_len_bytes.read_bytes()?;
let sstable_len = u64::deserialize(&mut sstable_len_bytes)?;
let (column_data, sstable) =
file_slice_without_sstable_len.split_from_end(sstable_len as usize);
let column_dictionary = Dictionary::open(sstable)?;
Ok(ColumnarReader {
column_dictionary,
column_data,
})
}
// 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(
&self,
field_name: &str,
) -> io::Result<Vec<(ColumnTypeAndCardinality, Range<u64>)>> {
let mut start_key = field_name.to_string();
start_key.push('\0');
let mut end_key = field_name.to_string();
end_key.push(1u8 as char);
let mut stream = self
.column_dictionary
.range()
.ge(start_key.as_bytes())
.lt(end_key.as_bytes())
.into_stream()?;
let mut results = Vec::new();
while stream.advance() {
let key_bytes: &[u8] = stream.key();
if !key_bytes.starts_with(start_key.as_bytes()) {
return Err(io_invalid_data(format!("Invalid key found. {key_bytes:?}")));
}
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();
results.push((column_type_and_cardinality, range));
}
Ok(results)
}
/// Return the number of columns in the columnar.
pub fn num_columns(&self) -> usize {
self.column_dictionary.num_terms()
}
}

View File

@@ -1,212 +0,0 @@
use std::net::Ipv6Addr;
use crate::column_values::MonotonicallyMappableToU128;
use crate::columnar::ColumnType;
use crate::dynamic_column::{DynamicColumn, DynamicColumnHandle};
use crate::value::NumericalValue;
use crate::{Cardinality, ColumnarReader, ColumnarWriter};
#[test]
fn test_dataframe_writer_str() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_str(1u32, "my_string", "hello");
dataframe_writer.record_str(3u32, "my_string", "helloeee");
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("my_string").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(cols[0].num_bytes(), 158);
}
#[test]
fn test_dataframe_writer_bytes() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_bytes(1u32, "my_string", b"hello");
dataframe_writer.record_bytes(3u32, "my_string", b"helloeee");
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("my_string").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(cols[0].num_bytes(), 158);
}
#[test]
fn test_dataframe_writer_bool() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_bool(1u32, "bool.value", false);
dataframe_writer.record_bool(3u32, "bool.value", true);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("bool.value").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(cols[0].num_bytes(), 22);
assert_eq!(cols[0].column_type(), ColumnType::Bool);
let dyn_bool_col = cols[0].open().unwrap();
let DynamicColumn::Bool(bool_col) = dyn_bool_col else { panic!(); };
let vals: Vec<Option<bool>> = (0..5).map(|row_id| bool_col.first(row_id)).collect();
assert_eq!(&vals, &[None, Some(false), None, Some(true), None,]);
}
#[test]
fn test_dataframe_writer_u64_multivalued() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_numerical(2u32, "divisor", 2u64);
dataframe_writer.record_numerical(3u32, "divisor", 3u64);
dataframe_writer.record_numerical(4u32, "divisor", 2u64);
dataframe_writer.record_numerical(5u32, "divisor", 5u64);
dataframe_writer.record_numerical(6u32, "divisor", 2u64);
dataframe_writer.record_numerical(6u32, "divisor", 3u64);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(7, &mut buffer).unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("divisor").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(cols[0].num_bytes(), 29);
let dyn_i64_col = cols[0].open().unwrap();
let DynamicColumn::I64(divisor_col) = dyn_i64_col else { panic!(); };
assert_eq!(
divisor_col.get_cardinality(),
crate::Cardinality::Multivalued
);
assert_eq!(divisor_col.num_rows(), 7);
}
#[test]
fn test_dataframe_writer_ip_addr() {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_ip_addr(1, "ip_addr", Ipv6Addr::from_u128(1001));
dataframe_writer.record_ip_addr(3, "ip_addr", Ipv6Addr::from_u128(1050));
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer.serialize(5, &mut buffer).unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("ip_addr").unwrap();
assert_eq!(cols.len(), 1);
assert_eq!(cols[0].num_bytes(), 42);
assert_eq!(cols[0].column_type(), ColumnType::IpAddr);
let dyn_bool_col = cols[0].open().unwrap();
let DynamicColumn::IpAddr(ip_col) = dyn_bool_col else { panic!(); };
let vals: Vec<Option<Ipv6Addr>> = (0..5).map(|row_id| ip_col.first(row_id)).collect();
assert_eq!(
&vals,
&[
None,
Some(Ipv6Addr::from_u128(1001)),
None,
Some(Ipv6Addr::from_u128(1050)),
None,
]
);
}
#[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(6, &mut buffer).unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("srical.value").unwrap();
assert_eq!(cols.len(), 1);
// Right now this 31 bytes are spent as follows
//
// - header 14 bytes
// - vals 8 //< due to padding? could have been 1byte?.
// - null footer 6 bytes
assert_eq!(cols[0].num_bytes(), 33);
let column = cols[0].open().unwrap();
let DynamicColumn::I64(column_i64) = column else { panic!(); };
assert_eq!(column_i64.idx.get_cardinality(), Cardinality::Optional);
assert_eq!(column_i64.first(0), None);
assert_eq!(column_i64.first(1), Some(12i64));
assert_eq!(column_i64.first(2), Some(13i64));
assert_eq!(column_i64.first(3), None);
assert_eq!(column_i64.first(4), Some(15i64));
assert_eq!(column_i64.first(5), None);
assert_eq!(column_i64.first(6), None); //< we can change the spec for that one.
}
#[test]
fn test_dictionary_encoded_str() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer.record_str(1, "my.column", "a");
columnar_writer.record_str(3, "my.column", "c");
columnar_writer.record_str(3, "my.column2", "different_column!");
columnar_writer.record_str(4, "my.column", "b");
columnar_writer.serialize(5, &mut buffer).unwrap();
let columnar_reader = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar_reader.num_columns(), 2);
let col_handles = columnar_reader.read_columns("my.column").unwrap();
assert_eq!(col_handles.len(), 1);
let DynamicColumn::Str(str_col) = col_handles[0].open().unwrap() else { panic!(); };
let index: Vec<Option<u64>> = (0..5).map(|row_id| str_col.ords().first(row_id)).collect();
assert_eq!(index, &[None, Some(0), None, Some(2), Some(1)]);
assert_eq!(str_col.num_rows(), 5);
let mut term_buffer = String::new();
let term_ords = str_col.ords();
assert_eq!(term_ords.first(0), None);
assert_eq!(term_ords.first(1), Some(0));
str_col.ord_to_str(0u64, &mut term_buffer).unwrap();
assert_eq!(term_buffer, "a");
assert_eq!(term_ords.first(2), None);
assert_eq!(term_ords.first(3), Some(2));
str_col.ord_to_str(2u64, &mut term_buffer).unwrap();
assert_eq!(term_buffer, "c");
assert_eq!(term_ords.first(4), Some(1));
str_col.ord_to_str(1u64, &mut term_buffer).unwrap();
assert_eq!(term_buffer, "b");
}
#[test]
fn test_dictionary_encoded_bytes() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer.record_bytes(1, "my.column", b"a");
columnar_writer.record_bytes(3, "my.column", b"c");
columnar_writer.record_bytes(3, "my.column2", b"different_column!");
columnar_writer.record_bytes(4, "my.column", b"b");
columnar_writer.serialize(5, &mut buffer).unwrap();
let columnar_reader = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar_reader.num_columns(), 2);
let col_handles = columnar_reader.read_columns("my.column").unwrap();
assert_eq!(col_handles.len(), 1);
let DynamicColumn::Bytes(bytes_col) = col_handles[0].open().unwrap() else { panic!(); };
let index: Vec<Option<u64>> = (0..5)
.map(|row_id| bytes_col.ords().first(row_id))
.collect();
assert_eq!(index, &[None, Some(0), None, Some(2), Some(1)]);
assert_eq!(bytes_col.num_rows(), 5);
let mut term_buffer = Vec::new();
let term_ords = bytes_col.ords();
assert_eq!(term_ords.first(0), None);
assert_eq!(term_ords.first(1), Some(0));
bytes_col
.dictionary
.ord_to_term(0u64, &mut term_buffer)
.unwrap();
assert_eq!(term_buffer, b"a");
assert_eq!(term_ords.first(2), None);
assert_eq!(term_ords.first(3), Some(2));
bytes_col
.dictionary
.ord_to_term(2u64, &mut term_buffer)
.unwrap();
assert_eq!(term_buffer, b"c");
assert_eq!(term_ords.first(4), Some(1));
bytes_col
.dictionary
.ord_to_term(1u64, &mut term_buffer)
.unwrap();
assert_eq!(term_buffer, b"b");
}

View File

@@ -1,22 +1,10 @@
use crate::InvalidData; #[derive(Copy, Clone, Debug, PartialEq)]
#[derive(Copy, Clone, PartialEq, Debug)]
pub enum NumericalValue { pub enum NumericalValue {
I64(i64), I64(i64),
U64(u64), U64(u64),
F64(f64), F64(f64),
} }
impl NumericalValue {
pub fn numerical_type(&self) -> NumericalType {
match self {
NumericalValue::I64(_) => NumericalType::I64,
NumericalValue::U64(_) => NumericalType::U64,
NumericalValue::F64(_) => NumericalType::F64,
}
}
}
impl From<u64> for NumericalValue { impl From<u64> for NumericalValue {
fn from(val: u64) -> NumericalValue { fn from(val: u64) -> NumericalValue {
NumericalValue::U64(val) NumericalValue::U64(val)
@@ -35,6 +23,18 @@ impl From<f64> for NumericalValue {
} }
} }
impl NumericalValue {
pub fn numerical_type(&self) -> NumericalType {
match self {
NumericalValue::F64(_) => NumericalType::F64,
NumericalValue::I64(_) => NumericalType::I64,
NumericalValue::U64(_) => NumericalType::U64,
}
}
}
impl Eq for NumericalValue {}
#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq)] #[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq)]
#[repr(u8)] #[repr(u8)]
pub enum NumericalType { pub enum NumericalType {
@@ -49,12 +49,12 @@ impl NumericalType {
self as u8 self as u8
} }
pub fn try_from_code(code: u8) -> Result<NumericalType, InvalidData> { pub fn try_from_code(code: u8) -> Option<NumericalType> {
match code { match code {
0 => Ok(NumericalType::I64), 0 => Some(NumericalType::I64),
1 => Ok(NumericalType::U64), 1 => Some(NumericalType::U64),
2 => Ok(NumericalType::F64), 2 => Some(NumericalType::F64),
_ => Err(InvalidData), _ => None,
} }
} }
} }
@@ -62,7 +62,6 @@ impl NumericalType {
/// We voluntarily avoid using `Into` here to keep this /// We voluntarily avoid using `Into` here to keep this
/// implementation quirk as private as possible. /// implementation quirk as private as possible.
/// ///
/// # Panics
/// This coercion trait actually panics if it is used /// This coercion trait actually panics if it is used
/// to convert a loose types to a stricter type. /// to convert a loose types to a stricter type.
/// ///
@@ -104,13 +103,6 @@ impl Coerce for f64 {
} }
} }
impl Coerce for crate::DateTime {
fn coerce(value: NumericalValue) -> Self {
let timestamp_micros = i64::coerce(value);
crate::DateTime { timestamp_micros }
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::NumericalType; use super::NumericalType;
@@ -119,7 +111,7 @@ mod tests {
fn test_numerical_type_code() { fn test_numerical_type_code() {
let mut num_numerical_type = 0; let mut num_numerical_type = 0;
for code in u8::MIN..=u8::MAX { for code in u8::MIN..=u8::MAX {
if let Ok(numerical_type) = NumericalType::try_from_code(code) { if let Some(numerical_type) = NumericalType::try_from_code(code) {
assert_eq!(numerical_type.to_code(), code); assert_eq!(numerical_type.to_code(), code);
num_numerical_type += 1; num_numerical_type += 1;
} }

View File

@@ -1,112 +1,84 @@
use std::net::Ipv6Addr;
use crate::dictionary::UnorderedId; use crate::dictionary::UnorderedId;
use crate::utils::{place_bits, pop_first_byte, select_bits}; use crate::utils::{place_bits, pop_first_byte, select_bits};
use crate::value::NumericalValue; use crate::value::NumericalValue;
use crate::{InvalidData, NumericalType, RowId}; use crate::{DocId, NumericalType};
/// 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 appends them in append-only buffer /// all mutations per column, and append them in append-only object.
/// in the stacker.
///
/// These ColumnOperation<T> are therefore serialize/deserialized
/// in memory.
/// ///
/// We represents all of these operations as `ColumnOperation`. /// We represents all of these operations as `ColumnOperation`.
#[derive(Eq, PartialEq, Debug, Clone, Copy)] #[derive(Eq, PartialEq, Debug, Clone, Copy)]
pub(super) enum ColumnOperation<T> { pub(crate) enum ColumnOperation<T> {
NewDoc(RowId), NewDoc(DocId),
Value(T), Value(T),
} }
#[derive(Copy, Clone, Eq, PartialEq, Debug)] #[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct ColumnOperationMetadata { struct ColumnOperationHeader {
op_type: ColumnOperationType, typ_code: u8,
len: u8, len: u8,
} }
impl ColumnOperationMetadata { impl ColumnOperationHeader {
fn to_code(self) -> u8 { fn to_code(self) -> u8 {
place_bits::<0, 6>(self.len) | place_bits::<6, 8>(self.op_type.to_code()) place_bits::<0, 4>(self.len) | place_bits::<4, 8>(self.typ_code)
} }
fn try_from_code(code: u8) -> Result<Self, InvalidData> { fn from_code(code: u8) -> Self {
let len = select_bits::<0, 6>(code); let len = select_bits::<0, 4>(code);
let typ_code = select_bits::<6, 8>(code); let typ_code = select_bits::<4, 8>(code);
let column_type = ColumnOperationType::try_from_code(typ_code)?; ColumnOperationHeader { typ_code, len }
Ok(ColumnOperationMetadata {
op_type: column_type,
len,
})
} }
} }
#[derive(Copy, Clone, Eq, PartialEq, Debug)] const NEW_DOC_CODE: u8 = 0u8;
#[repr(u8)] const NEW_VALUE_CODE: u8 = 1u8;
enum ColumnOperationType {
NewDoc = 0u8,
AddValue = 1u8,
}
impl ColumnOperationType {
pub fn to_code(self) -> u8 {
self as u8
}
pub fn try_from_code(code: u8) -> Result<Self, InvalidData> {
match code {
0 => Ok(Self::NewDoc),
1 => Ok(Self::AddValue),
_ => Err(InvalidData),
}
}
}
impl<V: SymbolValue> ColumnOperation<V> { impl<V: SymbolValue> ColumnOperation<V> {
pub(super) fn serialize(self) -> impl AsRef<[u8]> { pub fn serialize(self) -> impl AsRef<[u8]> {
let mut minibuf = MiniBuffer::default(); let mut minibuf = MiniBuffer::default();
let column_op_metadata = match self { let header = match self {
ColumnOperation::NewDoc(new_doc) => { ColumnOperation::NewDoc(new_doc) => {
let symbol_len = new_doc.serialize(&mut minibuf.bytes[1..]); let symbol_len = new_doc.serialize(&mut minibuf.bytes[1..]);
ColumnOperationMetadata { ColumnOperationHeader {
op_type: ColumnOperationType::NewDoc, typ_code: NEW_DOC_CODE,
len: symbol_len, len: symbol_len,
} }
} }
ColumnOperation::Value(val) => { ColumnOperation::Value(val) => {
let symbol_len = val.serialize(&mut minibuf.bytes[1..]); let symbol_len = val.serialize(&mut minibuf.bytes[1..]);
ColumnOperationMetadata { ColumnOperationHeader {
op_type: ColumnOperationType::AddValue, typ_code: NEW_VALUE_CODE,
len: symbol_len, len: symbol_len,
} }
} }
}; };
minibuf.bytes[0] = column_op_metadata.to_code(); minibuf.bytes[0] = header.to_code();
// +1 for the metadata minibuf.len = 1 + header.len;
minibuf.len = 1 + column_op_metadata.len;
minibuf minibuf
} }
/// Deserialize a colummn operation. /// Deserialize a colummn operation.
/// Returns None if the buffer is empty. /// Returns None if the buffer is empty.
/// ///
/// Panics if the payload is invalid: /// Panics if the payload is invalid.
/// this deserialize method is meant to target in memory. pub fn deserialize(bytes: &mut &[u8]) -> Option<Self> {
pub(super) fn deserialize(bytes: &mut &[u8]) -> Option<Self> { let header_byte = pop_first_byte(bytes)?;
let column_op_metadata_byte = pop_first_byte(bytes)?; let column_op_header = ColumnOperationHeader::from_code(header_byte);
let column_op_metadata = ColumnOperationMetadata::try_from_code(column_op_metadata_byte)
.expect("Invalid op metadata byte");
let symbol_bytes: &[u8]; let symbol_bytes: &[u8];
(symbol_bytes, *bytes) = bytes.split_at(column_op_metadata.len as usize); (symbol_bytes, *bytes) = bytes.split_at(column_op_header.len as usize);
match column_op_metadata.op_type { match column_op_header.typ_code {
ColumnOperationType::NewDoc => { NEW_DOC_CODE => {
let new_doc = u32::deserialize(symbol_bytes); let new_doc = u32::deserialize(symbol_bytes);
Some(ColumnOperation::NewDoc(new_doc)) Some(ColumnOperation::NewDoc(new_doc))
} }
ColumnOperationType::AddValue => { NEW_VALUE_CODE => {
let value = V::deserialize(symbol_bytes); let value = V::deserialize(symbol_bytes);
Some(ColumnOperation::Value(value)) Some(ColumnOperation::Value(value))
} }
_ => {
panic!("Unknown code {}", column_op_header.typ_code);
}
} }
} }
} }
@@ -117,25 +89,20 @@ impl<T> From<T> for ColumnOperation<T> {
} }
} }
// Serialization trait very local to the writer.
// As we write fast fields, we accumulate them in "in memory".
// In order to limit memory usage, and in order
// to benefit from the stacker, we do this by serialization our data
// as "Symbols".
#[allow(clippy::from_over_into)] #[allow(clippy::from_over_into)]
pub(super) trait SymbolValue: Clone + Copy { pub(crate) trait SymbolValue: Clone + Copy {
// Serializes the symbol into the given buffer.
// Returns the number of bytes written into the buffer.
/// # Panics
/// May not exceed 9bytes
fn serialize(self, buffer: &mut [u8]) -> u8; fn serialize(self, buffer: &mut [u8]) -> u8;
// Panics if invalid
// 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; fn deserialize(bytes: &[u8]) -> Self;
} }
impl SymbolValue for bool { impl SymbolValue for bool {
fn serialize(self, buffer: &mut [u8]) -> u8 { fn serialize(self, buffer: &mut [u8]) -> u8 {
buffer[0] = u8::from(self); buffer[0] = if self { 1u8 } else { 0u8 };
1u8 1u8
} }
@@ -144,21 +111,9 @@ impl SymbolValue for bool {
} }
} }
impl SymbolValue for Ipv6Addr {
fn serialize(self, buffer: &mut [u8]) -> u8 {
buffer[0..16].copy_from_slice(&self.octets());
16
}
fn deserialize(bytes: &[u8]) -> Self {
let octets: [u8; 16] = bytes[0..16].try_into().unwrap();
Ipv6Addr::from(octets)
}
}
#[derive(Default)] #[derive(Default)]
struct MiniBuffer { struct MiniBuffer {
pub bytes: [u8; 17], pub bytes: [u8; 10],
pub len: u8, pub len: u8,
} }
@@ -192,9 +147,6 @@ impl SymbolValue for NumericalValue {
} }
} }
/// F64: Serialize with a fixed size of 9 bytes
/// U64: Serialize without leading zeroes
/// I64: ZigZag encoded and serialize without leading zeroes
fn serialize(self, output: &mut [u8]) -> u8 { fn serialize(self, output: &mut [u8]) -> u8 {
match self { match self {
NumericalValue::F64(val) => { NumericalValue::F64(val) => {
@@ -291,14 +243,13 @@ mod tests {
} }
#[test] #[test]
fn test_column_op_metadata_byte_serialization() { fn test_header_byte_serialization() {
for len in 0..=15 { for len in 0..=15 {
for op_type in [ColumnOperationType::AddValue, ColumnOperationType::NewDoc] { for typ_code in 0..=15 {
let column_op_metadata = ColumnOperationMetadata { op_type, len }; let header = ColumnOperationHeader { typ_code, len };
let column_op_metadata_code = column_op_metadata.to_code(); let header_code = header.to_code();
let serdeser_metadata = let serdeser_header = ColumnOperationHeader::from_code(header_code);
ColumnOperationMetadata::try_from_code(column_op_metadata_code).unwrap(); assert_eq!(header, serdeser_header);
assert_eq!(column_op_metadata, serdeser_metadata);
} }
} }
} }

View File

@@ -2,25 +2,25 @@ use std::cmp::Ordering;
use stacker::{ExpUnrolledLinkedList, MemoryArena}; use stacker::{ExpUnrolledLinkedList, MemoryArena};
use crate::columnar::writer::column_operation::{ColumnOperation, SymbolValue};
use crate::dictionary::{DictionaryBuilder, UnorderedId}; use crate::dictionary::{DictionaryBuilder, UnorderedId};
use crate::{Cardinality, NumericalType, NumericalValue, RowId}; use crate::writer::column_operation::{ColumnOperation, SymbolValue};
use crate::{Cardinality, DocId, NumericalType, NumericalValue};
#[derive(Copy, Clone, Debug, Eq, PartialEq)] #[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[repr(u8)] #[repr(u8)]
enum DocumentStep { enum DocumentStep {
Same = 0, SameDoc = 0,
Next = 1, NextDoc = 1,
Skipped = 2, SkippedDoc = 2,
} }
#[inline(always)] #[inline(always)]
fn delta_with_last_doc(last_doc_opt: Option<u32>, doc: u32) -> DocumentStep { 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); let expected_next_doc = last_doc_opt.map(|last_doc| last_doc + 1).unwrap_or(0u32);
match doc.cmp(&expected_next_doc) { match doc.cmp(&expected_next_doc) {
Ordering::Less => DocumentStep::Same, Ordering::Less => DocumentStep::SameDoc,
Ordering::Equal => DocumentStep::Next, Ordering::Equal => DocumentStep::NextDoc,
Ordering::Greater => DocumentStep::Skipped, Ordering::Greater => DocumentStep::SkippedDoc,
} }
} }
@@ -38,7 +38,7 @@ pub struct ColumnWriter {
impl ColumnWriter { impl ColumnWriter {
/// Returns an iterator over the Symbol that have been recorded /// Returns an iterator over the Symbol that have been recorded
/// for the given column. /// for the given column.
pub(super) fn operation_iterator<'a, V: SymbolValue>( pub(crate) fn operation_iterator<'a, V: SymbolValue>(
&self, &self,
arena: &MemoryArena, arena: &MemoryArena,
buffer: &'a mut Vec<u8>, buffer: &'a mut Vec<u8>,
@@ -53,18 +53,18 @@ impl ColumnWriter {
/// ///
/// This function will also update the cardinality of the column /// This function will also update the cardinality of the column
/// if necessary. /// if necessary.
pub(super) fn record<S: SymbolValue>(&mut self, doc: RowId, value: S, arena: &mut MemoryArena) { pub(crate) fn record<S: SymbolValue>(&mut self, doc: DocId, value: S, arena: &mut MemoryArena) {
// Difference between `doc` and the last doc. // Difference between `doc` and the last doc.
match delta_with_last_doc(self.last_doc_opt, doc) { match delta_with_last_doc(self.last_doc_opt, doc) {
DocumentStep::Same => { DocumentStep::SameDoc => {
// This is the last encounterred document. // This is the last encounterred document.
self.cardinality = Cardinality::Multivalued; self.cardinality = Cardinality::Multivalued;
} }
DocumentStep::Next => { DocumentStep::NextDoc => {
self.last_doc_opt = Some(doc); self.last_doc_opt = Some(doc);
self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena); self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena);
} }
DocumentStep::Skipped => { DocumentStep::SkippedDoc => {
self.cardinality = self.cardinality.max(Cardinality::Optional); self.cardinality = self.cardinality.max(Cardinality::Optional);
self.last_doc_opt = Some(doc); self.last_doc_opt = Some(doc);
self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena); self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena);
@@ -77,10 +77,10 @@ impl ColumnWriter {
// The overall number of docs in the column is necessary to // 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 // deal with the case where the all docs contain 1 value, except some documents
// at the end of the column. // at the end of the column.
pub(crate) fn get_cardinality(&self, num_docs: RowId) -> Cardinality { pub fn get_cardinality(&self, num_docs: DocId) -> Cardinality {
match delta_with_last_doc(self.last_doc_opt, num_docs) { match delta_with_last_doc(self.last_doc_opt, num_docs) {
DocumentStep::Same | DocumentStep::Next => self.cardinality, DocumentStep::SameDoc | DocumentStep::NextDoc => self.cardinality,
DocumentStep::Skipped => self.cardinality.max(Cardinality::Optional), DocumentStep::SkippedDoc => self.cardinality.max(Cardinality::Optional),
} }
} }
@@ -102,29 +102,18 @@ pub(crate) struct NumericalColumnWriter {
column_writer: ColumnWriter, column_writer: ColumnWriter,
} }
impl NumericalColumnWriter {
pub fn force_numerical_type(&mut self, numerical_type: NumericalType) {
assert!(self
.compatible_numerical_types
.is_type_accepted(numerical_type));
self.compatible_numerical_types = CompatibleNumericalTypes::StaticType(numerical_type);
}
}
/// State used to store what types are still acceptable /// State used to store what types are still acceptable
/// after having seen a set of numerical values. /// after having seen a set of numerical values.
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
enum CompatibleNumericalTypes { pub(crate) struct CompatibleNumericalTypes {
Dynamic { all_values_within_i64_range: bool,
all_values_within_i64_range: bool, all_values_within_u64_range: bool,
all_values_within_u64_range: bool, // f64 is always acceptable.
},
StaticType(NumericalType),
} }
impl Default for CompatibleNumericalTypes { impl Default for CompatibleNumericalTypes {
fn default() -> CompatibleNumericalTypes { fn default() -> CompatibleNumericalTypes {
CompatibleNumericalTypes::Dynamic { CompatibleNumericalTypes {
all_values_within_i64_range: true, all_values_within_i64_range: true,
all_values_within_u64_range: true, all_values_within_u64_range: true,
} }
@@ -132,69 +121,43 @@ impl Default for CompatibleNumericalTypes {
} }
impl CompatibleNumericalTypes { impl CompatibleNumericalTypes {
fn is_type_accepted(&self, numerical_type: NumericalType) -> bool {
match self {
CompatibleNumericalTypes::Dynamic {
all_values_within_i64_range,
all_values_within_u64_range,
} => match numerical_type {
NumericalType::I64 => *all_values_within_i64_range,
NumericalType::U64 => *all_values_within_u64_range,
NumericalType::F64 => true,
},
CompatibleNumericalTypes::StaticType(static_numerical_type) => {
*static_numerical_type == numerical_type
}
}
}
fn accept_value(&mut self, numerical_value: NumericalValue) { fn accept_value(&mut self, numerical_value: NumericalValue) {
match self { match numerical_value {
CompatibleNumericalTypes::Dynamic { NumericalValue::I64(val_i64) => {
all_values_within_i64_range, let value_within_u64_range = val_i64 >= 0i64;
all_values_within_u64_range, self.all_values_within_u64_range &= value_within_u64_range;
} => match numerical_value { }
NumericalValue::I64(val_i64) => { NumericalValue::U64(val_u64) => {
let value_within_u64_range = val_i64 >= 0i64; let value_within_i64_range = val_u64 < i64::MAX as u64;
*all_values_within_u64_range &= value_within_u64_range; self.all_values_within_i64_range &= value_within_i64_range;
} }
NumericalValue::U64(val_u64) => { NumericalValue::F64(_) => {
let value_within_i64_range = val_u64 < i64::MAX as u64; self.all_values_within_i64_range = false;
*all_values_within_i64_range &= value_within_i64_range; self.all_values_within_u64_range = false;
}
NumericalValue::F64(_) => {
*all_values_within_i64_range = false;
*all_values_within_u64_range = false;
}
},
CompatibleNumericalTypes::StaticType(typ) => {
assert_eq!(numerical_value.numerical_type(), *typ);
} }
} }
} }
pub fn to_numerical_type(self) -> NumericalType { pub fn to_numerical_type(self) -> NumericalType {
for numerical_type in [NumericalType::I64, NumericalType::U64] { if self.all_values_within_i64_range {
if self.is_type_accepted(numerical_type) { NumericalType::I64
return numerical_type; } else if self.all_values_within_u64_range {
} NumericalType::U64
} else {
NumericalType::F64
} }
NumericalType::F64
} }
} }
impl NumericalColumnWriter { impl NumericalColumnWriter {
pub fn numerical_type(&self) -> NumericalType { pub fn column_type_and_cardinality(&self, num_docs: DocId) -> (NumericalType, Cardinality) {
self.compatible_numerical_types.to_numerical_type() let numerical_type = self.compatible_numerical_types.to_numerical_type();
let cardinality = self.column_writer.get_cardinality(num_docs);
(numerical_type, cardinality)
} }
pub fn cardinality(&self, num_docs: RowId) -> Cardinality {
self.column_writer.get_cardinality(num_docs)
}
pub fn record_numerical_value( pub fn record_numerical_value(
&mut self, &mut self,
doc: RowId, doc: DocId,
value: NumericalValue, value: NumericalValue,
arena: &mut MemoryArena, arena: &mut MemoryArena,
) { ) {
@@ -202,7 +165,7 @@ impl NumericalColumnWriter {
self.column_writer.record(doc, value, arena); self.column_writer.record(doc, value, arena);
} }
pub(super) fn operation_iterator<'a>( pub fn operation_iterator<'a>(
self, self,
arena: &MemoryArena, arena: &MemoryArena,
buffer: &'a mut Vec<u8>, buffer: &'a mut Vec<u8>,
@@ -211,15 +174,15 @@ impl NumericalColumnWriter {
} }
} }
#[derive(Copy, Clone)] #[derive(Copy, Clone, Default)]
pub(crate) struct StrOrBytesColumnWriter { pub struct StrColumnWriter {
pub(crate) dictionary_id: u32, pub(crate) dictionary_id: u32,
pub(crate) column_writer: ColumnWriter, pub(crate) column_writer: ColumnWriter,
} }
impl StrOrBytesColumnWriter { impl StrColumnWriter {
pub(crate) fn with_dictionary_id(dictionary_id: u32) -> StrOrBytesColumnWriter { pub fn with_dictionary_id(dictionary_id: u32) -> StrColumnWriter {
StrOrBytesColumnWriter { StrColumnWriter {
dictionary_id, dictionary_id,
column_writer: Default::default(), column_writer: Default::default(),
} }
@@ -227,7 +190,7 @@ impl StrOrBytesColumnWriter {
pub(crate) fn record_bytes( pub(crate) fn record_bytes(
&mut self, &mut self,
doc: RowId, doc: DocId,
bytes: &[u8], bytes: &[u8],
dictionaries: &mut [DictionaryBuilder], dictionaries: &mut [DictionaryBuilder],
arena: &mut MemoryArena, arena: &mut MemoryArena,
@@ -236,7 +199,7 @@ impl StrOrBytesColumnWriter {
self.column_writer.record(doc, unordered_id, arena); self.column_writer.record(doc, unordered_id, arena);
} }
pub(super) fn operation_iterator<'a>( pub(crate) fn operation_iterator<'a>(
&self, &self,
arena: &MemoryArena, arena: &MemoryArena,
byte_buffer: &'a mut Vec<u8>, byte_buffer: &'a mut Vec<u8>,
@@ -251,14 +214,20 @@ mod tests {
#[test] #[test]
fn test_delta_with_last_doc() { fn test_delta_with_last_doc() {
assert_eq!(delta_with_last_doc(None, 0u32), DocumentStep::Next); assert_eq!(delta_with_last_doc(None, 0u32), DocumentStep::NextDoc);
assert_eq!(delta_with_last_doc(None, 1u32), DocumentStep::Skipped); assert_eq!(delta_with_last_doc(None, 1u32), DocumentStep::SkippedDoc);
assert_eq!(delta_with_last_doc(None, 2u32), DocumentStep::Skipped); assert_eq!(delta_with_last_doc(None, 2u32), DocumentStep::SkippedDoc);
assert_eq!(delta_with_last_doc(Some(0u32), 0u32), DocumentStep::Same); assert_eq!(delta_with_last_doc(Some(0u32), 0u32), DocumentStep::SameDoc);
assert_eq!(delta_with_last_doc(Some(1u32), 1u32), DocumentStep::Same); assert_eq!(delta_with_last_doc(Some(1u32), 1u32), DocumentStep::SameDoc);
assert_eq!(delta_with_last_doc(Some(1u32), 2u32), DocumentStep::Next); assert_eq!(delta_with_last_doc(Some(1u32), 2u32), DocumentStep::NextDoc);
assert_eq!(delta_with_last_doc(Some(1u32), 3u32), DocumentStep::Skipped); assert_eq!(
assert_eq!(delta_with_last_doc(Some(1u32), 4u32), DocumentStep::Skipped); delta_with_last_doc(Some(1u32), 3u32),
DocumentStep::SkippedDoc
);
assert_eq!(
delta_with_last_doc(Some(1u32), 4u32),
DocumentStep::SkippedDoc
);
} }
#[track_caller] #[track_caller]
@@ -298,27 +267,4 @@ mod tests {
test_column_writer_coercion_aux(&[1i64.into(), 1u64.into()], NumericalType::I64); test_column_writer_coercion_aux(&[1i64.into(), 1u64.into()], NumericalType::I64);
test_column_writer_coercion_aux(&[u64::MAX.into(), (-1i64).into()], NumericalType::F64); test_column_writer_coercion_aux(&[u64::MAX.into(), (-1i64).into()], NumericalType::F64);
} }
#[test]
#[should_panic]
fn test_compatible_numerical_types_static_incompatible_type() {
let mut compatible_numerical_types =
CompatibleNumericalTypes::StaticType(NumericalType::U64);
compatible_numerical_types.accept_value(NumericalValue::I64(1i64));
}
#[test]
fn test_compatible_numerical_types_static_different_type_forbidden() {
let mut compatible_numerical_types =
CompatibleNumericalTypes::StaticType(NumericalType::U64);
compatible_numerical_types.accept_value(NumericalValue::U64(u64::MAX));
}
#[test]
fn test_compatible_numerical_types_static() {
for typ in [NumericalType::I64, NumericalType::I64, NumericalType::F64] {
let compatible_numerical_types = CompatibleNumericalTypes::StaticType(typ);
assert_eq!(compatible_numerical_types.to_numerical_type(), typ);
}
}
} }

526
columnar/src/writer/mod.rs Normal file
View File

@@ -0,0 +1,526 @@
mod column_operation;
mod column_writers;
mod serializer;
mod value_index;
use std::io::{self, Write};
use column_operation::ColumnOperation;
use fastfield_codecs::serialize::ValueIndexInfo;
use fastfield_codecs::{Column, MonotonicallyMappableToU64, VecColumn};
use serializer::ColumnarSerializer;
use stacker::{Addr, ArenaHashMap, MemoryArena};
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality, GeneralType};
use crate::dictionary::{DictionaryBuilder, IdMapping, UnorderedId};
use crate::value::{Coerce, NumericalType, NumericalValue};
use crate::writer::column_writers::{ColumnWriter, NumericalColumnWriter, StrColumnWriter};
use crate::writer::value_index::{IndexBuilder, SpareIndexBuilders};
use crate::{Cardinality, DocId};
/// Threshold above which a column data will be compressed
/// using ZSTD.
const COLUMN_COMPRESSION_THRESHOLD: usize = 100_000;
/// This is a set of buffers that are only here
/// to limit the amount of allocation.
#[derive(Default)]
struct SpareBuffers {
value_index_builders: SpareIndexBuilders,
i64_values: Vec<i64>,
u64_values: Vec<u64>,
f64_values: Vec<f64>,
bool_values: Vec<bool>,
column_buffer: Vec<u8>,
}
pub struct ColumnarWriter {
numerical_field_hash_map: ArenaHashMap,
bool_field_hash_map: ArenaHashMap,
bytes_field_hash_map: ArenaHashMap,
arena: MemoryArena,
// Dictionaries used to store dictionary-encoded values.
dictionaries: Vec<DictionaryBuilder>,
buffers: SpareBuffers,
}
impl Default for ColumnarWriter {
fn default() -> Self {
ColumnarWriter {
numerical_field_hash_map: ArenaHashMap::new(10_000),
bool_field_hash_map: ArenaHashMap::new(10_000),
bytes_field_hash_map: ArenaHashMap::new(10_000),
dictionaries: Vec::new(),
arena: MemoryArena::default(),
buffers: SpareBuffers::default(),
}
}
}
impl ColumnarWriter {
pub fn record_numerical(
&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);
hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<NumericalColumnWriter>| {
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
column.record_numerical_value(doc, numerical_value, arena);
column
},
);
}
pub fn record_bool(&mut self, doc: DocId, column_name: &str, val: bool) {
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) = (
&mut self.bytes_field_hash_map,
&mut self.arena,
&mut self.dictionaries,
);
hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<StrColumnWriter>| {
let mut column: StrColumnWriter = column_opt.unwrap_or_else(|| {
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrColumnWriter::with_dictionary_id(dictionary_id)
});
column.record_bytes(doc, value, dictionaries, arena);
column
},
);
}
pub fn serialize(&mut self, num_docs: DocId, wrt: &mut dyn io::Write) -> io::Result<()> {
let mut serializer = ColumnarSerializer::new(wrt);
let mut field_columns: Vec<(&[u8], GeneralType, Addr)> = self
.numerical_field_hash_map
.iter()
.map(|(term, addr, _)| (term, GeneralType::Numerical, addr))
.collect();
field_columns.extend(
self.bytes_field_hash_map
.iter()
.map(|(term, addr, _)| (term, GeneralType::Str, addr)),
);
field_columns.extend(
self.bool_field_hash_map
.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 mut symbol_byte_buffer: Vec<u8> = Vec::new();
for (column_name, bytes_or_numerical, addr) in field_columns {
match bytes_or_numerical {
GeneralType::Bool => {
let column_writer: ColumnWriter = self.bool_field_hash_map.read(addr);
let cardinality = column_writer.get_cardinality(num_docs);
let column_type_and_cardinality = ColumnTypeAndCardinality {
cardinality,
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 =
&dictionaries[str_column_writer.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(
cardinality,
num_docs,
dictionary_builder,
str_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
column_serializer,
)?;
}
GeneralType::Numerical => {
let numerical_column_writer: NumericalColumnWriter =
self.numerical_field_hash_map.read(addr);
let (numerical_type, cardinality) =
numerical_column_writer.column_type_and_cardinality(num_docs);
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(
cardinality,
num_docs,
numerical_type,
numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers,
column_serializer,
)?;
}
};
}
serializer.finalize()?;
Ok(())
}
}
fn compress_and_write_column<W: io::Write>(column_bytes: &[u8], wrt: &mut W) -> io::Result<()> {
if column_bytes.len() >= COLUMN_COMPRESSION_THRESHOLD {
wrt.write_all(&[1])?;
let mut encoder = zstd::Encoder::new(wrt, 3)?;
encoder.write_all(column_bytes)?;
encoder.finish()?;
} else {
wrt.write_all(&[0])?;
wrt.write_all(column_bytes)?;
}
Ok(())
}
fn serialize_bytes_column<W: io::Write>(
cardinality: Cardinality,
num_docs: DocId,
dictionary_builder: &DictionaryBuilder,
operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>,
buffers: &mut SpareBuffers,
mut wrt: W,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
u64_values,
column_buffer,
..
} = buffers;
column_buffer.clear();
let id_mapping: IdMapping = dictionary_builder.serialize(column_buffer)?;
let dictionary_num_bytes: u32 = column_buffer.len() as u32;
let operation_iterator = operation_it.map(|symbol: ColumnOperation<UnorderedId>| {
// We map unordered ids to ordered ids.
match symbol {
ColumnOperation::Value(unordered_id) => {
let ordered_id = id_mapping.to_ord(unordered_id);
ColumnOperation::Value(ordered_id.0 as u64)
}
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
}
});
serialize_column(
operation_iterator,
cardinality,
num_docs,
value_index_builders,
u64_values,
column_buffer,
)?;
column_buffer.write_all(&dictionary_num_bytes.to_le_bytes()[..])?;
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(())
}
fn serialize_numerical_column<W: io::Write>(
cardinality: Cardinality,
num_docs: DocId,
numerical_type: NumericalType,
op_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
buffers: &mut SpareBuffers,
mut wrt: W,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
u64_values,
i64_values,
f64_values,
column_buffer,
..
} = buffers;
column_buffer.clear();
match numerical_type {
NumericalType::I64 => {
serialize_column(
coerce_numerical_symbol::<i64>(op_iterator),
cardinality,
num_docs,
value_index_builders,
i64_values,
column_buffer,
)?;
}
NumericalType::U64 => {
serialize_column(
coerce_numerical_symbol::<u64>(op_iterator),
cardinality,
num_docs,
value_index_builders,
u64_values,
column_buffer,
)?;
}
NumericalType::F64 => {
serialize_column(
coerce_numerical_symbol::<f64>(op_iterator),
cardinality,
num_docs,
value_index_builders,
f64_values,
column_buffer,
)?;
}
};
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(())
}
fn serialize_column<
T: Copy + Default + std::fmt::Debug + Send + Sync + MonotonicallyMappableToU64 + PartialOrd,
>(
op_iterator: impl Iterator<Item = ColumnOperation<T>>,
cardinality: Cardinality,
num_docs: DocId,
value_index_builders: &mut SpareIndexBuilders,
values: &mut Vec<T>,
wrt: &mut Vec<u8>,
) -> io::Result<()>
where
for<'a> VecColumn<'a, T>: Column<T>,
{
values.clear();
match cardinality {
Cardinality::Required => {
consume_operation_iterator(
op_iterator,
value_index_builders.borrow_required_index_builder(),
values,
);
fastfield_codecs::serialize(
VecColumn::from(&values[..]),
wrt,
&fastfield_codecs::ALL_CODEC_TYPES[..],
)?;
}
Cardinality::Optional => {
let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_docs);
fastfield_codecs::serialize::serialize_new(
ValueIndexInfo::SingleValue(Box::new(optional_index)),
VecColumn::from(&values[..]),
wrt,
&fastfield_codecs::ALL_CODEC_TYPES[..],
)?;
}
Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_docs);
fastfield_codecs::serialize::serialize_new(
ValueIndexInfo::MultiValue(Box::new(multivalued_index)),
VecColumn::from(&values[..]),
wrt,
&fastfield_codecs::ALL_CODEC_TYPES[..],
)?;
}
}
Ok(())
}
fn coerce_numerical_symbol<T>(
operation_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
) -> impl Iterator<Item = ColumnOperation<T>>
where T: Coerce {
operation_iterator.map(|symbol| match symbol {
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
ColumnOperation::Value(numerical_value) => {
ColumnOperation::Value(Coerce::coerce(numerical_value))
}
})
}
fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>(
operation_iterator: impl Iterator<Item = ColumnOperation<T>>,
index_builder: &mut TIndexBuilder,
values: &mut Vec<T>,
) {
for symbol in operation_iterator {
match symbol {
ColumnOperation::NewDoc(doc) => {
index_builder.record_doc(doc);
}
ColumnOperation::Value(value) => {
index_builder.record_value();
values.push(value);
}
}
}
}
#[cfg(test)]
mod tests {
use column_operation::ColumnOperation;
use stacker::MemoryArena;
use super::*;
use crate::value::NumericalValue;
use crate::Cardinality;
#[test]
fn test_column_writer_required_simple() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(14i64), &mut arena);
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Required);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 6);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(14i64))
));
assert!(matches!(symbols[2], ColumnOperation::NewDoc(1u32)));
assert!(matches!(
symbols[3],
ColumnOperation::Value(NumericalValue::I64(15i64))
));
assert!(matches!(symbols[4], ColumnOperation::NewDoc(2u32)));
assert!(matches!(
symbols[5],
ColumnOperation::Value(NumericalValue::I64(-16i64))
));
}
#[test]
fn test_column_writer_optional_cardinality_missing_first() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Optional);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 4);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(1u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(15i64))
));
assert!(matches!(symbols[2], ColumnOperation::NewDoc(2u32)));
assert!(matches!(
symbols[3],
ColumnOperation::Value(NumericalValue::I64(-16i64))
));
}
#[test]
fn test_column_writer_optional_cardinality_missing_last() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(15i64), &mut arena);
assert_eq!(column_writer.get_cardinality(2), Cardinality::Optional);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 2);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(15i64))
));
}
#[test]
fn test_column_writer_multivalued() {
let mut arena = MemoryArena::default();
let mut column_writer = super::ColumnWriter::default();
column_writer.record(0u32, NumericalValue::from(16i64), &mut arena);
column_writer.record(0u32, NumericalValue::from(17i64), &mut arena);
assert_eq!(column_writer.get_cardinality(1), Cardinality::Multivalued);
let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer)
.collect();
assert_eq!(symbols.len(), 3);
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
assert!(matches!(
symbols[1],
ColumnOperation::Value(NumericalValue::I64(16i64))
));
assert!(matches!(
symbols[2],
ColumnOperation::Value(NumericalValue::I64(17i64))
));
}
}

View File

@@ -5,7 +5,7 @@ use common::CountingWriter;
use sstable::value::RangeValueWriter; use sstable::value::RangeValueWriter;
use sstable::RangeSSTable; use sstable::RangeSSTable;
use crate::columnar::ColumnType; use crate::column_type_header::ColumnTypeAndCardinality;
pub struct ColumnarSerializer<W: io::Write> { pub struct ColumnarSerializer<W: io::Write> {
wrt: CountingWriter<W>, wrt: CountingWriter<W>,
@@ -15,11 +15,15 @@ pub struct ColumnarSerializer<W: io::Write> {
/// Returns a key consisting of the concatenation of the key and the column_type_and_cardinality /// Returns a key consisting of the concatenation of the key and the column_type_and_cardinality
/// code. /// code.
fn prepare_key(key: &[u8], column_type: ColumnType, buffer: &mut Vec<u8>) { fn prepare_key<'a>(
key: &[u8],
column_type_cardinality: ColumnTypeAndCardinality,
buffer: &'a mut Vec<u8>,
) {
buffer.clear(); buffer.clear();
buffer.extend_from_slice(key); buffer.extend_from_slice(key);
buffer.push(0u8); buffer.push(0u8);
buffer.push(column_type.to_code()); buffer.push(column_type_cardinality.to_code());
} }
impl<W: io::Write> ColumnarSerializer<W> { impl<W: io::Write> ColumnarSerializer<W> {
@@ -36,10 +40,14 @@ impl<W: io::Write> ColumnarSerializer<W> {
pub fn serialize_column<'a>( pub fn serialize_column<'a>(
&'a mut self, &'a mut self,
column_name: &[u8], column_name: &[u8],
column_type: ColumnType, column_type_cardinality: ColumnTypeAndCardinality,
) -> impl io::Write + 'a { ) -> impl io::Write + 'a {
let start_offset = self.wrt.written_bytes(); let start_offset = self.wrt.written_bytes();
prepare_key(column_name, column_type, &mut self.prepare_key_buffer); prepare_key(
column_name,
column_type_cardinality,
&mut self.prepare_key_buffer,
);
ColumnSerializer { ColumnSerializer {
columnar_serializer: self, columnar_serializer: self,
start_offset, start_offset,
@@ -51,9 +59,6 @@ impl<W: io::Write> ColumnarSerializer<W> {
let sstable_num_bytes: u64 = sstable_bytes.len() as u64; let sstable_num_bytes: u64 = sstable_bytes.len() as u64;
self.wrt.write_all(&sstable_bytes)?; self.wrt.write_all(&sstable_bytes)?;
self.wrt.write_all(&sstable_num_bytes.to_le_bytes()[..])?; self.wrt.write_all(&sstable_num_bytes.to_le_bytes()[..])?;
self.wrt
.write_all(&super::super::format_version::footer())?;
self.wrt.flush()?;
Ok(()) Ok(())
} }
} }
@@ -92,15 +97,20 @@ impl<'a, W: io::Write> io::Write for ColumnSerializer<'a, W> {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::columnar::column_type::ColumnType; use crate::column_type_header::ColumnType;
use crate::Cardinality;
#[test] #[test]
fn test_prepare_key_bytes() { fn test_prepare_key_bytes() {
let mut buffer: Vec<u8> = b"somegarbage".to_vec(); let mut buffer: Vec<u8> = b"somegarbage".to_vec();
prepare_key(b"root\0child", ColumnType::Str, &mut buffer); 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.len(), 12);
assert_eq!(&buffer[..10], b"root\0child"); assert_eq!(&buffer[..10], b"root\0child");
assert_eq!(buffer[10], 0u8); assert_eq!(buffer[10], 0u8);
assert_eq!(buffer[11], ColumnType::Str.to_code()); assert_eq!(buffer[11], column_type_and_cardinality.to_code());
} }
} }

View File

@@ -1,61 +1,64 @@
use crate::column_index::SerializableOptionalIndex; use fastfield_codecs::serialize::{MultiValueIndexInfo, SingleValueIndexInfo};
use crate::column_values::{ColumnValues, VecColumn};
use crate::RowId; use crate::DocId;
/// The `IndexBuilder` interprets a sequence of /// The `IndexBuilder` interprets a sequence of
/// calls of the form: /// calls of the form:
/// (record_doc,record_value+)* /// (record_doc,record_value+)*
/// and can then serialize the results into an index to associate docids with their value[s]. /// and can then serialize the results into an index.
/// ///
/// It has different implementation depending on whether the /// It has different implementation depending on whether the
/// cardinality is required, optional, or multivalued. /// cardinality is required, optional, or multivalued.
pub(crate) trait IndexBuilder { pub(crate) trait IndexBuilder {
fn record_row(&mut self, doc: RowId); fn record_doc(&mut self, doc: DocId);
#[inline] #[inline]
fn record_value(&mut self) {} fn record_value(&mut self) {}
} }
/// The FullIndexBuilder does nothing. /// The RequiredIndexBuilder does nothing.
#[derive(Default)] #[derive(Default)]
pub struct FullIndexBuilder; pub struct RequiredIndexBuilder;
impl IndexBuilder for FullIndexBuilder { impl IndexBuilder for RequiredIndexBuilder {
#[inline(always)] #[inline(always)]
fn record_row(&mut self, _doc: RowId) {} fn record_doc(&mut self, _doc: DocId) {}
} }
#[derive(Default)] #[derive(Default)]
pub struct OptionalIndexBuilder { pub struct OptionalIndexBuilder {
docs: Vec<RowId>, docs: Vec<DocId>,
} }
struct SingleValueArrayIndex<'a> { struct SingleValueArrayIndex<'a> {
// RowIds with a value, in a strictly increasing order docs: &'a [DocId],
row_ids: &'a [RowId], num_docs: DocId,
num_rows: RowId,
} }
impl<'a> SerializableOptionalIndex<'a> for SingleValueArrayIndex<'a> { impl<'a> SingleValueIndexInfo for SingleValueArrayIndex<'a> {
fn num_rows(&self) -> RowId { fn num_vals(&self) -> u32 {
self.num_rows self.num_docs as u32
} }
fn non_null_rows(&self) -> Box<dyn Iterator<Item = RowId> + 'a> { fn num_non_nulls(&self) -> u32 {
Box::new(self.row_ids.iter().copied()) self.docs.len() as u32
}
fn iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
Box::new(self.docs.iter().copied())
} }
} }
impl OptionalIndexBuilder { impl OptionalIndexBuilder {
pub fn finish<'a>(&'a mut self, num_rows: RowId) -> impl SerializableOptionalIndex + 'a { pub fn finish(&mut self, num_docs: DocId) -> impl SingleValueIndexInfo + '_ {
debug_assert!(self debug_assert!(self
.docs .docs
.last() .last()
.copied() .copied()
.map(|last_doc| last_doc < num_rows) .map(|last_doc| last_doc < num_docs)
.unwrap_or(true)); .unwrap_or(true));
SingleValueArrayIndex { SingleValueArrayIndex {
row_ids: &self.docs[..], docs: &self.docs[..],
num_rows, num_docs,
} }
} }
@@ -66,7 +69,7 @@ impl OptionalIndexBuilder {
impl IndexBuilder for OptionalIndexBuilder { impl IndexBuilder for OptionalIndexBuilder {
#[inline(always)] #[inline(always)]
fn record_row(&mut self, doc: RowId) { fn record_doc(&mut self, doc: DocId) {
debug_assert!(self debug_assert!(self
.docs .docs
.last() .last()
@@ -79,32 +82,52 @@ impl IndexBuilder for OptionalIndexBuilder {
#[derive(Default)] #[derive(Default)]
pub struct MultivaluedIndexBuilder { pub struct MultivaluedIndexBuilder {
start_offsets: Vec<RowId>, // TODO should we switch to `start_offset`?
end_values: Vec<DocId>,
total_num_vals_seen: u32, total_num_vals_seen: u32,
} }
pub struct MultivaluedValueArrayIndex<'a> {
end_offsets: &'a [DocId],
}
impl<'a> MultiValueIndexInfo for MultivaluedValueArrayIndex<'a> {
fn num_docs(&self) -> u32 {
self.end_offsets.len() as u32
}
fn num_vals(&self) -> u32 {
self.end_offsets.last().copied().unwrap_or(0u32)
}
fn iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
if self.end_offsets.is_empty() {
return Box::new(std::iter::empty());
}
let n = self.end_offsets.len();
Box::new(std::iter::once(0u32).chain(self.end_offsets[..n - 1].iter().copied()))
}
}
impl MultivaluedIndexBuilder { impl MultivaluedIndexBuilder {
pub fn finish(&mut self, num_docs: RowId) -> impl ColumnValues<u32> + '_ { pub fn finish(&mut self, num_docs: DocId) -> impl MultiValueIndexInfo + '_ {
self.start_offsets self.end_values
.resize(num_docs as usize + 1, self.total_num_vals_seen); .resize(num_docs as usize, self.total_num_vals_seen);
VecColumn { MultivaluedValueArrayIndex {
values: &&self.start_offsets[..], end_offsets: &self.end_values[..],
min_value: 0,
max_value: self.start_offsets.last().copied().unwrap_or(0),
} }
} }
fn reset(&mut self) { fn reset(&mut self) {
self.start_offsets.clear(); self.end_values.clear();
self.start_offsets.push(0u32);
self.total_num_vals_seen = 0; self.total_num_vals_seen = 0;
} }
} }
impl IndexBuilder for MultivaluedIndexBuilder { impl IndexBuilder for MultivaluedIndexBuilder {
fn record_row(&mut self, row_id: RowId) { fn record_doc(&mut self, doc: DocId) {
self.start_offsets self.end_values
.resize(row_id as usize + 1, self.total_num_vals_seen); .resize(doc as usize, self.total_num_vals_seen);
} }
fn record_value(&mut self) { fn record_value(&mut self) {
@@ -115,14 +138,14 @@ impl IndexBuilder for MultivaluedIndexBuilder {
/// The `SpareIndexBuilders` is there to avoid allocating a /// The `SpareIndexBuilders` is there to avoid allocating a
/// new index builder for every single column. /// new index builder for every single column.
#[derive(Default)] #[derive(Default)]
pub struct PreallocatedIndexBuilders { pub struct SpareIndexBuilders {
required_index_builder: FullIndexBuilder, required_index_builder: RequiredIndexBuilder,
optional_index_builder: OptionalIndexBuilder, optional_index_builder: OptionalIndexBuilder,
multivalued_index_builder: MultivaluedIndexBuilder, multivalued_index_builder: MultivaluedIndexBuilder,
} }
impl PreallocatedIndexBuilders { impl SpareIndexBuilders {
pub fn borrow_required_index_builder(&mut self) -> &mut FullIndexBuilder { pub fn borrow_required_index_builder(&mut self) -> &mut RequiredIndexBuilder {
&mut self.required_index_builder &mut self.required_index_builder
} }
@@ -144,22 +167,22 @@ mod tests {
#[test] #[test]
fn test_optional_value_index_builder() { fn test_optional_value_index_builder() {
let mut opt_value_index_builder = OptionalIndexBuilder::default(); let mut opt_value_index_builder = OptionalIndexBuilder::default();
opt_value_index_builder.record_row(0u32); opt_value_index_builder.record_doc(0u32);
opt_value_index_builder.record_value(); opt_value_index_builder.record_value();
assert_eq!( assert_eq!(
&opt_value_index_builder &opt_value_index_builder
.finish(1u32) .finish(1u32)
.non_null_rows() .iter()
.collect::<Vec<u32>>(), .collect::<Vec<u32>>(),
&[0] &[0]
); );
opt_value_index_builder.reset(); opt_value_index_builder.reset();
opt_value_index_builder.record_row(1u32); opt_value_index_builder.record_doc(1u32);
opt_value_index_builder.record_value(); opt_value_index_builder.record_value();
assert_eq!( assert_eq!(
&opt_value_index_builder &opt_value_index_builder
.finish(2u32) .finish(2u32)
.non_null_rows() .iter()
.collect::<Vec<u32>>(), .collect::<Vec<u32>>(),
&[1] &[1]
); );
@@ -168,20 +191,20 @@ mod tests {
#[test] #[test]
fn test_multivalued_value_index_builder() { fn test_multivalued_value_index_builder() {
let mut multivalued_value_index_builder = MultivaluedIndexBuilder::default(); let mut multivalued_value_index_builder = MultivaluedIndexBuilder::default();
multivalued_value_index_builder.record_row(1u32); multivalued_value_index_builder.record_doc(1u32);
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
multivalued_value_index_builder.record_row(2u32); multivalued_value_index_builder.record_doc(2u32);
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
assert_eq!( assert_eq!(
multivalued_value_index_builder multivalued_value_index_builder
.finish(4u32) .finish(4u32)
.iter() .iter()
.collect::<Vec<u32>>(), .collect::<Vec<u32>>(),
vec![0, 0, 2, 3, 3] vec![0, 0, 2, 3]
); );
multivalued_value_index_builder.reset(); multivalued_value_index_builder.reset();
multivalued_value_index_builder.record_row(2u32); multivalued_value_index_builder.record_doc(2u32);
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
assert_eq!( assert_eq!(
@@ -189,7 +212,7 @@ mod tests {
.finish(4u32) .finish(4u32)
.iter() .iter()
.collect::<Vec<u32>>(), .collect::<Vec<u32>>(),
vec![0, 0, 0, 2, 2] vec![0, 0, 0, 2]
); );
} }
} }

View File

@@ -1,166 +0,0 @@
use std::cell::RefCell;
use std::iter::Peekable;
use std::rc::Rc;
pub trait GroupByIteratorExtended: Iterator {
/// Return an `Iterator` that groups iterator elements. Consecutive elements that map to the
/// same key are assigned to the same group.
///
/// The returned Iterator item is `(K, impl Iterator)`, where Iterator are the items of the
/// group.
///
/// ```
/// use tantivy_common::GroupByIteratorExtended;
///
/// // group data into blocks of larger than zero or not.
/// let data: Vec<i32> = vec![1, 3, -2, -2, 1, 0, 1, 2];
/// // groups: |---->|------>|--------->|
///
/// let mut data_grouped = Vec::new();
/// // Note: group is an iterator
/// for (key, group) in data.into_iter().group_by(|val| *val >= 0) {
/// data_grouped.push((key, group.collect()));
/// }
/// assert_eq!(data_grouped, vec![(true, vec![1, 3]), (false, vec![-2, -2]), (true, vec![1, 0, 1, 2])]);
/// ```
fn group_by<K, F>(self, key: F) -> GroupByIterator<Self, F, K>
where
Self: Sized,
F: FnMut(&Self::Item) -> K,
K: PartialEq + Copy,
Self::Item: Copy,
{
GroupByIterator::new(self, key)
}
}
impl<I: Iterator> GroupByIteratorExtended for I {}
pub struct GroupByIterator<I, F, K: Copy>
where
I: Iterator,
F: FnMut(&I::Item) -> K,
{
// I really would like to avoid the Rc<RefCell>, but the Iterator is shared between
// `GroupByIterator` and `GroupIter`. In practice they are used consecutive and
// `GroupByIter` is finished before calling next on `GroupByIterator`. I'm not sure there
// is a solution with lifetimes for that, because we would need to enforce it in the usage
// somehow.
//
// One potential solution would be to replace the iterator approach with something similar.
inner: Rc<RefCell<GroupByShared<I, F, K>>>,
}
struct GroupByShared<I, F, K: Copy>
where
I: Iterator,
F: FnMut(&I::Item) -> K,
{
iter: Peekable<I>,
group_by_fn: F,
}
impl<I, F, K> GroupByIterator<I, F, K>
where
I: Iterator,
F: FnMut(&I::Item) -> K,
K: Copy,
{
fn new(inner: I, group_by_fn: F) -> Self {
let inner = GroupByShared {
iter: inner.peekable(),
group_by_fn,
};
Self {
inner: Rc::new(RefCell::new(inner)),
}
}
}
impl<I, F, K> Iterator for GroupByIterator<I, F, K>
where
I: Iterator,
I::Item: Copy,
F: FnMut(&I::Item) -> K,
K: Copy,
{
type Item = (K, GroupIterator<I, F, K>);
fn next(&mut self) -> Option<Self::Item> {
let mut inner = self.inner.borrow_mut();
let value = *inner.iter.peek()?;
let key = (inner.group_by_fn)(&value);
let inner = self.inner.clone();
let group_iter = GroupIterator {
inner,
group_key: key,
};
Some((key, group_iter))
}
}
pub struct GroupIterator<I, F, K: Copy>
where
I: Iterator,
F: FnMut(&I::Item) -> K,
{
inner: Rc<RefCell<GroupByShared<I, F, K>>>,
group_key: K,
}
impl<I, F, K: PartialEq + Copy> Iterator for GroupIterator<I, F, K>
where
I: Iterator,
I::Item: Copy,
F: FnMut(&I::Item) -> K,
{
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
let mut inner = self.inner.borrow_mut();
// peek if next value is in group
let peek_val = *inner.iter.peek()?;
if (inner.group_by_fn)(&peek_val) == self.group_key {
inner.iter.next()
} else {
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn group_by_collect<I: Iterator<Item = u32>>(iter: I) -> Vec<(I::Item, Vec<I::Item>)> {
iter.group_by(|val| val / 10)
.map(|(el, iter)| (el, iter.collect::<Vec<_>>()))
.collect::<Vec<_>>()
}
#[test]
fn group_by_two_groups() {
let vals = vec![1u32, 4, 15];
let grouped_vals = group_by_collect(vals.into_iter());
assert_eq!(grouped_vals, vec![(0, vec![1, 4]), (1, vec![15])]);
}
#[test]
fn group_by_test_empty() {
let vals = vec![];
let grouped_vals = group_by_collect(vals.into_iter());
assert_eq!(grouped_vals, vec![]);
}
#[test]
fn group_by_three_groups() {
let vals = vec![1u32, 4, 15, 1];
let grouped_vals = group_by_collect(vals.into_iter());
assert_eq!(
grouped_vals,
vec![(0, vec![1, 4]), (1, vec![15]), (0, vec![1])]
);
}
}

View File

@@ -6,12 +6,10 @@ pub use byteorder::LittleEndian as Endianness;
mod bitset; mod bitset;
pub mod file_slice; pub mod file_slice;
mod group_by;
mod serialize; mod serialize;
mod vint; mod vint;
mod writer; mod writer;
pub use bitset::*; pub use bitset::*;
pub use group_by::GroupByIteratorExtended;
pub use ownedbytes::{OwnedBytes, StableDeref}; pub use ownedbytes::{OwnedBytes, StableDeref};
pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize}; pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize};
pub use vint::{ pub use vint::{

View File

@@ -27,7 +27,7 @@ fn main() -> tantivy::Result<()> {
let score_fieldtype = let score_fieldtype =
crate::schema::NumericOptions::default().set_fast(Cardinality::SingleValue); crate::schema::NumericOptions::default().set_fast(Cardinality::SingleValue);
let highscore_field = schema_builder.add_f64_field("highscore", score_fieldtype.clone()); let highscore_field = schema_builder.add_f64_field("highscore", score_fieldtype.clone());
let price_field = schema_builder.add_f64_field("price", score_fieldtype); let price_field = schema_builder.add_f64_field("price", score_fieldtype.clone());
let schema = schema_builder.build(); let schema = schema_builder.build();
@@ -112,7 +112,7 @@ fn main() -> tantivy::Result<()> {
], ],
..Default::default() ..Default::default()
}), }),
sub_aggregation: sub_agg_req_1, sub_aggregation: sub_agg_req_1.clone(),
}), }),
)] )]
.into_iter() .into_iter()
@@ -123,7 +123,7 @@ fn main() -> tantivy::Result<()> {
let searcher = reader.searcher(); let searcher = reader.searcher();
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap(); let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap();
let res: Value = serde_json::to_value(agg_res)?; let res: Value = serde_json::to_value(&agg_res)?;
println!("{}", serde_json::to_string_pretty(&res)?); println!("{}", serde_json::to_string_pretty(&res)?);
Ok(()) Ok(())

View File

@@ -14,7 +14,7 @@ use fastfield_codecs::Column;
// Importing tantivy... // Importing tantivy...
use tantivy::collector::{Collector, SegmentCollector}; use tantivy::collector::{Collector, SegmentCollector};
use tantivy::query::QueryParser; use tantivy::query::QueryParser;
use tantivy::schema::{Schema, FAST, INDEXED, TEXT}; use tantivy::schema::{Field, Schema, FAST, INDEXED, TEXT};
use tantivy::{doc, Index, Score, SegmentReader}; use tantivy::{doc, Index, Score, SegmentReader};
#[derive(Default)] #[derive(Default)]
@@ -52,11 +52,11 @@ impl Stats {
} }
struct StatsCollector { struct StatsCollector {
field: String, field: Field,
} }
impl StatsCollector { impl StatsCollector {
fn with_field(field: String) -> StatsCollector { fn with_field(field: Field) -> StatsCollector {
StatsCollector { field } StatsCollector { field }
} }
} }
@@ -73,7 +73,7 @@ impl Collector for StatsCollector {
_segment_local_id: u32, _segment_local_id: u32,
segment_reader: &SegmentReader, segment_reader: &SegmentReader,
) -> tantivy::Result<StatsSegmentCollector> { ) -> tantivy::Result<StatsSegmentCollector> {
let fast_field_reader = segment_reader.fast_fields().u64(&self.field)?; let fast_field_reader = segment_reader.fast_fields().u64(self.field)?;
Ok(StatsSegmentCollector { Ok(StatsSegmentCollector {
fast_field_reader, fast_field_reader,
stats: Stats::default(), stats: Stats::default(),
@@ -171,9 +171,7 @@ fn main() -> tantivy::Result<()> {
// here we want to get a hit on the 'ken' in Frankenstein // here we want to get a hit on the 'ken' in Frankenstein
let query = query_parser.parse_query("broom")?; let query = query_parser.parse_query("broom")?;
if let Some(stats) = if let Some(stats) = searcher.search(&query, &StatsCollector::with_field(price))? {
searcher.search(&query, &StatsCollector::with_field("price".to_string()))?
{
println!("count: {}", stats.count()); println!("count: {}", stats.count());
println!("mean: {}", stats.mean()); println!("mean: {}", stats.mean());
println!("standard deviation: {}", stats.standard_deviation()); println!("standard deviation: {}", stats.standard_deviation());

View File

@@ -1,17 +1,15 @@
// # Faceted Search // # Basic Example
// //
// This example covers the faceted search functionalities of // This example covers the basic functionalities of
// tantivy. // tantivy.
// //
// We will : // We will :
// - define a text field "name" in our schema // - define our schema
// - define a facet field "classification" in our schema // = create an index in a directory
// - create an index in memory // - index few documents in our index
// - index few documents with respective facets in our index // - search for the best document matchings "sea whale"
// - search and count the number of documents that the classifications start the facet "/Felidae" // - retrieve the best document original content.
// - Search the facet "/Felidae/Pantherinae" and count the number of documents that the
// classifications include the facet.
//
// --- // ---
// Importing tantivy... // Importing tantivy...
use tantivy::collector::FacetCollector; use tantivy::collector::FacetCollector;
@@ -23,7 +21,7 @@ fn main() -> tantivy::Result<()> {
// Let's create a temporary directory for the sake of this example // Let's create a temporary directory for the sake of this example
let mut schema_builder = Schema::builder(); let mut schema_builder = Schema::builder();
let name = schema_builder.add_text_field("name", TEXT | STORED); let name = schema_builder.add_text_field("felin_name", TEXT | STORED);
// this is our faceted field: its scientific classification // this is our faceted field: its scientific classification
let classification = schema_builder.add_facet_field("classification", FacetOptions::default()); let classification = schema_builder.add_facet_field("classification", FacetOptions::default());

View File

@@ -27,7 +27,7 @@ fn main() -> Result<()> {
reader.reload()?; reader.reload()?;
let searcher = reader.searcher(); let searcher = reader.searcher();
// The end is excluded i.e. here we are searching up to 1969 // The end is excluded i.e. here we are searching up to 1969
let docs_in_the_sixties = RangeQuery::new_u64("year".to_string(), 1960..1970); let docs_in_the_sixties = RangeQuery::new_u64(year_field, 1960..1970);
// Uses a Count collector to sum the total number of docs in the range // Uses a Count collector to sum the total number of docs in the range
let num_60s_books = searcher.search(&docs_in_the_sixties, &Count)?; let num_60s_books = searcher.search(&docs_in_the_sixties, &Count)?;
assert_eq!(num_60s_books, 10); assert_eq!(num_60s_books, 10);

View File

@@ -1,73 +0,0 @@
// # IP Address example
//
// This example shows how the ip field can be used
// with IpV6 and IpV4.
use tantivy::collector::{Count, TopDocs};
use tantivy::query::QueryParser;
use tantivy::schema::{Schema, FAST, INDEXED, STORED, STRING};
use tantivy::Index;
fn main() -> tantivy::Result<()> {
// # Defining the schema
let mut schema_builder = Schema::builder();
let event_type = schema_builder.add_text_field("event_type", STRING | STORED);
let ip = schema_builder.add_ip_addr_field("ip", STORED | INDEXED | FAST);
let schema = schema_builder.build();
// # Indexing documents
let index = Index::create_in_ram(schema.clone());
let mut index_writer = index.writer(50_000_000)?;
let doc = schema.parse_document(
r#"{
"ip": "192.168.0.33",
"event_type": "login"
}"#,
)?;
index_writer.add_document(doc)?;
let doc = schema.parse_document(
r#"{
"ip": "192.168.0.80",
"event_type": "checkout"
}"#,
)?;
index_writer.add_document(doc)?;
let doc = schema.parse_document(
r#"{
"ip": "2001:0db8:85a3:0000:0000:8a2e:0370:7334",
"event_type": "checkout"
}"#,
)?;
index_writer.add_document(doc)?;
index_writer.commit()?;
let reader = index.reader()?;
let searcher = reader.searcher();
let query_parser = QueryParser::for_index(&index, vec![event_type, ip]);
{
let query = query_parser.parse_query("ip:[192.168.0.0 TO 192.168.0.100]")?;
let count_docs = searcher.search(&*query, &TopDocs::with_limit(5))?;
assert_eq!(count_docs.len(), 2);
}
{
let query = query_parser.parse_query("ip:[192.168.1.0 TO 192.168.1.100]")?;
let count_docs = searcher.search(&*query, &TopDocs::with_limit(2))?;
assert_eq!(count_docs.len(), 0);
}
{
let query = query_parser.parse_query("ip:192.168.0.80")?;
let count_docs = searcher.search(&*query, &Count)?;
assert_eq!(count_docs, 1);
}
{
// IpV6 needs to be escaped because it contains `:`
let query = query_parser.parse_query("ip:\"2001:0db8:85a3:0000:0000:8a2e:0370:7334\"")?;
let count_docs = searcher.search(&*query, &Count)?;
assert_eq!(count_docs, 1);
}
Ok(())
}

View File

@@ -4,7 +4,7 @@ use std::sync::{Arc, RwLock, Weak};
use tantivy::collector::TopDocs; use tantivy::collector::TopDocs;
use tantivy::query::QueryParser; use tantivy::query::QueryParser;
use tantivy::schema::{Schema, FAST, TEXT}; use tantivy::schema::{Field, Schema, FAST, TEXT};
use tantivy::{ use tantivy::{
doc, DocAddress, DocId, Index, IndexReader, Opstamp, Searcher, SearcherGeneration, SegmentId, doc, DocAddress, DocId, Index, IndexReader, Opstamp, Searcher, SearcherGeneration, SegmentId,
SegmentReader, Warmer, SegmentReader, Warmer,
@@ -25,13 +25,13 @@ pub trait PriceFetcher: Send + Sync + 'static {
} }
struct DynamicPriceColumn { struct DynamicPriceColumn {
field: String, field: Field,
price_cache: RwLock<HashMap<(SegmentId, Option<Opstamp>), Arc<Vec<Price>>>>, price_cache: RwLock<HashMap<(SegmentId, Option<Opstamp>), Arc<Vec<Price>>>>,
price_fetcher: Box<dyn PriceFetcher>, price_fetcher: Box<dyn PriceFetcher>,
} }
impl DynamicPriceColumn { impl DynamicPriceColumn {
pub fn with_product_id_field<T: PriceFetcher>(field: String, price_fetcher: T) -> Self { pub fn with_product_id_field<T: PriceFetcher>(field: Field, price_fetcher: T) -> Self {
DynamicPriceColumn { DynamicPriceColumn {
field, field,
price_cache: Default::default(), price_cache: Default::default(),
@@ -48,7 +48,7 @@ impl Warmer for DynamicPriceColumn {
fn warm(&self, searcher: &Searcher) -> tantivy::Result<()> { fn warm(&self, searcher: &Searcher) -> tantivy::Result<()> {
for segment in searcher.segment_readers() { for segment in searcher.segment_readers() {
let key = (segment.segment_id(), segment.delete_opstamp()); let key = (segment.segment_id(), segment.delete_opstamp());
let product_id_reader = segment.fast_fields().u64(&self.field)?; let product_id_reader = segment.fast_fields().u64(self.field)?;
let product_ids: Vec<ProductId> = segment let product_ids: Vec<ProductId> = segment
.doc_ids_alive() .doc_ids_alive()
.map(|doc| product_id_reader.get_val(doc)) .map(|doc| product_id_reader.get_val(doc))
@@ -123,7 +123,7 @@ fn main() -> tantivy::Result<()> {
let price_table = ExternalPriceTable::default(); let price_table = ExternalPriceTable::default();
let price_dynamic_column = Arc::new(DynamicPriceColumn::with_product_id_field( let price_dynamic_column = Arc::new(DynamicPriceColumn::with_product_id_field(
"product_id".to_string(), product_id,
price_table.clone(), price_table.clone(),
)); ));
price_table.update_price(OLIVE_OIL, 12); price_table.update_price(OLIVE_OIL, 12);

View File

@@ -14,7 +14,7 @@ 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/" }
prettytable-rs = {version="0.10.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"

View File

@@ -4,7 +4,7 @@ extern crate test;
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::ops::RangeInclusive; use std::iter;
use std::sync::Arc; use std::sync::Arc;
use common::OwnedBytes; use common::OwnedBytes;
@@ -71,24 +71,27 @@ mod tests {
}); });
} }
const FIFTY_PERCENT_RANGE: RangeInclusive<u64> = 1..=50; fn get_exp_data() -> Vec<u64> {
const SINGLE_ITEM: u64 = 90;
const SINGLE_ITEM_RANGE: RangeInclusive<u64> = 90..=90;
const ONE_PERCENT_ITEM_RANGE: RangeInclusive<u64> = 49..=49;
fn get_data_50percent_item() -> Vec<u128> {
let mut rng = StdRng::from_seed([1u8; 32]);
let mut data = vec![]; let mut data = vec![];
for _ in 0..300_000 { for i in 0..100 {
let val = rng.gen_range(1..=100); let num = i * i;
data.push(val); data.extend(iter::repeat(i as u64).take(num));
} }
data.push(SINGLE_ITEM); data.shuffle(&mut StdRng::from_seed([1u8; 32]));
data.shuffle(&mut rng); // lengt = 328350
let data = data.iter().map(|el| *el as u128).collect::<Vec<_>>();
data data
} }
fn get_data_50percent_item() -> (u128, u128, Vec<u128>) {
let mut permutation = get_exp_data();
let major_item = 20;
let minor_item = 10;
permutation.extend(iter::repeat(major_item).take(permutation.len()));
permutation.shuffle(&mut StdRng::from_seed([1u8; 32]));
let permutation = permutation.iter().map(|el| *el as u128).collect::<Vec<_>>();
(major_item as u128, minor_item as u128, permutation)
}
fn get_u128_column_random() -> Arc<dyn Column<u128>> { fn get_u128_column_random() -> Arc<dyn Column<u128>> {
let permutation = generate_random(); let permutation = generate_random();
let permutation = permutation.iter().map(|el| *el as u128).collect::<Vec<_>>(); let permutation = permutation.iter().map(|el| *el as u128).collect::<Vec<_>>();
@@ -103,82 +106,15 @@ mod tests {
open_u128::<u128>(out).unwrap() open_u128::<u128>(out).unwrap()
} }
// U64 RANGE START
#[bench]
fn bench_intfastfield_getrange_u64_50percent_hit(b: &mut Bencher) {
let data = get_data_50percent_item();
let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
b.iter(|| {
let mut positions = Vec::new();
column.get_docids_for_value_range(
FIFTY_PERCENT_RANGE,
0..data.len() as u32,
&mut positions,
);
positions
});
}
#[bench]
fn bench_intfastfield_getrange_u64_1percent_hit(b: &mut Bencher) {
let data = get_data_50percent_item();
let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
b.iter(|| {
let mut positions = Vec::new();
column.get_docids_for_value_range(
ONE_PERCENT_ITEM_RANGE,
0..data.len() as u32,
&mut positions,
);
positions
});
}
#[bench]
fn bench_intfastfield_getrange_u64_single_hit(b: &mut Bencher) {
let data = get_data_50percent_item();
let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
b.iter(|| {
let mut positions = Vec::new();
column.get_docids_for_value_range(
SINGLE_ITEM_RANGE,
0..data.len() as u32,
&mut positions,
);
positions
});
}
#[bench]
fn bench_intfastfield_getrange_u64_hit_all(b: &mut Bencher) {
let data = get_data_50percent_item();
let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
b.iter(|| {
let mut positions = Vec::new();
column.get_docids_for_value_range(0..=u64::MAX, 0..data.len() as u32, &mut positions);
positions
});
}
// U64 RANGE END
// U128 RANGE START
#[bench] #[bench]
fn bench_intfastfield_getrange_u128_50percent_hit(b: &mut Bencher) { fn bench_intfastfield_getrange_u128_50percent_hit(b: &mut Bencher) {
let data = get_data_50percent_item(); let (major_item, _minor_item, data) = get_data_50percent_item();
let column = get_u128_column_from_data(&data); let column = get_u128_column_from_data(&data);
b.iter(|| { b.iter(|| {
let mut positions = Vec::new(); let mut positions = Vec::new();
column.get_docids_for_value_range( column.get_docids_for_value_range(
*FIFTY_PERCENT_RANGE.start() as u128..=*FIFTY_PERCENT_RANGE.end() as u128, major_item..=major_item,
0..data.len() as u32, 0..data.len() as u32,
&mut positions, &mut positions,
); );
@@ -188,13 +124,13 @@ mod tests {
#[bench] #[bench]
fn bench_intfastfield_getrange_u128_single_hit(b: &mut Bencher) { fn bench_intfastfield_getrange_u128_single_hit(b: &mut Bencher) {
let data = get_data_50percent_item(); let (_major_item, minor_item, data) = get_data_50percent_item();
let column = get_u128_column_from_data(&data); let column = get_u128_column_from_data(&data);
b.iter(|| { b.iter(|| {
let mut positions = Vec::new(); let mut positions = Vec::new();
column.get_docids_for_value_range( column.get_docids_for_value_range(
*SINGLE_ITEM_RANGE.start() as u128..=*SINGLE_ITEM_RANGE.end() as u128, minor_item..=minor_item,
0..data.len() as u32, 0..data.len() as u32,
&mut positions, &mut positions,
); );
@@ -204,7 +140,7 @@ mod tests {
#[bench] #[bench]
fn bench_intfastfield_getrange_u128_hit_all(b: &mut Bencher) { fn bench_intfastfield_getrange_u128_hit_all(b: &mut Bencher) {
let data = get_data_50percent_item(); let (_major_item, _minor_item, data) = get_data_50percent_item();
let column = get_u128_column_from_data(&data); let column = get_u128_column_from_data(&data);
b.iter(|| { b.iter(|| {
@@ -213,7 +149,6 @@ mod tests {
positions positions
}); });
} }
// U128 RANGE END
#[bench] #[bench]
fn bench_intfastfield_scan_all_fflookup_u128(b: &mut Bencher) { fn bench_intfastfield_scan_all_fflookup_u128(b: &mut Bencher) {

View File

@@ -1,4 +1,3 @@
use std::fmt::{self, Debug};
use std::marker::PhantomData; use std::marker::PhantomData;
use std::ops::{Range, RangeInclusive}; use std::ops::{Range, RangeInclusive};
@@ -7,7 +6,7 @@ use tantivy_bitpacker::minmax;
use crate::monotonic_mapping::StrictlyMonotonicFn; use crate::monotonic_mapping::StrictlyMonotonicFn;
/// `Column` provides columnar access on a field. /// `Column` provides columnar access on a field.
pub trait Column<T: PartialOrd + Debug = u64>: Send + Sync { pub trait Column<T: PartialOrd = u64>: Send + Sync {
/// Return the value associated with the given idx. /// Return the value associated with the given idx.
/// ///
/// This accessor should return as fast as possible. /// This accessor should return as fast as possible.
@@ -84,7 +83,7 @@ pub struct VecColumn<'a, T = u64> {
max_value: T, max_value: T,
} }
impl<'a, C: Column<T>, T: Copy + PartialOrd + fmt::Debug> Column<T> for &'a C { impl<'a, C: Column<T>, T: Copy + PartialOrd> Column<T> for &'a C {
fn get_val(&self, idx: u32) -> T { fn get_val(&self, idx: u32) -> T {
(*self).get_val(idx) (*self).get_val(idx)
} }
@@ -110,7 +109,7 @@ impl<'a, C: Column<T>, T: Copy + PartialOrd + fmt::Debug> Column<T> for &'a C {
} }
} }
impl<'a, T: Copy + PartialOrd + Send + Sync + Debug> Column<T> for VecColumn<'a, T> { impl<'a, T: Copy + PartialOrd + Send + Sync> Column<T> for VecColumn<'a, T> {
fn get_val(&self, position: u32) -> T { fn get_val(&self, position: u32) -> T {
self.values[position as usize] self.values[position as usize]
} }
@@ -178,8 +177,8 @@ pub fn monotonic_map_column<C, T, Input, Output>(
where where
C: Column<Input>, C: Column<Input>,
T: StrictlyMonotonicFn<Input, Output> + Send + Sync, T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
Input: PartialOrd + Send + Sync + Copy + Debug, Input: PartialOrd + Send + Sync + Clone,
Output: PartialOrd + Send + Sync + Copy + Debug, Output: PartialOrd + Send + Sync + Clone,
{ {
MonotonicMappingColumn { MonotonicMappingColumn {
from_column, from_column,
@@ -192,8 +191,8 @@ impl<C, T, Input, Output> Column<Output> for MonotonicMappingColumn<C, T, Input>
where where
C: Column<Input>, C: Column<Input>,
T: StrictlyMonotonicFn<Input, Output> + Send + Sync, T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
Input: PartialOrd + Send + Sync + Copy + Debug, Input: PartialOrd + Send + Sync + Clone,
Output: PartialOrd + Send + Sync + Copy + Debug, Output: PartialOrd + Send + Sync + Clone,
{ {
#[inline] #[inline]
fn get_val(&self, idx: u32) -> Output { fn get_val(&self, idx: u32) -> Output {
@@ -229,15 +228,12 @@ where
doc_id_range: Range<u32>, doc_id_range: Range<u32>,
positions: &mut Vec<u32>, positions: &mut Vec<u32>,
) { ) {
if range.start() > &self.max_value() || range.end() < &self.min_value() { self.from_column.get_docids_for_value_range(
return; self.monotonic_mapping.inverse(range.start().clone())
} ..=self.monotonic_mapping.inverse(range.end().clone()),
let range = self.monotonic_mapping.inverse_coerce(range); doc_id_range,
if range.start() > range.end() { positions,
return; )
}
self.from_column
.get_docids_for_value_range(range, doc_id_range, positions)
} }
// We voluntarily do not implement get_range as it yields a regression, // We voluntarily do not implement get_range as it yields a regression,
@@ -258,7 +254,7 @@ where T: Iterator + Clone + ExactSizeIterator
impl<T> Column<T::Item> for IterColumn<T> impl<T> Column<T::Item> for IterColumn<T>
where where
T: Iterator + Clone + ExactSizeIterator + Send + Sync, T: Iterator + Clone + ExactSizeIterator + Send + Sync,
T::Item: PartialOrd + fmt::Debug, T::Item: PartialOrd,
{ {
fn get_val(&self, idx: u32) -> T::Item { fn get_val(&self, idx: u32) -> T::Item {
self.0.clone().nth(idx as usize).unwrap() self.0.clone().nth(idx as usize).unwrap()

View File

@@ -453,8 +453,6 @@ impl CompactSpaceDecompressor {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::fmt;
use super::*; use super::*;
use crate::format_version::read_format_version; use crate::format_version::read_format_version;
use crate::null_index_footer::read_null_index_footer; use crate::null_index_footer::read_null_index_footer;
@@ -708,7 +706,7 @@ mod tests {
); );
} }
fn get_positions_for_value_range_helper<C: Column<T> + ?Sized, T: PartialOrd + fmt::Debug>( fn get_positions_for_value_range_helper<C: Column<T> + ?Sized, T: PartialOrd>(
column: &C, column: &C,
value_range: RangeInclusive<T>, value_range: RangeInclusive<T>,
doc_id_range: Range<u32>, doc_id_range: Range<u32>,

View File

@@ -14,9 +14,9 @@ extern crate more_asserts;
#[cfg(all(test, feature = "unstable"))] #[cfg(all(test, feature = "unstable"))]
extern crate test; extern crate test;
use std::io;
use std::io::Write; use std::io::Write;
use std::sync::Arc; use std::sync::Arc;
use std::{fmt, io};
use common::{BinarySerializable, OwnedBytes}; use common::{BinarySerializable, OwnedBytes};
use compact_space::CompactSpaceDecompressor; use compact_space::CompactSpaceDecompressor;
@@ -133,7 +133,7 @@ impl U128FastFieldCodecType {
} }
/// Returns the correct codec reader wrapped in the `Arc` for the data. /// Returns the correct codec reader wrapped in the `Arc` for the data.
pub fn open_u128<Item: MonotonicallyMappableToU128 + fmt::Debug>( pub fn open_u128<Item: MonotonicallyMappableToU128>(
bytes: OwnedBytes, bytes: OwnedBytes,
) -> io::Result<Arc<dyn Column<Item>>> { ) -> io::Result<Arc<dyn Column<Item>>> {
let (bytes, _format_version) = read_format_version(bytes)?; let (bytes, _format_version) = read_format_version(bytes)?;
@@ -147,9 +147,7 @@ pub fn open_u128<Item: MonotonicallyMappableToU128 + fmt::Debug>(
} }
/// Returns the correct codec reader wrapped in the `Arc` for the data. /// Returns the correct codec reader wrapped in the `Arc` for the data.
pub fn open<T: MonotonicallyMappableToU64 + fmt::Debug>( pub fn open<T: MonotonicallyMappableToU64>(bytes: OwnedBytes) -> io::Result<Arc<dyn Column<T>>> {
bytes: OwnedBytes,
) -> io::Result<Arc<dyn Column<T>>> {
let (bytes, _format_version) = read_format_version(bytes)?; let (bytes, _format_version) = read_format_version(bytes)?;
let (mut bytes, _null_index_footer) = read_null_index_footer(bytes)?; let (mut bytes, _null_index_footer) = read_null_index_footer(bytes)?;
let header = Header::deserialize(&mut bytes)?; let header = Header::deserialize(&mut bytes)?;
@@ -162,7 +160,7 @@ pub fn open<T: MonotonicallyMappableToU64 + fmt::Debug>(
} }
} }
fn open_specific_codec<C: FastFieldCodec, Item: MonotonicallyMappableToU64 + fmt::Debug>( fn open_specific_codec<C: FastFieldCodec, Item: MonotonicallyMappableToU64>(
bytes: OwnedBytes, bytes: OwnedBytes,
header: &Header, header: &Header,
) -> io::Result<Arc<dyn Column<Item>>> { ) -> io::Result<Arc<dyn Column<Item>>> {
@@ -323,9 +321,6 @@ mod tests {
pub fn get_codec_test_datasets() -> Vec<(Vec<u64>, &'static str)> { pub fn get_codec_test_datasets() -> Vec<(Vec<u64>, &'static str)> {
let mut data_and_names = vec![]; let mut data_and_names = vec![];
let data = vec![10];
data_and_names.push((data, "minimal test"));
let data = (10..=10_000_u64).collect::<Vec<_>>(); let data = (10..=10_000_u64).collect::<Vec<_>>();
data_and_names.push((data, "simple monotonically increasing")); data_and_names.push((data, "simple monotonically increasing"));
@@ -333,9 +328,6 @@ mod tests {
vec![5, 6, 7, 8, 9, 10, 99, 100], vec![5, 6, 7, 8, 9, 10, 99, 100],
"offset in linear interpol", "offset in linear interpol",
)); ));
data_and_names.push((vec![3, 18446744073709551613, 5], "docid range regression"));
data_and_names.push((vec![5, 50, 3, 13, 1, 1000, 35], "rand small")); data_and_names.push((vec![5, 50, 3, 13, 1, 1000, 35], "rand small"));
data_and_names.push((vec![10], "single value")); data_and_names.push((vec![10], "single value"));

View File

@@ -1,6 +1,4 @@
use std::fmt;
use std::marker::PhantomData; use std::marker::PhantomData;
use std::ops::RangeInclusive;
use fastdivide::DividerU64; use fastdivide::DividerU64;
@@ -8,9 +6,7 @@ use crate::MonotonicallyMappableToU128;
/// Monotonic maps a value to u64 value space. /// Monotonic maps a value to u64 value space.
/// Monotonic mapping enables `PartialOrd` on u64 space without conversion to original space. /// Monotonic mapping enables `PartialOrd` on u64 space without conversion to original space.
pub trait MonotonicallyMappableToU64: pub trait MonotonicallyMappableToU64: 'static + PartialOrd + Copy + Send + Sync {
'static + PartialOrd + Copy + Send + Sync + fmt::Debug
{
/// Converts a value to u64. /// Converts a value to u64.
/// ///
/// Internally all fast field values are encoded as u64. /// Internally all fast field values are encoded as u64.
@@ -33,29 +29,11 @@ pub trait MonotonicallyMappableToU64:
/// mapping from their range to their domain. The `inverse` method is required when opening a codec, /// mapping from their range to their domain. The `inverse` method is required when opening a codec,
/// so a value can be converted back to its original domain (e.g. ip address or f64) from its /// so a value can be converted back to its original domain (e.g. ip address or f64) from its
/// internal representation. /// internal representation.
pub trait StrictlyMonotonicFn<External: Copy, Internal: Copy> { pub trait StrictlyMonotonicFn<External, Internal> {
/// Strictly monotonically maps the value from External to Internal. /// Strictly monotonically maps the value from External to Internal.
fn mapping(&self, inp: External) -> Internal; fn mapping(&self, inp: External) -> Internal;
/// Inverse of `mapping`. Maps the value from Internal to External. /// Inverse of `mapping`. Maps the value from Internal to External.
fn inverse(&self, out: Internal) -> External; fn inverse(&self, out: Internal) -> External;
/// Maps a user provded value from External to Internal.
/// It may be necessary to coerce the value if it is outside the value space.
/// In that case it tries to find the next greater value in the value space.
///
/// Returns a bool to mark if a value was outside the value space and had to be coerced _up_.
/// With that information we can detect if two values in a range both map outside the same value
/// space.
///
/// coerce_up means the next valid upper value in the value space will be chosen if the value
/// has to be coerced.
fn mapping_coerce(&self, inp: RangeInclusive<External>) -> RangeInclusive<Internal> {
self.mapping(*inp.start())..=self.mapping(*inp.end())
}
/// Inverse of `mapping_coerce`.
fn inverse_coerce(&self, out: RangeInclusive<Internal>) -> RangeInclusive<External> {
self.inverse(*out.start())..=self.inverse(*out.end())
}
} }
/// Inverts a strictly monotonic mapping from `StrictlyMonotonicFn<A, B>` to /// Inverts a strictly monotonic mapping from `StrictlyMonotonicFn<A, B>` to
@@ -76,10 +54,7 @@ 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 where T: StrictlyMonotonicFn<From, To>
T: StrictlyMonotonicFn<From, To>,
From: Copy,
To: Copy,
{ {
#[inline(always)] #[inline(always)]
fn mapping(&self, val: To) -> From { fn mapping(&self, val: To) -> From {
@@ -90,15 +65,6 @@ where
fn inverse(&self, val: From) -> To { fn inverse(&self, val: From) -> To {
self.orig_mapping.mapping(val) self.orig_mapping.mapping(val)
} }
#[inline]
fn mapping_coerce(&self, inp: RangeInclusive<To>) -> RangeInclusive<From> {
self.orig_mapping.inverse_coerce(inp)
}
#[inline]
fn inverse_coerce(&self, out: RangeInclusive<From>) -> RangeInclusive<To> {
self.orig_mapping.mapping_coerce(out)
}
} }
/// Applies the strictly monotonic mapping from `T` without any additional changes. /// Applies the strictly monotonic mapping from `T` without any additional changes.
@@ -176,31 +142,6 @@ impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
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)
} }
#[inline]
#[allow(clippy::reversed_empty_ranges)]
fn mapping_coerce(&self, inp: RangeInclusive<External>) -> RangeInclusive<u64> {
let end = External::to_u64(*inp.end());
if end < self.min_value || inp.end() < inp.start() {
return 1..=0;
}
let map_coerce = |mut inp, coerce_up| {
let inp_lower_bound = self.inverse(0);
if inp < inp_lower_bound {
inp = inp_lower_bound;
}
let val = External::to_u64(inp);
let need_coercion = coerce_up && (val - self.min_value) % self.gcd != 0;
let mut mapped_val = self.mapping(inp);
if need_coercion {
mapped_val += 1;
}
mapped_val
};
let start = map_coerce(*inp.start(), true);
let end = map_coerce(*inp.end(), false);
start..=end
}
} }
/// Strictly monotonic mapping with a base value. /// Strictly monotonic mapping with a base value.
@@ -217,17 +158,6 @@ impl StrictlyMonotonicMappingToInternalBaseval {
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64> impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
for StrictlyMonotonicMappingToInternalBaseval for StrictlyMonotonicMappingToInternalBaseval
{ {
#[inline]
#[allow(clippy::reversed_empty_ranges)]
fn mapping_coerce(&self, inp: RangeInclusive<External>) -> RangeInclusive<u64> {
if External::to_u64(*inp.end()) < self.min_value {
return 1..=0;
}
let start = self.mapping(External::to_u64(*inp.start()).max(self.min_value));
let end = self.mapping(External::to_u64(*inp.end()));
start..=end
}
#[inline(always)] #[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
@@ -311,7 +241,7 @@ mod tests {
test_round_trip::<_, _, u64>(&mapping, 100u64); test_round_trip::<_, _, u64>(&mapping, 100u64);
} }
fn test_round_trip<T: StrictlyMonotonicFn<K, L>, K: std::fmt::Debug + Eq + Copy, L: Copy>( fn test_round_trip<T: StrictlyMonotonicFn<K, L>, K: std::fmt::Debug + Eq + Copy, L>(
mapping: &T, mapping: &T,
test_val: K, test_val: K,
) { ) {

View File

@@ -1,11 +1,8 @@
use std::fmt;
use std::net::Ipv6Addr; use std::net::Ipv6Addr;
/// Montonic maps a value to u128 value space /// Montonic maps a value to u128 value space
/// Monotonic mapping enables `PartialOrd` on u128 space without conversion to original space. /// Monotonic mapping enables `PartialOrd` on u128 space without conversion to original space.
pub trait MonotonicallyMappableToU128: pub trait MonotonicallyMappableToU128: 'static + PartialOrd + Copy + Send + Sync {
'static + PartialOrd + Copy + Send + Sync + fmt::Debug
{
/// Converts a value to u128. /// Converts a value to u128.
/// ///
/// Internally all fast field values are encoded as u64. /// Internally all fast field values are encoded as u64.

View File

@@ -31,16 +31,15 @@ const BLOCK_BITVEC_SIZE: usize = 8;
const BLOCK_OFFSET_SIZE: usize = 4; const BLOCK_OFFSET_SIZE: usize = 4;
const SERIALIZED_BLOCK_SIZE: usize = BLOCK_BITVEC_SIZE + BLOCK_OFFSET_SIZE; const SERIALIZED_BLOCK_SIZE: usize = BLOCK_BITVEC_SIZE + BLOCK_OFFSET_SIZE;
/// Interpreting the bitvec as a list of 64 bits from the low weight to the
/// high weight.
///
/// This function returns the number of bits set to 1 within
/// `[0..pos_in_vec)`.
#[inline] #[inline]
fn count_ones(bitvec: u64, pos_in_bitvec: u32) -> u32 { fn count_ones(bitvec: u64, pos_in_bitvec: u32) -> u32 {
let mask = (1u64 << pos_in_bitvec) - 1; if pos_in_bitvec == 63 {
let masked_bitvec = bitvec & mask; bitvec.count_ones()
masked_bitvec.count_ones() } else {
let mask = (1u64 << (pos_in_bitvec + 1)) - 1;
let masked_bitvec = bitvec & mask;
masked_bitvec.count_ones()
}
} }
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
@@ -67,7 +66,9 @@ impl DenseCodec {
pub fn exists(&self, idx: u32) -> bool { pub fn exists(&self, idx: u32) -> bool {
let block_pos = idx / ELEMENTS_PER_BLOCK; let block_pos = idx / ELEMENTS_PER_BLOCK;
let bitvec = self.dense_index_block(block_pos).bitvec; let bitvec = self.dense_index_block(block_pos).bitvec;
let pos_in_bitvec = idx % ELEMENTS_PER_BLOCK; let pos_in_bitvec = idx % ELEMENTS_PER_BLOCK;
get_bit_at(bitvec, pos_in_bitvec) get_bit_at(bitvec, pos_in_bitvec)
} }
#[inline] #[inline]
@@ -89,7 +90,8 @@ impl DenseCodec {
let pos_in_block_bit_vec = idx % ELEMENTS_PER_BLOCK; let pos_in_block_bit_vec = idx % ELEMENTS_PER_BLOCK;
let ones_in_block = count_ones(index_block.bitvec, pos_in_block_bit_vec); let ones_in_block = count_ones(index_block.bitvec, pos_in_block_bit_vec);
if get_bit_at(index_block.bitvec, pos_in_block_bit_vec) { if get_bit_at(index_block.bitvec, pos_in_block_bit_vec) {
Some(index_block.offset + ones_in_block) // -1 is ok, since idx does exist, so there's at least one
Some(index_block.offset + ones_in_block - 1)
} else { } else {
None None
} }
@@ -317,10 +319,9 @@ mod tests {
set_bit_at(&mut block, 0); set_bit_at(&mut block, 0);
set_bit_at(&mut block, 2); set_bit_at(&mut block, 2);
assert_eq!(count_ones(block, 0), 0); assert_eq!(count_ones(block, 0), 1);
assert_eq!(count_ones(block, 1), 1); assert_eq!(count_ones(block, 1), 1);
assert_eq!(count_ones(block, 2), 1); assert_eq!(count_ones(block, 2), 2);
assert_eq!(count_ones(block, 3), 2);
} }
} }
@@ -346,16 +347,11 @@ mod bench {
codec codec
} }
fn random_range_iterator( fn random_range_iterator(start: u32, end: u32, step_size: u32) -> impl Iterator<Item = u32> {
start: u32,
end: u32,
avg_step_size: u32,
avg_deviation: u32,
) -> impl Iterator<Item = u32> {
let mut rng: StdRng = StdRng::from_seed([1u8; 32]); let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
let mut current = start; let mut current = start;
std::iter::from_fn(move || { std::iter::from_fn(move || {
current += rng.gen_range(avg_step_size - avg_deviation..=avg_step_size + avg_deviation); current += rng.gen_range(1..step_size + 1);
if current >= end { if current >= end {
None None
} else { } else {
@@ -364,17 +360,10 @@ mod bench {
}) })
} }
fn n_percent_step_iterator(percent: f32, num_values: u32) -> impl Iterator<Item = u32> { fn walk_over_data(codec: &DenseCodec, max_step_size: u32) -> Option<u32> {
let ratio = percent as f32 / 100.0;
let step_size = (1f32 / ratio) as u32;
let deviation = step_size - 1;
random_range_iterator(0, num_values, step_size, deviation)
}
fn walk_over_data(codec: &DenseCodec, avg_step_size: u32) -> Option<u32> {
walk_over_data_from_positions( walk_over_data_from_positions(
codec, codec,
random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, 0), random_range_iterator(0, TOTAL_NUM_VALUES, max_step_size),
) )
} }
@@ -390,105 +379,69 @@ mod bench {
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_1percent_filled_10percent_hit(bench: &mut Bencher) { fn bench_dense_codec_translate_orig_to_codec_90percent_filled_random_stride(
let codec = gen_bools(0.01f64); bench: &mut Bencher,
bench.iter(|| walk_over_data(&codec, 100)); ) {
}
#[bench]
fn bench_translate_orig_to_codec_5percent_filled_10percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.05f64);
bench.iter(|| walk_over_data(&codec, 100));
}
#[bench]
fn bench_translate_orig_to_codec_5percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.05f64);
bench.iter(|| walk_over_data(&codec, 1000));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_1percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_10percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_90percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.9f64); let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_10percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_50percent_filled_1percent_hit(bench: &mut Bencher) { fn bench_dense_codec_translate_orig_to_codec_50percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.5f64); let codec = gen_bools(0.5f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_90percent_filled_1percent_hit(bench: &mut Bencher) { fn bench_dense_codec_translate_orig_to_codec_full_scan_10percent(bench: &mut Bencher) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_dense_codec_translate_orig_to_codec_full_scan_90percent(bench: &mut Bencher) {
let codec = gen_bools(0.9f64); let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_dense_codec_translate_orig_to_codec_10percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) { fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step(
let codec = gen_bools(0.01f64); bench: &mut Bencher,
let num_non_nulls = codec.num_non_nulls(); ) {
bench.iter(|| { let codec = gen_bools(0.9f64);
codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
.last()
});
}
#[bench]
fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
let num_non_nulls = codec.num_non_nulls();
bench.iter(|| {
codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(10.0, num_non_nulls))
.last()
});
}
#[bench]
fn bench_translate_codec_to_orig_1percent_filled_full_scan(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
let num_vals = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(0..num_vals) .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) { fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride(
let codec = gen_bools(0.90f64); bench: &mut Bencher,
let num_non_nulls = codec.num_non_nulls(); ) {
let codec = gen_bools(0.9f64);
let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls)) .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) { fn bench_dense_codec_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
let codec = gen_bools(0.9f64); let codec = gen_bools(0.9f64);
let num_vals = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {

View File

@@ -1,6 +1,6 @@
use std::io::{self, Write}; use std::io::{self, Write};
use common::{BitSet, GroupByIteratorExtended, OwnedBytes}; use common::{BitSet, OwnedBytes};
use super::{serialize_dense_codec, DenseCodec}; use super::{serialize_dense_codec, DenseCodec};
@@ -78,22 +78,12 @@ struct DenseBlock {
} }
impl DenseBlock { impl DenseBlock {
#[inline]
pub fn exists(&self, idx: u32) -> bool { pub fn exists(&self, idx: u32) -> bool {
self.codec.exists(idx) self.codec.exists(idx)
} }
#[inline]
pub fn translate_to_codec_idx(&self, idx: u32) -> Option<u32> { pub fn translate_to_codec_idx(&self, idx: u32) -> Option<u32> {
self.codec.translate_to_codec_idx(idx) self.codec.translate_to_codec_idx(idx)
} }
#[inline]
pub fn translate_codec_idx_to_original_idx_iter<'a>(
&'a self,
iter: impl Iterator<Item = u32> + 'a,
) -> impl Iterator<Item = u32> + 'a {
self.codec.translate_codec_idx_to_original_idx(iter)
}
#[inline]
pub fn translate_codec_idx_to_original_idx(&self, idx: u32) -> u32 { pub fn translate_codec_idx_to_original_idx(&self, idx: u32) -> u32 {
self.codec self.codec
.translate_codec_idx_to_original_idx(idx..=idx) .translate_codec_idx_to_original_idx(idx..=idx)
@@ -217,7 +207,6 @@ struct ValueAddr {
} }
/// Splits a idx into block index and value in the block /// Splits a idx into block index and value in the block
#[inline]
fn value_addr(idx: u32) -> ValueAddr { fn value_addr(idx: u32) -> ValueAddr {
/// Static assert number elements per block this method expects /// Static assert number elements per block this method expects
#[allow(clippy::assertions_on_constants)] #[allow(clippy::assertions_on_constants)]
@@ -284,7 +273,6 @@ impl SparseCodec {
} }
} }
#[inline]
fn find_block(&self, dense_idx: u32, mut block_pos: u32) -> u32 { fn find_block(&self, dense_idx: u32, mut block_pos: u32) -> u32 {
loop { loop {
let offset = self.blocks[block_pos as usize].offset(); let offset = self.blocks[block_pos as usize].offset();
@@ -296,7 +284,6 @@ impl SparseCodec {
} }
/// Translate positions from the codec index to the original index. /// Translate positions from the codec index to the original index.
/// Correctness: Provided values must be in increasing values
/// ///
/// # Panics /// # Panics
/// ///
@@ -305,41 +292,35 @@ impl SparseCodec {
&'a self, &'a self,
iter: impl Iterator<Item = u32> + 'a, iter: impl Iterator<Item = u32> + 'a,
) -> impl Iterator<Item = u32> + 'a { ) -> impl Iterator<Item = u32> + 'a {
// TODO: There's a big potential performance gain, by using iterators per block instead of
// random access for each element in a block
// group_by itertools won't help though, since it requires a temporary local variable
let mut block_pos = 0u32; let mut block_pos = 0u32;
iter.group_by(move |codec_idx| { iter.map(move |codec_idx| {
block_pos = self.find_block(*codec_idx, block_pos); // update block_pos to limit search scope
block_pos block_pos = self.find_block(codec_idx, block_pos);
})
.flat_map(move |(block_pos, block_iter)| {
let block_doc_idx_start = block_pos * ELEMENTS_PER_BLOCK; let block_doc_idx_start = block_pos * ELEMENTS_PER_BLOCK;
let block = &self.blocks[block_pos as usize]; let block = &self.blocks[block_pos as usize];
let offset = block.offset(); let idx_in_block = codec_idx - block.offset();
let indexes_in_block_iter = block_iter.map(move |codec_idx| codec_idx - offset);
match block { match block {
SparseCodecBlockVariant::Empty { offset: _ } => { SparseCodecBlockVariant::Empty { offset: _ } => {
panic!( panic!(
"invalid input, cannot translate to original index. associated empty \ "invalid input, cannot translate to original index. associated empty \
block with dense idx. block_pos {}, idx_in_block {:?}", block with dense idx. block_pos {}, idx_in_block {}",
block_pos, block_pos, idx_in_block
indexes_in_block_iter.collect::<Vec<_>>()
) )
} }
SparseCodecBlockVariant::Dense(dense) => { SparseCodecBlockVariant::Dense(dense) => {
Box::new(dense.translate_codec_idx_to_original_idx_iter(indexes_in_block_iter)) dense.translate_codec_idx_to_original_idx(idx_in_block) + block_doc_idx_start
as Box<dyn Iterator<Item = u32>>
} }
SparseCodecBlockVariant::Sparse(block) => { SparseCodecBlockVariant::Sparse(block) => {
Box::new(indexes_in_block_iter.map(move |idx_in_block| { block.value_at_idx(&self.data, idx_in_block as u16) as u32 + block_doc_idx_start
block.value_at_idx(&self.data, idx_in_block as u16) as u32
}))
} }
} }
.map(move |idx| idx + block_doc_idx_start)
}) })
} }
} }
#[inline]
fn is_sparse(num_elem_in_block: u32) -> bool { fn is_sparse(num_elem_in_block: u32) -> bool {
num_elem_in_block < DENSE_BLOCK_THRESHOLD num_elem_in_block < DENSE_BLOCK_THRESHOLD
} }
@@ -614,16 +595,11 @@ mod bench {
codec codec
} }
fn random_range_iterator( fn random_range_iterator(start: u32, end: u32, step_size: u32) -> impl Iterator<Item = u32> {
start: u32,
end: u32,
avg_step_size: u32,
avg_deviation: u32,
) -> impl Iterator<Item = u32> {
let mut rng: StdRng = StdRng::from_seed([1u8; 32]); let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
let mut current = start; let mut current = start;
std::iter::from_fn(move || { std::iter::from_fn(move || {
current += rng.gen_range(avg_step_size - avg_deviation..=avg_step_size + avg_deviation); current += rng.gen_range(1..step_size + 1);
if current >= end { if current >= end {
None None
} else { } else {
@@ -632,17 +608,10 @@ mod bench {
}) })
} }
fn n_percent_step_iterator(percent: f32, num_values: u32) -> impl Iterator<Item = u32> { fn walk_over_data(codec: &SparseCodec, max_step_size: u32) -> Option<u32> {
let ratio = percent as f32 / 100.0;
let step_size = (1f32 / ratio) as u32;
let deviation = step_size - 1;
random_range_iterator(0, num_values, step_size, deviation)
}
fn walk_over_data(codec: &SparseCodec, avg_step_size: u32) -> Option<u32> {
walk_over_data_from_positions( walk_over_data_from_positions(
codec, codec,
random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, 0), random_range_iterator(0, TOTAL_NUM_VALUES, max_step_size),
) )
} }
@@ -658,83 +627,83 @@ mod bench {
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_1percent_filled_10percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_orig_to_codec_1percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.01f64); let codec = gen_bools(0.01f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_5percent_filled_10percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_orig_to_codec_5percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.05f64); let codec = gen_bools(0.05f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_5percent_filled_1percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_orig_to_codec_full_scan_10percent(bench: &mut Bencher) {
let codec = gen_bools(0.05f64);
bench.iter(|| walk_over_data(&codec, 1000));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_1percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_translate_orig_to_codec_full_scan_10percent_filled(bench: &mut Bencher) {
let codec = gen_bools(0.1f64); let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES)); bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_full_scan_90percent_filled(bench: &mut Bencher) { fn bench_sparse_codec_translate_orig_to_codec_full_scan_90percent(bench: &mut Bencher) {
let codec = gen_bools(0.9f64); let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES)); bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_10percent_filled_1percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_orig_to_codec_full_scan_1percent(bench: &mut Bencher) {
let codec = gen_bools(0.01f64);
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
}
#[bench]
fn bench_sparse_codec_translate_orig_to_codec_10percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.1f64); let codec = gen_bools(0.1f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_orig_to_codec_50percent_filled_1percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_orig_to_codec_90percent_filled_random_stride(
let codec = gen_bools(0.5f64); bench: &mut Bencher,
bench.iter(|| walk_over_data(&codec, 100)); ) {
}
#[bench]
fn bench_translate_orig_to_codec_90percent_filled_1percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.9f64); let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_random_stride_big_step(
bench: &mut Bencher,
) {
let codec = gen_bools(0.01f64); let codec = gen_bools(0.01f64);
let num_non_nulls = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls)) .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.01f64); let codec = gen_bools(0.01f64);
let num_non_nulls = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(10.0, num_non_nulls)) .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_1percent_filled_full_scan(bench: &mut Bencher) { fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_full_scan(bench: &mut Bencher) {
let codec = gen_bools(0.01f64); let codec = gen_bools(0.01f64);
let num_vals = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
@@ -745,18 +714,33 @@ mod bench {
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) { fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step(
bench: &mut Bencher,
) {
let codec = gen_bools(0.90f64); let codec = gen_bools(0.90f64);
let num_non_nulls = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls)) .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) { fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_random_stride(
bench: &mut Bencher,
) {
let codec = gen_bools(0.9f64);
let num_vals = codec.num_non_nulls();
bench.iter(|| {
codec
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
.last()
});
}
#[bench]
fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
let codec = gen_bools(0.9f64); let codec = gen_bools(0.9f64);
let num_vals = codec.num_non_nulls(); let num_vals = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {

View File

@@ -17,9 +17,9 @@
// You should have received a copy of the GNU Affero General Public License // 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/>. // along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::io;
use std::num::NonZeroU64; use std::num::NonZeroU64;
use std::sync::Arc; use std::sync::Arc;
use std::{fmt, io};
use common::{BinarySerializable, OwnedBytes, VInt}; use common::{BinarySerializable, OwnedBytes, VInt};
use log::warn; use log::warn;
@@ -167,7 +167,7 @@ impl BinarySerializable for Header {
/// Return estimated compression for given codec in the value range [0.0..1.0], where 1.0 means no /// Return estimated compression for given codec in the value range [0.0..1.0], where 1.0 means no
/// compression. /// compression.
pub fn estimate<T: MonotonicallyMappableToU64 + fmt::Debug>( pub fn estimate<T: MonotonicallyMappableToU64>(
typed_column: impl Column<T>, typed_column: impl Column<T>,
codec_type: FastFieldCodecType, codec_type: FastFieldCodecType,
) -> Option<f32> { ) -> Option<f32> {
@@ -276,7 +276,7 @@ pub fn serialize_u128_new<F: Fn() -> I, I: Iterator<Item = u128>>(
} }
/// Serializes the column with the codec with the best estimate on the data. /// Serializes the column with the codec with the best estimate on the data.
pub fn serialize<T: MonotonicallyMappableToU64 + fmt::Debug>( pub fn serialize<T: MonotonicallyMappableToU64>(
typed_column: impl Column<T>, typed_column: impl Column<T>,
output: &mut impl io::Write, output: &mut impl io::Write,
codecs: &[FastFieldCodecType], codecs: &[FastFieldCodecType],
@@ -285,7 +285,7 @@ pub fn serialize<T: MonotonicallyMappableToU64 + fmt::Debug>(
} }
/// Serializes the column with the codec with the best estimate on the data. /// Serializes the column with the codec with the best estimate on the data.
pub fn serialize_new<T: MonotonicallyMappableToU64 + fmt::Debug>( pub fn serialize_new<T: MonotonicallyMappableToU64>(
value_index: ValueIndexInfo, value_index: ValueIndexInfo,
typed_column: impl Column<T>, typed_column: impl Column<T>,
output: &mut impl io::Write, output: &mut impl io::Write,
@@ -366,7 +366,7 @@ fn serialize_given_codec(
} }
/// Helper function to serialize a column (autodetect from all codecs) and then open it /// Helper function to serialize a column (autodetect from all codecs) and then open it
pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default + fmt::Debug>( pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default>(
column: &[T], column: &[T],
) -> Arc<dyn Column<T>> { ) -> Arc<dyn Column<T>> {
let mut buffer = Vec::new(); let mut buffer = Vec::new();
@@ -402,8 +402,8 @@ mod tests {
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let col = VecColumn::from(&[false, true][..]); let col = VecColumn::from(&[false, true][..]);
serialize(col, &mut buffer, &ALL_CODEC_TYPES).unwrap(); serialize(col, &mut buffer, &ALL_CODEC_TYPES).unwrap();
// 5 bytes of header, 1 byte of value // 5 bytes of header, 1 byte of value, 7 bytes of padding.
assert_eq!(buffer.len(), 3 + 5 + 1 + 4 + 2); assert_eq!(buffer.len(), 3 + 5 + 8 + 4 + 2);
} }
#[test] #[test]
@@ -411,8 +411,8 @@ mod tests {
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let col = VecColumn::from(&[true][..]); let col = VecColumn::from(&[true][..]);
serialize(col, &mut buffer, &ALL_CODEC_TYPES).unwrap(); serialize(col, &mut buffer, &ALL_CODEC_TYPES).unwrap();
// 5 bytes of header, 0 bytes of value // 5 bytes of header, 0 bytes of value, 7 bytes of padding.
assert_eq!(buffer.len(), 3 + 5 + 4 + 2); assert_eq!(buffer.len(), 3 + 5 + 7 + 4 + 2);
} }
#[test] #[test]
@@ -422,6 +422,6 @@ mod tests {
let col = VecColumn::from(&vals[..]); let col = VecColumn::from(&vals[..]);
serialize(col, &mut buffer, &[FastFieldCodecType::Bitpacked]).unwrap(); serialize(col, &mut buffer, &[FastFieldCodecType::Bitpacked]).unwrap();
// Values are stored over 3 bits. // Values are stored over 3 bits.
assert_eq!(buffer.len(), 3 + 7 + (3 * 80 / 8) + 4 + 2); assert_eq!(buffer.len(), 3 + 7 + (3 * 80 / 8) + 7 + 4 + 2);
} }
} }

2
run-tests.sh Executable file
View File

@@ -0,0 +1,2 @@
#!/bin/bash
cargo test

View File

@@ -51,10 +51,7 @@ use serde::{Deserialize, Serialize};
pub use super::bucket::RangeAggregation; pub use super::bucket::RangeAggregation;
use super::bucket::{HistogramAggregation, TermsAggregation}; use super::bucket::{HistogramAggregation, TermsAggregation};
use super::metric::{ use super::metric::{AverageAggregation, StatsAggregation};
AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, StatsAggregation,
SumAggregation,
};
use super::VecWithNames; use super::VecWithNames;
/// The top-level aggregation request structure, which contains [`Aggregation`] and their user /// The top-level aggregation request structure, which contains [`Aggregation`] and their user
@@ -240,38 +237,20 @@ impl BucketAggregationType {
/// called multi-value numeric metrics aggregation. /// called multi-value numeric metrics aggregation.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum MetricAggregation { pub enum MetricAggregation {
/// Computes the average of the extracted values. /// Calculates the average.
#[serde(rename = "avg")] #[serde(rename = "avg")]
Average(AverageAggregation), Average(AverageAggregation),
/// Counts the number of extracted values. /// Calculates stats sum, average, min, max, standard_deviation on a field.
#[serde(rename = "value_count")]
Count(CountAggregation),
/// Finds the maximum value.
#[serde(rename = "max")]
Max(MaxAggregation),
/// Finds the minimum value.
#[serde(rename = "min")]
Min(MinAggregation),
/// Computes a collection of statistics (`min`, `max`, `sum`, `count`, and `avg`) over the
/// extracted values.
#[serde(rename = "stats")] #[serde(rename = "stats")]
Stats(StatsAggregation), Stats(StatsAggregation),
/// Computes the sum of the extracted values.
#[serde(rename = "sum")]
Sum(SumAggregation),
} }
impl MetricAggregation { impl MetricAggregation {
fn get_fast_field_names(&self, fast_field_names: &mut HashSet<String>) { fn get_fast_field_names(&self, fast_field_names: &mut HashSet<String>) {
let fast_field_name = match self { match self {
MetricAggregation::Average(avg) => avg.field_name(), MetricAggregation::Average(avg) => fast_field_names.insert(avg.field.to_string()),
MetricAggregation::Count(count) => count.field_name(), MetricAggregation::Stats(stats) => fast_field_names.insert(stats.field.to_string()),
MetricAggregation::Max(max) => max.field_name(),
MetricAggregation::Min(min) => min.field_name(),
MetricAggregation::Stats(stats) => stats.field_name(),
MetricAggregation::Sum(sum) => sum.field_name(),
}; };
fast_field_names.insert(fast_field_name.to_string());
} }
} }
@@ -279,38 +258,6 @@ impl MetricAggregation {
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn test_metric_aggregations_deser() {
let agg_req_json = r#"{
"price_avg": { "avg": { "field": "price" } },
"price_count": { "value_count": { "field": "price" } },
"price_max": { "max": { "field": "price" } },
"price_min": { "min": { "field": "price" } },
"price_stats": { "stats": { "field": "price" } },
"price_sum": { "sum": { "field": "price" } }
}"#;
let agg_req: Aggregations = serde_json::from_str(agg_req_json).unwrap();
assert!(
matches!(agg_req.get("price_avg").unwrap(), Aggregation::Metric(MetricAggregation::Average(avg)) if avg.field == "price")
);
assert!(
matches!(agg_req.get("price_count").unwrap(), Aggregation::Metric(MetricAggregation::Count(count)) if count.field == "price")
);
assert!(
matches!(agg_req.get("price_max").unwrap(), Aggregation::Metric(MetricAggregation::Max(max)) if max.field == "price")
);
assert!(
matches!(agg_req.get("price_min").unwrap(), Aggregation::Metric(MetricAggregation::Min(min)) if min.field == "price")
);
assert!(
matches!(agg_req.get("price_stats").unwrap(), Aggregation::Metric(MetricAggregation::Stats(stats)) if stats.field == "price")
);
assert!(
matches!(agg_req.get("price_sum").unwrap(), Aggregation::Metric(MetricAggregation::Sum(sum)) if sum.field == "price")
);
}
#[test] #[test]
fn serialize_to_json_test() { fn serialize_to_json_test() {
let agg_req1: Aggregations = vec![( let agg_req1: Aggregations = vec![(

View File

@@ -8,10 +8,7 @@ use fastfield_codecs::Column;
use super::agg_req::{Aggregation, Aggregations, BucketAggregationType, MetricAggregation}; use super::agg_req::{Aggregation, Aggregations, BucketAggregationType, MetricAggregation};
use super::bucket::{HistogramAggregation, RangeAggregation, TermsAggregation}; use super::bucket::{HistogramAggregation, RangeAggregation, TermsAggregation};
use super::metric::{ use super::metric::{AverageAggregation, StatsAggregation};
AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, StatsAggregation,
SumAggregation,
};
use super::segment_agg_result::BucketCount; use super::segment_agg_result::BucketCount;
use super::VecWithNames; use super::VecWithNames;
use crate::fastfield::{type_and_cardinality, MultiValuedFastFieldReader}; use crate::fastfield::{type_and_cardinality, MultiValuedFastFieldReader};
@@ -94,7 +91,10 @@ impl BucketAggregationWithAccessor {
BucketAggregationType::Terms(TermsAggregation { BucketAggregationType::Terms(TermsAggregation {
field: field_name, .. field: field_name, ..
}) => { }) => {
let field = reader.schema().get_field(field_name)?; let field = reader
.schema()
.get_field(field_name)
.ok_or_else(|| TantivyError::FieldNotFound(field_name.to_string()))?;
inverted_index = Some(reader.inverted_index(field)?); inverted_index = Some(reader.inverted_index(field)?);
get_ff_reader_and_validate(reader, field_name, Cardinality::MultiValues)? get_ff_reader_and_validate(reader, field_name, Cardinality::MultiValues)?
} }
@@ -134,11 +134,7 @@ impl MetricAggregationWithAccessor {
) -> crate::Result<MetricAggregationWithAccessor> { ) -> crate::Result<MetricAggregationWithAccessor> {
match &metric { match &metric {
MetricAggregation::Average(AverageAggregation { field: field_name }) MetricAggregation::Average(AverageAggregation { field: field_name })
| MetricAggregation::Count(CountAggregation { field: field_name }) | MetricAggregation::Stats(StatsAggregation { field: field_name }) => {
| MetricAggregation::Max(MaxAggregation { field: field_name })
| MetricAggregation::Min(MinAggregation { field: field_name })
| MetricAggregation::Stats(StatsAggregation { field: field_name })
| MetricAggregation::Sum(SumAggregation { field: field_name }) => {
let (accessor, field_type) = let (accessor, field_type) =
get_ff_reader_and_validate(reader, field_name, Cardinality::SingleValue)?; get_ff_reader_and_validate(reader, field_name, Cardinality::SingleValue)?;
@@ -192,7 +188,10 @@ fn get_ff_reader_and_validate(
field_name: &str, field_name: &str,
cardinality: Cardinality, cardinality: Cardinality,
) -> crate::Result<(FastFieldAccessor, Type)> { ) -> crate::Result<(FastFieldAccessor, Type)> {
let field = reader.schema().get_field(field_name)?; let field = reader
.schema()
.get_field(field_name)
.ok_or_else(|| TantivyError::FieldNotFound(field_name.to_string()))?;
let field_type = reader.schema().get_field_entry(field).field_type(); let field_type = reader.schema().get_field_entry(field).field_type();
if let Some((_ff_type, field_cardinality)) = type_and_cardinality(field_type) { if let Some((_ff_type, field_cardinality)) = type_and_cardinality(field_type) {
@@ -212,10 +211,10 @@ fn get_ff_reader_and_validate(
let ff_fields = reader.fast_fields(); let ff_fields = reader.fast_fields();
match cardinality { match cardinality {
Cardinality::SingleValue => ff_fields Cardinality::SingleValue => ff_fields
.u64_lenient(field_name) .u64_lenient(field)
.map(|field| (FastFieldAccessor::Single(field), field_type.value_type())), .map(|field| (FastFieldAccessor::Single(field), field_type.value_type())),
Cardinality::MultiValues => ff_fields Cardinality::MultiValues => ff_fields
.u64s_lenient(field_name) .u64s_lenient(field)
.map(|field| (FastFieldAccessor::Multi(field), field_type.value_type())), .map(|field| (FastFieldAccessor::Multi(field), field_type.value_type())),
} }
} }

View File

@@ -30,7 +30,7 @@ impl AggregationResults {
} else { } else {
// Validation is be done during request parsing, so we can't reach this state. // Validation is be done during request parsing, so we can't reach this state.
Err(TantivyError::InternalError(format!( Err(TantivyError::InternalError(format!(
"Can't find aggregation {:?} in sub-aggregations", "Can't find aggregation {:?} in sub_aggregations",
name name
))) )))
} }
@@ -70,51 +70,27 @@ impl AggregationResult {
pub enum MetricResult { pub enum MetricResult {
/// Average metric result. /// Average metric result.
Average(SingleMetricResult), Average(SingleMetricResult),
/// Count metric result.
Count(SingleMetricResult),
/// Max metric result.
Max(SingleMetricResult),
/// Min metric result.
Min(SingleMetricResult),
/// Stats metric result. /// Stats metric result.
Stats(Stats), Stats(Stats),
/// Sum metric result.
Sum(SingleMetricResult),
} }
impl MetricResult { impl MetricResult {
fn get_value(&self, agg_property: &str) -> crate::Result<Option<f64>> { fn get_value(&self, agg_property: &str) -> crate::Result<Option<f64>> {
match self { match self {
MetricResult::Average(avg) => Ok(avg.value), MetricResult::Average(avg) => Ok(avg.value),
MetricResult::Count(count) => Ok(count.value),
MetricResult::Max(max) => Ok(max.value),
MetricResult::Min(min) => Ok(min.value),
MetricResult::Stats(stats) => stats.get_value(agg_property), MetricResult::Stats(stats) => stats.get_value(agg_property),
MetricResult::Sum(sum) => Ok(sum.value),
} }
} }
} }
impl From<IntermediateMetricResult> for MetricResult { impl From<IntermediateMetricResult> for MetricResult {
fn from(metric: IntermediateMetricResult) -> Self { fn from(metric: IntermediateMetricResult) -> Self {
match metric { match metric {
IntermediateMetricResult::Average(intermediate_avg) => { IntermediateMetricResult::Average(avg_data) => {
MetricResult::Average(intermediate_avg.finalize().into()) MetricResult::Average(avg_data.finalize().into())
}
IntermediateMetricResult::Count(intermediate_count) => {
MetricResult::Count(intermediate_count.finalize().into())
}
IntermediateMetricResult::Max(intermediate_max) => {
MetricResult::Max(intermediate_max.finalize().into())
}
IntermediateMetricResult::Min(intermediate_min) => {
MetricResult::Min(intermediate_min.finalize().into())
} }
IntermediateMetricResult::Stats(intermediate_stats) => { IntermediateMetricResult::Stats(intermediate_stats) => {
MetricResult::Stats(intermediate_stats.finalize()) MetricResult::Stats(intermediate_stats.finalize())
} }
IntermediateMetricResult::Sum(intermediate_sum) => {
MetricResult::Sum(intermediate_sum.finalize().into())
}
} }
} }
} }
@@ -124,13 +100,13 @@ impl From<IntermediateMetricResult> for MetricResult {
#[serde(untagged)] #[serde(untagged)]
pub enum BucketResult { pub enum BucketResult {
/// This is the range entry for a bucket, which contains a key, count, from, to, and optionally /// This is the range entry for a bucket, which contains a key, count, from, to, and optionally
/// sub-aggregations. /// sub_aggregations.
Range { Range {
/// The range buckets sorted by range. /// The range buckets sorted by range.
buckets: BucketEntries<RangeBucketEntry>, buckets: BucketEntries<RangeBucketEntry>,
}, },
/// This is the histogram entry for a bucket, which contains a key, count, and optionally /// This is the histogram entry for a bucket, which contains a key, count, and optionally
/// sub-aggregations. /// sub_aggregations.
Histogram { Histogram {
/// The buckets. /// The buckets.
/// ///
@@ -175,7 +151,7 @@ pub enum BucketEntries<T> {
} }
/// This is the default entry for a bucket, which contains a key, count, and optionally /// This is the default entry for a bucket, which contains a key, count, and optionally
/// sub-aggregations. /// sub_aggregations.
/// ///
/// # JSON Format /// # JSON Format
/// ```json /// ```json
@@ -225,7 +201,7 @@ impl GetDocCount for BucketEntry {
} }
/// This is the range entry for a bucket, which contains a key, count, and optionally /// This is the range entry for a bucket, which contains a key, count, and optionally
/// sub-aggregations. /// sub_aggregations.
/// ///
/// # JSON Format /// # JSON Format
/// ```json /// ```json
@@ -261,7 +237,7 @@ pub struct RangeBucketEntry {
/// Number of documents in the bucket. /// Number of documents in the bucket.
pub doc_count: u64, pub doc_count: u64,
#[serde(flatten)] #[serde(flatten)]
/// Sub-aggregations in this bucket. /// sub-aggregations in this bucket.
pub sub_aggregation: AggregationResults, pub sub_aggregation: AggregationResults,
/// The from range of the bucket. Equals `f64::MIN` when `None`. /// The from range of the bucket. Equals `f64::MIN` when `None`.
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]

View File

@@ -548,7 +548,9 @@ pub(crate) fn intermediate_histogram_buckets_to_final_buckets(
}; };
// If we have a date type on the histogram buckets, we add the `key_as_string` field as rfc339 // If we have a date type on the histogram buckets, we add the `key_as_string` field as rfc339
let field = schema.get_field(&histogram_req.field)?; let field = schema
.get_field(&histogram_req.field)
.ok_or_else(|| TantivyError::FieldNotFound(histogram_req.field.to_string()))?;
if schema.get_field_entry(field).field_type().is_date() { if schema.get_field_entry(field).field_type().is_date() {
for bucket in buckets.iter_mut() { for bucket in buckets.iter_mut() {
if let crate::aggregation::Key::F64(val) = bucket.key { if let crate::aggregation::Key::F64(val) = bucket.key {
@@ -1364,6 +1366,7 @@ mod tests {
"min": Value::Null, "min": Value::Null,
"max": Value::Null, "max": Value::Null,
"avg": Value::Null, "avg": Value::Null,
"standard_deviation": Value::Null,
} }
}) })
); );

View File

@@ -362,19 +362,13 @@ impl SegmentTermCollector {
let mut entries: Vec<(u32, TermBucketEntry)> = let mut entries: Vec<(u32, TermBucketEntry)> =
self.term_buckets.entries.into_iter().collect(); self.term_buckets.entries.into_iter().collect();
let order_by_key = self.req.order.target == OrderTarget::Key;
let order_by_sub_aggregation = let order_by_sub_aggregation =
matches!(self.req.order.target, OrderTarget::SubAggregation(_)); matches!(self.req.order.target, OrderTarget::SubAggregation(_));
match self.req.order.target { match self.req.order.target {
OrderTarget::Key => { OrderTarget::Key => {
// We rely on the fact, that term ordinals match the order of the strings // defer order and cut_off after loading the texts from the dictionary
// TODO: We could have a special collector, that keeps only TOP n results at any
// time.
if self.req.order.order == Order::Desc {
entries.sort_unstable_by_key(|bucket| std::cmp::Reverse(bucket.0));
} else {
entries.sort_unstable_by_key(|bucket| bucket.0);
}
} }
OrderTarget::SubAggregation(_name) => { OrderTarget::SubAggregation(_name) => {
// don't sort and cut off since it's hard to make assumptions on the quality of the // don't sort and cut off since it's hard to make assumptions on the quality of the
@@ -390,11 +384,12 @@ impl SegmentTermCollector {
} }
} }
let (term_doc_count_before_cutoff, sum_other_doc_count) = if order_by_sub_aggregation { let (term_doc_count_before_cutoff, mut sum_other_doc_count) =
(0, 0) if order_by_key || order_by_sub_aggregation {
} else { (0, 0)
cut_off_buckets(&mut entries, self.req.segment_size as usize) } else {
}; cut_off_buckets(&mut entries, self.req.segment_size as usize)
};
let inverted_index = agg_with_accessor let inverted_index = agg_with_accessor
.inverted_index .inverted_index
@@ -417,10 +412,6 @@ impl SegmentTermCollector {
if self.req.min_doc_count == 0 { if self.req.min_doc_count == 0 {
let mut stream = term_dict.stream()?; let mut stream = term_dict.stream()?;
while let Some((key, _ord)) = stream.next() { while let Some((key, _ord)) = stream.next() {
if dict.len() >= self.req.segment_size as usize {
break;
}
let key = std::str::from_utf8(key) let key = std::str::from_utf8(key)
.map_err(|utf8_err| DataCorruption::comment_only(utf8_err.to_string()))?; .map_err(|utf8_err| DataCorruption::comment_only(utf8_err.to_string()))?;
if !dict.contains_key(key) { if !dict.contains_key(key) {
@@ -429,6 +420,20 @@ impl SegmentTermCollector {
} }
} }
if order_by_key {
let mut dict_entries = dict.into_iter().collect_vec();
if self.req.order.order == Order::Desc {
dict_entries.sort_unstable_by(|(key1, _), (key2, _)| key1.cmp(key2));
} else {
dict_entries.sort_unstable_by(|(key1, _), (key2, _)| key2.cmp(key1));
}
let (_, sum_other_docs) =
cut_off_buckets(&mut dict_entries, self.req.segment_size as usize);
sum_other_doc_count += sum_other_docs;
dict = dict_entries.into_iter().collect();
}
Ok(IntermediateBucketResult::Terms( Ok(IntermediateBucketResult::Terms(
IntermediateTermBucketResult { IntermediateTermBucketResult {
entries: dict, entries: dict,
@@ -918,14 +923,14 @@ mod tests {
]; ];
let index = get_test_index_from_values_and_terms(merge_segments, &segment_and_terms)?; let index = get_test_index_from_values_and_terms(merge_segments, &segment_and_terms)?;
// key asc // key desc
let agg_req: Aggregations = vec![( let agg_req: Aggregations = vec![(
"my_texts".to_string(), "my_texts".to_string(),
Aggregation::Bucket(BucketAggregation { Aggregation::Bucket(BucketAggregation {
bucket_agg: BucketAggregationType::Terms(TermsAggregation { bucket_agg: BucketAggregationType::Terms(TermsAggregation {
field: "string_id".to_string(), field: "string_id".to_string(),
order: Some(CustomOrder { order: Some(CustomOrder {
order: Order::Asc, order: Order::Desc,
target: OrderTarget::Key, target: OrderTarget::Key,
}), }),
..Default::default() ..Default::default()
@@ -952,7 +957,7 @@ mod tests {
bucket_agg: BucketAggregationType::Terms(TermsAggregation { bucket_agg: BucketAggregationType::Terms(TermsAggregation {
field: "string_id".to_string(), field: "string_id".to_string(),
order: Some(CustomOrder { order: Some(CustomOrder {
order: Order::Asc, order: Order::Desc,
target: OrderTarget::Key, target: OrderTarget::Key,
}), }),
size: Some(2), size: Some(2),
@@ -976,14 +981,14 @@ mod tests {
assert_eq!(res["my_texts"]["sum_other_doc_count"], 3); assert_eq!(res["my_texts"]["sum_other_doc_count"], 3);
// key asc and segment_size cut_off // key desc and segment_size cut_off
let agg_req: Aggregations = vec![( let agg_req: Aggregations = vec![(
"my_texts".to_string(), "my_texts".to_string(),
Aggregation::Bucket(BucketAggregation { Aggregation::Bucket(BucketAggregation {
bucket_agg: BucketAggregationType::Terms(TermsAggregation { bucket_agg: BucketAggregationType::Terms(TermsAggregation {
field: "string_id".to_string(), field: "string_id".to_string(),
order: Some(CustomOrder { order: Some(CustomOrder {
order: Order::Asc, order: Order::Desc,
target: OrderTarget::Key, target: OrderTarget::Key,
}), }),
size: Some(2), size: Some(2),
@@ -1006,14 +1011,14 @@ mod tests {
serde_json::Value::Null serde_json::Value::Null
); );
// key desc // key asc
let agg_req: Aggregations = vec![( let agg_req: Aggregations = vec![(
"my_texts".to_string(), "my_texts".to_string(),
Aggregation::Bucket(BucketAggregation { Aggregation::Bucket(BucketAggregation {
bucket_agg: BucketAggregationType::Terms(TermsAggregation { bucket_agg: BucketAggregationType::Terms(TermsAggregation {
field: "string_id".to_string(), field: "string_id".to_string(),
order: Some(CustomOrder { order: Some(CustomOrder {
order: Order::Desc, order: Order::Asc,
target: OrderTarget::Key, target: OrderTarget::Key,
}), }),
..Default::default() ..Default::default()
@@ -1033,14 +1038,14 @@ mod tests {
assert_eq!(res["my_texts"]["buckets"][2]["doc_count"], 5); assert_eq!(res["my_texts"]["buckets"][2]["doc_count"], 5);
assert_eq!(res["my_texts"]["sum_other_doc_count"], 0); assert_eq!(res["my_texts"]["sum_other_doc_count"], 0);
// key desc, size cut_off // key asc, size cut_off
let agg_req: Aggregations = vec![( let agg_req: Aggregations = vec![(
"my_texts".to_string(), "my_texts".to_string(),
Aggregation::Bucket(BucketAggregation { Aggregation::Bucket(BucketAggregation {
bucket_agg: BucketAggregationType::Terms(TermsAggregation { bucket_agg: BucketAggregationType::Terms(TermsAggregation {
field: "string_id".to_string(), field: "string_id".to_string(),
order: Some(CustomOrder { order: Some(CustomOrder {
order: Order::Desc, order: Order::Asc,
target: OrderTarget::Key, target: OrderTarget::Key,
}), }),
size: Some(2), size: Some(2),
@@ -1063,14 +1068,14 @@ mod tests {
); );
assert_eq!(res["my_texts"]["sum_other_doc_count"], 5); assert_eq!(res["my_texts"]["sum_other_doc_count"], 5);
// key desc, segment_size cut_off // key asc, segment_size cut_off
let agg_req: Aggregations = vec![( let agg_req: Aggregations = vec![(
"my_texts".to_string(), "my_texts".to_string(),
Aggregation::Bucket(BucketAggregation { Aggregation::Bucket(BucketAggregation {
bucket_agg: BucketAggregationType::Terms(TermsAggregation { bucket_agg: BucketAggregationType::Terms(TermsAggregation {
field: "string_id".to_string(), field: "string_id".to_string(),
order: Some(CustomOrder { order: Some(CustomOrder {
order: Order::Desc, order: Order::Asc,
target: OrderTarget::Key, target: OrderTarget::Key,
}), }),
size: Some(2), size: Some(2),
@@ -1347,3 +1352,68 @@ mod tests {
Ok(()) Ok(())
} }
} }
#[cfg(all(test, feature = "unstable"))]
mod bench {
use itertools::Itertools;
use rand::seq::SliceRandom;
use rand::thread_rng;
use super::*;
fn get_collector_with_buckets(num_docs: u64) -> TermBuckets {
TermBuckets::from_req_and_validate(&Default::default(), num_docs as usize).unwrap()
}
fn get_rand_terms(total_terms: u64, num_terms_returned: u64) -> Vec<u64> {
let mut rng = thread_rng();
let all_terms = (0..total_terms - 1).collect_vec();
let mut vals = vec![];
for _ in 0..num_terms_returned {
let val = all_terms.as_slice().choose(&mut rng).unwrap();
vals.push(*val);
}
vals
}
fn bench_term_buckets(b: &mut test::Bencher, num_terms: u64, total_terms: u64) {
let mut collector = get_collector_with_buckets(total_terms);
let vals = get_rand_terms(total_terms, num_terms);
let aggregations_with_accessor: AggregationsWithAccessor = Default::default();
let bucket_count: BucketCount = BucketCount {
bucket_count: Default::default(),
max_bucket_count: 1_000_001u32,
};
b.iter(|| {
for &val in &vals {
collector
.increment_bucket(&[val], 0, &aggregations_with_accessor, &bucket_count, &None)
.unwrap();
}
})
}
#[bench]
fn bench_term_buckets_500_of_1_000_000(b: &mut test::Bencher) {
bench_term_buckets(b, 500u64, 1_000_000u64)
}
#[bench]
fn bench_term_buckets_1_000_000_of_50_000(b: &mut test::Bencher) {
bench_term_buckets(b, 1_000_000u64, 50_000u64)
}
#[bench]
fn bench_term_buckets_1_000_000_of_50(b: &mut test::Bencher) {
bench_term_buckets(b, 1_000_000u64, 50u64)
}
#[bench]
fn bench_term_buckets_1_000_000_of_1_000_000(b: &mut test::Bencher) {
bench_term_buckets(b, 1_000_000u64, 1_000_000u64)
}
}

View File

@@ -17,15 +17,13 @@ use super::bucket::{
cut_off_buckets, get_agg_name_and_property, intermediate_histogram_buckets_to_final_buckets, cut_off_buckets, get_agg_name_and_property, intermediate_histogram_buckets_to_final_buckets,
GetDocCount, Order, OrderTarget, SegmentHistogramBucketEntry, TermsAggregation, GetDocCount, Order, OrderTarget, SegmentHistogramBucketEntry, TermsAggregation,
}; };
use super::metric::{ use super::metric::{IntermediateAverage, IntermediateStats};
IntermediateAverage, IntermediateCount, IntermediateMax, IntermediateMin, IntermediateStats,
IntermediateSum,
};
use super::segment_agg_result::SegmentMetricResultCollector; use super::segment_agg_result::SegmentMetricResultCollector;
use super::{format_date, Key, SerializedKey, VecWithNames}; use super::{format_date, Key, SerializedKey, VecWithNames};
use crate::aggregation::agg_result::{AggregationResults, BucketEntries, BucketEntry}; use crate::aggregation::agg_result::{AggregationResults, BucketEntries, BucketEntry};
use crate::aggregation::bucket::TermsAggregationInternal; use crate::aggregation::bucket::TermsAggregationInternal;
use crate::schema::Schema; use crate::schema::Schema;
use crate::TantivyError;
/// Contains the intermediate aggregation result, which is optimized to be merged with other /// Contains the intermediate aggregation result, which is optimized to be merged with other
/// intermediate results. /// intermediate results.
@@ -206,43 +204,21 @@ pub enum IntermediateAggregationResult {
/// Holds the intermediate data for metric results /// Holds the intermediate data for metric results
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum IntermediateMetricResult { pub enum IntermediateMetricResult {
/// Intermediate average result. /// Average containing intermediate average data result
Average(IntermediateAverage), Average(IntermediateAverage),
/// Intermediate count result. /// AverageData variant
Count(IntermediateCount),
/// Intermediate max result.
Max(IntermediateMax),
/// Intermediate min result.
Min(IntermediateMin),
/// Intermediate stats result.
Stats(IntermediateStats), Stats(IntermediateStats),
/// Intermediate sum result.
Sum(IntermediateSum),
} }
impl From<SegmentMetricResultCollector> for IntermediateMetricResult { impl From<SegmentMetricResultCollector> for IntermediateMetricResult {
fn from(tree: SegmentMetricResultCollector) -> Self { fn from(tree: SegmentMetricResultCollector) -> Self {
match tree { match tree {
SegmentMetricResultCollector::Stats(collector) => match collector.collecting_for { SegmentMetricResultCollector::Average(collector) => {
super::metric::SegmentStatsType::Average => IntermediateMetricResult::Average( IntermediateMetricResult::Average(IntermediateAverage::from_collector(collector))
IntermediateAverage::from_collector(collector), }
), SegmentMetricResultCollector::Stats(collector) => {
super::metric::SegmentStatsType::Count => { IntermediateMetricResult::Stats(collector.stats)
IntermediateMetricResult::Count(IntermediateCount::from_collector(collector)) }
}
super::metric::SegmentStatsType::Max => {
IntermediateMetricResult::Max(IntermediateMax::from_collector(collector))
}
super::metric::SegmentStatsType::Min => {
IntermediateMetricResult::Min(IntermediateMin::from_collector(collector))
}
super::metric::SegmentStatsType::Stats => {
IntermediateMetricResult::Stats(collector.stats)
}
super::metric::SegmentStatsType::Sum => {
IntermediateMetricResult::Sum(IntermediateSum::from_collector(collector))
}
},
} }
} }
} }
@@ -253,36 +229,18 @@ impl IntermediateMetricResult {
MetricAggregation::Average(_) => { MetricAggregation::Average(_) => {
IntermediateMetricResult::Average(IntermediateAverage::default()) IntermediateMetricResult::Average(IntermediateAverage::default())
} }
MetricAggregation::Count(_) => {
IntermediateMetricResult::Count(IntermediateCount::default())
}
MetricAggregation::Max(_) => IntermediateMetricResult::Max(IntermediateMax::default()),
MetricAggregation::Min(_) => IntermediateMetricResult::Min(IntermediateMin::default()),
MetricAggregation::Stats(_) => { MetricAggregation::Stats(_) => {
IntermediateMetricResult::Stats(IntermediateStats::default()) IntermediateMetricResult::Stats(IntermediateStats::default())
} }
MetricAggregation::Sum(_) => IntermediateMetricResult::Sum(IntermediateSum::default()),
} }
} }
fn merge_fruits(&mut self, other: IntermediateMetricResult) { fn merge_fruits(&mut self, other: IntermediateMetricResult) {
match (self, other) { match (self, other) {
( (
IntermediateMetricResult::Average(avg_left), IntermediateMetricResult::Average(avg_data_left),
IntermediateMetricResult::Average(avg_right), IntermediateMetricResult::Average(avg_data_right),
) => { ) => {
avg_left.merge_fruits(avg_right); avg_data_left.merge_fruits(avg_data_right);
}
(
IntermediateMetricResult::Count(count_left),
IntermediateMetricResult::Count(count_right),
) => {
count_left.merge_fruits(count_right);
}
(IntermediateMetricResult::Max(max_left), IntermediateMetricResult::Max(max_right)) => {
max_left.merge_fruits(max_right);
}
(IntermediateMetricResult::Min(min_left), IntermediateMetricResult::Min(min_right)) => {
min_left.merge_fruits(min_right);
} }
( (
IntermediateMetricResult::Stats(stats_left), IntermediateMetricResult::Stats(stats_left),
@@ -290,9 +248,6 @@ impl IntermediateMetricResult {
) => { ) => {
stats_left.merge_fruits(stats_right); stats_left.merge_fruits(stats_right);
} }
(IntermediateMetricResult::Sum(sum_left), IntermediateMetricResult::Sum(sum_right)) => {
sum_left.merge_fruits(sum_right);
}
_ => { _ => {
panic!("incompatible fruit types in tree"); panic!("incompatible fruit types in tree");
} }
@@ -499,7 +454,7 @@ impl IntermediateTermBucketResult {
match req.order.target { match req.order.target {
OrderTarget::Key => { OrderTarget::Key => {
buckets.sort_by(|left, right| { buckets.sort_by(|left, right| {
if req.order.order == Order::Asc { if req.order.order == Order::Desc {
left.key.partial_cmp(&right.key) left.key.partial_cmp(&right.key)
} else { } else {
right.key.partial_cmp(&left.key) right.key.partial_cmp(&left.key)
@@ -657,7 +612,9 @@ impl IntermediateRangeBucketEntry {
// If we have a date type on the histogram buckets, we add the `key_as_string` field as // If we have a date type on the histogram buckets, we add the `key_as_string` field as
// rfc339 // rfc339
let field = schema.get_field(&range_req.field)?; let field = schema
.get_field(&range_req.field)
.ok_or_else(|| TantivyError::FieldNotFound(range_req.field.to_string()))?;
if schema.get_field_entry(field).field_type().is_date() { if schema.get_field_entry(field).field_type().is_date() {
if let Some(val) = range_bucket_entry.to { if let Some(val) = range_bucket_entry.to {
let key_as_string = format_date(val as i64)?; let key_as_string = format_date(val as i64)?;

View File

@@ -1,58 +1,114 @@
use std::fmt::Debug; use std::fmt::Debug;
use fastfield_codecs::Column;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use super::{IntermediateStats, SegmentStatsCollector}; use crate::aggregation::f64_from_fastfield_u64;
use crate::schema::Type;
use crate::DocId;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
/// A single-value metric aggregation that computes the average of numeric values that are /// A single-value metric aggregation that computes the average of numeric values that are
/// extracted from the aggregated documents. /// extracted from the aggregated documents.
/// Supported field types are u64, i64, and f64.
/// See [super::SingleMetricResult] for return value. /// See [super::SingleMetricResult] for return value.
/// ///
/// # JSON Format /// # JSON Format
/// ```json /// ```json
/// { /// {
/// "avg": { /// "avg": {
/// "field": "score" /// "field": "score",
/// } /// }
/// } /// }
/// ``` /// ```
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AverageAggregation { pub struct AverageAggregation {
/// The field name to compute the average on. /// The field name to compute the stats on.
pub field: String, pub field: String,
} }
impl AverageAggregation { impl AverageAggregation {
/// Creates a new [`AverageAggregation`] instance from a field name. /// Create new AverageAggregation from a field.
pub fn from_field_name(field_name: String) -> Self { pub fn from_field_name(field_name: String) -> Self {
Self { field: field_name } AverageAggregation { field: field_name }
} }
/// Returns the field name the aggregation is computed on. /// Return the field name.
pub fn field_name(&self) -> &str { pub fn field_name(&self) -> &str {
&self.field &self.field
} }
} }
/// Intermediate result of the average aggregation that can be combined with other intermediate #[derive(Clone, PartialEq)]
/// results. pub(crate) struct SegmentAverageCollector {
pub data: IntermediateAverage,
field_type: Type,
}
impl Debug for SegmentAverageCollector {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AverageCollector")
.field("data", &self.data)
.finish()
}
}
impl SegmentAverageCollector {
pub fn from_req(field_type: Type) -> Self {
Self {
field_type,
data: Default::default(),
}
}
pub(crate) fn collect_block(&mut self, doc: &[DocId], field: &dyn Column<u64>) {
let mut iter = doc.chunks_exact(4);
for docs in iter.by_ref() {
let val1 = field.get_val(docs[0]);
let val2 = field.get_val(docs[1]);
let val3 = field.get_val(docs[2]);
let val4 = field.get_val(docs[3]);
let val1 = f64_from_fastfield_u64(val1, &self.field_type);
let val2 = f64_from_fastfield_u64(val2, &self.field_type);
let val3 = f64_from_fastfield_u64(val3, &self.field_type);
let val4 = f64_from_fastfield_u64(val4, &self.field_type);
self.data.collect(val1);
self.data.collect(val2);
self.data.collect(val3);
self.data.collect(val4);
}
for &doc in iter.remainder() {
let val = field.get_val(doc);
let val = f64_from_fastfield_u64(val, &self.field_type);
self.data.collect(val);
}
}
}
/// Contains mergeable version of average data.
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateAverage { pub struct IntermediateAverage {
stats: IntermediateStats, pub(crate) sum: f64,
pub(crate) doc_count: u64,
} }
impl IntermediateAverage { impl IntermediateAverage {
/// Creates a new [`IntermediateAverage`] instance from a [`SegmentStatsCollector`]. pub(crate) fn from_collector(collector: SegmentAverageCollector) -> Self {
pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self { collector.data
Self { }
stats: collector.stats,
/// Merge average data into this instance.
pub fn merge_fruits(&mut self, other: IntermediateAverage) {
self.sum += other.sum;
self.doc_count += other.doc_count;
}
/// compute final result
pub fn finalize(&self) -> Option<f64> {
if self.doc_count == 0 {
None
} else {
Some(self.sum / self.doc_count as f64)
} }
} }
/// Merges the other intermediate result into self. #[inline]
pub fn merge_fruits(&mut self, other: IntermediateAverage) { fn collect(&mut self, val: f64) {
self.stats.merge_fruits(other.stats); self.doc_count += 1;
} self.sum += val;
/// Computes the final average value.
pub fn finalize(&self) -> Option<f64> {
self.stats.finalize().avg
} }
} }

View File

@@ -1,58 +0,0 @@
use std::fmt::Debug;
use serde::{Deserialize, Serialize};
use super::{IntermediateStats, SegmentStatsCollector};
/// A single-value metric aggregation that counts the number of values that are
/// extracted from the aggregated documents.
/// See [super::SingleMetricResult] for return value.
///
/// # JSON Format
/// ```json
/// {
/// "value_count": {
/// "field": "score"
/// }
/// }
/// ```
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CountAggregation {
/// The field name to compute the minimum on.
pub field: String,
}
impl CountAggregation {
/// Creates a new [`CountAggregation`] instance from a field name.
pub fn from_field_name(field_name: String) -> Self {
Self { field: field_name }
}
/// Returns the field name the aggregation is computed on.
pub fn field_name(&self) -> &str {
&self.field
}
}
/// Intermediate result of the count aggregation that can be combined with other intermediate
/// results.
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateCount {
stats: IntermediateStats,
}
impl IntermediateCount {
/// Creates a new [`IntermediateCount`] instance from a [`SegmentStatsCollector`].
pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self {
Self {
stats: collector.stats,
}
}
/// Merges the other intermediate result into self.
pub fn merge_fruits(&mut self, other: IntermediateCount) {
self.stats.merge_fruits(other.stats);
}
/// Computes the final minimum value.
pub fn finalize(&self) -> Option<f64> {
Some(self.stats.finalize().count as f64)
}
}

View File

@@ -1,58 +0,0 @@
use std::fmt::Debug;
use serde::{Deserialize, Serialize};
use super::{IntermediateStats, SegmentStatsCollector};
/// A single-value metric aggregation that computes the maximum of numeric values that are
/// extracted from the aggregated documents.
/// See [super::SingleMetricResult] for return value.
///
/// # JSON Format
/// ```json
/// {
/// "max": {
/// "field": "score"
/// }
/// }
/// ```
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct MaxAggregation {
/// The field name to compute the maximum on.
pub field: String,
}
impl MaxAggregation {
/// Creates a new [`MaxAggregation`] instance from a field name.
pub fn from_field_name(field_name: String) -> Self {
Self { field: field_name }
}
/// Returns the field name the aggregation is computed on.
pub fn field_name(&self) -> &str {
&self.field
}
}
/// Intermediate result of the maximum aggregation that can be combined with other intermediate
/// results.
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateMax {
stats: IntermediateStats,
}
impl IntermediateMax {
/// Creates a new [`IntermediateMax`] instance from a [`SegmentStatsCollector`].
pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self {
Self {
stats: collector.stats,
}
}
/// Merges the other intermediate result into self.
pub fn merge_fruits(&mut self, other: IntermediateMax) {
self.stats.merge_fruits(other.stats);
}
/// Computes the final maximum value.
pub fn finalize(&self) -> Option<f64> {
self.stats.finalize().max
}
}

View File

@@ -1,58 +0,0 @@
use std::fmt::Debug;
use serde::{Deserialize, Serialize};
use super::{IntermediateStats, SegmentStatsCollector};
/// A single-value metric aggregation that computes the minimum of numeric values that are
/// extracted from the aggregated documents.
/// See [super::SingleMetricResult] for return value.
///
/// # JSON Format
/// ```json
/// {
/// "min": {
/// "field": "score"
/// }
/// }
/// ```
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct MinAggregation {
/// The field name to compute the minimum on.
pub field: String,
}
impl MinAggregation {
/// Creates a new [`MinAggregation`] instance from a field name.
pub fn from_field_name(field_name: String) -> Self {
Self { field: field_name }
}
/// Returns the field name the aggregation is computed on.
pub fn field_name(&self) -> &str {
&self.field
}
}
/// Intermediate result of the minimum aggregation that can be combined with other intermediate
/// results.
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateMin {
stats: IntermediateStats,
}
impl IntermediateMin {
/// Creates a new [`IntermediateMin`] instance from a [`SegmentStatsCollector`].
pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self {
Self {
stats: collector.stats,
}
}
/// Merges the other intermediate result into self.
pub fn merge_fruits(&mut self, other: IntermediateMin) {
self.stats.merge_fruits(other.stats);
}
/// Computes the final minimum value.
pub fn finalize(&self) -> Option<f64> {
self.stats.finalize().min
}
}

View File

@@ -3,18 +3,10 @@
//! The aggregations in this family compute metrics, see [super::agg_req::MetricAggregation] for //! The aggregations in this family compute metrics, see [super::agg_req::MetricAggregation] for
//! details. //! details.
mod average; mod average;
mod count;
mod max;
mod min;
mod stats; mod stats;
mod sum;
pub use average::*; pub use average::*;
pub use count::*;
pub use max::*;
pub use min::*;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
pub use stats::*; pub use stats::*;
pub use sum::*;
/// Single-metric aggregations use this common result structure. /// Single-metric aggregations use this common result structure.
/// ///
@@ -36,61 +28,3 @@ impl From<Option<f64>> for SingleMetricResult {
Self { value } Self { value }
} }
} }
#[cfg(test)]
mod tests {
use crate::aggregation::agg_req::Aggregations;
use crate::aggregation::agg_result::AggregationResults;
use crate::aggregation::AggregationCollector;
use crate::query::AllQuery;
use crate::schema::{Cardinality, NumericOptions, Schema};
use crate::Index;
#[test]
fn test_metric_aggregations() {
let mut schema_builder = Schema::builder();
let field_options = NumericOptions::default().set_fast(Cardinality::SingleValue);
let field = schema_builder.add_f64_field("price", field_options);
let index = Index::create_in_ram(schema_builder.build());
let mut index_writer = index.writer_for_tests().unwrap();
for i in 0..3 {
index_writer
.add_document(doc!(
field => i as f64,
))
.unwrap();
}
index_writer.commit().unwrap();
for i in 3..6 {
index_writer
.add_document(doc!(
field => i as f64,
))
.unwrap();
}
index_writer.commit().unwrap();
let aggregations_json = r#"{
"price_avg": { "avg": { "field": "price" } },
"price_count": { "value_count": { "field": "price" } },
"price_max": { "max": { "field": "price" } },
"price_min": { "min": { "field": "price" } },
"price_stats": { "stats": { "field": "price" } },
"price_sum": { "sum": { "field": "price" } }
}"#;
let aggregations: Aggregations = serde_json::from_str(aggregations_json).unwrap();
let collector = AggregationCollector::from_aggs(aggregations, None, index.schema());
let reader = index.reader().unwrap();
let searcher = reader.searcher();
let aggregations_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
let aggregations_res_json = serde_json::to_value(aggregations_res).unwrap();
assert_eq!(aggregations_res_json["price_avg"]["value"], 2.5);
assert_eq!(aggregations_res_json["price_count"]["value"], 6.0);
assert_eq!(aggregations_res_json["price_max"]["value"], 5.0);
assert_eq!(aggregations_res_json["price_min"]["value"], 0.0);
assert_eq!(aggregations_res_json["price_sum"]["value"], 15.0);
}
}

View File

@@ -5,15 +5,16 @@ use crate::aggregation::f64_from_fastfield_u64;
use crate::schema::Type; use crate::schema::Type;
use crate::{DocId, TantivyError}; use crate::{DocId, TantivyError};
/// A multi-value metric aggregation that computes a collection of statistics on numeric values that /// A multi-value metric aggregation that computes stats of numeric values that are
/// are extracted from the aggregated documents. /// extracted from the aggregated documents.
/// Supported field types are `u64`, `i64`, and `f64`.
/// See [`Stats`] for returned statistics. /// See [`Stats`] for returned statistics.
/// ///
/// # JSON Format /// # JSON Format
/// ```json /// ```json
/// { /// {
/// "stats": { /// "stats": {
/// "field": "score" /// "field": "score",
/// } /// }
/// } /// }
/// ``` /// ```
@@ -25,11 +26,11 @@ pub struct StatsAggregation {
} }
impl StatsAggregation { impl StatsAggregation {
/// Creates a new [`StatsAggregation`] instance from a field name. /// Create new StatsAggregation from a field.
pub fn from_field_name(field_name: String) -> Self { pub fn from_field_name(field_name: String) -> Self {
StatsAggregation { field: field_name } StatsAggregation { field: field_name }
} }
/// Returns the field name the aggregation is computed on. /// Return the field name.
pub fn field_name(&self) -> &str { pub fn field_name(&self) -> &str {
&self.field &self.field
} }
@@ -39,14 +40,16 @@ impl StatsAggregation {
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Stats { pub struct Stats {
/// The number of documents. /// The number of documents.
pub count: u64, pub count: usize,
/// The sum of the fast field values. /// The sum of the fast field values.
pub sum: f64, pub sum: f64,
/// The standard deviation of the fast field values. `None` for count == 0.
pub standard_deviation: Option<f64>,
/// The min value of the fast field values. /// The min value of the fast field values.
pub min: Option<f64>, pub min: Option<f64>,
/// The max value of the fast field values. /// The max value of the fast field values.
pub max: Option<f64>, pub max: Option<f64>,
/// The average of the fast field values. `None` if count equals zero. /// The average of the values. `None` for count == 0.
pub avg: Option<f64>, pub avg: Option<f64>,
} }
@@ -55,36 +58,33 @@ impl Stats {
match agg_property { match agg_property {
"count" => Ok(Some(self.count as f64)), "count" => Ok(Some(self.count as f64)),
"sum" => Ok(Some(self.sum)), "sum" => Ok(Some(self.sum)),
"standard_deviation" => Ok(self.standard_deviation),
"min" => Ok(self.min), "min" => Ok(self.min),
"max" => Ok(self.max), "max" => Ok(self.max),
"avg" => Ok(self.avg), "avg" => Ok(self.avg),
_ => Err(TantivyError::InvalidArgument(format!( _ => Err(TantivyError::InvalidArgument(format!(
"Unknown property {} on stats metric aggregation", "unknown property {} on stats metric aggregation",
agg_property agg_property
))), ))),
} }
} }
} }
/// Intermediate result of the stats aggregation that can be combined with other intermediate /// `IntermediateStats` contains the mergeable version for stats.
/// results.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateStats { pub struct IntermediateStats {
/// The number of extracted values. count: usize,
count: u64,
/// The sum of the extracted values.
sum: f64, sum: f64,
/// The min value. squared_sum: f64,
min: f64, min: f64,
/// The max value.
max: f64, max: f64,
} }
impl Default for IntermediateStats { impl Default for IntermediateStats {
fn default() -> Self { fn default() -> Self {
Self { Self {
count: 0, count: 0,
sum: 0.0, sum: 0.0,
squared_sum: 0.0,
min: f64::MAX, min: f64::MAX,
max: f64::MIN, max: f64::MIN,
} }
@@ -92,15 +92,33 @@ impl Default for IntermediateStats {
} }
impl IntermediateStats { impl IntermediateStats {
/// Merges the other stats intermediate result into self. pub(crate) fn avg(&self) -> Option<f64> {
if self.count == 0 {
None
} else {
Some(self.sum / (self.count as f64))
}
}
fn square_mean(&self) -> f64 {
self.squared_sum / (self.count as f64)
}
pub(crate) fn standard_deviation(&self) -> Option<f64> {
self.avg()
.map(|average| (self.square_mean() - average * average).sqrt())
}
/// Merge data from other stats into this instance.
pub fn merge_fruits(&mut self, other: IntermediateStats) { pub fn merge_fruits(&mut self, other: IntermediateStats) {
self.count += other.count; self.count += other.count;
self.sum += other.sum; self.sum += other.sum;
self.squared_sum += other.squared_sum;
self.min = self.min.min(other.min); self.min = self.min.min(other.min);
self.max = self.max.max(other.max); self.max = self.max.max(other.max);
} }
/// Computes the final stats value. /// compute final resultimprove_docs
pub fn finalize(&self) -> Stats { pub fn finalize(&self) -> Stats {
let min = if self.count == 0 { let min = if self.count == 0 {
None None
@@ -112,17 +130,13 @@ impl IntermediateStats {
} else { } else {
Some(self.max) Some(self.max)
}; };
let avg = if self.count == 0 {
None
} else {
Some(self.sum / (self.count as f64))
};
Stats { Stats {
count: self.count, count: self.count,
sum: self.sum, sum: self.sum,
standard_deviation: self.standard_deviation(),
min, min,
max, max,
avg, avg: self.avg(),
} }
} }
@@ -130,33 +144,22 @@ impl IntermediateStats {
fn collect(&mut self, value: f64) { fn collect(&mut self, value: f64) {
self.count += 1; self.count += 1;
self.sum += value; self.sum += value;
self.squared_sum += value * value;
self.min = self.min.min(value); self.min = self.min.min(value);
self.max = self.max.max(value); self.max = self.max.max(value);
} }
} }
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum SegmentStatsType {
Average,
Count,
Max,
Min,
Stats,
Sum,
}
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
pub(crate) struct SegmentStatsCollector { pub(crate) struct SegmentStatsCollector {
field_type: Type,
pub(crate) collecting_for: SegmentStatsType,
pub(crate) stats: IntermediateStats, pub(crate) stats: IntermediateStats,
field_type: Type,
} }
impl SegmentStatsCollector { impl SegmentStatsCollector {
pub fn from_req(field_type: Type, collecting_for: SegmentStatsType) -> Self { pub fn from_req(field_type: Type) -> Self {
Self { Self {
field_type, field_type,
collecting_for,
stats: IntermediateStats::default(), stats: IntermediateStats::default(),
} }
} }
@@ -233,6 +236,7 @@ mod tests {
"count": 0, "count": 0,
"max": Value::Null, "max": Value::Null,
"min": Value::Null, "min": Value::Null,
"standard_deviation": Value::Null,
"sum": 0.0 "sum": 0.0
}) })
); );
@@ -309,6 +313,7 @@ mod tests {
"count": 7, "count": 7,
"max": 44.0, "max": 44.0,
"min": 1.0, "min": 1.0,
"standard_deviation": 13.65313748796613,
"sum": 85.0 "sum": 85.0
}) })
); );
@@ -320,6 +325,7 @@ mod tests {
"count": 7, "count": 7,
"max": 44.0, "max": 44.0,
"min": 1.0, "min": 1.0,
"standard_deviation": 13.65313748796613,
"sum": 85.0 "sum": 85.0
}) })
); );
@@ -331,6 +337,7 @@ mod tests {
"count": 7, "count": 7,
"max": 44.5, "max": 44.5,
"min": 1.0, "min": 1.0,
"standard_deviation": 13.819905785437443,
"sum": 85.5 "sum": 85.5
}) })
); );
@@ -342,6 +349,7 @@ mod tests {
"count": 3, "count": 3,
"max": 14.0, "max": 14.0,
"min": 7.0, "min": 7.0,
"standard_deviation": 2.867441755680877,
"sum": 32.0 "sum": 32.0
}) })
); );
@@ -353,6 +361,7 @@ mod tests {
"count": 0, "count": 0,
"max": serde_json::Value::Null, "max": serde_json::Value::Null,
"min": serde_json::Value::Null, "min": serde_json::Value::Null,
"standard_deviation": serde_json::Value::Null,
"sum": 0.0, "sum": 0.0,
}) })
); );

View File

@@ -1,58 +0,0 @@
use std::fmt::Debug;
use serde::{Deserialize, Serialize};
use super::{IntermediateStats, SegmentStatsCollector};
/// A single-value metric aggregation that sums up numeric values that are
/// extracted from the aggregated documents.
/// See [super::SingleMetricResult] for return value.
///
/// # JSON Format
/// ```json
/// {
/// "sum": {
/// "field": "score"
/// }
/// }
/// ```
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SumAggregation {
/// The field name to compute the minimum on.
pub field: String,
}
impl SumAggregation {
/// Creates a new [`SumAggregation`] instance from a field name.
pub fn from_field_name(field_name: String) -> Self {
Self { field: field_name }
}
/// Returns the field name the aggregation is computed on.
pub fn field_name(&self) -> &str {
&self.field
}
}
/// Intermediate result of the minimum aggregation that can be combined with other intermediate
/// results.
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateSum {
stats: IntermediateStats,
}
impl IntermediateSum {
/// Creates a new [`IntermediateSum`] instance from a [`SegmentStatsCollector`].
pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self {
Self {
stats: collector.stats,
}
}
/// Merges the other intermediate result into self.
pub fn merge_fruits(&mut self, other: IntermediateSum) {
self.stats.merge_fruits(other.stats);
}
/// Computes the final minimum value.
pub fn finalize(&self) -> Option<f64> {
Some(self.stats.finalize().sum)
}
}

View File

@@ -1,5 +1,6 @@
//! # Aggregations //! # Aggregations
//! //!
//!
//! An aggregation summarizes your data as statistics on buckets or metrics. //! An aggregation summarizes your data as statistics on buckets or metrics.
//! //!
//! Aggregations can provide answer to questions like: //! Aggregations can provide answer to questions like:
@@ -40,10 +41,6 @@
//! - [Metric](metric) //! - [Metric](metric)
//! - [Average](metric::AverageAggregation) //! - [Average](metric::AverageAggregation)
//! - [Stats](metric::StatsAggregation) //! - [Stats](metric::StatsAggregation)
//! - [Min](metric::MinAggregation)
//! - [Max](metric::MaxAggregation)
//! - [Sum](metric::SumAggregation)
//! - [Count](metric::CountAggregation)
//! //!
//! # Example //! # Example
//! Compute the average metric, by building [`agg_req::Aggregations`], which is built from an //! Compute the average metric, by building [`agg_req::Aggregations`], which is built from an
@@ -78,7 +75,7 @@
//! } //! }
//! ``` //! ```
//! # Example JSON //! # Example JSON
//! Requests are compatible with the elasticsearch JSON request format. //! Requests are compatible with the elasticsearch json request format.
//! //!
//! ``` //! ```
//! use tantivy::aggregation::agg_req::Aggregations; //! use tantivy::aggregation::agg_req::Aggregations;
@@ -219,8 +216,8 @@ impl<T: Clone> VecWithNames<T> {
fn from_entries(mut entries: Vec<(String, T)>) -> Self { fn from_entries(mut entries: Vec<(String, T)>) -> Self {
// Sort to ensure order of elements match across multiple instances // Sort to ensure order of elements match across multiple instances
entries.sort_by(|left, right| left.0.cmp(&right.0)); entries.sort_by(|left, right| left.0.cmp(&right.0));
let mut data = Vec::with_capacity(entries.len()); let mut data = vec![];
let mut data_names = Vec::with_capacity(entries.len()); let mut data_names = vec![];
for entry in entries { for entry in entries {
data_names.push(entry.0); data_names.push(entry.0);
data.push(entry.1); data.push(entry.1);
@@ -1156,6 +1153,12 @@ mod tests {
r#"FieldNotFound("not_exist_field")"# r#"FieldNotFound("not_exist_field")"#
); );
let agg_res = avg_on_field("scores_i64");
assert_eq!(
format!("{:?}", agg_res),
r#"InvalidArgument("Invalid field cardinality on field scores_i64 expected SingleValue, but got MultiValues")"#
);
Ok(()) Ok(())
} }
@@ -1205,7 +1208,7 @@ mod tests {
text_field_many_terms => many_terms_data.choose(&mut rng).unwrap().to_string(), text_field_many_terms => many_terms_data.choose(&mut rng).unwrap().to_string(),
text_field_few_terms => few_terms_data.choose(&mut rng).unwrap().to_string(), text_field_few_terms => few_terms_data.choose(&mut rng).unwrap().to_string(),
score_field => val as u64, score_field => val as u64,
score_field_f64 => val, score_field_f64 => val as f64,
score_field_i64 => val as i64, score_field_i64 => val as i64,
))?; ))?;
} }
@@ -1247,7 +1250,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&term_query, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1275,7 +1281,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&term_query, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1303,7 +1312,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&term_query, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1339,7 +1351,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&term_query, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1365,7 +1380,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&AllQuery, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1391,7 +1409,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&AllQuery, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1425,7 +1446,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&AllQuery, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1457,7 +1481,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&AllQuery, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1493,7 +1520,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&AllQuery, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1520,7 +1550,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&AllQuery, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1564,7 +1597,7 @@ mod tests {
], ],
..Default::default() ..Default::default()
}), }),
sub_aggregation: sub_agg_req_1, sub_aggregation: sub_agg_req_1.clone(),
}), }),
), ),
] ]
@@ -1574,7 +1607,10 @@ mod tests {
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema()); let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
let searcher = reader.searcher(); let searcher = reader.searcher();
searcher.search(&term_query, &collector).unwrap() let agg_res: AggregationResults =
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
} }

View File

@@ -15,8 +15,7 @@ use super::bucket::{SegmentHistogramCollector, SegmentRangeCollector, SegmentTer
use super::collector::MAX_BUCKET_COUNT; use super::collector::MAX_BUCKET_COUNT;
use super::intermediate_agg_result::{IntermediateAggregationResults, IntermediateBucketResult}; use super::intermediate_agg_result::{IntermediateAggregationResults, IntermediateBucketResult};
use super::metric::{ use super::metric::{
AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, SegmentStatsCollector, AverageAggregation, SegmentAverageCollector, SegmentStatsCollector, StatsAggregation,
SegmentStatsType, StatsAggregation, SumAggregation,
}; };
use super::VecWithNames; use super::VecWithNames;
use crate::aggregation::agg_req::BucketAggregationType; use crate::aggregation::agg_req::BucketAggregationType;
@@ -164,46 +163,30 @@ impl SegmentAggregationResultsCollector {
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
pub(crate) enum SegmentMetricResultCollector { pub(crate) enum SegmentMetricResultCollector {
Average(SegmentAverageCollector),
Stats(SegmentStatsCollector), Stats(SegmentStatsCollector),
} }
impl SegmentMetricResultCollector { impl SegmentMetricResultCollector {
pub fn from_req_and_validate(req: &MetricAggregationWithAccessor) -> crate::Result<Self> { pub fn from_req_and_validate(req: &MetricAggregationWithAccessor) -> crate::Result<Self> {
match &req.metric { match &req.metric {
MetricAggregation::Average(AverageAggregation { .. }) => { MetricAggregation::Average(AverageAggregation { field: _ }) => {
Ok(SegmentMetricResultCollector::Stats( Ok(SegmentMetricResultCollector::Average(
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Average), SegmentAverageCollector::from_req(req.field_type),
)) ))
} }
MetricAggregation::Count(CountAggregation { .. }) => { MetricAggregation::Stats(StatsAggregation { field: _ }) => {
Ok(SegmentMetricResultCollector::Stats( Ok(SegmentMetricResultCollector::Stats(
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Count), SegmentStatsCollector::from_req(req.field_type),
))
}
MetricAggregation::Max(MaxAggregation { .. }) => {
Ok(SegmentMetricResultCollector::Stats(
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Max),
))
}
MetricAggregation::Min(MinAggregation { .. }) => {
Ok(SegmentMetricResultCollector::Stats(
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Min),
))
}
MetricAggregation::Stats(StatsAggregation { .. }) => {
Ok(SegmentMetricResultCollector::Stats(
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Stats),
))
}
MetricAggregation::Sum(SumAggregation { .. }) => {
Ok(SegmentMetricResultCollector::Stats(
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Sum),
)) ))
} }
} }
} }
pub(crate) fn collect_block(&mut self, doc: &[DocId], metric: &MetricAggregationWithAccessor) { pub(crate) fn collect_block(&mut self, doc: &[DocId], metric: &MetricAggregationWithAccessor) {
match self { match self {
SegmentMetricResultCollector::Average(avg_collector) => {
avg_collector.collect_block(doc, &*metric.accessor);
}
SegmentMetricResultCollector::Stats(stats_collector) => { SegmentMetricResultCollector::Stats(stats_collector) => {
stats_collector.collect_block(doc, &*metric.accessor); stats_collector.collect_block(doc, &*metric.accessor);
} }

View File

@@ -130,7 +130,7 @@ where
let fast_field_reader = segment_reader let fast_field_reader = segment_reader
.fast_fields() .fast_fields()
.typed_fast_field_reader(schema.get_field_name(self.field))?; .typed_fast_field_reader(self.field)?;
let segment_collector = self let segment_collector = self
.collector .collector

View File

@@ -5,7 +5,7 @@ use fastfield_codecs::Column;
use crate::collector::{Collector, SegmentCollector}; use crate::collector::{Collector, SegmentCollector};
use crate::fastfield::FastValue; use crate::fastfield::FastValue;
use crate::schema::Type; use crate::schema::{Field, Type};
use crate::{DocId, Score}; use crate::{DocId, Score};
/// Histogram builds an histogram of the values of a fastfield for the /// Histogram builds an histogram of the values of a fastfield for the
@@ -28,7 +28,7 @@ pub struct HistogramCollector {
min_value: u64, min_value: u64,
num_buckets: usize, num_buckets: usize,
divider: DividerU64, divider: DividerU64,
field: String, field: Field,
} }
impl HistogramCollector { impl HistogramCollector {
@@ -46,7 +46,7 @@ impl HistogramCollector {
/// # Disclaimer /// # Disclaimer
/// This function panics if the field given is of type f64. /// This function panics if the field given is of type f64.
pub fn new<TFastValue: FastValue>( pub fn new<TFastValue: FastValue>(
field: String, field: Field,
min_value: TFastValue, min_value: TFastValue,
bucket_width: u64, bucket_width: u64,
num_buckets: usize, num_buckets: usize,
@@ -112,7 +112,7 @@ impl Collector for HistogramCollector {
_segment_local_id: crate::SegmentOrdinal, _segment_local_id: crate::SegmentOrdinal,
segment: &crate::SegmentReader, segment: &crate::SegmentReader,
) -> crate::Result<Self::Child> { ) -> crate::Result<Self::Child> {
let ff_reader = segment.fast_fields().u64_lenient(&self.field)?; let ff_reader = segment.fast_fields().u64_lenient(self.field)?;
Ok(SegmentHistogramCollector { Ok(SegmentHistogramCollector {
histogram_computer: HistogramComputer { histogram_computer: HistogramComputer {
counts: vec![0; self.num_buckets], counts: vec![0; self.num_buckets],
@@ -211,13 +211,13 @@ mod tests {
#[test] #[test]
fn test_no_segments() -> crate::Result<()> { fn test_no_segments() -> crate::Result<()> {
let mut schema_builder = Schema::builder(); let mut schema_builder = Schema::builder();
schema_builder.add_u64_field("val_field", FAST); let val_field = schema_builder.add_u64_field("val_field", FAST);
let schema = schema_builder.build(); let schema = schema_builder.build();
let index = Index::create_in_ram(schema); let index = Index::create_in_ram(schema);
let reader = index.reader()?; let reader = index.reader()?;
let searcher = reader.searcher(); let searcher = reader.searcher();
let all_query = AllQuery; let all_query = AllQuery;
let histogram_collector = HistogramCollector::new("val_field".to_string(), 0u64, 2, 5); let histogram_collector = HistogramCollector::new(val_field, 0u64, 2, 5);
let histogram = searcher.search(&all_query, &histogram_collector)?; let histogram = searcher.search(&all_query, &histogram_collector)?;
assert_eq!(histogram, vec![0; 5]); assert_eq!(histogram, vec![0; 5]);
Ok(()) Ok(())
@@ -238,8 +238,7 @@ mod tests {
let reader = index.reader()?; let reader = index.reader()?;
let searcher = reader.searcher(); let searcher = reader.searcher();
let all_query = AllQuery; let all_query = AllQuery;
let histogram_collector = let histogram_collector = HistogramCollector::new(val_field, -20i64, 10u64, 4);
HistogramCollector::new("val_field".to_string(), -20i64, 10u64, 4);
let histogram = searcher.search(&all_query, &histogram_collector)?; let histogram = searcher.search(&all_query, &histogram_collector)?;
assert_eq!(histogram, vec![1, 1, 0, 1]); assert_eq!(histogram, vec![1, 1, 0, 1]);
Ok(()) Ok(())
@@ -263,8 +262,7 @@ mod tests {
let reader = index.reader()?; let reader = index.reader()?;
let searcher = reader.searcher(); let searcher = reader.searcher();
let all_query = AllQuery; let all_query = AllQuery;
let histogram_collector = let histogram_collector = HistogramCollector::new(val_field, -20i64, 10u64, 4);
HistogramCollector::new("val_field".to_string(), -20i64, 10u64, 4);
let histogram = searcher.search(&all_query, &histogram_collector)?; let histogram = searcher.search(&all_query, &histogram_collector)?;
assert_eq!(histogram, vec![1, 1, 0, 1]); assert_eq!(histogram, vec![1, 1, 0, 1]);
Ok(()) Ok(())
@@ -287,7 +285,7 @@ mod tests {
let searcher = reader.searcher(); let searcher = reader.searcher();
let all_query = AllQuery; let all_query = AllQuery;
let week_histogram_collector = HistogramCollector::new( let week_histogram_collector = HistogramCollector::new(
"date_field".to_string(), date_field,
DateTime::from_primitive( DateTime::from_primitive(
Date::from_calendar_date(1980, Month::January, 1)?.with_hms(0, 0, 0)?, Date::from_calendar_date(1980, Month::January, 1)?.with_hms(0, 0, 0)?,
), ),

View File

@@ -155,7 +155,7 @@ impl SegmentCollector for TestSegmentCollector {
/// ///
/// This collector is mainly useful for tests. /// This collector is mainly useful for tests.
pub struct FastFieldTestCollector { pub struct FastFieldTestCollector {
field: String, field: Field,
} }
pub struct FastFieldSegmentCollector { pub struct FastFieldSegmentCollector {
@@ -164,7 +164,7 @@ pub struct FastFieldSegmentCollector {
} }
impl FastFieldTestCollector { impl FastFieldTestCollector {
pub fn for_field(field: String) -> FastFieldTestCollector { pub fn for_field(field: Field) -> FastFieldTestCollector {
FastFieldTestCollector { field } FastFieldTestCollector { field }
} }
} }
@@ -180,7 +180,7 @@ impl Collector for FastFieldTestCollector {
) -> crate::Result<FastFieldSegmentCollector> { ) -> crate::Result<FastFieldSegmentCollector> {
let reader = segment_reader let reader = segment_reader
.fast_fields() .fast_fields()
.u64(&self.field) .u64(self.field)
.expect("Requested field is not a fast field."); .expect("Requested field is not a fast field.");
Ok(FastFieldSegmentCollector { Ok(FastFieldSegmentCollector {
vals: Vec::new(), vals: Vec::new(),
@@ -238,9 +238,7 @@ impl Collector for BytesFastFieldTestCollector {
_segment_local_id: u32, _segment_local_id: u32,
segment_reader: &SegmentReader, segment_reader: &SegmentReader,
) -> crate::Result<BytesFastFieldSegmentCollector> { ) -> crate::Result<BytesFastFieldSegmentCollector> {
let reader = segment_reader let reader = segment_reader.fast_fields().bytes(self.field)?;
.fast_fields()
.bytes(segment_reader.schema().get_field_name(self.field))?;
Ok(BytesFastFieldSegmentCollector { Ok(BytesFastFieldSegmentCollector {
vals: Vec::new(), vals: Vec::new(),
reader, reader,

View File

@@ -156,7 +156,7 @@ impl CustomScorer<u64> for ScorerByField {
// The conversion will then happen only on the top-K docs. // The conversion will then happen only on the top-K docs.
let ff_reader = segment_reader let ff_reader = segment_reader
.fast_fields() .fast_fields()
.typed_fast_field_reader(segment_reader.schema().get_field_name(self.field))?; .typed_fast_field_reader(self.field)?;
Ok(ScorerByFastFieldReader { ff_reader }) Ok(ScorerByFastFieldReader { ff_reader })
} }
} }
@@ -454,7 +454,7 @@ impl TopDocs {
/// // In our case, we will get a reader for the popularity /// // In our case, we will get a reader for the popularity
/// // fast field. /// // fast field.
/// let popularity_reader = /// let popularity_reader =
/// segment_reader.fast_fields().u64("popularity").unwrap(); /// segment_reader.fast_fields().u64(popularity).unwrap();
/// ///
/// // We can now define our actual scoring function /// // We can now define our actual scoring function
/// move |doc: DocId, original_score: Score| { /// move |doc: DocId, original_score: Score| {
@@ -561,9 +561,9 @@ impl TopDocs {
/// // Note that this is implemented by using a `(u64, u64)` /// // Note that this is implemented by using a `(u64, u64)`
/// // as a score. /// // as a score.
/// let popularity_reader = /// let popularity_reader =
/// segment_reader.fast_fields().u64("popularity").unwrap(); /// segment_reader.fast_fields().u64(popularity).unwrap();
/// let boosted_reader = /// let boosted_reader =
/// segment_reader.fast_fields().u64("boosted").unwrap(); /// segment_reader.fast_fields().u64(boosted).unwrap();
/// ///
/// // We can now define our actual scoring function /// // We can now define our actual scoring function
/// move |doc: DocId| { /// move |doc: DocId| {

View File

@@ -231,7 +231,7 @@ impl IndexBuilder {
fn validate(&self) -> crate::Result<()> { fn validate(&self) -> crate::Result<()> {
if let Some(schema) = self.schema.as_ref() { if let Some(schema) = self.schema.as_ref() {
if let Some(sort_by_field) = self.index_settings.sort_by_field.as_ref() { if let Some(sort_by_field) = self.index_settings.sort_by_field.as_ref() {
let schema_field = schema.get_field(&sort_by_field.field).map_err(|_| { let schema_field = schema.get_field(&sort_by_field.field).ok_or_else(|| {
TantivyError::InvalidArgument(format!( TantivyError::InvalidArgument(format!(
"Field to sort index {} not found in schema", "Field to sort index {} not found in schema",
sort_by_field.field sort_by_field.field

View File

@@ -135,8 +135,6 @@ impl InvertedIndexReader {
term_info: &TermInfo, term_info: &TermInfo,
option: IndexRecordOption, option: IndexRecordOption,
) -> io::Result<SegmentPostings> { ) -> io::Result<SegmentPostings> {
let option = option.downgrade(self.record_option);
let block_postings = self.read_block_postings_from_terminfo(term_info, option)?; let block_postings = self.read_block_postings_from_terminfo(term_info, option)?;
let position_reader = { let position_reader = {
if option.has_positions() { if option.has_positions() {

View File

@@ -198,10 +198,11 @@ impl Searcher {
collector: &C, collector: &C,
executor: &Executor, executor: &Executor,
) -> crate::Result<C::Fruit> { ) -> crate::Result<C::Fruit> {
let enabled_scoring = if collector.requires_scoring() { let scoring_enabled = collector.requires_scoring();
EnableScoring::enabled_from_searcher(self) let enabled_scoring = if scoring_enabled {
EnableScoring::Enabled(self)
} else { } else {
EnableScoring::disabled_from_searcher(self) EnableScoring::Disabled(self.schema())
}; };
let weight = query.weight(enabled_scoring)?; let weight = query.weight(enabled_scoring)?;
let segment_readers = self.segment_readers(); let segment_readers = self.segment_readers();
@@ -249,7 +250,7 @@ impl SearcherInner {
index: Index, index: Index,
segment_readers: Vec<SegmentReader>, segment_readers: Vec<SegmentReader>,
generation: TrackedObject<SearcherGeneration>, generation: TrackedObject<SearcherGeneration>,
doc_store_cache_num_blocks: usize, doc_store_cache_size: usize,
) -> io::Result<SearcherInner> { ) -> io::Result<SearcherInner> {
assert_eq!( assert_eq!(
&segment_readers &segment_readers
@@ -261,7 +262,7 @@ impl SearcherInner {
); );
let store_readers: Vec<StoreReader> = segment_readers let store_readers: Vec<StoreReader> = segment_readers
.iter() .iter()
.map(|segment_reader| segment_reader.get_store_reader(doc_store_cache_num_blocks)) .map(|segment_reader| segment_reader.get_store_reader(doc_store_cache_size))
.collect::<io::Result<Vec<_>>>()?; .collect::<io::Result<Vec<_>>>()?;
Ok(SearcherInner { Ok(SearcherInner {

Some files were not shown because too many files have changed in this diff Show More