mirror of
https://github.com/quickwit-oss/tantivy.git
synced 2025-12-28 04:52:55 +00:00
Compare commits
25 Commits
column-han
...
use_stats
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e07f1970ea | ||
|
|
78273bfb0d | ||
|
|
2650111b76 | ||
|
|
1176555eff | ||
|
|
f8d111a75e | ||
|
|
e17996f2fd | ||
|
|
f3621c0487 | ||
|
|
14222a47a3 | ||
|
|
8312c882a5 | ||
|
|
7a8fce0ae7 | ||
|
|
196e42f33e | ||
|
|
82a183bc2d | ||
|
|
3090d49615 | ||
|
|
7c6cc818ae | ||
|
|
514d23a20c | ||
|
|
4f9efe654c | ||
|
|
1afa5bf3db | ||
|
|
07a51eb7c8 | ||
|
|
2080c370c2 | ||
|
|
b22f96624e | ||
|
|
b78dc5e313 | ||
|
|
3f915925af | ||
|
|
9c5fef5af7 | ||
|
|
9948a84ebe | ||
|
|
45156fd869 |
@@ -15,7 +15,7 @@ rust-version = "1.62"
|
||||
|
||||
[dependencies]
|
||||
oneshot = "0.1.5"
|
||||
base64 = "0.20.0"
|
||||
base64 = "0.21.0"
|
||||
byteorder = "1.4.3"
|
||||
crc32fast = "1.3.2"
|
||||
once_cell = "1.10.0"
|
||||
@@ -48,7 +48,7 @@ murmurhash32 = "0.2.0"
|
||||
time = { version = "0.3.10", features = ["serde-well-known"] }
|
||||
smallvec = "1.8.0"
|
||||
rayon = "1.5.2"
|
||||
lru = "0.7.5"
|
||||
lru = "0.9.0"
|
||||
fastdivide = "0.4.0"
|
||||
itertools = "0.10.3"
|
||||
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" }
|
||||
common = { version= "0.5", path = "./common/", package = "tantivy-common" }
|
||||
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]
|
||||
winapi = "0.3.9"
|
||||
@@ -106,7 +107,7 @@ unstable = [] # useful for benches.
|
||||
quickwit = ["sstable"]
|
||||
|
||||
[workspace]
|
||||
members = ["query-grammar", "bitpacker", "common", "fastfield_codecs", "ownedbytes", "stacker", "sstable"]
|
||||
members = ["query-grammar", "bitpacker", "common", "fastfield_codecs", "ownedbytes", "stacker", "sstable", "columnar", "tokenizer-api"]
|
||||
|
||||
# Following the "fail" crate best practises, we isolate
|
||||
# tests that define specific behavior in fail check points
|
||||
|
||||
10
README.md
10
README.md
@@ -29,7 +29,7 @@ Your mileage WILL vary depending on the nature of queries and their load.
|
||||
# Features
|
||||
|
||||
- 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:)
|
||||
- Tiny startup time (<10ms), perfect for command-line tools
|
||||
- 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)
|
||||
- `&[u8]` fast fields
|
||||
- Text, i64, u64, f64, dates, and hierarchical facet fields
|
||||
- LZ4 compressed document store
|
||||
- Compressed document store (LZ4, Zstd, None, Brotli, Snap)
|
||||
- Range queries
|
||||
- Faceted search
|
||||
- Configurable indexing (optional term frequency and position indexing)
|
||||
- JSON Field
|
||||
- Aggregation Collector: range buckets, average, and stats metrics
|
||||
- Aggregation Collector: histogram, range buckets, average, and stats metrics
|
||||
- LogMergePolicy with deletes
|
||||
- Searcher Warmer API
|
||||
- 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.
|
||||
|
||||
## Tokenizer
|
||||
|
||||
When implementing a tokenizer for tantivy depend on the `tantivy-tokenizer-api` crate.
|
||||
|
||||
## Minimum supported Rust version
|
||||
|
||||
Tantivy currently requires at least Rust 1.62 or later to compile.
|
||||
|
||||
18
columnar/Cargo.toml
Normal file
18
columnar/Cargo.toml
Normal file
@@ -0,0 +1,18 @@
|
||||
[package]
|
||||
name = "tantivy-columnar"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
|
||||
[dependencies]
|
||||
stacker = { path = "../stacker", package="tantivy-stacker"}
|
||||
serde_json = "1"
|
||||
thiserror = "1"
|
||||
fnv = "1"
|
||||
sstable = { path = "../sstable", package = "tantivy-sstable" }
|
||||
common = { path = "../common", package = "tantivy-common" }
|
||||
fastfield_codecs = { path = "../fastfield_codecs"}
|
||||
itertools = "0.10"
|
||||
|
||||
[dev-dependencies]
|
||||
proptest = "1"
|
||||
67
columnar/README.md
Normal file
67
columnar/README.md
Normal file
@@ -0,0 +1,67 @@
|
||||
# Columnar format
|
||||
|
||||
This crate describes columnar format used in tantivy.
|
||||
|
||||
## Goals
|
||||
|
||||
This format is special in the following way.
|
||||
- it needs to be compact
|
||||
- it does not required to be loaded in memory.
|
||||
- it is designed to fit well with quickwit's strange constraint:
|
||||
we need to be able to load columns rapidly.
|
||||
- columns of several types can be associated with the same column name.
|
||||
- it needs to support columns with different types `(str, u64, i64, f64)`
|
||||
and different cardinality `(required, optional, multivalued)`.
|
||||
- columns, once loaded, offer cheap random access.
|
||||
|
||||
# Coercion rules
|
||||
|
||||
Users can create a columnar by inserting rows to a `ColumnarWriter`,
|
||||
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:
|
||||
- JsonValues are grouped into 3 types (String, Number, bool).
|
||||
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
|
||||
|
||||
This columnar format may have more than one column (with different types) associated to the same `column_name` (see [Coercion rules](#coercion-rules) above).
|
||||
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:
|
||||
`[column_name][ZERO_BYTE][column_type_header: u8]`
|
||||
|
||||
```
|
||||
COLUMNAR:=
|
||||
[COLUMNAR_DATA]
|
||||
[COLUMNAR_KEY_TO_DATA_INDEX]
|
||||
[COLUMNAR_FOOTER];
|
||||
|
||||
|
||||
# Columns are sorted by their column key.
|
||||
COLUMNAR_DATA:=
|
||||
[COLUMN_DATA]+;
|
||||
|
||||
COLUMNAR_FOOTER := [RANGE_SSTABLE_BYTES_LEN: 8 bytes little endian]
|
||||
|
||||
```
|
||||
|
||||
The columnar file starts by the actual column data, concatenated one after the other,
|
||||
sorted by column key.
|
||||
|
||||
A sstable associates
|
||||
`(column name, column_cardinality, column_type) to range of bytes.
|
||||
|
||||
Column name may not contain the zero byte `\0`.
|
||||
|
||||
Listing all columns associated to `column_name` can therefore
|
||||
be done by listing all keys prefixed by
|
||||
`[column_name][ZERO_BYTE]`
|
||||
|
||||
The associated range of bytes refer to a range of bytes
|
||||
|
||||
201
columnar/src/column_type_header.rs
Normal file
201
columnar/src/column_type_header.rs
Normal file
@@ -0,0 +1,201 @@
|
||||
use crate::utils::{place_bits, select_bits};
|
||||
use crate::value::NumericalType;
|
||||
use crate::InvalidData;
|
||||
|
||||
/// 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.
|
||||
/// Required is the default for auto-detecting the Cardinality, since it is the most strict.
|
||||
#[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) -> Result<Cardinality, InvalidData> {
|
||||
match code {
|
||||
0 => Ok(Cardinality::Required),
|
||||
1 => Ok(Cardinality::Optional),
|
||||
2 => Ok(Cardinality::Multivalued),
|
||||
_ => Err(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,
|
||||
}
|
||||
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Represents the type and cardinality of a column.
|
||||
/// This is encoded over one-byte and added to a column key in the
|
||||
/// columnar sstable.
|
||||
///
|
||||
/// - [0..6] bits: encodes the column type
|
||||
/// - [6..8] bits: encodes the cardinality
|
||||
#[derive(Eq, Hash, PartialEq, Debug, Copy, Clone)]
|
||||
pub struct ColumnTypeAndCardinality {
|
||||
pub typ: ColumnType,
|
||||
pub cardinality: Cardinality,
|
||||
}
|
||||
|
||||
impl ColumnTypeAndCardinality {
|
||||
pub fn to_code(self) -> u8 {
|
||||
place_bits::<0, 6>(self.typ.to_code()) | place_bits::<6, 8>(self.cardinality.to_code())
|
||||
}
|
||||
|
||||
pub fn try_from_code(code: u8) -> Result<ColumnTypeAndCardinality, InvalidData> {
|
||||
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);
|
||||
Ok(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 Ok(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 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);
|
||||
}
|
||||
}
|
||||
84
columnar/src/dictionary.rs
Normal file
84
columnar/src/dictionary.rs
Normal file
@@ -0,0 +1,84 @@
|
||||
use std::io;
|
||||
|
||||
use fnv::FnvHashMap;
|
||||
use sstable::SSTable;
|
||||
|
||||
pub(crate) struct TermIdMapping {
|
||||
unordered_to_ord: Vec<OrderedId>,
|
||||
}
|
||||
|
||||
impl TermIdMapping {
|
||||
pub fn to_ord(&self, unordered: UnorderedId) -> OrderedId {
|
||||
self.unordered_to_ord[unordered.0 as usize]
|
||||
}
|
||||
}
|
||||
|
||||
/// When we add values, we cannot know their ordered id yet.
|
||||
/// For this reason, we temporarily assign them a `UnorderedId`
|
||||
/// that will be mapped to an `OrderedId` upon serialization.
|
||||
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
|
||||
pub struct UnorderedId(pub u32);
|
||||
|
||||
#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
|
||||
pub struct OrderedId(pub u32);
|
||||
|
||||
/// `DictionaryBuilder` for dictionary encoding.
|
||||
///
|
||||
/// It stores the different terms encounterred and assigns them a temporary value
|
||||
/// we call unordered id.
|
||||
///
|
||||
/// Upon serialization, we will sort the ids and hence build a `UnorderedId -> Term ordinal`
|
||||
/// mapping.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct DictionaryBuilder {
|
||||
dict: FnvHashMap<Vec<u8>, UnorderedId>,
|
||||
}
|
||||
|
||||
impl DictionaryBuilder {
|
||||
/// Get or allocate an unordered id.
|
||||
/// (This ID is simply an auto-incremented id.)
|
||||
pub fn get_or_allocate_id(&mut self, term: &[u8]) -> UnorderedId {
|
||||
if let Some(term_id) = self.dict.get(term) {
|
||||
return *term_id;
|
||||
}
|
||||
let new_id = UnorderedId(self.dict.len() as u32);
|
||||
self.dict.insert(term.to_vec(), new_id);
|
||||
new_id
|
||||
}
|
||||
|
||||
/// Serialize the dictionary into an fst, and returns the
|
||||
/// `UnorderedId -> TermOrdinal` map.
|
||||
pub fn serialize<'a, W: io::Write + 'a>(&self, wrt: &mut W) -> io::Result<TermIdMapping> {
|
||||
let mut terms: Vec<(&[u8], UnorderedId)> =
|
||||
self.dict.iter().map(|(k, v)| (k.as_slice(), *v)).collect();
|
||||
terms.sort_unstable_by_key(|(key, _)| *key);
|
||||
// TODO Remove the allocation.
|
||||
let mut unordered_to_ord: Vec<OrderedId> = vec![OrderedId(0u32); terms.len()];
|
||||
let mut sstable_builder = sstable::VoidSSTable::writer(wrt);
|
||||
for (ord, (key, unordered_id)) in terms.into_iter().enumerate() {
|
||||
let ordered_id = OrderedId(ord as u32);
|
||||
sstable_builder.insert(key, &())?;
|
||||
unordered_to_ord[unordered_id.0 as usize] = ordered_id;
|
||||
}
|
||||
sstable_builder.finish()?;
|
||||
Ok(TermIdMapping { unordered_to_ord })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_dictionary_builder() {
|
||||
let mut dictionary_builder = DictionaryBuilder::default();
|
||||
let hello_uid = dictionary_builder.get_or_allocate_id(b"hello");
|
||||
let happy_uid = dictionary_builder.get_or_allocate_id(b"happy");
|
||||
let tax_uid = dictionary_builder.get_or_allocate_id(b"tax");
|
||||
let mut buffer = Vec::new();
|
||||
let id_mapping = dictionary_builder.serialize(&mut buffer).unwrap();
|
||||
assert_eq!(id_mapping.to_ord(hello_uid), OrderedId(1));
|
||||
assert_eq!(id_mapping.to_ord(happy_uid), OrderedId(0));
|
||||
assert_eq!(id_mapping.to_ord(tax_uid), OrderedId(2));
|
||||
}
|
||||
}
|
||||
89
columnar/src/lib.rs
Normal file
89
columnar/src/lib.rs
Normal file
@@ -0,0 +1,89 @@
|
||||
mod column_type_header;
|
||||
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;
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub struct InvalidData;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ops::Range;
|
||||
|
||||
use common::file_slice::FileSlice;
|
||||
|
||||
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality};
|
||||
use crate::reader::ColumnarReader;
|
||||
use crate::value::NumericalValue;
|
||||
use crate::{Cardinality, ColumnarWriter};
|
||||
|
||||
#[test]
|
||||
fn test_dataframe_writer_bytes() {
|
||||
let mut dataframe_writer = ColumnarWriter::default();
|
||||
dataframe_writer.record_str(1u32, "my_string", "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_fileslice = FileSlice::from(buffer);
|
||||
let columnar = ColumnarReader::open(columnar_fileslice).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
|
||||
columnar.read_columns("my_string").unwrap();
|
||||
assert_eq!(cols.len(), 1);
|
||||
assert_eq!(cols[0].1, 0..158);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dataframe_writer_bool() {
|
||||
let mut dataframe_writer = ColumnarWriter::default();
|
||||
dataframe_writer.record_bool(1u32, "bool.value", false);
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar_fileslice = FileSlice::from(buffer);
|
||||
let columnar = ColumnarReader::open(columnar_fileslice).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
|
||||
columnar.read_columns("bool.value").unwrap();
|
||||
assert_eq!(cols.len(), 1);
|
||||
assert_eq!(
|
||||
cols[0].0,
|
||||
ColumnTypeAndCardinality {
|
||||
cardinality: Cardinality::Optional,
|
||||
typ: ColumnType::Bool
|
||||
}
|
||||
);
|
||||
assert_eq!(cols[0].1, 0..21);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dataframe_writer_numerical() {
|
||||
let mut dataframe_writer = ColumnarWriter::default();
|
||||
dataframe_writer.record_numerical(1u32, "srical.value", NumericalValue::U64(12u64));
|
||||
dataframe_writer.record_numerical(2u32, "srical.value", NumericalValue::U64(13u64));
|
||||
dataframe_writer.record_numerical(4u32, "srical.value", NumericalValue::U64(15u64));
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
dataframe_writer.serialize(5, &mut buffer).unwrap();
|
||||
let columnar_fileslice = FileSlice::from(buffer);
|
||||
let columnar = ColumnarReader::open(columnar_fileslice).unwrap();
|
||||
assert_eq!(columnar.num_columns(), 1);
|
||||
let cols: Vec<(ColumnTypeAndCardinality, Range<u64>)> =
|
||||
columnar.read_columns("srical.value").unwrap();
|
||||
assert_eq!(cols.len(), 1);
|
||||
// Right now this 31 bytes are spent as follows
|
||||
//
|
||||
// - header 14 bytes
|
||||
// - vals 8 //< due to padding? could have been 1byte?.
|
||||
// - null footer 6 bytes
|
||||
// - version footer 3 bytes // Should be file-wide
|
||||
assert_eq!(cols[0].1, 0..31);
|
||||
}
|
||||
}
|
||||
110
columnar/src/reader/mod.rs
Normal file
110
columnar/src/reader/mod.rs
Normal file
@@ -0,0 +1,110 @@
|
||||
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)
|
||||
}
|
||||
|
||||
/// 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)
|
||||
.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]);
|
||||
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 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<(ColumnTypeAndCardinality, Range<u64>)>> {
|
||||
// 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)`.
|
||||
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_and_cardinality = ColumnTypeAndCardinality::try_from_code(column_code)
|
||||
.map_err(|_| 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()
|
||||
}
|
||||
}
|
||||
76
columnar/src/utils.rs
Normal file
76
columnar/src/utils.rs
Normal file
@@ -0,0 +1,76 @@
|
||||
const fn compute_mask(num_bits: u8) -> u8 {
|
||||
if num_bits == 8 {
|
||||
u8::MAX
|
||||
} else {
|
||||
(1u8 << num_bits) - 1
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
#[must_use]
|
||||
pub(crate) fn select_bits<const START: u8, const END: u8>(code: u8) -> u8 {
|
||||
assert!(START <= END);
|
||||
assert!(END <= 8);
|
||||
let num_bits: u8 = END - START;
|
||||
let mask: u8 = compute_mask(num_bits);
|
||||
(code >> START) & mask
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
#[must_use]
|
||||
pub(crate) fn place_bits<const START: u8, const END: u8>(code: u8) -> u8 {
|
||||
assert!(START <= END);
|
||||
assert!(END <= 8);
|
||||
let num_bits: u8 = END - START;
|
||||
let mask: u8 = compute_mask(num_bits);
|
||||
assert!(code <= mask);
|
||||
code << START
|
||||
}
|
||||
|
||||
/// Pop-front one bytes from a slice of bytes.
|
||||
#[inline(always)]
|
||||
pub fn pop_first_byte(bytes: &mut &[u8]) -> Option<u8> {
|
||||
if bytes.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let first_byte = bytes[0];
|
||||
*bytes = &bytes[1..];
|
||||
Some(first_byte)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_select_bits() {
|
||||
assert_eq!(255u8, select_bits::<0, 8>(255u8));
|
||||
assert_eq!(0u8, select_bits::<0, 0>(255u8));
|
||||
assert_eq!(8u8, select_bits::<0, 4>(8u8));
|
||||
assert_eq!(4u8, select_bits::<1, 4>(8u8));
|
||||
assert_eq!(0u8, select_bits::<1, 3>(8u8));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_place_bits() {
|
||||
assert_eq!(255u8, place_bits::<0, 8>(255u8));
|
||||
assert_eq!(4u8, place_bits::<2, 3>(1u8));
|
||||
assert_eq!(0u8, place_bits::<2, 2>(0u8));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic]
|
||||
fn test_place_bits_overflows() {
|
||||
let _ = place_bits::<1, 4>(8u8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pop_first_byte() {
|
||||
let mut cursor: &[u8] = &b"abcd"[..];
|
||||
assert_eq!(pop_first_byte(&mut cursor), Some(b'a'));
|
||||
assert_eq!(pop_first_byte(&mut cursor), Some(b'b'));
|
||||
assert_eq!(pop_first_byte(&mut cursor), Some(b'c'));
|
||||
assert_eq!(pop_first_byte(&mut cursor), Some(b'd'));
|
||||
assert_eq!(pop_first_byte(&mut cursor), None);
|
||||
}
|
||||
}
|
||||
124
columnar/src/value.rs
Normal file
124
columnar/src/value.rs
Normal file
@@ -0,0 +1,124 @@
|
||||
use crate::InvalidData;
|
||||
|
||||
#[derive(Copy, Clone, Debug, PartialEq)]
|
||||
pub enum NumericalValue {
|
||||
I64(i64),
|
||||
U64(u64),
|
||||
F64(f64),
|
||||
}
|
||||
|
||||
impl From<u64> for NumericalValue {
|
||||
fn from(val: u64) -> NumericalValue {
|
||||
NumericalValue::U64(val)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for NumericalValue {
|
||||
fn from(val: i64) -> Self {
|
||||
NumericalValue::I64(val)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<f64> for NumericalValue {
|
||||
fn from(val: f64) -> Self {
|
||||
NumericalValue::F64(val)
|
||||
}
|
||||
}
|
||||
|
||||
impl NumericalValue {
|
||||
pub fn numerical_type(&self) -> NumericalType {
|
||||
match self {
|
||||
NumericalValue::F64(_) => NumericalType::F64,
|
||||
NumericalValue::I64(_) => NumericalType::I64,
|
||||
NumericalValue::U64(_) => NumericalType::U64,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for NumericalValue {}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq)]
|
||||
#[repr(u8)]
|
||||
pub enum NumericalType {
|
||||
#[default]
|
||||
I64 = 0,
|
||||
U64 = 1,
|
||||
F64 = 2,
|
||||
}
|
||||
|
||||
impl NumericalType {
|
||||
pub fn to_code(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
|
||||
pub fn try_from_code(code: u8) -> Result<NumericalType, InvalidData> {
|
||||
match code {
|
||||
0 => Ok(NumericalType::I64),
|
||||
1 => Ok(NumericalType::U64),
|
||||
2 => Ok(NumericalType::F64),
|
||||
_ => Err(InvalidData),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// We voluntarily avoid using `Into` here to keep this
|
||||
/// implementation quirk as private as possible.
|
||||
///
|
||||
/// # Panics
|
||||
/// This coercion trait actually panics if it is used
|
||||
/// to convert a loose types to a stricter type.
|
||||
///
|
||||
/// The level is strictness is somewhat arbitrary.
|
||||
/// - i64
|
||||
/// - u64
|
||||
/// - f64.
|
||||
pub(crate) trait Coerce {
|
||||
fn coerce(numerical_value: NumericalValue) -> Self;
|
||||
}
|
||||
|
||||
impl Coerce for i64 {
|
||||
fn coerce(value: NumericalValue) -> Self {
|
||||
match value {
|
||||
NumericalValue::I64(val) => val,
|
||||
NumericalValue::U64(val) => val as i64,
|
||||
NumericalValue::F64(_) => unreachable!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Coerce for u64 {
|
||||
fn coerce(value: NumericalValue) -> Self {
|
||||
match value {
|
||||
NumericalValue::I64(val) => val as u64,
|
||||
NumericalValue::U64(val) => val,
|
||||
NumericalValue::F64(_) => unreachable!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Coerce for f64 {
|
||||
fn coerce(value: NumericalValue) -> Self {
|
||||
match value {
|
||||
NumericalValue::I64(val) => val as f64,
|
||||
NumericalValue::U64(val) => val as f64,
|
||||
NumericalValue::F64(val) => val,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::NumericalType;
|
||||
|
||||
#[test]
|
||||
fn test_numerical_type_code() {
|
||||
let mut num_numerical_type = 0;
|
||||
for code in u8::MIN..=u8::MAX {
|
||||
if let Ok(numerical_type) = NumericalType::try_from_code(code) {
|
||||
assert_eq!(numerical_type.to_code(), code);
|
||||
num_numerical_type += 1;
|
||||
}
|
||||
}
|
||||
assert_eq!(num_numerical_type, 3);
|
||||
}
|
||||
}
|
||||
346
columnar/src/writer/column_operation.rs
Normal file
346
columnar/src/writer/column_operation.rs
Normal file
@@ -0,0 +1,346 @@
|
||||
use crate::dictionary::UnorderedId;
|
||||
use crate::utils::{place_bits, pop_first_byte, select_bits};
|
||||
use crate::value::NumericalValue;
|
||||
use crate::{DocId, InvalidData, NumericalType};
|
||||
|
||||
/// When we build a columnar dataframe, we first just group
|
||||
/// 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`.
|
||||
#[derive(Eq, PartialEq, Debug, Clone, Copy)]
|
||||
pub(super) enum ColumnOperation<T> {
|
||||
NewDoc(DocId),
|
||||
Value(T),
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
|
||||
struct ColumnOperationMetadata {
|
||||
op_type: ColumnOperationType,
|
||||
len: u8,
|
||||
}
|
||||
|
||||
impl ColumnOperationMetadata {
|
||||
fn to_code(self) -> u8 {
|
||||
place_bits::<0, 4>(self.len) | place_bits::<4, 8>(self.op_type.to_code())
|
||||
}
|
||||
|
||||
fn try_from_code(code: u8) -> Result<Self, InvalidData> {
|
||||
let len = select_bits::<0, 4>(code);
|
||||
let typ_code = select_bits::<4, 8>(code);
|
||||
let column_type = ColumnOperationType::try_from_code(typ_code)?;
|
||||
Ok(ColumnOperationMetadata {
|
||||
op_type: column_type,
|
||||
len,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
|
||||
#[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> {
|
||||
pub(super) fn serialize(self) -> impl AsRef<[u8]> {
|
||||
let mut minibuf = MiniBuffer::default();
|
||||
let column_op_metadata = match self {
|
||||
ColumnOperation::NewDoc(new_doc) => {
|
||||
let symbol_len = new_doc.serialize(&mut minibuf.bytes[1..]);
|
||||
ColumnOperationMetadata {
|
||||
op_type: ColumnOperationType::NewDoc,
|
||||
len: symbol_len,
|
||||
}
|
||||
}
|
||||
ColumnOperation::Value(val) => {
|
||||
let symbol_len = val.serialize(&mut minibuf.bytes[1..]);
|
||||
ColumnOperationMetadata {
|
||||
op_type: ColumnOperationType::AddValue,
|
||||
len: symbol_len,
|
||||
}
|
||||
}
|
||||
};
|
||||
minibuf.bytes[0] = column_op_metadata.to_code();
|
||||
// +1 for the metadata
|
||||
minibuf.len = 1 + column_op_metadata.len;
|
||||
minibuf
|
||||
}
|
||||
|
||||
/// Deserialize a colummn operation.
|
||||
/// Returns None if the buffer is empty.
|
||||
///
|
||||
/// Panics if the payload is invalid:
|
||||
/// this deserialize method is meant to target in memory.
|
||||
pub(super) fn deserialize(bytes: &mut &[u8]) -> Option<Self> {
|
||||
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];
|
||||
(symbol_bytes, *bytes) = bytes.split_at(column_op_metadata.len as usize);
|
||||
match column_op_metadata.op_type {
|
||||
ColumnOperationType::NewDoc => {
|
||||
let new_doc = u32::deserialize(symbol_bytes);
|
||||
Some(ColumnOperation::NewDoc(new_doc))
|
||||
}
|
||||
ColumnOperationType::AddValue => {
|
||||
let value = V::deserialize(symbol_bytes);
|
||||
Some(ColumnOperation::Value(value))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> From<T> for ColumnOperation<T> {
|
||||
fn from(value: T) -> Self {
|
||||
ColumnOperation::Value(value)
|
||||
}
|
||||
}
|
||||
|
||||
// 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)]
|
||||
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;
|
||||
// Panics if invalid
|
||||
fn deserialize(bytes: &[u8]) -> Self;
|
||||
}
|
||||
|
||||
impl SymbolValue for bool {
|
||||
fn serialize(self, buffer: &mut [u8]) -> u8 {
|
||||
buffer[0] = u8::from(self);
|
||||
1u8
|
||||
}
|
||||
|
||||
fn deserialize(bytes: &[u8]) -> Self {
|
||||
bytes[0] == 1u8
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct MiniBuffer {
|
||||
pub bytes: [u8; 10],
|
||||
pub len: u8,
|
||||
}
|
||||
|
||||
impl AsRef<[u8]> for MiniBuffer {
|
||||
fn as_ref(&self) -> &[u8] {
|
||||
&self.bytes[..self.len as usize]
|
||||
}
|
||||
}
|
||||
|
||||
impl SymbolValue for NumericalValue {
|
||||
fn deserialize(mut bytes: &[u8]) -> Self {
|
||||
let type_code = pop_first_byte(&mut bytes).unwrap();
|
||||
let symbol_type = NumericalType::try_from_code(type_code).unwrap();
|
||||
let mut octet: [u8; 8] = [0u8; 8];
|
||||
octet[..bytes.len()].copy_from_slice(bytes);
|
||||
match symbol_type {
|
||||
NumericalType::U64 => {
|
||||
let val: u64 = u64::from_le_bytes(octet);
|
||||
NumericalValue::U64(val)
|
||||
}
|
||||
NumericalType::I64 => {
|
||||
let encoded: u64 = u64::from_le_bytes(octet);
|
||||
let val: i64 = decode_zig_zag(encoded);
|
||||
NumericalValue::I64(val)
|
||||
}
|
||||
NumericalType::F64 => {
|
||||
debug_assert_eq!(bytes.len(), 8);
|
||||
let val: f64 = f64::from_le_bytes(octet);
|
||||
NumericalValue::F64(val)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
match self {
|
||||
NumericalValue::F64(val) => {
|
||||
output[0] = NumericalType::F64 as u8;
|
||||
output[1..9].copy_from_slice(&val.to_le_bytes());
|
||||
9u8
|
||||
}
|
||||
NumericalValue::U64(val) => {
|
||||
let len = compute_num_bytes_for_u64(val) as u8;
|
||||
output[0] = NumericalType::U64 as u8;
|
||||
output[1..9].copy_from_slice(&val.to_le_bytes());
|
||||
len + 1u8
|
||||
}
|
||||
NumericalValue::I64(val) => {
|
||||
let zig_zag_encoded = encode_zig_zag(val);
|
||||
let len = compute_num_bytes_for_u64(zig_zag_encoded) as u8;
|
||||
output[0] = NumericalType::I64 as u8;
|
||||
output[1..9].copy_from_slice(&zig_zag_encoded.to_le_bytes());
|
||||
len + 1u8
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SymbolValue for u32 {
|
||||
fn serialize(self, output: &mut [u8]) -> u8 {
|
||||
let len = compute_num_bytes_for_u64(self as u64);
|
||||
output[0..4].copy_from_slice(&self.to_le_bytes());
|
||||
len as u8
|
||||
}
|
||||
|
||||
fn deserialize(bytes: &[u8]) -> Self {
|
||||
let mut quartet: [u8; 4] = [0u8; 4];
|
||||
quartet[..bytes.len()].copy_from_slice(bytes);
|
||||
u32::from_le_bytes(quartet)
|
||||
}
|
||||
}
|
||||
|
||||
impl SymbolValue for UnorderedId {
|
||||
fn serialize(self, output: &mut [u8]) -> u8 {
|
||||
self.0.serialize(output)
|
||||
}
|
||||
|
||||
fn deserialize(bytes: &[u8]) -> Self {
|
||||
UnorderedId(u32::deserialize(bytes))
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_num_bytes_for_u64(val: u64) -> usize {
|
||||
let msb = (64u32 - val.leading_zeros()) as usize;
|
||||
(msb + 7) / 8
|
||||
}
|
||||
|
||||
fn encode_zig_zag(n: i64) -> u64 {
|
||||
((n << 1) ^ (n >> 63)) as u64
|
||||
}
|
||||
|
||||
fn decode_zig_zag(n: u64) -> i64 {
|
||||
((n >> 1) as i64) ^ (-((n & 1) as i64))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[track_caller]
|
||||
fn test_zig_zag_aux(val: i64) {
|
||||
let encoded = super::encode_zig_zag(val);
|
||||
assert_eq!(decode_zig_zag(encoded), val);
|
||||
if let Some(abs_val) = val.checked_abs() {
|
||||
let abs_val = abs_val as u64;
|
||||
assert!(encoded <= abs_val * 2);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_zig_zag() {
|
||||
assert_eq!(encode_zig_zag(0i64), 0u64);
|
||||
assert_eq!(encode_zig_zag(-1i64), 1u64);
|
||||
assert_eq!(encode_zig_zag(1i64), 2u64);
|
||||
test_zig_zag_aux(0i64);
|
||||
test_zig_zag_aux(i64::MIN);
|
||||
test_zig_zag_aux(i64::MAX);
|
||||
}
|
||||
|
||||
use proptest::prelude::any;
|
||||
use proptest::proptest;
|
||||
|
||||
proptest! {
|
||||
#[test]
|
||||
fn test_proptest_zig_zag(val in any::<i64>()) {
|
||||
test_zig_zag_aux(val);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_op_metadata_byte_serialization() {
|
||||
for len in 0..=15 {
|
||||
for op_type in [ColumnOperationType::AddValue, ColumnOperationType::NewDoc] {
|
||||
let column_op_metadata = ColumnOperationMetadata { op_type, len };
|
||||
let column_op_metadata_code = column_op_metadata.to_code();
|
||||
let serdeser_metadata =
|
||||
ColumnOperationMetadata::try_from_code(column_op_metadata_code).unwrap();
|
||||
assert_eq!(column_op_metadata, serdeser_metadata);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[track_caller]
|
||||
fn ser_deser_symbol(column_op: ColumnOperation<NumericalValue>) {
|
||||
let buf = column_op.serialize();
|
||||
let mut buffer = buf.as_ref().to_vec();
|
||||
buffer.extend_from_slice(b"234234");
|
||||
let mut bytes = &buffer[..];
|
||||
let serdeser_symbol = ColumnOperation::deserialize(&mut bytes).unwrap();
|
||||
assert_eq!(bytes.len() + buf.as_ref().len() as usize, buffer.len());
|
||||
assert_eq!(column_op, serdeser_symbol);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_num_bytes_for_u64() {
|
||||
assert_eq!(compute_num_bytes_for_u64(0), 0);
|
||||
assert_eq!(compute_num_bytes_for_u64(1), 1);
|
||||
assert_eq!(compute_num_bytes_for_u64(255), 1);
|
||||
assert_eq!(compute_num_bytes_for_u64(256), 2);
|
||||
assert_eq!(compute_num_bytes_for_u64((1 << 16) - 1), 2);
|
||||
assert_eq!(compute_num_bytes_for_u64(1 << 16), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_symbol_serialization() {
|
||||
ser_deser_symbol(ColumnOperation::NewDoc(0));
|
||||
ser_deser_symbol(ColumnOperation::NewDoc(3));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::I64(0i64)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::I64(1i64)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(257u64)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::I64(-257i64)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::I64(i64::MIN)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(0u64)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(u64::MIN)));
|
||||
ser_deser_symbol(ColumnOperation::Value(NumericalValue::U64(u64::MAX)));
|
||||
}
|
||||
|
||||
fn test_column_operation_unordered_aux(val: u32, expected_len: usize) {
|
||||
let column_op = ColumnOperation::Value(UnorderedId(val));
|
||||
let minibuf = column_op.serialize();
|
||||
assert_eq!(minibuf.as_ref().len() as usize, expected_len);
|
||||
let mut buf = minibuf.as_ref().to_vec();
|
||||
buf.extend_from_slice(&[2, 2, 2, 2, 2, 2]);
|
||||
let mut cursor = &buf[..];
|
||||
let column_op_serdeser: ColumnOperation<UnorderedId> =
|
||||
ColumnOperation::deserialize(&mut cursor).unwrap();
|
||||
assert_eq!(column_op_serdeser, ColumnOperation::Value(UnorderedId(val)));
|
||||
assert_eq!(cursor.len() + expected_len, buf.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_operation_unordered() {
|
||||
test_column_operation_unordered_aux(300u32, 3);
|
||||
test_column_operation_unordered_aux(1u32, 2);
|
||||
test_column_operation_unordered_aux(0u32, 1);
|
||||
}
|
||||
}
|
||||
265
columnar/src/writer/column_writers.rs
Normal file
265
columnar/src/writer/column_writers.rs
Normal file
@@ -0,0 +1,265 @@
|
||||
use std::cmp::Ordering;
|
||||
|
||||
use stacker::{ExpUnrolledLinkedList, MemoryArena};
|
||||
|
||||
use crate::dictionary::{DictionaryBuilder, UnorderedId};
|
||||
use crate::writer::column_operation::{ColumnOperation, SymbolValue};
|
||||
use crate::{Cardinality, DocId, NumericalType, NumericalValue};
|
||||
|
||||
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
|
||||
#[repr(u8)]
|
||||
enum DocumentStep {
|
||||
Same = 0,
|
||||
Next = 1,
|
||||
Skipped = 2,
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn delta_with_last_doc(last_doc_opt: Option<u32>, doc: u32) -> DocumentStep {
|
||||
let expected_next_doc = last_doc_opt.map(|last_doc| last_doc + 1).unwrap_or(0u32);
|
||||
match doc.cmp(&expected_next_doc) {
|
||||
Ordering::Less => DocumentStep::Same,
|
||||
Ordering::Equal => DocumentStep::Next,
|
||||
Ordering::Greater => DocumentStep::Skipped,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Default)]
|
||||
pub struct ColumnWriter {
|
||||
// Detected cardinality of the column so far.
|
||||
cardinality: Cardinality,
|
||||
// Last document inserted.
|
||||
// None if no doc has been added yet.
|
||||
last_doc_opt: Option<u32>,
|
||||
// Buffer containing the serialized values.
|
||||
values: ExpUnrolledLinkedList,
|
||||
}
|
||||
|
||||
impl ColumnWriter {
|
||||
/// Returns an iterator over the Symbol that have been recorded
|
||||
/// for the given column.
|
||||
pub(super) fn operation_iterator<'a, V: SymbolValue>(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
buffer: &'a mut Vec<u8>,
|
||||
) -> impl Iterator<Item = ColumnOperation<V>> + 'a {
|
||||
buffer.clear();
|
||||
self.values.read_to_end(arena, buffer);
|
||||
let mut cursor: &[u8] = &buffer[..];
|
||||
std::iter::from_fn(move || ColumnOperation::deserialize(&mut cursor))
|
||||
}
|
||||
|
||||
/// Records a change of the document being recorded.
|
||||
///
|
||||
/// This function will also update the cardinality of the column
|
||||
/// if necessary.
|
||||
pub(super) fn record<S: SymbolValue>(&mut self, doc: DocId, value: S, arena: &mut MemoryArena) {
|
||||
// Difference between `doc` and the last doc.
|
||||
match delta_with_last_doc(self.last_doc_opt, doc) {
|
||||
DocumentStep::Same => {
|
||||
// This is the last encounterred document.
|
||||
self.cardinality = Cardinality::Multivalued;
|
||||
}
|
||||
DocumentStep::Next => {
|
||||
self.last_doc_opt = Some(doc);
|
||||
self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena);
|
||||
}
|
||||
DocumentStep::Skipped => {
|
||||
self.cardinality = self.cardinality.max(Cardinality::Optional);
|
||||
self.last_doc_opt = Some(doc);
|
||||
self.write_symbol::<S>(ColumnOperation::NewDoc(doc), arena);
|
||||
}
|
||||
}
|
||||
self.write_symbol(ColumnOperation::Value(value), arena);
|
||||
}
|
||||
|
||||
// Get the cardinality.
|
||||
// The overall number of docs in the column is necessary to
|
||||
// deal with the case where the all docs contain 1 value, except some documents
|
||||
// at the end of the column.
|
||||
pub(crate) fn get_cardinality(&self, num_docs: DocId) -> Cardinality {
|
||||
match delta_with_last_doc(self.last_doc_opt, num_docs) {
|
||||
DocumentStep::Same | DocumentStep::Next => self.cardinality,
|
||||
DocumentStep::Skipped => self.cardinality.max(Cardinality::Optional),
|
||||
}
|
||||
}
|
||||
|
||||
/// Appends a new symbol to the `ColumnWriter`.
|
||||
fn write_symbol<V: SymbolValue>(
|
||||
&mut self,
|
||||
column_operation: ColumnOperation<V>,
|
||||
arena: &mut MemoryArena,
|
||||
) {
|
||||
self.values
|
||||
.writer(arena)
|
||||
.extend_from_slice(column_operation.serialize().as_ref());
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub(crate) struct NumericalColumnWriter {
|
||||
compatible_numerical_types: CompatibleNumericalTypes,
|
||||
column_writer: ColumnWriter,
|
||||
}
|
||||
|
||||
/// State used to store what types are still acceptable
|
||||
/// after having seen a set of numerical values.
|
||||
#[derive(Clone, Copy)]
|
||||
struct CompatibleNumericalTypes {
|
||||
all_values_within_i64_range: bool,
|
||||
all_values_within_u64_range: bool,
|
||||
// f64 is always acceptable.
|
||||
}
|
||||
|
||||
impl Default for CompatibleNumericalTypes {
|
||||
fn default() -> CompatibleNumericalTypes {
|
||||
CompatibleNumericalTypes {
|
||||
all_values_within_i64_range: true,
|
||||
all_values_within_u64_range: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CompatibleNumericalTypes {
|
||||
fn accept_value(&mut self, numerical_value: NumericalValue) {
|
||||
match numerical_value {
|
||||
NumericalValue::I64(val_i64) => {
|
||||
let value_within_u64_range = val_i64 >= 0i64;
|
||||
self.all_values_within_u64_range &= value_within_u64_range;
|
||||
}
|
||||
NumericalValue::U64(val_u64) => {
|
||||
let value_within_i64_range = val_u64 < i64::MAX as u64;
|
||||
self.all_values_within_i64_range &= value_within_i64_range;
|
||||
}
|
||||
NumericalValue::F64(_) => {
|
||||
self.all_values_within_i64_range = false;
|
||||
self.all_values_within_u64_range = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_numerical_type(self) -> NumericalType {
|
||||
if self.all_values_within_i64_range {
|
||||
NumericalType::I64
|
||||
} else if self.all_values_within_u64_range {
|
||||
NumericalType::U64
|
||||
} else {
|
||||
NumericalType::F64
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NumericalColumnWriter {
|
||||
pub fn column_type_and_cardinality(&self, num_docs: DocId) -> (NumericalType, Cardinality) {
|
||||
let numerical_type = self.compatible_numerical_types.to_numerical_type();
|
||||
let cardinality = self.column_writer.get_cardinality(num_docs);
|
||||
(numerical_type, cardinality)
|
||||
}
|
||||
|
||||
pub fn record_numerical_value(
|
||||
&mut self,
|
||||
doc: DocId,
|
||||
value: NumericalValue,
|
||||
arena: &mut MemoryArena,
|
||||
) {
|
||||
self.compatible_numerical_types.accept_value(value);
|
||||
self.column_writer.record(doc, value, arena);
|
||||
}
|
||||
|
||||
pub(super) fn operation_iterator<'a>(
|
||||
self,
|
||||
arena: &MemoryArena,
|
||||
buffer: &'a mut Vec<u8>,
|
||||
) -> impl Iterator<Item = ColumnOperation<NumericalValue>> + 'a {
|
||||
self.column_writer.operation_iterator(arena, buffer)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Default)]
|
||||
pub(crate) struct StrColumnWriter {
|
||||
pub(crate) dictionary_id: u32,
|
||||
pub(crate) column_writer: ColumnWriter,
|
||||
}
|
||||
|
||||
impl StrColumnWriter {
|
||||
pub(crate) fn with_dictionary_id(dictionary_id: u32) -> StrColumnWriter {
|
||||
StrColumnWriter {
|
||||
dictionary_id,
|
||||
column_writer: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn record_bytes(
|
||||
&mut self,
|
||||
doc: DocId,
|
||||
bytes: &[u8],
|
||||
dictionaries: &mut [DictionaryBuilder],
|
||||
arena: &mut MemoryArena,
|
||||
) {
|
||||
let unordered_id = dictionaries[self.dictionary_id as usize].get_or_allocate_id(bytes);
|
||||
self.column_writer.record(doc, unordered_id, arena);
|
||||
}
|
||||
|
||||
pub(super) fn operation_iterator<'a>(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
byte_buffer: &'a mut Vec<u8>,
|
||||
) -> impl Iterator<Item = ColumnOperation<UnorderedId>> + 'a {
|
||||
self.column_writer.operation_iterator(arena, byte_buffer)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_delta_with_last_doc() {
|
||||
assert_eq!(delta_with_last_doc(None, 0u32), DocumentStep::Next);
|
||||
assert_eq!(delta_with_last_doc(None, 1u32), DocumentStep::Skipped);
|
||||
assert_eq!(delta_with_last_doc(None, 2u32), DocumentStep::Skipped);
|
||||
assert_eq!(delta_with_last_doc(Some(0u32), 0u32), DocumentStep::Same);
|
||||
assert_eq!(delta_with_last_doc(Some(1u32), 1u32), DocumentStep::Same);
|
||||
assert_eq!(delta_with_last_doc(Some(1u32), 2u32), DocumentStep::Next);
|
||||
assert_eq!(delta_with_last_doc(Some(1u32), 3u32), DocumentStep::Skipped);
|
||||
assert_eq!(delta_with_last_doc(Some(1u32), 4u32), DocumentStep::Skipped);
|
||||
}
|
||||
|
||||
#[track_caller]
|
||||
fn test_column_writer_coercion_iter_aux(
|
||||
values: impl Iterator<Item = NumericalValue>,
|
||||
expected_numerical_type: NumericalType,
|
||||
) {
|
||||
let mut compatible_numerical_types = CompatibleNumericalTypes::default();
|
||||
for value in values {
|
||||
compatible_numerical_types.accept_value(value);
|
||||
}
|
||||
assert_eq!(
|
||||
compatible_numerical_types.to_numerical_type(),
|
||||
expected_numerical_type
|
||||
);
|
||||
}
|
||||
|
||||
#[track_caller]
|
||||
fn test_column_writer_coercion_aux(
|
||||
values: &[NumericalValue],
|
||||
expected_numerical_type: NumericalType,
|
||||
) {
|
||||
test_column_writer_coercion_iter_aux(values.iter().copied(), expected_numerical_type);
|
||||
test_column_writer_coercion_iter_aux(values.iter().rev().copied(), expected_numerical_type);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_writer_coercion() {
|
||||
test_column_writer_coercion_aux(&[], NumericalType::I64);
|
||||
test_column_writer_coercion_aux(&[1i64.into()], NumericalType::I64);
|
||||
test_column_writer_coercion_aux(&[1u64.into()], NumericalType::I64);
|
||||
// We don't detect exact integer at the moment. We could!
|
||||
test_column_writer_coercion_aux(&[1f64.into()], NumericalType::F64);
|
||||
test_column_writer_coercion_aux(&[u64::MAX.into()], NumericalType::U64);
|
||||
test_column_writer_coercion_aux(&[(i64::MAX as u64).into()], NumericalType::U64);
|
||||
test_column_writer_coercion_aux(&[(1u64 << 63).into()], NumericalType::U64);
|
||||
test_column_writer_coercion_aux(&[1i64.into(), 1u64.into()], NumericalType::I64);
|
||||
test_column_writer_coercion_aux(&[u64::MAX.into(), (-1i64).into()], NumericalType::F64);
|
||||
}
|
||||
}
|
||||
516
columnar/src/writer/mod.rs
Normal file
516
columnar/src/writer/mod.rs
Normal file
@@ -0,0 +1,516 @@
|
||||
mod column_operation;
|
||||
mod column_writers;
|
||||
mod serializer;
|
||||
mod value_index;
|
||||
|
||||
use std::io;
|
||||
|
||||
use column_operation::ColumnOperation;
|
||||
use common::CountingWriter;
|
||||
use fastfield_codecs::serialize::ValueIndexInfo;
|
||||
use fastfield_codecs::{Column, MonotonicallyMappableToU64, VecColumn};
|
||||
use serializer::ColumnarSerializer;
|
||||
use stacker::{Addr, ArenaHashMap, MemoryArena};
|
||||
|
||||
use crate::column_type_header::{ColumnType, ColumnTypeAndCardinality, ColumnTypeCategory};
|
||||
use crate::dictionary::{DictionaryBuilder, TermIdMapping, UnorderedId};
|
||||
use crate::value::{Coerce, NumericalType, NumericalValue};
|
||||
use crate::writer::column_writers::{ColumnWriter, NumericalColumnWriter, StrColumnWriter};
|
||||
use crate::writer::value_index::{IndexBuilder, SpareIndexBuilders};
|
||||
use crate::{Cardinality, DocId};
|
||||
|
||||
/// This is a set of buffers that are used to temporarily write the values into before passing them
|
||||
/// to the fast field codecs.
|
||||
#[derive(Default)]
|
||||
struct SpareBuffers {
|
||||
value_index_builders: SpareIndexBuilders,
|
||||
i64_values: Vec<i64>,
|
||||
u64_values: Vec<u64>,
|
||||
f64_values: Vec<f64>,
|
||||
bool_values: Vec<bool>,
|
||||
}
|
||||
|
||||
/// Makes it possible to create a new columnar.
|
||||
///
|
||||
/// ```rust
|
||||
/// use tantivy_columnar::ColumnarWriter;
|
||||
///
|
||||
/// let mut columnar_writer = ColumnarWriter::default();
|
||||
/// columnar_writer.record_str(0u32 /* doc id */, "product_name", "Red backpack");
|
||||
/// columnar_writer.record_numerical(0u32 /* doc id */, "price", 10u64);
|
||||
/// columnar_writer.record_str(1u32 /* doc id */, "product_name", "Apple");
|
||||
/// columnar_writer.record_numerical(0u32 /* doc id */, "price", 10.5f64); //< uh oh we ended up mixing integer and floats.
|
||||
/// let mut wrt: Vec<u8> = Vec::new();
|
||||
/// columnar_writer.serialize(2u32, &mut wrt).unwrap();
|
||||
/// ```
|
||||
pub struct ColumnarWriter {
|
||||
numerical_field_hash_map: ArenaHashMap,
|
||||
bool_field_hash_map: ArenaHashMap,
|
||||
bytes_field_hash_map: ArenaHashMap,
|
||||
arena: MemoryArena,
|
||||
// Dictionaries used to store dictionary-encoded values.
|
||||
dictionaries: Vec<DictionaryBuilder>,
|
||||
buffers: SpareBuffers,
|
||||
}
|
||||
|
||||
impl Default for ColumnarWriter {
|
||||
fn default() -> Self {
|
||||
ColumnarWriter {
|
||||
numerical_field_hash_map: ArenaHashMap::new(10_000),
|
||||
bool_field_hash_map: ArenaHashMap::new(10_000),
|
||||
bytes_field_hash_map: ArenaHashMap::new(10_000),
|
||||
dictionaries: Vec::new(),
|
||||
arena: MemoryArena::default(),
|
||||
buffers: SpareBuffers::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ColumnarWriter {
|
||||
pub fn record_numerical<T: Into<NumericalValue> + Copy>(
|
||||
&mut self,
|
||||
doc: DocId,
|
||||
column_name: &str,
|
||||
numerical_value: T,
|
||||
) {
|
||||
assert!(
|
||||
!column_name.as_bytes().contains(&0u8),
|
||||
"key may not contain the 0 byte"
|
||||
);
|
||||
let (hash_map, arena) = (&mut self.numerical_field_hash_map, &mut self.arena);
|
||||
hash_map.mutate_or_create(
|
||||
column_name.as_bytes(),
|
||||
|column_opt: Option<NumericalColumnWriter>| {
|
||||
let mut column: NumericalColumnWriter = column_opt.unwrap_or_default();
|
||||
column.record_numerical_value(doc, numerical_value.into(), arena);
|
||||
column
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn record_bool(&mut self, doc: DocId, column_name: &str, val: bool) {
|
||||
assert!(
|
||||
!column_name.as_bytes().contains(&0u8),
|
||||
"key may not contain the 0 byte"
|
||||
);
|
||||
let (hash_map, arena) = (&mut self.bool_field_hash_map, &mut self.arena);
|
||||
hash_map.mutate_or_create(
|
||||
column_name.as_bytes(),
|
||||
|column_opt: Option<ColumnWriter>| {
|
||||
let mut column: ColumnWriter = column_opt.unwrap_or_default();
|
||||
column.record(doc, val, arena);
|
||||
column
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn record_str(&mut self, doc: DocId, column_name: &str, value: &str) {
|
||||
assert!(
|
||||
!column_name.as_bytes().contains(&0u8),
|
||||
"key may not contain the 0 byte"
|
||||
);
|
||||
let (hash_map, arena, dictionaries) = (
|
||||
&mut self.bytes_field_hash_map,
|
||||
&mut self.arena,
|
||||
&mut self.dictionaries,
|
||||
);
|
||||
hash_map.mutate_or_create(
|
||||
column_name.as_bytes(),
|
||||
|column_opt: Option<StrColumnWriter>| {
|
||||
let mut column: StrColumnWriter = column_opt.unwrap_or_else(|| {
|
||||
// Each column has its own dictionary
|
||||
let dictionary_id = dictionaries.len() as u32;
|
||||
dictionaries.push(DictionaryBuilder::default());
|
||||
StrColumnWriter::with_dictionary_id(dictionary_id)
|
||||
});
|
||||
column.record_bytes(doc, value.as_bytes(), dictionaries, arena);
|
||||
column
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn serialize(&mut self, num_docs: DocId, wrt: &mut dyn io::Write) -> io::Result<()> {
|
||||
let mut serializer = ColumnarSerializer::new(wrt);
|
||||
let mut field_columns: Vec<(&[u8], ColumnTypeCategory, Addr)> = self
|
||||
.numerical_field_hash_map
|
||||
.iter()
|
||||
.map(|(term, addr, _)| (term, ColumnTypeCategory::Numerical, addr))
|
||||
.collect();
|
||||
field_columns.extend(
|
||||
self.bytes_field_hash_map
|
||||
.iter()
|
||||
.map(|(term, addr, _)| (term, ColumnTypeCategory::Str, addr)),
|
||||
);
|
||||
field_columns.extend(
|
||||
self.bool_field_hash_map
|
||||
.iter()
|
||||
.map(|(term, addr, _)| (term, ColumnTypeCategory::Bool, addr)),
|
||||
);
|
||||
field_columns.sort_unstable_by_key(|(column_name, col_type, _)| (*column_name, *col_type));
|
||||
let (arena, buffers, dictionaries) = (&self.arena, &mut self.buffers, &self.dictionaries);
|
||||
let mut symbol_byte_buffer: Vec<u8> = Vec::new();
|
||||
for (column_name, bytes_or_numerical, addr) in field_columns {
|
||||
match bytes_or_numerical {
|
||||
ColumnTypeCategory::Bool => {
|
||||
let column_writer: ColumnWriter = self.bool_field_hash_map.read(addr);
|
||||
let cardinality = column_writer.get_cardinality(num_docs);
|
||||
let column_type_and_cardinality = ColumnTypeAndCardinality {
|
||||
cardinality,
|
||||
typ: ColumnType::Bool,
|
||||
};
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name, column_type_and_cardinality);
|
||||
serialize_bool_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
ColumnTypeCategory::Str => {
|
||||
let str_column_writer: StrColumnWriter = self.bytes_field_hash_map.read(addr);
|
||||
let dictionary_builder =
|
||||
&dictionaries[str_column_writer.dictionary_id as usize];
|
||||
let cardinality = str_column_writer.column_writer.get_cardinality(num_docs);
|
||||
let column_type_and_cardinality = ColumnTypeAndCardinality {
|
||||
cardinality,
|
||||
typ: ColumnType::Bytes,
|
||||
};
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name, column_type_and_cardinality);
|
||||
serialize_bytes_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
dictionary_builder,
|
||||
str_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
ColumnTypeCategory::Numerical => {
|
||||
let numerical_column_writer: NumericalColumnWriter =
|
||||
self.numerical_field_hash_map.read(addr);
|
||||
let (numerical_type, cardinality) =
|
||||
numerical_column_writer.column_type_and_cardinality(num_docs);
|
||||
let column_type_and_cardinality = ColumnTypeAndCardinality {
|
||||
cardinality,
|
||||
typ: ColumnType::Numerical(numerical_type),
|
||||
};
|
||||
let mut column_serializer =
|
||||
serializer.serialize_column(column_name, column_type_and_cardinality);
|
||||
serialize_numerical_column(
|
||||
cardinality,
|
||||
num_docs,
|
||||
numerical_type,
|
||||
numerical_column_writer.operation_iterator(arena, &mut symbol_byte_buffer),
|
||||
buffers,
|
||||
&mut column_serializer,
|
||||
)?;
|
||||
}
|
||||
};
|
||||
}
|
||||
serializer.finalize()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn serialize_bytes_column(
|
||||
cardinality: Cardinality,
|
||||
num_docs: DocId,
|
||||
dictionary_builder: &DictionaryBuilder,
|
||||
operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>,
|
||||
buffers: &mut SpareBuffers,
|
||||
wrt: impl io::Write,
|
||||
) -> io::Result<()> {
|
||||
let SpareBuffers {
|
||||
value_index_builders,
|
||||
u64_values,
|
||||
..
|
||||
} = buffers;
|
||||
let mut counting_writer = CountingWriter::wrap(wrt);
|
||||
let term_id_mapping: TermIdMapping = dictionary_builder.serialize(&mut counting_writer)?;
|
||||
let dictionary_num_bytes: u32 = counting_writer.written_bytes() as u32;
|
||||
let mut wrt = counting_writer.finish();
|
||||
let operation_iterator = operation_it.map(|symbol: ColumnOperation<UnorderedId>| {
|
||||
// We map unordered ids to ordered ids.
|
||||
match symbol {
|
||||
ColumnOperation::Value(unordered_id) => {
|
||||
let ordered_id = term_id_mapping.to_ord(unordered_id);
|
||||
ColumnOperation::Value(ordered_id.0 as u64)
|
||||
}
|
||||
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
|
||||
}
|
||||
});
|
||||
serialize_column(
|
||||
operation_iterator,
|
||||
cardinality,
|
||||
num_docs,
|
||||
value_index_builders,
|
||||
u64_values,
|
||||
&mut wrt,
|
||||
)?;
|
||||
wrt.write_all(&dictionary_num_bytes.to_le_bytes()[..])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn serialize_numerical_column(
|
||||
cardinality: Cardinality,
|
||||
num_docs: DocId,
|
||||
numerical_type: NumericalType,
|
||||
op_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
|
||||
buffers: &mut SpareBuffers,
|
||||
wrt: &mut impl io::Write,
|
||||
) -> io::Result<()> {
|
||||
let SpareBuffers {
|
||||
value_index_builders,
|
||||
u64_values,
|
||||
i64_values,
|
||||
f64_values,
|
||||
..
|
||||
} = buffers;
|
||||
match numerical_type {
|
||||
NumericalType::I64 => {
|
||||
serialize_column(
|
||||
coerce_numerical_symbol::<i64>(op_iterator),
|
||||
cardinality,
|
||||
num_docs,
|
||||
value_index_builders,
|
||||
i64_values,
|
||||
wrt,
|
||||
)?;
|
||||
}
|
||||
NumericalType::U64 => {
|
||||
serialize_column(
|
||||
coerce_numerical_symbol::<u64>(op_iterator),
|
||||
cardinality,
|
||||
num_docs,
|
||||
value_index_builders,
|
||||
u64_values,
|
||||
wrt,
|
||||
)?;
|
||||
}
|
||||
NumericalType::F64 => {
|
||||
serialize_column(
|
||||
coerce_numerical_symbol::<f64>(op_iterator),
|
||||
cardinality,
|
||||
num_docs,
|
||||
value_index_builders,
|
||||
f64_values,
|
||||
wrt,
|
||||
)?;
|
||||
}
|
||||
};
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn serialize_bool_column(
|
||||
cardinality: Cardinality,
|
||||
num_docs: DocId,
|
||||
column_operations_it: impl Iterator<Item = ColumnOperation<bool>>,
|
||||
buffers: &mut SpareBuffers,
|
||||
wrt: &mut impl io::Write,
|
||||
) -> io::Result<()> {
|
||||
let SpareBuffers {
|
||||
value_index_builders,
|
||||
bool_values,
|
||||
..
|
||||
} = buffers;
|
||||
serialize_column(
|
||||
column_operations_it,
|
||||
cardinality,
|
||||
num_docs,
|
||||
value_index_builders,
|
||||
bool_values,
|
||||
wrt,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn serialize_column<
|
||||
T: Copy + Default + std::fmt::Debug + Send + Sync + MonotonicallyMappableToU64 + PartialOrd,
|
||||
>(
|
||||
op_iterator: impl Iterator<Item = ColumnOperation<T>>,
|
||||
cardinality: Cardinality,
|
||||
num_docs: DocId,
|
||||
value_index_builders: &mut SpareIndexBuilders,
|
||||
values: &mut Vec<T>,
|
||||
mut wrt: impl io::Write,
|
||||
) -> io::Result<()>
|
||||
where
|
||||
for<'a> VecColumn<'a, T>: Column<T>,
|
||||
{
|
||||
values.clear();
|
||||
match cardinality {
|
||||
Cardinality::Required => {
|
||||
consume_operation_iterator(
|
||||
op_iterator,
|
||||
value_index_builders.borrow_required_index_builder(),
|
||||
values,
|
||||
);
|
||||
fastfield_codecs::serialize(
|
||||
VecColumn::from(&values[..]),
|
||||
&mut wrt,
|
||||
&fastfield_codecs::ALL_CODEC_TYPES[..],
|
||||
)?;
|
||||
}
|
||||
Cardinality::Optional => {
|
||||
let optional_index_builder = value_index_builders.borrow_optional_index_builder();
|
||||
consume_operation_iterator(op_iterator, optional_index_builder, values);
|
||||
let optional_index = optional_index_builder.finish(num_docs);
|
||||
fastfield_codecs::serialize::serialize_new(
|
||||
ValueIndexInfo::SingleValue(Box::new(optional_index)),
|
||||
VecColumn::from(&values[..]),
|
||||
&mut wrt,
|
||||
&fastfield_codecs::ALL_CODEC_TYPES[..],
|
||||
)?;
|
||||
}
|
||||
Cardinality::Multivalued => {
|
||||
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
|
||||
consume_operation_iterator(op_iterator, multivalued_index_builder, values);
|
||||
let multivalued_index = multivalued_index_builder.finish(num_docs);
|
||||
fastfield_codecs::serialize::serialize_new(
|
||||
ValueIndexInfo::MultiValue(Box::new(multivalued_index)),
|
||||
VecColumn::from(&values[..]),
|
||||
&mut wrt,
|
||||
&fastfield_codecs::ALL_CODEC_TYPES[..],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn coerce_numerical_symbol<T>(
|
||||
operation_iterator: impl Iterator<Item = ColumnOperation<NumericalValue>>,
|
||||
) -> impl Iterator<Item = ColumnOperation<T>>
|
||||
where T: Coerce {
|
||||
operation_iterator.map(|symbol| match symbol {
|
||||
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
|
||||
ColumnOperation::Value(numerical_value) => {
|
||||
ColumnOperation::Value(Coerce::coerce(numerical_value))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn consume_operation_iterator<T: std::fmt::Debug, TIndexBuilder: IndexBuilder>(
|
||||
operation_iterator: impl Iterator<Item = ColumnOperation<T>>,
|
||||
index_builder: &mut TIndexBuilder,
|
||||
values: &mut Vec<T>,
|
||||
) {
|
||||
for symbol in operation_iterator {
|
||||
match symbol {
|
||||
ColumnOperation::NewDoc(doc) => {
|
||||
index_builder.record_doc(doc);
|
||||
}
|
||||
ColumnOperation::Value(value) => {
|
||||
index_builder.record_value();
|
||||
values.push(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use column_operation::ColumnOperation;
|
||||
use stacker::MemoryArena;
|
||||
|
||||
use super::*;
|
||||
use crate::value::NumericalValue;
|
||||
use crate::Cardinality;
|
||||
|
||||
#[test]
|
||||
fn test_column_writer_required_simple() {
|
||||
let mut arena = MemoryArena::default();
|
||||
let mut column_writer = super::ColumnWriter::default();
|
||||
column_writer.record(0u32, NumericalValue::from(14i64), &mut arena);
|
||||
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
|
||||
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
|
||||
assert_eq!(column_writer.get_cardinality(3), Cardinality::Required);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 6);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
|
||||
assert!(matches!(
|
||||
symbols[1],
|
||||
ColumnOperation::Value(NumericalValue::I64(14i64))
|
||||
));
|
||||
assert!(matches!(symbols[2], ColumnOperation::NewDoc(1u32)));
|
||||
assert!(matches!(
|
||||
symbols[3],
|
||||
ColumnOperation::Value(NumericalValue::I64(15i64))
|
||||
));
|
||||
assert!(matches!(symbols[4], ColumnOperation::NewDoc(2u32)));
|
||||
assert!(matches!(
|
||||
symbols[5],
|
||||
ColumnOperation::Value(NumericalValue::I64(-16i64))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_writer_optional_cardinality_missing_first() {
|
||||
let mut arena = MemoryArena::default();
|
||||
let mut column_writer = super::ColumnWriter::default();
|
||||
column_writer.record(1u32, NumericalValue::from(15i64), &mut arena);
|
||||
column_writer.record(2u32, NumericalValue::from(-16i64), &mut arena);
|
||||
assert_eq!(column_writer.get_cardinality(3), Cardinality::Optional);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 4);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(1u32)));
|
||||
assert!(matches!(
|
||||
symbols[1],
|
||||
ColumnOperation::Value(NumericalValue::I64(15i64))
|
||||
));
|
||||
assert!(matches!(symbols[2], ColumnOperation::NewDoc(2u32)));
|
||||
assert!(matches!(
|
||||
symbols[3],
|
||||
ColumnOperation::Value(NumericalValue::I64(-16i64))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_writer_optional_cardinality_missing_last() {
|
||||
let mut arena = MemoryArena::default();
|
||||
let mut column_writer = super::ColumnWriter::default();
|
||||
column_writer.record(0u32, NumericalValue::from(15i64), &mut arena);
|
||||
assert_eq!(column_writer.get_cardinality(2), Cardinality::Optional);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 2);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
|
||||
assert!(matches!(
|
||||
symbols[1],
|
||||
ColumnOperation::Value(NumericalValue::I64(15i64))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_column_writer_multivalued() {
|
||||
let mut arena = MemoryArena::default();
|
||||
let mut column_writer = super::ColumnWriter::default();
|
||||
column_writer.record(0u32, NumericalValue::from(16i64), &mut arena);
|
||||
column_writer.record(0u32, NumericalValue::from(17i64), &mut arena);
|
||||
assert_eq!(column_writer.get_cardinality(1), Cardinality::Multivalued);
|
||||
let mut buffer = Vec::new();
|
||||
let symbols: Vec<ColumnOperation<NumericalValue>> = column_writer
|
||||
.operation_iterator(&mut arena, &mut buffer)
|
||||
.collect();
|
||||
assert_eq!(symbols.len(), 3);
|
||||
assert!(matches!(symbols[0], ColumnOperation::NewDoc(0u32)));
|
||||
assert!(matches!(
|
||||
symbols[1],
|
||||
ColumnOperation::Value(NumericalValue::I64(16i64))
|
||||
));
|
||||
assert!(matches!(
|
||||
symbols[2],
|
||||
ColumnOperation::Value(NumericalValue::I64(17i64))
|
||||
));
|
||||
}
|
||||
}
|
||||
116
columnar/src/writer/serializer.rs
Normal file
116
columnar/src/writer/serializer.rs
Normal file
@@ -0,0 +1,116 @@
|
||||
use std::io;
|
||||
use std::io::Write;
|
||||
|
||||
use common::CountingWriter;
|
||||
use sstable::value::RangeValueWriter;
|
||||
use sstable::RangeSSTable;
|
||||
|
||||
use crate::column_type_header::ColumnTypeAndCardinality;
|
||||
|
||||
pub struct ColumnarSerializer<W: io::Write> {
|
||||
wrt: CountingWriter<W>,
|
||||
sstable_range: sstable::Writer<Vec<u8>, RangeValueWriter>,
|
||||
prepare_key_buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Returns a key consisting of the concatenation of the key and the column_type_and_cardinality
|
||||
/// code.
|
||||
fn prepare_key(
|
||||
key: &[u8],
|
||||
column_type_cardinality: ColumnTypeAndCardinality,
|
||||
buffer: &mut Vec<u8>,
|
||||
) {
|
||||
buffer.clear();
|
||||
buffer.extend_from_slice(key);
|
||||
buffer.push(0u8);
|
||||
buffer.push(column_type_cardinality.to_code());
|
||||
}
|
||||
|
||||
impl<W: io::Write> ColumnarSerializer<W> {
|
||||
pub(crate) fn new(wrt: W) -> ColumnarSerializer<W> {
|
||||
let sstable_range: sstable::Writer<Vec<u8>, RangeValueWriter> =
|
||||
sstable::Dictionary::<RangeSSTable>::builder(Vec::with_capacity(100_000)).unwrap();
|
||||
ColumnarSerializer {
|
||||
wrt: CountingWriter::wrap(wrt),
|
||||
sstable_range,
|
||||
prepare_key_buffer: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn serialize_column<'a>(
|
||||
&'a mut self,
|
||||
column_name: &[u8],
|
||||
column_type_cardinality: ColumnTypeAndCardinality,
|
||||
) -> impl io::Write + 'a {
|
||||
let start_offset = self.wrt.written_bytes();
|
||||
prepare_key(
|
||||
column_name,
|
||||
column_type_cardinality,
|
||||
&mut self.prepare_key_buffer,
|
||||
);
|
||||
ColumnSerializer {
|
||||
columnar_serializer: self,
|
||||
start_offset,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn finalize(mut self) -> io::Result<()> {
|
||||
let sstable_bytes: Vec<u8> = self.sstable_range.finish()?;
|
||||
let sstable_num_bytes: u64 = sstable_bytes.len() as u64;
|
||||
self.wrt.write_all(&sstable_bytes)?;
|
||||
self.wrt.write_all(&sstable_num_bytes.to_le_bytes()[..])?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
struct ColumnSerializer<'a, W: io::Write> {
|
||||
columnar_serializer: &'a mut ColumnarSerializer<W>,
|
||||
start_offset: u64,
|
||||
}
|
||||
|
||||
impl<'a, W: io::Write> Drop for ColumnSerializer<'a, W> {
|
||||
fn drop(&mut self) {
|
||||
let end_offset: u64 = self.columnar_serializer.wrt.written_bytes();
|
||||
let byte_range = self.start_offset..end_offset;
|
||||
self.columnar_serializer.sstable_range.insert_cannot_fail(
|
||||
&self.columnar_serializer.prepare_key_buffer[..],
|
||||
&byte_range,
|
||||
);
|
||||
self.columnar_serializer.prepare_key_buffer.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, W: io::Write> io::Write for ColumnSerializer<'a, W> {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.columnar_serializer.wrt.write(buf)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.columnar_serializer.wrt.flush()
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.columnar_serializer.wrt.write_all(buf)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::column_type_header::ColumnType;
|
||||
use crate::Cardinality;
|
||||
|
||||
#[test]
|
||||
fn test_prepare_key_bytes() {
|
||||
let mut buffer: Vec<u8> = b"somegarbage".to_vec();
|
||||
let column_type_and_cardinality = ColumnTypeAndCardinality {
|
||||
typ: ColumnType::Bytes,
|
||||
cardinality: Cardinality::Optional,
|
||||
};
|
||||
prepare_key(b"root\0child", column_type_and_cardinality, &mut buffer);
|
||||
assert_eq!(buffer.len(), 12);
|
||||
assert_eq!(&buffer[..10], b"root\0child");
|
||||
assert_eq!(buffer[10], 0u8);
|
||||
assert_eq!(buffer[11], column_type_and_cardinality.to_code());
|
||||
}
|
||||
}
|
||||
220
columnar/src/writer/value_index.rs
Normal file
220
columnar/src/writer/value_index.rs
Normal file
@@ -0,0 +1,220 @@
|
||||
use fastfield_codecs::serialize::{MultiValueIndexInfo, SingleValueIndexInfo};
|
||||
|
||||
use crate::DocId;
|
||||
|
||||
/// The `IndexBuilder` interprets a sequence of
|
||||
/// calls of the form:
|
||||
/// (record_doc,record_value+)*
|
||||
/// and can then serialize the results into an index to associate docids with their value[s].
|
||||
///
|
||||
/// It has different implementation depending on whether the
|
||||
/// cardinality is required, optional, or multivalued.
|
||||
pub(crate) trait IndexBuilder {
|
||||
fn record_doc(&mut self, doc: DocId);
|
||||
#[inline]
|
||||
fn record_value(&mut self) {}
|
||||
}
|
||||
|
||||
/// The RequiredIndexBuilder does nothing.
|
||||
#[derive(Default)]
|
||||
pub struct RequiredIndexBuilder;
|
||||
|
||||
impl IndexBuilder for RequiredIndexBuilder {
|
||||
#[inline(always)]
|
||||
fn record_doc(&mut self, _doc: DocId) {}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct OptionalIndexBuilder {
|
||||
docs: Vec<DocId>,
|
||||
}
|
||||
|
||||
struct SingleValueArrayIndex<'a> {
|
||||
// DocIds with a value. DocIds are strictly increasing
|
||||
docs: &'a [DocId],
|
||||
num_docs: DocId,
|
||||
}
|
||||
|
||||
impl<'a> SingleValueIndexInfo for SingleValueArrayIndex<'a> {
|
||||
fn num_vals(&self) -> u32 {
|
||||
self.num_docs as u32
|
||||
}
|
||||
|
||||
fn num_non_nulls(&self) -> u32 {
|
||||
self.docs.len() as u32
|
||||
}
|
||||
|
||||
fn iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
|
||||
Box::new(self.docs.iter().copied())
|
||||
}
|
||||
}
|
||||
|
||||
impl OptionalIndexBuilder {
|
||||
pub fn finish(&mut self, num_docs: DocId) -> impl SingleValueIndexInfo + '_ {
|
||||
debug_assert!(self
|
||||
.docs
|
||||
.last()
|
||||
.copied()
|
||||
.map(|last_doc| last_doc < num_docs)
|
||||
.unwrap_or(true));
|
||||
SingleValueArrayIndex {
|
||||
docs: &self.docs[..],
|
||||
num_docs,
|
||||
}
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.docs.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl IndexBuilder for OptionalIndexBuilder {
|
||||
#[inline(always)]
|
||||
fn record_doc(&mut self, doc: DocId) {
|
||||
debug_assert!(self
|
||||
.docs
|
||||
.last()
|
||||
.copied()
|
||||
.map(|prev_doc| doc > prev_doc)
|
||||
.unwrap_or(true));
|
||||
self.docs.push(doc);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct MultivaluedIndexBuilder {
|
||||
// TODO should we switch to `start_offset`?
|
||||
// contains the num values so far for each `DocId`.
|
||||
end_offsets: Vec<DocId>,
|
||||
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 {
|
||||
pub fn finish(&mut self, num_docs: DocId) -> impl MultiValueIndexInfo + '_ {
|
||||
self.end_offsets
|
||||
.resize(num_docs as usize, self.total_num_vals_seen);
|
||||
MultivaluedValueArrayIndex {
|
||||
end_offsets: &self.end_offsets[..],
|
||||
}
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.end_offsets.clear();
|
||||
self.total_num_vals_seen = 0;
|
||||
}
|
||||
}
|
||||
|
||||
impl IndexBuilder for MultivaluedIndexBuilder {
|
||||
fn record_doc(&mut self, doc: DocId) {
|
||||
self.end_offsets
|
||||
.resize(doc as usize, self.total_num_vals_seen);
|
||||
}
|
||||
|
||||
fn record_value(&mut self) {
|
||||
self.total_num_vals_seen += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// The `SpareIndexBuilders` is there to avoid allocating a
|
||||
/// new index builder for every single column.
|
||||
#[derive(Default)]
|
||||
pub struct SpareIndexBuilders {
|
||||
required_index_builder: RequiredIndexBuilder,
|
||||
optional_index_builder: OptionalIndexBuilder,
|
||||
multivalued_index_builder: MultivaluedIndexBuilder,
|
||||
}
|
||||
|
||||
impl SpareIndexBuilders {
|
||||
pub fn borrow_required_index_builder(&mut self) -> &mut RequiredIndexBuilder {
|
||||
&mut self.required_index_builder
|
||||
}
|
||||
|
||||
pub fn borrow_optional_index_builder(&mut self) -> &mut OptionalIndexBuilder {
|
||||
self.optional_index_builder.reset();
|
||||
&mut self.optional_index_builder
|
||||
}
|
||||
|
||||
pub fn borrow_multivalued_index_builder(&mut self) -> &mut MultivaluedIndexBuilder {
|
||||
self.multivalued_index_builder.reset();
|
||||
&mut self.multivalued_index_builder
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_optional_value_index_builder() {
|
||||
let mut opt_value_index_builder = OptionalIndexBuilder::default();
|
||||
opt_value_index_builder.record_doc(0u32);
|
||||
opt_value_index_builder.record_value();
|
||||
assert_eq!(
|
||||
&opt_value_index_builder
|
||||
.finish(1u32)
|
||||
.iter()
|
||||
.collect::<Vec<u32>>(),
|
||||
&[0]
|
||||
);
|
||||
opt_value_index_builder.reset();
|
||||
opt_value_index_builder.record_doc(1u32);
|
||||
opt_value_index_builder.record_value();
|
||||
assert_eq!(
|
||||
&opt_value_index_builder
|
||||
.finish(2u32)
|
||||
.iter()
|
||||
.collect::<Vec<u32>>(),
|
||||
&[1]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multivalued_value_index_builder() {
|
||||
let mut multivalued_value_index_builder = MultivaluedIndexBuilder::default();
|
||||
multivalued_value_index_builder.record_doc(1u32);
|
||||
multivalued_value_index_builder.record_value();
|
||||
multivalued_value_index_builder.record_value();
|
||||
multivalued_value_index_builder.record_doc(2u32);
|
||||
multivalued_value_index_builder.record_value();
|
||||
assert_eq!(
|
||||
multivalued_value_index_builder
|
||||
.finish(4u32)
|
||||
.iter()
|
||||
.collect::<Vec<u32>>(),
|
||||
vec![0, 0, 2, 3]
|
||||
);
|
||||
multivalued_value_index_builder.reset();
|
||||
multivalued_value_index_builder.record_doc(2u32);
|
||||
multivalued_value_index_builder.record_value();
|
||||
multivalued_value_index_builder.record_value();
|
||||
assert_eq!(
|
||||
multivalued_value_index_builder
|
||||
.finish(4u32)
|
||||
.iter()
|
||||
.collect::<Vec<u32>>(),
|
||||
vec![0, 0, 0, 2]
|
||||
);
|
||||
}
|
||||
}
|
||||
166
common/src/group_by.rs
Normal file
166
common/src/group_by.rs
Normal 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])]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -6,10 +6,12 @@ pub use byteorder::LittleEndian as Endianness;
|
||||
|
||||
mod bitset;
|
||||
pub mod file_slice;
|
||||
mod group_by;
|
||||
mod serialize;
|
||||
mod vint;
|
||||
mod writer;
|
||||
pub use bitset::*;
|
||||
pub use group_by::GroupByIteratorExtended;
|
||||
pub use ownedbytes::{OwnedBytes, StableDeref};
|
||||
pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize};
|
||||
pub use vint::{
|
||||
|
||||
@@ -1,15 +1,17 @@
|
||||
// # Basic Example
|
||||
// # Faceted Search
|
||||
//
|
||||
// This example covers the basic functionalities of
|
||||
// This example covers the faceted search functionalities of
|
||||
// tantivy.
|
||||
//
|
||||
// We will :
|
||||
// - define our schema
|
||||
// = create an index in a directory
|
||||
// - index few documents in our index
|
||||
// - search for the best document matchings "sea whale"
|
||||
// - retrieve the best document original content.
|
||||
|
||||
// - define a text field "name" in our schema
|
||||
// - define a facet field "classification" in our schema
|
||||
// - create an index in memory
|
||||
// - index few documents with respective facets in our index
|
||||
// - 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...
|
||||
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 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
|
||||
let classification = schema_builder.add_facet_field("classification", FacetOptions::default());
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ repository = "https://github.com/quickwit-oss/tantivy"
|
||||
[dependencies]
|
||||
common = { version = "0.5", path = "../common/", package = "tantivy-common" }
|
||||
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}
|
||||
fastdivide = "0.4"
|
||||
log = "0.4"
|
||||
|
||||
@@ -4,7 +4,7 @@ extern crate test;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::iter;
|
||||
use std::ops::RangeInclusive;
|
||||
use std::sync::Arc;
|
||||
|
||||
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![];
|
||||
for i in 0..100 {
|
||||
let num = i * i;
|
||||
data.extend(iter::repeat(i as u64).take(num));
|
||||
for _ in 0..300_000 {
|
||||
let val = rng.gen_range(1..=100);
|
||||
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
|
||||
}
|
||||
|
||||
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>> {
|
||||
let permutation = generate_random();
|
||||
let permutation = permutation.iter().map(|el| *el as u128).collect::<Vec<_>>();
|
||||
@@ -106,15 +103,82 @@ mod tests {
|
||||
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]
|
||||
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);
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
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,
|
||||
&mut positions,
|
||||
);
|
||||
@@ -124,13 +188,13 @@ mod tests {
|
||||
|
||||
#[bench]
|
||||
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);
|
||||
|
||||
b.iter(|| {
|
||||
let mut positions = Vec::new();
|
||||
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,
|
||||
&mut positions,
|
||||
);
|
||||
@@ -140,7 +204,7 @@ mod tests {
|
||||
|
||||
#[bench]
|
||||
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);
|
||||
|
||||
b.iter(|| {
|
||||
@@ -149,6 +213,7 @@ mod tests {
|
||||
positions
|
||||
});
|
||||
}
|
||||
// U128 RANGE END
|
||||
|
||||
#[bench]
|
||||
fn bench_intfastfield_scan_all_fflookup_u128(b: &mut Bencher) {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::fmt::{self, Debug};
|
||||
use std::marker::PhantomData;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
|
||||
@@ -6,7 +7,7 @@ use tantivy_bitpacker::minmax;
|
||||
use crate::monotonic_mapping::StrictlyMonotonicFn;
|
||||
|
||||
/// `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.
|
||||
///
|
||||
/// This accessor should return as fast as possible.
|
||||
@@ -83,7 +84,7 @@ pub struct VecColumn<'a, T = u64> {
|
||||
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 {
|
||||
(*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 {
|
||||
self.values[position as usize]
|
||||
}
|
||||
@@ -177,8 +178,8 @@ pub fn monotonic_map_column<C, T, Input, Output>(
|
||||
where
|
||||
C: Column<Input>,
|
||||
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
|
||||
Input: PartialOrd + Send + Sync + Clone,
|
||||
Output: PartialOrd + Send + Sync + Clone,
|
||||
Input: PartialOrd + Send + Sync + Copy + Debug,
|
||||
Output: PartialOrd + Send + Sync + Copy + Debug,
|
||||
{
|
||||
MonotonicMappingColumn {
|
||||
from_column,
|
||||
@@ -191,8 +192,8 @@ impl<C, T, Input, Output> Column<Output> for MonotonicMappingColumn<C, T, Input>
|
||||
where
|
||||
C: Column<Input>,
|
||||
T: StrictlyMonotonicFn<Input, Output> + Send + Sync,
|
||||
Input: PartialOrd + Send + Sync + Clone,
|
||||
Output: PartialOrd + Send + Sync + Clone,
|
||||
Input: PartialOrd + Send + Sync + Copy + Debug,
|
||||
Output: PartialOrd + Send + Sync + Copy + Debug,
|
||||
{
|
||||
#[inline]
|
||||
fn get_val(&self, idx: u32) -> Output {
|
||||
@@ -228,12 +229,15 @@ where
|
||||
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,
|
||||
)
|
||||
if range.start() > &self.max_value() || range.end() < &self.min_value() {
|
||||
return;
|
||||
}
|
||||
let range = self.monotonic_mapping.inverse_coerce(range);
|
||||
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,
|
||||
@@ -254,7 +258,7 @@ where T: Iterator + Clone + ExactSizeIterator
|
||||
impl<T> Column<T::Item> for IterColumn<T>
|
||||
where
|
||||
T: Iterator + Clone + ExactSizeIterator + Send + Sync,
|
||||
T::Item: PartialOrd,
|
||||
T::Item: PartialOrd + fmt::Debug,
|
||||
{
|
||||
fn get_val(&self, idx: u32) -> T::Item {
|
||||
self.0.clone().nth(idx as usize).unwrap()
|
||||
|
||||
@@ -453,6 +453,8 @@ impl CompactSpaceDecompressor {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use super::*;
|
||||
use crate::format_version::read_format_version;
|
||||
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,
|
||||
value_range: RangeInclusive<T>,
|
||||
doc_id_range: Range<u32>,
|
||||
|
||||
@@ -14,9 +14,9 @@ extern crate more_asserts;
|
||||
#[cfg(all(test, feature = "unstable"))]
|
||||
extern crate test;
|
||||
|
||||
use std::io;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
use std::{fmt, io};
|
||||
|
||||
use common::{BinarySerializable, OwnedBytes};
|
||||
use compact_space::CompactSpaceDecompressor;
|
||||
@@ -42,7 +42,7 @@ mod null_index_footer;
|
||||
|
||||
mod column;
|
||||
mod gcd;
|
||||
mod serialize;
|
||||
pub mod serialize;
|
||||
|
||||
use self::bitpacked::BitpackedCodec;
|
||||
use self::blockwise_linear::BlockwiseLinearCodec;
|
||||
@@ -133,7 +133,7 @@ impl U128FastFieldCodecType {
|
||||
}
|
||||
|
||||
/// 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,
|
||||
) -> io::Result<Arc<dyn Column<Item>>> {
|
||||
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.
|
||||
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 (mut bytes, _null_index_footer) = read_null_index_footer(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,
|
||||
header: &Header,
|
||||
) -> io::Result<Arc<dyn Column<Item>>> {
|
||||
@@ -321,6 +323,9 @@ mod tests {
|
||||
pub fn get_codec_test_datasets() -> Vec<(Vec<u64>, &'static str)> {
|
||||
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<_>>();
|
||||
data_and_names.push((data, "simple monotonically increasing"));
|
||||
|
||||
@@ -328,6 +333,9 @@ mod tests {
|
||||
vec![5, 6, 7, 8, 9, 10, 99, 100],
|
||||
"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![10], "single value"));
|
||||
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
use std::fmt;
|
||||
use std::marker::PhantomData;
|
||||
use std::ops::RangeInclusive;
|
||||
|
||||
use fastdivide::DividerU64;
|
||||
|
||||
@@ -6,7 +8,9 @@ use crate::MonotonicallyMappableToU128;
|
||||
|
||||
/// 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 {
|
||||
pub trait MonotonicallyMappableToU64:
|
||||
'static + PartialOrd + Copy + Send + Sync + fmt::Debug
|
||||
{
|
||||
/// Converts a value to 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,
|
||||
/// 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> {
|
||||
pub trait StrictlyMonotonicFn<External: Copy, Internal: Copy> {
|
||||
/// 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;
|
||||
|
||||
/// 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
|
||||
@@ -54,7 +76,10 @@ impl<T> From<T> 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)]
|
||||
fn mapping(&self, val: To) -> From {
|
||||
@@ -65,6 +90,15 @@ where T: StrictlyMonotonicFn<From, To>
|
||||
fn inverse(&self, val: From) -> To {
|
||||
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.
|
||||
@@ -142,6 +176,31 @@ impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
|
||||
fn inverse(&self, out: u64) -> External {
|
||||
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.
|
||||
@@ -158,6 +217,17 @@ impl StrictlyMonotonicMappingToInternalBaseval {
|
||||
impl<External: MonotonicallyMappableToU64> StrictlyMonotonicFn<External, u64>
|
||||
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)]
|
||||
fn mapping(&self, val: External) -> u64 {
|
||||
External::to_u64(val) - self.min_value
|
||||
@@ -241,7 +311,7 @@ mod tests {
|
||||
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,
|
||||
test_val: K,
|
||||
) {
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
use std::fmt;
|
||||
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 {
|
||||
pub trait MonotonicallyMappableToU128:
|
||||
'static + PartialOrd + Copy + Send + Sync + fmt::Debug
|
||||
{
|
||||
/// Converts a value to u128.
|
||||
///
|
||||
/// Internally all fast field values are encoded as u64.
|
||||
|
||||
@@ -31,15 +31,16 @@ const BLOCK_BITVEC_SIZE: usize = 8;
|
||||
const BLOCK_OFFSET_SIZE: usize = 4;
|
||||
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]
|
||||
fn count_ones(bitvec: u64, pos_in_bitvec: u32) -> u32 {
|
||||
if pos_in_bitvec == 63 {
|
||||
bitvec.count_ones()
|
||||
} else {
|
||||
let mask = (1u64 << (pos_in_bitvec + 1)) - 1;
|
||||
let masked_bitvec = bitvec & mask;
|
||||
masked_bitvec.count_ones()
|
||||
}
|
||||
let mask = (1u64 << pos_in_bitvec) - 1;
|
||||
let masked_bitvec = bitvec & mask;
|
||||
masked_bitvec.count_ones()
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
@@ -66,9 +67,7 @@ impl DenseCodec {
|
||||
pub fn exists(&self, idx: u32) -> bool {
|
||||
let block_pos = idx / ELEMENTS_PER_BLOCK;
|
||||
let bitvec = self.dense_index_block(block_pos).bitvec;
|
||||
|
||||
let pos_in_bitvec = idx % ELEMENTS_PER_BLOCK;
|
||||
|
||||
get_bit_at(bitvec, pos_in_bitvec)
|
||||
}
|
||||
#[inline]
|
||||
@@ -90,8 +89,7 @@ impl DenseCodec {
|
||||
let pos_in_block_bit_vec = idx % ELEMENTS_PER_BLOCK;
|
||||
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) {
|
||||
// -1 is ok, since idx does exist, so there's at least one
|
||||
Some(index_block.offset + ones_in_block - 1)
|
||||
Some(index_block.offset + ones_in_block)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
@@ -319,9 +317,10 @@ mod tests {
|
||||
set_bit_at(&mut block, 0);
|
||||
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, 2), 2);
|
||||
assert_eq!(count_ones(block, 2), 1);
|
||||
assert_eq!(count_ones(block, 3), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -347,11 +346,16 @@ mod bench {
|
||||
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 current = start;
|
||||
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 {
|
||||
None
|
||||
} 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(
|
||||
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]
|
||||
fn bench_dense_codec_translate_orig_to_codec_90percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
let codec = gen_bools(0.9f64);
|
||||
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_dense_codec_translate_orig_to_codec_50percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
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_dense_codec_translate_orig_to_codec_full_scan_10percent(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.1f64);
|
||||
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_dense_codec_translate_orig_to_codec_full_scan_90percent(bench: &mut Bencher) {
|
||||
fn bench_translate_orig_to_codec_90percent_filled_1percent_hit(bench: &mut Bencher) {
|
||||
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]
|
||||
fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
let codec = gen_bools(0.9f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.01f64);
|
||||
let num_non_nulls = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
codec
|
||||
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000))
|
||||
.translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
|
||||
.last()
|
||||
});
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
let codec = gen_bools(0.9f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.01f64);
|
||||
let num_non_nulls = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
codec
|
||||
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
|
||||
.translate_codec_idx_to_original_idx(n_percent_step_iterator(10.0, num_non_nulls))
|
||||
.last()
|
||||
});
|
||||
}
|
||||
|
||||
#[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 num_vals = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::io::{self, Write};
|
||||
|
||||
use common::{BitSet, OwnedBytes};
|
||||
use common::{BitSet, GroupByIteratorExtended, OwnedBytes};
|
||||
|
||||
use super::{serialize_dense_codec, DenseCodec};
|
||||
|
||||
@@ -78,12 +78,22 @@ struct DenseBlock {
|
||||
}
|
||||
|
||||
impl DenseBlock {
|
||||
#[inline]
|
||||
pub fn exists(&self, idx: u32) -> bool {
|
||||
self.codec.exists(idx)
|
||||
}
|
||||
#[inline]
|
||||
pub fn translate_to_codec_idx(&self, idx: u32) -> Option<u32> {
|
||||
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 {
|
||||
self.codec
|
||||
.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
|
||||
#[inline]
|
||||
fn value_addr(idx: u32) -> ValueAddr {
|
||||
/// Static assert number elements per block this method expects
|
||||
#[allow(clippy::assertions_on_constants)]
|
||||
@@ -273,6 +284,7 @@ impl SparseCodec {
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn find_block(&self, dense_idx: u32, mut block_pos: u32) -> u32 {
|
||||
loop {
|
||||
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.
|
||||
/// Correctness: Provided values must be in increasing values
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
@@ -292,35 +305,41 @@ impl SparseCodec {
|
||||
&'a self,
|
||||
iter: 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;
|
||||
iter.map(move |codec_idx| {
|
||||
// update block_pos to limit search scope
|
||||
block_pos = self.find_block(codec_idx, block_pos);
|
||||
iter.group_by(move |codec_idx| {
|
||||
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 = &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 {
|
||||
SparseCodecBlockVariant::Empty { offset: _ } => {
|
||||
panic!(
|
||||
"invalid input, cannot translate to original index. associated empty \
|
||||
block with dense idx. block_pos {}, idx_in_block {}",
|
||||
block_pos, idx_in_block
|
||||
block with dense idx. block_pos {}, idx_in_block {:?}",
|
||||
block_pos,
|
||||
indexes_in_block_iter.collect::<Vec<_>>()
|
||||
)
|
||||
}
|
||||
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) => {
|
||||
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 {
|
||||
num_elem_in_block < DENSE_BLOCK_THRESHOLD
|
||||
}
|
||||
@@ -595,11 +614,16 @@ mod bench {
|
||||
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 current = start;
|
||||
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 {
|
||||
None
|
||||
} 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(
|
||||
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]
|
||||
fn bench_sparse_codec_translate_orig_to_codec_1percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
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_sparse_codec_translate_orig_to_codec_5percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
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_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);
|
||||
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
|
||||
}
|
||||
|
||||
#[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);
|
||||
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_sparse_codec_translate_orig_to_codec_full_scan_1percent(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.01f64);
|
||||
bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_sparse_codec_translate_orig_to_codec_10percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
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_sparse_codec_translate_orig_to_codec_90percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
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_sparse_codec_translate_codec_to_orig_1percent_filled_random_stride_big_step(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.01f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
let num_non_nulls = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
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()
|
||||
});
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_sparse_codec_translate_codec_to_orig_1percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.01f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
let num_non_nulls = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
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()
|
||||
});
|
||||
}
|
||||
|
||||
#[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 num_vals = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
@@ -714,33 +745,18 @@ mod bench {
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.90f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
let num_non_nulls = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
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()
|
||||
});
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_random_stride(
|
||||
bench: &mut Bencher,
|
||||
) {
|
||||
let codec = gen_bools(0.9f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
codec
|
||||
.translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
|
||||
.last()
|
||||
});
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_sparse_codec_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
|
||||
fn bench_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
|
||||
let codec = gen_bools(0.9f64);
|
||||
let num_vals = codec.num_non_nulls();
|
||||
bench.iter(|| {
|
||||
|
||||
@@ -17,9 +17,9 @@
|
||||
// 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 std::{fmt, io};
|
||||
|
||||
use common::{BinarySerializable, OwnedBytes, VInt};
|
||||
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
|
||||
/// compression.
|
||||
pub fn estimate<T: MonotonicallyMappableToU64>(
|
||||
pub fn estimate<T: MonotonicallyMappableToU64 + fmt::Debug>(
|
||||
typed_column: impl Column<T>,
|
||||
codec_type: FastFieldCodecType,
|
||||
) -> 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.
|
||||
pub fn serialize<T: MonotonicallyMappableToU64>(
|
||||
pub fn serialize<T: MonotonicallyMappableToU64 + fmt::Debug>(
|
||||
typed_column: impl Column<T>,
|
||||
output: &mut impl io::Write,
|
||||
codecs: &[FastFieldCodecType],
|
||||
@@ -285,7 +285,7 @@ pub fn serialize<T: MonotonicallyMappableToU64>(
|
||||
}
|
||||
|
||||
/// 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,
|
||||
typed_column: impl Column<T>,
|
||||
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
|
||||
pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default>(
|
||||
pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default + fmt::Debug>(
|
||||
column: &[T],
|
||||
) -> Arc<dyn Column<T>> {
|
||||
let mut buffer = Vec::new();
|
||||
|
||||
@@ -204,21 +204,23 @@ pub enum IntermediateAggregationResult {
|
||||
/// Holds the intermediate data for metric results
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub enum IntermediateMetricResult {
|
||||
/// Average containing intermediate average data result
|
||||
/// Intermediate average result
|
||||
Average(IntermediateAverage),
|
||||
/// AverageData variant
|
||||
/// Intermediate stats result
|
||||
Stats(IntermediateStats),
|
||||
}
|
||||
|
||||
impl From<SegmentMetricResultCollector> for IntermediateMetricResult {
|
||||
fn from(tree: SegmentMetricResultCollector) -> Self {
|
||||
match tree {
|
||||
SegmentMetricResultCollector::Average(collector) => {
|
||||
IntermediateMetricResult::Average(IntermediateAverage::from_collector(collector))
|
||||
}
|
||||
SegmentMetricResultCollector::Stats(collector) => {
|
||||
IntermediateMetricResult::Stats(collector.stats)
|
||||
}
|
||||
SegmentMetricResultCollector::Stats(collector) => match collector.collecting_for {
|
||||
super::metric::SegmentStatsType::Stats => {
|
||||
IntermediateMetricResult::Stats(collector.stats)
|
||||
}
|
||||
super::metric::SegmentStatsType::Avg => IntermediateMetricResult::Average(
|
||||
IntermediateAverage::from_collector(collector),
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,11 +1,8 @@
|
||||
use std::fmt::Debug;
|
||||
|
||||
use fastfield_codecs::Column;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::aggregation::f64_from_fastfield_u64;
|
||||
use crate::schema::Type;
|
||||
use crate::DocId;
|
||||
use super::SegmentStatsCollector;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
/// A single-value metric aggregation that computes the average of numeric values that are
|
||||
@@ -36,51 +33,6 @@ impl AverageAggregation {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub(crate) struct SegmentAverageCollector {
|
||||
pub data: IntermediateAverage,
|
||||
field_type: Type,
|
||||
}
|
||||
|
||||
impl Debug for SegmentAverageCollector {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AverageCollector")
|
||||
.field("data", &self.data)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl SegmentAverageCollector {
|
||||
pub fn from_req(field_type: Type) -> Self {
|
||||
Self {
|
||||
field_type,
|
||||
data: Default::default(),
|
||||
}
|
||||
}
|
||||
pub(crate) fn collect_block(&mut self, doc: &[DocId], field: &dyn Column<u64>) {
|
||||
let mut iter = doc.chunks_exact(4);
|
||||
for docs in iter.by_ref() {
|
||||
let val1 = field.get_val(docs[0]);
|
||||
let val2 = field.get_val(docs[1]);
|
||||
let val3 = field.get_val(docs[2]);
|
||||
let val4 = field.get_val(docs[3]);
|
||||
let val1 = f64_from_fastfield_u64(val1, &self.field_type);
|
||||
let val2 = f64_from_fastfield_u64(val2, &self.field_type);
|
||||
let val3 = f64_from_fastfield_u64(val3, &self.field_type);
|
||||
let val4 = f64_from_fastfield_u64(val4, &self.field_type);
|
||||
self.data.collect(val1);
|
||||
self.data.collect(val2);
|
||||
self.data.collect(val3);
|
||||
self.data.collect(val4);
|
||||
}
|
||||
for &doc in iter.remainder() {
|
||||
let val = field.get_val(doc);
|
||||
let val = f64_from_fastfield_u64(val, &self.field_type);
|
||||
self.data.collect(val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Contains mergeable version of average data.
|
||||
#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct IntermediateAverage {
|
||||
@@ -89,8 +41,11 @@ pub struct IntermediateAverage {
|
||||
}
|
||||
|
||||
impl IntermediateAverage {
|
||||
pub(crate) fn from_collector(collector: SegmentAverageCollector) -> Self {
|
||||
collector.data
|
||||
pub(crate) fn from_collector(collector: SegmentStatsCollector) -> Self {
|
||||
Self {
|
||||
sum: collector.stats.sum,
|
||||
doc_count: collector.stats.count,
|
||||
}
|
||||
}
|
||||
|
||||
/// Merge average data into this instance.
|
||||
@@ -106,9 +61,4 @@ impl IntermediateAverage {
|
||||
Some(self.sum / self.doc_count as f64)
|
||||
}
|
||||
}
|
||||
#[inline]
|
||||
fn collect(&mut self, val: f64) {
|
||||
self.doc_count += 1;
|
||||
self.sum += val;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,7 +40,7 @@ impl StatsAggregation {
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Stats {
|
||||
/// The number of documents.
|
||||
pub count: usize,
|
||||
pub count: u64,
|
||||
/// The sum of the fast field values.
|
||||
pub sum: f64,
|
||||
/// The standard deviation of the fast field values. `None` for count == 0.
|
||||
@@ -73,11 +73,16 @@ impl Stats {
|
||||
/// `IntermediateStats` contains the mergeable version for stats.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct IntermediateStats {
|
||||
count: usize,
|
||||
sum: f64,
|
||||
squared_sum: f64,
|
||||
min: f64,
|
||||
max: f64,
|
||||
/// the number of values
|
||||
pub count: u64,
|
||||
/// the sum of the values
|
||||
pub sum: f64,
|
||||
/// the squared sum of the values
|
||||
pub squared_sum: f64,
|
||||
/// the min value of the values
|
||||
pub min: f64,
|
||||
/// the max value of the values
|
||||
pub max: f64,
|
||||
}
|
||||
impl Default for IntermediateStats {
|
||||
fn default() -> Self {
|
||||
@@ -150,17 +155,25 @@ impl IntermediateStats {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) enum SegmentStatsType {
|
||||
Stats,
|
||||
Avg,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) struct SegmentStatsCollector {
|
||||
pub(crate) stats: IntermediateStats,
|
||||
field_type: Type,
|
||||
pub(crate) collecting_for: SegmentStatsType,
|
||||
}
|
||||
|
||||
impl SegmentStatsCollector {
|
||||
pub fn from_req(field_type: Type) -> Self {
|
||||
pub fn from_req(field_type: Type, collecting_for: SegmentStatsType) -> Self {
|
||||
Self {
|
||||
field_type,
|
||||
stats: IntermediateStats::default(),
|
||||
collecting_for,
|
||||
}
|
||||
}
|
||||
pub(crate) fn collect_block(&mut self, doc: &[DocId], field: &dyn Column<u64>) {
|
||||
|
||||
@@ -1208,7 +1208,7 @@ mod tests {
|
||||
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(),
|
||||
score_field => val as u64,
|
||||
score_field_f64 => val as f64,
|
||||
score_field_f64 => val,
|
||||
score_field_i64 => val as i64,
|
||||
))?;
|
||||
}
|
||||
@@ -1250,10 +1250,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&term_query, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&term_query, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1281,10 +1278,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&term_query, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&term_query, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1312,10 +1306,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&term_query, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&term_query, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1351,10 +1342,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&term_query, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&term_query, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1380,10 +1368,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&AllQuery, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&AllQuery, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1409,10 +1394,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&AllQuery, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&AllQuery, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1446,10 +1428,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&AllQuery, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&AllQuery, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1481,10 +1460,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&AllQuery, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&AllQuery, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1520,10 +1496,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&AllQuery, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&AllQuery, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1550,10 +1523,7 @@ mod tests {
|
||||
let collector = AggregationCollector::from_aggs(agg_req_1, None, index.schema());
|
||||
|
||||
let searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&AllQuery, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&AllQuery, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1597,7 +1567,7 @@ mod tests {
|
||||
],
|
||||
..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 searcher = reader.searcher();
|
||||
let agg_res: AggregationResults =
|
||||
searcher.search(&term_query, &collector).unwrap().into();
|
||||
|
||||
agg_res
|
||||
searcher.search(&term_query, &collector).unwrap()
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,7 +15,7 @@ use super::bucket::{SegmentHistogramCollector, SegmentRangeCollector, SegmentTer
|
||||
use super::collector::MAX_BUCKET_COUNT;
|
||||
use super::intermediate_agg_result::{IntermediateAggregationResults, IntermediateBucketResult};
|
||||
use super::metric::{
|
||||
AverageAggregation, SegmentAverageCollector, SegmentStatsCollector, StatsAggregation,
|
||||
AverageAggregation, SegmentStatsCollector, SegmentStatsType, StatsAggregation,
|
||||
};
|
||||
use super::VecWithNames;
|
||||
use crate::aggregation::agg_req::BucketAggregationType;
|
||||
@@ -163,7 +163,6 @@ impl SegmentAggregationResultsCollector {
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) enum SegmentMetricResultCollector {
|
||||
Average(SegmentAverageCollector),
|
||||
Stats(SegmentStatsCollector),
|
||||
}
|
||||
|
||||
@@ -171,22 +170,19 @@ impl SegmentMetricResultCollector {
|
||||
pub fn from_req_and_validate(req: &MetricAggregationWithAccessor) -> crate::Result<Self> {
|
||||
match &req.metric {
|
||||
MetricAggregation::Average(AverageAggregation { field: _ }) => {
|
||||
Ok(SegmentMetricResultCollector::Average(
|
||||
SegmentAverageCollector::from_req(req.field_type),
|
||||
Ok(SegmentMetricResultCollector::Stats(
|
||||
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Avg),
|
||||
))
|
||||
}
|
||||
MetricAggregation::Stats(StatsAggregation { field: _ }) => {
|
||||
Ok(SegmentMetricResultCollector::Stats(
|
||||
SegmentStatsCollector::from_req(req.field_type),
|
||||
SegmentStatsCollector::from_req(req.field_type, SegmentStatsType::Stats),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
pub(crate) fn collect_block(&mut self, doc: &[DocId], metric: &MetricAggregationWithAccessor) {
|
||||
match self {
|
||||
SegmentMetricResultCollector::Average(avg_collector) => {
|
||||
avg_collector.collect_block(doc, &*metric.accessor);
|
||||
}
|
||||
SegmentMetricResultCollector::Stats(stats_collector) => {
|
||||
stats_collector.collect_block(doc, &*metric.accessor);
|
||||
}
|
||||
|
||||
@@ -198,11 +198,10 @@ impl Searcher {
|
||||
collector: &C,
|
||||
executor: &Executor,
|
||||
) -> crate::Result<C::Fruit> {
|
||||
let scoring_enabled = collector.requires_scoring();
|
||||
let enabled_scoring = if scoring_enabled {
|
||||
EnableScoring::Enabled(self)
|
||||
let enabled_scoring = if collector.requires_scoring() {
|
||||
EnableScoring::enabled_from_searcher(self)
|
||||
} else {
|
||||
EnableScoring::Disabled(self.schema())
|
||||
EnableScoring::disabled_from_searcher(self)
|
||||
};
|
||||
let weight = query.weight(enabled_scoring)?;
|
||||
let segment_readers = self.segment_readers();
|
||||
|
||||
@@ -32,7 +32,7 @@ impl LockError {
|
||||
/// Error that may occur when opening a directory
|
||||
#[derive(Debug, Clone, Error)]
|
||||
pub enum OpenDirectoryError {
|
||||
/// The underlying directory does not exists.
|
||||
/// The underlying directory does not exist.
|
||||
#[error("Directory does not exist: '{0}'.")]
|
||||
DoesNotExist(PathBuf),
|
||||
/// 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
|
||||
#[derive(Debug, Clone, Error)]
|
||||
pub enum OpenReadError {
|
||||
/// The file does not exists.
|
||||
#[error("Files does not exists: {0:?}")]
|
||||
/// The file does not exist.
|
||||
#[error("Files does not exist: {0:?}")]
|
||||
FileDoesNotExist(PathBuf),
|
||||
/// Any kind of io::Error.
|
||||
#[error(
|
||||
@@ -181,8 +181,8 @@ impl OpenReadError {
|
||||
/// Error that may occur when trying to delete a file
|
||||
#[derive(Debug, Clone, Error)]
|
||||
pub enum DeleteError {
|
||||
/// The file does not exists.
|
||||
#[error("File does not exists: '{0}'.")]
|
||||
/// The file does not exist.
|
||||
#[error("File does not exist: '{0}'.")]
|
||||
FileDoesNotExist(PathBuf),
|
||||
/// Any kind of IO error that happens when
|
||||
/// interacting with the underlying IO device.
|
||||
|
||||
@@ -232,7 +232,7 @@ impl Directory for RamDirectory {
|
||||
let path_buf = PathBuf::from(path);
|
||||
self.fs.write().unwrap().write(path_buf, data);
|
||||
if path == *META_FILEPATH {
|
||||
let _ = self.fs.write().unwrap().watch_router.broadcast();
|
||||
drop(self.fs.write().unwrap().watch_router.broadcast());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -168,7 +168,7 @@ mod tests {
|
||||
watch_event_router.broadcast().wait().unwrap();
|
||||
assert_eq!(2, counter.load(Ordering::SeqCst));
|
||||
mem::drop(handle_a);
|
||||
let _ = watch_event_router.broadcast();
|
||||
drop(watch_event_router.broadcast());
|
||||
watch_event_router.broadcast().wait().unwrap();
|
||||
assert_eq!(2, counter.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
@@ -175,7 +175,7 @@ mod bench {
|
||||
|
||||
fn get_alive() -> 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);
|
||||
}
|
||||
data
|
||||
|
||||
@@ -96,7 +96,7 @@ mod tests {
|
||||
let term = Term::from_field_bytes(field, b"lucene".as_ref());
|
||||
let term_query = TermQuery::new(term, IndexRecordOption::Basic);
|
||||
let term_weight_err =
|
||||
term_query.specialized_weight(EnableScoring::Disabled(searcher.schema()));
|
||||
term_query.specialized_weight(EnableScoring::disabled_from_schema(searcher.schema()));
|
||||
assert!(matches!(
|
||||
term_weight_err,
|
||||
Err(crate::TantivyError::SchemaError(_))
|
||||
|
||||
@@ -12,13 +12,15 @@
|
||||
//!
|
||||
//!
|
||||
//! 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
|
||||
//! automatically, when serializing.
|
||||
//!
|
||||
//! Read access performance is comparable to that of an array lookup.
|
||||
|
||||
use std::net::Ipv6Addr;
|
||||
|
||||
use fastfield_codecs::MonotonicallyMappableToU64;
|
||||
|
||||
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 use self::multivalued::{
|
||||
MultiValueIndex, MultiValueU128FastFieldWriter, MultiValuedFastFieldReader,
|
||||
MultiValuedFastFieldWriter, MultiValuedU128FastFieldReader,
|
||||
MultiValuedFastFieldWriter,
|
||||
};
|
||||
pub(crate) use self::readers::type_and_cardinality;
|
||||
pub use self::readers::FastFieldReaders;
|
||||
@@ -47,6 +49,33 @@ mod readers;
|
||||
mod serializer;
|
||||
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:
|
||||
/// (u64, i64 and f64, bool, DateTime).
|
||||
pub trait FastValue:
|
||||
@@ -54,12 +83,6 @@ pub trait FastValue:
|
||||
{
|
||||
/// Returns the `schema::Type` for this FastValue.
|
||||
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 {
|
||||
@@ -101,12 +124,6 @@ impl FastValue for DateTime {
|
||||
fn to_type() -> Type {
|
||||
Type::Date
|
||||
}
|
||||
|
||||
fn make_zero() -> Self {
|
||||
DateTime {
|
||||
timestamp_micros: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn value_to_u64(value: &Value) -> crate::Result<u64> {
|
||||
@@ -145,7 +162,7 @@ impl FastFieldType {
|
||||
mod tests {
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::ops::Range;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -159,7 +176,9 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::directory::{CompositeFile, Directory, RamDirectory, WritePtr};
|
||||
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::{DateOptions, DatePrecision, Index, SegmentId, SegmentReader};
|
||||
|
||||
@@ -520,11 +539,6 @@ mod tests {
|
||||
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> {
|
||||
let mut all = vec![];
|
||||
|
||||
@@ -969,4 +983,117 @@ mod tests {
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ mod writer;
|
||||
use fastfield_codecs::FastFieldCodecType;
|
||||
pub use index::MultiValueIndex;
|
||||
|
||||
pub use self::reader::{MultiValuedFastFieldReader, MultiValuedU128FastFieldReader};
|
||||
pub use self::reader::MultiValuedFastFieldReader;
|
||||
pub(crate) use self::writer::MultivalueStartIndex;
|
||||
pub use self::writer::{MultiValueU128FastFieldWriter, MultiValuedFastFieldWriter};
|
||||
|
||||
@@ -525,7 +525,7 @@ mod bench {
|
||||
serializer.close().unwrap();
|
||||
field
|
||||
};
|
||||
let file = directory.open_read(&path).unwrap();
|
||||
let file = directory.open_read(path).unwrap();
|
||||
{
|
||||
let fast_fields_composite = CompositeFile::open(&file).unwrap();
|
||||
let data_idx = fast_fields_composite
|
||||
|
||||
@@ -1,107 +1,31 @@
|
||||
use core::fmt;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
use std::sync::Arc;
|
||||
|
||||
use fastfield_codecs::{Column, MonotonicallyMappableToU128};
|
||||
use fastfield_codecs::Column;
|
||||
|
||||
use super::MultiValueIndex;
|
||||
use crate::fastfield::FastValue;
|
||||
use crate::fastfield::MakeZero;
|
||||
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
|
||||
/// values for all reader.
|
||||
/// The `idx_reader` associated, for each document, the index of its first value.
|
||||
/// Stores the start position for each document.
|
||||
/// The `vals_reader` will access the concatenated list of all values.
|
||||
/// The `idx_reader` associates, for each document, the index of its first value.
|
||||
#[derive(Clone)]
|
||||
pub struct MultiValuedFastFieldReader<Item: FastValue> {
|
||||
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> {
|
||||
pub struct MultiValuedFastFieldReader<T> {
|
||||
idx_reader: MultiValueIndex,
|
||||
vals_reader: Arc<dyn Column<T>>,
|
||||
}
|
||||
|
||||
impl<T: MonotonicallyMappableToU128> MultiValuedU128FastFieldReader<T> {
|
||||
impl<T: PartialOrd + MakeZero + Copy + fmt::Debug> MultiValuedFastFieldReader<T> {
|
||||
pub(crate) fn open(
|
||||
idx_reader: Arc<dyn Column<u64>>,
|
||||
vals_reader: Arc<dyn Column<T>>,
|
||||
) -> MultiValuedU128FastFieldReader<T> {
|
||||
) -> MultiValuedFastFieldReader<T> {
|
||||
Self {
|
||||
idx_reader: MultiValueIndex::new(idx_reader),
|
||||
vals_reader,
|
||||
@@ -122,7 +46,7 @@ impl<T: MonotonicallyMappableToU128> MultiValuedU128FastFieldReader<T> {
|
||||
#[inline]
|
||||
fn get_vals_for_range(&self, range: Range<u32>, vals: &mut Vec<T>) {
|
||||
let len = (range.end - range.start) as usize;
|
||||
vals.resize(len, T::from_u128(0));
|
||||
vals.resize(len, T::make_zero());
|
||||
self.vals_reader
|
||||
.get_range(range.start as u64, &mut vals[..]);
|
||||
}
|
||||
@@ -199,8 +123,131 @@ impl<T: MonotonicallyMappableToU128> MultiValuedU128FastFieldReader<T> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use time::{Duration, OffsetDateTime};
|
||||
|
||||
use crate::collector::Count;
|
||||
use crate::core::Index;
|
||||
use crate::query::RangeQuery;
|
||||
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(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(
|
||||
date_field,
|
||||
DateTime::from_utc(first_time_stamp)..DateTime::from_utc(two_secs_ahead)
|
||||
)),
|
||||
2
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multifastfield_reader() -> crate::Result<()> {
|
||||
|
||||
@@ -3,11 +3,9 @@ use std::sync::Arc;
|
||||
|
||||
use fastfield_codecs::{open, open_u128, Column};
|
||||
|
||||
use super::multivalued::MultiValuedU128FastFieldReader;
|
||||
use super::multivalued::MultiValuedFastFieldReader;
|
||||
use crate::directory::{CompositeFile, FileSlice};
|
||||
use crate::fastfield::{
|
||||
BytesFastFieldReader, FastFieldNotAvailableError, FastValue, MultiValuedFastFieldReader,
|
||||
};
|
||||
use crate::fastfield::{BytesFastFieldReader, FastFieldNotAvailableError, FastValue};
|
||||
use crate::schema::{Cardinality, Field, FieldType, Schema};
|
||||
use crate::space_usage::PerFieldSpaceUsage;
|
||||
use crate::{DateTime, TantivyError};
|
||||
@@ -161,20 +159,14 @@ impl FastFieldReaders {
|
||||
/// Returns the `ip` fast field reader reader associated to `field`.
|
||||
///
|
||||
/// If `field` is not a u128 fast field, this method returns an Error.
|
||||
pub fn ip_addrs(
|
||||
&self,
|
||||
field: Field,
|
||||
) -> crate::Result<MultiValuedU128FastFieldReader<Ipv6Addr>> {
|
||||
pub fn ip_addrs(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<Ipv6Addr>> {
|
||||
self.check_type(field, FastType::U128, Cardinality::MultiValues)?;
|
||||
let idx_reader: Arc<dyn Column<u64>> = self.typed_fast_field_reader(field)?;
|
||||
|
||||
let bytes = self.fast_field_data(field, 1)?.read_bytes()?;
|
||||
let vals_reader = open_u128::<Ipv6Addr>(bytes)?;
|
||||
|
||||
Ok(MultiValuedU128FastFieldReader::open(
|
||||
idx_reader,
|
||||
vals_reader,
|
||||
))
|
||||
Ok(MultiValuedFastFieldReader::open(idx_reader, vals_reader))
|
||||
}
|
||||
|
||||
/// Returns the `u128` fast field reader reader associated to `field`.
|
||||
@@ -189,17 +181,14 @@ impl FastFieldReaders {
|
||||
/// 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.
|
||||
pub fn u128s(&self, field: Field) -> crate::Result<MultiValuedU128FastFieldReader<u128>> {
|
||||
pub fn u128s(&self, field: Field) -> crate::Result<MultiValuedFastFieldReader<u128>> {
|
||||
self.check_type(field, FastType::U128, Cardinality::MultiValues)?;
|
||||
let idx_reader: Arc<dyn Column<u64>> = self.typed_fast_field_reader(field)?;
|
||||
|
||||
let bytes = self.fast_field_data(field, 1)?.read_bytes()?;
|
||||
let vals_reader = open_u128::<u128>(bytes)?;
|
||||
|
||||
Ok(MultiValuedU128FastFieldReader::open(
|
||||
idx_reader,
|
||||
vals_reader,
|
||||
))
|
||||
Ok(MultiValuedFastFieldReader::open(idx_reader, vals_reader))
|
||||
}
|
||||
|
||||
/// Returns the `u64` fast field reader reader associated with `field`, regardless of whether
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::fmt;
|
||||
use std::io::{self, Write};
|
||||
|
||||
pub use fastfield_codecs::Column;
|
||||
@@ -49,7 +50,7 @@ impl CompositeFastFieldSerializer {
|
||||
|
||||
/// Serialize data into a new u64 fast field. The best compression codec will be chosen
|
||||
/// automatically.
|
||||
pub fn create_auto_detect_u64_fast_field<T: MonotonicallyMappableToU64>(
|
||||
pub fn create_auto_detect_u64_fast_field<T: MonotonicallyMappableToU64 + fmt::Debug>(
|
||||
&mut self,
|
||||
field: Field,
|
||||
fastfield_accessor: impl Column<T>,
|
||||
@@ -59,7 +60,9 @@ impl CompositeFastFieldSerializer {
|
||||
|
||||
/// Serialize data into a new u64 fast field. The best compression codec will be chosen
|
||||
/// automatically.
|
||||
pub fn create_auto_detect_u64_fast_field_with_idx<T: MonotonicallyMappableToU64>(
|
||||
pub fn create_auto_detect_u64_fast_field_with_idx<
|
||||
T: MonotonicallyMappableToU64 + fmt::Debug,
|
||||
>(
|
||||
&mut self,
|
||||
field: Field,
|
||||
fastfield_accessor: impl Column<T>,
|
||||
@@ -72,7 +75,9 @@ impl CompositeFastFieldSerializer {
|
||||
|
||||
/// Serialize data into a new u64 fast field. The best compression codec of the the provided
|
||||
/// will be chosen.
|
||||
pub fn create_auto_detect_u64_fast_field_with_idx_and_codecs<T: MonotonicallyMappableToU64>(
|
||||
pub fn create_auto_detect_u64_fast_field_with_idx_and_codecs<
|
||||
T: MonotonicallyMappableToU64 + fmt::Debug,
|
||||
>(
|
||||
&mut self,
|
||||
field: Field,
|
||||
fastfield_accessor: impl Column<T>,
|
||||
|
||||
@@ -678,7 +678,7 @@ impl IndexWriter {
|
||||
/// only after calling `commit()`.
|
||||
#[doc(hidden)]
|
||||
pub fn delete_query(&self, query: Box<dyn Query>) -> crate::Result<Opstamp> {
|
||||
let weight = query.weight(EnableScoring::Disabled(&self.index.schema()))?;
|
||||
let weight = query.weight(EnableScoring::disabled_from_schema(&self.index.schema()))?;
|
||||
let opstamp = self.stamper.stamp();
|
||||
let delete_operation = DeleteOperation {
|
||||
opstamp,
|
||||
@@ -759,7 +759,8 @@ impl IndexWriter {
|
||||
match user_op {
|
||||
UserOperation::Delete(term) => {
|
||||
let query = TermQuery::new(term, IndexRecordOption::Basic);
|
||||
let weight = query.weight(EnableScoring::Disabled(&self.index.schema()))?;
|
||||
let weight =
|
||||
query.weight(EnableScoring::disabled_from_schema(&self.index.schema()))?;
|
||||
let delete_operation = DeleteOperation {
|
||||
opstamp,
|
||||
target: weight,
|
||||
|
||||
@@ -89,11 +89,11 @@ pub(crate) fn index_json_values<'a>(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn index_json_object<'a>(
|
||||
fn index_json_object(
|
||||
doc: DocId,
|
||||
json_value: &serde_json::Map<String, serde_json::Value>,
|
||||
text_analyzer: &TextAnalyzer,
|
||||
json_term_writer: &mut JsonTermWriter<'a>,
|
||||
json_term_writer: &mut JsonTermWriter,
|
||||
postings_writer: &mut dyn PostingsWriter,
|
||||
ctx: &mut IndexingContext,
|
||||
positions_per_path: &mut IndexingPositionsPerPath,
|
||||
@@ -113,11 +113,11 @@ fn index_json_object<'a>(
|
||||
}
|
||||
}
|
||||
|
||||
fn index_json_value<'a>(
|
||||
fn index_json_value(
|
||||
doc: DocId,
|
||||
json_value: &serde_json::Value,
|
||||
text_analyzer: &TextAnalyzer,
|
||||
json_term_writer: &mut JsonTermWriter<'a>,
|
||||
json_term_writer: &mut JsonTermWriter,
|
||||
postings_writer: &mut dyn PostingsWriter,
|
||||
ctx: &mut IndexingContext,
|
||||
positions_per_path: &mut IndexingPositionsPerPath,
|
||||
|
||||
@@ -13,7 +13,7 @@ use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::error::DataCorruption;
|
||||
use crate::fastfield::{
|
||||
get_fastfield_codecs_for_multivalue, AliveBitSet, Column, CompositeFastFieldSerializer,
|
||||
MultiValueIndex, MultiValuedFastFieldReader, MultiValuedU128FastFieldReader,
|
||||
MultiValueIndex, MultiValuedFastFieldReader,
|
||||
};
|
||||
use crate::fieldnorm::{FieldNormReader, FieldNormReaders, FieldNormsSerializer, FieldNormsWriter};
|
||||
use crate::indexer::doc_id_mapping::{expect_field_id_for_sort_field, SegmentDocIdMapping};
|
||||
@@ -331,18 +331,18 @@ impl IndexMerger {
|
||||
fast_field_serializer: &mut CompositeFastFieldSerializer,
|
||||
doc_id_mapping: &SegmentDocIdMapping,
|
||||
) -> crate::Result<()> {
|
||||
let segment_and_ff_readers: Vec<(&SegmentReader, MultiValuedU128FastFieldReader<u128>)> =
|
||||
self.readers
|
||||
.iter()
|
||||
.map(|segment_reader| {
|
||||
let ff_reader: MultiValuedU128FastFieldReader<u128> =
|
||||
segment_reader.fast_fields().u128s(field).expect(
|
||||
"Failed to find index for multivalued field. This is a bug in \
|
||||
tantivy, please report.",
|
||||
);
|
||||
(segment_reader, ff_reader)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let segment_and_ff_readers: Vec<(&SegmentReader, MultiValuedFastFieldReader<u128>)> = self
|
||||
.readers
|
||||
.iter()
|
||||
.map(|segment_reader| {
|
||||
let ff_reader: MultiValuedFastFieldReader<u128> =
|
||||
segment_reader.fast_fields().u128s(field).expect(
|
||||
"Failed to find index for multivalued field. This is a bug in tantivy, \
|
||||
please report.",
|
||||
);
|
||||
(segment_reader, ff_reader)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Self::write_1_n_fast_field_idx_generic(
|
||||
field,
|
||||
|
||||
@@ -577,7 +577,7 @@ impl SegmentUpdater {
|
||||
for merge_operation in merge_candidates {
|
||||
// If a merge cannot be started this is not a fatal error.
|
||||
// We do log a warning in `start_merge`.
|
||||
let _ = self.start_merge(merge_operation);
|
||||
drop(self.start_merge(merge_operation));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
19
src/lib.rs
19
src/lib.rs
@@ -1,17 +1,14 @@
|
||||
#![doc(html_logo_url = "http://fulmicoton.com/tantivy-logo/tantivy-logo.png")]
|
||||
#![cfg_attr(all(feature = "unstable", test), feature(test))]
|
||||
#![cfg_attr(
|
||||
feature = "cargo-clippy",
|
||||
allow(
|
||||
clippy::module_inception,
|
||||
clippy::needless_range_loop,
|
||||
clippy::bool_assert_comparison
|
||||
)
|
||||
)]
|
||||
#![doc(test(attr(allow(unused_variables), deny(warnings))))]
|
||||
#![warn(missing_docs)]
|
||||
#![allow(clippy::len_without_is_empty)]
|
||||
#![allow(clippy::derive_partial_eq_without_eq)]
|
||||
#![allow(
|
||||
clippy::len_without_is_empty,
|
||||
clippy::derive_partial_eq_without_eq,
|
||||
clippy::module_inception,
|
||||
clippy::needless_range_loop,
|
||||
clippy::bool_assert_comparison
|
||||
)]
|
||||
|
||||
//! # `tantivy`
|
||||
//!
|
||||
@@ -144,7 +141,7 @@ use crate::time::{OffsetDateTime, PrimitiveDateTime, UtcOffset};
|
||||
/// All constructors and conversions are provided as explicit
|
||||
/// functions and not by implementing any `From`/`Into` traits
|
||||
/// to prevent unintended usage.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
#[derive(Clone, Default, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub struct DateTime {
|
||||
// Timestamp in microseconds.
|
||||
pub(crate) timestamp_micros: i64,
|
||||
|
||||
@@ -291,7 +291,7 @@ pub mod tests {
|
||||
const PADDING_VALUE: u32 = 234_234_345u32;
|
||||
let expected_length = 154;
|
||||
let mut encoder = BlockEncoder::new();
|
||||
let input: Vec<u32> = (0u32..123u32).map(|i| 4 + i * 7 / 2).into_iter().collect();
|
||||
let input: Vec<u32> = (0u32..123u32).map(|i| 4 + i * 7 / 2).collect();
|
||||
for offset in &[0u32, 1u32, 2u32] {
|
||||
let encoded_data = encoder.compress_vint_sorted(&input, *offset);
|
||||
assert!(encoded_data.len() <= expected_length);
|
||||
|
||||
@@ -631,7 +631,7 @@ mod bench {
|
||||
let mut segment_postings = segment_reader
|
||||
.inverted_index(TERM_A.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_A, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_A, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
while segment_postings.advance() != TERMINATED {}
|
||||
@@ -647,25 +647,25 @@ mod bench {
|
||||
let segment_postings_a = segment_reader
|
||||
.inverted_index(TERM_A.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_A, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_A, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let segment_postings_b = segment_reader
|
||||
.inverted_index(TERM_B.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_B, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_B, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let segment_postings_c = segment_reader
|
||||
.inverted_index(TERM_C.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_C, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_C, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let segment_postings_d = segment_reader
|
||||
.inverted_index(TERM_D.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_D, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_D, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut intersection = Intersection::new(vec![
|
||||
@@ -687,7 +687,7 @@ mod bench {
|
||||
let mut segment_postings = segment_reader
|
||||
.inverted_index(TERM_A.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_A, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_A, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
@@ -705,7 +705,7 @@ mod bench {
|
||||
let mut segment_postings = segment_reader
|
||||
.inverted_index(TERM_A.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_A, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_A, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
for doc in &existing_docs {
|
||||
@@ -746,7 +746,7 @@ mod bench {
|
||||
let mut segment_postings = segment_reader
|
||||
.inverted_index(TERM_A.field())
|
||||
.unwrap()
|
||||
.read_postings(&*TERM_A, IndexRecordOption::Basic)
|
||||
.read_postings(&TERM_A, IndexRecordOption::Basic)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut s = 0u32;
|
||||
|
||||
@@ -95,7 +95,7 @@ mod tests {
|
||||
let index = create_test_index()?;
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
let weight = AllQuery.weight(EnableScoring::Disabled(&index.schema()))?;
|
||||
let weight = AllQuery.weight(EnableScoring::disabled_from_schema(&index.schema()))?;
|
||||
{
|
||||
let reader = searcher.segment_reader(0);
|
||||
let mut scorer = weight.scorer(reader, 1.0)?;
|
||||
@@ -118,7 +118,7 @@ mod tests {
|
||||
let index = create_test_index()?;
|
||||
let reader = index.reader()?;
|
||||
let searcher = reader.searcher();
|
||||
let weight = AllQuery.weight(EnableScoring::Disabled(searcher.schema()))?;
|
||||
let weight = AllQuery.weight(EnableScoring::disabled_from_schema(searcher.schema()))?;
|
||||
let reader = searcher.segment_reader(0);
|
||||
{
|
||||
let mut scorer = weight.scorer(reader, 2.0)?;
|
||||
|
||||
@@ -146,7 +146,7 @@ impl Query for BooleanQuery {
|
||||
let sub_weights = self
|
||||
.subqueries
|
||||
.iter()
|
||||
.map(|&(ref occur, ref subquery)| Ok((*occur, subquery.weight(enable_scoring)?)))
|
||||
.map(|(occur, subquery)| Ok((*occur, subquery.weight(enable_scoring)?)))
|
||||
.collect::<crate::Result<_>>()?;
|
||||
Ok(Box::new(BooleanWeight::new(
|
||||
sub_weights,
|
||||
|
||||
@@ -91,7 +91,7 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
boost: Score,
|
||||
) -> crate::Result<HashMap<Occur, Vec<Box<dyn Scorer>>>> {
|
||||
let mut per_occur_scorers: HashMap<Occur, Vec<Box<dyn Scorer>>> = HashMap::new();
|
||||
for &(ref occur, ref subweight) in &self.weights {
|
||||
for (occur, subweight) in &self.weights {
|
||||
let sub_scorer: Box<dyn Scorer> = subweight.scorer(reader, boost)?;
|
||||
per_occur_scorers
|
||||
.entry(*occur)
|
||||
@@ -191,7 +191,7 @@ impl<TScoreCombiner: ScoreCombiner + Sync> Weight for BooleanWeight<TScoreCombin
|
||||
}
|
||||
|
||||
let mut explanation = Explanation::new("BooleanClause. Sum of ...", scorer.score());
|
||||
for &(ref occur, ref subweight) in &self.weights {
|
||||
for (occur, subweight) in &self.weights {
|
||||
if is_positive_occur(*occur) {
|
||||
if let Ok(child_explanation) = subweight.explain(reader, doc) {
|
||||
explanation.add_detail(child_explanation);
|
||||
|
||||
@@ -98,7 +98,7 @@ mod tests {
|
||||
}
|
||||
{
|
||||
let query = query_parser.parse_query("+a b")?;
|
||||
let weight = query.weight(EnableScoring::Disabled(searcher.schema()))?;
|
||||
let weight = query.weight(EnableScoring::disabled_from_schema(searcher.schema()))?;
|
||||
let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?;
|
||||
assert!(scorer.is::<TermScorer>());
|
||||
}
|
||||
|
||||
@@ -1,8 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
use std::ops::Range;
|
||||
|
||||
use levenshtein_automata::{Distance, LevenshteinAutomatonBuilder, DFA};
|
||||
use once_cell::sync::Lazy;
|
||||
use once_cell::sync::OnceCell;
|
||||
use tantivy_fst::Automaton;
|
||||
|
||||
use crate::query::{AutomatonWeight, EnableScoring, Query, Weight};
|
||||
@@ -34,22 +31,6 @@ impl Automaton for DfaWrapper {
|
||||
}
|
||||
}
|
||||
|
||||
/// A range of Levenshtein distances that we will build DFAs for our terms
|
||||
/// The computation is exponential, so best keep it to low single digits
|
||||
const VALID_LEVENSHTEIN_DISTANCE_RANGE: Range<u8> = 0..3;
|
||||
|
||||
static LEV_BUILDER: Lazy<HashMap<(u8, bool), LevenshteinAutomatonBuilder>> = Lazy::new(|| {
|
||||
let mut lev_builder_cache = HashMap::new();
|
||||
// TODO make population lazy on a `(distance, val)` basis
|
||||
for distance in VALID_LEVENSHTEIN_DISTANCE_RANGE {
|
||||
for &transposition in &[false, true] {
|
||||
let lev_automaton_builder = LevenshteinAutomatonBuilder::new(distance, transposition);
|
||||
lev_builder_cache.insert((distance, transposition), lev_automaton_builder);
|
||||
}
|
||||
}
|
||||
lev_builder_cache
|
||||
});
|
||||
|
||||
/// A Fuzzy Query matches all of the documents
|
||||
/// containing a specific term that is within
|
||||
/// Levenshtein distance
|
||||
@@ -129,30 +110,39 @@ impl FuzzyTermQuery {
|
||||
}
|
||||
|
||||
fn specialized_weight(&self) -> crate::Result<AutomatonWeight<DfaWrapper>> {
|
||||
// LEV_BUILDER is a HashMap, whose `get` method returns an Option
|
||||
match LEV_BUILDER.get(&(self.distance, self.transposition_cost_one)) {
|
||||
// Unwrap the option and build the Ok(AutomatonWeight)
|
||||
Some(automaton_builder) => {
|
||||
let term_text = self.term.as_str().ok_or_else(|| {
|
||||
crate::TantivyError::InvalidArgument(
|
||||
"The fuzzy term query requires a string term.".to_string(),
|
||||
)
|
||||
})?;
|
||||
let automaton = if self.prefix {
|
||||
automaton_builder.build_prefix_dfa(term_text)
|
||||
} else {
|
||||
automaton_builder.build_dfa(term_text)
|
||||
};
|
||||
Ok(AutomatonWeight::new(
|
||||
self.term.field(),
|
||||
DfaWrapper(automaton),
|
||||
static AUTOMATON_BUILDER: [[OnceCell<LevenshteinAutomatonBuilder>; 2]; 3] = [
|
||||
[OnceCell::new(), OnceCell::new()],
|
||||
[OnceCell::new(), OnceCell::new()],
|
||||
[OnceCell::new(), OnceCell::new()],
|
||||
];
|
||||
|
||||
let automaton_builder = AUTOMATON_BUILDER
|
||||
.get(self.distance as usize)
|
||||
.ok_or_else(|| {
|
||||
InvalidArgument(format!(
|
||||
"Levenshtein distance of {} is not allowed. Choose a value less than {}",
|
||||
self.distance,
|
||||
AUTOMATON_BUILDER.len()
|
||||
))
|
||||
}
|
||||
None => Err(InvalidArgument(format!(
|
||||
"Levenshtein distance of {} is not allowed. Choose a value in the {:?} range",
|
||||
self.distance, VALID_LEVENSHTEIN_DISTANCE_RANGE
|
||||
))),
|
||||
}
|
||||
})?
|
||||
.get(self.transposition_cost_one as usize)
|
||||
.unwrap()
|
||||
.get_or_init(|| {
|
||||
LevenshteinAutomatonBuilder::new(self.distance, self.transposition_cost_one)
|
||||
});
|
||||
|
||||
let term_text = self.term.as_str().ok_or_else(|| {
|
||||
InvalidArgument("The fuzzy term query requires a string term.".to_string())
|
||||
})?;
|
||||
let automaton = if self.prefix {
|
||||
automaton_builder.build_prefix_dfa(term_text)
|
||||
} else {
|
||||
automaton_builder.build_dfa(term_text)
|
||||
};
|
||||
Ok(AutomatonWeight::new(
|
||||
self.term.field(),
|
||||
DfaWrapper(automaton),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,7 +16,6 @@ mod phrase_query;
|
||||
mod query;
|
||||
mod query_parser;
|
||||
mod range_query;
|
||||
mod range_query_ip_fastfield;
|
||||
mod regex_query;
|
||||
mod reqopt_scorer;
|
||||
mod scorer;
|
||||
|
||||
@@ -45,7 +45,7 @@ impl Query for MoreLikeThisQuery {
|
||||
fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
|
||||
let searcher = match enable_scoring {
|
||||
EnableScoring::Enabled(searcher) => searcher,
|
||||
EnableScoring::Disabled(_) => {
|
||||
EnableScoring::Disabled { .. } => {
|
||||
let err = "MoreLikeThisQuery requires to enable scoring.".to_string();
|
||||
return Err(crate::TantivyError::InvalidArgument(err));
|
||||
}
|
||||
|
||||
@@ -80,7 +80,7 @@ pub mod tests {
|
||||
.collect();
|
||||
let phrase_query = PhraseQuery::new(terms);
|
||||
let phrase_weight =
|
||||
phrase_query.phrase_weight(EnableScoring::Disabled(searcher.schema()))?;
|
||||
phrase_query.phrase_weight(EnableScoring::disabled_from_schema(searcher.schema()))?;
|
||||
let mut phrase_scorer = phrase_weight.scorer(searcher.segment_reader(0), 1.0)?;
|
||||
assert_eq!(phrase_scorer.doc(), 1);
|
||||
assert_eq!(phrase_scorer.advance(), TERMINATED);
|
||||
@@ -361,7 +361,7 @@ pub mod tests {
|
||||
let query_parser = QueryParser::for_index(&index, vec![json_field]);
|
||||
let phrase_query = query_parser.parse_query(query).unwrap();
|
||||
let phrase_weight = phrase_query
|
||||
.weight(EnableScoring::Disabled(searcher.schema()))
|
||||
.weight(EnableScoring::disabled_from_schema(searcher.schema()))
|
||||
.unwrap();
|
||||
let mut phrase_scorer = phrase_weight
|
||||
.scorer(searcher.segment_reader(0), 1.0f32)
|
||||
|
||||
@@ -109,7 +109,7 @@ impl PhraseQuery {
|
||||
let terms = self.phrase_terms();
|
||||
let bm25_weight_opt = match enable_scoring {
|
||||
EnableScoring::Enabled(searcher) => Some(Bm25Weight::for_terms(searcher, &terms)?),
|
||||
EnableScoring::Disabled(_) => None,
|
||||
EnableScoring::Disabled { .. } => None,
|
||||
};
|
||||
let mut weight = PhraseWeight::new(self.phrase_terms.clone(), bm25_weight_opt);
|
||||
if self.slop > 0 {
|
||||
|
||||
@@ -15,24 +15,55 @@ pub enum EnableScoring<'a> {
|
||||
Enabled(&'a Searcher),
|
||||
/// Pass this to disable scoring.
|
||||
/// This can improve performance.
|
||||
Disabled(&'a Schema),
|
||||
Disabled {
|
||||
/// Schema is required.
|
||||
schema: &'a Schema,
|
||||
/// Searcher should be provided if available.
|
||||
searcher_opt: Option<&'a Searcher>,
|
||||
},
|
||||
}
|
||||
|
||||
impl<'a> EnableScoring<'a> {
|
||||
/// Create using [Searcher] with scoring enabled.
|
||||
pub fn enabled_from_searcher(searcher: &'a Searcher) -> EnableScoring<'a> {
|
||||
EnableScoring::Enabled(searcher)
|
||||
}
|
||||
|
||||
/// Create using [Searcher] with scoring disabled.
|
||||
pub fn disabled_from_searcher(searcher: &'a Searcher) -> EnableScoring<'a> {
|
||||
EnableScoring::Disabled {
|
||||
schema: searcher.schema(),
|
||||
searcher_opt: Some(searcher),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create using [Schema] with scoring disabled.
|
||||
pub fn disabled_from_schema(schema: &'a Schema) -> EnableScoring<'a> {
|
||||
Self::Disabled {
|
||||
schema,
|
||||
searcher_opt: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the searcher if available.
|
||||
pub fn searcher(&self) -> Option<&Searcher> {
|
||||
match self {
|
||||
EnableScoring::Enabled(searcher) => Some(searcher),
|
||||
EnableScoring::Disabled { searcher_opt, .. } => searcher_opt.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the schema.
|
||||
pub fn schema(&self) -> &Schema {
|
||||
match self {
|
||||
EnableScoring::Enabled(searcher) => searcher.schema(),
|
||||
EnableScoring::Disabled(schema) => schema,
|
||||
EnableScoring::Disabled { schema, .. } => schema,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if the scoring is enabled.
|
||||
pub fn is_scoring_enabled(&self) -> bool {
|
||||
match self {
|
||||
EnableScoring::Enabled(_) => true,
|
||||
EnableScoring::Disabled(_) => false,
|
||||
}
|
||||
matches!(self, EnableScoring::Enabled(..))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -81,14 +112,14 @@ pub trait Query: QueryClone + Send + Sync + downcast_rs::Downcast + fmt::Debug {
|
||||
|
||||
/// Returns an `Explanation` for the score of the document.
|
||||
fn explain(&self, searcher: &Searcher, doc_address: DocAddress) -> crate::Result<Explanation> {
|
||||
let weight = self.weight(EnableScoring::Enabled(searcher))?;
|
||||
let weight = self.weight(EnableScoring::enabled_from_searcher(searcher))?;
|
||||
let reader = searcher.segment_reader(doc_address.segment_ord);
|
||||
weight.explain(reader, doc_address.doc_id)
|
||||
}
|
||||
|
||||
/// Returns the number of documents matching the query.
|
||||
fn count(&self, searcher: &Searcher) -> crate::Result<usize> {
|
||||
let weight = self.weight(EnableScoring::Disabled(searcher.schema()))?;
|
||||
let weight = self.weight(EnableScoring::disabled_from_searcher(searcher))?;
|
||||
let mut result = 0;
|
||||
for reader in searcher.segment_readers() {
|
||||
result += weight.count(reader)? as usize;
|
||||
|
||||
@@ -54,9 +54,9 @@ impl fmt::Debug for LogicalAst {
|
||||
if clause.is_empty() {
|
||||
write!(formatter, "<emptyclause>")?;
|
||||
} else {
|
||||
let (ref occur, ref subquery) = clause[0];
|
||||
let (occur, subquery) = &clause[0];
|
||||
write!(formatter, "({}{:?}", occur_letter(*occur), subquery)?;
|
||||
for &(ref occur, ref subquery) in &clause[1..] {
|
||||
for (occur, subquery) in &clause[1..] {
|
||||
write!(formatter, " {}{:?}", occur_letter(*occur), subquery)?;
|
||||
}
|
||||
formatter.write_str(")")?;
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
use std::collections::HashMap;
|
||||
use std::net::{AddrParseError, IpAddr};
|
||||
use std::num::{ParseFloatError, ParseIntError};
|
||||
use std::ops::Bound;
|
||||
use std::str::{FromStr, ParseBoolError};
|
||||
|
||||
use base64::engine::general_purpose::STANDARD as BASE64;
|
||||
use base64::Engine;
|
||||
use rustc_hash::FxHashMap;
|
||||
use tantivy_query_grammar::{UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral};
|
||||
|
||||
use super::logical_ast::*;
|
||||
@@ -11,9 +13,10 @@ use crate::core::Index;
|
||||
use crate::indexer::{
|
||||
convert_to_fast_value_and_get_term, set_string_and_get_terms, JsonTermWriter,
|
||||
};
|
||||
use crate::query::range_query::is_type_valid_for_fastfield_range_query;
|
||||
use crate::query::{
|
||||
AllQuery, BooleanQuery, BoostQuery, EmptyQuery, Occur, PhraseQuery, Query, RangeQuery,
|
||||
TermQuery, TermSetQuery,
|
||||
AllQuery, BooleanQuery, BoostQuery, EmptyQuery, FuzzyTermQuery, Occur, PhraseQuery, Query,
|
||||
RangeQuery, TermQuery, TermSetQuery,
|
||||
};
|
||||
use crate::schema::{
|
||||
Facet, FacetParseError, Field, FieldType, IndexRecordOption, IntoIpv6Addr, JsonObjectOptions,
|
||||
@@ -34,7 +37,7 @@ pub enum QueryParserError {
|
||||
#[error("Unsupported query: {0}")]
|
||||
UnsupportedQuery(String),
|
||||
/// The query references a field that is not in the schema
|
||||
#[error("Field does not exists: '{0}'")]
|
||||
#[error("Field does not exist: '{0}'")]
|
||||
FieldDoesNotExist(String),
|
||||
/// The query contains a term for a `u64` or `i64`-field, but the value
|
||||
/// is neither.
|
||||
@@ -160,6 +163,10 @@ fn trim_ast(logical_ast: LogicalAst) -> Option<LogicalAst> {
|
||||
/// word lexicographically between `a` and `c` (inclusive lower bound, exclusive upper bound).
|
||||
/// Inclusive bounds are `[]`, exclusive are `{}`.
|
||||
///
|
||||
/// * set terms: Using the `IN` operator, a field can be matched against a set of literals, e.g.
|
||||
/// `title: IN [a b cd]` will match documents where `title` is either `a`, `b` or `cd`, but do so
|
||||
/// more efficiently than the alternative query `title:a OR title:b OR title:c` does.
|
||||
///
|
||||
/// * date values: The query parser supports rfc3339 formatted dates. For example
|
||||
/// `"2002-10-02T15:00:00.05Z"` or `some_date_field:[2002-10-02T15:00:00Z TO
|
||||
/// 2002-10-02T18:00:00Z}`
|
||||
@@ -174,6 +181,9 @@ fn trim_ast(logical_ast: LogicalAst) -> Option<LogicalAst> {
|
||||
/// (See [`set_field_boost(...)`](QueryParser::set_field_boost)). Typically you may want to boost a
|
||||
/// title field.
|
||||
///
|
||||
/// Additionally, specific fields can be marked to use fuzzy term queries for each literal
|
||||
/// via the [`QueryParser::set_field_fuzzy`] method.
|
||||
///
|
||||
/// Phrase terms support the `~` slop operator which allows to set the phrase's matching
|
||||
/// distance in words. `"big wolf"~1` will return documents containing the phrase `"big bad wolf"`.
|
||||
#[derive(Clone)]
|
||||
@@ -182,7 +192,15 @@ pub struct QueryParser {
|
||||
default_fields: Vec<Field>,
|
||||
conjunction_by_default: bool,
|
||||
tokenizer_manager: TokenizerManager,
|
||||
boost: HashMap<Field, Score>,
|
||||
boost: FxHashMap<Field, Score>,
|
||||
fuzzy: FxHashMap<Field, Fuzzy>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Fuzzy {
|
||||
prefix: bool,
|
||||
distance: u8,
|
||||
transpose_cost_one: bool,
|
||||
}
|
||||
|
||||
fn all_negative(ast: &LogicalAst) -> bool {
|
||||
@@ -210,6 +228,7 @@ impl QueryParser {
|
||||
tokenizer_manager,
|
||||
conjunction_by_default: false,
|
||||
boost: Default::default(),
|
||||
fuzzy: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -247,6 +266,30 @@ impl QueryParser {
|
||||
self.boost.insert(field, boost);
|
||||
}
|
||||
|
||||
/// Sets the given [field][`Field`] to use [fuzzy term queries][`FuzzyTermQuery`]
|
||||
///
|
||||
/// If set, the parse will produce queries using fuzzy term queries
|
||||
/// with the given parameters for each literal matched against the given field.
|
||||
///
|
||||
/// See the [`FuzzyTermQuery::new`] and [`FuzzyTermQuery::new_prefix`] methods
|
||||
/// for the meaning of the individual parameters.
|
||||
pub fn set_field_fuzzy(
|
||||
&mut self,
|
||||
field: Field,
|
||||
prefix: bool,
|
||||
distance: u8,
|
||||
transpose_cost_one: bool,
|
||||
) {
|
||||
self.fuzzy.insert(
|
||||
field,
|
||||
Fuzzy {
|
||||
prefix,
|
||||
distance,
|
||||
transpose_cost_one,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/// Parse a query
|
||||
///
|
||||
/// Note that `parse_query` returns an error if the input
|
||||
@@ -259,7 +302,7 @@ impl QueryParser {
|
||||
/// in [Issue 5](https://github.com/fulmicoton/tantivy/issues/5)
|
||||
pub fn parse_query(&self, query: &str) -> Result<Box<dyn Query>, QueryParserError> {
|
||||
let logical_ast = self.parse_query_to_logical_ast(query)?;
|
||||
Ok(convert_to_query(logical_ast))
|
||||
Ok(convert_to_query(&self.fuzzy, logical_ast))
|
||||
}
|
||||
|
||||
/// Parse the user query into an AST.
|
||||
@@ -293,9 +336,10 @@ impl QueryParser {
|
||||
) -> Result<Term, QueryParserError> {
|
||||
let field_entry = self.schema.get_field_entry(field);
|
||||
let field_type = field_entry.field_type();
|
||||
let field_supports_ff_range_queries = field_type.is_fast()
|
||||
&& is_type_valid_for_fastfield_range_query(field_type.value_type());
|
||||
|
||||
let is_ip_and_fast = field_type.is_ip_addr() && field_type.is_fast();
|
||||
if !field_type.is_indexed() && !is_ip_and_fast {
|
||||
if !field_type.is_indexed() && !field_supports_ff_range_queries {
|
||||
return Err(QueryParserError::FieldNotIndexed(
|
||||
field_entry.name().to_string(),
|
||||
));
|
||||
@@ -363,7 +407,9 @@ impl QueryParser {
|
||||
Err(e) => Err(QueryParserError::from(e)),
|
||||
},
|
||||
FieldType::Bytes(_) => {
|
||||
let bytes = base64::decode(phrase).map_err(QueryParserError::ExpectedBase64)?;
|
||||
let bytes = BASE64
|
||||
.decode(phrase)
|
||||
.map_err(QueryParserError::ExpectedBase64)?;
|
||||
Ok(Term::from_field_bytes(field, &bytes))
|
||||
}
|
||||
FieldType::IpAddr(_) => {
|
||||
@@ -458,7 +504,9 @@ impl QueryParser {
|
||||
Err(e) => Err(QueryParserError::from(e)),
|
||||
},
|
||||
FieldType::Bytes(_) => {
|
||||
let bytes = base64::decode(phrase).map_err(QueryParserError::ExpectedBase64)?;
|
||||
let bytes = BASE64
|
||||
.decode(phrase)
|
||||
.map_err(QueryParserError::ExpectedBase64)?;
|
||||
let bytes_term = Term::from_field_bytes(field, &bytes);
|
||||
Ok(vec![LogicalLiteral::Term(bytes_term)])
|
||||
}
|
||||
@@ -660,9 +708,30 @@ impl QueryParser {
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_literal_to_query(logical_literal: LogicalLiteral) -> Box<dyn Query> {
|
||||
fn convert_literal_to_query(
|
||||
fuzzy: &FxHashMap<Field, Fuzzy>,
|
||||
logical_literal: LogicalLiteral,
|
||||
) -> Box<dyn Query> {
|
||||
match logical_literal {
|
||||
LogicalLiteral::Term(term) => Box::new(TermQuery::new(term, IndexRecordOption::WithFreqs)),
|
||||
LogicalLiteral::Term(term) => {
|
||||
if let Some(fuzzy) = fuzzy.get(&term.field()) {
|
||||
if fuzzy.prefix {
|
||||
Box::new(FuzzyTermQuery::new_prefix(
|
||||
term,
|
||||
fuzzy.distance,
|
||||
fuzzy.transpose_cost_one,
|
||||
))
|
||||
} else {
|
||||
Box::new(FuzzyTermQuery::new(
|
||||
term,
|
||||
fuzzy.distance,
|
||||
fuzzy.transpose_cost_one,
|
||||
))
|
||||
}
|
||||
} else {
|
||||
Box::new(TermQuery::new(term, IndexRecordOption::WithFreqs))
|
||||
}
|
||||
}
|
||||
LogicalLiteral::Phrase(term_with_offsets, slop) => Box::new(
|
||||
PhraseQuery::new_with_offset_and_slop(term_with_offsets, slop),
|
||||
),
|
||||
@@ -755,12 +824,12 @@ fn generate_literals_for_json_object(
|
||||
Ok(logical_literals)
|
||||
}
|
||||
|
||||
fn convert_to_query(logical_ast: LogicalAst) -> Box<dyn Query> {
|
||||
fn convert_to_query(fuzzy: &FxHashMap<Field, Fuzzy>, logical_ast: LogicalAst) -> Box<dyn Query> {
|
||||
match trim_ast(logical_ast) {
|
||||
Some(LogicalAst::Clause(trimmed_clause)) => {
|
||||
let occur_subqueries = trimmed_clause
|
||||
.into_iter()
|
||||
.map(|(occur, subquery)| (occur, convert_to_query(subquery)))
|
||||
.map(|(occur, subquery)| (occur, convert_to_query(fuzzy, subquery)))
|
||||
.collect::<Vec<_>>();
|
||||
assert!(
|
||||
!occur_subqueries.is_empty(),
|
||||
@@ -769,10 +838,10 @@ fn convert_to_query(logical_ast: LogicalAst) -> Box<dyn Query> {
|
||||
Box::new(BooleanQuery::new(occur_subqueries))
|
||||
}
|
||||
Some(LogicalAst::Leaf(trimmed_logical_literal)) => {
|
||||
convert_literal_to_query(*trimmed_logical_literal)
|
||||
convert_literal_to_query(fuzzy, *trimmed_logical_literal)
|
||||
}
|
||||
Some(LogicalAst::Boost(ast, boost)) => {
|
||||
let query = convert_to_query(*ast);
|
||||
let query = convert_to_query(fuzzy, *ast);
|
||||
let boosted_query = BoostQuery::new(query, boost);
|
||||
Box::new(boosted_query)
|
||||
}
|
||||
@@ -788,7 +857,7 @@ mod test {
|
||||
use super::{QueryParser, QueryParserError};
|
||||
use crate::query::Query;
|
||||
use crate::schema::{
|
||||
FacetOptions, Field, IndexRecordOption, Schema, Term, TextFieldIndexing, TextOptions,
|
||||
FacetOptions, Field, IndexRecordOption, Schema, Term, TextFieldIndexing, TextOptions, FAST,
|
||||
INDEXED, STORED, STRING, TEXT,
|
||||
};
|
||||
use crate::tokenizer::{
|
||||
@@ -822,6 +891,7 @@ mod test {
|
||||
schema_builder.add_json_field("json_not_indexed", STORED);
|
||||
schema_builder.add_bool_field("bool", INDEXED);
|
||||
schema_builder.add_bool_field("notindexed_bool", STORED);
|
||||
schema_builder.add_u64_field("u64_ff", FAST);
|
||||
schema_builder.build()
|
||||
}
|
||||
|
||||
@@ -1277,6 +1347,11 @@ mod test {
|
||||
r#"(Excluded(Term(type=F64, field=10, -1.5)) TO Excluded(Term(type=F64, field=10, 1.5)))"#,
|
||||
false,
|
||||
);
|
||||
test_parse_query_to_logical_ast_helper(
|
||||
"u64_ff:[7 TO 77]",
|
||||
r#"(Included(Term(type=U64, field=18, 7)) TO Included(Term(type=U64, field=18, 77)))"#,
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1568,4 +1643,41 @@ mod test {
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
pub fn test_set_field_fuzzy() {
|
||||
{
|
||||
let mut query_parser = make_query_parser();
|
||||
query_parser.set_field_fuzzy(
|
||||
query_parser.schema.get_field("title").unwrap(),
|
||||
false,
|
||||
1,
|
||||
true,
|
||||
);
|
||||
let query = query_parser.parse_query("abc").unwrap();
|
||||
assert_eq!(
|
||||
format!("{:?}", query),
|
||||
"BooleanQuery { subqueries: [(Should, FuzzyTermQuery { term: Term(type=Str, \
|
||||
field=0, \"abc\"), distance: 1, transposition_cost_one: true, prefix: false }), \
|
||||
(Should, TermQuery(Term(type=Str, field=1, \"abc\")))] }"
|
||||
);
|
||||
}
|
||||
|
||||
{
|
||||
let mut query_parser = make_query_parser();
|
||||
query_parser.set_field_fuzzy(
|
||||
query_parser.schema.get_field("text").unwrap(),
|
||||
true,
|
||||
2,
|
||||
false,
|
||||
);
|
||||
let query = query_parser.parse_query("abc").unwrap();
|
||||
assert_eq!(
|
||||
format!("{:?}", query),
|
||||
"BooleanQuery { subqueries: [(Should, TermQuery(Term(type=Str, field=0, \
|
||||
\"abc\"))), (Should, FuzzyTermQuery { term: Term(type=Str, field=1, \"abc\"), \
|
||||
distance: 2, transposition_cost_one: false, prefix: true })] }"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
210
src/query/range_query/fast_field_range_query.rs
Normal file
210
src/query/range_query/fast_field_range_query.rs
Normal file
@@ -0,0 +1,210 @@
|
||||
use core::fmt;
|
||||
use std::ops::RangeInclusive;
|
||||
use std::sync::Arc;
|
||||
|
||||
use fastfield_codecs::Column;
|
||||
|
||||
use crate::fastfield::{MakeZero, MultiValuedFastFieldReader};
|
||||
use crate::{DocId, DocSet, TERMINATED};
|
||||
|
||||
/// Helper to have a cursor over a vec of docids
|
||||
struct VecCursor {
|
||||
docs: Vec<u32>,
|
||||
current_pos: usize,
|
||||
}
|
||||
impl VecCursor {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
docs: Vec::with_capacity(32),
|
||||
current_pos: 0,
|
||||
}
|
||||
}
|
||||
fn next(&mut self) -> Option<u32> {
|
||||
self.current_pos += 1;
|
||||
self.current()
|
||||
}
|
||||
#[inline]
|
||||
fn current(&self) -> Option<u32> {
|
||||
self.docs.get(self.current_pos).copied()
|
||||
}
|
||||
fn get_cleared_data(&mut self) -> &mut Vec<u32> {
|
||||
self.docs.clear();
|
||||
self.current_pos = 0;
|
||||
&mut self.docs
|
||||
}
|
||||
fn last_value(&self) -> Option<u32> {
|
||||
self.docs.iter().last().cloned()
|
||||
}
|
||||
fn is_empty(&self) -> bool {
|
||||
self.current().is_none()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum FastFieldCardinality<T: MakeZero> {
|
||||
SingleValue(Arc<dyn Column<T>>),
|
||||
MultiValue(MultiValuedFastFieldReader<T>),
|
||||
}
|
||||
|
||||
impl<T: MakeZero + PartialOrd + Copy + fmt::Debug> FastFieldCardinality<T> {
|
||||
fn num_docs(&self) -> u32 {
|
||||
match self {
|
||||
FastFieldCardinality::SingleValue(single_value) => single_value.num_vals(),
|
||||
FastFieldCardinality::MultiValue(multi_value) => {
|
||||
multi_value.get_index_reader().num_docs()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct RangeDocSet<T: MakeZero> {
|
||||
/// The range filter on the values.
|
||||
value_range: RangeInclusive<T>,
|
||||
fast_field: FastFieldCardinality<T>,
|
||||
/// The next docid start range to fetch (inclusive).
|
||||
next_fetch_start: u32,
|
||||
/// Number of docs range checked in a batch.
|
||||
///
|
||||
/// There are two patterns.
|
||||
/// - We do a full scan. => We can load large chunks. We don't know in advance if seek call
|
||||
/// will come, so we start with small chunks
|
||||
/// - We load docs, interspersed with seek calls. When there are big jumps in the seek, we
|
||||
/// should load small chunks. When the seeks are small, we can employ the same strategy as on a
|
||||
/// full scan.
|
||||
fetch_horizon: u32,
|
||||
/// Current batch of loaded docs.
|
||||
loaded_docs: VecCursor,
|
||||
last_seek_pos_opt: Option<u32>,
|
||||
}
|
||||
|
||||
const DEFAULT_FETCH_HORIZON: u32 = 128;
|
||||
impl<T: MakeZero + Send + PartialOrd + Copy + fmt::Debug> RangeDocSet<T> {
|
||||
pub(crate) fn new(value_range: RangeInclusive<T>, fast_field: FastFieldCardinality<T>) -> Self {
|
||||
let mut range_docset = Self {
|
||||
value_range,
|
||||
fast_field,
|
||||
loaded_docs: VecCursor::new(),
|
||||
next_fetch_start: 0,
|
||||
fetch_horizon: DEFAULT_FETCH_HORIZON,
|
||||
last_seek_pos_opt: None,
|
||||
};
|
||||
range_docset.reset_fetch_range();
|
||||
range_docset.fetch_block();
|
||||
range_docset
|
||||
}
|
||||
|
||||
fn reset_fetch_range(&mut self) {
|
||||
self.fetch_horizon = DEFAULT_FETCH_HORIZON;
|
||||
}
|
||||
|
||||
/// Returns true if more data could be fetched
|
||||
fn fetch_block(&mut self) {
|
||||
const MAX_HORIZON: u32 = 100_000;
|
||||
while self.loaded_docs.is_empty() {
|
||||
let finished_to_end = self.fetch_horizon(self.fetch_horizon);
|
||||
if finished_to_end {
|
||||
break;
|
||||
}
|
||||
// Fetch more data, increase horizon. Horizon only gets reset when doing a seek.
|
||||
self.fetch_horizon = (self.fetch_horizon * 2).min(MAX_HORIZON);
|
||||
}
|
||||
}
|
||||
|
||||
/// check if the distance between the seek calls is large
|
||||
fn is_last_seek_distance_large(&self, new_seek: DocId) -> bool {
|
||||
if let Some(last_seek_pos) = self.last_seek_pos_opt {
|
||||
(new_seek - last_seek_pos) >= 128
|
||||
} else {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetches a block for docid range [next_fetch_start .. next_fetch_start + HORIZON]
|
||||
fn fetch_horizon(&mut self, horizon: u32) -> bool {
|
||||
let mut finished_to_end = false;
|
||||
|
||||
let limit = self.fast_field.num_docs();
|
||||
let mut end = self.next_fetch_start + horizon;
|
||||
if end >= limit {
|
||||
end = limit;
|
||||
finished_to_end = true;
|
||||
}
|
||||
|
||||
match &self.fast_field {
|
||||
FastFieldCardinality::MultiValue(multi) => {
|
||||
let last_value = self.loaded_docs.last_value();
|
||||
|
||||
multi.get_docids_for_value_range(
|
||||
self.value_range.clone(),
|
||||
self.next_fetch_start..end,
|
||||
self.loaded_docs.get_cleared_data(),
|
||||
);
|
||||
// In case of multivalues, we may have an overlap of the same docid between fetching
|
||||
// blocks
|
||||
if let Some(last_value) = last_value {
|
||||
while self.loaded_docs.current() == Some(last_value) {
|
||||
self.loaded_docs.next();
|
||||
}
|
||||
}
|
||||
}
|
||||
FastFieldCardinality::SingleValue(single) => {
|
||||
single.get_docids_for_value_range(
|
||||
self.value_range.clone(),
|
||||
self.next_fetch_start..end,
|
||||
self.loaded_docs.get_cleared_data(),
|
||||
);
|
||||
}
|
||||
}
|
||||
self.next_fetch_start = end;
|
||||
|
||||
finished_to_end
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: MakeZero + Send + PartialOrd + Copy + fmt::Debug> DocSet for RangeDocSet<T> {
|
||||
#[inline]
|
||||
fn advance(&mut self) -> DocId {
|
||||
if let Some(docid) = self.loaded_docs.next() {
|
||||
docid
|
||||
} else {
|
||||
if self.next_fetch_start >= self.fast_field.num_docs() {
|
||||
return TERMINATED;
|
||||
}
|
||||
self.fetch_block();
|
||||
self.loaded_docs.current().unwrap_or(TERMINATED)
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn doc(&self) -> DocId {
|
||||
self.loaded_docs.current().unwrap_or(TERMINATED)
|
||||
}
|
||||
|
||||
/// Advances the `DocSet` forward until reaching the target, or going to the
|
||||
/// lowest [`DocId`] greater than the target.
|
||||
///
|
||||
/// If the end of the `DocSet` is reached, [`TERMINATED`] is returned.
|
||||
///
|
||||
/// Calling `.seek(target)` on a terminated `DocSet` is legal. Implementation
|
||||
/// of `DocSet` should support it.
|
||||
///
|
||||
/// Calling `seek(TERMINATED)` is also legal and is the normal way to consume a `DocSet`.
|
||||
fn seek(&mut self, target: DocId) -> DocId {
|
||||
if self.is_last_seek_distance_large(target) {
|
||||
self.reset_fetch_range();
|
||||
}
|
||||
if target > self.next_fetch_start {
|
||||
self.next_fetch_start = target;
|
||||
}
|
||||
let mut doc = self.doc();
|
||||
debug_assert!(doc <= target);
|
||||
while doc < target {
|
||||
doc = self.advance();
|
||||
}
|
||||
self.last_seek_pos_opt = Some(target);
|
||||
doc
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> u32 {
|
||||
0 // heuristic possible by checking number of hits when fetching a block
|
||||
}
|
||||
}
|
||||
8
src/query/range_query/mod.rs
Normal file
8
src/query/range_query/mod.rs
Normal file
@@ -0,0 +1,8 @@
|
||||
mod fast_field_range_query;
|
||||
mod range_query;
|
||||
mod range_query_ip_fastfield;
|
||||
mod range_query_u64_fastfield;
|
||||
|
||||
pub(crate) use range_query::is_type_valid_for_fastfield_range_query;
|
||||
|
||||
pub use self::range_query::RangeQuery;
|
||||
@@ -1,16 +1,17 @@
|
||||
use std::io;
|
||||
use std::ops::{Bound, Range};
|
||||
|
||||
use common::BitSet;
|
||||
use common::{BinarySerializable, BitSet};
|
||||
|
||||
use super::range_query_u64_fastfield::FastFieldRangeWeight;
|
||||
use crate::core::SegmentReader;
|
||||
use crate::error::TantivyError;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::range_query_ip_fastfield::IPFastFieldRangeWeight;
|
||||
use crate::query::range_query::range_query_ip_fastfield::IPFastFieldRangeWeight;
|
||||
use crate::query::{BitSetDocSet, ConstScorer, EnableScoring, Explanation, Query, Scorer, Weight};
|
||||
use crate::schema::{Field, IndexRecordOption, Term, Type};
|
||||
use crate::termdict::{TermDictionary, TermStreamer};
|
||||
use crate::{DocId, Score};
|
||||
use crate::{DateTime, DocId, Score};
|
||||
|
||||
pub(crate) fn map_bound<TFrom, TTo, Transform: Fn(&TFrom) -> TTo>(
|
||||
bound: &Bound<TFrom>,
|
||||
@@ -203,6 +204,40 @@ impl RangeQuery {
|
||||
)
|
||||
}
|
||||
|
||||
/// Create a new `RangeQuery` over a `date` field.
|
||||
///
|
||||
/// The two `Bound` arguments make it possible to create more complex
|
||||
/// ranges than semi-inclusive range.
|
||||
///
|
||||
/// If the field is not of the type `date`, tantivy
|
||||
/// will panic when the `Weight` object is created.
|
||||
pub fn new_date_bounds(
|
||||
field: Field,
|
||||
left_bound: Bound<DateTime>,
|
||||
right_bound: Bound<DateTime>,
|
||||
) -> RangeQuery {
|
||||
let make_term_val =
|
||||
|val: &DateTime| Term::from_field_date(field, *val).value_bytes().to_owned();
|
||||
RangeQuery {
|
||||
field,
|
||||
value_type: Type::Date,
|
||||
left_bound: map_bound(&left_bound, &make_term_val),
|
||||
right_bound: map_bound(&right_bound, &make_term_val),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new `RangeQuery` over a `date` field.
|
||||
///
|
||||
/// If the field is not of the type `date`, tantivy
|
||||
/// will panic when the `Weight` object is created.
|
||||
pub fn new_date(field: Field, range: Range<DateTime>) -> RangeQuery {
|
||||
RangeQuery::new_date_bounds(
|
||||
field,
|
||||
Bound::Included(range.start),
|
||||
Bound::Excluded(range.end),
|
||||
)
|
||||
}
|
||||
|
||||
/// Create a new `RangeQuery` over a `Str` field.
|
||||
///
|
||||
/// The two `Bound` arguments make it possible to create more complex
|
||||
@@ -252,6 +287,23 @@ impl RangeQuery {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_type_valid_for_fastfield_range_query(typ: Type) -> bool {
|
||||
match typ {
|
||||
Type::U64 | Type::I64 | Type::F64 | Type::Bool | Type::Date => true,
|
||||
Type::IpAddr => true,
|
||||
Type::Str | Type::Facet | Type::Bytes | Type::Json => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if the type maps to a u64 fast field
|
||||
pub(crate) fn maps_to_u64_fastfield(typ: Type) -> bool {
|
||||
match typ {
|
||||
Type::U64 | Type::I64 | Type::F64 | Type::Bool | Type::Date => true,
|
||||
Type::IpAddr => false,
|
||||
Type::Str | Type::Facet | Type::Bytes | Type::Json => false,
|
||||
}
|
||||
}
|
||||
|
||||
impl Query for RangeQuery {
|
||||
fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
|
||||
let schema = enable_scoring.schema();
|
||||
@@ -265,12 +317,29 @@ impl Query for RangeQuery {
|
||||
return Err(TantivyError::SchemaError(err_msg));
|
||||
}
|
||||
|
||||
if field_type.is_ip_addr() && field_type.is_fast() {
|
||||
Ok(Box::new(IPFastFieldRangeWeight::new(
|
||||
self.field,
|
||||
&self.left_bound,
|
||||
&self.right_bound,
|
||||
)))
|
||||
if field_type.is_fast() && is_type_valid_for_fastfield_range_query(self.value_type) {
|
||||
if field_type.is_ip_addr() {
|
||||
Ok(Box::new(IPFastFieldRangeWeight::new(
|
||||
self.field,
|
||||
&self.left_bound,
|
||||
&self.right_bound,
|
||||
)))
|
||||
} else {
|
||||
// We run the range query on u64 value space for performance reasons and simpicity
|
||||
// assert the type maps to u64
|
||||
assert!(maps_to_u64_fastfield(self.value_type));
|
||||
let parse_from_bytes = |data: &Vec<u8>| {
|
||||
u64::from_be(BinarySerializable::deserialize(&mut &data[..]).unwrap())
|
||||
};
|
||||
|
||||
let left_bound = map_bound(&self.left_bound, &parse_from_bytes);
|
||||
let right_bound = map_bound(&self.right_bound, &parse_from_bytes);
|
||||
Ok(Box::new(FastFieldRangeWeight::new(
|
||||
self.field,
|
||||
left_bound,
|
||||
right_bound,
|
||||
)))
|
||||
}
|
||||
} else {
|
||||
Ok(Box::new(RangeWeight {
|
||||
field: self.field,
|
||||
@@ -4,16 +4,15 @@
|
||||
|
||||
use std::net::Ipv6Addr;
|
||||
use std::ops::{Bound, RangeInclusive};
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::BinarySerializable;
|
||||
use fastfield_codecs::{Column, MonotonicallyMappableToU128};
|
||||
use fastfield_codecs::MonotonicallyMappableToU128;
|
||||
|
||||
use super::fast_field_range_query::{FastFieldCardinality, RangeDocSet};
|
||||
use super::range_query::map_bound;
|
||||
use super::{ConstScorer, Explanation, Scorer, Weight};
|
||||
use crate::fastfield::MultiValuedU128FastFieldReader;
|
||||
use crate::query::{ConstScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::{Cardinality, Field};
|
||||
use crate::{DocId, DocSet, Score, SegmentReader, TantivyError, TERMINATED};
|
||||
use crate::{DocId, DocSet, Score, SegmentReader, TantivyError};
|
||||
|
||||
/// `IPFastFieldRangeWeight` uses the ip address fast field to execute range queries.
|
||||
pub struct IPFastFieldRangeWeight {
|
||||
@@ -24,13 +23,13 @@ pub struct IPFastFieldRangeWeight {
|
||||
|
||||
impl IPFastFieldRangeWeight {
|
||||
pub fn new(field: Field, left_bound: &Bound<Vec<u8>>, right_bound: &Bound<Vec<u8>>) -> Self {
|
||||
let ip_from_bound_raw_data = |data: &Vec<u8>| {
|
||||
let left_ip_u128: u128 =
|
||||
let parse_ip_from_bytes = |data: &Vec<u8>| {
|
||||
let ip_u128: u128 =
|
||||
u128::from_be(BinarySerializable::deserialize(&mut &data[..]).unwrap());
|
||||
Ipv6Addr::from_u128(left_ip_u128)
|
||||
Ipv6Addr::from_u128(ip_u128)
|
||||
};
|
||||
let left_bound = map_bound(left_bound, &ip_from_bound_raw_data);
|
||||
let right_bound = map_bound(right_bound, &ip_from_bound_raw_data);
|
||||
let left_bound = map_bound(left_bound, &parse_ip_from_bytes);
|
||||
let right_bound = map_bound(right_bound, &parse_ip_from_bytes);
|
||||
Self {
|
||||
field,
|
||||
left_bound,
|
||||
@@ -51,9 +50,9 @@ impl Weight for IPFastFieldRangeWeight {
|
||||
ip_addr_fast_field.min_value(),
|
||||
ip_addr_fast_field.max_value(),
|
||||
);
|
||||
let docset = IpRangeDocSet::new(
|
||||
let docset = RangeDocSet::new(
|
||||
value_range,
|
||||
IpFastFieldCardinality::SingleValue(ip_addr_fast_field),
|
||||
FastFieldCardinality::SingleValue(ip_addr_fast_field),
|
||||
);
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
}
|
||||
@@ -65,9 +64,9 @@ impl Weight for IPFastFieldRangeWeight {
|
||||
ip_addr_fast_field.min_value(),
|
||||
ip_addr_fast_field.max_value(),
|
||||
);
|
||||
let docset = IpRangeDocSet::new(
|
||||
let docset = RangeDocSet::new(
|
||||
value_range,
|
||||
IpFastFieldCardinality::MultiValue(ip_addr_fast_field),
|
||||
FastFieldCardinality::MultiValue(ip_addr_fast_field),
|
||||
);
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
}
|
||||
@@ -108,211 +107,6 @@ fn bound_to_value_range(
|
||||
start_value..=end_value
|
||||
}
|
||||
|
||||
/// Helper to have a cursor over a vec of docids
|
||||
struct VecCursor {
|
||||
docs: Vec<u32>,
|
||||
current_pos: usize,
|
||||
}
|
||||
impl VecCursor {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
docs: Vec::with_capacity(32),
|
||||
current_pos: 0,
|
||||
}
|
||||
}
|
||||
fn next(&mut self) -> Option<u32> {
|
||||
self.current_pos += 1;
|
||||
self.current()
|
||||
}
|
||||
#[inline]
|
||||
fn current(&self) -> Option<u32> {
|
||||
self.docs.get(self.current_pos).copied()
|
||||
}
|
||||
fn get_cleared_data(&mut self) -> &mut Vec<u32> {
|
||||
self.docs.clear();
|
||||
self.current_pos = 0;
|
||||
&mut self.docs
|
||||
}
|
||||
fn last_value(&self) -> Option<u32> {
|
||||
self.docs.iter().last().cloned()
|
||||
}
|
||||
fn is_empty(&self) -> bool {
|
||||
self.current_pos >= self.docs.len()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum IpFastFieldCardinality {
|
||||
SingleValue(Arc<dyn Column<Ipv6Addr>>),
|
||||
MultiValue(MultiValuedU128FastFieldReader<Ipv6Addr>),
|
||||
}
|
||||
|
||||
impl IpFastFieldCardinality {
|
||||
fn num_docs(&self) -> u32 {
|
||||
match self {
|
||||
IpFastFieldCardinality::SingleValue(single_value) => single_value.num_vals(),
|
||||
IpFastFieldCardinality::MultiValue(multi_value) => {
|
||||
multi_value.get_index_reader().num_docs()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct IpRangeDocSet {
|
||||
/// The range filter on the values.
|
||||
value_range: RangeInclusive<Ipv6Addr>,
|
||||
ip_addr_fast_field: IpFastFieldCardinality,
|
||||
/// The next docid start range to fetch (inclusive).
|
||||
next_fetch_start: u32,
|
||||
/// Number of docs range checked in a batch.
|
||||
///
|
||||
/// There are two patterns.
|
||||
/// - We do a full scan. => We can load large chunks. We don't know in advance if seek call
|
||||
/// will come, so we start with small chunks
|
||||
/// - We load docs, interspersed with seek calls. When there are big jumps in the seek, we
|
||||
/// should load small chunks. When the seeks are small, we can employ the same strategy as on a
|
||||
/// full scan.
|
||||
fetch_horizon: u32,
|
||||
/// Current batch of loaded docs.
|
||||
loaded_docs: VecCursor,
|
||||
last_seek_pos_opt: Option<u32>,
|
||||
}
|
||||
|
||||
const DEFAULT_FETCH_HORIZON: u32 = 128;
|
||||
impl IpRangeDocSet {
|
||||
fn new(
|
||||
value_range: RangeInclusive<Ipv6Addr>,
|
||||
ip_addr_fast_field: IpFastFieldCardinality,
|
||||
) -> Self {
|
||||
let mut ip_range_docset = Self {
|
||||
value_range,
|
||||
ip_addr_fast_field,
|
||||
loaded_docs: VecCursor::new(),
|
||||
next_fetch_start: 0,
|
||||
fetch_horizon: DEFAULT_FETCH_HORIZON,
|
||||
last_seek_pos_opt: None,
|
||||
};
|
||||
ip_range_docset.reset_fetch_range();
|
||||
ip_range_docset.fetch_block();
|
||||
ip_range_docset
|
||||
}
|
||||
|
||||
fn reset_fetch_range(&mut self) {
|
||||
self.fetch_horizon = DEFAULT_FETCH_HORIZON;
|
||||
}
|
||||
|
||||
/// Returns true if more data could be fetched
|
||||
fn fetch_block(&mut self) {
|
||||
const MAX_HORIZON: u32 = 100_000;
|
||||
while self.loaded_docs.is_empty() {
|
||||
let finished_to_end = self.fetch_horizon(self.fetch_horizon);
|
||||
if finished_to_end {
|
||||
break;
|
||||
}
|
||||
// Fetch more data, increase horizon. Horizon only gets reset when doing a seek.
|
||||
self.fetch_horizon = (self.fetch_horizon * 2).min(MAX_HORIZON);
|
||||
}
|
||||
}
|
||||
|
||||
/// check if the distance between the seek calls is large
|
||||
fn is_last_seek_distance_large(&self, new_seek: DocId) -> bool {
|
||||
if let Some(last_seek_pos) = self.last_seek_pos_opt {
|
||||
(new_seek - last_seek_pos) >= 128
|
||||
} else {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetches a block for docid range [next_fetch_start .. next_fetch_start + HORIZON]
|
||||
fn fetch_horizon(&mut self, horizon: u32) -> bool {
|
||||
let mut finished_to_end = false;
|
||||
|
||||
let limit = self.ip_addr_fast_field.num_docs();
|
||||
let mut end = self.next_fetch_start + horizon;
|
||||
if end >= limit {
|
||||
end = limit;
|
||||
finished_to_end = true;
|
||||
}
|
||||
|
||||
match &self.ip_addr_fast_field {
|
||||
IpFastFieldCardinality::MultiValue(multi) => {
|
||||
let last_value = self.loaded_docs.last_value();
|
||||
|
||||
multi.get_docids_for_value_range(
|
||||
self.value_range.clone(),
|
||||
self.next_fetch_start..end,
|
||||
self.loaded_docs.get_cleared_data(),
|
||||
);
|
||||
// In case of multivalues, we may have an overlap of the same docid between fetching
|
||||
// blocks
|
||||
if let Some(last_value) = last_value {
|
||||
while self.loaded_docs.current() == Some(last_value) {
|
||||
self.loaded_docs.next();
|
||||
}
|
||||
}
|
||||
}
|
||||
IpFastFieldCardinality::SingleValue(single) => {
|
||||
single.get_docids_for_value_range(
|
||||
self.value_range.clone(),
|
||||
self.next_fetch_start..end,
|
||||
self.loaded_docs.get_cleared_data(),
|
||||
);
|
||||
}
|
||||
}
|
||||
self.next_fetch_start = end;
|
||||
|
||||
finished_to_end
|
||||
}
|
||||
}
|
||||
|
||||
impl DocSet for IpRangeDocSet {
|
||||
#[inline]
|
||||
fn advance(&mut self) -> DocId {
|
||||
if let Some(docid) = self.loaded_docs.next() {
|
||||
docid
|
||||
} else {
|
||||
if self.next_fetch_start >= self.ip_addr_fast_field.num_docs() {
|
||||
return TERMINATED;
|
||||
}
|
||||
self.fetch_block();
|
||||
self.loaded_docs.current().unwrap_or(TERMINATED)
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn doc(&self) -> DocId {
|
||||
self.loaded_docs.current().unwrap_or(TERMINATED)
|
||||
}
|
||||
|
||||
/// Advances the `DocSet` forward until reaching the target, or going to the
|
||||
/// lowest [`DocId`] greater than the target.
|
||||
///
|
||||
/// If the end of the `DocSet` is reached, [`TERMINATED`] is returned.
|
||||
///
|
||||
/// Calling `.seek(target)` on a terminated `DocSet` is legal. Implementation
|
||||
/// of `DocSet` should support it.
|
||||
///
|
||||
/// Calling `seek(TERMINATED)` is also legal and is the normal way to consume a `DocSet`.
|
||||
fn seek(&mut self, target: DocId) -> DocId {
|
||||
if self.is_last_seek_distance_large(target) {
|
||||
self.reset_fetch_range();
|
||||
}
|
||||
if target > self.next_fetch_start {
|
||||
self.next_fetch_start = target;
|
||||
}
|
||||
let mut doc = self.doc();
|
||||
debug_assert!(doc <= target);
|
||||
while doc < target {
|
||||
doc = self.advance();
|
||||
}
|
||||
self.last_seek_pos_opt = Some(target);
|
||||
doc
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> u32 {
|
||||
0 // heuristic possible by checking number of hits when fetching a block
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use proptest::prelude::ProptestConfig;
|
||||
@@ -401,7 +195,7 @@ mod tests {
|
||||
let index = Index::create_in_ram(schema);
|
||||
|
||||
{
|
||||
let mut index_writer = index.writer(3_000_000).unwrap();
|
||||
let mut index_writer = index.writer(10_000_000).unwrap();
|
||||
for doc in docs.iter() {
|
||||
index_writer
|
||||
.add_document(doc!(
|
||||
@@ -519,8 +313,7 @@ mod bench {
|
||||
})
|
||||
.collect();
|
||||
|
||||
let index = create_index_from_docs(&docs);
|
||||
index
|
||||
create_index_from_docs(&docs)
|
||||
}
|
||||
|
||||
fn get_90_percent() -> RangeInclusive<Ipv6Addr> {
|
||||
@@ -559,7 +352,7 @@ mod bench {
|
||||
|
||||
let query = gen_query_inclusive(ip_range.start(), ip_range.end());
|
||||
let query_from_text = |text: &str| {
|
||||
QueryParser::for_index(&index, vec![])
|
||||
QueryParser::for_index(index, vec![])
|
||||
.parse_query(text)
|
||||
.unwrap()
|
||||
};
|
||||
557
src/query/range_query/range_query_u64_fastfield.rs
Normal file
557
src/query/range_query/range_query_u64_fastfield.rs
Normal file
@@ -0,0 +1,557 @@
|
||||
//! Fastfields support efficient scanning for range queries.
|
||||
//! We use this variant only if the fastfield exists, otherwise the default in `range_query` is
|
||||
//! used, which uses the term dictionary + postings.
|
||||
|
||||
use std::ops::{Bound, RangeInclusive};
|
||||
|
||||
use fastfield_codecs::MonotonicallyMappableToU64;
|
||||
|
||||
use super::fast_field_range_query::{FastFieldCardinality, RangeDocSet};
|
||||
use super::range_query::map_bound;
|
||||
use crate::query::{ConstScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::{Cardinality, Field};
|
||||
use crate::{DocId, DocSet, Score, SegmentReader, TantivyError};
|
||||
|
||||
/// `FastFieldRangeWeight` uses the fast field to execute range queries.
|
||||
pub struct FastFieldRangeWeight {
|
||||
field: Field,
|
||||
left_bound: Bound<u64>,
|
||||
right_bound: Bound<u64>,
|
||||
}
|
||||
|
||||
impl FastFieldRangeWeight {
|
||||
pub fn new(field: Field, left_bound: Bound<u64>, right_bound: Bound<u64>) -> Self {
|
||||
let left_bound = map_bound(&left_bound, &|val| *val);
|
||||
let right_bound = map_bound(&right_bound, &|val| *val);
|
||||
Self {
|
||||
field,
|
||||
left_bound,
|
||||
right_bound,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Weight for FastFieldRangeWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
let field_type = reader.schema().get_field_entry(self.field).field_type();
|
||||
match field_type.fastfield_cardinality().unwrap() {
|
||||
Cardinality::SingleValue => {
|
||||
let fast_field = reader.fast_fields().u64_lenient(self.field)?;
|
||||
let value_range = bound_to_value_range(
|
||||
&self.left_bound,
|
||||
&self.right_bound,
|
||||
fast_field.min_value(),
|
||||
fast_field.max_value(),
|
||||
);
|
||||
let docset =
|
||||
RangeDocSet::new(value_range, FastFieldCardinality::SingleValue(fast_field));
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
}
|
||||
Cardinality::MultiValues => {
|
||||
let fast_field = reader.fast_fields().u64s_lenient(self.field)?;
|
||||
let value_range = bound_to_value_range(
|
||||
&self.left_bound,
|
||||
&self.right_bound,
|
||||
fast_field.min_value(),
|
||||
fast_field.max_value(),
|
||||
);
|
||||
let docset =
|
||||
RangeDocSet::new(value_range, FastFieldCardinality::MultiValue(fast_field));
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
|
||||
let mut scorer = self.scorer(reader, 1.0)?;
|
||||
if scorer.seek(doc) != doc {
|
||||
return Err(TantivyError::InvalidArgument(format!(
|
||||
"Document #({}) does not match",
|
||||
doc
|
||||
)));
|
||||
}
|
||||
let explanation = Explanation::new("Const", scorer.score());
|
||||
|
||||
Ok(explanation)
|
||||
}
|
||||
}
|
||||
|
||||
fn bound_to_value_range<T: MonotonicallyMappableToU64>(
|
||||
left_bound: &Bound<T>,
|
||||
right_bound: &Bound<T>,
|
||||
min_value: T,
|
||||
max_value: T,
|
||||
) -> RangeInclusive<T> {
|
||||
let start_value = match left_bound {
|
||||
Bound::Included(val) => *val,
|
||||
Bound::Excluded(val) => T::from_u64(val.to_u64() + 1),
|
||||
Bound::Unbounded => min_value,
|
||||
};
|
||||
|
||||
let end_value = match right_bound {
|
||||
Bound::Included(val) => *val,
|
||||
Bound::Excluded(val) => T::from_u64(val.to_u64() - 1),
|
||||
Bound::Unbounded => max_value,
|
||||
};
|
||||
start_value..=end_value
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use proptest::prelude::ProptestConfig;
|
||||
use proptest::strategy::Strategy;
|
||||
use proptest::{prop_oneof, proptest};
|
||||
use rand::rngs::StdRng;
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::SeedableRng;
|
||||
|
||||
use super::*;
|
||||
use crate::collector::Count;
|
||||
use crate::query::QueryParser;
|
||||
use crate::schema::{NumericOptions, Schema, FAST, INDEXED, STORED, STRING};
|
||||
use crate::Index;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Doc {
|
||||
pub id_name: String,
|
||||
pub id: u64,
|
||||
}
|
||||
|
||||
fn operation_strategy() -> impl Strategy<Value = Doc> {
|
||||
prop_oneof![
|
||||
(0u64..10_000u64).prop_map(doc_from_id_1),
|
||||
(1u64..10_000u64).prop_map(doc_from_id_2),
|
||||
]
|
||||
}
|
||||
|
||||
pub fn doc_from_id_1(id: u64) -> Doc {
|
||||
let id = id * 1000;
|
||||
Doc {
|
||||
id_name: id.to_string(),
|
||||
id,
|
||||
}
|
||||
}
|
||||
fn doc_from_id_2(id: u64) -> Doc {
|
||||
let id = id * 1000;
|
||||
Doc {
|
||||
id_name: (id - 1).to_string(),
|
||||
id,
|
||||
}
|
||||
}
|
||||
|
||||
proptest! {
|
||||
#![proptest_config(ProptestConfig::with_cases(10))]
|
||||
#[test]
|
||||
fn test_range_for_docs_prop(ops in proptest::collection::vec(operation_strategy(), 1..1000)) {
|
||||
assert!(test_id_range_for_docs(ops).is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_regression1_test() {
|
||||
let ops = vec![doc_from_id_1(0)];
|
||||
assert!(test_id_range_for_docs(ops).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_regression2_test() {
|
||||
let ops = vec![
|
||||
doc_from_id_1(52),
|
||||
doc_from_id_1(63),
|
||||
doc_from_id_1(12),
|
||||
doc_from_id_2(91),
|
||||
doc_from_id_2(33),
|
||||
];
|
||||
assert!(test_id_range_for_docs(ops).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_regression3_test() {
|
||||
let ops = vec![doc_from_id_1(1), doc_from_id_1(2), doc_from_id_1(3)];
|
||||
assert!(test_id_range_for_docs(ops).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_regression4_test() {
|
||||
let ops = vec![doc_from_id_2(100)];
|
||||
assert!(test_id_range_for_docs(ops).is_ok());
|
||||
}
|
||||
|
||||
pub fn create_index_from_docs(docs: &[Doc]) -> Index {
|
||||
let mut schema_builder = Schema::builder();
|
||||
let id_u64_field = schema_builder.add_u64_field("id", INDEXED | STORED | FAST);
|
||||
let ids_u64_field = schema_builder.add_u64_field(
|
||||
"ids",
|
||||
NumericOptions::default()
|
||||
.set_fast(Cardinality::MultiValues)
|
||||
.set_indexed(),
|
||||
);
|
||||
|
||||
let id_f64_field = schema_builder.add_f64_field("id_f64", INDEXED | STORED | FAST);
|
||||
let ids_f64_field = schema_builder.add_f64_field(
|
||||
"ids_f64",
|
||||
NumericOptions::default()
|
||||
.set_fast(Cardinality::MultiValues)
|
||||
.set_indexed(),
|
||||
);
|
||||
|
||||
let id_i64_field = schema_builder.add_i64_field("id_i64", INDEXED | STORED | FAST);
|
||||
let ids_i64_field = schema_builder.add_i64_field(
|
||||
"ids_i64",
|
||||
NumericOptions::default()
|
||||
.set_fast(Cardinality::MultiValues)
|
||||
.set_indexed(),
|
||||
);
|
||||
|
||||
let text_field = schema_builder.add_text_field("id_name", STRING | STORED);
|
||||
let schema = schema_builder.build();
|
||||
let index = Index::create_in_ram(schema);
|
||||
|
||||
{
|
||||
let mut index_writer = index.writer(3_000_000).unwrap();
|
||||
for doc in docs.iter() {
|
||||
index_writer
|
||||
.add_document(doc!(
|
||||
ids_i64_field => doc.id as i64,
|
||||
ids_i64_field => doc.id as i64,
|
||||
ids_f64_field => doc.id as f64,
|
||||
ids_f64_field => doc.id as f64,
|
||||
ids_u64_field => doc.id,
|
||||
ids_u64_field => doc.id,
|
||||
id_u64_field => doc.id,
|
||||
id_f64_field => doc.id as f64,
|
||||
id_i64_field => doc.id as i64,
|
||||
text_field => doc.id_name.to_string(),
|
||||
))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
index_writer.commit().unwrap();
|
||||
}
|
||||
index
|
||||
}
|
||||
|
||||
fn test_id_range_for_docs(docs: Vec<Doc>) -> crate::Result<()> {
|
||||
let index = create_index_from_docs(&docs);
|
||||
let reader = index.reader().unwrap();
|
||||
let searcher = reader.searcher();
|
||||
|
||||
let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
|
||||
|
||||
let get_num_hits = |query| searcher.search(&query, &(Count)).unwrap();
|
||||
let query_from_text = |text: &str| {
|
||||
QueryParser::for_index(&index, vec![])
|
||||
.parse_query(text)
|
||||
.unwrap()
|
||||
};
|
||||
|
||||
let gen_query_inclusive = |field: &str, from: u64, to: u64| {
|
||||
format!("{}:[{} TO {}]", field, &from.to_string(), &to.to_string())
|
||||
};
|
||||
|
||||
let test_sample = |sample_docs: Vec<Doc>| {
|
||||
let mut ids: Vec<u64> = sample_docs.iter().map(|doc| doc.id).collect();
|
||||
ids.sort();
|
||||
let expected_num_hits = docs
|
||||
.iter()
|
||||
.filter(|doc| (ids[0]..=ids[1]).contains(&doc.id))
|
||||
.count();
|
||||
|
||||
let query = gen_query_inclusive("id", ids[0], ids[1]);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
|
||||
let query = gen_query_inclusive("ids", ids[0], ids[1]);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
|
||||
// Intersection search
|
||||
let id_filter = sample_docs[0].id_name.to_string();
|
||||
let expected_num_hits = docs
|
||||
.iter()
|
||||
.filter(|doc| (ids[0]..=ids[1]).contains(&doc.id) && doc.id_name == id_filter)
|
||||
.count();
|
||||
let query = format!(
|
||||
"{} AND id_name:{}",
|
||||
gen_query_inclusive("id", ids[0], ids[1]),
|
||||
&id_filter
|
||||
);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
let query = format!(
|
||||
"{} AND id_name:{}",
|
||||
gen_query_inclusive("id_f64", ids[0], ids[1]),
|
||||
&id_filter
|
||||
);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
let query = format!(
|
||||
"{} AND id_name:{}",
|
||||
gen_query_inclusive("id_i64", ids[0], ids[1]),
|
||||
&id_filter
|
||||
);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
|
||||
// Intersection search on multivalue id field
|
||||
let id_filter = sample_docs[0].id_name.to_string();
|
||||
let query = format!(
|
||||
"{} AND id_name:{}",
|
||||
gen_query_inclusive("ids", ids[0], ids[1]),
|
||||
&id_filter
|
||||
);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
let query = format!(
|
||||
"{} AND id_name:{}",
|
||||
gen_query_inclusive("ids_f64", ids[0], ids[1]),
|
||||
&id_filter
|
||||
);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
let query = format!(
|
||||
"{} AND id_name:{}",
|
||||
gen_query_inclusive("ids_i64", ids[0], ids[1]),
|
||||
&id_filter
|
||||
);
|
||||
assert_eq!(get_num_hits(query_from_text(&query)), expected_num_hits);
|
||||
};
|
||||
|
||||
test_sample(vec![docs[0].clone(), docs[0].clone()]);
|
||||
|
||||
let samples: Vec<_> = docs.choose_multiple(&mut rng, 3).collect();
|
||||
|
||||
if samples.len() > 1 {
|
||||
test_sample(vec![samples[0].clone(), samples[1].clone()]);
|
||||
test_sample(vec![samples[1].clone(), samples[1].clone()]);
|
||||
}
|
||||
if samples.len() > 2 {
|
||||
test_sample(vec![samples[1].clone(), samples[2].clone()]);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "unstable"))]
|
||||
mod bench {
|
||||
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use test::Bencher;
|
||||
|
||||
use super::tests::*;
|
||||
use super::*;
|
||||
use crate::collector::Count;
|
||||
use crate::query::QueryParser;
|
||||
use crate::Index;
|
||||
|
||||
fn get_index_0_to_100() -> Index {
|
||||
let mut rng = StdRng::from_seed([1u8; 32]);
|
||||
let num_vals = 100_000;
|
||||
let docs: Vec<_> = (0..num_vals)
|
||||
.map(|_i| {
|
||||
let id_name = if rng.gen_bool(0.01) {
|
||||
"veryfew".to_string() // 1%
|
||||
} else if rng.gen_bool(0.1) {
|
||||
"few".to_string() // 9%
|
||||
} else {
|
||||
"many".to_string() // 90%
|
||||
};
|
||||
Doc {
|
||||
id_name,
|
||||
id: rng.gen_range(0..100),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
create_index_from_docs(&docs)
|
||||
}
|
||||
|
||||
fn get_90_percent() -> RangeInclusive<u64> {
|
||||
0..=90
|
||||
}
|
||||
|
||||
fn get_10_percent() -> RangeInclusive<u64> {
|
||||
0..=10
|
||||
}
|
||||
|
||||
fn get_1_percent() -> RangeInclusive<u64> {
|
||||
10..=10
|
||||
}
|
||||
|
||||
fn excute_query(
|
||||
field: &str,
|
||||
id_range: RangeInclusive<u64>,
|
||||
suffix: &str,
|
||||
index: &Index,
|
||||
) -> usize {
|
||||
let gen_query_inclusive = |from: &u64, to: &u64| {
|
||||
format!(
|
||||
"{}:[{} TO {}] {}",
|
||||
field,
|
||||
&from.to_string(),
|
||||
&to.to_string(),
|
||||
suffix
|
||||
)
|
||||
};
|
||||
|
||||
let query = gen_query_inclusive(id_range.start(), id_range.end());
|
||||
let query_from_text = |text: &str| {
|
||||
QueryParser::for_index(index, vec![])
|
||||
.parse_query(text)
|
||||
.unwrap()
|
||||
};
|
||||
let query = query_from_text(&query);
|
||||
let reader = index.reader().unwrap();
|
||||
let searcher = reader.searcher();
|
||||
searcher.search(&query, &(Count)).unwrap()
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_90_percent(), "", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_10_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_10_percent(), "", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_1_percent(), "", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_10_percent_intersect_with_10_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_10_percent(), "AND id_name:few", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_intersect_with_10_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_1_percent(), "AND id_name:few", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_intersect_with_90_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_1_percent(), "AND id_name:many", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_intersect_with_1_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_1_percent(), "AND id_name:veryfew", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_10_percent_intersect_with_90_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_10_percent(), "AND id_name:many", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_intersect_with_90_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_90_percent(), "AND id_name:many", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_intersect_with_10_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_90_percent(), "AND id_name:few", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_intersect_with_1_percent(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("id", get_90_percent(), "AND id_name:veryfew", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_90_percent(), "", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_10_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_10_percent(), "", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_1_percent(), "", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_10_percent_intersect_with_10_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_10_percent(), "AND id_name:few", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_intersect_with_10_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_1_percent(), "AND id_name:few", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_intersect_with_90_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_1_percent(), "AND id_name:many", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_1_percent_intersect_with_1_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_1_percent(), "AND id_name:veryfew", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_10_percent_intersect_with_90_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_10_percent(), "AND id_name:many", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_intersect_with_90_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_90_percent(), "AND id_name:many", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_intersect_with_10_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_90_percent(), "AND id_name:few", &index));
|
||||
}
|
||||
|
||||
#[bench]
|
||||
fn bench_id_range_hit_90_percent_intersect_with_1_percent_multi(bench: &mut Bencher) {
|
||||
let index = get_index_0_to_100();
|
||||
|
||||
bench.iter(|| excute_query("ids", get_90_percent(), "AND id_name:veryfew", &index));
|
||||
}
|
||||
}
|
||||
@@ -158,7 +158,8 @@ mod tests {
|
||||
let term_a = Term::from_field_text(text_field, "a");
|
||||
let term_query = TermQuery::new(term_a, IndexRecordOption::Basic);
|
||||
let searcher = index.reader()?.searcher();
|
||||
let term_weight = term_query.weight(EnableScoring::Disabled(searcher.schema()))?;
|
||||
let term_weight =
|
||||
term_query.weight(EnableScoring::disabled_from_schema(searcher.schema()))?;
|
||||
let mut term_scorer = term_weight.scorer(searcher.segment_reader(0u32), 1.0)?;
|
||||
assert_eq!(term_scorer.doc(), 0u32);
|
||||
term_scorer.seek(1u32);
|
||||
|
||||
@@ -99,7 +99,7 @@ impl TermQuery {
|
||||
EnableScoring::Enabled(searcher) => {
|
||||
Bm25Weight::for_terms(searcher, &[self.term.clone()])?
|
||||
}
|
||||
EnableScoring::Disabled(_schema) => {
|
||||
EnableScoring::Disabled { .. } => {
|
||||
Bm25Weight::new(Explanation::new("<no score>".to_string(), 1.0f32), 1.0f32)
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use std::net::IpAddr;
|
||||
use std::str::FromStr;
|
||||
|
||||
use base64::engine::general_purpose::STANDARD as BASE64;
|
||||
use base64::Engine;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value as JsonValue;
|
||||
use thiserror::Error;
|
||||
@@ -358,7 +360,8 @@ impl FieldType {
|
||||
json: JsonValue::String(field_text),
|
||||
}),
|
||||
FieldType::Facet(_) => Ok(Value::Facet(Facet::from(&field_text))),
|
||||
FieldType::Bytes(_) => base64::decode(&field_text)
|
||||
FieldType::Bytes(_) => BASE64
|
||||
.decode(&field_text)
|
||||
.map(Value::Bytes)
|
||||
.map_err(|_| ValueParsingError::InvalidBase64 { base64: field_text }),
|
||||
FieldType::JsonObject(_) => Err(ValueParsingError::TypeError {
|
||||
|
||||
@@ -258,7 +258,7 @@ mod tests {
|
||||
let field = schema.get_field("body").unwrap();
|
||||
let field_entry = schema.get_field_entry(field);
|
||||
assert!(matches!(field_entry.field_type(),
|
||||
&FieldType::Str(ref text_options)
|
||||
FieldType::Str(text_options)
|
||||
if text_options.get_indexing_options().unwrap().tokenizer() == "default"));
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use std::fmt;
|
||||
use std::net::Ipv6Addr;
|
||||
|
||||
use base64::engine::general_purpose::STANDARD as BASE64;
|
||||
use base64::Engine;
|
||||
use serde::de::Visitor;
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
use serde_json::Map;
|
||||
@@ -51,7 +53,7 @@ impl Serialize for Value {
|
||||
Value::Bool(b) => serializer.serialize_bool(b),
|
||||
Value::Date(ref date) => time::serde::rfc3339::serialize(&date.into_utc(), serializer),
|
||||
Value::Facet(ref facet) => facet.serialize(serializer),
|
||||
Value::Bytes(ref bytes) => serializer.serialize_str(&base64::encode(bytes)),
|
||||
Value::Bytes(ref bytes) => serializer.serialize_str(&BASE64.encode(bytes)),
|
||||
Value::JsonObject(ref obj) => obj.serialize(serializer),
|
||||
Value::IpAddr(ref obj) => {
|
||||
// Ensure IpV4 addresses get serialized as IpV4, but excluding IpV6 loopback.
|
||||
|
||||
@@ -124,9 +124,9 @@ impl Snippet {
|
||||
///
|
||||
/// Fragments must be valid in the sense that `&text[fragment.start..fragment.stop]`\
|
||||
/// has to be a valid string.
|
||||
fn search_fragments<'a>(
|
||||
fn search_fragments(
|
||||
tokenizer: &TextAnalyzer,
|
||||
text: &'a str,
|
||||
text: &str,
|
||||
terms: &BTreeMap<String, Score>,
|
||||
max_num_chars: usize,
|
||||
) -> Vec<FragmentCandidate> {
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::io;
|
||||
use std::iter::Sum;
|
||||
use std::num::NonZeroUsize;
|
||||
use std::ops::{AddAssign, Range};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
@@ -33,23 +34,29 @@ pub struct StoreReader {
|
||||
|
||||
/// The cache for decompressed blocks.
|
||||
struct BlockCache {
|
||||
cache: Mutex<LruCache<usize, Block>>,
|
||||
cache_hits: Arc<AtomicUsize>,
|
||||
cache_misses: Arc<AtomicUsize>,
|
||||
cache: Option<Mutex<LruCache<usize, Block>>>,
|
||||
cache_hits: AtomicUsize,
|
||||
cache_misses: AtomicUsize,
|
||||
}
|
||||
|
||||
impl BlockCache {
|
||||
fn get_from_cache(&self, pos: usize) -> Option<Block> {
|
||||
if let Some(block) = self.cache.lock().unwrap().get(&pos) {
|
||||
if let Some(block) = self
|
||||
.cache
|
||||
.as_ref()
|
||||
.and_then(|cache| cache.lock().unwrap().get(&pos).cloned())
|
||||
{
|
||||
self.cache_hits.fetch_add(1, Ordering::SeqCst);
|
||||
return Some(block.clone());
|
||||
return Some(block);
|
||||
}
|
||||
self.cache_misses.fetch_add(1, Ordering::SeqCst);
|
||||
None
|
||||
}
|
||||
|
||||
fn put_into_cache(&self, pos: usize, data: Block) {
|
||||
self.cache.lock().unwrap().put(pos, data);
|
||||
if let Some(cache) = self.cache.as_ref() {
|
||||
cache.lock().unwrap().put(pos, data);
|
||||
}
|
||||
}
|
||||
|
||||
fn stats(&self) -> CacheStats {
|
||||
@@ -59,13 +66,18 @@ impl BlockCache {
|
||||
num_entries: self.len(),
|
||||
}
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.cache.lock().unwrap().len()
|
||||
self.cache
|
||||
.as_ref()
|
||||
.map_or(0, |cache| cache.lock().unwrap().len())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn peek_lru(&self) -> Option<usize> {
|
||||
self.cache.lock().unwrap().peek_lru().map(|(&k, _)| k)
|
||||
self.cache
|
||||
.as_ref()
|
||||
.and_then(|cache| cache.lock().unwrap().peek_lru().map(|(&k, _)| k))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -113,7 +125,8 @@ impl StoreReader {
|
||||
decompressor: footer.decompressor,
|
||||
data: data_file,
|
||||
cache: BlockCache {
|
||||
cache: Mutex::new(LruCache::new(cache_size)),
|
||||
cache: NonZeroUsize::new(cache_size)
|
||||
.map(|cache_size| Mutex::new(LruCache::new(cache_size))),
|
||||
cache_hits: Default::default(),
|
||||
cache_misses: Default::default(),
|
||||
},
|
||||
|
||||
@@ -113,7 +113,7 @@ mod bench {
|
||||
}
|
||||
|
||||
/// Create a dictionary of random strings.
|
||||
fn rand_dict(num_terms: usize) -> crate::Result<TermDictionary> {
|
||||
fn rand_dict(num_terms: usize) -> std::io::Result<TermDictionary> {
|
||||
let buffer: Vec<u8> = {
|
||||
let mut terms = vec![];
|
||||
for _i in 0..num_terms {
|
||||
|
||||
@@ -99,7 +99,7 @@ fn test_term_dictionary_stream() -> crate::Result<()> {
|
||||
.collect();
|
||||
let buffer: Vec<u8> = {
|
||||
let mut term_dictionary_builder = TermDictionaryBuilder::create(vec![]).unwrap();
|
||||
for &(ref id, ref i) in &ids {
|
||||
for (id, i) in &ids {
|
||||
term_dictionary_builder
|
||||
.insert(id.as_bytes(), &make_term_info(*i as u64))
|
||||
.unwrap();
|
||||
@@ -112,14 +112,14 @@ fn test_term_dictionary_stream() -> crate::Result<()> {
|
||||
let mut streamer = term_dictionary.stream()?;
|
||||
let mut i = 0;
|
||||
while let Some((streamer_k, streamer_v)) = streamer.next() {
|
||||
let &(ref key, ref v) = &ids[i];
|
||||
let (key, v) = &ids[i];
|
||||
assert_eq!(streamer_k, key.as_bytes());
|
||||
assert_eq!(streamer_v, &make_term_info(*v as u64));
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let &(ref key, ref val) = &ids[2047];
|
||||
let (key, val) = &ids[2047];
|
||||
assert_eq!(
|
||||
term_dictionary.get(key.as_bytes())?,
|
||||
Some(make_term_info(*val as u64))
|
||||
@@ -160,7 +160,7 @@ fn test_stream_range() -> crate::Result<()> {
|
||||
.collect();
|
||||
let buffer: Vec<u8> = {
|
||||
let mut term_dictionary_builder = TermDictionaryBuilder::create(vec![]).unwrap();
|
||||
for &(ref id, ref i) in &ids {
|
||||
for (id, i) in &ids {
|
||||
term_dictionary_builder
|
||||
.insert(id.as_bytes(), &make_term_info(*i as u64))
|
||||
.unwrap();
|
||||
@@ -173,14 +173,14 @@ fn test_stream_range() -> crate::Result<()> {
|
||||
let term_dictionary: TermDictionary = TermDictionary::open(file)?;
|
||||
{
|
||||
for i in (0..20).chain(6000..8_000) {
|
||||
let &(ref target_key, _) = &ids[i];
|
||||
let (target_key, _) = &ids[i];
|
||||
let mut streamer = term_dictionary
|
||||
.range()
|
||||
.ge(target_key.as_bytes())
|
||||
.into_stream()?;
|
||||
for j in 0..3 {
|
||||
let (streamer_k, streamer_v) = streamer.next().unwrap();
|
||||
let &(ref key, ref v) = &ids[i + j];
|
||||
let (key, v) = &ids[i + j];
|
||||
assert_eq!(str::from_utf8(streamer_k).unwrap(), key);
|
||||
assert_eq!(streamer_v.doc_freq, *v);
|
||||
assert_eq!(streamer_v, &make_term_info(*v as u64));
|
||||
@@ -190,14 +190,14 @@ fn test_stream_range() -> crate::Result<()> {
|
||||
|
||||
{
|
||||
for i in (0..20).chain(BLOCK_SIZE - 10..BLOCK_SIZE + 10) {
|
||||
let &(ref target_key, _) = &ids[i];
|
||||
let (target_key, _) = &ids[i];
|
||||
let mut streamer = term_dictionary
|
||||
.range()
|
||||
.gt(target_key.as_bytes())
|
||||
.into_stream()?;
|
||||
for j in 0..3 {
|
||||
let (streamer_k, streamer_v) = streamer.next().unwrap();
|
||||
let &(ref key, ref v) = &ids[i + j + 1];
|
||||
let (key, v) = &ids[i + j + 1];
|
||||
assert_eq!(streamer_k, key.as_bytes());
|
||||
assert_eq!(streamer_v.doc_freq, *v);
|
||||
}
|
||||
@@ -207,8 +207,8 @@ fn test_stream_range() -> crate::Result<()> {
|
||||
{
|
||||
for i in (0..20).chain(BLOCK_SIZE - 10..BLOCK_SIZE + 10) {
|
||||
for j in 0..3 {
|
||||
let &(ref fst_key, _) = &ids[i];
|
||||
let &(ref last_key, _) = &ids[i + j];
|
||||
let (fst_key, _) = &ids[i];
|
||||
let (last_key, _) = &ids[i + j];
|
||||
let mut streamer = term_dictionary
|
||||
.range()
|
||||
.ge(fst_key.as_bytes())
|
||||
|
||||
@@ -52,6 +52,8 @@
|
||||
//! remove their inflection. This tokenizer is slower than the default one,
|
||||
//! but is recommended to improve recall.
|
||||
//!
|
||||
//! # Custom tokenizer Library
|
||||
//! Avoid using tantivy as dependency and prefer `tantivy-tokenizer-api` instead.
|
||||
//!
|
||||
//! # Custom tokenizers
|
||||
//!
|
||||
@@ -124,6 +126,7 @@ mod facet_tokenizer;
|
||||
mod lower_caser;
|
||||
mod ngram_tokenizer;
|
||||
mod raw_tokenizer;
|
||||
mod regex_tokenizer;
|
||||
mod remove_long;
|
||||
mod simple_tokenizer;
|
||||
mod split_compound_words;
|
||||
@@ -134,21 +137,24 @@ mod tokenizer;
|
||||
mod tokenizer_manager;
|
||||
mod whitespace_tokenizer;
|
||||
|
||||
pub use tokenizer_api::{
|
||||
BoxTokenFilter, BoxTokenStream, Token, TokenFilter, TokenStream, Tokenizer,
|
||||
};
|
||||
|
||||
pub use self::alphanum_only::AlphaNumOnlyFilter;
|
||||
pub use self::ascii_folding_filter::AsciiFoldingFilter;
|
||||
pub use self::facet_tokenizer::FacetTokenizer;
|
||||
pub use self::lower_caser::LowerCaser;
|
||||
pub use self::ngram_tokenizer::NgramTokenizer;
|
||||
pub use self::raw_tokenizer::RawTokenizer;
|
||||
pub use self::regex_tokenizer::RegexTokenizer;
|
||||
pub use self::remove_long::RemoveLongFilter;
|
||||
pub use self::simple_tokenizer::SimpleTokenizer;
|
||||
pub use self::split_compound_words::SplitCompoundWords;
|
||||
pub use self::stemmer::{Language, Stemmer};
|
||||
pub use self::stop_word_filter::StopWordFilter;
|
||||
pub use self::tokenized_string::{PreTokenizedStream, PreTokenizedString};
|
||||
pub use self::tokenizer::{
|
||||
BoxTokenFilter, BoxTokenStream, TextAnalyzer, Token, TokenFilter, TokenStream, Tokenizer,
|
||||
};
|
||||
pub use self::tokenizer::TextAnalyzer;
|
||||
pub use self::tokenizer_manager::TokenizerManager;
|
||||
pub use self::whitespace_tokenizer::WhitespaceTokenizer;
|
||||
|
||||
|
||||
@@ -303,8 +303,7 @@ mod tests {
|
||||
|
||||
use super::{utf8_codepoint_width, CodepointFrontiers, NgramTokenizer, StutteringIterator};
|
||||
use crate::tokenizer::tests::assert_token;
|
||||
use crate::tokenizer::tokenizer::Tokenizer;
|
||||
use crate::tokenizer::{BoxTokenStream, Token};
|
||||
use crate::tokenizer::{BoxTokenStream, Token, Tokenizer};
|
||||
|
||||
fn test_helper(mut tokenizer: BoxTokenStream) -> Vec<Token> {
|
||||
let mut tokens: Vec<Token> = vec![];
|
||||
|
||||
158
src/tokenizer/regex_tokenizer.rs
Normal file
158
src/tokenizer/regex_tokenizer.rs
Normal file
@@ -0,0 +1,158 @@
|
||||
use regex::Regex;
|
||||
|
||||
use super::{BoxTokenStream, Token, TokenStream, Tokenizer};
|
||||
use crate::TantivyError;
|
||||
|
||||
/// Tokenize the text by using a regex pattern to split.
|
||||
/// Each match of the regex emits a distinct token, empty tokens will not be emitted. Anchors such
|
||||
/// as `\A` will match the text from the part where the last token was emitted or the beginning of
|
||||
/// the complete text if no token was emitted yet.
|
||||
///
|
||||
/// Example: `` 'aaa' bbb 'ccc' 'ddd' `` with the pattern `` '(?:\w*)' `` will be tokenized as
|
||||
/// followed:
|
||||
///
|
||||
/// | Term | aaa | ccc | ddd |
|
||||
/// |----------|------|--------|-------|
|
||||
/// | Position | 1 | 2 | 3 |
|
||||
/// | Offsets |0,5 | 10,15 | 16,21 |
|
||||
///
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use tantivy::tokenizer::*;
|
||||
///
|
||||
/// let tokenizer = RegexTokenizer::new(r"'(?:\w*)'").unwrap();
|
||||
/// let mut stream = tokenizer.token_stream("'aaa' bbb 'ccc' 'ddd'");
|
||||
/// {
|
||||
/// let token = stream.next().unwrap();
|
||||
/// assert_eq!(token.text, "'aaa'");
|
||||
/// assert_eq!(token.offset_from, 0);
|
||||
/// assert_eq!(token.offset_to, 5);
|
||||
/// }
|
||||
/// {
|
||||
/// let token = stream.next().unwrap();
|
||||
/// assert_eq!(token.text, "'ccc'");
|
||||
/// assert_eq!(token.offset_from, 10);
|
||||
/// assert_eq!(token.offset_to, 15);
|
||||
/// }
|
||||
/// {
|
||||
/// let token = stream.next().unwrap();
|
||||
/// assert_eq!(token.text, "'ddd'");
|
||||
/// assert_eq!(token.offset_from, 16);
|
||||
/// assert_eq!(token.offset_to, 21);
|
||||
/// }
|
||||
/// assert!(stream.next().is_none());
|
||||
/// ```
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RegexTokenizer {
|
||||
regex: Regex,
|
||||
}
|
||||
|
||||
impl RegexTokenizer {
|
||||
/// Creates a new RegexTokenizer.
|
||||
pub fn new(regex_pattern: &str) -> crate::Result<RegexTokenizer> {
|
||||
Regex::new(regex_pattern)
|
||||
.map_err(|_| TantivyError::InvalidArgument(regex_pattern.to_owned()))
|
||||
.map(|regex| Self { regex })
|
||||
}
|
||||
}
|
||||
|
||||
impl Tokenizer for RegexTokenizer {
|
||||
fn token_stream<'a>(&self, text: &'a str) -> BoxTokenStream<'a> {
|
||||
BoxTokenStream::from(RegexTokenStream {
|
||||
regex: self.regex.clone(),
|
||||
text,
|
||||
token: Token::default(),
|
||||
cursor: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RegexTokenStream<'a> {
|
||||
regex: Regex,
|
||||
text: &'a str,
|
||||
token: Token,
|
||||
cursor: usize,
|
||||
}
|
||||
|
||||
impl<'a> TokenStream for RegexTokenStream<'a> {
|
||||
fn advance(&mut self) -> bool {
|
||||
let Some(regex_match) = self.regex.find(self.text) else {
|
||||
return false;
|
||||
};
|
||||
if regex_match.as_str().is_empty() {
|
||||
return false;
|
||||
}
|
||||
self.token.text.clear();
|
||||
self.token.text.push_str(regex_match.as_str());
|
||||
|
||||
self.token.offset_from = self.cursor + regex_match.start();
|
||||
self.cursor += regex_match.end();
|
||||
self.token.offset_to = self.cursor;
|
||||
|
||||
self.token.position = self.token.position.wrapping_add(1);
|
||||
|
||||
self.text = &self.text[regex_match.end()..];
|
||||
true
|
||||
}
|
||||
|
||||
fn token(&self) -> &Token {
|
||||
&self.token
|
||||
}
|
||||
|
||||
fn token_mut(&mut self) -> &mut Token {
|
||||
&mut self.token
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tokenizer::regex_tokenizer::RegexTokenizer;
|
||||
use crate::tokenizer::tests::assert_token;
|
||||
use crate::tokenizer::{TextAnalyzer, Token};
|
||||
|
||||
#[test]
|
||||
fn test_regex_tokenizer() {
|
||||
let tokens = token_stream_helper("'aaa' bbb 'ccc' 'ddd'", r"'(?:\w*)'");
|
||||
assert_eq!(tokens.len(), 3);
|
||||
assert_token(&tokens[0], 0, "'aaa'", 0, 5);
|
||||
assert_token(&tokens[1], 1, "'ccc'", 10, 15);
|
||||
assert_token(&tokens[2], 2, "'ddd'", 16, 21);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regexp_tokenizer_no_match_on_input_data() {
|
||||
let tokens = token_stream_helper("aaa", r"'(?:\w*)'");
|
||||
assert_eq!(tokens.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regexp_tokenizer_no_input_data() {
|
||||
let tokens = token_stream_helper("", r"'(?:\w*)'");
|
||||
assert_eq!(tokens.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regexp_tokenizer_error_on_invalid_regex() {
|
||||
let tokenizer = RegexTokenizer::new(r"\@");
|
||||
assert_eq!(tokenizer.is_err(), true);
|
||||
assert_eq!(
|
||||
tokenizer.err().unwrap().to_string(),
|
||||
"An invalid argument was passed: '\\@'"
|
||||
);
|
||||
}
|
||||
|
||||
fn token_stream_helper(text: &str, pattern: &str) -> Vec<Token> {
|
||||
let r = RegexTokenizer::new(pattern).unwrap();
|
||||
let a = TextAnalyzer::from(r);
|
||||
let mut token_stream = a.token_stream(text);
|
||||
let mut tokens: Vec<Token> = vec![];
|
||||
let mut add_token = |token: &Token| {
|
||||
tokens.push(token.clone());
|
||||
};
|
||||
token_stream.process(&mut add_token);
|
||||
tokens
|
||||
}
|
||||
}
|
||||
@@ -26,7 +26,7 @@ impl<'a> SimpleTokenStream<'a> {
|
||||
// search for the end of the current token.
|
||||
fn search_token_end(&mut self) -> usize {
|
||||
(&mut self.chars)
|
||||
.filter(|&(_, ref c)| !c.is_alphanumeric())
|
||||
.filter(|(_, c)| !c.is_alphanumeric())
|
||||
.map(|(offset, _)| offset)
|
||||
.next()
|
||||
.unwrap_or(self.text.len())
|
||||
|
||||
@@ -1,42 +1,9 @@
|
||||
/// The tokenizer module contains all of the tools used to process
|
||||
/// text in `tantivy`.
|
||||
use std::borrow::{Borrow, BorrowMut};
|
||||
use std::ops::{Deref, DerefMut};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokenizer_api::{BoxTokenFilter, BoxTokenStream, Tokenizer};
|
||||
|
||||
use crate::tokenizer::empty_tokenizer::EmptyTokenizer;
|
||||
|
||||
/// Token
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
|
||||
pub struct Token {
|
||||
/// Offset (byte index) of the first character of the token.
|
||||
/// Offsets shall not be modified by token filters.
|
||||
pub offset_from: usize,
|
||||
/// Offset (byte index) of the last character of the token + 1.
|
||||
/// The text that generated the token should be obtained by
|
||||
/// &text[token.offset_from..token.offset_to]
|
||||
pub offset_to: usize,
|
||||
/// Position, expressed in number of tokens.
|
||||
pub position: usize,
|
||||
/// Actual text content of the token.
|
||||
pub text: String,
|
||||
/// Is the length expressed in term of number of original tokens.
|
||||
pub position_length: usize,
|
||||
}
|
||||
|
||||
impl Default for Token {
|
||||
fn default() -> Token {
|
||||
Token {
|
||||
offset_from: 0,
|
||||
offset_to: 0,
|
||||
position: usize::MAX,
|
||||
text: String::with_capacity(200),
|
||||
position_length: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `TextAnalyzer` tokenizes an input text into tokens and modifies the resulting `TokenStream`.
|
||||
///
|
||||
/// It simply wraps a `Tokenizer` and a list of `TokenFilter` that are applied sequentially.
|
||||
@@ -112,200 +79,3 @@ impl Clone for TextAnalyzer {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `Tokenizer` are in charge of splitting text into a stream of token
|
||||
/// before indexing.
|
||||
///
|
||||
/// See the [module documentation](crate::tokenizer) for more detail.
|
||||
///
|
||||
/// # Warning
|
||||
///
|
||||
/// This API may change to use associated types.
|
||||
pub trait Tokenizer: 'static + Send + Sync + TokenizerClone {
|
||||
/// Creates a token stream for a given `str`.
|
||||
fn token_stream<'a>(&self, text: &'a str) -> BoxTokenStream<'a>;
|
||||
}
|
||||
|
||||
pub trait TokenizerClone {
|
||||
fn box_clone(&self) -> Box<dyn Tokenizer>;
|
||||
}
|
||||
|
||||
impl<T: Tokenizer + Clone> TokenizerClone for T {
|
||||
fn box_clone(&self) -> Box<dyn Tokenizer> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> TokenStream for Box<dyn TokenStream + 'a> {
|
||||
fn advance(&mut self) -> bool {
|
||||
let token_stream: &mut dyn TokenStream = self.borrow_mut();
|
||||
token_stream.advance()
|
||||
}
|
||||
|
||||
fn token<'b>(&'b self) -> &'b Token {
|
||||
let token_stream: &'b (dyn TokenStream + 'a) = self.borrow();
|
||||
token_stream.token()
|
||||
}
|
||||
|
||||
fn token_mut<'b>(&'b mut self) -> &'b mut Token {
|
||||
let token_stream: &'b mut (dyn TokenStream + 'a) = self.borrow_mut();
|
||||
token_stream.token_mut()
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple wrapper of `Box<dyn TokenStream + 'a>`.
|
||||
///
|
||||
/// See [`TokenStream`] for more information.
|
||||
pub struct BoxTokenStream<'a>(Box<dyn TokenStream + 'a>);
|
||||
|
||||
impl<'a, T> From<T> for BoxTokenStream<'a>
|
||||
where T: TokenStream + 'a
|
||||
{
|
||||
fn from(token_stream: T) -> BoxTokenStream<'a> {
|
||||
BoxTokenStream(Box::new(token_stream))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Deref for BoxTokenStream<'a> {
|
||||
type Target = dyn TokenStream + 'a;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&*self.0
|
||||
}
|
||||
}
|
||||
impl<'a> DerefMut for BoxTokenStream<'a> {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
&mut *self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple wrapper of `Box<dyn TokenFilter + 'a>`.
|
||||
///
|
||||
/// See [`TokenFilter`] for more information.
|
||||
pub struct BoxTokenFilter(Box<dyn TokenFilter>);
|
||||
|
||||
impl Deref for BoxTokenFilter {
|
||||
type Target = dyn TokenFilter;
|
||||
|
||||
fn deref(&self) -> &dyn TokenFilter {
|
||||
&*self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: TokenFilter> From<T> for BoxTokenFilter {
|
||||
fn from(tokenizer: T) -> BoxTokenFilter {
|
||||
BoxTokenFilter(Box::new(tokenizer))
|
||||
}
|
||||
}
|
||||
|
||||
/// `TokenStream` is the result of the tokenization.
|
||||
///
|
||||
/// It consists consumable stream of `Token`s.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```
|
||||
/// use tantivy::tokenizer::*;
|
||||
///
|
||||
/// let tokenizer = TextAnalyzer::from(SimpleTokenizer)
|
||||
/// .filter(RemoveLongFilter::limit(40))
|
||||
/// .filter(LowerCaser);
|
||||
/// let mut token_stream = tokenizer.token_stream("Hello, happy tax payer");
|
||||
/// {
|
||||
/// let token = token_stream.next().unwrap();
|
||||
/// assert_eq!(&token.text, "hello");
|
||||
/// assert_eq!(token.offset_from, 0);
|
||||
/// assert_eq!(token.offset_to, 5);
|
||||
/// assert_eq!(token.position, 0);
|
||||
/// }
|
||||
/// {
|
||||
/// let token = token_stream.next().unwrap();
|
||||
/// assert_eq!(&token.text, "happy");
|
||||
/// assert_eq!(token.offset_from, 7);
|
||||
/// assert_eq!(token.offset_to, 12);
|
||||
/// assert_eq!(token.position, 1);
|
||||
/// }
|
||||
/// ```
|
||||
pub trait TokenStream {
|
||||
/// Advance to the next token
|
||||
///
|
||||
/// Returns false if there are no other tokens.
|
||||
fn advance(&mut self) -> bool;
|
||||
|
||||
/// Returns a reference to the current token.
|
||||
fn token(&self) -> &Token;
|
||||
|
||||
/// Returns a mutable reference to the current token.
|
||||
fn token_mut(&mut self) -> &mut Token;
|
||||
|
||||
/// Helper to iterate over tokens. It
|
||||
/// simply combines a call to `.advance()`
|
||||
/// and `.token()`.
|
||||
///
|
||||
/// ```
|
||||
/// use tantivy::tokenizer::*;
|
||||
///
|
||||
/// let tokenizer = TextAnalyzer::from(SimpleTokenizer)
|
||||
/// .filter(RemoveLongFilter::limit(40))
|
||||
/// .filter(LowerCaser);
|
||||
/// let mut token_stream = tokenizer.token_stream("Hello, happy tax payer");
|
||||
/// while let Some(token) = token_stream.next() {
|
||||
/// println!("Token {:?}", token.text);
|
||||
/// }
|
||||
/// ```
|
||||
fn next(&mut self) -> Option<&Token> {
|
||||
if self.advance() {
|
||||
Some(self.token())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper function to consume the entire `TokenStream`
|
||||
/// and push the tokens to a sink function.
|
||||
///
|
||||
/// Remove this.
|
||||
fn process(&mut self, sink: &mut dyn FnMut(&Token)) {
|
||||
while self.advance() {
|
||||
sink(self.token());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait TokenFilterClone {
|
||||
fn box_clone(&self) -> BoxTokenFilter;
|
||||
}
|
||||
|
||||
/// Trait for the pluggable components of `Tokenizer`s.
|
||||
pub trait TokenFilter: 'static + Send + Sync + TokenFilterClone {
|
||||
/// Wraps a token stream and returns the modified one.
|
||||
fn transform<'a>(&self, token_stream: BoxTokenStream<'a>) -> BoxTokenStream<'a>;
|
||||
}
|
||||
|
||||
impl<T: TokenFilter + Clone> TokenFilterClone for T {
|
||||
fn box_clone(&self) -> BoxTokenFilter {
|
||||
BoxTokenFilter::from(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::Token;
|
||||
|
||||
#[test]
|
||||
fn clone() {
|
||||
let t1 = Token {
|
||||
position: 1,
|
||||
offset_from: 2,
|
||||
offset_to: 3,
|
||||
text: "abc".to_string(),
|
||||
position_length: 1,
|
||||
};
|
||||
let t2 = t1.clone();
|
||||
|
||||
assert_eq!(t1.position, t2.position);
|
||||
assert_eq!(t1.offset_from, t2.offset_from);
|
||||
assert_eq!(t1.offset_to, t2.offset_to);
|
||||
assert_eq!(t1.text, t2.text);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,7 +26,7 @@ impl<'a> WhitespaceTokenStream<'a> {
|
||||
// search for the end of the current token.
|
||||
fn search_token_end(&mut self) -> usize {
|
||||
(&mut self.chars)
|
||||
.filter(|&(_, ref c)| c.is_ascii_whitespace())
|
||||
.filter(|(_, c)| c.is_ascii_whitespace())
|
||||
.map(|(offset, _)| offset)
|
||||
.next()
|
||||
.unwrap_or(self.text.len())
|
||||
|
||||
@@ -209,8 +209,11 @@ where
|
||||
}
|
||||
|
||||
/// Inserts a `(key, value)` pair in the term dictionary.
|
||||
/// Keys have to be inserted in order.
|
||||
///
|
||||
/// *Keys have to be inserted in order.*
|
||||
/// # Panics
|
||||
///
|
||||
/// Will panics if keys are inserted in an invalid order.
|
||||
#[inline]
|
||||
pub fn insert<K: AsRef<[u8]>>(
|
||||
&mut self,
|
||||
@@ -295,6 +298,17 @@ where
|
||||
Ok(wrt.into_inner()?)
|
||||
}
|
||||
}
|
||||
|
||||
impl<TValueWriter> Writer<Vec<u8>, TValueWriter>
|
||||
where TValueWriter: value::ValueWriter
|
||||
{
|
||||
#[inline]
|
||||
pub fn insert_cannot_fail<K: AsRef<[u8]>>(&mut self, key: K, value: &TValueWriter::Value) {
|
||||
self.insert(key, value)
|
||||
.expect("SSTable over a Vec should never fail");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use std::io;
|
||||
@@ -457,8 +471,8 @@ mod test {
|
||||
fn bound_strategy() -> impl Strategy<Value = Bound<String>> {
|
||||
prop_oneof![
|
||||
Just(Bound::<String>::Unbounded),
|
||||
"[a-d]*".prop_map(|key| Bound::Included(key)),
|
||||
"[a-d]*".prop_map(|key| Bound::Excluded(key)),
|
||||
"[a-c]{0,5}".prop_map(|key| Bound::Included(key)),
|
||||
"[a-c]{0,5}".prop_map(|key| Bound::Excluded(key)),
|
||||
]
|
||||
}
|
||||
|
||||
@@ -476,14 +490,14 @@ mod test {
|
||||
|(left, right)| match (extract_key(left.as_ref()), extract_key(right.as_ref())) {
|
||||
(None, _) => true,
|
||||
(_, None) => true,
|
||||
(left, right) => left <= right,
|
||||
(left, right) => left < right,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
proptest! {
|
||||
#[test]
|
||||
fn test_prop_test_ranges(words in prop::collection::btree_set("[a-d]*", 1..100),
|
||||
fn test_proptest_sstable_ranges(words in prop::collection::btree_set("[a-c]{0,6}", 1..100),
|
||||
(lower_bound, upper_bound) in bounds_strategy(),
|
||||
) {
|
||||
// TODO tweak block size.
|
||||
|
||||
@@ -209,7 +209,7 @@ impl ArenaHashMap {
|
||||
}
|
||||
}
|
||||
|
||||
/// `update` create a new entry for a given key if it does not exists
|
||||
/// `update` create a new entry for a given key if it does not exist
|
||||
/// or updates the existing entry.
|
||||
///
|
||||
/// The actual logic for this update is define in the `updater`
|
||||
@@ -237,7 +237,7 @@ impl ArenaHashMap {
|
||||
let bucket = probe.next_probe();
|
||||
let kv: KeyValue = self.table[bucket];
|
||||
if kv.is_empty() {
|
||||
// The key does not exists yet.
|
||||
// The key does not exist yet.
|
||||
let val = updater(None);
|
||||
let num_bytes = std::mem::size_of::<u16>() + key.len() + std::mem::size_of::<V>();
|
||||
let key_addr = self.memory_arena.allocate_space(num_bytes);
|
||||
|
||||
11
tokenizer-api/Cargo.toml
Normal file
11
tokenizer-api/Cargo.toml
Normal file
@@ -0,0 +1,11 @@
|
||||
[package]
|
||||
name = "tantivy-tokenizer-api"
|
||||
version = "0.1.0"
|
||||
license = "MIT"
|
||||
edition = "2021"
|
||||
description = "Tokenizer API of tantivy"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1.0.152", features = ["derive"] }
|
||||
6
tokenizer-api/README.md
Normal file
6
tokenizer-api/README.md
Normal file
@@ -0,0 +1,6 @@
|
||||
|
||||
#Tokenizer-API
|
||||
|
||||
An API to interface a tokenizer with tantivy.
|
||||
|
||||
The API will be kept stable in order to not break support for existing tokenizers.
|
||||
197
tokenizer-api/src/lib.rs
Normal file
197
tokenizer-api/src/lib.rs
Normal file
@@ -0,0 +1,197 @@
|
||||
//! Tokenizer are in charge of chopping text into a stream of tokens
|
||||
//! ready for indexing. This is an seperate crate from tantivy, so implementors don't need to update
|
||||
//! for each new tantivy version.
|
||||
//!
|
||||
//! To add support for a tokenizer, implement the [`Tokenizer`](crate::Tokenizer) trait.
|
||||
//! Checkout the [tantivy repo](https://github.com/quickwit-oss/tantivy/tree/main/src/tokenizer) for some examples.
|
||||
|
||||
use std::borrow::{Borrow, BorrowMut};
|
||||
use std::ops::{Deref, DerefMut};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Token
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
|
||||
pub struct Token {
|
||||
/// Offset (byte index) of the first character of the token.
|
||||
/// Offsets shall not be modified by token filters.
|
||||
pub offset_from: usize,
|
||||
/// Offset (byte index) of the last character of the token + 1.
|
||||
/// The text that generated the token should be obtained by
|
||||
/// &text[token.offset_from..token.offset_to]
|
||||
pub offset_to: usize,
|
||||
/// Position, expressed in number of tokens.
|
||||
pub position: usize,
|
||||
/// Actual text content of the token.
|
||||
pub text: String,
|
||||
/// Is the length expressed in term of number of original tokens.
|
||||
pub position_length: usize,
|
||||
}
|
||||
|
||||
impl Default for Token {
|
||||
fn default() -> Token {
|
||||
Token {
|
||||
offset_from: 0,
|
||||
offset_to: 0,
|
||||
position: usize::MAX,
|
||||
text: String::with_capacity(200),
|
||||
position_length: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `Tokenizer` are in charge of splitting text into a stream of token
|
||||
/// before indexing.
|
||||
///
|
||||
/// # Warning
|
||||
///
|
||||
/// This API may change to use associated types.
|
||||
pub trait Tokenizer: 'static + Send + Sync + TokenizerClone {
|
||||
/// Creates a token stream for a given `str`.
|
||||
fn token_stream<'a>(&self, text: &'a str) -> BoxTokenStream<'a>;
|
||||
}
|
||||
|
||||
pub trait TokenizerClone {
|
||||
fn box_clone(&self) -> Box<dyn Tokenizer>;
|
||||
}
|
||||
|
||||
impl<T: Tokenizer + Clone> TokenizerClone for T {
|
||||
fn box_clone(&self) -> Box<dyn Tokenizer> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple wrapper of `Box<dyn TokenStream + 'a>`.
|
||||
///
|
||||
/// See [`TokenStream`] for more information.
|
||||
pub struct BoxTokenStream<'a>(Box<dyn TokenStream + 'a>);
|
||||
|
||||
impl<'a, T> From<T> for BoxTokenStream<'a>
|
||||
where T: TokenStream + 'a
|
||||
{
|
||||
fn from(token_stream: T) -> BoxTokenStream<'a> {
|
||||
BoxTokenStream(Box::new(token_stream))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Deref for BoxTokenStream<'a> {
|
||||
type Target = dyn TokenStream + 'a;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&*self.0
|
||||
}
|
||||
}
|
||||
impl<'a> DerefMut for BoxTokenStream<'a> {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
&mut *self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> TokenStream for Box<dyn TokenStream + 'a> {
|
||||
fn advance(&mut self) -> bool {
|
||||
let token_stream: &mut dyn TokenStream = self.borrow_mut();
|
||||
token_stream.advance()
|
||||
}
|
||||
|
||||
fn token<'b>(&'b self) -> &'b Token {
|
||||
let token_stream: &'b (dyn TokenStream + 'a) = self.borrow();
|
||||
token_stream.token()
|
||||
}
|
||||
|
||||
fn token_mut<'b>(&'b mut self) -> &'b mut Token {
|
||||
let token_stream: &'b mut (dyn TokenStream + 'a) = self.borrow_mut();
|
||||
token_stream.token_mut()
|
||||
}
|
||||
}
|
||||
|
||||
/// `TokenStream` is the result of the tokenization.
|
||||
///
|
||||
/// It consists consumable stream of `Token`s.
|
||||
pub trait TokenStream {
|
||||
/// Advance to the next token
|
||||
///
|
||||
/// Returns false if there are no other tokens.
|
||||
fn advance(&mut self) -> bool;
|
||||
|
||||
/// Returns a reference to the current token.
|
||||
fn token(&self) -> &Token;
|
||||
|
||||
/// Returns a mutable reference to the current token.
|
||||
fn token_mut(&mut self) -> &mut Token;
|
||||
|
||||
/// Helper to iterate over tokens. It
|
||||
/// simply combines a call to `.advance()`
|
||||
/// and `.token()`.
|
||||
fn next(&mut self) -> Option<&Token> {
|
||||
if self.advance() {
|
||||
Some(self.token())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper function to consume the entire `TokenStream`
|
||||
/// and push the tokens to a sink function.
|
||||
fn process(&mut self, sink: &mut dyn FnMut(&Token)) {
|
||||
while self.advance() {
|
||||
sink(self.token());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple wrapper of `Box<dyn TokenFilter + 'a>`.
|
||||
///
|
||||
/// See [`TokenFilter`] for more information.
|
||||
pub struct BoxTokenFilter(Box<dyn TokenFilter>);
|
||||
|
||||
impl Deref for BoxTokenFilter {
|
||||
type Target = dyn TokenFilter;
|
||||
|
||||
fn deref(&self) -> &dyn TokenFilter {
|
||||
&*self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: TokenFilter> From<T> for BoxTokenFilter {
|
||||
fn from(tokenizer: T) -> BoxTokenFilter {
|
||||
BoxTokenFilter(Box::new(tokenizer))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait TokenFilterClone {
|
||||
fn box_clone(&self) -> BoxTokenFilter;
|
||||
}
|
||||
|
||||
/// Trait for the pluggable components of `Tokenizer`s.
|
||||
pub trait TokenFilter: 'static + Send + Sync + TokenFilterClone {
|
||||
/// Wraps a token stream and returns the modified one.
|
||||
fn transform<'a>(&self, token_stream: BoxTokenStream<'a>) -> BoxTokenStream<'a>;
|
||||
}
|
||||
|
||||
impl<T: TokenFilter + Clone> TokenFilterClone for T {
|
||||
fn box_clone(&self) -> BoxTokenFilter {
|
||||
BoxTokenFilter::from(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn clone() {
|
||||
let t1 = Token {
|
||||
position: 1,
|
||||
offset_from: 2,
|
||||
offset_to: 3,
|
||||
text: "abc".to_string(),
|
||||
position_length: 1,
|
||||
};
|
||||
let t2 = t1.clone();
|
||||
|
||||
assert_eq!(t1.position, t2.position);
|
||||
assert_eq!(t1.offset_from, t2.offset_from);
|
||||
assert_eq!(t1.offset_to, t2.offset_to);
|
||||
assert_eq!(t1.text, t2.text);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user