Compare commits

..

33 Commits

Author SHA1 Message Date
Paul Masurel
cd0af6d6b1 first commit to show how to impl u128 2023-01-17 13:44:34 +09:00
Pascal Seitz
49baa15f0f start migrate Field to &str
start migrate Field to &str in preparation of columnar
return Result for get_field
2023-01-17 13:34:21 +09:00
Adrien Guillo
c9cb3d04bf Merge pull request #1788 from quickwit-oss/guilload/remove-std-dev-from-stats-agg
Remove standard deviation from stats aggregation
2023-01-16 23:16:36 -05:00
Adrien Guillo
0caaf13a90 Remove standard deviation from stats aggregation 2023-01-16 22:58:23 -05:00
Adrien Guillo
a59bd965cc Merge pull request #1794 from quickwit-oss/guilload/count-min-max-sum-aggs
Add count, min, max, and sum aggregations
2023-01-16 22:45:01 -05:00
Adrien Guillo
f2dad194ea Add count, min, max, and sum aggregations 2023-01-16 12:22:20 -05:00
Paul Masurel
25bad784ad Integrated fastfield codecs into columnar. (#1782)
Introduced asymetric OptionalCodec / SerializableOptionalCodec
Removed cardinality from the columnar sstable.
Added DynamicColumn
Reorganized all files
Change DenseCodec serialization logic.
Renamed methods to rank/select
Moved versioning footer to the columnar level
2023-01-16 17:24:49 +09:00
PSeitz
4bac945709 add ip field example (#1775) 2023-01-16 06:06:11 +01:00
trinity-1686a
16b704e190 make file_slice_for_range on sstable public (#1784) 2023-01-16 13:59:57 +09:00
PSeitz
6ca9a477f3 reuse stats for average (#1785)
* reuse stats for average

* fix count type
2023-01-13 23:32:27 +08:00
Shikhar Bhushan
2650111b76 EnableScoring::Disabled - optional Searcher (#1780) 2023-01-12 09:26:50 -05:00
PSeitz
1176555eff handle user input on get_docid_for_value_range (#1760)
* handle user input on get_docid_for_value_range

fixes #1757

* pass range as parameter
2023-01-12 14:20:16 +01:00
Adrien Guillo
f8d111a75e Merge pull request #1777 from quickwit-oss/guilload/ff-range-query-on-not-indexed-fields
Allow range queries via fast fields on non-indexed fields
2023-01-11 10:14:32 -05:00
Adrien Guillo
e17996f2fd Allow range queries via fast fields on non-indexed fields 2023-01-11 09:56:13 -05:00
Adrien Guillo
f3621c0487 Add license to tokenizer-api crate (#1778) 2023-01-11 05:26:41 +01:00
Adrien Guillo
14222a47a3 Fix typo (#1776) 2023-01-11 00:49:13 +09:00
Adam Reichold
8312c882a5 More cosmetic fixes for upcoming Clippy lints. (#1771) 2023-01-10 10:32:45 +01:00
Paul Masurel
7a8fce0ae7 Minor mini fixes 2023-01-10 14:15:30 +09:00
Michael Kleen
196e42f33e Add regex tokenizer (#1759)
This adds a regex tokenizer which tokenizes the text by using a
regex pattern to split.

Co-authored-by: Michael Kleen <mkleen@gmailw.com>
2023-01-10 13:38:37 +09:00
Adam Reichold
82a183bc2d Bump dependency on lru to from version 0.7.5 to version 0.9.0. (#1755) 2023-01-10 13:35:37 +09:00
dependabot[bot]
3090d49615 Update base64 requirement from 0.20.0 to 0.21.0 (#1769)
Updates the requirements on [base64](https://github.com/marshallpierce/rust-base64) to permit the latest version.
- [Release notes](https://github.com/marshallpierce/rust-base64/releases)
- [Changelog](https://github.com/marshallpierce/rust-base64/blob/master/RELEASE-NOTES.md)
- [Commits](https://github.com/marshallpierce/rust-base64/compare/v0.20.0...v0.21.0)

---
updated-dependencies:
- dependency-name: base64
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2023-01-10 13:35:05 +09:00
PSeitz
7c6cc818ae enable range query on fast field for u64 compatible types (#1762)
* enable range query on fast field for u64 compatible types

* rename, update benches
2023-01-10 04:08:26 +01:00
PSeitz
514d23a20c move tokenizer API to seperate crate (#1767)
closes #1766

Finding tantivy tokenizers is a frustrating experience currently, since
they need be updated for each tantivy version. That's unnecessary since
the API is rather stable anyway.
2023-01-09 06:37:38 +01:00
Paul Masurel
4f9efe654c Support for columnar (#1734)
* Added support for dynamic fast field.

See README for more information.

* Apply suggestions from code review

Co-authored-by: PSeitz <PSeitz@users.noreply.github.com>
2023-01-07 17:37:00 +09:00
Adam Reichold
1afa5bf3db Make construction of LevenshteinAutomatonBuilder for FuzzyTermQuery instances lazy. (#1756) 2023-01-06 12:44:49 +09:00
PSeitz
07a51eb7c8 refactor multivalue fastfield, refactor range query (#1749)
Introduce MakeZero trait, remove make_zero from FastValue
Merge two multivalue fastfield implementations into one
prepare range query on fastfield for different types
2023-01-05 12:09:50 +01:00
Adam Reichold
2080c370c2 Enable usage of FuzzyTermQuery for specific fields via QueryParser (#1750)
* Make nightly Clippy mostly happy.

* Document how to produce TermSetQuery queries using QueryParser.

* Enable construction of queries using FuzzyTermQuery via the QueryParser

* Use FxHashMap instead of HashMap in the QueryParser as these hash tables are not exposed to DoS attacks.

* Use a struct instead of a tuple to improve readability.
2023-01-04 18:11:27 +09:00
Daw-Chih Liou
b22f96624e doc: update comments in the faceted search example (#1737)
* doc: update comments in the faceted search example

* chore: format
2023-01-02 11:07:30 +01:00
pinkforest(she/her)
b78dc5e313 Bump prettytables (#1746) 2022-12-31 15:01:39 +01:00
Paul Masurel
3f915925af Fixing unit tests 2022-12-27 12:02:16 +09:00
Paul Masurel
9c5fef5af7 Fixing sstable proptest (#1743) 2022-12-26 16:29:33 +09:00
Paul Masurel
9948a84ebe Simplifies the count_ones definition. (#1742) 2022-12-26 16:08:01 +09:00
PSeitz
45156fd869 use group_by in translate_codec_idx_to_original_id (#1736) 2022-12-26 06:13:29 +01:00
151 changed files with 9818 additions and 2218 deletions

2
.gitignore vendored
View File

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

View File

@@ -15,7 +15,7 @@ rust-version = "1.62"
[dependencies] [dependencies]
oneshot = "0.1.5" oneshot = "0.1.5"
base64 = "0.20.0" base64 = "0.21.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"
@@ -48,7 +48,7 @@ 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.7.5" lru = "0.9.0"
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"
@@ -61,6 +61,7 @@ tantivy-query-grammar = { version= "0.19.0", path="./query-grammar" }
tantivy-bitpacker = { version= "0.3", path="./bitpacker" } tantivy-bitpacker = { version= "0.3", path="./bitpacker" }
common = { version= "0.5", path = "./common/", package = "tantivy-common" } common = { version= "0.5", path = "./common/", package = "tantivy-common" }
fastfield_codecs = { version= "0.3", path="./fastfield_codecs", default-features = false } fastfield_codecs = { version= "0.3", path="./fastfield_codecs", default-features = false }
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"
@@ -106,7 +107,7 @@ unstable = [] # useful for benches.
quickwit = ["sstable"] quickwit = ["sstable"]
[workspace] [workspace]
members = ["query-grammar", "bitpacker", "common", "fastfield_codecs", "ownedbytes", "stacker", "sstable", "columnar"] members = ["query-grammar", "bitpacker", "common", "fastfield_codecs", "ownedbytes", "stacker", "sstable", "tokenizer-api"]
# 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)
@@ -42,12 +42,12 @@ Your mileage WILL vary depending on the nature of queries and their load.
- 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, and hierarchical facet fields - Text, i64, u64, f64, dates, and hierarchical facet fields
- LZ4 compressed document store - Compressed document store (LZ4, Zstd, None, Brotli, Snap)
- 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: range buckets, average, and stats metrics - Aggregation Collector: histogram, 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
@@ -81,6 +81,10 @@ There are many ways to support this project.
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.
## Tokenizer
When implementing a tokenizer for tantivy depend on the `tantivy-tokenizer-api` crate.
## Minimum supported Rust version ## Minimum supported Rust version
Tantivy currently requires at least Rust 1.62 or later to compile. Tantivy currently requires at least Rust 1.62 or later to compile.

View File

@@ -10,10 +10,23 @@ serde_json = "1"
thiserror = "1" thiserror = "1"
fnv = "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" }
fastfield_codecs = { path = "../fastfield_codecs"}
itertools = "0.10" itertools = "0.10"
log = "0.4"
tantivy-bitpacker = { version= "0.3", path = "../bitpacker/" }
prettytable-rs = {version="0.10.0", optional= true}
rand = {version="0.8.3", optional= true}
fastdivide = "0.4"
measure_time = { version="0.8.2", optional=true}
[dev-dependencies] [dev-dependencies]
proptest = "1" proptest = "1"
more-asserts = "0.3.0"
rand = "0.8.3"
# temporary
[workspace]
members = []
[features]
unstable = []

View File

@@ -6,52 +6,45 @@ 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
- it does not required to be loaded in memory. - accessing a specific column does not require to load the entire columnar. It can be done in 2 to 3 random access.
- 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 appending rows to a writer. Users can create a columnar by inserting rows to a `ColumnarWriter`,
Nothing prevents a user from recording values with different to a same `column_key`. and serializing it into a `Write` object.
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:
- 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. - JsonValues are grouped into 3 types (String, Number, bool).
- 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. Values that corresponds to different groups are mapped to different columns. For instance, String values are treated independently
from Number or boolean values. `tantivy-columnar` will simply emit several columns associated to a given column_name.
- Only one column for a given json value type is emitted. If number values with different number types are recorded (e.g. u64, i64, f64),
`tantivy-columnar` will pick the first type that can represents the set of appended value, with the following prioriy order (`i64`, `u64`, `f64`).
`i64` is picked over `u64` as it is likely to yield less change of types. Most use cases strictly requiring `u64` show the
restriction on 50% of the values (e.g. a 64-bit hash). On the other hand, a lot of use cases can show rare negative value.
# Columnar format # Columnar format
Because this columnar format tries to avoid some coercion. This columnar format may have more than one column (with different types) associated to the same `column_name` (see [Coercion rules](#coercion-rules) above).
There can be several columns (with different type) associated to a single `column_name`. The `(column_name, columne_type)` couple however uniquely identifies a column.
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_INDEX] [COLUMNAR_KEY_TO_DATA_INDEX]
[COLUMNAR_FOOTER]; [COLUMNAR_FOOTER];
# Columns are sorted by their column key. # Columns are sorted by their column key.
COLUMNAR_DATA:= COLUMNAR_DATA:=
[COLUMN]+; [COLUMN_DATA]+;
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]
@@ -60,10 +53,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 quickwit/tantivy style sstable associates A sstable associates
`(column names, column_cardinality, column_type) to range of bytes. `(column name, column_cardinality, column_type) to range of bytes.
Column name may not contain the zero byte. Column name may not contain the zero byte `\0`.
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
@@ -71,3 +64,46 @@ 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

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

View File

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

45
columnar/src/TODO.md Normal file
View File

@@ -0,0 +1,45 @@
# 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
# Other
fix enhance column-cli
# Santa claus
autodetect datetime ipaddr, plug customizable tokenizer.

View File

@@ -0,0 +1,40 @@
use std::io;
use std::ops::Deref;
use std::sync::Arc;
use sstable::{Dictionary, VoidSSTable};
use crate::column::Column;
use crate::column_index::ColumnIndex;
/// Dictionary encoded column.
#[derive(Clone)]
pub struct BytesColumn {
pub(crate) dictionary: Arc<Dictionary<VoidSSTable>>,
pub(crate) term_ord_column: Column<u64>,
}
impl BytesColumn {
/// Returns `false` if the term does not exist (e.g. `term_ord` is greater or equal to the
/// overll number of terms).
pub fn term_ord_to_str(&self, term_ord: u64, output: &mut Vec<u8>) -> io::Result<bool> {
self.dictionary.ord_to_term(term_ord, output)
}
pub fn term_ords(&self) -> &Column<u64> {
&self.term_ord_column
}
}
impl Deref for BytesColumn {
type Target = ColumnIndex<'static>;
fn deref(&self) -> &Self::Target {
&**self.term_ords()
}
}
#[cfg(test)]
mod tests {
use crate::{ColumnarReader, ColumnarWriter};
}

View File

@@ -0,0 +1,56 @@
mod dictionary_encoded;
mod serialize;
use std::ops::Deref;
use std::sync::Arc;
use common::BinarySerializable;
pub use dictionary_encoded::BytesColumn;
pub use serialize::{open_column_bytes, open_column_u64, serialize_column_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>>,
}
use crate::column_index::Set;
impl<T: PartialOrd> Column<T> {
pub fn first(&self, row_id: RowId) -> Option<T> {
match &self.idx {
ColumnIndex::Full => Some(self.values.get_val(row_id)),
ColumnIndex::Optional(opt_idx) => {
let value_row_idx = opt_idx.rank_if_exists(row_id)?;
Some(self.values.get_val(value_row_idx))
}
ColumnIndex::Multivalued(_multivalued_index) => {
todo!();
}
}
}
}
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)
}
}

View File

@@ -0,0 +1,54 @@
use std::io;
use std::io::Write;
use std::sync::Arc;
use common::{CountingWriter, OwnedBytes};
use sstable::Dictionary;
use crate::column::{BytesColumn, Column};
use crate::column_index::{serialize_column_index, SerializableColumnIndex};
use crate::column_values::{
serialize_column_values, ColumnValues, MonotonicallyMappableToU64, ALL_CODEC_TYPES,
};
pub fn serialize_column_u64<T: MonotonicallyMappableToU64>(
column_index: SerializableColumnIndex<'_>,
column_values: &impl ColumnValues<T>,
output: &mut impl Write,
) -> io::Result<()> {
let mut counting_writer = CountingWriter::wrap(output);
serialize_column_index(column_index, &mut counting_writer)?;
let column_index_num_bytes = counting_writer.written_bytes() as u32;
let output = counting_writer.finish();
serialize_column_values(column_values, &ALL_CODEC_TYPES[..], 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_bytes(data: OwnedBytes) -> io::Result<BytesColumn> {
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)?;
Ok(BytesColumn {
dictionary,
term_ord_column,
})
}

View File

@@ -0,0 +1,40 @@
mod multivalued_index;
mod optional_index;
mod serialize;
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 Arc<dyn> apart from serialization this is not
// dynamic at all.
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 num_rows(&self) -> RowId {
match self {
ColumnIndex::Full => {
todo!()
}
ColumnIndex::Optional(optional_index) => optional_index.num_rows(),
ColumnIndex::Multivalued(multivalued_index) => multivalued_index.num_vals() - 1,
}
}
}

View File

@@ -0,0 +1,27 @@
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: MultivaluedIndex,
output: &mut impl Write,
) -> io::Result<()> {
crate::column_values::serialize_column_values(
&*multivalued_index.0,
&[FastFieldCodecType::Bitpacked, FastFieldCodecType::Linear],
output,
)?;
Ok(())
}
pub fn open_multivalued_index(bytes: OwnedBytes) -> io::Result<Arc<dyn ColumnValues<RowId>>> {
todo!();
}

View File

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

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

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

View File

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

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

@@ -0,0 +1,110 @@
use std::collections::HashMap;
use crate::column_index::optional_index::set_block::set_block::DENSE_BLOCK_NUM_BYTES;
use crate::column_index::optional_index::set_block::{DenseBlockCodec, SparseBlockCodec};
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! {
#[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

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

@@ -0,0 +1,70 @@
use std::io;
use std::io::Write;
use common::OwnedBytes;
use crate::column_index::multivalued_index::{serialize_multivalued_index, MultivaluedIndex};
use crate::column_index::optional_index::serialize_optional_index;
use crate::column_index::{ColumnIndex, SerializableOptionalIndex};
use crate::Cardinality;
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(MultivaluedIndex),
}
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<()> {
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, output)?
}
SerializableColumnIndex::Multivalued(multivalued_index) => {
serialize_multivalued_index(multivalued_index, output)?
}
}
Ok(())
}
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

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

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

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

@@ -0,0 +1,350 @@
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<'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

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

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

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

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

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

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

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

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

@@ -0,0 +1,333 @@
#![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;
// 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;
pub use self::serialize::{serialize_and_load, 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<Item: MonotonicallyMappableToU128>(
// bytes: OwnedBytes,
// ) -> io::Result<Arc<dyn Column<Item>>> {
// todo!();
// // let (bytes, _format_version) = read_format_version(bytes)?;
// // let (mut bytes, _null_index_footer) = read_null_index_footer(bytes)?;
// // let header = U128Header::deserialize(&mut bytes)?;
// // assert_eq!(header.codec_type, U128FastFieldCodecType::CompactSpace);
// // let reader = CompactSpaceDecompressor::open(bytes)?;
// // let inverted: StrictlyMonotonicMappingInverter<StrictlyMonotonicMappingToInternal<Item>> =
// // StrictlyMonotonicMappingToInternal::<Item>::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>;
}
/// The list of all available codecs for u64 convertible data.
pub const ALL_CODEC_TYPES: [FastFieldCodecType; 3] = [
FastFieldCodecType::Bitpacked,
FastFieldCodecType::BlockwiseLinear,
FastFieldCodecType::Linear,
];
#[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

@@ -0,0 +1,264 @@
use std::marker::PhantomData;
use fastdivide::DividerU64;
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,
}
}
}
// TODO
// 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 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

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

@@ -0,0 +1,343 @@
// 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 std::sync::Arc;
use common::{BinarySerializable, OwnedBytes, 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, VecColumn, ALL_CODEC_TYPES,
};
/// 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,
})
}
}
/// Return estimated compression for given codec in the value range [0.0..1.0], where 1.0 means no
/// compression.
pub(crate) fn estimate<T: MonotonicallyMappableToU64>(
typed_column: impl ColumnValues<T>,
codec_type: FastFieldCodecType,
) -> Option<f32> {
let column = monotonic_map_column(typed_column, StrictlyMonotonicMappingToInternal::<T>::new());
let min_value = column.min_value();
let gcd = super::gcd::find_gcd(column.iter().map(|val| val - min_value))
.filter(|gcd| gcd.get() > 1u64);
let mapping = StrictlyMonotonicMappingToInternalGCDBaseval::new(
gcd.map(|gcd| gcd.get()).unwrap_or(1u64),
min_value,
);
let normalized_column = monotonic_map_column(&column, mapping);
match codec_type {
FastFieldCodecType::Bitpacked => BitpackedCodec::estimate(&normalized_column),
FastFieldCodecType::Linear => LinearCodec::estimate(&normalized_column),
FastFieldCodecType::BlockwiseLinear => BlockwiseLinearCodec::estimate(&normalized_column),
}
}
// TODO
/// Serializes u128 values with the compact space codec.
// pub fn serialize_u128_new<F: Fn() -> I, I: Iterator<Item = u128>>(
// value_index: ColumnIndex,
// 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).unwrap();
// let null_index_footer = ColumnFooter {
// cardinality: value_index.get_cardinality(),
// null_index_codec: NullIndexCodec::Full,
// null_index_byte_range: 0..0,
// };
// append_null_index_footer(output, null_index_footer)?;
// append_format_version(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(())
}
/// 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();
super::serialize_column_values(&VecColumn::from(&column), &ALL_CODEC_TYPES, &mut buffer)
.unwrap();
super::open_u64_mapped(OwnedBytes::new(buffer)).unwrap()
}
#[cfg(test)]
mod tests {
use super::*;
#[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 + 7);
}
#[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 + 7);
}
#[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) + 7);
}
}

View File

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

@@ -0,0 +1,127 @@
use crate::utils::{place_bits, select_bits};
use crate::value::NumericalType;
use crate::InvalidData;
/// The column type represents the column type and can fit on 6-bits.
///
/// - bits[0..3]: Column category type.
/// - bits[3..6]: Numerical type if necessary.
#[derive(Hash, Eq, PartialEq, Debug, Clone, Copy)]
pub enum ColumnType {
Bytes,
Numerical(NumericalType),
Bool,
IpAddr,
}
impl ColumnType {
/// Encoded over 6 bits.
pub(crate) fn to_code(self) -> u8 {
let column_type_category;
let numerical_type_code: u8;
match self {
ColumnType::Bytes => {
column_type_category = ColumnTypeCategory::Str;
numerical_type_code = 0u8;
}
ColumnType::Numerical(numerical_type) => {
column_type_category = ColumnTypeCategory::Numerical;
numerical_type_code = numerical_type.to_code();
}
ColumnType::Bool => {
column_type_category = ColumnTypeCategory::Bool;
numerical_type_code = 0u8;
}
}
place_bits::<0, 3>(column_type_category.to_code()) | place_bits::<3, 6>(numerical_type_code)
}
pub(crate) fn try_from_code(code: u8) -> Result<ColumnType, InvalidData> {
if select_bits::<6, 8>(code) != 0u8 {
return Err(InvalidData);
}
let column_type_category_code = select_bits::<0, 3>(code);
let numerical_type_code = select_bits::<3, 6>(code);
let column_type_category = ColumnTypeCategory::try_from_code(column_type_category_code)?;
match column_type_category {
ColumnTypeCategory::Bool => {
if numerical_type_code != 0u8 {
return Err(InvalidData);
}
Ok(ColumnType::Bool)
}
ColumnTypeCategory::Str => {
if numerical_type_code != 0u8 {
return Err(InvalidData);
}
Ok(ColumnType::Bytes)
}
ColumnTypeCategory::Numerical => {
let numerical_type = NumericalType::try_from_code(numerical_type_code)?;
Ok(ColumnType::Numerical(numerical_type))
}
}
}
}
/// Column types are grouped into different categories that
/// corresponds to the different types of `JsonValue` types.
///
/// The columnar writer will apply coercion rules to make sure that
/// at most one column exist per `ColumnTypeCategory`.
///
/// See also [README.md].
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
#[repr(u8)]
pub(crate) enum ColumnTypeCategory {
Bool = 0u8,
Str = 1u8,
Numerical = 2u8,
}
impl ColumnTypeCategory {
pub fn to_code(self) -> u8 {
self as u8
}
pub fn try_from_code(code: u8) -> Result<Self, InvalidData> {
match code {
0u8 => Ok(Self::Bool),
1u8 => Ok(Self::Str),
2u8 => Ok(Self::Numerical),
_ => Err(InvalidData),
}
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::*;
use crate::Cardinality;
#[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 Ok(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 {
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

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

@@ -0,0 +1,28 @@
// Copyright (C) 2022 Quickwit, Inc.
//
// Quickwit is offered under the AGPL v3.0 and as commercial software.
// For commercial licensing, contact us at hello@quickwit.io.
//
// AGPL:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
//
mod column_type;
mod format_version;
mod reader;
mod writer;
pub use column_type::ColumnType;
pub use reader::ColumnarReader;
pub use writer::ColumnarWriter;

View File

@@ -0,0 +1,121 @@
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 fix ugly API
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 =
String::from_utf8_lossy(&key_bytes[..key_bytes.len() - 1]).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,84 +1,112 @@
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::{DocId, NumericalType}; use crate::{InvalidData, NumericalType, RowId};
/// When we build a columnar dataframe, we first just group /// When we build a columnar dataframe, we first just group
/// all mutations per column, and append them in append-only object. /// all mutations per column, and appends them in append-only buffer
/// 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(crate) enum ColumnOperation<T> { pub(super) enum ColumnOperation<T> {
NewDoc(DocId), NewDoc(RowId),
Value(T), Value(T),
} }
#[derive(Copy, Clone, Debug, Eq, PartialEq)] #[derive(Copy, Clone, Eq, PartialEq, Debug)]
struct ColumnOperationHeader { struct ColumnOperationMetadata {
typ_code: u8, op_type: ColumnOperationType,
len: u8, len: u8,
} }
impl ColumnOperationHeader { impl ColumnOperationMetadata {
fn to_code(self) -> u8 { fn to_code(self) -> u8 {
place_bits::<0, 4>(self.len) | place_bits::<4, 8>(self.typ_code) place_bits::<0, 4>(self.len) | place_bits::<4, 8>(self.op_type.to_code())
} }
fn from_code(code: u8) -> Self { fn try_from_code(code: u8) -> Result<Self, InvalidData> {
let len = select_bits::<0, 4>(code); let len = select_bits::<0, 4>(code);
let typ_code = select_bits::<4, 8>(code); let typ_code = select_bits::<4, 8>(code);
ColumnOperationHeader { typ_code, len } let column_type = ColumnOperationType::try_from_code(typ_code)?;
Ok(ColumnOperationMetadata {
op_type: column_type,
len,
})
} }
} }
const NEW_DOC_CODE: u8 = 0u8; #[derive(Copy, Clone, Eq, PartialEq, Debug)]
const NEW_VALUE_CODE: u8 = 1u8; #[repr(u8)]
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 fn serialize(self) -> impl AsRef<[u8]> { pub(super) fn serialize(self) -> impl AsRef<[u8]> {
let mut minibuf = MiniBuffer::default(); let mut minibuf = MiniBuffer::default();
let header = match self { let column_op_metadata = 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..]);
ColumnOperationHeader { ColumnOperationMetadata {
typ_code: NEW_DOC_CODE, op_type: ColumnOperationType::NewDoc,
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..]);
ColumnOperationHeader { ColumnOperationMetadata {
typ_code: NEW_VALUE_CODE, op_type: ColumnOperationType::AddValue,
len: symbol_len, len: symbol_len,
} }
} }
}; };
minibuf.bytes[0] = header.to_code(); minibuf.bytes[0] = column_op_metadata.to_code();
minibuf.len = 1 + header.len; // +1 for the metadata
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:
pub fn deserialize(bytes: &mut &[u8]) -> Option<Self> { /// this deserialize method is meant to target in memory.
let header_byte = pop_first_byte(bytes)?; pub(super) fn deserialize(bytes: &mut &[u8]) -> Option<Self> {
let column_op_header = ColumnOperationHeader::from_code(header_byte); let column_op_metadata_byte = pop_first_byte(bytes)?;
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_header.len as usize); (symbol_bytes, *bytes) = bytes.split_at(column_op_metadata.len as usize);
match column_op_header.typ_code { match column_op_metadata.op_type {
NEW_DOC_CODE => { ColumnOperationType::NewDoc => {
let new_doc = u32::deserialize(symbol_bytes); let new_doc = u32::deserialize(symbol_bytes);
Some(ColumnOperation::NewDoc(new_doc)) Some(ColumnOperation::NewDoc(new_doc))
} }
NEW_VALUE_CODE => { ColumnOperationType::AddValue => {
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);
}
} }
} }
} }
@@ -89,20 +117,25 @@ 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(crate) trait SymbolValue: Clone + Copy { pub(super) 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] = if self { 1u8 } else { 0u8 }; buffer[0] = u8::from(self);
1u8 1u8
} }
@@ -111,6 +144,18 @@ impl SymbolValue for bool {
} }
} }
impl SymbolValue for Ipv6Addr {
fn serialize(self, buffer: &mut [u8]) -> u8 {
// maybe not ueseful to use VIntEncoding for the mooment since we only use it for IP addr.
// We could roll our own RLE compression but it is overkill. let's stick to 8 bytes.
todo!();
}
fn deserialize(bytes: &[u8]) -> Self {
todo!();
}
}
#[derive(Default)] #[derive(Default)]
struct MiniBuffer { struct MiniBuffer {
pub bytes: [u8; 10], pub bytes: [u8; 10],
@@ -147,6 +192,9 @@ 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) => {
@@ -243,13 +291,14 @@ mod tests {
} }
#[test] #[test]
fn test_header_byte_serialization() { fn test_column_op_metadata_byte_serialization() {
for len in 0..=15 { for len in 0..=15 {
for typ_code in 0..=15 { for op_type in [ColumnOperationType::AddValue, ColumnOperationType::NewDoc] {
let header = ColumnOperationHeader { typ_code, len }; let column_op_metadata = ColumnOperationMetadata { op_type, len };
let header_code = header.to_code(); let column_op_metadata_code = column_op_metadata.to_code();
let serdeser_header = ColumnOperationHeader::from_code(header_code); let serdeser_metadata =
assert_eq!(header, serdeser_header); ColumnOperationMetadata::try_from_code(column_op_metadata_code).unwrap();
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::writer::column_operation::{ColumnOperation, SymbolValue}; use crate::{Cardinality, NumericalType, NumericalValue, RowId};
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 {
SameDoc = 0, Same = 0,
NextDoc = 1, Next = 1,
SkippedDoc = 2, Skipped = 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::SameDoc, Ordering::Less => DocumentStep::Same,
Ordering::Equal => DocumentStep::NextDoc, Ordering::Equal => DocumentStep::Next,
Ordering::Greater => DocumentStep::SkippedDoc, Ordering::Greater => DocumentStep::Skipped,
} }
} }
@@ -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(crate) fn operation_iterator<'a, V: SymbolValue>( pub(super) 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(crate) fn record<S: SymbolValue>(&mut self, doc: DocId, value: S, arena: &mut MemoryArena) { pub(super) fn record<S: SymbolValue>(&mut self, doc: RowId, 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::SameDoc => { DocumentStep::Same => {
// This is the last encounterred document. // This is the last encounterred document.
self.cardinality = Cardinality::Multivalued; self.cardinality = Cardinality::Multivalued;
} }
DocumentStep::NextDoc => { DocumentStep::Next => {
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::SkippedDoc => { DocumentStep::Skipped => {
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 fn get_cardinality(&self, num_docs: DocId) -> Cardinality { pub(crate) fn get_cardinality(&self, num_docs: RowId) -> Cardinality {
match delta_with_last_doc(self.last_doc_opt, num_docs) { match delta_with_last_doc(self.last_doc_opt, num_docs) {
DocumentStep::SameDoc | DocumentStep::NextDoc => self.cardinality, DocumentStep::Same | DocumentStep::Next => self.cardinality,
DocumentStep::SkippedDoc => self.cardinality.max(Cardinality::Optional), DocumentStep::Skipped => self.cardinality.max(Cardinality::Optional),
} }
} }
@@ -105,7 +105,7 @@ pub(crate) struct NumericalColumnWriter {
/// 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)]
pub(crate) struct CompatibleNumericalTypes { struct CompatibleNumericalTypes {
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. // f64 is always acceptable.
@@ -150,14 +150,15 @@ impl CompatibleNumericalTypes {
} }
impl NumericalColumnWriter { impl NumericalColumnWriter {
pub fn column_type_and_cardinality(&self, num_docs: DocId) -> (NumericalType, Cardinality) { pub fn column_type_and_cardinality(&self, num_docs: RowId) -> (NumericalType, Cardinality) {
let numerical_type = 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); let cardinality = self.column_writer.get_cardinality(num_docs);
(numerical_type, cardinality) (numerical_type, cardinality)
} }
pub fn record_numerical_value( pub fn record_numerical_value(
&mut self, &mut self,
doc: DocId, doc: RowId,
value: NumericalValue, value: NumericalValue,
arena: &mut MemoryArena, arena: &mut MemoryArena,
) { ) {
@@ -165,7 +166,7 @@ impl NumericalColumnWriter {
self.column_writer.record(doc, value, arena); self.column_writer.record(doc, value, arena);
} }
pub fn operation_iterator<'a>( pub(super) fn operation_iterator<'a>(
self, self,
arena: &MemoryArena, arena: &MemoryArena,
buffer: &'a mut Vec<u8>, buffer: &'a mut Vec<u8>,
@@ -175,13 +176,13 @@ impl NumericalColumnWriter {
} }
#[derive(Copy, Clone, Default)] #[derive(Copy, Clone, Default)]
pub struct StrColumnWriter { pub(crate) struct StrColumnWriter {
pub(crate) dictionary_id: u32, pub(crate) dictionary_id: u32,
pub(crate) column_writer: ColumnWriter, pub(crate) column_writer: ColumnWriter,
} }
impl StrColumnWriter { impl StrColumnWriter {
pub fn with_dictionary_id(dictionary_id: u32) -> StrColumnWriter { pub(crate) fn with_dictionary_id(dictionary_id: u32) -> StrColumnWriter {
StrColumnWriter { StrColumnWriter {
dictionary_id, dictionary_id,
column_writer: Default::default(), column_writer: Default::default(),
@@ -190,7 +191,7 @@ impl StrColumnWriter {
pub(crate) fn record_bytes( pub(crate) fn record_bytes(
&mut self, &mut self,
doc: DocId, doc: RowId,
bytes: &[u8], bytes: &[u8],
dictionaries: &mut [DictionaryBuilder], dictionaries: &mut [DictionaryBuilder],
arena: &mut MemoryArena, arena: &mut MemoryArena,
@@ -199,7 +200,7 @@ impl StrColumnWriter {
self.column_writer.record(doc, unordered_id, arena); self.column_writer.record(doc, unordered_id, arena);
} }
pub(crate) fn operation_iterator<'a>( pub(super) fn operation_iterator<'a>(
&self, &self,
arena: &MemoryArena, arena: &MemoryArena,
byte_buffer: &'a mut Vec<u8>, byte_buffer: &'a mut Vec<u8>,
@@ -214,20 +215,14 @@ 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::NextDoc); assert_eq!(delta_with_last_doc(None, 0u32), DocumentStep::Next);
assert_eq!(delta_with_last_doc(None, 1u32), DocumentStep::SkippedDoc); assert_eq!(delta_with_last_doc(None, 1u32), DocumentStep::Skipped);
assert_eq!(delta_with_last_doc(None, 2u32), DocumentStep::SkippedDoc); assert_eq!(delta_with_last_doc(None, 2u32), DocumentStep::Skipped);
assert_eq!(delta_with_last_doc(Some(0u32), 0u32), DocumentStep::SameDoc); assert_eq!(delta_with_last_doc(Some(0u32), 0u32), DocumentStep::Same);
assert_eq!(delta_with_last_doc(Some(1u32), 1u32), DocumentStep::SameDoc); assert_eq!(delta_with_last_doc(Some(1u32), 1u32), DocumentStep::Same);
assert_eq!(delta_with_last_doc(Some(1u32), 2u32), DocumentStep::NextDoc); assert_eq!(delta_with_last_doc(Some(1u32), 2u32), DocumentStep::Next);
assert_eq!( assert_eq!(delta_with_last_doc(Some(1u32), 3u32), DocumentStep::Skipped);
delta_with_last_doc(Some(1u32), 3u32), assert_eq!(delta_with_last_doc(Some(1u32), 4u32), DocumentStep::Skipped);
DocumentStep::SkippedDoc
);
assert_eq!(
delta_with_last_doc(Some(1u32), 4u32),
DocumentStep::SkippedDoc
);
} }
#[track_caller] #[track_caller]

View File

@@ -3,40 +3,53 @@ mod column_writers;
mod serializer; mod serializer;
mod value_index; mod value_index;
use std::io::{self, Write}; use std::io;
use std::net::Ipv6Addr;
use column_operation::ColumnOperation; use column_operation::ColumnOperation;
use fastfield_codecs::serialize::ValueIndexInfo; use common::CountingWriter;
use fastfield_codecs::{Column, MonotonicallyMappableToU64, VecColumn};
use serializer::ColumnarSerializer; use serializer::ColumnarSerializer;
use stacker::{Addr, ArenaHashMap, MemoryArena}; use stacker::{Addr, ArenaHashMap, MemoryArena};
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality, GeneralType}; use crate::column_index::SerializableColumnIndex;
use crate::dictionary::{DictionaryBuilder, IdMapping, UnorderedId}; use crate::column_values::{ColumnValues, MonotonicallyMappableToU64, VecColumn};
use crate::columnar::column_type::{ColumnType, ColumnTypeCategory};
use crate::columnar::writer::column_writers::{
ColumnWriter, NumericalColumnWriter, StrColumnWriter,
};
use crate::columnar::writer::value_index::{IndexBuilder, PreallocatedIndexBuilders};
use crate::dictionary::{DictionaryBuilder, TermIdMapping, UnorderedId};
use crate::value::{Coerce, NumericalType, NumericalValue}; use crate::value::{Coerce, NumericalType, NumericalValue};
use crate::writer::column_writers::{ColumnWriter, NumericalColumnWriter, StrColumnWriter}; use crate::{Cardinality, RowId};
use crate::writer::value_index::{IndexBuilder, SpareIndexBuilders};
use crate::{Cardinality, DocId};
/// Threshold above which a column data will be compressed /// This is a set of buffers that are used to temporarily write the values into before passing them
/// using ZSTD. /// to the fast field codecs.
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)] #[derive(Default)]
struct SpareBuffers { struct SpareBuffers {
value_index_builders: SpareIndexBuilders, value_index_builders: PreallocatedIndexBuilders,
i64_values: Vec<i64>, i64_values: Vec<i64>,
u64_values: Vec<u64>, u64_values: Vec<u64>,
f64_values: Vec<f64>, f64_values: Vec<f64>,
bool_values: Vec<bool>, bool_values: Vec<bool>,
column_buffer: Vec<u8>,
} }
/// 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 { pub struct ColumnarWriter {
numerical_field_hash_map: ArenaHashMap, numerical_field_hash_map: ArenaHashMap,
bool_field_hash_map: ArenaHashMap, bool_field_hash_map: ArenaHashMap,
ip_addr_field_hash_map: ArenaHashMap,
bytes_field_hash_map: ArenaHashMap, bytes_field_hash_map: ArenaHashMap,
arena: MemoryArena, arena: MemoryArena,
// Dictionaries used to store dictionary-encoded values. // Dictionaries used to store dictionary-encoded values.
@@ -58,11 +71,11 @@ impl Default for ColumnarWriter {
} }
impl ColumnarWriter { impl ColumnarWriter {
pub fn record_numerical( pub fn record_numerical<T: Into<NumericalValue> + Copy>(
&mut self, &mut self,
doc: DocId, doc: RowId,
column_name: &str, column_name: &str,
numerical_value: NumericalValue, numerical_value: T,
) { ) {
assert!( assert!(
!column_name.as_bytes().contains(&0u8), !column_name.as_bytes().contains(&0u8),
@@ -73,13 +86,29 @@ impl ColumnarWriter {
column_name.as_bytes(), column_name.as_bytes(),
|column_opt: Option<NumericalColumnWriter>| { |column_opt: Option<NumericalColumnWriter>| {
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default(); let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
column.record_numerical_value(doc, numerical_value, arena); column.record_numerical_value(doc, numerical_value.into(), arena);
column column
}, },
); );
} }
pub fn record_bool(&mut self, doc: DocId, column_name: &str, val: bool) { 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) {
assert!( assert!(
!column_name.as_bytes().contains(&0u8), !column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte" "key may not contain the 0 byte"
@@ -95,7 +124,7 @@ impl ColumnarWriter {
); );
} }
pub fn record_str(&mut self, doc: DocId, column_name: &str, value: &[u8]) { pub fn record_str(&mut self, doc: RowId, column_name: &str, value: &str) {
assert!( assert!(
!column_name.as_bytes().contains(&0u8), !column_name.as_bytes().contains(&0u8),
"key may not contain the 0 byte" "key may not contain the 0 byte"
@@ -109,93 +138,82 @@ impl ColumnarWriter {
column_name.as_bytes(), column_name.as_bytes(),
|column_opt: Option<StrColumnWriter>| { |column_opt: Option<StrColumnWriter>| {
let mut column: StrColumnWriter = column_opt.unwrap_or_else(|| { let mut column: StrColumnWriter = column_opt.unwrap_or_else(|| {
// Each column has its own dictionary
let dictionary_id = dictionaries.len() as u32; let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default()); dictionaries.push(DictionaryBuilder::default());
StrColumnWriter::with_dictionary_id(dictionary_id) StrColumnWriter::with_dictionary_id(dictionary_id)
}); });
column.record_bytes(doc, value, dictionaries, arena); column.record_bytes(doc, value.as_bytes(), dictionaries, arena);
column column
}, },
); );
} }
pub fn serialize(&mut self, num_docs: DocId, wrt: &mut dyn io::Write) -> io::Result<()> { pub fn serialize(&mut self, num_docs: RowId, wrt: &mut dyn io::Write) -> io::Result<()> {
let mut serializer = ColumnarSerializer::new(wrt); let mut serializer = ColumnarSerializer::new(wrt);
let mut field_columns: Vec<(&[u8], GeneralType, Addr)> = self let mut field_columns: Vec<(&[u8], ColumnTypeCategory, Addr)> = self
.numerical_field_hash_map .numerical_field_hash_map
.iter() .iter()
.map(|(term, addr, _)| (term, GeneralType::Numerical, addr)) .map(|(term, addr, _)| (term, ColumnTypeCategory::Numerical, addr))
.collect(); .collect();
field_columns.extend( field_columns.extend(
self.bytes_field_hash_map self.bytes_field_hash_map
.iter() .iter()
.map(|(term, addr, _)| (term, GeneralType::Str, addr)), .map(|(term, addr, _)| (term, ColumnTypeCategory::Str, addr)),
); );
field_columns.extend( field_columns.extend(
self.bool_field_hash_map self.bool_field_hash_map
.iter() .iter()
.map(|(term, addr, _)| (term, GeneralType::Bool, addr)), .map(|(term, addr, _)| (term, ColumnTypeCategory::Bool, addr)),
); );
field_columns.sort_unstable_by_key(|(column_name, col_type, _)| (*column_name, *col_type)); field_columns.sort_unstable_by_key(|(column_name, col_type, _)| (*column_name, *col_type));
let (arena, buffers, dictionaries) = (&self.arena, &mut self.buffers, &self.dictionaries); let (arena, buffers, dictionaries) = (&self.arena, &mut self.buffers, &self.dictionaries);
let mut symbol_byte_buffer: Vec<u8> = Vec::new(); let mut symbol_byte_buffer: Vec<u8> = Vec::new();
for (column_name, bytes_or_numerical, addr) in field_columns { for (column_name, bytes_or_numerical, addr) in field_columns {
match bytes_or_numerical { match bytes_or_numerical {
GeneralType::Bool => { ColumnTypeCategory::Bool => {
let column_writer: ColumnWriter = self.bool_field_hash_map.read(addr); let column_writer: ColumnWriter = self.bool_field_hash_map.read(addr);
let cardinality = column_writer.get_cardinality(num_docs); let cardinality = column_writer.get_cardinality(num_docs);
let column_type_and_cardinality = ColumnTypeAndCardinality { let mut column_serializer =
cardinality, serializer.serialize_column(column_name, ColumnType::Bool);
typ: ColumnType::Bool,
};
let column_serializer =
serializer.serialize_column(column_name, column_type_and_cardinality);
serialize_bool_column( serialize_bool_column(
cardinality, cardinality,
num_docs, num_docs,
column_writer.operation_iterator(arena, &mut symbol_byte_buffer), column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers, buffers,
column_serializer, &mut column_serializer,
)?; )?;
} }
GeneralType::Str => { ColumnTypeCategory::Str => {
let str_column_writer: StrColumnWriter = self.bytes_field_hash_map.read(addr); let str_column_writer: StrColumnWriter = self.bytes_field_hash_map.read(addr);
let dictionary_builder = let dictionary_builder =
&dictionaries[str_column_writer.dictionary_id as usize]; &dictionaries[str_column_writer.dictionary_id as usize];
let cardinality = str_column_writer.column_writer.get_cardinality(num_docs); let cardinality = str_column_writer.column_writer.get_cardinality(num_docs);
let column_type_and_cardinality = ColumnTypeAndCardinality { let mut column_serializer =
cardinality, serializer.serialize_column(column_name, ColumnType::Bytes);
typ: ColumnType::Bytes,
};
let column_serializer =
serializer.serialize_column(column_name, column_type_and_cardinality);
serialize_bytes_column( serialize_bytes_column(
cardinality, cardinality,
num_docs, num_docs,
dictionary_builder, dictionary_builder,
str_column_writer.operation_iterator(arena, &mut symbol_byte_buffer), str_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers, buffers,
column_serializer, &mut column_serializer,
)?; )?;
} }
GeneralType::Numerical => { ColumnTypeCategory::Numerical => {
let numerical_column_writer: NumericalColumnWriter = let numerical_column_writer: NumericalColumnWriter =
self.numerical_field_hash_map.read(addr); self.numerical_field_hash_map.read(addr);
let (numerical_type, cardinality) = let (numerical_type, cardinality) =
numerical_column_writer.column_type_and_cardinality(num_docs); numerical_column_writer.column_type_and_cardinality(num_docs);
let column_type_and_cardinality = ColumnTypeAndCardinality { let mut column_serializer = serializer
cardinality, .serialize_column(column_name, ColumnType::Numerical(numerical_type));
typ: ColumnType::Numerical(numerical_type),
};
let column_serializer =
serializer.serialize_column(column_name, column_type_and_cardinality);
serialize_numerical_column( serialize_numerical_column(
cardinality, cardinality,
num_docs, num_docs,
numerical_type, numerical_type,
numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer), numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
buffers, buffers,
column_serializer, &mut column_serializer,
)?; )?;
} }
}; };
@@ -205,41 +223,28 @@ impl ColumnarWriter {
} }
} }
fn compress_and_write_column<W: io::Write>(column_bytes: &[u8], wrt: &mut W) -> io::Result<()> { fn serialize_bytes_column(
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, cardinality: Cardinality,
num_docs: DocId, num_docs: RowId,
dictionary_builder: &DictionaryBuilder, dictionary_builder: &DictionaryBuilder,
operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>, operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>,
buffers: &mut SpareBuffers, buffers: &mut SpareBuffers,
mut wrt: W, wrt: impl io::Write,
) -> io::Result<()> { ) -> io::Result<()> {
let SpareBuffers { let SpareBuffers {
value_index_builders, value_index_builders,
u64_values, u64_values,
column_buffer,
.. ..
} = buffers; } = buffers;
column_buffer.clear(); let mut counting_writer = CountingWriter::wrap(wrt);
let id_mapping: IdMapping = dictionary_builder.serialize(column_buffer)?; let term_id_mapping: TermIdMapping = dictionary_builder.serialize(&mut counting_writer)?;
let dictionary_num_bytes: u32 = column_buffer.len() as u32; 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>| { let operation_iterator = operation_it.map(|symbol: ColumnOperation<UnorderedId>| {
// We map unordered ids to ordered ids. // We map unordered ids to ordered ids.
match symbol { match symbol {
ColumnOperation::Value(unordered_id) => { ColumnOperation::Value(unordered_id) => {
let ordered_id = id_mapping.to_ord(unordered_id); let ordered_id = term_id_mapping.to_ord(unordered_id);
ColumnOperation::Value(ordered_id.0 as u64) ColumnOperation::Value(ordered_id.0 as u64)
} }
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc), ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
@@ -251,30 +256,27 @@ fn serialize_bytes_column<W: io::Write>(
num_docs, num_docs,
value_index_builders, value_index_builders,
u64_values, u64_values,
column_buffer, &mut wrt,
)?; )?;
column_buffer.write_all(&dictionary_num_bytes.to_le_bytes()[..])?; wrt.write_all(&dictionary_num_bytes.to_le_bytes()[..])?;
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(()) Ok(())
} }
fn serialize_numerical_column<W: io::Write>( fn serialize_numerical_column(
cardinality: Cardinality, cardinality: Cardinality,
num_docs: DocId, num_docs: RowId,
numerical_type: NumericalType, numerical_type: NumericalType,
op_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>, op_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
buffers: &mut SpareBuffers, buffers: &mut SpareBuffers,
mut wrt: W, wrt: &mut impl io::Write,
) -> io::Result<()> { ) -> io::Result<()> {
let SpareBuffers { let SpareBuffers {
value_index_builders, value_index_builders,
u64_values, u64_values,
i64_values, i64_values,
f64_values, f64_values,
column_buffer,
.. ..
} = buffers; } = buffers;
column_buffer.clear();
match numerical_type { match numerical_type {
NumericalType::I64 => { NumericalType::I64 => {
serialize_column( serialize_column(
@@ -283,7 +285,7 @@ fn serialize_numerical_column<W: io::Write>(
num_docs, num_docs,
value_index_builders, value_index_builders,
i64_values, i64_values,
column_buffer, wrt,
)?; )?;
} }
NumericalType::U64 => { NumericalType::U64 => {
@@ -293,7 +295,7 @@ fn serialize_numerical_column<W: io::Write>(
num_docs, num_docs,
value_index_builders, value_index_builders,
u64_values, u64_values,
column_buffer, wrt,
)?; )?;
} }
NumericalType::F64 => { NumericalType::F64 => {
@@ -303,37 +305,33 @@ fn serialize_numerical_column<W: io::Write>(
num_docs, num_docs,
value_index_builders, value_index_builders,
f64_values, f64_values,
column_buffer, wrt,
)?; )?;
} }
}; };
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(()) Ok(())
} }
fn serialize_bool_column<W: io::Write>( fn serialize_bool_column(
cardinality: Cardinality, cardinality: Cardinality,
num_docs: DocId, num_docs: RowId,
column_operations_it: impl Iterator<Item = ColumnOperation<bool>>, column_operations_it: impl Iterator<Item = ColumnOperation<bool>>,
buffers: &mut SpareBuffers, buffers: &mut SpareBuffers,
mut wrt: W, wrt: &mut impl io::Write,
) -> io::Result<()> { ) -> io::Result<()> {
let SpareBuffers { let SpareBuffers {
value_index_builders, value_index_builders,
bool_values, bool_values,
column_buffer,
.. ..
} = buffers; } = buffers;
column_buffer.clear();
serialize_column( serialize_column(
column_operations_it, column_operations_it,
cardinality, cardinality,
num_docs, num_docs,
value_index_builders, value_index_builders,
bool_values, bool_values,
column_buffer, wrt,
)?; )?;
compress_and_write_column(column_buffer, &mut wrt)?;
Ok(()) Ok(())
} }
@@ -342,51 +340,43 @@ fn serialize_column<
>( >(
op_iterator: impl Iterator<Item = ColumnOperation<T>>, op_iterator: impl Iterator<Item = ColumnOperation<T>>,
cardinality: Cardinality, cardinality: Cardinality,
num_docs: DocId, num_docs: RowId,
value_index_builders: &mut SpareIndexBuilders, value_index_builders: &mut PreallocatedIndexBuilders,
values: &mut Vec<T>, values: &mut Vec<T>,
wrt: &mut Vec<u8>, mut wrt: impl io::Write,
) -> io::Result<()> ) -> io::Result<()>
where where
for<'a> VecColumn<'a, T>: Column<T>, for<'a> VecColumn<'a, T>: ColumnValues<T>,
{ {
values.clear(); values.clear();
match cardinality { let serializable_column_index = match cardinality {
Cardinality::Required => { Cardinality::Full => {
consume_operation_iterator( consume_operation_iterator(
op_iterator, op_iterator,
value_index_builders.borrow_required_index_builder(), value_index_builders.borrow_required_index_builder(),
values, values,
); );
fastfield_codecs::serialize( SerializableColumnIndex::Full
VecColumn::from(&values[..]),
wrt,
&fastfield_codecs::ALL_CODEC_TYPES[..],
)?;
} }
Cardinality::Optional => { Cardinality::Optional => {
let optional_index_builder = value_index_builders.borrow_optional_index_builder(); let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values); consume_operation_iterator(op_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_docs); let optional_index = optional_index_builder.finish(num_docs);
fastfield_codecs::serialize::serialize_new( SerializableColumnIndex::Optional(Box::new(optional_index))
ValueIndexInfo::SingleValue(Box::new(optional_index)),
VecColumn::from(&values[..]),
wrt,
&fastfield_codecs::ALL_CODEC_TYPES[..],
)?;
} }
Cardinality::Multivalued => { Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder(); let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values); consume_operation_iterator(op_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_docs); let multivalued_index = multivalued_index_builder.finish(num_docs);
fastfield_codecs::serialize::serialize_new( todo!();
ValueIndexInfo::MultiValue(Box::new(multivalued_index)), // SerializableColumnIndex::Multivalued(Box::new(multivalued_index))
VecColumn::from(&values[..]),
wrt,
&fastfield_codecs::ALL_CODEC_TYPES[..],
)?;
} }
} };
crate::column::serialize_column_u64(
serializable_column_index,
&VecColumn::from(&values[..]),
&mut wrt,
)?;
Ok(()) Ok(())
} }
@@ -410,7 +400,7 @@ fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>(
for symbol in operation_iterator { for symbol in operation_iterator {
match symbol { match symbol {
ColumnOperation::NewDoc(doc) => { ColumnOperation::NewDoc(doc) => {
index_builder.record_doc(doc); index_builder.record_row(doc);
} }
ColumnOperation::Value(value) => { ColumnOperation::Value(value) => {
index_builder.record_value(); index_builder.record_value();
@@ -420,6 +410,52 @@ fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>(
} }
} }
// /// Serializes the column with the codec with the best estimate on the data.
// fn serialize_numerical<T: MonotonicallyMappableToU64>(
// value_index: ValueIndexInfo,
// typed_column: impl Column<T>,
// output: &mut impl io::Write,
// codecs: &[FastFieldCodecType],
// ) -> io::Result<()> {
// let counting_writer = CountingWriter::wrap(output);
// serialize_value_index(value_index, output)?;
// let value_index_len = counting_writer.written_bytes();
// let output = counting_writer.finish();
// serialize_column(value_index, output)?;
// let column = monotonic_map_column(
// typed_column,
// crate::column::monotonic_mapping::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)?;
// let column_header = ColumnFooter {
// value_index_len: todo!(),
// cardinality: todo!(),
// };
// let null_index_footer = NullIndexFooter {
// cardinality: value_index.get_cardinality(),
// null_index_codec: NullIndexCodec::Full,
// null_index_byte_range: 0..0,
// };
// append_null_index_footer(output, null_index_footer)?;
// Ok(())
// }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use column_operation::ColumnOperation; use column_operation::ColumnOperation;
@@ -427,7 +463,6 @@ mod tests {
use super::*; use super::*;
use crate::value::NumericalValue; use crate::value::NumericalValue;
use crate::Cardinality;
#[test] #[test]
fn test_column_writer_required_simple() { fn test_column_writer_required_simple() {
@@ -436,7 +471,7 @@ mod tests {
column_writer.record(0u32, NumericalValue::from(14i64), &mut arena); column_writer.record(0u32, NumericalValue::from(14i64), &mut arena);
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena); column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena); column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
assert_eq!(column_writer.get_cardinality(3), Cardinality::Required); assert_eq!(column_writer.get_cardinality(3), Cardinality::Full);
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
.operation_iterator(&mut arena, &mut buffer) .operation_iterator(&mut arena, &mut buffer)

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::column_type_header::ColumnTypeAndCardinality; use crate::columnar::ColumnType;
pub struct ColumnarSerializer<W: io::Write> { pub struct ColumnarSerializer<W: io::Write> {
wrt: CountingWriter<W>, wrt: CountingWriter<W>,
@@ -15,15 +15,11 @@ 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<'a>( fn prepare_key(key: &[u8], column_type: ColumnType, buffer: &mut Vec<u8>) {
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_cardinality.to_code()); buffer.push(column_type.to_code());
} }
impl<W: io::Write> ColumnarSerializer<W> { impl<W: io::Write> ColumnarSerializer<W> {
@@ -40,14 +36,10 @@ 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_cardinality: ColumnTypeAndCardinality, column_type: ColumnType,
) -> impl io::Write + 'a { ) -> impl io::Write + 'a {
let start_offset = self.wrt.written_bytes(); let start_offset = self.wrt.written_bytes();
prepare_key( prepare_key(column_name, column_type, &mut self.prepare_key_buffer);
column_name,
column_type_cardinality,
&mut self.prepare_key_buffer,
);
ColumnSerializer { ColumnSerializer {
columnar_serializer: self, columnar_serializer: self,
start_offset, start_offset,
@@ -59,6 +51,9 @@ 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(())
} }
} }
@@ -97,20 +92,15 @@ 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::column_type_header::ColumnType; use crate::columnar::column_type::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();
let column_type_and_cardinality = ColumnTypeAndCardinality { prepare_key(b"root\0child", ColumnType::Bytes, &mut buffer);
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], column_type_and_cardinality.to_code()); assert_eq!(buffer[11], ColumnType::Bytes.to_code());
} }
} }

View File

@@ -1,44 +1,51 @@
use fastfield_codecs::serialize::{MultiValueIndexInfo, SingleValueIndexInfo}; use crate::column_index::SerializableOptionalIndex;
use crate::column_values::{ColumnValues, VecColumn};
use crate::DocId; use crate::RowId;
/// 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. /// and can then serialize the results into an index to associate docids with their value[s].
/// ///
/// 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_doc(&mut self, doc: DocId); fn record_row(&mut self, doc: RowId);
#[inline] #[inline]
fn record_value(&mut self) {} fn record_value(&mut self) {}
} }
/// The RequiredIndexBuilder does nothing. /// The FullIndexBuilder does nothing.
#[derive(Default)] #[derive(Default)]
pub struct RequiredIndexBuilder; pub struct FullIndexBuilder;
impl IndexBuilder for RequiredIndexBuilder { impl IndexBuilder for FullIndexBuilder {
#[inline(always)] #[inline(always)]
fn record_doc(&mut self, _doc: DocId) {} fn record_row(&mut self, _doc: RowId) {}
} }
#[derive(Default)] #[derive(Default)]
pub struct OptionalIndexBuilder { pub struct OptionalIndexBuilder {
docs: Vec<DocId>, docs: Vec<RowId>,
} }
struct SingleValueArrayIndex<'a> { struct SingleValueArrayIndex<'a> {
docs: &'a [DocId], // RowIds with a value, in a strictly increasing order
num_docs: DocId, row_ids: &'a [RowId],
num_rows: RowId,
} }
impl<'a> SingleValueIndexInfo for SingleValueArrayIndex<'a> { impl<'a> SerializableOptionalIndex<'a> for SingleValueArrayIndex<'a> {
fn num_vals(&self) -> u32 { fn num_rows(&self) -> RowId {
self.num_docs as u32 self.num_rows
} }
fn non_null_rows(&self) -> Box<dyn Iterator<Item = RowId> + 'a> {
Box::new(self.row_ids.iter().copied())
}
}
impl OptionalIndexBuilder {
fn num_non_nulls(&self) -> u32 { fn num_non_nulls(&self) -> u32 {
self.docs.len() as u32 self.docs.len() as u32
} }
@@ -49,16 +56,16 @@ impl<'a> SingleValueIndexInfo for SingleValueArrayIndex<'a> {
} }
impl OptionalIndexBuilder { impl OptionalIndexBuilder {
pub fn finish(&mut self, num_docs: DocId) -> impl SingleValueIndexInfo + '_ { pub fn finish<'a>(&'a mut self, num_rows: RowId) -> impl SerializableOptionalIndex + 'a {
debug_assert!(self debug_assert!(self
.docs .docs
.last() .last()
.copied() .copied()
.map(|last_doc| last_doc < num_docs) .map(|last_doc| last_doc < num_rows)
.unwrap_or(true)); .unwrap_or(true));
SingleValueArrayIndex { SingleValueArrayIndex {
docs: &self.docs[..], row_ids: &self.docs[..],
num_docs, num_rows,
} }
} }
@@ -69,7 +76,7 @@ impl OptionalIndexBuilder {
impl IndexBuilder for OptionalIndexBuilder { impl IndexBuilder for OptionalIndexBuilder {
#[inline(always)] #[inline(always)]
fn record_doc(&mut self, doc: DocId) { fn record_row(&mut self, doc: RowId) {
debug_assert!(self debug_assert!(self
.docs .docs
.last() .last()
@@ -82,52 +89,32 @@ impl IndexBuilder for OptionalIndexBuilder {
#[derive(Default)] #[derive(Default)]
pub struct MultivaluedIndexBuilder { pub struct MultivaluedIndexBuilder {
// TODO should we switch to `start_offset`? start_offsets: Vec<RowId>,
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: DocId) -> impl MultiValueIndexInfo + '_ { pub fn finish(&mut self, num_docs: RowId) -> impl ColumnValues<u32> + '_ {
self.end_values self.start_offsets
.resize(num_docs as usize, self.total_num_vals_seen); .resize(num_docs as usize, self.total_num_vals_seen);
MultivaluedValueArrayIndex { VecColumn {
end_offsets: &self.end_values[..], values: &&self.start_offsets[..],
min_value: 0,
max_value: self.start_offsets.last().copied().unwrap_or(0),
} }
} }
fn reset(&mut self) { fn reset(&mut self) {
self.end_values.clear(); self.start_offsets.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_doc(&mut self, doc: DocId) { fn record_row(&mut self, row_id: RowId) {
self.end_values self.start_offsets
.resize(doc as usize, self.total_num_vals_seen); .resize(row_id as usize + 1, self.total_num_vals_seen);
} }
fn record_value(&mut self) { fn record_value(&mut self) {
@@ -138,14 +125,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 SpareIndexBuilders { pub struct PreallocatedIndexBuilders {
required_index_builder: RequiredIndexBuilder, required_index_builder: FullIndexBuilder,
optional_index_builder: OptionalIndexBuilder, optional_index_builder: OptionalIndexBuilder,
multivalued_index_builder: MultivaluedIndexBuilder, multivalued_index_builder: MultivaluedIndexBuilder,
} }
impl SpareIndexBuilders { impl PreallocatedIndexBuilders {
pub fn borrow_required_index_builder(&mut self) -> &mut RequiredIndexBuilder { pub fn borrow_required_index_builder(&mut self) -> &mut FullIndexBuilder {
&mut self.required_index_builder &mut self.required_index_builder
} }
@@ -167,22 +154,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_doc(0u32); opt_value_index_builder.record_row(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)
.iter() .non_null_rows()
.collect::<Vec<u32>>(), .collect::<Vec<u32>>(),
&[0] &[0]
); );
opt_value_index_builder.reset(); opt_value_index_builder.reset();
opt_value_index_builder.record_doc(1u32); opt_value_index_builder.record_row(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)
.iter() .non_null_rows()
.collect::<Vec<u32>>(), .collect::<Vec<u32>>(),
&[1] &[1]
); );
@@ -191,10 +178,10 @@ 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_doc(1u32); multivalued_value_index_builder.record_row(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_doc(2u32); multivalued_value_index_builder.record_row(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
@@ -204,7 +191,7 @@ mod tests {
vec![0, 0, 2, 3] vec![0, 0, 2, 3]
); );
multivalued_value_index_builder.reset(); multivalued_value_index_builder.reset();
multivalued_value_index_builder.record_doc(2u32); multivalued_value_index_builder.record_row(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!(

View File

@@ -3,11 +3,11 @@ use std::io;
use fnv::FnvHashMap; use fnv::FnvHashMap;
use sstable::SSTable; use sstable::SSTable;
pub(crate) struct IdMapping { pub(crate) struct TermIdMapping {
unordered_to_ord: Vec<OrderedId>, unordered_to_ord: Vec<OrderedId>,
} }
impl IdMapping { impl TermIdMapping {
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<IdMapping> { pub fn serialize<'a, W: io::Write + 'a>(&self, wrt: &mut W) -> io::Result<TermIdMapping> {
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(IdMapping { unordered_to_ord }) Ok(TermIdMapping { unordered_to_ord })
} }
} }

View File

@@ -0,0 +1,95 @@
use std::io;
use std::net::IpAddr;
use common::file_slice::FileSlice;
use common::{HasLen, OwnedBytes};
use crate::column::{BytesColumn, Column};
use crate::columnar::ColumnType;
use crate::DateTime;
#[derive(Clone)]
pub enum DynamicColumn {
Bool(Column<bool>),
I64(Column<i64>),
U64(Column<u64>),
F64(Column<f64>),
IpAddr(Column<IpAddr>),
DateTime(Column<DateTime>),
Str(BytesColumn),
}
impl From<Column<i64>> for DynamicColumn {
fn from(column_i64: Column<i64>) -> Self {
DynamicColumn::I64(column_i64)
}
}
impl From<Column<u64>> for DynamicColumn {
fn from(column_u64: Column<u64>) -> Self {
DynamicColumn::U64(column_u64)
}
}
impl From<Column<f64>> for DynamicColumn {
fn from(column_f64: Column<f64>) -> Self {
DynamicColumn::F64(column_f64)
}
}
impl From<Column<bool>> for DynamicColumn {
fn from(bool_column: Column<bool>) -> Self {
DynamicColumn::Bool(bool_column)
}
}
impl From<BytesColumn> for DynamicColumn {
fn from(dictionary_encoded_col: BytesColumn) -> Self {
DynamicColumn::Str(dictionary_encoded_col)
}
}
#[derive(Clone)]
pub struct DynamicColumnHandle {
pub(crate) file_slice: FileSlice,
pub(crate) column_type: ColumnType,
}
impl DynamicColumnHandle {
pub fn open(&self) -> io::Result<DynamicColumn> {
let column_bytes: OwnedBytes = self.file_slice.read_bytes()?;
self.open_internal(column_bytes)
}
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)
}
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(column_bytes)?.into(),
ColumnType::Numerical(numerical_type) => match numerical_type {
crate::NumericalType::I64 => {
crate::column::open_column_u64::<i64>(column_bytes)?.into()
}
crate::NumericalType::U64 => {
crate::column::open_column_u64::<u64>(column_bytes)?.into()
}
crate::NumericalType::F64 => {
crate::column::open_column_u64::<f64>(column_bytes)?.into()
}
},
ColumnType::Bool => crate::column::open_column_u64::<bool>(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,86 +1,75 @@
mod column_type_header; #![cfg_attr(all(feature = "unstable", test), feature(test))]
mod dictionary;
mod reader;
pub(crate) mod utils;
mod value;
mod writer;
pub use column_type_header::Cardinality;
pub use reader::ColumnarReader;
pub use value::{NumericalType, NumericalValue};
pub use writer::ColumnarWriter;
pub type DocId = u32;
#[cfg(test)] #[cfg(test)]
mod tests { #[macro_use]
use std::ops::Range; extern crate more_asserts;
use common::file_slice::FileSlice; #[cfg(all(test, feature = "unstable"))]
extern crate test;
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality}; use std::io;
use crate::reader::ColumnarReader;
use crate::value::NumericalValue;
use crate::{Cardinality, ColumnarWriter};
#[test] mod column;
fn test_dataframe_writer_bytes() { mod column_index;
let mut dataframe_writer = ColumnarWriter::default(); mod column_values;
dataframe_writer.record_str(1u32, "my_string", b"hello"); mod columnar;
dataframe_writer.record_str(3u32, "my_string", b"helloeee"); mod dictionary;
let mut buffer: Vec<u8> = Vec::new(); mod dynamic_column;
dataframe_writer.serialize(5, &mut buffer).unwrap(); pub(crate) mod utils;
let columnar_fileslice = FileSlice::from(buffer); mod value;
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] pub use columnar::{ColumnarReader, ColumnarWriter};
fn test_dataframe_writer_bool() { pub use value::{NumericalType, NumericalValue};
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] // pub use self::dynamic_column::DynamicColumnHandle;
fn test_dataframe_writer_numerical() {
let mut dataframe_writer = ColumnarWriter::default(); pub type RowId = u32;
dataframe_writer.record_numerical(1u32, "srical.value", NumericalValue::U64(12u64));
dataframe_writer.record_numerical(2u32, "srical.value", NumericalValue::U64(13u64)); #[derive(Clone, Copy)]
dataframe_writer.record_numerical(4u32, "srical.value", NumericalValue::U64(15u64)); pub struct DateTime {
let mut buffer: Vec<u8> = Vec::new(); timestamp_micros: i64,
dataframe_writer.serialize(5, &mut buffer).unwrap(); }
let columnar_fileslice = FileSlice::from(buffer);
let columnar = ColumnarReader::open(columnar_fileslice).unwrap(); #[derive(Copy, Clone, Debug)]
assert_eq!(columnar.num_columns(), 1); pub struct InvalidData;
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
columnar.read_columns("srical.value").unwrap(); impl From<InvalidData> for io::Error {
assert_eq!(cols.len(), 1); fn from(_: InvalidData) -> Self {
// Right now this 31 bytes are spent as follows io::Error::new(io::ErrorKind::InvalidData, "Invalid data")
//
// - 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);
} }
} }
/// 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)]
mod tests;

View File

@@ -1,102 +0,0 @@
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()
}
}

84
columnar/src/tests.rs Normal file
View File

@@ -0,0 +1,84 @@
use crate::columnar::ColumnType;
use crate::dynamic_column::{DynamicColumn, DynamicColumnHandle};
use crate::value::NumericalValue;
use crate::{Cardinality, ColumnarReader, ColumnarWriter};
#[test]
fn test_dataframe_writer_bytes() {
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(), 165);
}
#[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(), 29);
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_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(), 40);
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() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer.record_str(1, "my.column", "my.key");
columnar_writer.record_str(3, "my.column", "my.key2");
columnar_writer.record_str(3, "my.column2", "different_column!");
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!(); };
assert_eq!(str_col.num_rows(), 5);
// let term_ords = (0..)
}

View File

@@ -1,3 +1,5 @@
use crate::InvalidData;
#[derive(Copy, Clone, Debug, PartialEq)] #[derive(Copy, Clone, Debug, PartialEq)]
pub enum NumericalValue { pub enum NumericalValue {
I64(i64), I64(i64),
@@ -49,12 +51,12 @@ impl NumericalType {
self as u8 self as u8
} }
pub fn try_from_code(code: u8) -> Option<NumericalType> { pub fn try_from_code(code: u8) -> Result<NumericalType, InvalidData> {
match code { match code {
0 => Some(NumericalType::I64), 0 => Ok(NumericalType::I64),
1 => Some(NumericalType::U64), 1 => Ok(NumericalType::U64),
2 => Some(NumericalType::F64), 2 => Ok(NumericalType::F64),
_ => None, _ => Err(InvalidData),
} }
} }
} }
@@ -62,6 +64,7 @@ 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.
/// ///
@@ -111,7 +114,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 Some(numerical_type) = NumericalType::try_from_code(code) { if let Ok(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;
} }

166
common/src/group_by.rs Normal file
View File

@@ -0,0 +1,166 @@
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,10 +6,12 @@ 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

@@ -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::{Field, Schema, FAST, INDEXED, TEXT}; use tantivy::schema::{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: Field, field: String,
} }
impl StatsCollector { impl StatsCollector {
fn with_field(field: Field) -> StatsCollector { fn with_field(field: String) -> 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,7 +171,9 @@ 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) = searcher.search(&query, &StatsCollector::with_field(price))? { if let Some(stats) =
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,15 +1,17 @@
// # Basic Example // # Faceted Search
// //
// This example covers the basic functionalities of // This example covers the faceted search functionalities of
// tantivy. // tantivy.
// //
// We will : // We will :
// - define our schema // - define a text field "name" in our schema
// = create an index in a directory // - define a facet field "classification" in our schema
// - index few documents in our index // - create an index in memory
// - search for the best document matchings "sea whale" // - index few documents with respective facets in our index
// - retrieve the best document original content. // - search and count the number of documents that the classifications start the facet "/Felidae"
// - 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;
@@ -21,7 +23,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("felin_name", TEXT | STORED); let name = schema_builder.add_text_field("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_field, 1960..1970); let docs_in_the_sixties = RangeQuery::new_u64("year".to_string(), 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);

73
examples/ip_field.rs Normal file
View File

@@ -0,0 +1,73 @@
// # 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::{Field, Schema, FAST, TEXT}; use tantivy::schema::{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: Field, field: String,
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: Field, price_fetcher: T) -> Self { pub fn with_product_id_field<T: PriceFetcher>(field: String, 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, "product_id".to_string(),
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.9.0", optional= true} prettytable-rs = {version="0.10.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::iter; use std::ops::RangeInclusive;
use std::sync::Arc; use std::sync::Arc;
use common::OwnedBytes; use common::OwnedBytes;
@@ -71,27 +71,24 @@ mod tests {
}); });
} }
fn get_exp_data() -> Vec<u64> { const FIFTY_PERCENT_RANGE: RangeInclusive<u64> = 1..=50;
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 i in 0..100 { for _ in 0..300_000 {
let num = i * i; let val = rng.gen_range(1..=100);
data.extend(iter::repeat(i as u64).take(num)); data.push(val);
} }
data.shuffle(&mut StdRng::from_seed([1u8; 32])); data.push(SINGLE_ITEM);
// lengt = 328350 data.shuffle(&mut rng);
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<_>>();
@@ -106,15 +103,82 @@ 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 (major_item, _minor_item, data) = get_data_50percent_item(); let 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(
major_item..=major_item, *FIFTY_PERCENT_RANGE.start() as u128..=*FIFTY_PERCENT_RANGE.end() as u128,
0..data.len() as u32, 0..data.len() as u32,
&mut positions, &mut positions,
); );
@@ -124,13 +188,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 (_major_item, minor_item, data) = get_data_50percent_item(); let 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(
minor_item..=minor_item, *SINGLE_ITEM_RANGE.start() as u128..=*SINGLE_ITEM_RANGE.end() as u128,
0..data.len() as u32, 0..data.len() as u32,
&mut positions, &mut positions,
); );
@@ -140,7 +204,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 (_major_item, _minor_item, data) = get_data_50percent_item(); let 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(|| {
@@ -149,6 +213,7 @@ 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,3 +1,4 @@
use std::fmt::{self, Debug};
use std::marker::PhantomData; use std::marker::PhantomData;
use std::ops::{Range, RangeInclusive}; use std::ops::{Range, RangeInclusive};
@@ -6,7 +7,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 = u64>: Send + Sync { pub trait Column<T: PartialOrd + Debug = 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.
@@ -83,7 +84,7 @@ pub struct VecColumn<'a, T = u64> {
max_value: T, max_value: T,
} }
impl<'a, C: Column<T>, T: Copy + PartialOrd> Column<T> for &'a C { impl<'a, C: Column<T>, T: Copy + PartialOrd + fmt::Debug> 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)
} }
@@ -109,7 +110,7 @@ impl<'a, C: Column<T>, T: Copy + PartialOrd> Column<T> for &'a C {
} }
} }
impl<'a, T: Copy + PartialOrd + Send + Sync> Column<T> for VecColumn<'a, T> { impl<'a, T: Copy + PartialOrd + Send + Sync + Debug> 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]
} }
@@ -177,8 +178,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 + Clone, Input: PartialOrd + Send + Sync + Copy + Debug,
Output: PartialOrd + Send + Sync + Clone, Output: PartialOrd + Send + Sync + Copy + Debug,
{ {
MonotonicMappingColumn { MonotonicMappingColumn {
from_column, from_column,
@@ -191,8 +192,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 + Clone, Input: PartialOrd + Send + Sync + Copy + Debug,
Output: PartialOrd + Send + Sync + Clone, Output: PartialOrd + Send + Sync + Copy + Debug,
{ {
#[inline] #[inline]
fn get_val(&self, idx: u32) -> Output { fn get_val(&self, idx: u32) -> Output {
@@ -228,12 +229,15 @@ where
doc_id_range: Range<u32>, doc_id_range: Range<u32>,
positions: &mut Vec<u32>, positions: &mut Vec<u32>,
) { ) {
self.from_column.get_docids_for_value_range( if range.start() > &self.max_value() || range.end() < &self.min_value() {
self.monotonic_mapping.inverse(range.start().clone()) return;
..=self.monotonic_mapping.inverse(range.end().clone()), }
doc_id_range, let range = self.monotonic_mapping.inverse_coerce(range);
positions, if range.start() > range.end() {
) 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,
@@ -254,7 +258,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, T::Item: PartialOrd + fmt::Debug,
{ {
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,6 +453,8 @@ 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;
@@ -706,7 +708,7 @@ mod tests {
); );
} }
fn get_positions_for_value_range_helper<C: Column<T> + ?Sized, T: PartialOrd>( fn get_positions_for_value_range_helper<C: Column<T> + ?Sized, T: PartialOrd + fmt::Debug>(
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>( pub fn open_u128<Item: MonotonicallyMappableToU128 + fmt::Debug>(
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,7 +147,9 @@ pub fn open_u128<Item: MonotonicallyMappableToU128>(
} }
/// 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>(bytes: OwnedBytes) -> io::Result<Arc<dyn Column<T>>> { pub fn open<T: MonotonicallyMappableToU64 + fmt::Debug>(
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)?;
@@ -160,7 +162,7 @@ pub fn open<T: MonotonicallyMappableToU64>(bytes: OwnedBytes) -> io::Result<Arc<
} }
} }
fn open_specific_codec<C: FastFieldCodec, Item: MonotonicallyMappableToU64>( fn open_specific_codec<C: FastFieldCodec, Item: MonotonicallyMappableToU64 + fmt::Debug>(
bytes: OwnedBytes, bytes: OwnedBytes,
header: &Header, header: &Header,
) -> io::Result<Arc<dyn Column<Item>>> { ) -> io::Result<Arc<dyn Column<Item>>> {
@@ -321,6 +323,9 @@ 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"));
@@ -328,6 +333,9 @@ 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,4 +1,6 @@
use std::fmt;
use std::marker::PhantomData; use std::marker::PhantomData;
use std::ops::RangeInclusive;
use fastdivide::DividerU64; use fastdivide::DividerU64;
@@ -6,7 +8,9 @@ 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: 'static + PartialOrd + Copy + Send + Sync { pub trait MonotonicallyMappableToU64:
'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.
@@ -29,11 +33,29 @@ pub trait MonotonicallyMappableToU64: 'static + PartialOrd + Copy + Send + Sync
/// 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, Internal> { pub trait StrictlyMonotonicFn<External: Copy, Internal: Copy> {
/// 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
@@ -54,7 +76,10 @@ impl<T> From<T> for StrictlyMonotonicMappingInverter<T> {
} }
impl<From, To, T> StrictlyMonotonicFn<To, From> for StrictlyMonotonicMappingInverter<T> impl<From, To, T> StrictlyMonotonicFn<To, From> for StrictlyMonotonicMappingInverter<T>
where T: StrictlyMonotonicFn<From, To> where
T: StrictlyMonotonicFn<From, To>,
From: Copy,
To: Copy,
{ {
#[inline(always)] #[inline(always)]
fn mapping(&self, val: To) -> From { fn mapping(&self, val: To) -> From {
@@ -65,6 +90,15 @@ where T: StrictlyMonotonicFn<From, To>
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.
@@ -142,6 +176,31 @@ 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.
@@ -158,6 +217,17 @@ 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
@@ -241,7 +311,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>( fn test_round_trip<T: StrictlyMonotonicFn<K, L>, K: std::fmt::Debug + Eq + Copy, L: Copy>(
mapping: &T, mapping: &T,
test_val: K, test_val: K,
) { ) {

View File

@@ -1,8 +1,11 @@
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: 'static + PartialOrd + Copy + Send + Sync { pub trait MonotonicallyMappableToU128:
'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,15 +31,16 @@ 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 {
if pos_in_bitvec == 63 { let mask = (1u64 << pos_in_bitvec) - 1;
bitvec.count_ones() let masked_bitvec = bitvec & mask;
} else { masked_bitvec.count_ones()
let mask = (1u64 << (pos_in_bitvec + 1)) - 1;
let masked_bitvec = bitvec & mask;
masked_bitvec.count_ones()
}
} }
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
@@ -66,9 +67,7 @@ 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]
@@ -90,8 +89,7 @@ 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) {
// -1 is ok, since idx does exist, so there's at least one Some(index_block.offset + ones_in_block)
Some(index_block.offset + ones_in_block - 1)
} else { } else {
None None
} }
@@ -319,9 +317,10 @@ 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), 1); assert_eq!(count_ones(block, 0), 0);
assert_eq!(count_ones(block, 1), 1); assert_eq!(count_ones(block, 1), 1);
assert_eq!(count_ones(block, 2), 2); assert_eq!(count_ones(block, 2), 1);
assert_eq!(count_ones(block, 3), 2);
} }
} }
@@ -347,11 +346,16 @@ mod bench {
codec codec
} }
fn random_range_iterator(start: u32, end: u32, step_size: u32) -> impl Iterator<Item = u32> { 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 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(1..step_size + 1); current += rng.gen_range(avg_step_size - avg_deviation..=avg_step_size + avg_deviation);
if current >= end { if current >= end {
None None
} else { } else {
@@ -360,10 +364,17 @@ mod bench {
}) })
} }
fn walk_over_data(codec: &DenseCodec, max_step_size: u32) -> Option<u32> { 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: &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, max_step_size), random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, 0),
) )
} }
@@ -379,69 +390,105 @@ mod bench {
} }
#[bench] #[bench]
fn bench_dense_codec_translate_orig_to_codec_90percent_filled_random_stride( fn bench_translate_orig_to_codec_1percent_filled_10percent_hit(bench: &mut Bencher) {
bench: &mut Bencher, let codec = gen_bools(0.01f64);
) {
let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_dense_codec_translate_orig_to_codec_50percent_filled_random_stride( fn bench_translate_orig_to_codec_5percent_filled_10percent_hit(bench: &mut Bencher) {
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); let codec = gen_bools(0.5f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_dense_codec_translate_orig_to_codec_full_scan_10percent(bench: &mut Bencher) { fn bench_translate_orig_to_codec_90percent_filled_1percent_hit(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_dense_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step( fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) {
bench: &mut Bencher, let codec = gen_bools(0.01f64);
) { 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(random_range_iterator(0, num_vals, 50_000)) .translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride( fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
bench: &mut Bencher, let codec = gen_bools(0.01f64);
) { 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(random_range_iterator(0, num_vals, 100)) .translate_codec_idx_to_original_idx(n_percent_step_iterator(10.0, num_non_nulls))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_dense_codec_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) { 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();
bench.iter(|| {
codec
.translate_codec_idx_to_original_idx(0..num_vals)
.last()
});
}
#[bench]
fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) {
let codec = gen_bools(0.90f64);
let num_non_nulls = codec.num_non_nulls();
bench.iter(|| {
codec
.translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
.last()
});
}
#[bench]
fn bench_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, OwnedBytes}; use common::{BitSet, GroupByIteratorExtended, OwnedBytes};
use super::{serialize_dense_codec, DenseCodec}; use super::{serialize_dense_codec, DenseCodec};
@@ -78,12 +78,22 @@ 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)
@@ -207,6 +217,7 @@ 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)]
@@ -273,6 +284,7 @@ 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();
@@ -284,6 +296,7 @@ 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
/// ///
@@ -292,35 +305,41 @@ 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.map(move |codec_idx| { iter.group_by(move |codec_idx| {
// update block_pos to limit search scope block_pos = self.find_block(*codec_idx, block_pos);
block_pos = self.find_block(codec_idx, block_pos); 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 idx_in_block = codec_idx - block.offset(); let offset = 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, idx_in_block block_pos,
indexes_in_block_iter.collect::<Vec<_>>()
) )
} }
SparseCodecBlockVariant::Dense(dense) => { SparseCodecBlockVariant::Dense(dense) => {
dense.translate_codec_idx_to_original_idx(idx_in_block) + block_doc_idx_start Box::new(dense.translate_codec_idx_to_original_idx_iter(indexes_in_block_iter))
as Box<dyn Iterator<Item = u32>>
} }
SparseCodecBlockVariant::Sparse(block) => { SparseCodecBlockVariant::Sparse(block) => {
block.value_at_idx(&self.data, idx_in_block as u16) as u32 + block_doc_idx_start Box::new(indexes_in_block_iter.map(move |idx_in_block| {
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
} }
@@ -595,11 +614,16 @@ mod bench {
codec codec
} }
fn random_range_iterator(start: u32, end: u32, step_size: u32) -> impl Iterator<Item = u32> { 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 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(1..step_size + 1); current += rng.gen_range(avg_step_size - avg_deviation..=avg_step_size + avg_deviation);
if current >= end { if current >= end {
None None
} else { } else {
@@ -608,10 +632,17 @@ mod bench {
}) })
} }
fn walk_over_data(codec: &SparseCodec, max_step_size: u32) -> Option<u32> { 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: &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, max_step_size), random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, 0),
) )
} }
@@ -627,83 +658,83 @@ mod bench {
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_orig_to_codec_1percent_filled_random_stride( fn bench_translate_orig_to_codec_1percent_filled_10percent_hit(bench: &mut Bencher) {
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_sparse_codec_translate_orig_to_codec_5percent_filled_random_stride( fn bench_translate_orig_to_codec_5percent_filled_10percent_hit(bench: &mut Bencher) {
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_sparse_codec_translate_orig_to_codec_full_scan_10percent(bench: &mut Bencher) { 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); 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_sparse_codec_translate_orig_to_codec_full_scan_90percent(bench: &mut Bencher) { 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.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_orig_to_codec_full_scan_1percent(bench: &mut Bencher) { fn bench_translate_orig_to_codec_10percent_filled_1percent_hit(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_sparse_codec_translate_orig_to_codec_90percent_filled_random_stride( fn bench_translate_orig_to_codec_50percent_filled_1percent_hit(bench: &mut Bencher) {
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); let codec = gen_bools(0.9f64);
bench.iter(|| walk_over_data(&codec, 100)); bench.iter(|| walk_over_data(&codec, 100));
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_random_stride_big_step( fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) {
bench: &mut Bencher,
) {
let codec = gen_bools(0.01f64); let codec = gen_bools(0.01f64);
let num_vals = codec.num_non_nulls(); let num_non_nulls = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000)) .translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_random_stride( fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
bench: &mut Bencher,
) {
let codec = gen_bools(0.01f64); let codec = gen_bools(0.01f64);
let num_vals = codec.num_non_nulls(); let num_non_nulls = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100)) .translate_codec_idx_to_original_idx(n_percent_step_iterator(10.0, num_non_nulls))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_full_scan(bench: &mut Bencher) { fn bench_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(|| {
@@ -714,33 +745,18 @@ mod bench {
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step( fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) {
bench: &mut Bencher,
) {
let codec = gen_bools(0.90f64); let codec = gen_bools(0.90f64);
let num_vals = codec.num_non_nulls(); let num_non_nulls = codec.num_non_nulls();
bench.iter(|| { bench.iter(|| {
codec codec
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000)) .translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
.last() .last()
}); });
} }
#[bench] #[bench]
fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_random_stride( fn bench_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
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>( pub fn estimate<T: MonotonicallyMappableToU64 + fmt::Debug>(
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>( pub fn serialize<T: MonotonicallyMappableToU64 + fmt::Debug>(
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>(
} }
/// 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>( pub fn serialize_new<T: MonotonicallyMappableToU64 + fmt::Debug>(
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>( pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default + fmt::Debug>(
column: &[T], column: &[T],
) -> Arc<dyn Column<T>> { ) -> Arc<dyn Column<T>> {
let mut buffer = Vec::new(); let mut buffer = Vec::new();

View File

@@ -51,7 +51,10 @@ 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::{AverageAggregation, StatsAggregation}; use super::metric::{
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
@@ -237,20 +240,38 @@ 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 {
/// Calculates the average. /// Computes the average of the extracted values.
#[serde(rename = "avg")] #[serde(rename = "avg")]
Average(AverageAggregation), Average(AverageAggregation),
/// Calculates stats sum, average, min, max, standard_deviation on a field. /// Counts the number of extracted values.
#[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>) {
match self { let fast_field_name = match self {
MetricAggregation::Average(avg) => fast_field_names.insert(avg.field.to_string()), MetricAggregation::Average(avg) => avg.field_name(),
MetricAggregation::Stats(stats) => fast_field_names.insert(stats.field.to_string()), MetricAggregation::Count(count) => count.field_name(),
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());
} }
} }
@@ -258,6 +279,38 @@ 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,7 +8,10 @@ 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::{AverageAggregation, StatsAggregation}; use super::metric::{
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};
@@ -91,10 +94,7 @@ impl BucketAggregationWithAccessor {
BucketAggregationType::Terms(TermsAggregation { BucketAggregationType::Terms(TermsAggregation {
field: field_name, .. field: field_name, ..
}) => { }) => {
let field = reader let field = reader.schema().get_field(field_name)?;
.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,7 +134,11 @@ 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::Stats(StatsAggregation { field: field_name }) => { | MetricAggregation::Count(CountAggregation { 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)?;
@@ -188,10 +192,7 @@ 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 let field = reader.schema().get_field(field_name)?;
.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) {
@@ -211,10 +212,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) .u64_lenient(field_name)
.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) .u64s_lenient(field_name)
.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,27 +70,51 @@ 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(avg_data) => { IntermediateMetricResult::Average(intermediate_avg) => {
MetricResult::Average(avg_data.finalize().into()) MetricResult::Average(intermediate_avg.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())
}
} }
} }
} }
@@ -100,13 +124,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.
/// ///
@@ -151,7 +175,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
@@ -201,7 +225,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
@@ -237,7 +261,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,9 +548,7 @@ 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 let field = schema.get_field(&histogram_req.field)?;
.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 {
@@ -1366,7 +1364,6 @@ 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

@@ -17,13 +17,15 @@ 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::{IntermediateAverage, IntermediateStats}; use super::metric::{
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.
@@ -204,21 +206,43 @@ 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 {
/// Average containing intermediate average data result /// Intermediate average result.
Average(IntermediateAverage), Average(IntermediateAverage),
/// AverageData variant /// Intermediate count result.
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::Average(collector) => { SegmentMetricResultCollector::Stats(collector) => match collector.collecting_for {
IntermediateMetricResult::Average(IntermediateAverage::from_collector(collector)) super::metric::SegmentStatsType::Average => IntermediateMetricResult::Average(
} IntermediateAverage::from_collector(collector),
SegmentMetricResultCollector::Stats(collector) => { ),
IntermediateMetricResult::Stats(collector.stats) super::metric::SegmentStatsType::Count => {
} 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))
}
},
} }
} }
} }
@@ -229,18 +253,36 @@ 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_data_left), IntermediateMetricResult::Average(avg_left),
IntermediateMetricResult::Average(avg_data_right), IntermediateMetricResult::Average(avg_right),
) => { ) => {
avg_data_left.merge_fruits(avg_data_right); avg_left.merge_fruits(avg_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),
@@ -248,6 +290,9 @@ 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");
} }
@@ -612,9 +657,7 @@ 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 let field = schema.get_field(&range_req.field)?;
.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,13 +1,9 @@
use std::fmt::Debug; use std::fmt::Debug;
use fastfield_codecs::Column;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::aggregation::f64_from_fastfield_u64; use super::{IntermediateStats, SegmentStatsCollector};
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. /// Supported field types are u64, i64, and f64.
@@ -21,94 +17,43 @@ use crate::DocId;
/// } /// }
/// } /// }
/// ``` /// ```
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AverageAggregation { pub struct AverageAggregation {
/// The field name to compute the stats on. /// The field name to compute the average on.
pub field: String, pub field: String,
} }
impl AverageAggregation { impl AverageAggregation {
/// Create new AverageAggregation from a field. /// Creates a new [`AverageAggregation`] instance from a field name.
pub fn from_field_name(field_name: String) -> Self { pub fn from_field_name(field_name: String) -> Self {
AverageAggregation { field: field_name } Self { field: field_name }
} }
/// Return the field name. /// Returns the field name the aggregation is computed on.
pub fn field_name(&self) -> &str { pub fn field_name(&self) -> &str {
&self.field &self.field
} }
} }
#[derive(Clone, PartialEq)] /// Intermediate result of the average aggregation that can be combined with other intermediate
pub(crate) struct SegmentAverageCollector { /// results.
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 {
pub(crate) sum: f64, stats: IntermediateStats,
pub(crate) doc_count: u64,
} }
impl IntermediateAverage { impl IntermediateAverage {
pub(crate) fn from_collector(collector: SegmentAverageCollector) -> Self { /// Creates a new [`IntermediateAverage`] instance from a [`SegmentStatsCollector`].
collector.data pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self {
} 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)
} }
} }
#[inline] /// Merges the other intermediate result into self.
fn collect(&mut self, val: f64) { pub fn merge_fruits(&mut self, other: IntermediateAverage) {
self.doc_count += 1; self.stats.merge_fruits(other.stats);
self.sum += val; }
/// Computes the final average value.
pub fn finalize(&self) -> Option<f64> {
self.stats.finalize().avg
} }
} }

View File

@@ -0,0 +1,59 @@
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.
/// Supported field types are u64, i64, and f64.
/// 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

@@ -0,0 +1,59 @@
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.
/// Supported field types are u64, i64, and f64.
/// 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

@@ -0,0 +1,59 @@
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.
/// Supported field types are u64, i64, and f64.
/// 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,10 +3,18 @@
//! 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.
/// ///
@@ -28,3 +36,61 @@ 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,8 +5,8 @@ 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 stats of numeric values that are /// A multi-value metric aggregation that computes a collection of statistics on numeric values that
/// extracted from the aggregated documents. /// are extracted from the aggregated documents.
/// Supported field types are `u64`, `i64`, and `f64`. /// Supported field types are `u64`, `i64`, and `f64`.
/// See [`Stats`] for returned statistics. /// See [`Stats`] for returned statistics.
/// ///
@@ -26,11 +26,11 @@ pub struct StatsAggregation {
} }
impl StatsAggregation { impl StatsAggregation {
/// Create new StatsAggregation from a field. /// Creates a new [`StatsAggregation`] instance from a field name.
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 }
} }
/// Return the field name. /// Returns the field name the aggregation is computed on.
pub fn field_name(&self) -> &str { pub fn field_name(&self) -> &str {
&self.field &self.field
} }
@@ -40,16 +40,14 @@ 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: usize, pub count: u64,
/// 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 values. `None` for count == 0. /// The average of the fast field values. `None` if count equals zero.
pub avg: Option<f64>, pub avg: Option<f64>,
} }
@@ -58,33 +56,36 @@ 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
))), ))),
} }
} }
} }
/// `IntermediateStats` contains the mergeable version for stats. /// Intermediate result of the stats aggregation that can be combined with other intermediate
/// results.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct IntermediateStats { pub struct IntermediateStats {
count: usize, /// The number of extracted values.
count: u64,
/// The sum of the extracted values.
sum: f64, sum: f64,
squared_sum: f64, /// The min value.
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,33 +93,15 @@ impl Default for IntermediateStats {
} }
impl IntermediateStats { impl IntermediateStats {
pub(crate) fn avg(&self) -> Option<f64> { /// Merges the other stats intermediate result into self.
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);
} }
/// compute final resultimprove_docs /// Computes the final stats value.
pub fn finalize(&self) -> Stats { pub fn finalize(&self) -> Stats {
let min = if self.count == 0 { let min = if self.count == 0 {
None None
@@ -130,13 +113,17 @@ 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: self.avg(), avg,
} }
} }
@@ -144,22 +131,33 @@ 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 {
pub(crate) stats: IntermediateStats,
field_type: Type, field_type: Type,
pub(crate) collecting_for: SegmentStatsType,
pub(crate) stats: IntermediateStats,
} }
impl SegmentStatsCollector { impl SegmentStatsCollector {
pub fn from_req(field_type: Type) -> Self { pub fn from_req(field_type: Type, collecting_for: SegmentStatsType) -> Self {
Self { Self {
field_type, field_type,
collecting_for,
stats: IntermediateStats::default(), stats: IntermediateStats::default(),
} }
} }
@@ -236,7 +234,6 @@ 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
}) })
); );
@@ -313,7 +310,6 @@ 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
}) })
); );
@@ -325,7 +321,6 @@ 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
}) })
); );
@@ -337,7 +332,6 @@ 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
}) })
); );
@@ -349,7 +343,6 @@ 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
}) })
); );
@@ -361,7 +354,6 @@ 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

@@ -0,0 +1,59 @@
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.
/// Supported field types are u64, i64, and f64.
/// 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

@@ -216,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![]; let mut data = Vec::with_capacity(entries.len());
let mut data_names = vec![]; let mut data_names = Vec::with_capacity(entries.len());
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);
@@ -1208,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 as f64, score_field_f64 => val,
score_field_i64 => val as i64, score_field_i64 => val as i64,
))?; ))?;
} }
@@ -1250,10 +1250,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap()
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1281,10 +1278,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap()
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1312,10 +1306,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap()
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1351,10 +1342,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap()
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1380,10 +1368,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap()
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1409,10 +1394,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap()
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1446,10 +1428,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap()
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1481,10 +1460,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap()
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1520,10 +1496,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap()
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1550,10 +1523,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap()
searcher.search(&AllQuery, &collector).unwrap().into();
agg_res
}); });
} }
@@ -1597,7 +1567,7 @@ mod tests {
], ],
..Default::default() ..Default::default()
}), }),
sub_aggregation: sub_agg_req_1.clone(), sub_aggregation: sub_agg_req_1,
}), }),
), ),
] ]
@@ -1607,10 +1577,7 @@ 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();
let agg_res: AggregationResults = searcher.search(&term_query, &collector).unwrap()
searcher.search(&term_query, &collector).unwrap().into();
agg_res
}); });
} }
} }

View File

@@ -15,7 +15,8 @@ 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, SegmentAverageCollector, SegmentStatsCollector, StatsAggregation, AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, SegmentStatsCollector,
SegmentStatsType, StatsAggregation, SumAggregation,
}; };
use super::VecWithNames; use super::VecWithNames;
use crate::aggregation::agg_req::BucketAggregationType; use crate::aggregation::agg_req::BucketAggregationType;
@@ -163,30 +164,46 @@ 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 { field: _ }) => { MetricAggregation::Average(AverageAggregation { .. }) => {
Ok(SegmentMetricResultCollector::Average( Ok(SegmentMetricResultCollector::Stats(
SegmentAverageCollector::from_req(req.field_type), SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Average),
)) ))
} }
MetricAggregation::Stats(StatsAggregation { field: _ }) => { MetricAggregation::Count(CountAggregation { .. }) => {
Ok(SegmentMetricResultCollector::Stats( Ok(SegmentMetricResultCollector::Stats(
SegmentStatsCollector::from_req(req.field_type), SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Count),
))
}
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(self.field)?; .typed_fast_field_reader(schema.get_field_name(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::{Field, Type}; use crate::schema::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: Field, field: String,
} }
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: Field, field: String,
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();
let val_field = schema_builder.add_u64_field("val_field", FAST); 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, 0u64, 2, 5); let histogram_collector = HistogramCollector::new("val_field".to_string(), 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,7 +238,8 @@ 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 = HistogramCollector::new(val_field, -20i64, 10u64, 4); let histogram_collector =
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(())
@@ -262,7 +263,8 @@ 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 = HistogramCollector::new(val_field, -20i64, 10u64, 4); let histogram_collector =
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(())
@@ -285,7 +287,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, "date_field".to_string(),
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: Field, field: String,
} }
pub struct FastFieldSegmentCollector { pub struct FastFieldSegmentCollector {
@@ -164,7 +164,7 @@ pub struct FastFieldSegmentCollector {
} }
impl FastFieldTestCollector { impl FastFieldTestCollector {
pub fn for_field(field: Field) -> FastFieldTestCollector { pub fn for_field(field: String) -> 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,7 +238,9 @@ 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.fast_fields().bytes(self.field)?; let reader = segment_reader
.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(self.field)?; .typed_fast_field_reader(segment_reader.schema().get_field_name(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).ok_or_else(|| { let schema_field = schema.get_field(&sort_by_field.field).map_err(|_| {
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

@@ -198,11 +198,10 @@ impl Searcher {
collector: &C, collector: &C,
executor: &Executor, executor: &Executor,
) -> crate::Result<C::Fruit> { ) -> crate::Result<C::Fruit> {
let scoring_enabled = collector.requires_scoring(); let enabled_scoring = if collector.requires_scoring() {
let enabled_scoring = if scoring_enabled { EnableScoring::enabled_from_searcher(self)
EnableScoring::Enabled(self)
} else { } else {
EnableScoring::Disabled(self.schema()) EnableScoring::disabled_from_searcher(self)
}; };
let weight = query.weight(enabled_scoring)?; let weight = query.weight(enabled_scoring)?;
let segment_readers = self.segment_readers(); let segment_readers = self.segment_readers();

View File

@@ -95,7 +95,8 @@ impl SegmentReader {
match field_entry.field_type() { match field_entry.field_type() {
FieldType::Facet(_) => { FieldType::Facet(_) => {
let term_ords_reader = self.fast_fields().u64s(field)?; let term_ords_reader =
self.fast_fields().u64s(self.schema.get_field_name(field))?;
let termdict = self let termdict = self
.termdict_composite .termdict_composite
.open_read(field) .open_read(field)

View File

@@ -32,7 +32,7 @@ impl LockError {
/// Error that may occur when opening a directory /// Error that may occur when opening a directory
#[derive(Debug, Clone, Error)] #[derive(Debug, Clone, Error)]
pub enum OpenDirectoryError { pub enum OpenDirectoryError {
/// The underlying directory does not exists. /// The underlying directory does not exist.
#[error("Directory does not exist: '{0}'.")] #[error("Directory does not exist: '{0}'.")]
DoesNotExist(PathBuf), DoesNotExist(PathBuf),
/// The path exists but is not a directory. /// The path exists but is not a directory.
@@ -151,8 +151,8 @@ impl fmt::Debug for Incompatibility {
/// Error that may occur when accessing a file read /// Error that may occur when accessing a file read
#[derive(Debug, Clone, Error)] #[derive(Debug, Clone, Error)]
pub enum OpenReadError { pub enum OpenReadError {
/// The file does not exists. /// The file does not exist.
#[error("Files does not exists: {0:?}")] #[error("Files does not exist: {0:?}")]
FileDoesNotExist(PathBuf), FileDoesNotExist(PathBuf),
/// Any kind of io::Error. /// Any kind of io::Error.
#[error( #[error(
@@ -181,8 +181,8 @@ impl OpenReadError {
/// Error that may occur when trying to delete a file /// Error that may occur when trying to delete a file
#[derive(Debug, Clone, Error)] #[derive(Debug, Clone, Error)]
pub enum DeleteError { pub enum DeleteError {
/// The file does not exists. /// The file does not exist.
#[error("File does not exists: '{0}'.")] #[error("File does not exist: '{0}'.")]
FileDoesNotExist(PathBuf), FileDoesNotExist(PathBuf),
/// Any kind of IO error that happens when /// Any kind of IO error that happens when
/// interacting with the underlying IO device. /// interacting with the underlying IO device.

View File

@@ -232,7 +232,7 @@ impl Directory for RamDirectory {
let path_buf = PathBuf::from(path); let path_buf = PathBuf::from(path);
self.fs.write().unwrap().write(path_buf, data); self.fs.write().unwrap().write(path_buf, data);
if path == *META_FILEPATH { if path == *META_FILEPATH {
let _ = self.fs.write().unwrap().watch_router.broadcast(); drop(self.fs.write().unwrap().watch_router.broadcast());
} }
Ok(()) Ok(())
} }

View File

@@ -168,7 +168,7 @@ mod tests {
watch_event_router.broadcast().wait().unwrap(); watch_event_router.broadcast().wait().unwrap();
assert_eq!(2, counter.load(Ordering::SeqCst)); assert_eq!(2, counter.load(Ordering::SeqCst));
mem::drop(handle_a); mem::drop(handle_a);
let _ = watch_event_router.broadcast(); drop(watch_event_router.broadcast());
watch_event_router.broadcast().wait().unwrap(); watch_event_router.broadcast().wait().unwrap();
assert_eq!(2, counter.load(Ordering::SeqCst)); assert_eq!(2, counter.load(Ordering::SeqCst));
} }

View File

@@ -175,7 +175,7 @@ mod bench {
fn get_alive() -> Vec<u32> { fn get_alive() -> Vec<u32> {
let mut data = (0..1_000_000_u32).collect::<Vec<u32>>(); let mut data = (0..1_000_000_u32).collect::<Vec<u32>>();
for _ in 0..(1_000_000) * 1 / 8 { for _ in 0..1_000_000 / 8 {
remove_rand(&mut data); remove_rand(&mut data);
} }
data data

View File

@@ -25,7 +25,7 @@ mod tests {
index_writer.commit()?; index_writer.commit()?;
let searcher = index.reader()?.searcher(); let searcher = index.reader()?.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let bytes_reader = segment_reader.fast_fields().bytes(bytes_field).unwrap(); let bytes_reader = segment_reader.fast_fields().bytes("bytesfield").unwrap();
assert_eq!(bytes_reader.get_bytes(0), &[0u8, 1, 2, 3]); assert_eq!(bytes_reader.get_bytes(0), &[0u8, 1, 2, 3]);
assert!(bytes_reader.get_bytes(1).is_empty()); assert!(bytes_reader.get_bytes(1).is_empty());
assert_eq!(bytes_reader.get_bytes(2), &[255u8]); assert_eq!(bytes_reader.get_bytes(2), &[255u8]);
@@ -96,7 +96,7 @@ mod tests {
let term = Term::from_field_bytes(field, b"lucene".as_ref()); let term = Term::from_field_bytes(field, b"lucene".as_ref());
let term_query = TermQuery::new(term, IndexRecordOption::Basic); let term_query = TermQuery::new(term, IndexRecordOption::Basic);
let term_weight_err = let term_weight_err =
term_query.specialized_weight(EnableScoring::Disabled(searcher.schema())); term_query.specialized_weight(EnableScoring::disabled_from_schema(searcher.schema()));
assert!(matches!( assert!(matches!(
term_weight_err, term_weight_err,
Err(crate::TantivyError::SchemaError(_)) Err(crate::TantivyError::SchemaError(_))
@@ -109,8 +109,7 @@ mod tests {
let searcher = create_index_for_test(FAST)?; let searcher = create_index_for_test(FAST)?;
assert_eq!(searcher.num_docs(), 1); assert_eq!(searcher.num_docs(), 1);
let fast_fields = searcher.segment_reader(0u32).fast_fields(); let fast_fields = searcher.segment_reader(0u32).fast_fields();
let field = searcher.schema().get_field("string_bytes").unwrap(); let fast_field_reader = fast_fields.bytes("string_bytes").unwrap();
let fast_field_reader = fast_fields.bytes(field).unwrap();
assert_eq!(fast_field_reader.get_bytes(0u32), b"tantivy"); assert_eq!(fast_field_reader.get_bytes(0u32), b"tantivy");
Ok(()) Ok(())
} }

View File

@@ -12,13 +12,15 @@
//! //!
//! //!
//! Fields have to be declared as `FAST` in the schema. //! Fields have to be declared as `FAST` in the schema.
//! Currently supported fields are: u64, i64, f64, bytes and text. //! Currently supported fields are: u64, i64, f64, bytes, ip and text.
//! //!
//! Fast fields are stored in with [different codecs](fastfield_codecs). The best codec is detected //! Fast fields are stored in with [different codecs](fastfield_codecs). The best codec is detected
//! automatically, when serializing. //! automatically, when serializing.
//! //!
//! Read access performance is comparable to that of an array lookup. //! Read access performance is comparable to that of an array lookup.
use std::net::Ipv6Addr;
use fastfield_codecs::MonotonicallyMappableToU64; use fastfield_codecs::MonotonicallyMappableToU64;
pub use self::alive_bitset::{intersect_alive_bitsets, write_alive_bitset, AliveBitSet}; pub use self::alive_bitset::{intersect_alive_bitsets, write_alive_bitset, AliveBitSet};
@@ -28,7 +30,7 @@ pub use self::facet_reader::FacetReader;
pub(crate) use self::multivalued::{get_fastfield_codecs_for_multivalue, MultivalueStartIndex}; pub(crate) use self::multivalued::{get_fastfield_codecs_for_multivalue, MultivalueStartIndex};
pub use self::multivalued::{ pub use self::multivalued::{
MultiValueIndex, MultiValueU128FastFieldWriter, MultiValuedFastFieldReader, MultiValueIndex, MultiValueU128FastFieldWriter, MultiValuedFastFieldReader,
MultiValuedFastFieldWriter, MultiValuedU128FastFieldReader, MultiValuedFastFieldWriter,
}; };
pub(crate) use self::readers::type_and_cardinality; pub(crate) use self::readers::type_and_cardinality;
pub use self::readers::FastFieldReaders; pub use self::readers::FastFieldReaders;
@@ -47,6 +49,33 @@ mod readers;
mod serializer; mod serializer;
mod writer; mod writer;
/// Trait for types that provide a zero value.
///
/// The resulting value is never used, just as placeholder, e.g. for `vec.resize()`.
pub trait MakeZero {
/// Build a default value. This default value is never used, so the value does not
/// really matter.
fn make_zero() -> Self;
}
impl<T: FastValue> MakeZero for T {
fn make_zero() -> Self {
T::from_u64(0)
}
}
impl MakeZero for u128 {
fn make_zero() -> Self {
0
}
}
impl MakeZero for Ipv6Addr {
fn make_zero() -> Self {
Ipv6Addr::from(0u128.to_be_bytes())
}
}
/// Trait for types that are allowed for fast fields: /// Trait for types that are allowed for fast fields:
/// (u64, i64 and f64, bool, DateTime). /// (u64, i64 and f64, bool, DateTime).
pub trait FastValue: pub trait FastValue:
@@ -54,12 +83,6 @@ pub trait FastValue:
{ {
/// Returns the `schema::Type` for this FastValue. /// Returns the `schema::Type` for this FastValue.
fn to_type() -> Type; fn to_type() -> Type;
/// Build a default value. This default value is never used, so the value does not
/// really matter.
fn make_zero() -> Self {
Self::from_u64(0u64)
}
} }
impl FastValue for u64 { impl FastValue for u64 {
@@ -101,12 +124,6 @@ impl FastValue for DateTime {
fn to_type() -> Type { fn to_type() -> Type {
Type::Date Type::Date
} }
fn make_zero() -> Self {
DateTime {
timestamp_micros: 0,
}
}
} }
fn value_to_u64(value: &Value) -> crate::Result<u64> { fn value_to_u64(value: &Value) -> crate::Result<u64> {
@@ -145,7 +162,7 @@ impl FastFieldType {
mod tests { mod tests {
use std::collections::HashMap; use std::collections::HashMap;
use std::ops::Range; use std::ops::{Range, RangeInclusive};
use std::path::Path; use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
@@ -159,7 +176,9 @@ mod tests {
use super::*; use super::*;
use crate::directory::{CompositeFile, Directory, RamDirectory, WritePtr}; use crate::directory::{CompositeFile, Directory, RamDirectory, WritePtr};
use crate::merge_policy::NoMergePolicy; use crate::merge_policy::NoMergePolicy;
use crate::schema::{Cardinality, Document, Field, Schema, SchemaBuilder, FAST, STRING, TEXT}; use crate::schema::{
Cardinality, Document, Field, Schema, SchemaBuilder, FAST, INDEXED, STRING, TEXT,
};
use crate::time::OffsetDateTime; use crate::time::OffsetDateTime;
use crate::{DateOptions, DatePrecision, Index, SegmentId, SegmentReader}; use crate::{DateOptions, DatePrecision, Index, SegmentId, SegmentReader};
@@ -520,11 +539,6 @@ mod tests {
Ok(()) Ok(())
} }
#[test]
fn test_default_date() {
assert_eq!(0, DateTime::make_zero().into_timestamp_secs());
}
fn get_vals_for_docs(ff: &MultiValuedFastFieldReader<u64>, docs: Range<u32>) -> Vec<u64> { fn get_vals_for_docs(ff: &MultiValuedFastFieldReader<u64>, docs: Range<u32>) -> Vec<u64> {
let mut all = vec![]; let mut all = vec![];
@@ -569,7 +583,7 @@ mod tests {
assert_eq!(searcher.segment_readers().len(), 1); assert_eq!(searcher.segment_readers().len(), 1);
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let text_fast_field = fast_fields.u64s(text_field).unwrap(); let text_fast_field = fast_fields.u64s("text").unwrap();
assert_eq!( assert_eq!(
get_vals_for_docs(&text_fast_field, 0..5), get_vals_for_docs(&text_fast_field, 0..5),
@@ -608,7 +622,7 @@ mod tests {
assert_eq!(searcher.segment_readers().len(), 2); assert_eq!(searcher.segment_readers().len(), 2);
let segment_reader = searcher.segment_reader(1); let segment_reader = searcher.segment_reader(1);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let text_fast_field = fast_fields.u64s(text_field).unwrap(); let text_fast_field = fast_fields.u64s("text").unwrap();
assert_eq!(get_vals_for_docs(&text_fast_field, 0..3), vec![0, 1, 0]); assert_eq!(get_vals_for_docs(&text_fast_field, 0..3), vec![0, 1, 0]);
} }
@@ -624,7 +638,7 @@ mod tests {
let searcher = reader.searcher(); let searcher = reader.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let text_fast_field = fast_fields.u64s(text_field).unwrap(); let text_fast_field = fast_fields.u64s("text").unwrap();
assert_eq!( assert_eq!(
get_vals_for_docs(&text_fast_field, 0..8), get_vals_for_docs(&text_fast_field, 0..8),
@@ -667,7 +681,7 @@ mod tests {
assert_eq!(searcher.segment_readers().len(), 1); assert_eq!(searcher.segment_readers().len(), 1);
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let text_fast_field = fast_fields.u64s(text_field).unwrap(); let text_fast_field = fast_fields.u64s("text").unwrap();
assert_eq!(get_vals_for_docs(&text_fast_field, 0..6), vec![1, 0, 0, 2]); assert_eq!(get_vals_for_docs(&text_fast_field, 0..6), vec![1, 0, 0, 2]);
@@ -698,7 +712,7 @@ mod tests {
assert_eq!(searcher.segment_readers().len(), 2); assert_eq!(searcher.segment_readers().len(), 2);
let segment_reader = searcher.segment_reader(1); let segment_reader = searcher.segment_reader(1);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let text_fast_field = fast_fields.u64s(text_field).unwrap(); let text_fast_field = fast_fields.u64s("text").unwrap();
assert_eq!(get_vals_for_docs(&text_fast_field, 0..2), vec![0, 1]); assert_eq!(get_vals_for_docs(&text_fast_field, 0..2), vec![0, 1]);
} }
@@ -714,7 +728,7 @@ mod tests {
let searcher = reader.searcher(); let searcher = reader.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let text_fast_field = fast_fields.u64s(text_field).unwrap(); let text_fast_field = fast_fields.u64s("text").unwrap();
assert_eq!( assert_eq!(
get_vals_for_docs(&text_fast_field, 0..9), get_vals_for_docs(&text_fast_field, 0..9),
@@ -759,8 +773,8 @@ mod tests {
assert_eq!(searcher.segment_readers().len(), 1); assert_eq!(searcher.segment_readers().len(), 1);
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let fast_fields = segment_reader.fast_fields(); let fast_fields = segment_reader.fast_fields();
let date_fast_field = fast_fields.date(date_field).unwrap(); let date_fast_field = fast_fields.date("date").unwrap();
let dates_fast_field = fast_fields.dates(multi_date_field).unwrap(); let dates_fast_field = fast_fields.dates("multi_date").unwrap();
let mut dates = vec![]; let mut dates = vec![];
{ {
assert_eq!(date_fast_field.get_val(0).into_timestamp_micros(), 1i64); assert_eq!(date_fast_field.get_val(0).into_timestamp_micros(), 1i64);
@@ -969,4 +983,117 @@ mod tests {
} }
Ok(len) Ok(len)
} }
#[test]
fn test_gcd_bug_regression_1757() {
let mut schema_builder = Schema::builder();
let num_field = schema_builder.add_u64_field("url_norm_hash", FAST | INDEXED);
let schema = schema_builder.build();
let index = Index::create_in_ram(schema);
{
let mut writer = index.writer_for_tests().unwrap();
writer
.add_document(doc! {
num_field => 100u64,
})
.unwrap();
writer
.add_document(doc! {
num_field => 200u64,
})
.unwrap();
writer
.add_document(doc! {
num_field => 300u64,
})
.unwrap();
writer.commit().unwrap();
}
let reader = index.reader().unwrap();
let searcher = reader.searcher();
let segment = &searcher.segment_readers()[0];
let field = segment.fast_fields().u64(num_field).unwrap();
let numbers = vec![100, 200, 300];
let test_range = |range: RangeInclusive<u64>| {
let expexted_count = numbers.iter().filter(|num| range.contains(num)).count();
let mut vec = vec![];
field.get_docids_for_value_range(range, 0..u32::MAX, &mut vec);
assert_eq!(vec.len(), expexted_count);
};
test_range(50..=50);
test_range(150..=150);
test_range(350..=350);
test_range(100..=250);
test_range(101..=200);
test_range(101..=199);
test_range(100..=300);
test_range(100..=299);
}
#[test]
fn test_mapping_bug_docids_for_value_range() {
let mut schema_builder = Schema::builder();
let num_field = schema_builder.add_u64_field("url_norm_hash", FAST | INDEXED);
let schema = schema_builder.build();
let index = Index::create_in_ram(schema);
{
// Values without gcd, but with min_value
let mut writer = index.writer_for_tests().unwrap();
writer
.add_document(doc! {
num_field => 1000u64,
})
.unwrap();
writer
.add_document(doc! {
num_field => 1001u64,
})
.unwrap();
writer
.add_document(doc! {
num_field => 1003u64,
})
.unwrap();
writer.commit().unwrap();
}
let reader = index.reader().unwrap();
let searcher = reader.searcher();
let segment = &searcher.segment_readers()[0];
let field = segment.fast_fields().u64(num_field).unwrap();
let numbers = vec![1000, 1001, 1003];
let test_range = |range: RangeInclusive<u64>| {
let expexted_count = numbers.iter().filter(|num| range.contains(num)).count();
let mut vec = vec![];
field.get_docids_for_value_range(range, 0..u32::MAX, &mut vec);
assert_eq!(vec.len(), expexted_count);
};
let test_range_variant = |start, stop| {
let start_range = start..=stop;
test_range(start_range);
let start_range = start..=(stop - 1);
test_range(start_range);
let start_range = start..=(stop + 1);
test_range(start_range);
let start_range = (start - 1)..=stop;
test_range(start_range);
let start_range = (start - 1)..=(stop - 1);
test_range(start_range);
let start_range = (start - 1)..=(stop + 1);
test_range(start_range);
let start_range = (start + 1)..=stop;
test_range(start_range);
let start_range = (start + 1)..=(stop - 1);
test_range(start_range);
let start_range = (start + 1)..=(stop + 1);
test_range(start_range);
};
test_range_variant(50, 50);
test_range_variant(1000, 1000);
test_range_variant(1000, 1002);
}
} }

View File

@@ -5,7 +5,7 @@ mod writer;
use fastfield_codecs::FastFieldCodecType; use fastfield_codecs::FastFieldCodecType;
pub use index::MultiValueIndex; pub use index::MultiValueIndex;
pub use self::reader::{MultiValuedFastFieldReader, MultiValuedU128FastFieldReader}; pub use self::reader::MultiValuedFastFieldReader;
pub(crate) use self::writer::MultivalueStartIndex; pub(crate) use self::writer::MultivalueStartIndex;
pub use self::writer::{MultiValueU128FastFieldWriter, MultiValuedFastFieldWriter}; pub use self::writer::{MultiValueU128FastFieldWriter, MultiValuedFastFieldWriter};
@@ -52,7 +52,7 @@ mod tests {
let searcher = index.reader()?.searcher(); let searcher = index.reader()?.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let mut vals = Vec::new(); let mut vals = Vec::new();
let multi_value_reader = segment_reader.fast_fields().u64s(field)?; let multi_value_reader = segment_reader.fast_fields().u64s("multifield")?;
{ {
multi_value_reader.get_vals(2, &mut vals); multi_value_reader.get_vals(2, &mut vals);
assert_eq!(&vals, &[4u64]); assert_eq!(&vals, &[4u64]);
@@ -229,7 +229,7 @@ mod tests {
let searcher = index.reader()?.searcher(); let searcher = index.reader()?.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let mut vals = Vec::new(); let mut vals = Vec::new();
let multi_value_reader = segment_reader.fast_fields().i64s(field).unwrap(); let multi_value_reader = segment_reader.fast_fields().i64s("multifield").unwrap();
multi_value_reader.get_vals(2, &mut vals); multi_value_reader.get_vals(2, &mut vals);
assert_eq!(&vals, &[-4i64]); assert_eq!(&vals, &[-4i64]);
multi_value_reader.get_vals(0, &mut vals); multi_value_reader.get_vals(0, &mut vals);
@@ -261,7 +261,7 @@ mod tests {
let searcher = index.reader()?.searcher(); let searcher = index.reader()?.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let mut vals = Vec::new(); let mut vals = Vec::new();
let multi_value_reader = segment_reader.fast_fields().bools(bool_field).unwrap(); let multi_value_reader = segment_reader.fast_fields().bools("multifield").unwrap();
multi_value_reader.get_vals(2, &mut vals); multi_value_reader.get_vals(2, &mut vals);
assert_eq!(&vals, &[false]); assert_eq!(&vals, &[false]);
multi_value_reader.get_vals(0, &mut vals); multi_value_reader.get_vals(0, &mut vals);
@@ -525,7 +525,7 @@ mod bench {
serializer.close().unwrap(); serializer.close().unwrap();
field field
}; };
let file = directory.open_read(&path).unwrap(); let file = directory.open_read(path).unwrap();
{ {
let fast_fields_composite = CompositeFile::open(&file).unwrap(); let fast_fields_composite = CompositeFile::open(&file).unwrap();
let data_idx = fast_fields_composite let data_idx = fast_fields_composite

View File

@@ -1,107 +1,31 @@
use core::fmt;
use std::ops::{Range, RangeInclusive}; use std::ops::{Range, RangeInclusive};
use std::sync::Arc; use std::sync::Arc;
use fastfield_codecs::{Column, MonotonicallyMappableToU128}; use fastfield_codecs::Column;
use super::MultiValueIndex; use super::MultiValueIndex;
use crate::fastfield::FastValue; use crate::fastfield::MakeZero;
use crate::DocId; use crate::DocId;
/// Reader for a multivalued `u64` fast field. /// Reader for a multivalued fast field.
/// ///
/// The reader is implemented as two `u64` fast field. /// The reader is implemented as two fast fields, one u64 fast field for the index and one for the
/// values.
/// ///
/// The `vals_reader` will access the concatenated list of all /// The `vals_reader` will access the concatenated list of all values.
/// values for all reader. /// The `idx_reader` associates, for each document, the index of its first value.
/// The `idx_reader` associated, for each document, the index of its first value.
/// Stores the start position for each document.
#[derive(Clone)] #[derive(Clone)]
pub struct MultiValuedFastFieldReader<Item: FastValue> { pub struct MultiValuedFastFieldReader<T> {
idx_reader: MultiValueIndex,
vals_reader: Arc<dyn Column<Item>>,
}
impl<Item: FastValue> MultiValuedFastFieldReader<Item> {
pub(crate) fn open(
idx_reader: Arc<dyn Column<u64>>,
vals_reader: Arc<dyn Column<Item>>,
) -> MultiValuedFastFieldReader<Item> {
MultiValuedFastFieldReader {
idx_reader: MultiValueIndex::new(idx_reader),
vals_reader,
}
}
/// Returns the array of values associated with the given `doc`.
#[inline]
fn get_vals_for_range(&self, range: Range<u32>, vals: &mut Vec<Item>) {
let len = (range.end - range.start) as usize;
vals.resize(len, Item::make_zero());
self.vals_reader
.get_range(range.start as u64, &mut vals[..]);
}
/// Returns the array of values associated with the given `doc`.
#[inline]
pub fn get_vals(&self, doc: DocId, vals: &mut Vec<Item>) {
let range = self.idx_reader.range(doc);
self.get_vals_for_range(range, vals);
}
/// returns the multivalue index
pub fn get_index_reader(&self) -> &MultiValueIndex {
&self.idx_reader
}
/// Returns the minimum value for this fast field.
///
/// The min value does not take in account of possible
/// deleted document, and should be considered as a lower bound
/// of the actual minimum value.
pub fn min_value(&self) -> Item {
self.vals_reader.min_value()
}
/// Returns the maximum value for this fast field.
///
/// The max value does not take in account of possible
/// deleted document, and should be considered as an upper bound
/// of the actual maximum value.
pub fn max_value(&self) -> Item {
self.vals_reader.max_value()
}
/// Returns the number of values associated with the document `DocId`.
#[inline]
pub fn num_vals(&self, doc: DocId) -> u32 {
self.idx_reader.num_vals_for_doc(doc)
}
/// Returns the overall number of values in this field.
#[inline]
pub fn total_num_vals(&self) -> u32 {
self.idx_reader.total_num_vals()
}
}
/// Reader for a multivalued `u128` fast field.
///
/// The reader is implemented as a `u64` fast field for the index and a `u128` fast field.
///
/// The `vals_reader` will access the concatenated list of all
/// values for all reader.
/// The `idx_reader` associated, for each document, the index of its first value.
#[derive(Clone)]
pub struct MultiValuedU128FastFieldReader<T: MonotonicallyMappableToU128> {
idx_reader: MultiValueIndex, idx_reader: MultiValueIndex,
vals_reader: Arc<dyn Column<T>>, vals_reader: Arc<dyn Column<T>>,
} }
impl<T: MonotonicallyMappableToU128> MultiValuedU128FastFieldReader<T> { impl<T: PartialOrd + MakeZero + Copy + fmt::Debug> MultiValuedFastFieldReader<T> {
pub(crate) fn open( pub(crate) fn open(
idx_reader: Arc<dyn Column<u64>>, idx_reader: Arc<dyn Column<u64>>,
vals_reader: Arc<dyn Column<T>>, vals_reader: Arc<dyn Column<T>>,
) -> MultiValuedU128FastFieldReader<T> { ) -> MultiValuedFastFieldReader<T> {
Self { Self {
idx_reader: MultiValueIndex::new(idx_reader), idx_reader: MultiValueIndex::new(idx_reader),
vals_reader, vals_reader,
@@ -122,7 +46,7 @@ impl<T: MonotonicallyMappableToU128> MultiValuedU128FastFieldReader<T> {
#[inline] #[inline]
fn get_vals_for_range(&self, range: Range<u32>, vals: &mut Vec<T>) { fn get_vals_for_range(&self, range: Range<u32>, vals: &mut Vec<T>) {
let len = (range.end - range.start) as usize; let len = (range.end - range.start) as usize;
vals.resize(len, T::from_u128(0)); vals.resize(len, T::make_zero());
self.vals_reader self.vals_reader
.get_range(range.start as u64, &mut vals[..]); .get_range(range.start as u64, &mut vals[..]);
} }
@@ -199,8 +123,131 @@ impl<T: MonotonicallyMappableToU128> MultiValuedU128FastFieldReader<T> {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use time::{Duration, OffsetDateTime};
use crate::collector::Count;
use crate::core::Index; use crate::core::Index;
use crate::query::RangeQuery;
use crate::schema::{Cardinality, Facet, FacetOptions, NumericOptions, Schema}; use crate::schema::{Cardinality, Facet, FacetOptions, NumericOptions, Schema};
use crate::{DateOptions, DatePrecision, DateTime};
#[test]
fn test_multivalued_date_docids_for_value_range_1() -> crate::Result<()> {
let mut schema_builder = Schema::builder();
let date_field = schema_builder.add_date_field(
"multi_date_field",
DateOptions::default()
.set_fast(Cardinality::MultiValues)
.set_indexed()
.set_fieldnorm()
.set_precision(DatePrecision::Microseconds)
.set_stored(),
);
let schema = schema_builder.build();
let index = Index::create_in_ram(schema);
let mut index_writer = index.writer_for_tests()?;
let first_time_stamp = OffsetDateTime::now_utc();
index_writer.add_document(doc!(
date_field => DateTime::from_utc(first_time_stamp),
date_field => DateTime::from_utc(first_time_stamp),
))?;
// add another second
let two_secs_ahead = first_time_stamp + Duration::seconds(2);
index_writer.commit()?;
let reader = index.reader()?;
let searcher = reader.searcher();
let reader = searcher.segment_reader(0);
let date_ff_reader = reader.fast_fields().dates(date_field).unwrap();
let mut docids = vec![];
date_ff_reader.get_docids_for_value_range(
DateTime::from_utc(first_time_stamp)..=DateTime::from_utc(two_secs_ahead),
0..5,
&mut docids,
);
assert_eq!(docids, vec![0]);
let count_multiples =
|range_query: RangeQuery| searcher.search(&range_query, &Count).unwrap();
assert_eq!(
count_multiples(RangeQuery::new_date(
date_field,
DateTime::from_utc(first_time_stamp)..DateTime::from_utc(two_secs_ahead)
)),
1
);
Ok(())
}
#[test]
fn test_multivalued_date_docids_for_value_range_2() -> crate::Result<()> {
let mut schema_builder = Schema::builder();
let date_field = schema_builder.add_date_field(
"multi_date_field",
DateOptions::default()
.set_fast(Cardinality::MultiValues)
// TODO: Test different precision after fixing https://github.com/quickwit-oss/tantivy/issues/1783
.set_precision(DatePrecision::Microseconds)
.set_indexed()
.set_fieldnorm()
.set_stored(),
);
let schema = schema_builder.build();
let index = Index::create_in_ram(schema);
let mut index_writer = index.writer_for_tests()?;
let first_time_stamp = OffsetDateTime::now_utc();
index_writer.add_document(doc!(
date_field => DateTime::from_utc(first_time_stamp),
date_field => DateTime::from_utc(first_time_stamp),
))?;
index_writer.add_document(doc!())?;
// add one second
index_writer.add_document(doc!(
date_field => DateTime::from_utc(first_time_stamp + Duration::seconds(1)),
))?;
// add another second
let two_secs_ahead = first_time_stamp + Duration::seconds(2);
index_writer.add_document(doc!(
date_field => DateTime::from_utc(two_secs_ahead),
date_field => DateTime::from_utc(two_secs_ahead),
date_field => DateTime::from_utc(two_secs_ahead),
))?;
// add three seconds
index_writer.add_document(doc!(
date_field => DateTime::from_utc(first_time_stamp + Duration::seconds(3)),
))?;
index_writer.commit()?;
let reader = index.reader()?;
let searcher = reader.searcher();
let reader = searcher.segment_reader(0);
assert_eq!(reader.num_docs(), 5);
let date_ff_reader = reader.fast_fields().dates("multi_date_field").unwrap();
let mut docids = vec![];
date_ff_reader.get_docids_for_value_range(
DateTime::from_utc(first_time_stamp)..=DateTime::from_utc(two_secs_ahead),
0..5,
&mut docids,
);
assert_eq!(docids, vec![0, 2, 3]);
let count_multiples =
|range_query: RangeQuery| searcher.search(&range_query, &Count).unwrap();
assert_eq!(
count_multiples(RangeQuery::new_date(
"multi_date_field".to_string(),
DateTime::from_utc(first_time_stamp)..DateTime::from_utc(two_secs_ahead)
)),
2
);
Ok(())
}
#[test] #[test]
fn test_multifastfield_reader() -> crate::Result<()> { fn test_multifastfield_reader() -> crate::Result<()> {
@@ -277,7 +324,7 @@ mod tests {
index_writer.commit()?; index_writer.commit()?;
let searcher = index.reader()?.searcher(); let searcher = index.reader()?.searcher();
let segment_reader = searcher.segment_reader(0); let segment_reader = searcher.segment_reader(0);
let field_reader = segment_reader.fast_fields().i64s(item_field)?; let field_reader = segment_reader.fast_fields().i64s("items")?;
assert_eq!(field_reader.min_value(), -2); assert_eq!(field_reader.min_value(), -2);
assert_eq!(field_reader.max_value(), 6); assert_eq!(field_reader.max_value(), 6);

View File

@@ -3,11 +3,9 @@ use std::sync::Arc;
use fastfield_codecs::{open, open_u128, Column}; use fastfield_codecs::{open, open_u128, Column};
use super::multivalued::MultiValuedU128FastFieldReader; use super::multivalued::MultiValuedFastFieldReader;
use crate::directory::{CompositeFile, FileSlice}; use crate::directory::{CompositeFile, FileSlice};
use crate::fastfield::{ use crate::fastfield::{BytesFastFieldReader, FastFieldNotAvailableError, FastValue};
BytesFastFieldReader, FastFieldNotAvailableError, FastValue, MultiValuedFastFieldReader,
};
use crate::schema::{Cardinality, Field, FieldType, Schema}; use crate::schema::{Cardinality, Field, FieldType, Schema};
use crate::space_usage::PerFieldSpaceUsage; use crate::space_usage::PerFieldSpaceUsage;
use crate::{DateTime, TantivyError}; use crate::{DateTime, TantivyError};
@@ -116,9 +114,11 @@ impl FastFieldReaders {
pub(crate) fn typed_fast_field_reader_with_idx<TFastValue: FastValue>( pub(crate) fn typed_fast_field_reader_with_idx<TFastValue: FastValue>(
&self, &self,
field: Field, field_name: &str,
index: usize, index: usize,
) -> crate::Result<Arc<dyn Column<TFastValue>>> { ) -> crate::Result<Arc<dyn Column<TFastValue>>> {
let field = self.schema.get_field(field_name)?;
let fast_field_slice = self.fast_field_data(field, index)?; let fast_field_slice = self.fast_field_data(field, index)?;
let bytes = fast_field_slice.read_bytes()?; let bytes = fast_field_slice.read_bytes()?;
let column = fastfield_codecs::open(bytes)?; let column = fastfield_codecs::open(bytes)?;
@@ -127,32 +127,37 @@ impl FastFieldReaders {
pub(crate) fn typed_fast_field_reader<TFastValue: FastValue>( pub(crate) fn typed_fast_field_reader<TFastValue: FastValue>(
&self, &self,
field: Field, field_name: &str,
) -> crate::Result<Arc<dyn Column<TFastValue>>> { ) -> crate::Result<Arc<dyn Column<TFastValue>>> {
self.typed_fast_field_reader_with_idx(field, 0) self.typed_fast_field_reader_with_idx(field_name, 0)
} }
pub(crate) fn typed_fast_field_multi_reader<TFastValue: FastValue>( pub(crate) fn typed_fast_field_multi_reader<TFastValue: FastValue>(
&self, &self,
field: Field, field_name: &str,
) -> crate::Result<MultiValuedFastFieldReader<TFastValue>> { ) -> crate::Result<MultiValuedFastFieldReader<TFastValue>> {
let idx_reader = self.typed_fast_field_reader(field)?; let idx_reader = self.typed_fast_field_reader(field_name)?;
let vals_reader = self.typed_fast_field_reader_with_idx(field, 1)?; let vals_reader = self.typed_fast_field_reader_with_idx(field_name, 1)?;
Ok(MultiValuedFastFieldReader::open(idx_reader, vals_reader)) Ok(MultiValuedFastFieldReader::open(idx_reader, vals_reader))
} }
/// Returns the `u64` fast field reader reader associated with `field`. /// Returns the `u64` fast field reader reader associated with `field`.
/// ///
/// If `field` is not a u64 fast field, this method returns an Error. /// If `field` is not a u64 fast field, this method returns an Error.
pub fn u64(&self, field: Field) -> crate::Result<Arc<dyn Column<u64>>> { pub fn u64(&self, field_name: &str) -> crate::Result<Arc<dyn Column<u64>>> {
self.check_type(field, FastType::U64, Cardinality::SingleValue)?; self.check_type(
self.typed_fast_field_reader(field) self.schema.get_field(field_name)?,
FastType::U64,
Cardinality::SingleValue,
)?;
self.typed_fast_field_reader(field_name)
} }
/// Returns the `ip` fast field reader reader associated to `field`. /// Returns the `ip` fast field reader reader associated to `field`.
/// ///
/// If `field` is not a u128 fast field, this method returns an Error. /// If `field` is not a u128 fast field, this method returns an Error.
pub fn ip_addr(&self, field: Field) -> crate::Result<Arc<dyn Column<Ipv6Addr>>> { pub fn ip_addr(&self, field_name: &str) -> crate::Result<Arc<dyn Column<Ipv6Addr>>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::U128, Cardinality::SingleValue)?; self.check_type(field, FastType::U128, Cardinality::SingleValue)?;
let bytes = self.fast_field_data(field, 0)?.read_bytes()?; let bytes = self.fast_field_data(field, 0)?.read_bytes()?;
Ok(open_u128::<Ipv6Addr>(bytes)?) Ok(open_u128::<Ipv6Addr>(bytes)?)
@@ -163,24 +168,23 @@ impl FastFieldReaders {
/// If `field` is not a u128 fast field, this method returns an Error. /// If `field` is not a u128 fast field, this method returns an Error.
pub fn ip_addrs( pub fn ip_addrs(
&self, &self,
field: Field, field_name: &str,
) -> crate::Result<MultiValuedU128FastFieldReader<Ipv6Addr>> { ) -> crate::Result<MultiValuedFastFieldReader<Ipv6Addr>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::U128, Cardinality::MultiValues)?; self.check_type(field, FastType::U128, Cardinality::MultiValues)?;
let idx_reader: Arc<dyn Column<u64>> = self.typed_fast_field_reader(field)?; let idx_reader: Arc<dyn Column<u64>> = self.typed_fast_field_reader(field_name)?;
let bytes = self.fast_field_data(field, 1)?.read_bytes()?; let bytes = self.fast_field_data(field, 1)?.read_bytes()?;
let vals_reader = open_u128::<Ipv6Addr>(bytes)?; let vals_reader = open_u128::<Ipv6Addr>(bytes)?;
Ok(MultiValuedU128FastFieldReader::open( Ok(MultiValuedFastFieldReader::open(idx_reader, vals_reader))
idx_reader,
vals_reader,
))
} }
/// Returns the `u128` fast field reader reader associated to `field`. /// Returns the `u128` fast field reader reader associated to `field`.
/// ///
/// If `field` is not a u128 fast field, this method returns an Error. /// If `field` is not a u128 fast field, this method returns an Error.
pub(crate) fn u128(&self, field: Field) -> crate::Result<Arc<dyn Column<u128>>> { pub(crate) fn u128(&self, field_name: &str) -> crate::Result<Arc<dyn Column<u128>>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::U128, Cardinality::SingleValue)?; self.check_type(field, FastType::U128, Cardinality::SingleValue)?;
let bytes = self.fast_field_data(field, 0)?.read_bytes()?; let bytes = self.fast_field_data(field, 0)?.read_bytes()?;
Ok(open_u128::<u128>(bytes)?) Ok(open_u128::<u128>(bytes)?)
@@ -189,17 +193,16 @@ impl FastFieldReaders {
/// Returns the `u128` multi-valued fast field reader reader associated to `field`. /// Returns the `u128` multi-valued fast field reader reader associated to `field`.
/// ///
/// If `field` is not a u128 multi-valued fast field, this method returns an Error. /// If `field` is not a u128 multi-valued fast field, this method returns an Error.
pub fn u128s(&self, field: Field) -> crate::Result<MultiValuedU128FastFieldReader<u128>> { pub fn u128s(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<u128>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::U128, Cardinality::MultiValues)?; self.check_type(field, FastType::U128, Cardinality::MultiValues)?;
let idx_reader: Arc<dyn Column<u64>> = self.typed_fast_field_reader(field)?; let idx_reader: Arc<dyn Column<u64>> =
self.typed_fast_field_reader(self.schema.get_field_name(field))?;
let bytes = self.fast_field_data(field, 1)?.read_bytes()?; let bytes = self.fast_field_data(field, 1)?.read_bytes()?;
let vals_reader = open_u128::<u128>(bytes)?; let vals_reader = open_u128::<u128>(bytes)?;
Ok(MultiValuedU128FastFieldReader::open( Ok(MultiValuedFastFieldReader::open(idx_reader, vals_reader))
idx_reader,
vals_reader,
))
} }
/// Returns the `u64` fast field reader reader associated with `field`, regardless of whether /// Returns the `u64` fast field reader reader associated with `field`, regardless of whether
@@ -207,80 +210,88 @@ impl FastFieldReaders {
/// ///
/// If not, the fastfield reader will returns the u64-value associated with the original /// If not, the fastfield reader will returns the u64-value associated with the original
/// FastValue. /// FastValue.
pub fn u64_lenient(&self, field: Field) -> crate::Result<Arc<dyn Column<u64>>> { pub fn u64_lenient(&self, field_name: &str) -> crate::Result<Arc<dyn Column<u64>>> {
self.typed_fast_field_reader(field) self.typed_fast_field_reader(field_name)
} }
/// Returns the `i64` fast field reader reader associated with `field`. /// Returns the `i64` fast field reader reader associated with `field`.
/// ///
/// If `field` is not a i64 fast field, this method returns an Error. /// If `field` is not a i64 fast field, this method returns an Error.
pub fn i64(&self, field: Field) -> crate::Result<Arc<dyn Column<i64>>> { pub fn i64(&self, field_name: &str) -> crate::Result<Arc<dyn Column<i64>>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::I64, Cardinality::SingleValue)?; self.check_type(field, FastType::I64, Cardinality::SingleValue)?;
self.typed_fast_field_reader(field) self.typed_fast_field_reader(self.schema.get_field_name(field))
} }
/// Returns the `date` fast field reader reader associated with `field`. /// Returns the `date` fast field reader reader associated with `field`.
/// ///
/// If `field` is not a date fast field, this method returns an Error. /// If `field` is not a date fast field, this method returns an Error.
pub fn date(&self, field: Field) -> crate::Result<Arc<dyn Column<DateTime>>> { pub fn date(&self, field_name: &str) -> crate::Result<Arc<dyn Column<DateTime>>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::Date, Cardinality::SingleValue)?; self.check_type(field, FastType::Date, Cardinality::SingleValue)?;
self.typed_fast_field_reader(field) self.typed_fast_field_reader(field_name)
} }
/// Returns the `f64` fast field reader reader associated with `field`. /// Returns the `f64` fast field reader reader associated with `field`.
/// ///
/// If `field` is not a f64 fast field, this method returns an Error. /// If `field` is not a f64 fast field, this method returns an Error.
pub fn f64(&self, field: Field) -> crate::Result<Arc<dyn Column<f64>>> { pub fn f64(&self, field_name: &str) -> crate::Result<Arc<dyn Column<f64>>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::F64, Cardinality::SingleValue)?; self.check_type(field, FastType::F64, Cardinality::SingleValue)?;
self.typed_fast_field_reader(field) self.typed_fast_field_reader(field_name)
} }
/// Returns the `bool` fast field reader reader associated with `field`. /// Returns the `bool` fast field reader reader associated with `field`.
/// ///
/// If `field` is not a bool fast field, this method returns an Error. /// If `field` is not a bool fast field, this method returns an Error.
pub fn bool(&self, field: Field) -> crate::Result<Arc<dyn Column<bool>>> { pub fn bool(&self, field_name: &str) -> crate::Result<Arc<dyn Column<bool>>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::Bool, Cardinality::SingleValue)?; self.check_type(field, FastType::Bool, Cardinality::SingleValue)?;
self.typed_fast_field_reader(field) self.typed_fast_field_reader(field_name)
} }
/// Returns a `u64s` multi-valued fast field reader reader associated with `field`. /// Returns a `u64s` multi-valued fast field reader reader associated with `field`.
/// ///
/// If `field` is not a u64 multi-valued fast field, this method returns an Error. /// If `field` is not a u64 multi-valued fast field, this method returns an Error.
pub fn u64s(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<u64>> { pub fn u64s(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<u64>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::U64, Cardinality::MultiValues)?; self.check_type(field, FastType::U64, Cardinality::MultiValues)?;
self.typed_fast_field_multi_reader(field) self.typed_fast_field_multi_reader(field_name)
} }
/// Returns a `u64s` multi-valued fast field reader reader associated with `field`, regardless /// Returns a `u64s` multi-valued fast field reader reader associated with `field`, regardless
/// of whether the given field is effectively of type `u64` or not. /// of whether the given field is effectively of type `u64` or not.
/// ///
/// If `field` is not a u64 multi-valued fast field, this method returns an Error. /// If `field` is not a u64 multi-valued fast field, this method returns an Error.
pub fn u64s_lenient(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<u64>> { pub fn u64s_lenient(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<u64>> {
self.typed_fast_field_multi_reader(field) self.typed_fast_field_multi_reader(field_name)
} }
/// Returns a `i64s` multi-valued fast field reader reader associated with `field`. /// Returns a `i64s` multi-valued fast field reader reader associated with `field`.
/// ///
/// If `field` is not a i64 multi-valued fast field, this method returns an Error. /// If `field` is not a i64 multi-valued fast field, this method returns an Error.
pub fn i64s(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<i64>> { pub fn i64s(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<i64>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::I64, Cardinality::MultiValues)?; self.check_type(field, FastType::I64, Cardinality::MultiValues)?;
self.typed_fast_field_multi_reader(field) self.typed_fast_field_multi_reader(self.schema.get_field_name(field))
} }
/// Returns a `f64s` multi-valued fast field reader reader associated with `field`. /// Returns a `f64s` multi-valued fast field reader reader associated with `field`.
/// ///
/// If `field` is not a f64 multi-valued fast field, this method returns an Error. /// If `field` is not a f64 multi-valued fast field, this method returns an Error.
pub fn f64s(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<f64>> { pub fn f64s(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<f64>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::F64, Cardinality::MultiValues)?; self.check_type(field, FastType::F64, Cardinality::MultiValues)?;
self.typed_fast_field_multi_reader(field) self.typed_fast_field_multi_reader(self.schema.get_field_name(field))
} }
/// Returns a `bools` multi-valued fast field reader reader associated with `field`. /// Returns a `bools` multi-valued fast field reader reader associated with `field`.
/// ///
/// If `field` is not a bool multi-valued fast field, this method returns an Error. /// If `field` is not a bool multi-valued fast field, this method returns an Error.
pub fn bools(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<bool>> { pub fn bools(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<bool>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::Bool, Cardinality::MultiValues)?; self.check_type(field, FastType::Bool, Cardinality::MultiValues)?;
self.typed_fast_field_multi_reader(field) self.typed_fast_field_multi_reader(self.schema.get_field_name(field))
} }
/// Returns a `time::OffsetDateTime` multi-valued fast field reader reader associated with /// Returns a `time::OffsetDateTime` multi-valued fast field reader reader associated with
@@ -288,15 +299,17 @@ impl FastFieldReaders {
/// ///
/// If `field` is not a `time::OffsetDateTime` multi-valued fast field, this method returns an /// If `field` is not a `time::OffsetDateTime` multi-valued fast field, this method returns an
/// Error. /// Error.
pub fn dates(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<DateTime>> { pub fn dates(&self, field_name: &str) -> crate::Result<MultiValuedFastFieldReader<DateTime>> {
let field = self.schema.get_field(field_name)?;
self.check_type(field, FastType::Date, Cardinality::MultiValues)?; self.check_type(field, FastType::Date, Cardinality::MultiValues)?;
self.typed_fast_field_multi_reader(field) self.typed_fast_field_multi_reader(self.schema.get_field_name(field))
} }
/// Returns the `bytes` fast field reader associated with `field`. /// Returns the `bytes` fast field reader associated with `field`.
/// ///
/// If `field` is not a bytes fast field, returns an Error. /// If `field` is not a bytes fast field, returns an Error.
pub fn bytes(&self, field: Field) -> crate::Result<BytesFastFieldReader> { pub fn bytes(&self, field_name: &str) -> crate::Result<BytesFastFieldReader> {
let field = self.schema.get_field(field_name)?;
let field_entry = self.schema.get_field_entry(field); let field_entry = self.schema.get_field_entry(field);
if let FieldType::Bytes(bytes_option) = field_entry.field_type() { if let FieldType::Bytes(bytes_option) = field_entry.field_type() {
if !bytes_option.is_fast() { if !bytes_option.is_fast() {

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