mirror of
https://github.com/quickwit-oss/tantivy.git
synced 2026-06-17 16:00:42 +00:00
Compare commits
26 Commits
postings-w
...
seek_dange
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7577a0bdda | ||
|
|
c096b2ad89 | ||
|
|
ac7a3d347c | ||
|
|
03520a0719 | ||
|
|
86a4c47bed | ||
|
|
fb23e8908f | ||
|
|
3ca510dff0 | ||
|
|
3cb400c300 | ||
|
|
ef13489d63 | ||
|
|
9f7aea4765 | ||
|
|
2c8536ab11 | ||
|
|
05f4c02ac5 | ||
|
|
d137779219 | ||
|
|
8f9846ac80 | ||
|
|
52e24a9757 | ||
|
|
00714326af | ||
|
|
799f7b4646 | ||
|
|
fc88d80726 | ||
|
|
6a684e7c38 | ||
|
|
94fe52cc67 | ||
|
|
2ff39f6f7f | ||
|
|
1d06328cb3 | ||
|
|
7fd1dbe9f5 | ||
|
|
b19f0ddc77 | ||
|
|
b4acfcf881 | ||
|
|
3a8240b123 |
2
.github/workflows/coverage.yml
vendored
2
.github/workflows/coverage.yml
vendored
@@ -28,7 +28,7 @@ jobs:
|
||||
- name: Generate code coverage
|
||||
run: cargo +nightly-2025-12-01 llvm-cov --all-features --workspace --doctests --lcov --output-path lcov.info
|
||||
- name: Upload coverage to Codecov
|
||||
uses: codecov/codecov-action@57e3a136b779b570ffcdbf80b3bdc90e7fab3de2 # v6.0.0
|
||||
uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0
|
||||
continue-on-error: true
|
||||
with:
|
||||
token: ${{ secrets.CODECOV_TOKEN }} # not required for public repos
|
||||
|
||||
@@ -66,6 +66,9 @@ fn bench_agg(mut group: InputGroup<Index>) {
|
||||
register!(group, terms_status_with_terms_zipf_1000_sub_agg);
|
||||
register!(group, terms_zipf_1000_with_terms_status_sub_agg);
|
||||
register!(group, terms_status_with_histogram);
|
||||
register!(group, terms_status_with_date_histogram);
|
||||
register!(group, terms_status_with_date_histogram_hard_bounds);
|
||||
register!(group, terms_status_with_date_histogram_and_sibling_terms);
|
||||
register!(group, terms_zipf_1000);
|
||||
register!(group, terms_zipf_1000_with_histogram);
|
||||
register!(group, terms_zipf_1000_with_avg_sub_agg);
|
||||
@@ -390,6 +393,57 @@ fn terms_status_with_histogram(index: &Index) {
|
||||
execute_agg(index, agg_req);
|
||||
}
|
||||
|
||||
fn terms_status_with_date_histogram(index: &Index) {
|
||||
let agg_req = json!({
|
||||
"my_texts": {
|
||||
"terms": { "field": "text_few_terms_status" },
|
||||
"aggs": {
|
||||
"over_time": { "date_histogram": { "field": "timestamp", "fixed_interval": "1h" } }
|
||||
}
|
||||
}
|
||||
});
|
||||
execute_agg(index, agg_req);
|
||||
}
|
||||
|
||||
/// Same fused terms × date_histogram, but with `hard_bounds`. The timestamps span 0..120h; the
|
||||
/// bounds drop only the first and last hour (ms: 1h=3_600_000, 119h=428_400_000), so almost every
|
||||
/// doc is in-bounds. This exercises the collector's hard-bounds path: `bounds.contains` runs per
|
||||
/// doc (the `all_docs_in_bounds` short-circuit is off) and the rare out-of-bounds doc takes the
|
||||
/// `term_counts` branch.
|
||||
fn terms_status_with_date_histogram_hard_bounds(index: &Index) {
|
||||
let agg_req = json!({
|
||||
"my_texts": {
|
||||
"terms": { "field": "text_few_terms_status" },
|
||||
"aggs": {
|
||||
"over_time": {
|
||||
"date_histogram": {
|
||||
"field": "timestamp",
|
||||
"fixed_interval": "1h",
|
||||
"hard_bounds": { "min": 3_600_000, "max": 428_400_000 }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
execute_agg(index, agg_req);
|
||||
}
|
||||
|
||||
/// Same fused terms × date_histogram, but with a sibling terms aggregation next to it. The fused
|
||||
/// fast path should still trigger for `my_texts` (sibling aggregations are independent top-level
|
||||
/// aggregations, so they don't change its eligibility).
|
||||
fn terms_status_with_date_histogram_and_sibling_terms(index: &Index) {
|
||||
let agg_req = json!({
|
||||
"my_texts": {
|
||||
"terms": { "field": "text_few_terms_status" },
|
||||
"aggs": {
|
||||
"over_time": { "date_histogram": { "field": "timestamp", "fixed_interval": "1h" } }
|
||||
}
|
||||
},
|
||||
"other_texts": { "terms": { "field": "text_few_terms" } }
|
||||
});
|
||||
execute_agg(index, agg_req);
|
||||
}
|
||||
|
||||
fn terms_zipf_1000_with_histogram(index: &Index) {
|
||||
let agg_req = json!({
|
||||
"my_texts": {
|
||||
@@ -783,7 +837,9 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result<Index> {
|
||||
doc_with_value /= 20;
|
||||
}
|
||||
let _val_max = 1_000_000.0;
|
||||
for _ in 0..doc_with_value {
|
||||
const SPAN_MS: i64 = 120 * 3600 * 1000; // 120 hours in ms
|
||||
const NOISE_MS: i64 = 2 * 3600 * 1000; // ±2h noise
|
||||
for i in 0..doc_with_value {
|
||||
let val: f64 = rng.random_range(0.0..1_000_000.0);
|
||||
let json = if rng.random_bool(0.1) {
|
||||
// 10% are numeric values
|
||||
@@ -791,6 +847,9 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result<Index> {
|
||||
} else {
|
||||
json!({"mixed_type": many_terms_data.choose(&mut rng).unwrap().to_string()})
|
||||
};
|
||||
let base_ms = (i as i64 * SPAN_MS) / doc_with_value as i64;
|
||||
let noise_ms = rng.random_range(-NOISE_MS..NOISE_MS);
|
||||
let ts_ms = (base_ms + noise_ms).clamp(0, SPAN_MS);
|
||||
index_writer.add_document(doc!(
|
||||
single_term => "single_term",
|
||||
text_field => "cool",
|
||||
@@ -803,7 +862,7 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result<Index> {
|
||||
score_field => val as u64,
|
||||
score_field_f64 => lg_norm.sample(&mut rng),
|
||||
score_field_i64 => val as i64,
|
||||
date_field => DateTime::from_timestamp_millis((val * 1_000_000.) as i64),
|
||||
date_field => DateTime::from_timestamp_millis(ts_ms),
|
||||
))?;
|
||||
if cardinality == Cardinality::OptionalSparse {
|
||||
for _ in 0..20 {
|
||||
|
||||
@@ -110,43 +110,31 @@ fn main() {
|
||||
// Prepare corpora with varying scenarios
|
||||
let scenarios = vec![
|
||||
(
|
||||
"dense and 99% a".to_string(),
|
||||
10_000_000,
|
||||
0.99,
|
||||
"dense and 0.1% a".to_string(),
|
||||
5_000_000,
|
||||
0.001,
|
||||
"dense",
|
||||
0,
|
||||
9,
|
||||
),
|
||||
("dense and 1% a".to_string(), 5_000_000, 0.01, "dense", 0, 9),
|
||||
("dense and 10% a".to_string(), 5_000_000, 0.1, "dense", 0, 9),
|
||||
(
|
||||
"dense and 99% a".to_string(),
|
||||
10_000_000,
|
||||
0.99,
|
||||
"dense",
|
||||
990,
|
||||
999,
|
||||
),
|
||||
(
|
||||
"sparse and 99% a".to_string(),
|
||||
10_000_000,
|
||||
"sparse and 50% a".to_string(),
|
||||
5_000_000,
|
||||
0.99,
|
||||
"sparse",
|
||||
0,
|
||||
9,
|
||||
),
|
||||
(
|
||||
"sparse and 99% a".to_string(),
|
||||
10_000_000,
|
||||
0.99,
|
||||
"sparse",
|
||||
9_999_990,
|
||||
9_999_999,
|
||||
),
|
||||
];
|
||||
|
||||
let mut runner = BenchRunner::new();
|
||||
for (scenario_id, n, p_title_a, num_rand_distribution, range_low, range_high) in scenarios {
|
||||
for (scenario_id, num_docs, p_title_a, num_rand_distribution, range_low, range_high) in
|
||||
scenarios
|
||||
{
|
||||
// Build index for this scenario
|
||||
let bench_index = build_shared_indices(n, p_title_a, num_rand_distribution);
|
||||
let bench_index = build_shared_indices(num_docs, p_title_a, num_rand_distribution);
|
||||
|
||||
// Create benchmark group
|
||||
let mut group = runner.new_group();
|
||||
@@ -158,7 +146,7 @@ fn main() {
|
||||
let field_names = ["num_rand", "num_asc", "num_rand_fast", "num_asc_fast"];
|
||||
|
||||
// Define the three terms we want to test with
|
||||
let terms = ["a", "b", "z"];
|
||||
let terms = ["a"];
|
||||
|
||||
// Generate all combinations of terms and field names
|
||||
let mut queries = Vec::new();
|
||||
@@ -203,7 +191,7 @@ fn run_benchmark_tasks(
|
||||
bench_index,
|
||||
query_str,
|
||||
DocSetCollector,
|
||||
"all results",
|
||||
"all_results",
|
||||
);
|
||||
|
||||
// Test top 100 by the field (if it's a FAST field)
|
||||
|
||||
@@ -15,9 +15,37 @@ impl<T: PartialOrd + Copy + std::fmt::Debug + Send + Sync + 'static + Default>
|
||||
{
|
||||
#[inline]
|
||||
pub fn fetch_block<'a>(&'a mut self, docs: &'a [u32], accessor: &Column<T>) {
|
||||
if accessor.index.get_cardinality().is_full() {
|
||||
self.val_cache.resize(docs.len(), T::default());
|
||||
accessor.values.get_vals(docs, &mut self.val_cache);
|
||||
self.fetch_block_with_is_full(docs, accessor, accessor.index.get_cardinality().is_full());
|
||||
}
|
||||
|
||||
/// Like [`Self::fetch_block`] but takes the column's fullness instead of querying
|
||||
/// `accessor.index.get_cardinality()` each call — for callers that know it up front (e.g.
|
||||
/// checked once at construction). `is_full` must equal
|
||||
/// `accessor.index.get_cardinality().is_full()`.
|
||||
#[inline]
|
||||
pub fn fetch_block_with_is_full<'a>(
|
||||
&'a mut self,
|
||||
docs: &'a [u32],
|
||||
accessor: &Column<T>,
|
||||
is_full: bool,
|
||||
) {
|
||||
if is_full {
|
||||
// Skip the resize when already the right length (common case: fixed-size blocks).
|
||||
if self.val_cache.len() != docs.len() {
|
||||
self.val_cache.resize(docs.len(), T::default());
|
||||
}
|
||||
// When the docs form a contiguous ascending run we can fetch the values
|
||||
// as a single range. This lets codecs (e.g. bitpacked) bulk-decode the
|
||||
// slice instead of gathering value-by-value, and avoids per-value dynamic
|
||||
// dispatch. `docs` is always sorted ascending and free of duplicates here,
|
||||
// so comparing the endpoints is enough to detect contiguity.
|
||||
if is_contiguous(docs) {
|
||||
accessor
|
||||
.values
|
||||
.get_range(docs[0] as u64, &mut self.val_cache);
|
||||
} else {
|
||||
accessor.values.get_vals(docs, &mut self.val_cache);
|
||||
}
|
||||
} else {
|
||||
self.docid_cache.clear();
|
||||
self.row_id_cache.clear();
|
||||
@@ -158,6 +186,22 @@ impl<T: PartialOrd + Copy + std::fmt::Debug + Send + Sync + 'static + Default>
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if `docs` is a contiguous ascending run `[d, d + 1, ..., d + n - 1]`.
|
||||
///
|
||||
/// Assumes `docs` is sorted ascending and free of duplicates (the invariant for the
|
||||
/// doc blocks passed to `fetch_block`), so comparing the endpoints is sufficient.
|
||||
#[inline]
|
||||
fn is_contiguous(docs: &[u32]) -> bool {
|
||||
let (Some(&first), Some(&last)) = (docs.first(), docs.last()) else {
|
||||
return false;
|
||||
};
|
||||
debug_assert!(
|
||||
docs.windows(2).all(|w| w[0] < w[1]),
|
||||
"fetch_block requires docs sorted ascending without duplicates"
|
||||
);
|
||||
(last - first) as usize + 1 == docs.len()
|
||||
}
|
||||
|
||||
/// Given two sorted lists of docids `docs` and `hits`, hits is a subset of `docs`.
|
||||
/// Return all docs that are not in `hits`.
|
||||
fn find_missing_docs<F>(docs: &[u32], hits: &[u32], mut callback: F)
|
||||
@@ -288,4 +332,46 @@ mod tests {
|
||||
assert_eq!(accessor.docid_cache, vec![0]);
|
||||
assert_eq!(accessor.val_cache, vec![1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_contiguous() {
|
||||
assert!(!is_contiguous(&[]));
|
||||
assert!(is_contiguous(&[5]));
|
||||
assert!(is_contiguous(&[5, 6, 7, 8]));
|
||||
assert!(is_contiguous(&[0, 1, 2]));
|
||||
assert!(!is_contiguous(&[5, 7, 8]));
|
||||
assert!(!is_contiguous(&[0, 1, 3]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fetch_block_contiguous_and_gather_match() {
|
||||
use crate::column_index::ColumnIndex;
|
||||
use crate::column_values::{
|
||||
ALL_U64_CODEC_TYPES, serialize_and_load_u64_based_column_values,
|
||||
};
|
||||
|
||||
let vals: Vec<u64> = (0..200u64).map(|i| i * 7 + 3).collect();
|
||||
let values =
|
||||
serialize_and_load_u64_based_column_values::<u64>(&&vals[..], &ALL_U64_CODEC_TYPES);
|
||||
let column = Column {
|
||||
index: ColumnIndex::Full,
|
||||
values,
|
||||
};
|
||||
|
||||
let check = |accessor: &mut ColumnBlockAccessor<u64>, docs: &[u32]| {
|
||||
accessor.fetch_block(docs, &column);
|
||||
let got: Vec<(u32, u64)> = accessor.iter_docid_vals(docs, &column).collect();
|
||||
let expected: Vec<(u32, u64)> = docs.iter().map(|&d| (d, vals[d as usize])).collect();
|
||||
assert_eq!(got, expected);
|
||||
};
|
||||
|
||||
let mut accessor = ColumnBlockAccessor::<u64>::default();
|
||||
// Contiguous block -> get_range fast path.
|
||||
check(&mut accessor, &(10..74).collect::<Vec<u32>>());
|
||||
// Non-contiguous block -> get_vals gather path.
|
||||
check(&mut accessor, &[0, 5, 9, 100, 199]);
|
||||
// Single doc and full span.
|
||||
check(&mut accessor, &[42]);
|
||||
check(&mut accessor, &(0..200).collect::<Vec<u32>>());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -119,8 +119,18 @@ pub trait ColumnValues<T: PartialOrd = u64>: Send + Sync + DowncastSync {
|
||||
/// the segment's `maxdoc`.
|
||||
#[inline(always)]
|
||||
fn get_range(&self, start: u64, output: &mut [T]) {
|
||||
for (out, idx) in output.iter_mut().zip(start..) {
|
||||
let mut out_chunks = output.chunks_exact_mut(4);
|
||||
let mut idx = start;
|
||||
for out_x4 in out_chunks.by_ref() {
|
||||
out_x4[0] = self.get_val(idx as u32);
|
||||
out_x4[1] = self.get_val((idx + 1) as u32);
|
||||
out_x4[2] = self.get_val((idx + 2) as u32);
|
||||
out_x4[3] = self.get_val((idx + 3) as u32);
|
||||
idx += 4;
|
||||
}
|
||||
for out in out_chunks.into_remainder() {
|
||||
*out = self.get_val(idx as u32);
|
||||
idx += 1;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -121,6 +121,22 @@ pub(crate) fn create_and_validate<TColumnCodec: ColumnCodec>(
|
||||
reader.get_vals(&all_docs, &mut buffer);
|
||||
assert_eq!(vals, buffer);
|
||||
|
||||
// Validate `get_range` over the full column and a sub-range. The sub-range starts
|
||||
// at a non-zero offset to exercise the entrance-ramp alignment of the batch decode.
|
||||
buffer.resize(all_docs.len(), 0);
|
||||
reader.get_range(0, &mut buffer);
|
||||
assert_eq!(vals, buffer, "get_range (full) mismatch in data set {name}");
|
||||
if vals.len() >= 2 {
|
||||
let start = 1usize;
|
||||
buffer.resize(vals.len() - start, 0);
|
||||
reader.get_range(start as u64, &mut buffer);
|
||||
assert_eq!(
|
||||
&vals[start..],
|
||||
&buffer[..],
|
||||
"get_range (sub-range) mismatch in data set {name}"
|
||||
);
|
||||
}
|
||||
|
||||
if !vals.is_empty() {
|
||||
let test_rand_idx = rand::rng().random_range(0..=vals.len() - 1);
|
||||
let expected_positions: Vec<u32> = vals
|
||||
|
||||
@@ -196,11 +196,13 @@ impl TinySet {
|
||||
#[derive(Clone)]
|
||||
pub struct BitSet {
|
||||
tinysets: Box<[TinySet]>,
|
||||
len: u64,
|
||||
max_value: u32,
|
||||
}
|
||||
impl std::fmt::Debug for BitSet {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("BitSet")
|
||||
.field("len", &self.len)
|
||||
.field("max_value", &self.max_value)
|
||||
.finish()
|
||||
}
|
||||
@@ -228,6 +230,7 @@ impl BitSet {
|
||||
let tinybitsets = vec![TinySet::empty(); num_buckets as usize].into_boxed_slice();
|
||||
BitSet {
|
||||
tinysets: tinybitsets,
|
||||
len: 0,
|
||||
max_value,
|
||||
}
|
||||
}
|
||||
@@ -245,6 +248,7 @@ impl BitSet {
|
||||
}
|
||||
BitSet {
|
||||
tinysets: tinybitsets,
|
||||
len: max_value as u64,
|
||||
max_value,
|
||||
}
|
||||
}
|
||||
@@ -263,19 +267,17 @@ impl BitSet {
|
||||
|
||||
/// Intersect with tinysets
|
||||
fn intersect_update_with_iter(&mut self, other: impl Iterator<Item = TinySet>) {
|
||||
self.len = 0;
|
||||
for (left, right) in self.tinysets.iter_mut().zip(other) {
|
||||
*left = left.intersect(right);
|
||||
self.len += left.len() as u64;
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the number of elements in the `BitSet`.
|
||||
#[inline]
|
||||
pub fn len(&self) -> usize {
|
||||
self.tinysets
|
||||
.iter()
|
||||
.copied()
|
||||
.map(|tinyset| tinyset.len())
|
||||
.sum::<u32>() as usize
|
||||
self.len as usize
|
||||
}
|
||||
|
||||
/// Inserts an element in the `BitSet`
|
||||
@@ -284,7 +286,7 @@ impl BitSet {
|
||||
// we do not check saturated els.
|
||||
let higher = el / 64u32;
|
||||
let lower = el % 64u32;
|
||||
self.tinysets[higher as usize].insert_mut(lower);
|
||||
self.len += u64::from(self.tinysets[higher as usize].insert_mut(lower));
|
||||
}
|
||||
|
||||
/// Inserts an element in the `BitSet`
|
||||
@@ -293,7 +295,7 @@ impl BitSet {
|
||||
// we do not check saturated els.
|
||||
let higher = el / 64u32;
|
||||
let lower = el % 64u32;
|
||||
self.tinysets[higher as usize].remove_mut(lower);
|
||||
self.len -= u64::from(self.tinysets[higher as usize].remove_mut(lower));
|
||||
}
|
||||
|
||||
/// Returns true iff the elements is in the `BitSet`.
|
||||
@@ -315,9 +317,6 @@ impl BitSet {
|
||||
.map(|delta_bucket| bucket + delta_bucket as u32)
|
||||
}
|
||||
|
||||
/// Returns the maximum number of elements in the bitset.
|
||||
///
|
||||
/// Warning: The largest element the bitset can contain is `max_value - 1`.
|
||||
#[inline]
|
||||
pub fn max_value(&self) -> u32 {
|
||||
self.max_value
|
||||
|
||||
@@ -91,10 +91,46 @@ fn main() -> tantivy::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// Some other powerful operations (especially `.seek`) may be useful to consume these
|
||||
// A `Term` is a text token associated with a field.
|
||||
// Let's go through all docs containing the term `title:the` and access their position
|
||||
let term_the = Term::from_field_text(title, "the");
|
||||
|
||||
// Some other powerful operations (especially `.skip_to`) may be useful to consume these
|
||||
// posting lists rapidly.
|
||||
// You can check for them in the [`DocSet`](https://docs.rs/tantivy/~0/tantivy/trait.DocSet.html) trait
|
||||
// and the [`Postings`](https://docs.rs/tantivy/~0/tantivy/trait.Postings.html) trait
|
||||
|
||||
// Also, for some VERY specific high performance use case like an OLAP analysis of logs,
|
||||
// you can get better performance by accessing directly the blocks of doc ids.
|
||||
for segment_reader in searcher.segment_readers() {
|
||||
// A segment contains different data structure.
|
||||
// Inverted index stands for the combination of
|
||||
// - the term dictionary
|
||||
// - the inverted lists associated with each terms and their positions
|
||||
let inverted_index = segment_reader.inverted_index(title)?;
|
||||
|
||||
// This segment posting object is like a cursor over the documents matching the term.
|
||||
// The `IndexRecordOption` arguments tells tantivy we will be interested in both term
|
||||
// frequencies and positions.
|
||||
//
|
||||
// If you don't need all this information, you may get better performance by decompressing
|
||||
// less information.
|
||||
if let Some(mut block_segment_postings) =
|
||||
inverted_index.read_block_postings(&term_the, IndexRecordOption::Basic)?
|
||||
{
|
||||
loop {
|
||||
let docs = block_segment_postings.docs();
|
||||
if docs.is_empty() {
|
||||
break;
|
||||
}
|
||||
// Once again these docs MAY contains deleted documents as well.
|
||||
let docs = block_segment_postings.docs();
|
||||
// Prints `Docs [0, 2].`
|
||||
println!("Docs {docs:?}");
|
||||
block_segment_postings.advance();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -327,7 +327,9 @@ fn exists(inp: &str) -> IResult<&str, UserInputLeaf> {
|
||||
peek(alt((
|
||||
value(
|
||||
"",
|
||||
satisfy(|c: char| c.is_whitespace() || ESCAPE_IN_WORD.contains(&c)),
|
||||
satisfy(|c: char| {
|
||||
c.is_whitespace() || (ESCAPE_IN_WORD.contains(&c) && c != '\\')
|
||||
}),
|
||||
),
|
||||
eof,
|
||||
))),
|
||||
@@ -345,7 +347,9 @@ fn exists_precond(inp: &str) -> IResult<&str, (), ()> {
|
||||
peek(alt((
|
||||
value(
|
||||
"",
|
||||
satisfy(|c: char| c.is_whitespace() || ESCAPE_IN_WORD.contains(&c)),
|
||||
satisfy(|c: char| {
|
||||
c.is_whitespace() || (ESCAPE_IN_WORD.contains(&c) && c != '\\')
|
||||
}),
|
||||
),
|
||||
eof,
|
||||
))), // we need to check this isn't a wildcard query
|
||||
@@ -707,6 +711,7 @@ fn regex(inp: &str) -> IResult<&str, UserInputLeaf> {
|
||||
peek(alt((
|
||||
value((), multispace1),
|
||||
value((), char(')')),
|
||||
value((), char('^')),
|
||||
value((), eof),
|
||||
))),
|
||||
),
|
||||
@@ -728,9 +733,10 @@ fn regex_infallible(inp: &str) -> JResult<&str, UserInputLeaf> {
|
||||
peek(alt((
|
||||
value((), multispace1),
|
||||
value((), char(')')),
|
||||
value((), char('^')),
|
||||
value((), eof),
|
||||
))),
|
||||
"expected whitespace, closing parenthesis, or end of input",
|
||||
"expected whitespace, closing parenthesis, boost, or end of input",
|
||||
),
|
||||
)(inp)
|
||||
{
|
||||
@@ -773,6 +779,10 @@ fn leaf(inp: &str) -> IResult<&str, UserInputAst> {
|
||||
value((), multispace1),
|
||||
value((), char(')')),
|
||||
value((), eof),
|
||||
value(
|
||||
(),
|
||||
satisfy(|c: char| ESCAPE_IN_WORD.contains(&c) && c != '\\'),
|
||||
),
|
||||
))),
|
||||
),
|
||||
|_| UserInputAst::from(UserInputLeaf::All),
|
||||
@@ -805,6 +815,10 @@ fn leaf_infallible(inp: &str) -> JResult<&str, Option<UserInputAst>> {
|
||||
value((), multispace1),
|
||||
value((), char(')')),
|
||||
value((), eof),
|
||||
value(
|
||||
(),
|
||||
satisfy(|c: char| ESCAPE_IN_WORD.contains(&c) && c != '\\'),
|
||||
),
|
||||
))),
|
||||
),
|
||||
),
|
||||
@@ -1751,6 +1765,8 @@ mod test {
|
||||
test_parse_query_to_ast_helper("*", "*");
|
||||
test_parse_query_to_ast_helper("(*)", "*");
|
||||
test_parse_query_to_ast_helper("(* )", "*");
|
||||
// All query with boost
|
||||
test_parse_query_to_ast_helper("*^2", "(*)^2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1813,6 +1829,7 @@ mod test {
|
||||
test_parse_query_to_ast_helper("a:b*", "\"a\":b*");
|
||||
test_parse_query_to_ast_helper("a:*b", "\"a\":*b");
|
||||
test_parse_query_to_ast_helper(r#"a:*def*"#, "\"a\":*def*");
|
||||
test_parse_query_to_ast_helper("a:*\\:foo", "\"a\":*:foo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1877,6 +1894,8 @@ mod test {
|
||||
},
|
||||
_ => panic!("Expected a leaf"),
|
||||
}
|
||||
// Regex followed by `^boost`
|
||||
test_parse_query_to_ast_helper(r#"foo:/bar/^2"#, r#"("foo":/bar/)^2"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -10,11 +10,11 @@ use crate::aggregation::accessor_helpers::{
|
||||
};
|
||||
use crate::aggregation::agg_req::{Aggregation, AggregationVariants, Aggregations};
|
||||
use crate::aggregation::bucket::{
|
||||
build_segment_filter_collector, build_segment_range_collector, CompositeAggReqData,
|
||||
CompositeAggregation, CompositeSourceAccessors, FilterAggReqData, HistogramAggReqData,
|
||||
HistogramBounds, IncludeExcludeParam, MissingTermAggReqData, RangeAggReqData,
|
||||
SegmentHistogramCollector, TermMissingAgg, TermsAggReqData, TermsAggregation,
|
||||
TermsAggregationInternal,
|
||||
build_segment_filter_collector, build_segment_histogram_collector,
|
||||
build_segment_range_collector, CompositeAggReqData, CompositeAggregation,
|
||||
CompositeSourceAccessors, FilterAggReqData, HistogramAggReqData, HistogramBounds,
|
||||
IncludeExcludeParam, MissingTermAggReqData, RangeAggReqData, TermMissingAgg, TermsAggReqData,
|
||||
TermsAggregation, TermsAggregationInternal,
|
||||
};
|
||||
use crate::aggregation::metric::{
|
||||
build_segment_stats_collector, AverageAggregation, CardinalityAggReqData,
|
||||
@@ -41,7 +41,7 @@ pub struct AggregationsSegmentCtx {
|
||||
|
||||
impl AggregationsSegmentCtx {
|
||||
pub(crate) fn push_term_req_data(&mut self, data: TermsAggReqData) -> usize {
|
||||
self.per_request.term_req_data.push(Some(Box::new(data)));
|
||||
self.per_request.term_req_data.push(data);
|
||||
self.per_request.term_req_data.len() - 1
|
||||
}
|
||||
pub(crate) fn push_cardinality_req_data(&mut self, data: CardinalityAggReqData) -> usize {
|
||||
@@ -61,31 +61,25 @@ impl AggregationsSegmentCtx {
|
||||
self.per_request.missing_term_req_data.len() - 1
|
||||
}
|
||||
pub(crate) fn push_histogram_req_data(&mut self, data: HistogramAggReqData) -> usize {
|
||||
self.per_request
|
||||
.histogram_req_data
|
||||
.push(Some(Box::new(data)));
|
||||
self.per_request.histogram_req_data.push(data);
|
||||
self.per_request.histogram_req_data.len() - 1
|
||||
}
|
||||
pub(crate) fn push_range_req_data(&mut self, data: RangeAggReqData) -> usize {
|
||||
self.per_request.range_req_data.push(Some(Box::new(data)));
|
||||
self.per_request.range_req_data.push(data);
|
||||
self.per_request.range_req_data.len() - 1
|
||||
}
|
||||
pub(crate) fn push_filter_req_data(&mut self, data: FilterAggReqData) -> usize {
|
||||
self.per_request.filter_req_data.push(Some(Box::new(data)));
|
||||
self.per_request.filter_req_data.push(data);
|
||||
self.per_request.filter_req_data.len() - 1
|
||||
}
|
||||
pub(crate) fn push_composite_req_data(&mut self, data: CompositeAggReqData) -> usize {
|
||||
self.per_request
|
||||
.composite_req_data
|
||||
.push(Some(Box::new(data)));
|
||||
self.per_request.composite_req_data.push(data);
|
||||
self.per_request.composite_req_data.len() - 1
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn get_term_req_data(&self, idx: usize) -> &TermsAggReqData {
|
||||
self.per_request.term_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("term_req_data slot is empty (taken)")
|
||||
&self.per_request.term_req_data[idx]
|
||||
}
|
||||
#[inline]
|
||||
pub(crate) fn get_cardinality_req_data(&self, idx: usize) -> &CardinalityAggReqData {
|
||||
@@ -103,116 +97,6 @@ impl AggregationsSegmentCtx {
|
||||
pub(crate) fn get_missing_term_req_data(&self, idx: usize) -> &MissingTermAggReqData {
|
||||
&self.per_request.missing_term_req_data[idx]
|
||||
}
|
||||
#[inline]
|
||||
pub(crate) fn get_histogram_req_data(&self, idx: usize) -> &HistogramAggReqData {
|
||||
self.per_request.histogram_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("histogram_req_data slot is empty (taken)")
|
||||
}
|
||||
#[inline]
|
||||
pub(crate) fn get_range_req_data(&self, idx: usize) -> &RangeAggReqData {
|
||||
self.per_request.range_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("range_req_data slot is empty (taken)")
|
||||
}
|
||||
#[inline]
|
||||
pub(crate) fn get_composite_req_data(&self, idx: usize) -> &CompositeAggReqData {
|
||||
self.per_request.composite_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("composite_req_data slot is empty (taken)")
|
||||
}
|
||||
|
||||
// ---------- mutable getters ----------
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn get_metric_req_data_mut(&mut self, idx: usize) -> &mut MetricAggReqData {
|
||||
&mut self.per_request.stats_metric_req_data[idx]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn get_cardinality_req_data_mut(
|
||||
&mut self,
|
||||
idx: usize,
|
||||
) -> &mut CardinalityAggReqData {
|
||||
&mut self.per_request.cardinality_req_data[idx]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn get_histogram_req_data_mut(&mut self, idx: usize) -> &mut HistogramAggReqData {
|
||||
self.per_request.histogram_req_data[idx]
|
||||
.as_deref_mut()
|
||||
.expect("histogram_req_data slot is empty (taken)")
|
||||
}
|
||||
|
||||
// ---------- take / put (terms, histogram, range) ----------
|
||||
|
||||
/// Move out the boxed Histogram request at `idx`, leaving `None`.
|
||||
#[inline]
|
||||
pub(crate) fn take_histogram_req_data(&mut self, idx: usize) -> Box<HistogramAggReqData> {
|
||||
self.per_request.histogram_req_data[idx]
|
||||
.take()
|
||||
.expect("histogram_req_data slot is empty (taken)")
|
||||
}
|
||||
|
||||
/// Put back a Histogram request into an empty slot at `idx`.
|
||||
#[inline]
|
||||
pub(crate) fn put_back_histogram_req_data(
|
||||
&mut self,
|
||||
idx: usize,
|
||||
value: Box<HistogramAggReqData>,
|
||||
) {
|
||||
debug_assert!(self.per_request.histogram_req_data[idx].is_none());
|
||||
self.per_request.histogram_req_data[idx] = Some(value);
|
||||
}
|
||||
|
||||
/// Move out the boxed Range request at `idx`, leaving `None`.
|
||||
#[inline]
|
||||
pub(crate) fn take_range_req_data(&mut self, idx: usize) -> Box<RangeAggReqData> {
|
||||
self.per_request.range_req_data[idx]
|
||||
.take()
|
||||
.expect("range_req_data slot is empty (taken)")
|
||||
}
|
||||
|
||||
/// Put back a Range request into an empty slot at `idx`.
|
||||
#[inline]
|
||||
pub(crate) fn put_back_range_req_data(&mut self, idx: usize, value: Box<RangeAggReqData>) {
|
||||
debug_assert!(self.per_request.range_req_data[idx].is_none());
|
||||
self.per_request.range_req_data[idx] = Some(value);
|
||||
}
|
||||
|
||||
/// Move out the boxed Filter request at `idx`, leaving `None`.
|
||||
#[inline]
|
||||
pub(crate) fn take_filter_req_data(&mut self, idx: usize) -> Box<FilterAggReqData> {
|
||||
self.per_request.filter_req_data[idx]
|
||||
.take()
|
||||
.expect("filter_req_data slot is empty (taken)")
|
||||
}
|
||||
|
||||
/// Put back a Filter request into an empty slot at `idx`.
|
||||
#[inline]
|
||||
pub(crate) fn put_back_filter_req_data(&mut self, idx: usize, value: Box<FilterAggReqData>) {
|
||||
debug_assert!(self.per_request.filter_req_data[idx].is_none());
|
||||
self.per_request.filter_req_data[idx] = Some(value);
|
||||
}
|
||||
|
||||
/// Move out the Composite request at `idx`.
|
||||
#[inline]
|
||||
pub(crate) fn take_composite_req_data(&mut self, idx: usize) -> Box<CompositeAggReqData> {
|
||||
self.per_request.composite_req_data[idx]
|
||||
.take()
|
||||
.expect("composite_req_data slot is empty (taken)")
|
||||
}
|
||||
|
||||
/// Put back a Composite request into an empty slot at `idx`.
|
||||
#[inline]
|
||||
pub(crate) fn put_back_composite_req_data(
|
||||
&mut self,
|
||||
idx: usize,
|
||||
value: Box<CompositeAggReqData>,
|
||||
) {
|
||||
debug_assert!(self.per_request.composite_req_data[idx].is_none());
|
||||
self.per_request.composite_req_data[idx] = Some(value);
|
||||
}
|
||||
}
|
||||
|
||||
/// Each type of aggregation has its own request data struct. This struct holds
|
||||
@@ -223,15 +107,14 @@ impl AggregationsSegmentCtx {
|
||||
/// for a node with [AggKind::Terms]).
|
||||
#[derive(Default)]
|
||||
pub struct PerRequestAggSegCtx {
|
||||
// Box for cheap take/put - Only necessary for bucket aggs that have sub-aggregations
|
||||
/// TermsAggReqData contains the request data for a terms aggregation.
|
||||
pub term_req_data: Vec<Option<Box<TermsAggReqData>>>,
|
||||
pub term_req_data: Vec<TermsAggReqData>,
|
||||
/// HistogramAggReqData contains the request data for a histogram aggregation.
|
||||
pub histogram_req_data: Vec<Option<Box<HistogramAggReqData>>>,
|
||||
pub histogram_req_data: Vec<HistogramAggReqData>,
|
||||
/// RangeAggReqData contains the request data for a range aggregation.
|
||||
pub range_req_data: Vec<Option<Box<RangeAggReqData>>>,
|
||||
pub range_req_data: Vec<RangeAggReqData>,
|
||||
/// FilterAggReqData contains the request data for a filter aggregation.
|
||||
pub filter_req_data: Vec<Option<Box<FilterAggReqData>>>,
|
||||
pub filter_req_data: Vec<FilterAggReqData>,
|
||||
/// Shared by avg, min, max, sum, stats, extended_stats, count
|
||||
pub stats_metric_req_data: Vec<MetricAggReqData>,
|
||||
/// CardinalityAggReqData contains the request data for a cardinality aggregation.
|
||||
@@ -241,7 +124,7 @@ pub struct PerRequestAggSegCtx {
|
||||
/// MissingTermAggReqData contains the request data for a missing term aggregation.
|
||||
pub missing_term_req_data: Vec<MissingTermAggReqData>,
|
||||
/// CompositeAggReqData contains the request data for a composite aggregation.
|
||||
pub composite_req_data: Vec<Option<Box<CompositeAggReqData>>>,
|
||||
pub composite_req_data: Vec<CompositeAggReqData>,
|
||||
|
||||
/// Request tree used to build collectors.
|
||||
pub agg_tree: Vec<AggRefNode>,
|
||||
@@ -252,22 +135,22 @@ impl PerRequestAggSegCtx {
|
||||
fn get_memory_consumption(&self) -> usize {
|
||||
self.term_req_data
|
||||
.iter()
|
||||
.map(|b| b.as_ref().unwrap().get_memory_consumption())
|
||||
.map(|t| t.get_memory_consumption())
|
||||
.sum::<usize>()
|
||||
+ self
|
||||
.histogram_req_data
|
||||
.iter()
|
||||
.map(|b| b.as_ref().unwrap().get_memory_consumption())
|
||||
.map(|t| t.get_memory_consumption())
|
||||
.sum::<usize>()
|
||||
+ self
|
||||
.range_req_data
|
||||
.iter()
|
||||
.map(|b| b.as_ref().unwrap().get_memory_consumption())
|
||||
.map(|t| t.get_memory_consumption())
|
||||
.sum::<usize>()
|
||||
+ self
|
||||
.filter_req_data
|
||||
.iter()
|
||||
.map(|b| b.as_ref().unwrap().get_memory_consumption())
|
||||
.map(|t| t.get_memory_consumption())
|
||||
.sum::<usize>()
|
||||
+ self
|
||||
.stats_metric_req_data
|
||||
@@ -292,7 +175,7 @@ impl PerRequestAggSegCtx {
|
||||
+ self
|
||||
.composite_req_data
|
||||
.iter()
|
||||
.map(|b| b.as_ref().map(|d| d.get_memory_consumption()).unwrap_or(0))
|
||||
.map(|t| t.get_memory_consumption())
|
||||
.sum::<usize>()
|
||||
+ self.agg_tree.len() * std::mem::size_of::<AggRefNode>()
|
||||
}
|
||||
@@ -301,40 +184,16 @@ impl PerRequestAggSegCtx {
|
||||
let idx = node.idx_in_req_data;
|
||||
let kind = node.kind;
|
||||
match kind {
|
||||
AggKind::Terms => self.term_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("term_req_data slot is empty (taken)")
|
||||
.name
|
||||
.as_str(),
|
||||
AggKind::Terms => self.term_req_data[idx].name.as_str(),
|
||||
AggKind::Cardinality => &self.cardinality_req_data[idx].name,
|
||||
AggKind::StatsKind(_) => &self.stats_metric_req_data[idx].name,
|
||||
AggKind::TopHits => &self.top_hits_req_data[idx].name,
|
||||
AggKind::MissingTerm => &self.missing_term_req_data[idx].name,
|
||||
AggKind::Histogram => self.histogram_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("histogram_req_data slot is empty (taken)")
|
||||
.name
|
||||
.as_str(),
|
||||
AggKind::DateHistogram => self.histogram_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("histogram_req_data slot is empty (taken)")
|
||||
.name
|
||||
.as_str(),
|
||||
AggKind::Range => self.range_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("range_req_data slot is empty (taken)")
|
||||
.name
|
||||
.as_str(),
|
||||
AggKind::Filter => self.filter_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("filter_req_data slot is empty (taken)")
|
||||
.name
|
||||
.as_str(),
|
||||
AggKind::Composite => self.composite_req_data[idx]
|
||||
.as_deref()
|
||||
.expect("composite_req_data slot is empty (taken)")
|
||||
.name
|
||||
.as_str(),
|
||||
AggKind::Histogram => self.histogram_req_data[idx].name.as_str(),
|
||||
AggKind::DateHistogram => self.histogram_req_data[idx].name.as_str(),
|
||||
AggKind::Range => self.range_req_data[idx].name.as_str(),
|
||||
AggKind::Filter => self.filter_req_data[idx].name.as_str(),
|
||||
AggKind::Composite => self.composite_req_data[idx].name.as_str(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -412,7 +271,7 @@ pub(crate) fn build_segment_agg_collector(
|
||||
Ok(Box::new(TermMissingAgg::new(req, node)?))
|
||||
}
|
||||
AggKind::Cardinality => {
|
||||
let req_data = &mut req.get_cardinality_req_data_mut(node.idx_in_req_data);
|
||||
let req_data = req.get_cardinality_req_data(node.idx_in_req_data);
|
||||
// For str columns, choose the per-bucket entries representation
|
||||
// based on the segment's column.max_value():
|
||||
// * small (< BITSET_MAX_TERM_ORD): `BitSet`, pre-allocated, no promotion machinery.
|
||||
@@ -459,7 +318,7 @@ pub(crate) fn build_segment_agg_collector(
|
||||
SegmentExtendedStatsCollector::from_req(req_data, sigma),
|
||||
)),
|
||||
StatsType::Percentiles => {
|
||||
let req_data = req.get_metric_req_data_mut(node.idx_in_req_data);
|
||||
let req_data = req.get_metric_req_data(node.idx_in_req_data);
|
||||
Ok(Box::new(
|
||||
SegmentPercentilesCollector::from_req_and_validate(
|
||||
req_data.field_type,
|
||||
@@ -479,12 +338,8 @@ pub(crate) fn build_segment_agg_collector(
|
||||
req_data.segment_ordinal,
|
||||
)))
|
||||
}
|
||||
AggKind::Histogram => Ok(Box::new(SegmentHistogramCollector::from_req_and_validate(
|
||||
req, node,
|
||||
)?)),
|
||||
AggKind::DateHistogram => Ok(Box::new(SegmentHistogramCollector::from_req_and_validate(
|
||||
req, node,
|
||||
)?)),
|
||||
AggKind::Histogram => build_segment_histogram_collector(req, node),
|
||||
AggKind::DateHistogram => build_segment_histogram_collector(req, node),
|
||||
AggKind::Range => Ok(build_segment_range_collector(req, node)?),
|
||||
AggKind::Filter => build_segment_filter_collector(req, node),
|
||||
AggKind::Composite => Ok(Box::new(
|
||||
@@ -799,23 +654,18 @@ fn build_nodes(
|
||||
let schema = reader.schema();
|
||||
let tokenizers = &data.context.tokenizers;
|
||||
let query = filter_req.parse_query(schema, tokenizers)?;
|
||||
let evaluator = crate::aggregation::bucket::DocumentQueryEvaluator::new(
|
||||
query,
|
||||
schema.clone(),
|
||||
reader,
|
||||
)?;
|
||||
|
||||
// Pre-allocate buffer for batch filtering
|
||||
let max_doc = reader.max_doc();
|
||||
let buffer_capacity = crate::docset::COLLECT_BLOCK_BUFFER_LEN.min(max_doc as usize);
|
||||
let matching_docs_buffer = Vec::with_capacity(buffer_capacity);
|
||||
let evaluator =
|
||||
std::rc::Rc::new(crate::aggregation::bucket::DocumentQueryEvaluator::new(
|
||||
query,
|
||||
schema.clone(),
|
||||
reader,
|
||||
)?);
|
||||
|
||||
let idx_in_req_data = data.push_filter_req_data(FilterAggReqData {
|
||||
name: agg_name.to_string(),
|
||||
req: filter_req.clone(),
|
||||
segment_reader: reader.clone(),
|
||||
evaluator,
|
||||
matching_docs_buffer,
|
||||
is_top_level,
|
||||
});
|
||||
let children = build_children(&req.sub_aggregation, reader, segment_ordinal, data)?;
|
||||
|
||||
@@ -299,6 +299,12 @@ impl AggregationVariants {
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
pub(crate) fn as_sum(&self) -> Option<&SumAggregation> {
|
||||
match &self {
|
||||
AggregationVariants::Sum(sum) => Some(sum),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -16,6 +16,7 @@ use crate::{SegmentReader, TantivyError};
|
||||
|
||||
/// Contains all information required by the SegmentCompositeCollector to perform the
|
||||
/// composite aggregation on a segment.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompositeAggReqData {
|
||||
/// The name of the aggregation.
|
||||
pub name: String,
|
||||
@@ -34,6 +35,7 @@ impl CompositeAggReqData {
|
||||
}
|
||||
|
||||
/// Accessors for a single column in a composite source.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompositeAccessor {
|
||||
/// The fast field column
|
||||
pub column: Column<u64>,
|
||||
@@ -48,6 +50,7 @@ pub struct CompositeAccessor {
|
||||
}
|
||||
|
||||
/// Accessors to all the columns that belong to the field of a composite source.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompositeSourceAccessors {
|
||||
/// The accessors for this source
|
||||
pub accessors: Vec<CompositeAccessor>,
|
||||
@@ -358,7 +361,7 @@ impl PrecomputedDateInterval {
|
||||
///
|
||||
/// Some column types (term, IP) might not have an exact representation of the
|
||||
/// specified after key
|
||||
#[derive(Debug)]
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PrecomputedAfterKey {
|
||||
/// The after key could be exactly represented in the column space.
|
||||
Exact(u64),
|
||||
|
||||
@@ -118,7 +118,7 @@ impl InternalValueRepr {
|
||||
pub struct SegmentCompositeCollector {
|
||||
/// One DynArrayHeapMap per parent bucket.
|
||||
parent_buckets: Vec<DynArrayHeapMap<InternalValueRepr, CompositeBucketCollector>>,
|
||||
accessor_idx: usize,
|
||||
req_data: CompositeAggReqData,
|
||||
sub_agg: Option<BufferedSubAggs<HighCardSubAggBuffer>>,
|
||||
bucket_id_provider: BucketIdProvider,
|
||||
/// Number of sources, needed when creating new DynArrayHeapMaps.
|
||||
@@ -132,10 +132,7 @@ impl SegmentAggregationCollector for SegmentCompositeCollector {
|
||||
results: &mut IntermediateAggregationResults,
|
||||
parent_bucket_id: BucketId,
|
||||
) -> crate::Result<()> {
|
||||
let name = agg_data
|
||||
.get_composite_req_data(self.accessor_idx)
|
||||
.name
|
||||
.clone();
|
||||
let name = self.req_data.name.clone();
|
||||
|
||||
let buckets = self.add_intermediate_bucket_result(agg_data, parent_bucket_id)?;
|
||||
results.push(
|
||||
@@ -153,12 +150,11 @@ impl SegmentAggregationCollector for SegmentCompositeCollector {
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
let mem_pre = self.get_memory_consumption(parent_bucket_id);
|
||||
let composite_agg_data = agg_data.take_composite_req_data(self.accessor_idx);
|
||||
|
||||
for doc in docs {
|
||||
let mut visitor = CompositeKeyVisitor {
|
||||
doc_id: *doc,
|
||||
composite_agg_data: &composite_agg_data,
|
||||
composite_agg_data: &self.req_data,
|
||||
buckets: &mut self.parent_buckets[parent_bucket_id as usize],
|
||||
sub_agg: &mut self.sub_agg,
|
||||
bucket_id_provider: &mut self.bucket_id_provider,
|
||||
@@ -166,7 +162,6 @@ impl SegmentAggregationCollector for SegmentCompositeCollector {
|
||||
};
|
||||
visitor.visit(0, true)?;
|
||||
}
|
||||
agg_data.put_back_composite_req_data(self.accessor_idx, composite_agg_data);
|
||||
|
||||
if let Some(sub_agg) = &mut self.sub_agg {
|
||||
sub_agg.check_flush_local(agg_data)?;
|
||||
@@ -221,7 +216,13 @@ impl SegmentCompositeCollector {
|
||||
req_data: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
) -> crate::Result<Self> {
|
||||
validate_req(req_data, node.idx_in_req_data)?;
|
||||
let composite_req_data =
|
||||
req_data.per_request.composite_req_data[node.idx_in_req_data].clone();
|
||||
validate_req(&composite_req_data)?;
|
||||
req_data
|
||||
.context
|
||||
.limits
|
||||
.add_memory_consumed(composite_req_data.get_memory_consumption() as u64)?;
|
||||
|
||||
let has_sub_aggregations = !node.children.is_empty();
|
||||
let sub_agg = if has_sub_aggregations {
|
||||
@@ -231,12 +232,11 @@ impl SegmentCompositeCollector {
|
||||
None
|
||||
};
|
||||
|
||||
let composite_req_data = req_data.get_composite_req_data(node.idx_in_req_data);
|
||||
let num_sources = composite_req_data.req.sources.len();
|
||||
|
||||
Ok(SegmentCompositeCollector {
|
||||
parent_buckets: vec![DynArrayHeapMap::try_new(num_sources)?],
|
||||
accessor_idx: node.idx_in_req_data,
|
||||
req_data: composite_req_data,
|
||||
sub_agg,
|
||||
bucket_id_provider: BucketIdProvider::default(),
|
||||
num_sources,
|
||||
@@ -258,7 +258,7 @@ impl SegmentCompositeCollector {
|
||||
let mut dict: FxHashMap<Vec<CompositeIntermediateKey>, IntermediateCompositeBucketEntry> =
|
||||
Default::default();
|
||||
dict.reserve(heap_map.size());
|
||||
let composite_data = agg_data.get_composite_req_data(self.accessor_idx);
|
||||
let composite_data = &self.req_data;
|
||||
for (key_internal_repr, agg) in heap_map.into_iter() {
|
||||
let key = resolve_key(&key_internal_repr, composite_data)?;
|
||||
let mut sub_aggregation_res = IntermediateAggregationResults::default();
|
||||
@@ -298,8 +298,7 @@ impl SegmentCompositeCollector {
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_req(req_data: &mut AggregationsSegmentCtx, accessor_idx: usize) -> crate::Result<()> {
|
||||
let composite_data = req_data.get_composite_req_data(accessor_idx);
|
||||
fn validate_req(composite_data: &CompositeAggReqData) -> crate::Result<()> {
|
||||
let req = &composite_data.req;
|
||||
if req.sources.is_empty() {
|
||||
return Err(TantivyError::InvalidArgument(
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::fmt::Debug;
|
||||
use std::rc::Rc;
|
||||
|
||||
use common::BitSet;
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
@@ -396,6 +397,7 @@ impl PartialEq for FilterAggregation {
|
||||
|
||||
/// Request data for filter aggregation
|
||||
/// This struct holds the per-segment data needed to execute a filter aggregation
|
||||
#[derive(Clone)]
|
||||
pub struct FilterAggReqData {
|
||||
/// The name of the filter aggregation
|
||||
pub name: String,
|
||||
@@ -403,22 +405,20 @@ pub struct FilterAggReqData {
|
||||
pub req: FilterAggregation,
|
||||
/// The segment reader
|
||||
pub segment_reader: SegmentReader,
|
||||
/// Document evaluator for the filter query (precomputed BitSet)
|
||||
/// This is built once when the request data is created
|
||||
pub evaluator: DocumentQueryEvaluator,
|
||||
/// Reusable buffer for matching documents to minimize allocations during collection
|
||||
pub matching_docs_buffer: Vec<DocId>,
|
||||
/// Document evaluator for the filter query (precomputed BitSet).
|
||||
/// Wrapped in `Rc` so cloning the request data does not duplicate the (potentially large)
|
||||
/// underlying BitSet.
|
||||
pub evaluator: Rc<DocumentQueryEvaluator>,
|
||||
/// True if this filter aggregation is at the top level of the aggregation tree (not nested).
|
||||
pub is_top_level: bool,
|
||||
}
|
||||
|
||||
impl FilterAggReqData {
|
||||
pub(crate) fn get_memory_consumption(&self) -> usize {
|
||||
// Estimate: name + segment reader reference + bitset + buffer capacity
|
||||
// Estimate: name + segment reader reference + bitset
|
||||
self.name.len()
|
||||
+ std::mem::size_of::<SegmentReader>()
|
||||
+ self.evaluator.bitset.len() / 8 // BitSet memory (bits to bytes)
|
||||
+ self.matching_docs_buffer.capacity() * std::mem::size_of::<DocId>()
|
||||
+ std::mem::size_of::<bool>()
|
||||
}
|
||||
}
|
||||
@@ -509,8 +509,10 @@ pub struct SegmentFilterCollector<B: SubAggBuffer> {
|
||||
/// Sub-aggregation collectors
|
||||
sub_aggregations: Option<BufferedSubAggs<B>>,
|
||||
bucket_id_provider: BucketIdProvider,
|
||||
/// Accessor index for this filter aggregation (to access FilterAggReqData)
|
||||
accessor_idx: usize,
|
||||
/// Per-segment filter request data, owned by this collector.
|
||||
req_data: FilterAggReqData,
|
||||
/// Reusable buffer for matching documents to minimize allocations during collection.
|
||||
matching_docs_buffer: Vec<DocId>,
|
||||
}
|
||||
|
||||
impl<B: SubAggBuffer> SegmentFilterCollector<B> {
|
||||
@@ -518,6 +520,7 @@ impl<B: SubAggBuffer> SegmentFilterCollector<B> {
|
||||
pub(crate) fn from_req_and_validate(
|
||||
req: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
req_data: FilterAggReqData,
|
||||
) -> crate::Result<Self> {
|
||||
// Build sub-aggregation collectors if any
|
||||
let sub_agg_collector = if !node.children.is_empty() {
|
||||
@@ -527,11 +530,15 @@ impl<B: SubAggBuffer> SegmentFilterCollector<B> {
|
||||
};
|
||||
let sub_agg_collector = sub_agg_collector.map(BufferedSubAggs::new);
|
||||
|
||||
let max_doc = req_data.segment_reader.max_doc();
|
||||
let buffer_capacity = crate::docset::COLLECT_BLOCK_BUFFER_LEN.min(max_doc as usize);
|
||||
|
||||
Ok(SegmentFilterCollector {
|
||||
parent_buckets: Vec::new(),
|
||||
sub_aggregations: sub_agg_collector,
|
||||
accessor_idx: node.idx_in_req_data,
|
||||
req_data,
|
||||
bucket_id_provider: BucketIdProvider::default(),
|
||||
matching_docs_buffer: Vec::with_capacity(buffer_capacity),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -540,18 +547,23 @@ pub(crate) fn build_segment_filter_collector(
|
||||
req: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
|
||||
let is_top_level = req.per_request.filter_req_data[node.idx_in_req_data]
|
||||
.as_ref()
|
||||
.expect("filter_req_data slot is empty")
|
||||
.is_top_level;
|
||||
let req_data = req.per_request.filter_req_data[node.idx_in_req_data].clone();
|
||||
req.context
|
||||
.limits
|
||||
.add_memory_consumed(req_data.get_memory_consumption() as u64)?;
|
||||
let is_top_level = req_data.is_top_level;
|
||||
|
||||
if is_top_level {
|
||||
Ok(Box::new(
|
||||
SegmentFilterCollector::<LowCardSubAggBuffer>::from_req_and_validate(req, node)?,
|
||||
SegmentFilterCollector::<LowCardSubAggBuffer>::from_req_and_validate(
|
||||
req, node, req_data,
|
||||
)?,
|
||||
))
|
||||
} else {
|
||||
Ok(Box::new(
|
||||
SegmentFilterCollector::<HighCardSubAggBuffer>::from_req_and_validate(req, node)?,
|
||||
SegmentFilterCollector::<HighCardSubAggBuffer>::from_req_and_validate(
|
||||
req, node, req_data,
|
||||
)?,
|
||||
))
|
||||
}
|
||||
}
|
||||
@@ -561,7 +573,7 @@ impl<B: SubAggBuffer> Debug for SegmentFilterCollector<B> {
|
||||
f.debug_struct("SegmentFilterCollector")
|
||||
.field("buckets", &self.parent_buckets)
|
||||
.field("has_sub_aggs", &self.sub_aggregations.is_some())
|
||||
.field("accessor_idx", &self.accessor_idx)
|
||||
.field("name", &self.req_data.name)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -598,11 +610,7 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentFilterCollector<B>
|
||||
};
|
||||
|
||||
// Get the name of this filter aggregation
|
||||
let name = agg_data.per_request.filter_req_data[self.accessor_idx]
|
||||
.as_ref()
|
||||
.expect("filter_req_data slot is empty")
|
||||
.name
|
||||
.clone();
|
||||
let name = self.req_data.name.clone();
|
||||
|
||||
results.push(
|
||||
name,
|
||||
@@ -623,27 +631,24 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentFilterCollector<B>
|
||||
}
|
||||
|
||||
let mut bucket = self.parent_buckets[parent_bucket_id as usize];
|
||||
// Take the request data to avoid borrow checker issues with sub-aggregations
|
||||
let mut req = agg_data.take_filter_req_data(self.accessor_idx);
|
||||
|
||||
// Use batch filtering with O(1) BitSet lookups
|
||||
req.matching_docs_buffer.clear();
|
||||
req.evaluator
|
||||
.filter_batch(docs, &mut req.matching_docs_buffer);
|
||||
self.matching_docs_buffer.clear();
|
||||
self.req_data
|
||||
.evaluator
|
||||
.filter_batch(docs, &mut self.matching_docs_buffer);
|
||||
|
||||
bucket.doc_count += req.matching_docs_buffer.len() as u64;
|
||||
bucket.doc_count += self.matching_docs_buffer.len() as u64;
|
||||
|
||||
// Batch process sub-aggregations if we have matches
|
||||
if !req.matching_docs_buffer.is_empty() {
|
||||
if !self.matching_docs_buffer.is_empty() {
|
||||
if let Some(sub_aggs) = &mut self.sub_aggregations {
|
||||
for &doc_id in &req.matching_docs_buffer {
|
||||
for &doc_id in &self.matching_docs_buffer {
|
||||
sub_aggs.push(bucket.bucket_id, doc_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Put the request data back
|
||||
agg_data.put_back_filter_req_data(self.accessor_idx, req);
|
||||
if let Some(sub_aggs) = &mut self.sub_aggregations {
|
||||
sub_aggs.check_flush_local(agg_data)?;
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ use crate::TantivyError;
|
||||
|
||||
/// Contains all information required by the SegmentHistogramCollector to perform the
|
||||
/// histogram or date_histogram aggregation on a segment.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HistogramAggReqData {
|
||||
/// The column accessor to access the fast field values.
|
||||
pub accessor: Column<u64>,
|
||||
@@ -243,19 +244,52 @@ impl Display for HistogramBounds {
|
||||
}
|
||||
|
||||
impl HistogramBounds {
|
||||
fn contains(&self, val: f64) -> bool {
|
||||
pub(crate) fn contains(&self, val: f64) -> bool {
|
||||
val >= self.min && val <= self.max
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default, Clone, Debug, PartialEq)]
|
||||
pub(crate) struct SegmentHistogramBucketEntry {
|
||||
pub key: f64,
|
||||
pub doc_count: u64,
|
||||
pub bucket_id: BucketId,
|
||||
/// The per-bucket identifier stored in a [`SegmentHistogramBucketEntry`].
|
||||
///
|
||||
/// It is [`BucketId`] when the histogram has sub aggregations (which key their state by it), and
|
||||
/// the zero-sized `()` when it does not. Without sub aggregations the id is never read, so storing
|
||||
/// `()` drops 8 bytes per bucket (24 -> 16) and turns id assignment into a no-op.
|
||||
pub trait BucketIdSlot: Copy + Default + std::fmt::Debug + PartialEq {
|
||||
/// Assigns the next id from the provider, called once when a bucket is first filled.
|
||||
fn assign(provider: &mut BucketIdProvider) -> Self;
|
||||
/// Resolves to the `BucketId` for sub-aggregation bookkeeping.
|
||||
///
|
||||
/// Only ever called for the [`BucketId`] slot: the `()` slot is used exactly when there are no
|
||||
/// sub aggregations, so every call site is guarded by `sub_agg.is_some()` and is dead for `()`.
|
||||
fn to_bucket_id(self) -> BucketId;
|
||||
}
|
||||
impl BucketIdSlot for BucketId {
|
||||
#[inline(always)]
|
||||
fn assign(provider: &mut BucketIdProvider) -> Self {
|
||||
provider.next_bucket_id()
|
||||
}
|
||||
#[inline(always)]
|
||||
fn to_bucket_id(self) -> BucketId {
|
||||
self
|
||||
}
|
||||
}
|
||||
impl BucketIdSlot for () {
|
||||
#[inline(always)]
|
||||
fn assign(_provider: &mut BucketIdProvider) -> Self {}
|
||||
#[inline(always)]
|
||||
fn to_bucket_id(self) -> BucketId {
|
||||
unreachable!("bucket ids are only resolved when sub aggregations are present")
|
||||
}
|
||||
}
|
||||
|
||||
impl SegmentHistogramBucketEntry {
|
||||
#[derive(Default, Clone, Debug, PartialEq)]
|
||||
pub(crate) struct SegmentHistogramBucketEntry<B> {
|
||||
pub key: f64,
|
||||
pub doc_count: u64,
|
||||
pub bucket_id: B,
|
||||
}
|
||||
|
||||
impl<B: BucketIdSlot> SegmentHistogramBucketEntry<B> {
|
||||
pub(crate) fn into_intermediate_bucket_entry(
|
||||
self,
|
||||
sub_aggregation: &mut Option<HighCardBufferedSubAggs>,
|
||||
@@ -268,7 +302,7 @@ impl SegmentHistogramBucketEntry {
|
||||
.add_intermediate_aggregation_result(
|
||||
agg_data,
|
||||
&mut sub_aggregation_res,
|
||||
self.bucket_id,
|
||||
self.bucket_id.to_bucket_id(),
|
||||
)?;
|
||||
}
|
||||
Ok(IntermediateHistogramBucketEntry {
|
||||
@@ -279,39 +313,147 @@ impl SegmentHistogramBucketEntry {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
struct HistogramBuckets {
|
||||
pub buckets: FxHashMap<i64, SegmentHistogramBucketEntry>,
|
||||
/// The contiguous bucket range a histogram can span, derived from the column min/max (clamped to
|
||||
/// the histogram bounds). Buckets in `[base_pos, base_pos + len)` can be stored in a flat `Vec`
|
||||
/// indexed by `bucket_pos - base_pos`, avoiding the hash map on the hot path.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub(crate) struct DenseRange {
|
||||
/// `bucket_pos` mapped to index 0 of the dense `Vec`.
|
||||
pub(crate) base_pos: i64,
|
||||
/// Number of bucket positions in the range.
|
||||
pub(crate) len: usize,
|
||||
}
|
||||
impl HistogramBuckets {
|
||||
|
||||
/// Storage for the histogram buckets of a single parent bucket.
|
||||
///
|
||||
/// Starts out sparse (a hash map keyed by `bucket_pos`). Once enough distinct buckets have been
|
||||
/// filled that we are clearly going to cover most of the column's theoretical range, it switches
|
||||
/// to a dense `Vec` indexed by `bucket_pos - base_pos`, which removes hashing from the hot loop.
|
||||
#[derive(Clone, Debug)]
|
||||
enum HistogramBuckets<B> {
|
||||
Sparse(FxHashMap<i64, SegmentHistogramBucketEntry<B>>),
|
||||
Dense {
|
||||
base_pos: i64,
|
||||
/// One slot per bucket position; a slot with `doc_count == 0` has not been hit yet.
|
||||
buckets: Vec<SegmentHistogramBucketEntry<B>>,
|
||||
},
|
||||
}
|
||||
impl<B> Default for HistogramBuckets<B> {
|
||||
fn default() -> Self {
|
||||
HistogramBuckets::Sparse(FxHashMap::default())
|
||||
}
|
||||
}
|
||||
impl<B: BucketIdSlot> HistogramBuckets<B> {
|
||||
fn memory_consumption(&self) -> u64 {
|
||||
self.buckets.capacity() as u64 * std::mem::size_of::<SegmentHistogramBucketEntry>() as u64
|
||||
let num_slots = match self {
|
||||
HistogramBuckets::Sparse(map) => map.capacity(),
|
||||
HistogramBuckets::Dense { buckets, .. } => buckets.capacity(),
|
||||
};
|
||||
num_slots as u64 * std::mem::size_of::<SegmentHistogramBucketEntry<B>>() as u64
|
||||
}
|
||||
|
||||
/// Switches from sparse to dense storage once the dense `Vec` would use no more memory than the
|
||||
/// hash map does now, so the switch never increases memory. Called at block boundaries.
|
||||
///
|
||||
/// The `Vec` holds one `Entry` per bucket position in the range. The map additionally stores
|
||||
/// the key and a control byte per slot, at a load factor of 7/16..7/8, so for a dense histogram
|
||||
/// its footprint grows past the `Vec` well before full coverage. And since the `Vec` never
|
||||
/// grows afterwards while the map would keep growing, dense only gets relatively cheaper — so
|
||||
/// no upper bound on the range is needed: a large but sparse range simply never crosses over.
|
||||
#[inline]
|
||||
fn maybe_densify(&mut self, dense_range: Option<DenseRange>) {
|
||||
let Some(range) = dense_range else { return };
|
||||
let HistogramBuckets::Sparse(map) = self else {
|
||||
return;
|
||||
};
|
||||
let dense_bytes = range
|
||||
.len
|
||||
.saturating_mul(std::mem::size_of::<SegmentHistogramBucketEntry<B>>());
|
||||
let sparse_bytes = map
|
||||
.capacity()
|
||||
.saturating_mul(std::mem::size_of::<(i64, SegmentHistogramBucketEntry<B>)>() + 1);
|
||||
if dense_bytes > sparse_bytes {
|
||||
return;
|
||||
}
|
||||
let map = std::mem::take(map);
|
||||
let mut buckets = vec![SegmentHistogramBucketEntry::<B>::default(); range.len];
|
||||
for (bucket_pos, entry) in map {
|
||||
buckets[(bucket_pos - range.base_pos) as usize] = entry;
|
||||
}
|
||||
*self = HistogramBuckets::Dense {
|
||||
base_pos: range.base_pos,
|
||||
buckets,
|
||||
};
|
||||
}
|
||||
|
||||
/// Returns the bucket entry for `bucket_pos`, setting its key (and `bucket_id`, when `B` is
|
||||
/// [`BucketId`]) on first use.
|
||||
///
|
||||
/// For the dense variant `bucket_pos` is guaranteed to be inside the range, since it is
|
||||
/// derived from the column min/max that bounds every value (see [`compute_dense_range`]).
|
||||
#[inline]
|
||||
fn get_or_create(
|
||||
&mut self,
|
||||
bucket_pos: i64,
|
||||
bucket_id_provider: &mut BucketIdProvider,
|
||||
key_from_pos: impl FnOnce(i64) -> f64,
|
||||
) -> &mut SegmentHistogramBucketEntry<B> {
|
||||
match self {
|
||||
HistogramBuckets::Sparse(map) => {
|
||||
map.entry(bucket_pos)
|
||||
.or_insert_with(|| SegmentHistogramBucketEntry {
|
||||
key: key_from_pos(bucket_pos),
|
||||
doc_count: 0,
|
||||
bucket_id: B::assign(bucket_id_provider),
|
||||
})
|
||||
}
|
||||
HistogramBuckets::Dense { base_pos, buckets } => {
|
||||
let idx = (bucket_pos - *base_pos) as usize;
|
||||
debug_assert!(idx < buckets.len(), "bucket_pos outside the dense range");
|
||||
let entry = &mut buckets[idx];
|
||||
if entry.doc_count == 0 {
|
||||
entry.key = key_from_pos(bucket_pos);
|
||||
entry.bucket_id = B::assign(bucket_id_provider);
|
||||
}
|
||||
entry
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Consumes the storage, yielding all non-empty bucket entries.
|
||||
fn into_filled_entries(self) -> Vec<SegmentHistogramBucketEntry<B>> {
|
||||
match self {
|
||||
HistogramBuckets::Sparse(map) => map.into_values().collect(),
|
||||
HistogramBuckets::Dense { buckets, .. } => {
|
||||
buckets.into_iter().filter(|b| b.doc_count > 0).collect()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The collector puts values from the fast field into the correct buckets and does a conversion to
|
||||
/// the correct datatype.
|
||||
#[derive(Debug)]
|
||||
pub struct SegmentHistogramCollector {
|
||||
pub struct SegmentHistogramCollector<B> {
|
||||
/// The buckets containing the aggregation data.
|
||||
/// One Histogram bucket per parent bucket id.
|
||||
parent_buckets: Vec<HistogramBuckets>,
|
||||
parent_buckets: Vec<HistogramBuckets<B>>,
|
||||
sub_agg: Option<HighCardBufferedSubAggs>,
|
||||
accessor_idx: usize,
|
||||
req_data: HistogramAggReqData,
|
||||
bucket_id_provider: BucketIdProvider,
|
||||
/// Theoretical bucket range derived from the column min/max, if dense `Vec` storage is
|
||||
/// viable. `None` keeps every parent bucket in the sparse hash map.
|
||||
dense_range: Option<DenseRange>,
|
||||
}
|
||||
|
||||
impl SegmentAggregationCollector for SegmentHistogramCollector {
|
||||
impl<B: BucketIdSlot> SegmentAggregationCollector for SegmentHistogramCollector<B> {
|
||||
fn add_intermediate_aggregation_result(
|
||||
&mut self,
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
results: &mut IntermediateAggregationResults,
|
||||
parent_bucket_id: BucketId,
|
||||
) -> crate::Result<()> {
|
||||
let name = agg_data
|
||||
.get_histogram_req_data(self.accessor_idx)
|
||||
.name
|
||||
.clone();
|
||||
let name = self.req_data.name.clone();
|
||||
// TODO: avoid prepare_max_bucket here and handle empty buckets.
|
||||
self.prepare_max_bucket(parent_bucket_id, agg_data)?;
|
||||
let histogram = std::mem::take(&mut self.parent_buckets[parent_bucket_id as usize]);
|
||||
@@ -328,10 +470,13 @@ impl SegmentAggregationCollector for SegmentHistogramCollector {
|
||||
docs: &[crate::DocId],
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
let req = agg_data.take_histogram_req_data(self.accessor_idx);
|
||||
let mem_pre = self.get_memory_consumption(parent_bucket_id);
|
||||
let buckets = &mut self.parent_buckets[parent_bucket_id as usize].buckets;
|
||||
let dense_range = self.dense_range;
|
||||
let store = &mut self.parent_buckets[parent_bucket_id as usize];
|
||||
// Upgrade to dense storage before processing the block if the buckets are dense enough.
|
||||
store.maybe_densify(dense_range);
|
||||
|
||||
let req = &self.req_data;
|
||||
let bounds = req.bounds;
|
||||
let interval = req.req.interval;
|
||||
let offset = req.offset;
|
||||
@@ -340,31 +485,42 @@ impl SegmentAggregationCollector for SegmentHistogramCollector {
|
||||
agg_data
|
||||
.column_block_accessor
|
||||
.fetch_block(docs, &req.accessor);
|
||||
for (doc, val) in agg_data
|
||||
.column_block_accessor
|
||||
.iter_docid_vals(docs, &req.accessor)
|
||||
{
|
||||
let val = f64_from_fastfield_u64(val, req.field_type);
|
||||
let bucket_pos = get_bucket_pos(val);
|
||||
if bounds.contains(val) {
|
||||
let bucket = buckets.entry(bucket_pos).or_insert_with(|| {
|
||||
let key = get_bucket_key_from_pos(bucket_pos as f64, interval, offset);
|
||||
SegmentHistogramBucketEntry {
|
||||
key,
|
||||
doc_count: 0,
|
||||
bucket_id: self.bucket_id_provider.next_bucket_id(),
|
||||
}
|
||||
});
|
||||
bucket.doc_count += 1;
|
||||
if let Some(sub_agg) = &mut self.sub_agg {
|
||||
sub_agg.push(bucket.bucket_id, doc);
|
||||
// special path for nested buckets
|
||||
if let Some(sub_agg) = &mut self.sub_agg {
|
||||
for (doc, val) in agg_data
|
||||
.column_block_accessor
|
||||
.iter_docid_vals(docs, &req.accessor)
|
||||
{
|
||||
let val = f64_from_fastfield_u64(val, req.field_type);
|
||||
if bounds.contains(val) {
|
||||
let bucket = store.get_or_create(
|
||||
get_bucket_pos(val),
|
||||
&mut self.bucket_id_provider,
|
||||
|pos| get_bucket_key_from_pos(pos as f64, interval, offset),
|
||||
);
|
||||
bucket.doc_count += 1;
|
||||
sub_agg.push(bucket.bucket_id.to_bucket_id(), doc);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for val in agg_data.column_block_accessor.iter_vals() {
|
||||
let val = f64_from_fastfield_u64(val, req.field_type);
|
||||
if bounds.contains(val) {
|
||||
let bucket = store.get_or_create(
|
||||
get_bucket_pos(val),
|
||||
&mut self.bucket_id_provider,
|
||||
|pos| get_bucket_key_from_pos(pos as f64, interval, offset),
|
||||
);
|
||||
bucket.doc_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
agg_data.put_back_histogram_req_data(self.accessor_idx, req);
|
||||
|
||||
let mem_delta = self.get_memory_consumption(parent_bucket_id) - mem_pre;
|
||||
if mem_delta > 0 {
|
||||
// `checked_sub` is `None` when densifying shrank the accounted memory; only account growth.
|
||||
if let Some(mem_delta) = self
|
||||
.get_memory_consumption(parent_bucket_id)
|
||||
.checked_sub(mem_pre)
|
||||
{
|
||||
agg_data.context.limits.add_memory_consumed(mem_delta)?;
|
||||
}
|
||||
|
||||
@@ -388,9 +544,7 @@ impl SegmentAggregationCollector for SegmentHistogramCollector {
|
||||
_agg_data: &AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
while self.parent_buckets.len() <= max_bucket as usize {
|
||||
self.parent_buckets.push(HistogramBuckets {
|
||||
buckets: FxHashMap::default(),
|
||||
});
|
||||
self.parent_buckets.push(HistogramBuckets::default());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -407,7 +561,7 @@ impl SegmentAggregationCollector for SegmentHistogramCollector {
|
||||
}
|
||||
}
|
||||
|
||||
impl SegmentHistogramCollector {
|
||||
impl<B: BucketIdSlot> SegmentHistogramCollector<B> {
|
||||
fn get_memory_consumption(&self, parent_bucket_id: BucketId) -> u64 {
|
||||
self.parent_buckets[parent_bucket_id as usize].memory_consumption()
|
||||
}
|
||||
@@ -416,21 +570,19 @@ impl SegmentHistogramCollector {
|
||||
fn add_intermediate_bucket_result(
|
||||
&mut self,
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
histogram: HistogramBuckets,
|
||||
histogram: HistogramBuckets<B>,
|
||||
) -> crate::Result<IntermediateBucketResult> {
|
||||
let mut buckets = Vec::with_capacity(histogram.buckets.len());
|
||||
let filled = histogram.into_filled_entries();
|
||||
let mut buckets = Vec::with_capacity(filled.len());
|
||||
|
||||
for bucket in histogram.buckets.into_values() {
|
||||
for bucket in filled {
|
||||
let bucket_res = bucket.into_intermediate_bucket_entry(&mut self.sub_agg, agg_data);
|
||||
|
||||
buckets.push(bucket_res?);
|
||||
}
|
||||
buckets.sort_unstable_by(|b1, b2| b1.key.total_cmp(&b2.key));
|
||||
|
||||
let is_date_agg = agg_data
|
||||
.get_histogram_req_data(self.accessor_idx)
|
||||
.field_type
|
||||
== ColumnType::DateTime;
|
||||
let is_date_agg = self.req_data.field_type == ColumnType::DateTime;
|
||||
Ok(IntermediateBucketResult::Histogram {
|
||||
buckets,
|
||||
is_date_agg,
|
||||
@@ -446,32 +598,175 @@ impl SegmentHistogramCollector {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let req_data = agg_data.get_histogram_req_data_mut(node.idx_in_req_data);
|
||||
req_data.req.validate()?;
|
||||
if req_data.field_type == ColumnType::DateTime && !req_data.is_date_histogram {
|
||||
req_data.req.normalize_date_time();
|
||||
}
|
||||
req_data.bounds = req_data.req.hard_bounds.unwrap_or(HistogramBounds {
|
||||
min: f64::MIN,
|
||||
max: f64::MAX,
|
||||
});
|
||||
req_data.offset = req_data.req.offset.unwrap_or(0.0);
|
||||
let mut req_data = agg_data.per_request.histogram_req_data[node.idx_in_req_data].clone();
|
||||
normalize_histogram_req(&mut req_data)?;
|
||||
agg_data
|
||||
.context
|
||||
.limits
|
||||
.add_memory_consumed(req_data.get_memory_consumption() as u64)?;
|
||||
let dense_range = compute_dense_range(
|
||||
&req_data.accessor,
|
||||
req_data.field_type,
|
||||
req_data.req.interval,
|
||||
req_data.offset,
|
||||
req_data.bounds,
|
||||
);
|
||||
let sub_agg = sub_agg.map(BufferedSubAggs::new);
|
||||
|
||||
Ok(Self {
|
||||
parent_buckets: Default::default(),
|
||||
sub_agg,
|
||||
accessor_idx: node.idx_in_req_data,
|
||||
req_data,
|
||||
bucket_id_provider: BucketIdProvider::default(),
|
||||
dense_range,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl SegmentHistogramCollector<()> {
|
||||
/// Builds a histogram collector whose parent `t` is a dense histogram filled from
|
||||
/// `counts[t * num_time_buckets .. (t + 1) * num_time_buckets]` (row-major). Used by the fused
|
||||
/// terms×histogram collector to turn its flat 2D counters into the regular intermediate result,
|
||||
/// so cross-segment merging is shared with the general path.
|
||||
pub(crate) fn from_dense_rows(
|
||||
req_data: HistogramAggReqData,
|
||||
base_pos: i64,
|
||||
num_time_buckets: usize,
|
||||
counts: &[u32],
|
||||
) -> Self {
|
||||
let interval = req_data.req.interval;
|
||||
let offset = req_data.offset;
|
||||
let num_parents = counts.len().checked_div(num_time_buckets).unwrap_or(0);
|
||||
let parent_buckets = (0..num_parents)
|
||||
.map(|t| {
|
||||
let row = &counts[t * num_time_buckets..(t + 1) * num_time_buckets];
|
||||
let buckets = row
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(b, &doc_count)| SegmentHistogramBucketEntry {
|
||||
key: get_bucket_key_from_pos(
|
||||
(base_pos + b as i64) as f64,
|
||||
interval,
|
||||
offset,
|
||||
),
|
||||
doc_count: doc_count as u64,
|
||||
bucket_id: (),
|
||||
})
|
||||
.collect();
|
||||
HistogramBuckets::Dense { base_pos, buckets }
|
||||
})
|
||||
.collect();
|
||||
Self {
|
||||
parent_buckets,
|
||||
sub_agg: None,
|
||||
req_data,
|
||||
bucket_id_provider: BucketIdProvider::default(),
|
||||
dense_range: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validates and normalizes a histogram request in place: applies date ns-normalization (for a
|
||||
/// `histogram` on a date column) and resolves `bounds`/`offset` from the request.
|
||||
fn normalize_histogram_req(req_data: &mut HistogramAggReqData) -> crate::Result<()> {
|
||||
req_data.req.validate()?;
|
||||
if req_data.field_type == ColumnType::DateTime && !req_data.is_date_histogram {
|
||||
req_data.req.normalize_date_time();
|
||||
}
|
||||
req_data.bounds = req_data.req.hard_bounds.unwrap_or(HistogramBounds {
|
||||
min: f64::MIN,
|
||||
max: f64::MAX,
|
||||
});
|
||||
req_data.offset = req_data.req.offset.unwrap_or(0.0);
|
||||
// Drop `hard_bounds` that can't exclude any value (the column's range already sits inside
|
||||
// them): the per-doc `bounds.contains` check is then a no-op, so collapsing to the unbounded
|
||||
// sentinel lets the histogram hot loop skip it and the fused term×histogram path derive
|
||||
// per-term counts from the grid. Only this collect-time filter is touched — empty-bucket
|
||||
// emission reads `req.hard_bounds` directly (see `get_req_min_max`), and `hard_bounds` only
|
||||
// ever clips that range, so a wider-than-data bound leaves the result unchanged.
|
||||
if req_data.req.hard_bounds.is_some() {
|
||||
let col_min = f64_from_fastfield_u64(req_data.accessor.min_value(), req_data.field_type);
|
||||
let col_max = f64_from_fastfield_u64(req_data.accessor.max_value(), req_data.field_type);
|
||||
if col_min >= req_data.bounds.min && col_max <= req_data.bounds.max {
|
||||
req_data.bounds = HistogramBounds {
|
||||
min: f64::MIN,
|
||||
max: f64::MAX,
|
||||
};
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clones and normalizes (resolving interval/offset/bounds) the histogram request at `node`, and
|
||||
/// returns it together with its dense bucket range — or `None` if the column has no usable range.
|
||||
/// Used by the fused terms×histogram collector, which then owns the normalized request.
|
||||
pub(crate) fn prepare_histogram_dense_range(
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
) -> crate::Result<Option<(HistogramAggReqData, DenseRange)>> {
|
||||
let mut req_data = agg_data.per_request.histogram_req_data[node.idx_in_req_data].clone();
|
||||
normalize_histogram_req(&mut req_data)?;
|
||||
let dense_range = compute_dense_range(
|
||||
&req_data.accessor,
|
||||
req_data.field_type,
|
||||
req_data.req.interval,
|
||||
req_data.offset,
|
||||
req_data.bounds,
|
||||
);
|
||||
Ok(dense_range.map(|range| (req_data, range)))
|
||||
}
|
||||
|
||||
/// Builds a boxed histogram (or date histogram) segment collector, picking the bucket-id storage
|
||||
/// based on whether there are sub aggregations: `()` (no id stored) when there are none, otherwise
|
||||
/// [`BucketId`].
|
||||
pub(crate) fn build_segment_histogram_collector(
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
|
||||
if node.children.is_empty() {
|
||||
Ok(Box::new(
|
||||
SegmentHistogramCollector::<()>::from_req_and_validate(agg_data, node)?,
|
||||
))
|
||||
} else {
|
||||
Ok(Box::new(
|
||||
SegmentHistogramCollector::<BucketId>::from_req_and_validate(agg_data, node)?,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get_bucket_pos_f64(val: f64, interval: f64, offset: f64) -> f64 {
|
||||
pub(crate) fn get_bucket_pos_f64(val: f64, interval: f64, offset: f64) -> f64 {
|
||||
((val - offset) / interval).floor()
|
||||
}
|
||||
|
||||
/// Computes the dense bucket range for a column from its min/max value (clamped to the histogram
|
||||
/// bounds), or `None` if there are no values within bounds (or the range overflows `usize`).
|
||||
///
|
||||
/// There is no upper bound on the range: whether dense storage is actually used is decided later,
|
||||
/// per parent bucket, by [`HistogramBuckets::maybe_densify`] based on the memory it would save.
|
||||
///
|
||||
/// The column min/max bound every value the collector can see, so a `Vec` sized to this range can
|
||||
/// be indexed by `bucket_pos - base_pos` without any out-of-bounds check on the hot path.
|
||||
fn compute_dense_range(
|
||||
accessor: &Column<u64>,
|
||||
field_type: ColumnType,
|
||||
interval: f64,
|
||||
offset: f64,
|
||||
bounds: HistogramBounds,
|
||||
) -> Option<DenseRange> {
|
||||
let col_min = f64_from_fastfield_u64(accessor.min_value(), field_type);
|
||||
let col_max = f64_from_fastfield_u64(accessor.max_value(), field_type);
|
||||
let lo = col_min.max(bounds.min);
|
||||
let hi = col_max.min(bounds.max);
|
||||
if lo > hi {
|
||||
return None;
|
||||
}
|
||||
let base_pos = get_bucket_pos_f64(lo, interval, offset) as i64;
|
||||
let top_pos = get_bucket_pos_f64(hi, interval, offset) as i64;
|
||||
let len = usize::try_from(top_pos.checked_sub(base_pos)?.checked_add(1)?).ok()?;
|
||||
(len > 0).then_some(DenseRange { base_pos, len })
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get_bucket_key_from_pos(bucket_pos: f64, interval: f64, offset: f64) -> f64 {
|
||||
bucket_pos * interval + offset
|
||||
@@ -776,6 +1071,62 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn histogram_dense_storage_test() -> crate::Result<()> {
|
||||
histogram_dense_storage_test_with_opt(false)?;
|
||||
histogram_dense_storage_test_with_opt(true)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Exercises the switch from sparse hash map to dense `Vec` storage. The switch happens at a
|
||||
/// block boundary (a block is `COLLECT_BLOCK_BUFFER_LEN` = 64 docs), so we need many docs in a
|
||||
/// single segment, densely covering the bucket range. `with_sub_agg` toggles the `iter_vals`
|
||||
/// fast path vs. the `iter_docid_vals` path used when there is a sub aggregation.
|
||||
fn histogram_dense_storage_test_with_opt(with_sub_agg: bool) -> crate::Result<()> {
|
||||
let num_buckets = 50usize;
|
||||
let docs_per_bucket = 10usize;
|
||||
// Value `k` repeated `docs_per_bucket` times for each bucket `k`, so every value in bucket
|
||||
// `k` equals `k` and the per-bucket average is exactly `k`.
|
||||
let values: Vec<f64> = (0..num_buckets * docs_per_bucket)
|
||||
.map(|i| (i % num_buckets) as f64)
|
||||
.collect();
|
||||
// `merge_segments = true` collapses the per-value segments into a single segment with all
|
||||
// the docs, which is collected in 64-doc blocks and therefore switches to dense storage.
|
||||
let index = get_test_index_from_values(true, &values)?;
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_value(if with_sub_agg {
|
||||
json!({
|
||||
"histogram": {
|
||||
"histogram": { "field": "score_f64", "interval": 1.0 },
|
||||
"aggs": { "avg": { "avg": { "field": "score_f64" } } }
|
||||
}
|
||||
})
|
||||
} else {
|
||||
json!({
|
||||
"histogram": {
|
||||
"histogram": { "field": "score_f64", "interval": 1.0 }
|
||||
}
|
||||
})
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
|
||||
for k in 0..num_buckets {
|
||||
assert_eq!(res["histogram"]["buckets"][k]["key"], k as f64);
|
||||
assert_eq!(
|
||||
res["histogram"]["buckets"][k]["doc_count"],
|
||||
docs_per_bucket as u64
|
||||
);
|
||||
if with_sub_agg {
|
||||
assert_eq!(res["histogram"]["buckets"][k]["avg"]["value"], k as f64);
|
||||
}
|
||||
}
|
||||
assert_eq!(res["histogram"]["buckets"][num_buckets], Value::Null);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn histogram_memory_limit() -> crate::Result<()> {
|
||||
let index = get_test_index_with_num_docs(true, 100)?;
|
||||
@@ -1070,6 +1421,55 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn histogram_non_binding_hard_bounds_test_multi_segment() -> crate::Result<()> {
|
||||
histogram_non_binding_hard_bounds_test_with_opt(false)
|
||||
}
|
||||
#[test]
|
||||
fn histogram_non_binding_hard_bounds_test_single_segment() -> crate::Result<()> {
|
||||
histogram_non_binding_hard_bounds_test_with_opt(true)
|
||||
}
|
||||
/// `hard_bounds` wider than the data (here with mid-interval edges, to cover the "bound cuts a
|
||||
/// bucket" case) can't exclude any value, so the result must be identical to the same request
|
||||
/// without bounds. Guards the normalization that collapses such bounds to the unbounded
|
||||
/// sentinel so the hot loop / fused path can skip the per-doc bounds check.
|
||||
fn histogram_non_binding_hard_bounds_test_with_opt(merge_segments: bool) -> crate::Result<()> {
|
||||
let values = vec![10.0, 12.0, 14.0, 16.0, 10.0, 13.0, 10.0, 12.0];
|
||||
let index = get_test_index_from_values(merge_segments, &values)?;
|
||||
|
||||
// Mid-interval edges, but wider than the data range [10, 16] -> they exclude nothing.
|
||||
let with_bounds: Aggregations = serde_json::from_value(json!({
|
||||
"histogram": {
|
||||
"histogram": {
|
||||
"field": "score_f64",
|
||||
"interval": 1.0,
|
||||
"hard_bounds": { "min": 9.5, "max": 16.5 }
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
let no_bounds: Aggregations = serde_json::from_value(json!({
|
||||
"histogram": {
|
||||
"histogram": { "field": "score_f64", "interval": 1.0 }
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res_bounds = exec_request(with_bounds, &index)?;
|
||||
let res_plain = exec_request(no_bounds, &index)?;
|
||||
// Dropping a non-binding bound must not change anything.
|
||||
assert_eq!(res_bounds, res_plain);
|
||||
|
||||
// Sanity: buckets span the data range with gaps filled (min_doc_count defaults to 0).
|
||||
assert_eq!(res_bounds["histogram"]["buckets"][0]["key"], 10.0);
|
||||
assert_eq!(res_bounds["histogram"]["buckets"][0]["doc_count"], 3);
|
||||
assert_eq!(res_bounds["histogram"]["buckets"][6]["key"], 16.0);
|
||||
assert_eq!(res_bounds["histogram"]["buckets"][6]["doc_count"], 1);
|
||||
assert_eq!(res_bounds["histogram"]["buckets"][7], Value::Null);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn histogram_empty_result_behaviour_test_single_segment() -> crate::Result<()> {
|
||||
histogram_empty_result_behaviour_test_with_opt(true)
|
||||
|
||||
@@ -23,6 +23,7 @@ use crate::TantivyError;
|
||||
|
||||
/// Contains all information required by the SegmentRangeCollector to perform the
|
||||
/// range aggregation on a segment.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RangeAggReqData {
|
||||
/// The column accessor to access the fast field values.
|
||||
pub accessor: Column<u64>,
|
||||
@@ -161,7 +162,7 @@ pub struct SegmentRangeCollector<B: SubAggBuffer> {
|
||||
/// One for each ParentBucketId
|
||||
parent_buckets: Vec<Vec<SegmentRangeAndBucketEntry>>,
|
||||
column_type: ColumnType,
|
||||
pub(crate) accessor_idx: usize,
|
||||
pub(crate) req_data: RangeAggReqData,
|
||||
sub_agg: Option<BufferedSubAggs<B>>,
|
||||
/// Here things get a bit weird. We need to assign unique bucket ids across all
|
||||
/// parent buckets. So we keep track of the next available bucket id here.
|
||||
@@ -184,7 +185,7 @@ impl<B: SubAggBuffer> Debug for SegmentRangeCollector<B> {
|
||||
f.debug_struct("SegmentRangeCollector")
|
||||
.field("parent_buckets_len", &self.parent_buckets.len())
|
||||
.field("column_type", &self.column_type)
|
||||
.field("accessor_idx", &self.accessor_idx)
|
||||
.field("name", &self.req_data.name)
|
||||
.field("has_sub_agg", &self.sub_agg.is_some())
|
||||
.finish()
|
||||
}
|
||||
@@ -239,10 +240,7 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
|
||||
) -> crate::Result<()> {
|
||||
self.prepare_max_bucket(parent_bucket_id, agg_data)?;
|
||||
let field_type = self.column_type;
|
||||
let name = agg_data
|
||||
.get_range_req_data(self.accessor_idx)
|
||||
.name
|
||||
.to_string();
|
||||
let name = self.req_data.name.to_string();
|
||||
|
||||
let buckets = std::mem::take(&mut self.parent_buckets[parent_bucket_id as usize]);
|
||||
|
||||
@@ -281,17 +279,15 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
|
||||
docs: &[crate::DocId],
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
let req = agg_data.take_range_req_data(self.accessor_idx);
|
||||
|
||||
agg_data
|
||||
.column_block_accessor
|
||||
.fetch_block(docs, &req.accessor);
|
||||
.fetch_block(docs, &self.req_data.accessor);
|
||||
|
||||
let buckets = &mut self.parent_buckets[parent_bucket_id as usize];
|
||||
|
||||
for (doc, val) in agg_data
|
||||
.column_block_accessor
|
||||
.iter_docid_vals(docs, &req.accessor)
|
||||
.iter_docid_vals(docs, &self.req_data.accessor)
|
||||
{
|
||||
let bucket_pos = get_bucket_pos(val, buckets);
|
||||
let bucket = &mut buckets[bucket_pos];
|
||||
@@ -301,7 +297,6 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
|
||||
}
|
||||
}
|
||||
|
||||
agg_data.put_back_range_req_data(self.accessor_idx, req);
|
||||
if let Some(sub_agg) = self.sub_agg.as_mut() {
|
||||
sub_agg.check_flush_local(agg_data)?;
|
||||
}
|
||||
@@ -319,10 +314,10 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
|
||||
fn prepare_max_bucket(
|
||||
&mut self,
|
||||
max_bucket: BucketId,
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
_agg_data: &AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
while self.parent_buckets.len() <= max_bucket as usize {
|
||||
let new_buckets = self.create_new_buckets(agg_data)?;
|
||||
let new_buckets = self.create_new_buckets()?;
|
||||
self.parent_buckets.push(new_buckets);
|
||||
}
|
||||
|
||||
@@ -346,8 +341,11 @@ pub(crate) fn build_segment_range_collector(
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
) -> crate::Result<Box<dyn SegmentAggregationCollector>> {
|
||||
let accessor_idx = node.idx_in_req_data;
|
||||
let req_data = agg_data.get_range_req_data(node.idx_in_req_data);
|
||||
let req_data = agg_data.per_request.range_req_data[node.idx_in_req_data].clone();
|
||||
agg_data
|
||||
.context
|
||||
.limits
|
||||
.add_memory_consumed(req_data.get_memory_consumption() as u64)?;
|
||||
let field_type = req_data.field_type;
|
||||
|
||||
// TODO: A better metric instead of is_top_level would be the number of buckets expected.
|
||||
@@ -365,7 +363,7 @@ pub(crate) fn build_segment_range_collector(
|
||||
Ok(Box::new(SegmentRangeCollector::<LowCardSubAggBuffer> {
|
||||
sub_agg: sub_agg.map(LowCardBufferedSubAggs::new),
|
||||
column_type: field_type,
|
||||
accessor_idx,
|
||||
req_data,
|
||||
parent_buckets: Vec::new(),
|
||||
bucket_id_provider: BucketIdProvider::default(),
|
||||
limits: agg_data.context.limits.clone(),
|
||||
@@ -374,7 +372,7 @@ pub(crate) fn build_segment_range_collector(
|
||||
Ok(Box::new(SegmentRangeCollector::<HighCardSubAggBuffer> {
|
||||
sub_agg: sub_agg.map(BufferedSubAggs::new),
|
||||
column_type: field_type,
|
||||
accessor_idx,
|
||||
req_data,
|
||||
parent_buckets: Vec::new(),
|
||||
bucket_id_provider: BucketIdProvider::default(),
|
||||
limits: agg_data.context.limits.clone(),
|
||||
@@ -383,12 +381,9 @@ pub(crate) fn build_segment_range_collector(
|
||||
}
|
||||
|
||||
impl<B: SubAggBuffer> SegmentRangeCollector<B> {
|
||||
pub(crate) fn create_new_buckets(
|
||||
&mut self,
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
) -> crate::Result<Vec<SegmentRangeAndBucketEntry>> {
|
||||
pub(crate) fn create_new_buckets(&mut self) -> crate::Result<Vec<SegmentRangeAndBucketEntry>> {
|
||||
let field_type = self.column_type;
|
||||
let req_data = agg_data.get_range_req_data(self.accessor_idx);
|
||||
let req_data = &self.req_data;
|
||||
// The range input on the request is f64.
|
||||
// We need to convert to u64 ranges, because we read the values as u64.
|
||||
// The mapping from the conversion is monotonic so ordering is preserved.
|
||||
@@ -563,17 +558,16 @@ mod tests {
|
||||
get_test_index_with_num_docs,
|
||||
};
|
||||
|
||||
pub fn get_collector_from_ranges(
|
||||
ranges: Vec<RangeAggregationRange>,
|
||||
pub fn build_test_buckets(
|
||||
ranges: &[RangeAggregationRange],
|
||||
field_type: ColumnType,
|
||||
) -> SegmentRangeCollector<HighCardSubAggBuffer> {
|
||||
) -> Vec<SegmentRangeAndBucketEntry> {
|
||||
let req = RangeAggregation {
|
||||
field: "dummy".to_string(),
|
||||
ranges,
|
||||
ranges: ranges.to_vec(),
|
||||
..Default::default()
|
||||
};
|
||||
// Build buckets directly as in from_req_and_validate without AggregationsData
|
||||
let buckets: Vec<_> = extend_validate_ranges(&req.ranges, &field_type)
|
||||
extend_validate_ranges(&req.ranges, &field_type)
|
||||
.expect("unexpected error in extend_validate_ranges")
|
||||
.iter()
|
||||
.map(|range| {
|
||||
@@ -604,16 +598,7 @@ mod tests {
|
||||
},
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
SegmentRangeCollector {
|
||||
parent_buckets: vec![buckets],
|
||||
column_type: field_type,
|
||||
accessor_idx: 0,
|
||||
sub_agg: None,
|
||||
bucket_id_provider: Default::default(),
|
||||
limits: AggregationLimitsGuard::default(),
|
||||
}
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -856,10 +841,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn bucket_test_extend_range_hole() {
|
||||
let buckets = vec![(10f64..20f64).into(), (30f64..40f64).into()];
|
||||
let collector = get_collector_from_ranges(buckets, ColumnType::F64);
|
||||
let buckets = [(10f64..20f64).into(), (30f64..40f64).into()];
|
||||
let parent_buckets = [build_test_buckets(&buckets, ColumnType::F64)];
|
||||
|
||||
let buckets = collector.parent_buckets[0].clone();
|
||||
let buckets = parent_buckets[0].clone();
|
||||
assert_eq!(buckets[0].range.start, u64::MIN);
|
||||
assert_eq!(buckets[0].range.end, 10f64.to_u64());
|
||||
assert_eq!(buckets[1].range.start, 10f64.to_u64());
|
||||
@@ -875,14 +860,14 @@ mod tests {
|
||||
fn bucket_test_range_conversion_special_case() {
|
||||
// the monotonic conversion between f64 and u64, does not map f64::MIN.to_u64() ==
|
||||
// u64::MIN, but the into trait converts f64::MIN/MAX to None
|
||||
let buckets = vec![
|
||||
let buckets = [
|
||||
(f64::MIN..10f64).into(),
|
||||
(10f64..20f64).into(),
|
||||
(20f64..f64::MAX).into(),
|
||||
];
|
||||
let collector = get_collector_from_ranges(buckets, ColumnType::F64);
|
||||
let parent_buckets = [build_test_buckets(&buckets, ColumnType::F64)];
|
||||
|
||||
let buckets = collector.parent_buckets[0].clone();
|
||||
let buckets = parent_buckets[0].clone();
|
||||
assert_eq!(buckets[0].range.start, u64::MIN);
|
||||
assert_eq!(buckets[0].range.end, 10f64.to_u64());
|
||||
assert_eq!(buckets[1].range.start, 10f64.to_u64());
|
||||
@@ -894,28 +879,28 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn bucket_range_test_negative_vals() {
|
||||
let buckets = vec![(-10f64..-1f64).into()];
|
||||
let collector = get_collector_from_ranges(buckets, ColumnType::F64);
|
||||
let buckets = [(-10f64..-1f64).into()];
|
||||
let parent_buckets = [build_test_buckets(&buckets, ColumnType::F64)];
|
||||
|
||||
let buckets = collector.parent_buckets[0].clone();
|
||||
let buckets = parent_buckets[0].clone();
|
||||
assert_eq!(&buckets[0].bucket.key.to_string(), "*--10");
|
||||
assert_eq!(&buckets[buckets.len() - 1].bucket.key.to_string(), "-1-*");
|
||||
}
|
||||
#[test]
|
||||
fn bucket_range_test_positive_vals() {
|
||||
let buckets = vec![(0f64..10f64).into()];
|
||||
let collector = get_collector_from_ranges(buckets, ColumnType::F64);
|
||||
let buckets = [(0f64..10f64).into()];
|
||||
let parent_buckets = [build_test_buckets(&buckets, ColumnType::F64)];
|
||||
|
||||
let buckets = collector.parent_buckets[0].clone();
|
||||
let buckets = parent_buckets[0].clone();
|
||||
assert_eq!(&buckets[0].bucket.key.to_string(), "*-0");
|
||||
assert_eq!(&buckets[buckets.len() - 1].bucket.key.to_string(), "10-*");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_binary_search_test_u64() {
|
||||
let check_ranges = |ranges: Vec<RangeAggregationRange>| {
|
||||
let collector = get_collector_from_ranges(ranges, ColumnType::U64);
|
||||
let search = |val: u64| get_bucket_pos(val, &collector.parent_buckets[0]);
|
||||
let check_ranges = |ranges: &[RangeAggregationRange]| {
|
||||
let parent_buckets = [build_test_buckets(ranges, ColumnType::U64)];
|
||||
let search = |val: u64| get_bucket_pos(val, &parent_buckets[0]);
|
||||
|
||||
assert_eq!(search(u64::MIN), 0);
|
||||
assert_eq!(search(9), 0);
|
||||
@@ -928,7 +913,7 @@ mod tests {
|
||||
};
|
||||
|
||||
let ranges = vec![(10.0..100.0).into()];
|
||||
check_ranges(ranges);
|
||||
check_ranges(&ranges);
|
||||
|
||||
let ranges = vec![
|
||||
RangeAggregationRange {
|
||||
@@ -938,7 +923,7 @@ mod tests {
|
||||
},
|
||||
(10.0..100.0).into(),
|
||||
];
|
||||
check_ranges(ranges);
|
||||
check_ranges(&ranges);
|
||||
|
||||
let ranges = vec![
|
||||
RangeAggregationRange {
|
||||
@@ -953,15 +938,15 @@ mod tests {
|
||||
from: Some(100.0),
|
||||
},
|
||||
];
|
||||
check_ranges(ranges);
|
||||
check_ranges(&ranges);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_binary_search_test_f64() {
|
||||
let ranges = vec![(10.0..100.0).into()];
|
||||
let ranges = [(10.0..100.0).into()];
|
||||
|
||||
let collector = get_collector_from_ranges(ranges, ColumnType::F64);
|
||||
let search = |val: u64| get_bucket_pos(val, &collector.parent_buckets[0]);
|
||||
let parent_buckets = [build_test_buckets(&ranges, ColumnType::F64)];
|
||||
let search = |val: u64| get_bucket_pos(val, &parent_buckets[0]);
|
||||
|
||||
assert_eq!(search(u64::MIN), 0);
|
||||
assert_eq!(search(9f64.to_u64()), 0);
|
||||
|
||||
@@ -29,6 +29,8 @@ use crate::aggregation::{format_date, BucketId, Key};
|
||||
use crate::error::DataCorruption;
|
||||
use crate::TantivyError;
|
||||
|
||||
mod term_histogram;
|
||||
|
||||
/// Contains all information required by the SegmentTermCollector to perform the
|
||||
/// terms aggregation on a segment.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -374,9 +376,21 @@ pub(crate) fn build_segment_term_collector(
|
||||
// Let's see if we can use a vec to aggregate our data
|
||||
// instead of a hashmap.
|
||||
let col_max_value = terms_req_data.accessor.max_value();
|
||||
let max_term_id: u64 =
|
||||
let max_column_val: u64 =
|
||||
col_max_value.max(terms_req_data.missing_value_for_accessor.unwrap_or(0u64));
|
||||
|
||||
// Fused fast path: low-cardinality terms × a single `histogram`/`date_histogram` leaf over full
|
||||
// columns with a small enough bucket grid. Anything else falls through to the general path.
|
||||
if let Some(collector) = term_histogram::maybe_build_collector(
|
||||
req_data,
|
||||
node,
|
||||
&terms_req_data,
|
||||
max_column_val,
|
||||
is_top_level,
|
||||
)? {
|
||||
return Ok(collector);
|
||||
}
|
||||
|
||||
let sub_agg_collector = if has_sub_aggregations {
|
||||
Some(build_segment_agg_collectors(req_data, &node.children)?)
|
||||
} else {
|
||||
@@ -385,30 +399,30 @@ pub(crate) fn build_segment_term_collector(
|
||||
|
||||
let mut bucket_id_provider = BucketIdProvider::default();
|
||||
// Decide which bucket storage is best suited for this aggregation.
|
||||
if is_top_level && max_term_id < MAX_NUM_TERMS_FOR_VEC && !has_sub_aggregations {
|
||||
let term_buckets = VecTermBucketsNoAgg::new(max_term_id + 1, &mut bucket_id_provider);
|
||||
if is_top_level && max_column_val < MAX_NUM_TERMS_FOR_VEC && !has_sub_aggregations {
|
||||
let term_buckets = VecTermBucketsNoAgg::new(max_column_val + 1, &mut bucket_id_provider);
|
||||
let collector: SegmentTermCollector<_, HighCardSubAggBuffer> = SegmentTermCollector {
|
||||
parent_buckets: vec![term_buckets],
|
||||
sub_agg: None,
|
||||
bucket_id_provider,
|
||||
max_term_id,
|
||||
max_term_id: max_column_val,
|
||||
terms_req_data,
|
||||
};
|
||||
Ok(Box::new(collector))
|
||||
} else if is_top_level && max_term_id < MAX_NUM_TERMS_FOR_VEC {
|
||||
let term_buckets = VecTermBuckets::new(max_term_id + 1, &mut bucket_id_provider);
|
||||
} else if is_top_level && max_column_val < MAX_NUM_TERMS_FOR_VEC {
|
||||
let term_buckets = VecTermBuckets::new(max_column_val + 1, &mut bucket_id_provider);
|
||||
let sub_agg = sub_agg_collector.map(LowCardBufferedSubAggs::new);
|
||||
let collector: SegmentTermCollector<_, LowCardSubAggBuffer> = SegmentTermCollector {
|
||||
parent_buckets: vec![term_buckets],
|
||||
sub_agg,
|
||||
bucket_id_provider,
|
||||
max_term_id,
|
||||
max_term_id: max_column_val,
|
||||
terms_req_data,
|
||||
};
|
||||
Ok(Box::new(collector))
|
||||
} else if max_term_id < 8_000_000 && is_top_level {
|
||||
} else if max_column_val < 8_000_000 && is_top_level {
|
||||
let term_buckets: PagedTermMap =
|
||||
PagedTermMap::new(max_term_id + 1, &mut bucket_id_provider);
|
||||
PagedTermMap::new(max_column_val + 1, &mut bucket_id_provider);
|
||||
// Build sub-aggregation blueprint (flat pairs)
|
||||
let sub_agg = sub_agg_collector.map(BufferedSubAggs::new);
|
||||
let collector: SegmentTermCollector<PagedTermMap, HighCardSubAggBuffer> =
|
||||
@@ -416,7 +430,7 @@ pub(crate) fn build_segment_term_collector(
|
||||
parent_buckets: vec![term_buckets],
|
||||
sub_agg,
|
||||
bucket_id_provider,
|
||||
max_term_id,
|
||||
max_term_id: max_column_val,
|
||||
terms_req_data,
|
||||
};
|
||||
Ok(Box::new(collector))
|
||||
@@ -429,7 +443,7 @@ pub(crate) fn build_segment_term_collector(
|
||||
parent_buckets: vec![term_buckets],
|
||||
sub_agg,
|
||||
bucket_id_provider,
|
||||
max_term_id,
|
||||
max_term_id: max_column_val,
|
||||
terms_req_data,
|
||||
};
|
||||
Ok(Box::new(collector))
|
||||
585
src/aggregation/bucket/term_agg/term_histogram.rs
Normal file
585
src/aggregation/bucket/term_agg/term_histogram.rs
Normal file
@@ -0,0 +1,585 @@
|
||||
//! Fused collector for the very common shape `terms` (low cardinality) × a single
|
||||
//! `histogram`/`date_histogram` sub-aggregation with nothing nested below it.
|
||||
//!
|
||||
//! See [`SegmentTermHistogramCollector`] for the approach and [`maybe_build_collector`] for the
|
||||
//! conditions under which it is used.
|
||||
|
||||
use columnar::ColumnBlockAccessor;
|
||||
|
||||
use super::{Bucket, SegmentTermCollector, TermsAggReqData, VecTermBuckets};
|
||||
use crate::aggregation::agg_data::{AggKind, AggRefNode, AggregationsSegmentCtx};
|
||||
use crate::aggregation::bucket::{
|
||||
get_bucket_pos_f64, prepare_histogram_dense_range, HistogramAggReqData,
|
||||
SegmentHistogramCollector,
|
||||
};
|
||||
use crate::aggregation::buffered_sub_aggs::LowCardSubAggBuffer;
|
||||
use crate::aggregation::intermediate_agg_result::{
|
||||
IntermediateAggregationResult, IntermediateAggregationResults,
|
||||
};
|
||||
use crate::aggregation::segment_agg_result::{BucketIdProvider, SegmentAggregationCollector};
|
||||
use crate::aggregation::{f64_from_fastfield_u64, BucketId};
|
||||
|
||||
/// Maximum number of cells (`num_terms × num_time_buckets`) in the fused flat 2D grid. Above this
|
||||
/// the grid would be too large/cache-unfriendly, so we fall back to the general buffered path.
|
||||
/// `1 << 14` cells = 128 KB of `u64` counters, comfortably L2-resident.
|
||||
///
|
||||
/// Since we are only at the top-level, this won't be multiplied by any parent buckets.
|
||||
const MAX_FUSED_GRID_BUCKETS: usize = 16384;
|
||||
|
||||
/// Fused collector for `terms` (low cardinality) × a single `histogram`/`date_histogram` leaf with
|
||||
/// nothing nested below it, when the resulting `num_terms × num_time_buckets` grid is small (see
|
||||
/// [`MAX_FUSED_GRID_BUCKETS`]).
|
||||
///
|
||||
/// It keeps a flat, fully dense 2D counter grid (`counts[term * num_time_buckets + bucket]`) and a
|
||||
/// per-term total. A single pass reads both the term and histogram columns in document order and
|
||||
/// bumps the counters directly — no doc-id buffering, no per-term scattered re-fetch, no dynamic
|
||||
/// dispatch on flush, no per-bucket key/id storage during collection (keys are derived from the
|
||||
/// index at the end).
|
||||
///
|
||||
/// At result time the flat grid is expanded back into the regular term map + histogram storage and
|
||||
/// handed to the shared intermediate-result builders, so cross-segment merging is identical to the
|
||||
/// general path.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct SegmentTermHistogramCollector {
|
||||
/// Per-term count of docs *outside* `hard_bounds` (still in `doc_count`, but in no bucket).
|
||||
/// Per-term total = this + the term's `counts` row-sum; left empty when there are no hard
|
||||
/// bounds (every doc is in-bounds, so there's no remainder to track).
|
||||
term_counts: Vec<u32>,
|
||||
/// Flattened `[num_terms * num_time_buckets]` histogram counters (`u32`, see
|
||||
/// `term_counts`).
|
||||
///
|
||||
/// Each term id get its own contiguous slice of `num_time_buckets` histogram counter.
|
||||
/// When we count all docs (#nofilter), we can derive the per-term total as the sum over that
|
||||
/// term's slice.
|
||||
counts: Vec<u32>,
|
||||
/// Histogram buckets per term (the dense time-range length).
|
||||
num_time_buckets: usize,
|
||||
/// `bucket_pos` mapped to time-bucket index 0.
|
||||
base_pos: i64,
|
||||
terms_req_data: TermsAggReqData,
|
||||
/// The (cloned, normalized) histogram request: its column + interval/offset/bounds.
|
||||
hist_req_data: HistogramAggReqData,
|
||||
/// Private block accessors for both columns. We read them together, so each needs its own
|
||||
/// (the shared `agg_data` scratch accessor only holds one block at a time). Owning them keeps
|
||||
/// `collect` independent of `agg_data`.
|
||||
term_block: ColumnBlockAccessor<u64>,
|
||||
hist_block: ColumnBlockAccessor<u64>,
|
||||
/// No hard bounds, so every doc is in-bounds.
|
||||
all_docs_in_bounds: bool,
|
||||
/// Both columns are full (fused-path precondition); cached so `collect` skips the per-block
|
||||
/// cardinality lookup in `fetch_block`.
|
||||
is_full: bool,
|
||||
}
|
||||
|
||||
impl SegmentAggregationCollector for SegmentTermHistogramCollector {
|
||||
fn add_intermediate_aggregation_result(
|
||||
&mut self,
|
||||
agg_data: &AggregationsSegmentCtx,
|
||||
results: &mut IntermediateAggregationResults,
|
||||
parent_bucket_id: BucketId,
|
||||
) -> crate::Result<()> {
|
||||
debug_assert_eq!(
|
||||
parent_bucket_id, 0,
|
||||
"fused term-histogram collector is top-level only"
|
||||
);
|
||||
// Expand the flat grid back into the regular structures and reuse the shared builders, so
|
||||
// ordering/cut-off/dict handling and cross-segment merging match the general path exactly.
|
||||
let mut bucket_id_provider = BucketIdProvider::default();
|
||||
// Per-term total = histogram row-sum (in-bounds) + `term_counts` (out-of-bounds remainder,
|
||||
// empty when there are no hard bounds).
|
||||
let term_buckets = VecTermBuckets {
|
||||
buckets: self
|
||||
.counts
|
||||
.chunks_exact(self.num_time_buckets)
|
||||
.enumerate()
|
||||
.map(|(term_id, row)| {
|
||||
let in_bounds: u32 = row.iter().sum();
|
||||
let out_of_bounds = self.term_counts.get(term_id).copied().unwrap_or(0);
|
||||
Bucket {
|
||||
count: in_bounds + out_of_bounds,
|
||||
bucket_id: bucket_id_provider.next_bucket_id(),
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
let mut histogram = SegmentHistogramCollector::<()>::from_dense_rows(
|
||||
self.hist_req_data.clone(),
|
||||
self.base_pos,
|
||||
self.num_time_buckets,
|
||||
&self.counts,
|
||||
);
|
||||
let name = self.terms_req_data.name.clone();
|
||||
let bucket = SegmentTermCollector::<VecTermBuckets, LowCardSubAggBuffer>::into_intermediate_bucket_result(
|
||||
&self.terms_req_data,
|
||||
Some(&mut histogram as &mut dyn SegmentAggregationCollector),
|
||||
term_buckets,
|
||||
agg_data,
|
||||
)?;
|
||||
results.push(name, IntermediateAggregationResult::Bucket(bucket))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn collect(
|
||||
&mut self,
|
||||
parent_bucket_id: BucketId,
|
||||
docs: &[crate::DocId],
|
||||
_agg_data: &mut AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
debug_assert_eq!(
|
||||
parent_bucket_id, 0,
|
||||
"fused term-histogram collector is top-level only"
|
||||
);
|
||||
|
||||
// Fetch both columns into our own accessors (we read them together, so they can't share the
|
||||
// single `agg_data` scratch accessor). The collector owns all its inputs, so `collect`
|
||||
// doesn't touch `agg_data`.
|
||||
self.term_block
|
||||
.fetch_block_with_is_full(docs, &self.terms_req_data.accessor, self.is_full);
|
||||
self.hist_block
|
||||
.fetch_block_with_is_full(docs, &self.hist_req_data.accessor, self.is_full);
|
||||
|
||||
// Hoist the loop-invariant fields into locals: the optimizer can't prove the
|
||||
// `self.counts`/`self.term_counts` writes don't alias these `self` fields, so it can't keep
|
||||
// them in registers and re-reads them from memory every iteration — ~15% slower on
|
||||
// `terms_status_with_date_histogram` when read straight from `self`.
|
||||
// Note: check which are actually relevant.
|
||||
let field_type = self.hist_req_data.field_type;
|
||||
let bounds = self.hist_req_data.bounds;
|
||||
let interval = self.hist_req_data.req.interval;
|
||||
let offset = self.hist_req_data.offset;
|
||||
let base_pos = self.base_pos;
|
||||
let num_time_buckets = self.num_time_buckets;
|
||||
let all_docs_in_bounds = self.all_docs_in_bounds;
|
||||
let term_counts = &mut self.term_counts;
|
||||
let counts = &mut self.counts;
|
||||
|
||||
// Both columns are full (checked at construction), so values align with `docs` positionally
|
||||
// and are read together in one pass.
|
||||
// In-bounds docs bump the `counts` grid, out-of-bounds bump `term_counts`; deriving the
|
||||
// total at flush avoids a per-doc `term_counts` RMW that serializes on
|
||||
// store-to-load forwarding.
|
||||
for (term_id, hist_raw) in self.term_block.iter_vals().zip(self.hist_block.iter_vals()) {
|
||||
let term_id = term_id as usize;
|
||||
let val = f64_from_fastfield_u64(hist_raw, field_type);
|
||||
if all_docs_in_bounds || bounds.contains(val) {
|
||||
let bucket = (get_bucket_pos_f64(val, interval, offset) as i64 - base_pos) as usize;
|
||||
debug_assert!(
|
||||
bucket < num_time_buckets,
|
||||
"histogram bucket outside dense range"
|
||||
);
|
||||
counts[term_id * num_time_buckets + bucket] += 1;
|
||||
} else {
|
||||
term_counts[term_id] += 1;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush(&mut self, _agg_data: &mut AggregationsSegmentCtx) -> crate::Result<()> {
|
||||
// Nothing is buffered: `collect` writes the flat grid directly.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prepare_max_bucket(
|
||||
&mut self,
|
||||
_max_bucket: BucketId,
|
||||
_agg_data: &AggregationsSegmentCtx,
|
||||
) -> crate::Result<()> {
|
||||
// Top-level: the flat grid is allocated up front.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn compute_metric_value(
|
||||
&self,
|
||||
_bucket_id: BucketId,
|
||||
_sub_agg_name: &str,
|
||||
_sub_agg_property: &str,
|
||||
_agg_data: &AggregationsSegmentCtx,
|
||||
) -> Option<f64> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds the fused terms×histogram collector for a single top-level parent, when the shape is
|
||||
/// eligible. Returns `Ok(None)` to fall back to the general buffered terms path.
|
||||
///
|
||||
/// Eligibility: top-level, low-cardinality terms over a full column with no missing/include-exclude
|
||||
/// handling; a single `histogram`/`date_histogram` leaf (no nesting below it) over a full column;
|
||||
/// and a `num_terms × num_time_buckets` grid no larger than [`MAX_FUSED_GRID_BUCKETS`].
|
||||
pub(super) fn maybe_build_collector(
|
||||
agg_data: &mut AggregationsSegmentCtx,
|
||||
node: &AggRefNode,
|
||||
terms_req_data: &TermsAggReqData,
|
||||
col_max_val: u64,
|
||||
is_top_level: bool,
|
||||
) -> crate::Result<Option<Box<dyn SegmentAggregationCollector>>> {
|
||||
// Both columns must be full (one value per doc) so their values align positionally with `docs`
|
||||
// and we can zip them. Requiring full columns also makes the terms agg's `missing` config a
|
||||
// no-op (`fetch_block_with_missing` early-returns on full columns), so we needn't check for it.
|
||||
//
|
||||
// We don't cap the term cardinality here: the flat grid is bounded by the total cell count
|
||||
// (`num_terms * num_time_buckets <= MAX_FUSED_GRID_BUCKETS`) checked below, which subsumes it.
|
||||
//
|
||||
// We only allow this at the top-level, since we don't know how many buckets are created. We
|
||||
// are less likely to get enough docs for the preallocation to be worth and there's a risk of
|
||||
// using too much memory. We could check the maximum theoretical buckets up-front and pass
|
||||
// them down.
|
||||
let fuseable = is_top_level
|
||||
// TODO: We can easily support this
|
||||
&& terms_req_data.allowed_term_ids.is_none()
|
||||
&& terms_req_data.accessor.get_cardinality().is_full()
|
||||
// The flat counters are `u32`, bumped once per value, so no count can exceed the column's
|
||||
// value count. (Essentially always true here: the column is full, so its value count
|
||||
// equals the doc count, and `DocId` is `u32`.)
|
||||
&& terms_req_data.accessor.values.num_vals() < u32::MAX
|
||||
&& node.children.len() == 1
|
||||
&& matches!(
|
||||
node.children[0].kind,
|
||||
AggKind::Histogram | AggKind::DateHistogram
|
||||
)
|
||||
&& node.children[0].children.is_empty()
|
||||
&& agg_data.per_request.histogram_req_data[node.children[0].idx_in_req_data]
|
||||
.accessor
|
||||
.get_cardinality()
|
||||
.is_full();
|
||||
if !fuseable {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Clone + normalize the histogram request and get its dense bucket range; only take the fused
|
||||
// path when the flat `num_terms × num_time_buckets` grid is small enough.
|
||||
let Some((hist_req_data, range)) = prepare_histogram_dense_range(agg_data, &node.children[0])?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let num_terms = col_max_val.saturating_add(1) as usize;
|
||||
if num_terms.saturating_mul(range.len) > MAX_FUSED_GRID_BUCKETS {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// No hard bounds means every doc is in-bounds, letting `collect` short-circuit the bounds
|
||||
// check — and leaving `term_counts` (the out-of-bounds remainder) unused, so we skip allocating
|
||||
// it.
|
||||
let all_docs_in_bounds =
|
||||
hist_req_data.bounds.min == f64::MIN && hist_req_data.bounds.max == f64::MAX;
|
||||
let counts = vec![0u32; num_terms * range.len];
|
||||
let term_counts = if all_docs_in_bounds {
|
||||
Vec::new()
|
||||
} else {
|
||||
vec![0u32; num_terms]
|
||||
};
|
||||
// Charge both grids to the aggregation memory limit.
|
||||
agg_data.context.limits.add_memory_consumed(
|
||||
((counts.len() + term_counts.len()) * std::mem::size_of::<u32>()) as u64,
|
||||
)?;
|
||||
Ok(Some(Box::new(SegmentTermHistogramCollector {
|
||||
term_counts,
|
||||
counts,
|
||||
num_time_buckets: range.len,
|
||||
base_pos: range.base_pos,
|
||||
terms_req_data: terms_req_data.clone(),
|
||||
hist_req_data,
|
||||
term_block: ColumnBlockAccessor::default(),
|
||||
hist_block: ColumnBlockAccessor::default(),
|
||||
all_docs_in_bounds,
|
||||
is_full: terms_req_data.accessor.get_cardinality().is_full(),
|
||||
})))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::aggregation::agg_req::Aggregations;
|
||||
use crate::aggregation::tests::{
|
||||
exec_request, exec_request_with_query_and_memory_limit,
|
||||
get_test_index_from_values_and_terms,
|
||||
};
|
||||
use crate::aggregation::AggregationLimitsGuard;
|
||||
|
||||
/// Hand-computed correctness check for the fused terms×histogram fast path
|
||||
/// ([`super::SegmentTermHistogramCollector`]): low-cardinality terms × a histogram leaf over
|
||||
/// full columns, exercised single- and multi-segment.
|
||||
#[test]
|
||||
fn fused_term_histogram_test() -> crate::Result<()> {
|
||||
fused_term_histogram_with_opt(false)?;
|
||||
fused_term_histogram_with_opt(true)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn fused_term_histogram_with_opt(merge_segments: bool) -> crate::Result<()> {
|
||||
// 300 docs: term = {a, b, c} by i % 3, histogram value = i % 20 (interval 1 => buckets
|
||||
// 0..19). gcd(3, 20) = 1, so every (term, bucket) pair occurs exactly 300 / 60 = 5 times.
|
||||
let docs: Vec<(f64, String)> = (0..300u64)
|
||||
.map(|i| {
|
||||
(
|
||||
(i % 20) as f64,
|
||||
["a", "b", "c"][(i % 3) as usize].to_string(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
// Two segments, to also exercise cross-segment merging of the fused per-term histograms.
|
||||
let segments = vec![docs[..150].to_vec(), docs[150..].to_vec()];
|
||||
let index = get_test_index_from_values_and_terms(merge_segments, &segments)?;
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id", "order": { "_key": "asc" } },
|
||||
"aggs": {
|
||||
"histo": { "histogram": { "field": "score_f64", "interval": 1.0 } }
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
|
||||
for (term_idx, term) in ["a", "b", "c"].iter().enumerate() {
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["key"], *term);
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["doc_count"], 100);
|
||||
let histo = &res["by_term"]["buckets"][term_idx]["histo"]["buckets"];
|
||||
for b in 0..20usize {
|
||||
assert_eq!(histo[b]["key"], b as f64, "term {term} bucket {b}");
|
||||
assert_eq!(histo[b]["doc_count"], 5, "term {term} bucket {b}");
|
||||
}
|
||||
assert_eq!(histo[20], serde_json::Value::Null);
|
||||
}
|
||||
assert_eq!(res["by_term"]["buckets"][3], serde_json::Value::Null);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A `missing` config on a *full* term column still takes the fused path (the string sentinel
|
||||
/// is just `col_max + 1`, so the column stays low-cardinality). Since no doc is missing, the
|
||||
/// real term buckets must be exactly as without `missing`.
|
||||
#[test]
|
||||
fn fused_term_histogram_with_missing_on_full_column() -> crate::Result<()> {
|
||||
let docs: Vec<(f64, String)> = (0..300u64)
|
||||
.map(|i| {
|
||||
(
|
||||
(i % 20) as f64,
|
||||
["a", "b", "c"][(i % 3) as usize].to_string(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let index = get_test_index_from_values_and_terms(true, &[docs])?;
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id", "missing": "MISSING", "order": { "_key": "asc" } },
|
||||
"aggs": {
|
||||
"histo": { "histogram": { "field": "score_f64", "interval": 1.0 } }
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
|
||||
// Column is full, so "MISSING" never applies: a, b, c are unchanged (100 docs, 5 per
|
||||
// bucket).
|
||||
for (term_idx, term) in ["a", "b", "c"].iter().enumerate() {
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["key"], *term);
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["doc_count"], 100);
|
||||
let histo = &res["by_term"]["buckets"][term_idx]["histo"]["buckets"];
|
||||
for b in 0..20usize {
|
||||
assert_eq!(histo[b]["doc_count"], 5, "term {term} bucket {b}");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Term cardinality above the general path's `MAX_NUM_TERMS_FOR_VEC` (100) still fuses: the
|
||||
/// flat grid is bounded by the total cell count (`num_terms * num_time_buckets`), not the
|
||||
/// term count.
|
||||
#[test]
|
||||
fn fused_term_histogram_many_terms() -> crate::Result<()> {
|
||||
let num_terms = 150usize;
|
||||
let docs_per_term = 2usize;
|
||||
// All docs share histogram value 0 (a single bucket), so the grid is 150 x 1 = 150 cells.
|
||||
let docs: Vec<(f64, String)> = (0..num_terms * docs_per_term)
|
||||
.map(|i| (0.0, format!("t{:03}", i % num_terms)))
|
||||
.collect();
|
||||
let index = get_test_index_from_values_and_terms(true, &[docs])?;
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id", "size": 1000, "order": { "_key": "asc" } },
|
||||
"aggs": {
|
||||
"histo": { "histogram": { "field": "score_f64", "interval": 1.0 } }
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
|
||||
let buckets = res["by_term"]["buckets"].as_array().unwrap();
|
||||
assert_eq!(buckets.len(), num_terms);
|
||||
for (i, bucket) in buckets.iter().enumerate() {
|
||||
assert_eq!(bucket["key"], format!("t{i:03}"));
|
||||
assert_eq!(bucket["doc_count"], docs_per_term as u64);
|
||||
assert_eq!(bucket["histo"]["buckets"][0]["key"], 0.0);
|
||||
assert_eq!(
|
||||
bucket["histo"]["buckets"][0]["doc_count"],
|
||||
docs_per_term as u64
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `hard_bounds` exercises the non-derived `term_counts` branch: a term's `doc_count` must
|
||||
/// count *every* doc with that term, including docs whose histogram value is outside the
|
||||
/// bounds (those are excluded from the histogram buckets but still counted for the term). This
|
||||
/// is the case where the per-doc `term_counts` increment cannot be replaced by the grid
|
||||
/// row-sum.
|
||||
#[test]
|
||||
fn fused_term_histogram_with_hard_bounds() -> crate::Result<()> {
|
||||
// 300 docs: term = {a, b, c} by i % 3, value = i % 20. Per term: 100 docs, each value in
|
||||
// 0..=19 occurring 5 times.
|
||||
let docs: Vec<(f64, String)> = (0..300u64)
|
||||
.map(|i| {
|
||||
(
|
||||
(i % 20) as f64,
|
||||
["a", "b", "c"][(i % 3) as usize].to_string(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let index = get_test_index_from_values_and_terms(true, &[docs])?;
|
||||
|
||||
// hard_bounds [5, 14] (inclusive) keeps only values 5..=14 in the histogram (10 buckets);
|
||||
// values 0..=4 and 15..=19 are out of bounds.
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id", "order": { "_key": "asc" } },
|
||||
"aggs": {
|
||||
"histo": {
|
||||
"histogram": {
|
||||
"field": "score_f64",
|
||||
"interval": 1.0,
|
||||
"hard_bounds": { "min": 5.0, "max": 14.0 }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
|
||||
for (term_idx, term) in ["a", "b", "c"].iter().enumerate() {
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["key"], *term);
|
||||
// doc_count includes the 50 per-term docs whose value is outside [5, 14].
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["doc_count"], 100);
|
||||
let histo = &res["by_term"]["buckets"][term_idx]["histo"]["buckets"];
|
||||
for b in 0..10usize {
|
||||
let key = 5 + b;
|
||||
assert_eq!(histo[b]["key"], key as f64, "term {term} bucket key {key}");
|
||||
assert_eq!(histo[b]["doc_count"], 5, "term {term} bucket {key}");
|
||||
}
|
||||
// Only the 10 in-bounds buckets exist.
|
||||
assert_eq!(histo[10], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Non-binding `hard_bounds` (wider than the data, with mid-interval edges) must still produce
|
||||
/// exact results via the derive-from-grid path: since no doc is out of bounds, normalization
|
||||
/// drops the bound, every doc lands in the dense range, and each term's total equals its
|
||||
/// histogram row-sum. This is the case that previously fell back to the per-doc counter only
|
||||
/// because `bounds != [MIN, MAX]`.
|
||||
#[test]
|
||||
fn fused_term_histogram_with_non_binding_hard_bounds() -> crate::Result<()> {
|
||||
// 300 docs: term = {a, b, c} by i % 3, value = i % 20. Data values span [0, 19].
|
||||
let docs: Vec<(f64, String)> = (0..300u64)
|
||||
.map(|i| {
|
||||
(
|
||||
(i % 20) as f64,
|
||||
["a", "b", "c"][(i % 3) as usize].to_string(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let index = get_test_index_from_values_and_terms(true, &[docs])?;
|
||||
|
||||
// Bounds wider than [0, 19], with mid-interval edges -> they exclude nothing.
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id", "order": { "_key": "asc" } },
|
||||
"aggs": {
|
||||
"histo": {
|
||||
"histogram": {
|
||||
"field": "score_f64",
|
||||
"interval": 1.0,
|
||||
"hard_bounds": { "min": -0.5, "max": 19.5 }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
|
||||
for (term_idx, term) in ["a", "b", "c"].iter().enumerate() {
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["key"], *term);
|
||||
// Every doc is in-bounds, so the per-term total is the full 100 (as without bounds).
|
||||
assert_eq!(res["by_term"]["buckets"][term_idx]["doc_count"], 100);
|
||||
let histo = &res["by_term"]["buckets"][term_idx]["histo"]["buckets"];
|
||||
for b in 0..20usize {
|
||||
assert_eq!(histo[b]["key"], b as f64, "term {term} bucket {b}");
|
||||
assert_eq!(histo[b]["doc_count"], 5, "term {term} bucket {b}");
|
||||
}
|
||||
assert_eq!(histo[20], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Regression: with hard bounds the fused path allocates `term_counts` (one `u32`/term) on top
|
||||
/// of the grid, and that allocation must be charged to the memory limit. With many terms and a
|
||||
/// single time bucket the two are equal in size, so a limit admitting the grid alone but not
|
||||
/// grid + `term_counts` must fail.
|
||||
#[test]
|
||||
fn fused_term_histogram_hard_bounds_charges_term_counts() -> crate::Result<()> {
|
||||
// 16k distinct terms, one doc each; values alternate in/out of the single-bucket bounds
|
||||
// [5, 5] so the bounds bind and `term_counts` is allocated. num_terms=16000,
|
||||
// num_time_buckets=1 => `counts` and `term_counts` are ~64 KB each.
|
||||
let docs: Vec<(f64, String)> = (0..16_000u64)
|
||||
.map(|i| (if i % 2 == 0 { 5.0 } else { 10.0 }, format!("t{i:05}")))
|
||||
.collect();
|
||||
let index = get_test_index_from_values_and_terms(true, &[docs])?;
|
||||
|
||||
let agg_req: Aggregations = serde_json::from_value(serde_json::json!({
|
||||
"by_term": {
|
||||
"terms": { "field": "string_id" },
|
||||
"aggs": {
|
||||
"histo": {
|
||||
"histogram": {
|
||||
"field": "score_f64",
|
||||
"interval": 1.0,
|
||||
"hard_bounds": { "min": 5.0, "max": 5.0 }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
// ~96 KB admits the grid (~64 KB) but not grid + `term_counts` (~128 KB).
|
||||
let err = exec_request_with_query_and_memory_limit(
|
||||
agg_req,
|
||||
&index,
|
||||
None,
|
||||
AggregationLimitsGuard::new(Some(96_000), None),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("memory limit was exceeded"),
|
||||
"expected a memory-limit error, got: {err}"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -138,6 +138,7 @@ impl SubAggBuffer for HighCardSubAggBuffer {
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn push(&mut self, bucket_id: BucketId, doc_id: DocId) {
|
||||
let idx = bucket_id % NUM_PARTITIONS as u32;
|
||||
let slot = &mut self.partitions[idx as usize];
|
||||
@@ -196,6 +197,7 @@ impl SubAggBuffer for LowCardSubAggBuffer {
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn push(&mut self, bucket_id: BucketId, doc_id: DocId) {
|
||||
let idx = bucket_id as usize;
|
||||
if self.per_bucket_docs.len() <= idx {
|
||||
|
||||
@@ -377,7 +377,22 @@ impl IntermediateMetricResult {
|
||||
MetricResult::ExtendedStats(intermediate_stats.finalize())
|
||||
}
|
||||
IntermediateMetricResult::Sum(intermediate_sum) => {
|
||||
MetricResult::Sum(intermediate_sum.finalize().into())
|
||||
// By default match Elasticsearch: empty / all-missing sum
|
||||
// buckets serialize as `"value": 0`, not `"value": null`.
|
||||
// The non-ES `none_if_no_match` flag on `SumAggregation`
|
||||
// opts into SQL-style `null` for downstream consumers.
|
||||
let none_if_no_match = req
|
||||
.agg
|
||||
.as_sum()
|
||||
.and_then(|sum| sum.none_if_no_match)
|
||||
.unwrap_or(false);
|
||||
let value = intermediate_sum.finalize();
|
||||
if none_if_no_match {
|
||||
MetricResult::Sum(value.into())
|
||||
} else {
|
||||
let value = Some(value.unwrap_or(0.0));
|
||||
MetricResult::Sum(value.into())
|
||||
}
|
||||
}
|
||||
IntermediateMetricResult::Percentiles(percentiles) => MetricResult::Percentiles(
|
||||
percentiles
|
||||
|
||||
@@ -171,6 +171,7 @@ impl CouponCache {
|
||||
let uninitialized_coupon = Coupon::from_hash(0);
|
||||
let mut coupon_map: Vec<Coupon> =
|
||||
vec![uninitialized_coupon; highest_term_ord as usize + 1];
|
||||
|
||||
for (term_ord, coupon) in term_ords.into_iter().zip(coupons) {
|
||||
coupon_map[term_ord as usize] = coupon;
|
||||
}
|
||||
|
||||
@@ -27,6 +27,16 @@ pub struct SumAggregation {
|
||||
/// { "field": "my_numbers", "missing": "10.0" }
|
||||
#[serde(default, deserialize_with = "deserialize_option_f64")]
|
||||
pub missing: Option<f64>,
|
||||
/// Non-Elasticsearch extension. When `Some(true)`, the serialized result
|
||||
/// returns `"value": null` if no values were collected (all documents had
|
||||
/// missing/NULL values for the field), matching the behavior of `min`,
|
||||
/// `max`, and `avg`. When `None` or `Some(false)` (the default) the
|
||||
/// result returns `"value": 0`, matching Elasticsearch.
|
||||
///
|
||||
/// Intended for SQL-style consumers where `SUM` of zero rows is `NULL`
|
||||
/// and must be distinguishable from a bucket that genuinely sums to `0`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub none_if_no_match: Option<bool>,
|
||||
}
|
||||
|
||||
impl SumAggregation {
|
||||
@@ -35,6 +45,7 @@ impl SumAggregation {
|
||||
Self {
|
||||
field: field_name,
|
||||
missing: None,
|
||||
none_if_no_match: None,
|
||||
}
|
||||
}
|
||||
/// Returns the field name the aggregation is computed on.
|
||||
@@ -59,8 +70,104 @@ impl IntermediateSum {
|
||||
pub fn merge_fruits(&mut self, other: IntermediateSum) {
|
||||
self.stats.merge_fruits(other.stats);
|
||||
}
|
||||
/// Computes the final minimum value.
|
||||
/// Computes the final sum value.
|
||||
///
|
||||
/// Returns `None` when no values were collected, matching the Rust-side
|
||||
/// behavior of `IntermediateMin`, `IntermediateMax`, and
|
||||
/// `IntermediateAvg`. The Elasticsearch-vs-SQL choice for the
|
||||
/// user-visible result is made at the boundary in
|
||||
/// [`IntermediateMetricResult::into_final_metric_result`]: by default
|
||||
/// `None` is coerced to `Some(0.0)` to match Elasticsearch
|
||||
/// (`"value": 0`), and the [`SumAggregation::none_if_no_match`] flag
|
||||
/// opts out of that coercion for SQL-style consumers.
|
||||
pub fn finalize(&self) -> Option<f64> {
|
||||
Some(self.stats.finalize().sum)
|
||||
let stats = self.stats.finalize();
|
||||
if stats.count == 0 {
|
||||
None
|
||||
} else {
|
||||
Some(stats.sum)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_sum_finalize_returns_none_when_no_values() {
|
||||
// Default IntermediateSum has count=0 — finalize should return None,
|
||||
// matching MIN/MAX/AVG behavior for all-NULL groups.
|
||||
let sum = IntermediateSum::default();
|
||||
assert_eq!(sum.finalize(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sum_finalize_returns_value_when_has_values() {
|
||||
let mut sum = IntermediateSum::default();
|
||||
// Merge in a result that has actual values
|
||||
let stats = IntermediateStats {
|
||||
count: 3,
|
||||
sum: 42.0,
|
||||
min: 10.0,
|
||||
max: 20.0,
|
||||
..Default::default()
|
||||
};
|
||||
let other = IntermediateSum::from_stats(stats);
|
||||
sum.merge_fruits(other);
|
||||
assert_eq!(sum.finalize(), Some(42.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sum_merge_two_empty_still_none() {
|
||||
let mut a = IntermediateSum::default();
|
||||
let b = IntermediateSum::default();
|
||||
a.merge_fruits(b);
|
||||
assert_eq!(a.finalize(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sum_aggregation_empty_index_default_matches_es() -> crate::Result<()> {
|
||||
use serde_json::json;
|
||||
|
||||
use crate::aggregation::agg_req::Aggregations;
|
||||
use crate::aggregation::tests::{exec_request, get_test_index_from_terms};
|
||||
|
||||
// Empty index — sum has no values to collect.
|
||||
let values: Vec<Vec<&str>> = vec![];
|
||||
let index = get_test_index_from_terms(false, &values)?;
|
||||
let agg_req: Aggregations = serde_json::from_value(json!({
|
||||
"score_sum": { "sum": { "field": "score" } }
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
// Default: match Elasticsearch — empty sum serializes as 0, not null.
|
||||
assert_eq!(res["score_sum"]["value"], 0.0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sum_aggregation_empty_index_none_if_no_match_opt_in() -> crate::Result<()> {
|
||||
use serde_json::json;
|
||||
|
||||
use crate::aggregation::agg_req::Aggregations;
|
||||
use crate::aggregation::tests::{exec_request, get_test_index_from_terms};
|
||||
|
||||
let values: Vec<Vec<&str>> = vec![];
|
||||
let index = get_test_index_from_terms(false, &values)?;
|
||||
let agg_req: Aggregations = serde_json::from_value(json!({
|
||||
"score_sum": { "sum": { "field": "score", "none_if_no_match": true } }
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let res = exec_request(agg_req, &index)?;
|
||||
// Opt-in non-ES extension — empty sum serializes as null.
|
||||
assert!(
|
||||
res["score_sum"]["value"].is_null(),
|
||||
"expected null, got {:?}",
|
||||
res["score_sum"]["value"]
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
229
src/codec/mod.rs
229
src/codec/mod.rs
@@ -1,229 +0,0 @@
|
||||
/// Codec specific to postings data.
|
||||
pub mod postings;
|
||||
|
||||
/// Standard tantivy codec. This is the codec you use by default.
|
||||
pub mod standard;
|
||||
|
||||
use std::io;
|
||||
|
||||
pub use standard::StandardCodec;
|
||||
|
||||
use crate::codec::postings::PostingsCodec;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::{Postings, TermInfo};
|
||||
use crate::query::score_combiner::DoNothingCombiner;
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::{box_scorer, Bm25Weight, BufferedUnionScorer, Scorer, SumCombiner};
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::{DocId, InvertedIndexReader, Score};
|
||||
|
||||
/// Codecs describes how data is layed out on disk.
|
||||
///
|
||||
/// For the moment, only postings codec can be custom.
|
||||
pub trait Codec: Clone + std::fmt::Debug + Send + Sync + 'static {
|
||||
/// The specific postings type used by this codec.
|
||||
type PostingsCodec: PostingsCodec;
|
||||
|
||||
/// ID of the codec. It should be unique to your codec.
|
||||
/// Make it human-readable, descriptive, short and unique.
|
||||
const ID: &'static str;
|
||||
|
||||
/// Load codec based on the codec configuration.
|
||||
fn from_json_props(json_value: &serde_json::Value) -> crate::Result<Self>;
|
||||
|
||||
/// Get codec configuration.
|
||||
fn to_json_props(&self) -> serde_json::Value;
|
||||
|
||||
/// Returns the postings codec.
|
||||
fn postings_codec(&self) -> &Self::PostingsCodec;
|
||||
}
|
||||
|
||||
/// Object-safe codec is a Codec that can be used in a trait object.
|
||||
///
|
||||
/// The point of it is to offer a way to use a codec without a proliferation of generics.
|
||||
pub trait ObjectSafeCodec: 'static + Send + Sync {
|
||||
/// Loads a type-erased Postings object for the given term.
|
||||
///
|
||||
/// If the schema used to build the index did not provide enough
|
||||
/// information to match the requested `option`, a Postings is still
|
||||
/// returned in a best-effort manner.
|
||||
fn load_postings_type_erased(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
inverted_index_reader: &InvertedIndexReader,
|
||||
) -> io::Result<Box<dyn Postings>>;
|
||||
|
||||
/// Loads a type-erased TermScorer object for the given term.
|
||||
///
|
||||
/// If the schema used to build the index did not provide enough
|
||||
/// information to match the requested `option`, a TermScorer is still
|
||||
/// returned in a best-effort manner.
|
||||
///
|
||||
/// The point of this contraption is that the return TermScorer is backed,
|
||||
/// not by Box<dyn Postings> but by the codec's concrete Postings type.
|
||||
fn load_term_scorer_type_erased(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
inverted_index_reader: &InvertedIndexReader,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
similarity_weight: Bm25Weight,
|
||||
) -> io::Result<Box<dyn Scorer>>;
|
||||
|
||||
/// Loads a type-erased PhraseScorer object for the given term.
|
||||
///
|
||||
/// If the schema used to build the index did not provide enough
|
||||
/// information to match the requested `option`, a TermScorer is still
|
||||
/// returned in a best-effort manner.
|
||||
///
|
||||
/// The point of this contraption is that the return PhraseScorer is backed,
|
||||
/// not by Box<dyn Postings> but by the codec's concrete Postings type.
|
||||
fn new_phrase_scorer_type_erased(
|
||||
&self,
|
||||
term_infos: &[(usize, TermInfo)],
|
||||
similarity_weight: Option<Bm25Weight>,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
slop: u32,
|
||||
inverted_index_reader: &InvertedIndexReader,
|
||||
) -> io::Result<Box<dyn Scorer>>;
|
||||
|
||||
/// Performs a for_each_pruning operation on the given scorer.
|
||||
///
|
||||
/// The function will go through matching documents and call the callback
|
||||
/// function for all docs with a score exceeding the threshold.
|
||||
///
|
||||
/// The function itself will return a larger threshold value,
|
||||
/// meant to update the threshold value.
|
||||
///
|
||||
/// If the codec and the scorer allow it, this function can rely on
|
||||
/// optimizations like the block-max wand.
|
||||
fn for_each_pruning(
|
||||
&self,
|
||||
threshold: Score,
|
||||
scorer: Box<dyn Scorer>,
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
);
|
||||
|
||||
/// Builds a union scorer possibly specialized if
|
||||
/// all scorers are `Term<Self::Postings>`.
|
||||
fn build_union_scorer_with_sum_combiner(
|
||||
&self,
|
||||
scorers: Vec<Box<dyn Scorer>>,
|
||||
num_docs: DocId,
|
||||
score_combiner_type: SumOrDoNothingCombiner,
|
||||
) -> Box<dyn Scorer>;
|
||||
}
|
||||
|
||||
impl<TCodec: Codec> ObjectSafeCodec for TCodec {
|
||||
fn load_postings_type_erased(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
inverted_index_reader: &InvertedIndexReader,
|
||||
) -> io::Result<Box<dyn Postings>> {
|
||||
let postings = inverted_index_reader
|
||||
.read_postings_from_terminfo_specialized(term_info, option, self)?;
|
||||
Ok(Box::new(postings))
|
||||
}
|
||||
|
||||
fn load_term_scorer_type_erased(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
inverted_index_reader: &InvertedIndexReader,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
similarity_weight: Bm25Weight,
|
||||
) -> io::Result<Box<dyn Scorer>> {
|
||||
let scorer = inverted_index_reader.new_term_scorer_specialized(
|
||||
term_info,
|
||||
option,
|
||||
fieldnorm_reader,
|
||||
similarity_weight,
|
||||
self,
|
||||
)?;
|
||||
Ok(box_scorer(scorer))
|
||||
}
|
||||
|
||||
fn new_phrase_scorer_type_erased(
|
||||
&self,
|
||||
term_infos: &[(usize, TermInfo)],
|
||||
similarity_weight: Option<Bm25Weight>,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
slop: u32,
|
||||
inverted_index_reader: &InvertedIndexReader,
|
||||
) -> io::Result<Box<dyn Scorer>> {
|
||||
let scorer = inverted_index_reader.new_phrase_scorer_type_specialized(
|
||||
term_infos,
|
||||
similarity_weight,
|
||||
fieldnorm_reader,
|
||||
slop,
|
||||
self,
|
||||
)?;
|
||||
Ok(box_scorer(scorer))
|
||||
}
|
||||
|
||||
fn build_union_scorer_with_sum_combiner(
|
||||
&self,
|
||||
scorers: Vec<Box<dyn Scorer>>,
|
||||
num_docs: DocId,
|
||||
sum_or_do_nothing_combiner: SumOrDoNothingCombiner,
|
||||
) -> Box<dyn Scorer> {
|
||||
if !scorers.iter().all(|scorer| {
|
||||
scorer.is::<TermScorer<<<Self as Codec>::PostingsCodec as PostingsCodec>::Postings>>()
|
||||
}) {
|
||||
return box_scorer(BufferedUnionScorer::build(
|
||||
scorers,
|
||||
SumCombiner::default,
|
||||
num_docs,
|
||||
));
|
||||
}
|
||||
let specialized_scorers: Vec<
|
||||
TermScorer<<<Self as Codec>::PostingsCodec as PostingsCodec>::Postings>,
|
||||
> = scorers
|
||||
.into_iter()
|
||||
.map(|scorer| {
|
||||
*scorer.downcast::<TermScorer<_>>().ok().expect(
|
||||
"Downcast failed despite the fact we already checked the type was correct",
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
match sum_or_do_nothing_combiner {
|
||||
SumOrDoNothingCombiner::Sum => box_scorer(BufferedUnionScorer::build(
|
||||
specialized_scorers,
|
||||
SumCombiner::default,
|
||||
num_docs,
|
||||
)),
|
||||
SumOrDoNothingCombiner::DoNothing => box_scorer(BufferedUnionScorer::build(
|
||||
specialized_scorers,
|
||||
DoNothingCombiner::default,
|
||||
num_docs,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn for_each_pruning(
|
||||
&self,
|
||||
threshold: Score,
|
||||
scorer: Box<dyn Scorer>,
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) {
|
||||
let accerelerated_foreach_pruning_res =
|
||||
<TCodec as Codec>::PostingsCodec::try_accelerated_for_each_pruning(
|
||||
threshold, scorer, callback,
|
||||
);
|
||||
if let Err(mut scorer) = accerelerated_foreach_pruning_res {
|
||||
// No acceleration available. We need to do things manually.
|
||||
scorer.for_each_pruning(threshold, callback);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SumCombiner or DoNothingCombiner
|
||||
#[derive(Copy, Clone)]
|
||||
pub enum SumOrDoNothingCombiner {
|
||||
/// Sum scores together
|
||||
Sum,
|
||||
/// Do not track any score.
|
||||
DoNothing,
|
||||
}
|
||||
@@ -1,122 +0,0 @@
|
||||
use std::io;
|
||||
|
||||
/// Block-max WAND algorithm.
|
||||
pub mod block_wand;
|
||||
use common::OwnedBytes;
|
||||
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::Postings;
|
||||
use crate::query::{Bm25Weight, Scorer};
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::{DocId, Score};
|
||||
|
||||
/// Postings codec.
|
||||
pub trait PostingsCodec: Send + Sync + 'static {
|
||||
/// Serializer type for the postings codec.
|
||||
type PostingsSerializer: PostingsSerializer;
|
||||
/// Postings type for the postings codec.
|
||||
type Postings: Postings + Clone;
|
||||
/// Creates a new postings serializer.
|
||||
fn new_serializer(
|
||||
&self,
|
||||
avg_fieldnorm: Score,
|
||||
mode: IndexRecordOption,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
) -> Self::PostingsSerializer;
|
||||
|
||||
/// Loads postings
|
||||
///
|
||||
/// Record option is the option that was passed at indexing time.
|
||||
/// Requested option is the option that is requested.
|
||||
///
|
||||
/// For instance, we may have term_freq in the posting list
|
||||
/// but we can skip decompressing as we read the posting list.
|
||||
///
|
||||
/// If record option does not support the requested option,
|
||||
/// this method does NOT return an error and will in fact restrict
|
||||
/// requested_option to what is available.
|
||||
fn load_postings(
|
||||
&self,
|
||||
doc_freq: u32,
|
||||
postings_data: OwnedBytes,
|
||||
record_option: IndexRecordOption,
|
||||
requested_option: IndexRecordOption,
|
||||
positions_data: Option<OwnedBytes>,
|
||||
) -> io::Result<Self::Postings>;
|
||||
|
||||
/// If your codec supports different ways to accelerate `for_each_pruning` that's
|
||||
/// where you should implement it.
|
||||
///
|
||||
/// Returning `Err(scorer)` without mutating the scorer nor calling the callback function,
|
||||
/// is never "wrong". It just leaves the responsability to the caller to call a fallback
|
||||
/// implementation on the scorer.
|
||||
///
|
||||
/// If your codec supports BlockMax-Wand, you just need to have your
|
||||
/// postings implement `PostingsWithBlockMax` and copy what is done in the StandardPostings
|
||||
/// codec to enable it.
|
||||
fn try_accelerated_for_each_pruning(
|
||||
_threshold: Score,
|
||||
scorer: Box<dyn Scorer>,
|
||||
_callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) -> Result<(), Box<dyn Scorer>> {
|
||||
Err(scorer)
|
||||
}
|
||||
}
|
||||
|
||||
/// A postings serializer is a listener that is in charge of serializing postings
|
||||
///
|
||||
/// IO is done only once per postings, once all of the data has been received.
|
||||
/// A serializer will therefore contain internal buffers.
|
||||
///
|
||||
/// A serializer is created once and recycled for all postings.
|
||||
///
|
||||
/// Clients should use PostingsSerializer as follows.
|
||||
/// ```text
|
||||
/// // First postings list
|
||||
/// serializer.new_term(2, true);
|
||||
/// serializer.write_doc(2, 1);
|
||||
/// serializer.write_doc(6, 2);
|
||||
/// serializer.close_term(3, &mut wrt)?;
|
||||
/// // Second postings list
|
||||
/// serializer.new_term(1, true);
|
||||
/// serializer.write_doc(3, 1);
|
||||
/// serializer.close_term(1, &mut wrt)?;
|
||||
/// ```
|
||||
pub trait PostingsSerializer {
|
||||
/// The term_doc_freq here is the number of documents
|
||||
/// in the postings lists.
|
||||
///
|
||||
/// It can be used to compute the idf that will be used for the
|
||||
/// blockmax parameters.
|
||||
///
|
||||
/// If not available (e.g. if we do not collect `term_frequencies`
|
||||
/// blockwand is disabled), the term_doc_freq passed will be set 0.
|
||||
fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool);
|
||||
|
||||
/// Records a new document id for the current term.
|
||||
/// The serializer may ignore it.
|
||||
fn write_doc(&mut self, doc_id: DocId, term_freq: u32);
|
||||
|
||||
/// Closes the current term and writes the postings list associated.
|
||||
fn close_term(&mut self, doc_freq: u32, wrt: &mut impl io::Write) -> io::Result<()>;
|
||||
}
|
||||
|
||||
/// A light complement interface to Postings to allow block-max wand acceleration.
|
||||
pub trait PostingsWithBlockMax: Postings {
|
||||
/// Moves the postings to the block containign `target_doc` and returns
|
||||
/// an upperbound of the score for documents in the block.
|
||||
///
|
||||
/// `Warning`: Calling this method may leave the postings in an invalid state.
|
||||
/// callers are required to call seek before calling any other of the
|
||||
/// `Postings` method (like doc / advance etc.).
|
||||
fn seek_block_max(
|
||||
&mut self,
|
||||
target_doc: crate::DocId,
|
||||
fieldnorm_reader: &FieldNormReader,
|
||||
similarity_weight: &Bm25Weight,
|
||||
) -> Score;
|
||||
|
||||
/// Returns the last document in the current block (or Terminated if this
|
||||
/// is the last block).
|
||||
fn last_doc_in_block(&self) -> crate::DocId;
|
||||
}
|
||||
@@ -1,35 +0,0 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::codec::standard::postings::StandardPostingsCodec;
|
||||
use crate::codec::Codec;
|
||||
|
||||
/// Tantivy's default postings codec.
|
||||
pub mod postings;
|
||||
|
||||
/// Tantivy's default codec.
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
|
||||
pub struct StandardCodec;
|
||||
|
||||
impl Codec for StandardCodec {
|
||||
type PostingsCodec = StandardPostingsCodec;
|
||||
|
||||
const ID: &'static str = "tantivy-default";
|
||||
|
||||
fn from_json_props(json_value: &serde_json::Value) -> crate::Result<Self> {
|
||||
if !json_value.is_null() {
|
||||
return Err(crate::TantivyError::InvalidArgument(format!(
|
||||
"Codec property for the StandardCodec are unexpected. expected null, got {}",
|
||||
json_value.as_str().unwrap_or("null")
|
||||
)));
|
||||
}
|
||||
Ok(StandardCodec)
|
||||
}
|
||||
|
||||
fn to_json_props(&self) -> serde_json::Value {
|
||||
serde_json::Value::Null
|
||||
}
|
||||
|
||||
fn postings_codec(&self) -> &Self::PostingsCodec {
|
||||
&StandardPostingsCodec
|
||||
}
|
||||
}
|
||||
@@ -1,50 +0,0 @@
|
||||
use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
|
||||
use crate::DocId;
|
||||
|
||||
pub struct Block {
|
||||
doc_ids: [DocId; COMPRESSION_BLOCK_SIZE],
|
||||
term_freqs: [u32; COMPRESSION_BLOCK_SIZE],
|
||||
len: usize,
|
||||
}
|
||||
|
||||
impl Block {
|
||||
pub fn new() -> Self {
|
||||
Block {
|
||||
doc_ids: [0u32; COMPRESSION_BLOCK_SIZE],
|
||||
term_freqs: [0u32; COMPRESSION_BLOCK_SIZE],
|
||||
len: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn doc_ids(&self) -> &[DocId] {
|
||||
&self.doc_ids[..self.len]
|
||||
}
|
||||
|
||||
pub fn term_freqs(&self) -> &[u32] {
|
||||
&self.term_freqs[..self.len]
|
||||
}
|
||||
|
||||
pub fn clear(&mut self) {
|
||||
self.len = 0;
|
||||
}
|
||||
|
||||
pub fn append_doc(&mut self, doc: DocId, term_freq: u32) {
|
||||
let len = self.len;
|
||||
self.doc_ids[len] = doc;
|
||||
self.term_freqs[len] = term_freq;
|
||||
self.len = len + 1;
|
||||
}
|
||||
|
||||
pub fn is_full(&self) -> bool {
|
||||
self.len == COMPRESSION_BLOCK_SIZE
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len == 0
|
||||
}
|
||||
|
||||
pub fn last_doc(&self) -> DocId {
|
||||
assert_eq!(self.len, COMPRESSION_BLOCK_SIZE);
|
||||
self.doc_ids[COMPRESSION_BLOCK_SIZE - 1]
|
||||
}
|
||||
}
|
||||
@@ -1,164 +0,0 @@
|
||||
use std::io;
|
||||
|
||||
use crate::codec::postings::block_wand::{block_wand, block_wand_single_scorer};
|
||||
use crate::codec::postings::PostingsCodec;
|
||||
use crate::codec::standard::postings::block_segment_postings::BlockSegmentPostings;
|
||||
pub use crate::codec::standard::postings::segment_postings::SegmentPostings;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::positions::PositionReader;
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::{BufferedUnionScorer, Scorer, SumCombiner};
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::{DocSet as _, Score, TERMINATED};
|
||||
|
||||
mod block;
|
||||
mod block_segment_postings;
|
||||
mod segment_postings;
|
||||
mod skip;
|
||||
mod standard_postings_serializer;
|
||||
|
||||
pub use segment_postings::SegmentPostings as StandardPostings;
|
||||
pub use standard_postings_serializer::StandardPostingsSerializer;
|
||||
|
||||
/// The default postings codec for tantivy.
|
||||
pub struct StandardPostingsCodec;
|
||||
|
||||
#[expect(clippy::enum_variant_names)]
|
||||
#[derive(Debug, PartialEq, Clone, Copy, Eq)]
|
||||
pub(crate) enum FreqReadingOption {
|
||||
NoFreq,
|
||||
SkipFreq,
|
||||
ReadFreq,
|
||||
}
|
||||
|
||||
impl PostingsCodec for StandardPostingsCodec {
|
||||
type PostingsSerializer = StandardPostingsSerializer;
|
||||
type Postings = SegmentPostings;
|
||||
|
||||
fn new_serializer(
|
||||
&self,
|
||||
avg_fieldnorm: Score,
|
||||
mode: IndexRecordOption,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
) -> Self::PostingsSerializer {
|
||||
StandardPostingsSerializer::new(avg_fieldnorm, mode, fieldnorm_reader)
|
||||
}
|
||||
|
||||
fn load_postings(
|
||||
&self,
|
||||
doc_freq: u32,
|
||||
postings_data: common::OwnedBytes,
|
||||
record_option: IndexRecordOption,
|
||||
requested_option: IndexRecordOption,
|
||||
positions_data_opt: Option<common::OwnedBytes>,
|
||||
) -> io::Result<Self::Postings> {
|
||||
// Rationalize record_option/requested_option.
|
||||
let requested_option = requested_option.downgrade(record_option);
|
||||
let block_segment_postings =
|
||||
BlockSegmentPostings::open(doc_freq, postings_data, record_option, requested_option)?;
|
||||
let position_reader = positions_data_opt.map(PositionReader::open).transpose()?;
|
||||
Ok(SegmentPostings::from_block_postings(
|
||||
block_segment_postings,
|
||||
position_reader,
|
||||
))
|
||||
}
|
||||
|
||||
fn try_accelerated_for_each_pruning(
|
||||
mut threshold: Score,
|
||||
mut scorer: Box<dyn Scorer>,
|
||||
callback: &mut dyn FnMut(crate::DocId, Score) -> Score,
|
||||
) -> Result<(), Box<dyn Scorer>> {
|
||||
scorer = match scorer.downcast::<TermScorer<Self::Postings>>() {
|
||||
Ok(term_scorer) => {
|
||||
block_wand_single_scorer(*term_scorer, threshold, callback);
|
||||
return Ok(());
|
||||
}
|
||||
Err(scorer) => scorer,
|
||||
};
|
||||
let mut union_scorer =
|
||||
scorer.downcast::<BufferedUnionScorer<Box<dyn Scorer>, SumCombiner>>()?;
|
||||
if !union_scorer
|
||||
.scorers()
|
||||
.iter()
|
||||
.all(|scorer| scorer.is::<TermScorer<Self::Postings>>())
|
||||
{
|
||||
return Err(union_scorer);
|
||||
}
|
||||
let doc = union_scorer.doc();
|
||||
if doc == TERMINATED {
|
||||
return Ok(());
|
||||
}
|
||||
let score = union_scorer.score();
|
||||
if score > threshold {
|
||||
threshold = callback(doc, score);
|
||||
}
|
||||
let boxed_scorers: Vec<Box<dyn Scorer>> = union_scorer.into_scorers();
|
||||
let scorers: Vec<TermScorer<Self::Postings>> = boxed_scorers
|
||||
.into_iter()
|
||||
.map(|scorer| {
|
||||
*scorer.downcast::<TermScorer<Self::Postings>>().ok().expect(
|
||||
"Downcast failed despite the fact we already checked the type was correct",
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
block_wand(scorers, threshold, callback);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use common::OwnedBytes;
|
||||
|
||||
use super::*;
|
||||
use crate::codec::postings::PostingsSerializer as _;
|
||||
use crate::postings::Postings as _;
|
||||
|
||||
fn test_segment_postings_tf_aux(num_docs: u32, include_term_freq: bool) -> SegmentPostings {
|
||||
let mut postings_serializer =
|
||||
StandardPostingsCodec.new_serializer(1.0f32, IndexRecordOption::WithFreqs, None);
|
||||
let mut buffer = Vec::new();
|
||||
postings_serializer.new_term(num_docs, include_term_freq);
|
||||
for i in 0..num_docs {
|
||||
postings_serializer.write_doc(i, 2);
|
||||
}
|
||||
postings_serializer
|
||||
.close_term(num_docs, &mut buffer)
|
||||
.unwrap();
|
||||
StandardPostingsCodec
|
||||
.load_postings(
|
||||
num_docs,
|
||||
OwnedBytes::new(buffer),
|
||||
IndexRecordOption::WithFreqs,
|
||||
IndexRecordOption::WithFreqs,
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_segment_postings_small_block_with_and_without_freq() {
|
||||
let small_block_without_term_freq = test_segment_postings_tf_aux(1, false);
|
||||
assert!(!small_block_without_term_freq.has_freq());
|
||||
assert_eq!(small_block_without_term_freq.doc(), 0);
|
||||
assert_eq!(small_block_without_term_freq.term_freq(), 1);
|
||||
|
||||
let small_block_with_term_freq = test_segment_postings_tf_aux(1, true);
|
||||
assert!(small_block_with_term_freq.has_freq());
|
||||
assert_eq!(small_block_with_term_freq.doc(), 0);
|
||||
assert_eq!(small_block_with_term_freq.term_freq(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_segment_postings_large_block_with_and_without_freq() {
|
||||
let large_block_without_term_freq = test_segment_postings_tf_aux(128, false);
|
||||
assert!(!large_block_without_term_freq.has_freq());
|
||||
assert_eq!(large_block_without_term_freq.doc(), 0);
|
||||
assert_eq!(large_block_without_term_freq.term_freq(), 1);
|
||||
|
||||
let large_block_with_term_freq = test_segment_postings_tf_aux(128, true);
|
||||
assert!(large_block_with_term_freq.has_freq());
|
||||
assert_eq!(large_block_with_term_freq.doc(), 0);
|
||||
assert_eq!(large_block_with_term_freq.term_freq(), 2);
|
||||
}
|
||||
}
|
||||
@@ -1,184 +0,0 @@
|
||||
use std::cmp::Ordering;
|
||||
use std::io::{self, Write as _};
|
||||
|
||||
use common::{BinarySerializable as _, VInt};
|
||||
|
||||
use crate::codec::postings::PostingsSerializer;
|
||||
use crate::codec::standard::postings::block::Block;
|
||||
use crate::codec::standard::postings::skip::SkipSerializer;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::compression::{BlockEncoder, VIntEncoder as _, COMPRESSION_BLOCK_SIZE};
|
||||
use crate::query::Bm25Weight;
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::{DocId, Score};
|
||||
|
||||
/// Serializer object for tantivy's default postings format.
|
||||
pub struct StandardPostingsSerializer {
|
||||
last_doc_id_encoded: u32,
|
||||
|
||||
block_encoder: BlockEncoder,
|
||||
block: Box<Block>,
|
||||
|
||||
postings_write: Vec<u8>,
|
||||
skip_write: SkipSerializer,
|
||||
|
||||
mode: IndexRecordOption,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
|
||||
bm25_weight: Option<Bm25Weight>,
|
||||
avg_fieldnorm: Score, /* Average number of term in the field for that segment.
|
||||
* this value is used to compute the block wand information. */
|
||||
term_has_freq: bool,
|
||||
}
|
||||
|
||||
impl StandardPostingsSerializer {
|
||||
pub(crate) fn new(
|
||||
avg_fieldnorm: Score,
|
||||
mode: IndexRecordOption,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
) -> StandardPostingsSerializer {
|
||||
Self {
|
||||
last_doc_id_encoded: 0,
|
||||
block_encoder: BlockEncoder::new(),
|
||||
block: Box::new(Block::new()),
|
||||
postings_write: Vec::new(),
|
||||
skip_write: SkipSerializer::new(),
|
||||
mode,
|
||||
fieldnorm_reader,
|
||||
bm25_weight: None,
|
||||
avg_fieldnorm,
|
||||
term_has_freq: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PostingsSerializer for StandardPostingsSerializer {
|
||||
fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool) {
|
||||
self.clear();
|
||||
|
||||
self.term_has_freq = self.mode.has_freq() && record_term_freq;
|
||||
if !self.term_has_freq {
|
||||
return;
|
||||
}
|
||||
|
||||
let num_docs_in_segment: u64 =
|
||||
if let Some(fieldnorm_reader) = self.fieldnorm_reader.as_ref() {
|
||||
fieldnorm_reader.num_docs() as u64
|
||||
} else {
|
||||
return;
|
||||
};
|
||||
|
||||
if num_docs_in_segment == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
self.bm25_weight = Some(Bm25Weight::for_one_term_without_explain(
|
||||
term_doc_freq as u64,
|
||||
num_docs_in_segment,
|
||||
self.avg_fieldnorm,
|
||||
));
|
||||
}
|
||||
|
||||
fn write_doc(&mut self, doc_id: DocId, term_freq: u32) {
|
||||
self.block.append_doc(doc_id, term_freq);
|
||||
if self.block.is_full() {
|
||||
self.write_block();
|
||||
}
|
||||
}
|
||||
|
||||
fn close_term(&mut self, doc_freq: u32, output_write: &mut impl io::Write) -> io::Result<()> {
|
||||
if !self.block.is_empty() {
|
||||
// we have doc ids waiting to be written
|
||||
// this happens when the number of doc ids is
|
||||
// not a perfect multiple of our block size.
|
||||
//
|
||||
// In that case, the remaining part is encoded
|
||||
// using variable int encoding.
|
||||
{
|
||||
let block_encoded = self
|
||||
.block_encoder
|
||||
.compress_vint_sorted(self.block.doc_ids(), self.last_doc_id_encoded);
|
||||
self.postings_write.write_all(block_encoded)?;
|
||||
}
|
||||
// ... Idem for term frequencies
|
||||
if self.term_has_freq {
|
||||
let block_encoded = self
|
||||
.block_encoder
|
||||
.compress_vint_unsorted(self.block.term_freqs());
|
||||
self.postings_write.write_all(block_encoded)?;
|
||||
}
|
||||
self.block.clear();
|
||||
}
|
||||
if doc_freq >= COMPRESSION_BLOCK_SIZE as u32 {
|
||||
let skip_data = self.skip_write.data();
|
||||
VInt(skip_data.len() as u64).serialize(output_write)?;
|
||||
output_write.write_all(skip_data)?;
|
||||
}
|
||||
output_write.write_all(&self.postings_write[..])?;
|
||||
self.skip_write.clear();
|
||||
self.postings_write.clear();
|
||||
self.bm25_weight = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl StandardPostingsSerializer {
|
||||
fn clear(&mut self) {
|
||||
self.bm25_weight = None;
|
||||
self.block.clear();
|
||||
self.last_doc_id_encoded = 0;
|
||||
}
|
||||
|
||||
fn write_block(&mut self) {
|
||||
{
|
||||
// encode the doc ids
|
||||
let (num_bits, block_encoded): (u8, &[u8]) = self
|
||||
.block_encoder
|
||||
.compress_block_sorted(self.block.doc_ids(), self.last_doc_id_encoded);
|
||||
self.last_doc_id_encoded = self.block.last_doc();
|
||||
self.skip_write
|
||||
.write_doc(self.last_doc_id_encoded, num_bits);
|
||||
// last el block 0, offset block 1,
|
||||
self.postings_write.extend(block_encoded);
|
||||
}
|
||||
if self.term_has_freq {
|
||||
let (num_bits, block_encoded): (u8, &[u8]) = self
|
||||
.block_encoder
|
||||
.compress_block_unsorted(self.block.term_freqs(), true);
|
||||
self.postings_write.extend(block_encoded);
|
||||
self.skip_write.write_term_freq(num_bits);
|
||||
if self.mode.has_positions() {
|
||||
// We serialize the sum of term freqs within the skip information
|
||||
// in order to navigate through positions.
|
||||
let sum_freq = self.block.term_freqs().iter().cloned().sum();
|
||||
self.skip_write.write_total_term_freq(sum_freq);
|
||||
}
|
||||
let mut blockwand_params = (0u8, 0u32);
|
||||
if let Some(bm25_weight) = self.bm25_weight.as_ref() {
|
||||
if let Some(fieldnorm_reader) = self.fieldnorm_reader.as_ref() {
|
||||
let docs = self.block.doc_ids().iter().cloned();
|
||||
let term_freqs = self.block.term_freqs().iter().cloned();
|
||||
let fieldnorms = docs.map(|doc| fieldnorm_reader.fieldnorm_id(doc));
|
||||
blockwand_params = fieldnorms
|
||||
.zip(term_freqs)
|
||||
.max_by(
|
||||
|(left_fieldnorm_id, left_term_freq),
|
||||
(right_fieldnorm_id, right_term_freq)| {
|
||||
let left_score =
|
||||
bm25_weight.tf_factor(*left_fieldnorm_id, *left_term_freq);
|
||||
let right_score =
|
||||
bm25_weight.tf_factor(*right_fieldnorm_id, *right_term_freq);
|
||||
left_score
|
||||
.partial_cmp(&right_score)
|
||||
.unwrap_or(Ordering::Equal)
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
let (fieldnorm_id, term_freq) = blockwand_params;
|
||||
self.skip_write.write_blockwand_max(fieldnorm_id, term_freq);
|
||||
}
|
||||
self.block.clear();
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,7 @@ use common::{replace_in_place, JsonPathWriter};
|
||||
use rustc_hash::FxHashMap;
|
||||
|
||||
use crate::indexer::indexing_term::IndexingTerm;
|
||||
use crate::postings::{IndexingContext, IndexingPosition, PostingsWriter as _, PostingsWriterEnum};
|
||||
use crate::postings::{IndexingContext, IndexingPosition, PostingsWriter};
|
||||
use crate::schema::document::{ReferenceValue, ReferenceValueLeaf, Value};
|
||||
use crate::schema::{Type, DATE_TIME_PRECISION_INDEXED};
|
||||
use crate::time::format_description::well_known::Rfc3339;
|
||||
@@ -80,7 +80,7 @@ fn index_json_object<'a, V: Value<'a>>(
|
||||
text_analyzer: &mut TextAnalyzer,
|
||||
term_buffer: &mut IndexingTerm,
|
||||
json_path_writer: &mut JsonPathWriter,
|
||||
postings_writer: &mut PostingsWriterEnum,
|
||||
postings_writer: &mut dyn PostingsWriter,
|
||||
ctx: &mut IndexingContext,
|
||||
positions_per_path: &mut IndexingPositionsPerPath,
|
||||
) {
|
||||
@@ -110,7 +110,7 @@ pub(crate) fn index_json_value<'a, V: Value<'a>>(
|
||||
text_analyzer: &mut TextAnalyzer,
|
||||
term_buffer: &mut IndexingTerm,
|
||||
json_path_writer: &mut JsonPathWriter,
|
||||
postings_writer: &mut PostingsWriterEnum,
|
||||
postings_writer: &mut dyn PostingsWriter,
|
||||
ctx: &mut IndexingContext,
|
||||
positions_per_path: &mut IndexingPositionsPerPath,
|
||||
) {
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
use std::ops::{Deref as _, DerefMut as _};
|
||||
|
||||
use common::BitSet;
|
||||
use std::borrow::{Borrow, BorrowMut};
|
||||
|
||||
use common::TinySet;
|
||||
|
||||
@@ -140,19 +138,6 @@ pub trait DocSet: Send {
|
||||
buffer.len()
|
||||
}
|
||||
|
||||
/// Fills the given bitset with the documents in the docset.
|
||||
///
|
||||
/// If the docset max_doc is smaller than the largest doc, this function might not consume the
|
||||
/// docset entirely.
|
||||
fn fill_bitset(&mut self, bitset: &mut BitSet) {
|
||||
let bitset_max_value: u32 = bitset.max_value();
|
||||
let mut doc = self.doc();
|
||||
while doc < bitset_max_value {
|
||||
bitset.insert(doc);
|
||||
doc = self.advance();
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the current document
|
||||
/// Right after creating a new `DocSet`, the docset points to the first document.
|
||||
///
|
||||
@@ -293,30 +278,27 @@ impl DocSet for &mut dyn DocSet {
|
||||
fn count_including_deleted(&mut self) -> u32 {
|
||||
(**self).count_including_deleted()
|
||||
}
|
||||
|
||||
fn fill_bitset(&mut self, bitset: &mut BitSet) {
|
||||
(**self).fill_bitset(bitset);
|
||||
}
|
||||
}
|
||||
|
||||
impl<TDocSet: DocSet + ?Sized> DocSet for Box<TDocSet> {
|
||||
#[inline]
|
||||
fn advance(&mut self) -> DocId {
|
||||
self.deref_mut().advance()
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.advance()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn seek(&mut self, target: DocId) -> DocId {
|
||||
self.deref_mut().seek(target)
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.seek(target)
|
||||
}
|
||||
|
||||
fn seek_danger(&mut self, target: DocId) -> SeekDangerResult {
|
||||
self.deref_mut().seek_danger(target)
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.seek_danger(target)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn fill_buffer(&mut self, buffer: &mut [DocId; COLLECT_BLOCK_BUFFER_LEN]) -> usize {
|
||||
self.deref_mut().fill_buffer(buffer)
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.fill_buffer(buffer)
|
||||
}
|
||||
|
||||
fn fill_bitset_block(
|
||||
@@ -324,35 +306,32 @@ impl<TDocSet: DocSet + ?Sized> DocSet for Box<TDocSet> {
|
||||
min_doc: DocId,
|
||||
mask: &mut [TinySet; BLOCK_NUM_TINYBITSETS],
|
||||
) -> DocId {
|
||||
let unboxed: &mut TDocSet = &mut **self;
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.fill_bitset_block(min_doc, mask)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn doc(&self) -> DocId {
|
||||
self.deref().doc()
|
||||
let unboxed: &TDocSet = self.borrow();
|
||||
unboxed.doc()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn size_hint(&self) -> u32 {
|
||||
self.deref().size_hint()
|
||||
let unboxed: &TDocSet = self.borrow();
|
||||
unboxed.size_hint()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn cost(&self) -> u64 {
|
||||
self.deref().cost()
|
||||
let unboxed: &TDocSet = self.borrow();
|
||||
unboxed.cost()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn count(&mut self, alive_bitset: &AliveBitSet) -> u32 {
|
||||
self.deref_mut().count(alive_bitset)
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.count(alive_bitset)
|
||||
}
|
||||
|
||||
fn count_including_deleted(&mut self) -> u32 {
|
||||
self.deref_mut().count_including_deleted()
|
||||
}
|
||||
|
||||
fn fill_bitset(&mut self, bitset: &mut BitSet) {
|
||||
self.deref_mut().fill_bitset(bitset);
|
||||
let unboxed: &mut TDocSet = self.borrow_mut();
|
||||
unboxed.count_including_deleted()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
use std::borrow::Cow;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::codec::{Codec, StandardCodec};
|
||||
|
||||
/// A Codec configuration is just a serializable object.
|
||||
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||
pub struct CodecConfiguration {
|
||||
codec_id: Cow<'static, str>,
|
||||
#[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
|
||||
props: serde_json::Value,
|
||||
}
|
||||
|
||||
impl CodecConfiguration {
|
||||
/// Returns true if the codec is the standard codec.
|
||||
pub fn is_standard(&self) -> bool {
|
||||
self.codec_id == StandardCodec::ID && self.props.is_null()
|
||||
}
|
||||
|
||||
/// Creates a codec instance from the configuration.
|
||||
///
|
||||
/// If the codec id does not match the code's name, an error is returned.
|
||||
pub fn to_codec<C: Codec>(&self) -> crate::Result<C> {
|
||||
if self.codec_id != C::ID {
|
||||
return Err(crate::TantivyError::InvalidArgument(format!(
|
||||
"Codec id mismatch: expected {}, got {}",
|
||||
C::ID,
|
||||
self.codec_id
|
||||
)));
|
||||
}
|
||||
C::from_json_props(&self.props)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, C: Codec> From<&'a C> for CodecConfiguration {
|
||||
fn from(codec: &'a C) -> Self {
|
||||
CodecConfiguration {
|
||||
codec_id: Cow::Borrowed(C::ID),
|
||||
props: codec.to_json_props(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for CodecConfiguration {
|
||||
fn default() -> Self {
|
||||
CodecConfiguration::from(&StandardCodec)
|
||||
}
|
||||
}
|
||||
@@ -8,14 +8,12 @@ use std::thread::available_parallelism;
|
||||
use super::segment::Segment;
|
||||
use super::segment_reader::merge_field_meta_data;
|
||||
use super::{FieldMetadata, IndexSettings};
|
||||
use crate::codec::StandardCodec;
|
||||
use crate::core::{Executor, META_FILEPATH};
|
||||
use crate::directory::error::OpenReadError;
|
||||
#[cfg(feature = "mmap")]
|
||||
use crate::directory::MmapDirectory;
|
||||
use crate::directory::{Directory, ManagedDirectory, RamDirectory, INDEX_WRITER_LOCK};
|
||||
use crate::error::{DataCorruption, TantivyError};
|
||||
use crate::index::codec_configuration::CodecConfiguration;
|
||||
use crate::index::{IndexMeta, SegmentId, SegmentMeta, SegmentMetaInventory};
|
||||
use crate::indexer::index_writer::{
|
||||
IndexWriterOptions, MAX_NUM_THREAD, MEMORY_BUDGET_NUM_BYTES_MIN,
|
||||
@@ -61,7 +59,6 @@ fn save_new_metas(
|
||||
schema: Schema,
|
||||
index_settings: IndexSettings,
|
||||
directory: &dyn Directory,
|
||||
codec: CodecConfiguration,
|
||||
) -> crate::Result<()> {
|
||||
save_metas(
|
||||
&IndexMeta {
|
||||
@@ -70,7 +67,6 @@ fn save_new_metas(
|
||||
schema,
|
||||
opstamp: 0u64,
|
||||
payload: None,
|
||||
codec,
|
||||
},
|
||||
directory,
|
||||
)?;
|
||||
@@ -105,21 +101,18 @@ fn save_new_metas(
|
||||
/// };
|
||||
/// let index = Index::builder().schema(schema).settings(settings).create_in_ram();
|
||||
/// ```
|
||||
pub struct IndexBuilder<Codec: crate::codec::Codec = StandardCodec> {
|
||||
pub struct IndexBuilder {
|
||||
schema: Option<Schema>,
|
||||
index_settings: IndexSettings,
|
||||
tokenizer_manager: TokenizerManager,
|
||||
fast_field_tokenizer_manager: TokenizerManager,
|
||||
codec: Codec,
|
||||
}
|
||||
|
||||
impl Default for IndexBuilder<StandardCodec> {
|
||||
impl Default for IndexBuilder {
|
||||
fn default() -> Self {
|
||||
IndexBuilder::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl IndexBuilder<StandardCodec> {
|
||||
impl IndexBuilder {
|
||||
/// Creates a new `IndexBuilder`
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
@@ -127,21 +120,6 @@ impl IndexBuilder<StandardCodec> {
|
||||
index_settings: IndexSettings::default(),
|
||||
tokenizer_manager: TokenizerManager::default(),
|
||||
fast_field_tokenizer_manager: TokenizerManager::default(),
|
||||
codec: StandardCodec,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
/// Set the codec
|
||||
#[must_use]
|
||||
pub fn codec<NewCodec: crate::codec::Codec>(self, codec: NewCodec) -> IndexBuilder<NewCodec> {
|
||||
IndexBuilder {
|
||||
schema: self.schema,
|
||||
index_settings: self.index_settings,
|
||||
tokenizer_manager: self.tokenizer_manager,
|
||||
fast_field_tokenizer_manager: self.fast_field_tokenizer_manager,
|
||||
codec,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -176,7 +154,7 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
/// The index will be allocated in anonymous memory.
|
||||
/// This is useful for indexing small set of documents
|
||||
/// for instances like unit test or temporary in memory index.
|
||||
pub fn create_in_ram(self) -> Result<Index<Codec>, TantivyError> {
|
||||
pub fn create_in_ram(self) -> Result<Index, TantivyError> {
|
||||
let ram_directory = RamDirectory::create();
|
||||
self.create(ram_directory)
|
||||
}
|
||||
@@ -187,7 +165,7 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
/// If a previous index was in this directory, it returns an
|
||||
/// [`TantivyError::IndexAlreadyExists`] error.
|
||||
#[cfg(feature = "mmap")]
|
||||
pub fn create_in_dir<P: AsRef<Path>>(self, directory_path: P) -> crate::Result<Index<Codec>> {
|
||||
pub fn create_in_dir<P: AsRef<Path>>(self, directory_path: P) -> crate::Result<Index> {
|
||||
let mmap_directory: Box<dyn Directory> = Box::new(MmapDirectory::open(directory_path)?);
|
||||
if Index::exists(&*mmap_directory)? {
|
||||
return Err(TantivyError::IndexAlreadyExists);
|
||||
@@ -208,7 +186,7 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
self,
|
||||
dir: impl Into<Box<dyn Directory>>,
|
||||
mem_budget: usize,
|
||||
) -> crate::Result<SingleSegmentIndexWriter<Codec, D>> {
|
||||
) -> crate::Result<SingleSegmentIndexWriter<D>> {
|
||||
let index = self.create(dir)?;
|
||||
let index_simple_writer = SingleSegmentIndexWriter::new(index, mem_budget)?;
|
||||
Ok(index_simple_writer)
|
||||
@@ -224,7 +202,7 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
/// For other unit tests, prefer the [`RamDirectory`], see:
|
||||
/// [`IndexBuilder::create_in_ram()`].
|
||||
#[cfg(feature = "mmap")]
|
||||
pub fn create_from_tempdir(self) -> crate::Result<Index<Codec>> {
|
||||
pub fn create_from_tempdir(self) -> crate::Result<Index> {
|
||||
let mmap_directory: Box<dyn Directory> = Box::new(MmapDirectory::create_from_tempdir()?);
|
||||
self.create(mmap_directory)
|
||||
}
|
||||
@@ -237,15 +215,12 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
}
|
||||
|
||||
/// Opens or creates a new index in the provided directory
|
||||
pub fn open_or_create<T: Into<Box<dyn Directory>>>(
|
||||
self,
|
||||
dir: T,
|
||||
) -> crate::Result<Index<Codec>> {
|
||||
pub fn open_or_create<T: Into<Box<dyn Directory>>>(self, dir: T) -> crate::Result<Index> {
|
||||
let dir: Box<dyn Directory> = dir.into();
|
||||
if !Index::exists(&*dir)? {
|
||||
return self.create(dir);
|
||||
}
|
||||
let mut index: Index<Codec> = Index::<Codec>::open_with_codec(dir)?;
|
||||
let mut index = Index::open(dir)?;
|
||||
index.set_tokenizers(self.tokenizer_manager.clone());
|
||||
if index.schema() == self.get_expect_schema()? {
|
||||
Ok(index)
|
||||
@@ -269,25 +244,18 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
/// Creates a new index given an implementation of the trait `Directory`.
|
||||
///
|
||||
/// If a directory previously existed, it will be erased.
|
||||
pub fn create<T: Into<Box<dyn Directory>>>(self, dir: T) -> crate::Result<Index<Codec>> {
|
||||
self.create_avoid_monomorphization(dir.into())
|
||||
}
|
||||
|
||||
fn create_avoid_monomorphization(self, dir: Box<dyn Directory>) -> crate::Result<Index<Codec>> {
|
||||
fn create<T: Into<Box<dyn Directory>>>(self, dir: T) -> crate::Result<Index> {
|
||||
self.validate()?;
|
||||
let dir = dir.into();
|
||||
let directory = ManagedDirectory::wrap(dir)?;
|
||||
let codec: CodecConfiguration = CodecConfiguration::from(&self.codec);
|
||||
save_new_metas(
|
||||
self.get_expect_schema()?,
|
||||
self.index_settings.clone(),
|
||||
&directory,
|
||||
codec,
|
||||
)?;
|
||||
let schema = self.get_expect_schema()?;
|
||||
let mut metas = IndexMeta::with_schema_and_codec(schema, &self.codec);
|
||||
let mut metas = IndexMeta::with_schema(self.get_expect_schema()?);
|
||||
metas.index_settings = self.index_settings;
|
||||
let mut index: Index<Codec> =
|
||||
Index::<Codec>::open_from_metas(directory, &metas, SegmentMetaInventory::default())?;
|
||||
let mut index = Index::open_from_metas(directory, &metas, SegmentMetaInventory::default());
|
||||
index.set_tokenizers(self.tokenizer_manager);
|
||||
index.set_fast_field_tokenizers(self.fast_field_tokenizer_manager);
|
||||
Ok(index)
|
||||
@@ -296,7 +264,7 @@ impl<Codec: crate::codec::Codec> IndexBuilder<Codec> {
|
||||
|
||||
/// Search Index
|
||||
#[derive(Clone)]
|
||||
pub struct Index<Codec: crate::codec::Codec = crate::codec::StandardCodec> {
|
||||
pub struct Index {
|
||||
directory: ManagedDirectory,
|
||||
schema: Schema,
|
||||
settings: IndexSettings,
|
||||
@@ -304,7 +272,6 @@ pub struct Index<Codec: crate::codec::Codec = crate::codec::StandardCodec> {
|
||||
tokenizers: TokenizerManager,
|
||||
fast_field_tokenizers: TokenizerManager,
|
||||
inventory: SegmentMetaInventory,
|
||||
codec: Codec,
|
||||
}
|
||||
|
||||
impl Index {
|
||||
@@ -312,6 +279,41 @@ impl Index {
|
||||
pub fn builder() -> IndexBuilder {
|
||||
IndexBuilder::new()
|
||||
}
|
||||
/// Examines the directory to see if it contains an index.
|
||||
///
|
||||
/// Effectively, it only checks for the presence of the `meta.json` file.
|
||||
pub fn exists(dir: &dyn Directory) -> Result<bool, OpenReadError> {
|
||||
dir.exists(&META_FILEPATH)
|
||||
}
|
||||
|
||||
/// Accessor to the search executor.
|
||||
///
|
||||
/// This pool is used by default when calling `searcher.search(...)`
|
||||
/// to perform search on the individual segments.
|
||||
///
|
||||
/// By default the executor is single thread, and simply runs in the calling thread.
|
||||
pub fn search_executor(&self) -> &Executor {
|
||||
&self.executor
|
||||
}
|
||||
|
||||
/// Replace the default single thread search executor pool
|
||||
/// by a thread pool with a given number of threads.
|
||||
pub fn set_multithread_executor(&mut self, num_threads: usize) -> crate::Result<()> {
|
||||
self.executor = Executor::multi_thread(num_threads, "tantivy-search-")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Custom thread pool by a outer thread pool.
|
||||
pub fn set_executor(&mut self, executor: Executor) {
|
||||
self.executor = executor;
|
||||
}
|
||||
|
||||
/// Replace the default single thread search executor pool
|
||||
/// by a thread pool with as many threads as there are CPUs on the system.
|
||||
pub fn set_default_multithread_executor(&mut self) -> crate::Result<()> {
|
||||
let default_num_threads = available_parallelism()?.get();
|
||||
self.set_multithread_executor(default_num_threads)
|
||||
}
|
||||
|
||||
/// Creates a new index using the [`RamDirectory`].
|
||||
///
|
||||
@@ -322,13 +324,6 @@ impl Index {
|
||||
IndexBuilder::new().schema(schema).create_in_ram().unwrap()
|
||||
}
|
||||
|
||||
/// Examines the directory to see if it contains an index.
|
||||
///
|
||||
/// Effectively, it only checks for the presence of the `meta.json` file.
|
||||
pub fn exists(directory: &dyn Directory) -> Result<bool, OpenReadError> {
|
||||
directory.exists(&META_FILEPATH)
|
||||
}
|
||||
|
||||
/// Creates a new index in a given filepath.
|
||||
/// The index will use the [`MmapDirectory`].
|
||||
///
|
||||
@@ -375,108 +370,20 @@ impl Index {
|
||||
schema: Schema,
|
||||
settings: IndexSettings,
|
||||
) -> crate::Result<Index> {
|
||||
Self::create_to_avoid_monomorphization(dir.into(), schema, settings)
|
||||
}
|
||||
|
||||
fn create_to_avoid_monomorphization(
|
||||
dir: Box<dyn Directory>,
|
||||
schema: Schema,
|
||||
settings: IndexSettings,
|
||||
) -> crate::Result<Index> {
|
||||
let dir: Box<dyn Directory> = dir.into();
|
||||
let mut builder = IndexBuilder::new().schema(schema);
|
||||
builder = builder.settings(settings);
|
||||
builder.create(dir)
|
||||
}
|
||||
|
||||
/// Opens a new directory from an index path.
|
||||
#[cfg(feature = "mmap")]
|
||||
pub fn open_in_dir<P: AsRef<Path>>(directory_path: P) -> crate::Result<Index> {
|
||||
Self::open_in_dir_to_avoid_monomorphization(directory_path.as_ref())
|
||||
}
|
||||
|
||||
#[cfg(feature = "mmap")]
|
||||
#[inline(never)]
|
||||
fn open_in_dir_to_avoid_monomorphization(directory_path: &Path) -> crate::Result<Index> {
|
||||
let mmap_directory = MmapDirectory::open(directory_path)?;
|
||||
Index::open(mmap_directory)
|
||||
}
|
||||
|
||||
/// Open the index using the provided directory
|
||||
pub fn open<T: Into<Box<dyn Directory>>>(directory: T) -> crate::Result<Index> {
|
||||
Index::<StandardCodec>::open_with_codec(directory.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
/// Returns a version of this index with the standard codec.
|
||||
/// This is useful when you need to pass the index to APIs that
|
||||
/// don't care about the codec (e.g., for reading).
|
||||
pub(crate) fn with_standard_codec(&self) -> Index<StandardCodec> {
|
||||
Index {
|
||||
directory: self.directory.clone(),
|
||||
schema: self.schema.clone(),
|
||||
settings: self.settings.clone(),
|
||||
executor: self.executor.clone(),
|
||||
tokenizers: self.tokenizers.clone(),
|
||||
fast_field_tokenizers: self.fast_field_tokenizers.clone(),
|
||||
inventory: self.inventory.clone(),
|
||||
codec: StandardCodec,
|
||||
}
|
||||
}
|
||||
|
||||
/// Open the index using the provided directory
|
||||
#[inline(never)]
|
||||
pub fn open_with_codec(directory: Box<dyn Directory>) -> crate::Result<Index<Codec>> {
|
||||
let directory = ManagedDirectory::wrap(directory)?;
|
||||
let inventory = SegmentMetaInventory::default();
|
||||
let metas = load_metas(&directory, &inventory)?;
|
||||
let index: Index<Codec> = Index::<Codec>::open_from_metas(directory, &metas, inventory)?;
|
||||
Ok(index)
|
||||
}
|
||||
|
||||
/// Accessor to the codec.
|
||||
pub fn codec(&self) -> &Codec {
|
||||
&self.codec
|
||||
}
|
||||
|
||||
/// Accessor to the search executor.
|
||||
///
|
||||
/// This pool is used by default when calling `searcher.search(...)`
|
||||
/// to perform search on the individual segments.
|
||||
///
|
||||
/// By default the executor is single thread, and simply runs in the calling thread.
|
||||
pub fn search_executor(&self) -> &Executor {
|
||||
&self.executor
|
||||
}
|
||||
|
||||
/// Replace the default single thread search executor pool
|
||||
/// by a thread pool with a given number of threads.
|
||||
pub fn set_multithread_executor(&mut self, num_threads: usize) -> crate::Result<()> {
|
||||
self.executor = Executor::multi_thread(num_threads, "tantivy-search-")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Custom thread pool by a outer thread pool.
|
||||
pub fn set_executor(&mut self, executor: Executor) {
|
||||
self.executor = executor;
|
||||
}
|
||||
|
||||
/// Replace the default single thread search executor pool
|
||||
/// by a thread pool with as many threads as there are CPUs on the system.
|
||||
pub fn set_default_multithread_executor(&mut self) -> crate::Result<()> {
|
||||
let default_num_threads = available_parallelism()?.get();
|
||||
self.set_multithread_executor(default_num_threads)
|
||||
}
|
||||
|
||||
/// Creates a new index given a directory and an [`IndexMeta`].
|
||||
fn open_from_metas<C: crate::codec::Codec>(
|
||||
fn open_from_metas(
|
||||
directory: ManagedDirectory,
|
||||
metas: &IndexMeta,
|
||||
inventory: SegmentMetaInventory,
|
||||
) -> crate::Result<Index<C>> {
|
||||
) -> Index {
|
||||
let schema = metas.schema.clone();
|
||||
let codec = metas.codec.to_codec::<C>()?;
|
||||
Ok(Index {
|
||||
Index {
|
||||
settings: metas.index_settings.clone(),
|
||||
directory,
|
||||
schema,
|
||||
@@ -484,8 +391,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
fast_field_tokenizers: TokenizerManager::default(),
|
||||
executor: Executor::single_thread(),
|
||||
inventory,
|
||||
codec,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Setter for the tokenizer manager.
|
||||
@@ -541,7 +447,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
/// Create a default [`IndexReader`] for the given index.
|
||||
///
|
||||
/// See [`Index.reader_builder()`].
|
||||
pub fn reader(&self) -> crate::Result<IndexReader<Codec>> {
|
||||
pub fn reader(&self) -> crate::Result<IndexReader> {
|
||||
self.reader_builder().try_into()
|
||||
}
|
||||
|
||||
@@ -549,10 +455,17 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
///
|
||||
/// Most project should create at most one reader for a given index.
|
||||
/// This method is typically called only once per `Index` instance.
|
||||
pub fn reader_builder(&self) -> IndexReaderBuilder<Codec> {
|
||||
pub fn reader_builder(&self) -> IndexReaderBuilder {
|
||||
IndexReaderBuilder::new(self.clone())
|
||||
}
|
||||
|
||||
/// Opens a new directory from an index path.
|
||||
#[cfg(feature = "mmap")]
|
||||
pub fn open_in_dir<P: AsRef<Path>>(directory_path: P) -> crate::Result<Index> {
|
||||
let mmap_directory = MmapDirectory::open(directory_path)?;
|
||||
Index::open(mmap_directory)
|
||||
}
|
||||
|
||||
/// Returns the list of the segment metas tracked by the index.
|
||||
///
|
||||
/// Such segments can of course be part of the index,
|
||||
@@ -593,6 +506,16 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
self.inventory.new_segment_meta(segment_id, max_doc)
|
||||
}
|
||||
|
||||
/// Open the index using the provided directory
|
||||
pub fn open<T: Into<Box<dyn Directory>>>(directory: T) -> crate::Result<Index> {
|
||||
let directory = directory.into();
|
||||
let directory = ManagedDirectory::wrap(directory)?;
|
||||
let inventory = SegmentMetaInventory::default();
|
||||
let metas = load_metas(&directory, &inventory)?;
|
||||
let index = Index::open_from_metas(directory, &metas, inventory);
|
||||
Ok(index)
|
||||
}
|
||||
|
||||
/// Reads the index meta file from the directory.
|
||||
pub fn load_metas(&self) -> crate::Result<IndexMeta> {
|
||||
load_metas(self.directory(), &self.inventory)
|
||||
@@ -616,7 +539,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
pub fn writer_with_options<D: Document>(
|
||||
&self,
|
||||
options: IndexWriterOptions,
|
||||
) -> crate::Result<IndexWriter<Codec, D>> {
|
||||
) -> crate::Result<IndexWriter<D>> {
|
||||
let directory_lock = self
|
||||
.directory
|
||||
.acquire_lock(&INDEX_WRITER_LOCK)
|
||||
@@ -658,7 +581,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
&self,
|
||||
num_threads: usize,
|
||||
overall_memory_budget_in_bytes: usize,
|
||||
) -> crate::Result<IndexWriter<Codec, D>> {
|
||||
) -> crate::Result<IndexWriter<D>> {
|
||||
let memory_arena_in_bytes_per_thread = overall_memory_budget_in_bytes / num_threads;
|
||||
let options = IndexWriterOptions::builder()
|
||||
.num_worker_threads(num_threads)
|
||||
@@ -672,7 +595,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
/// That index writer only simply has a single thread and a memory budget of 15 MB.
|
||||
/// Using a single thread gives us a deterministic allocation of DocId.
|
||||
#[cfg(test)]
|
||||
pub fn writer_for_tests<D: Document>(&self) -> crate::Result<IndexWriter<Codec, D>> {
|
||||
pub fn writer_for_tests<D: Document>(&self) -> crate::Result<IndexWriter<D>> {
|
||||
self.writer_with_num_threads(1, MEMORY_BUDGET_NUM_BYTES_MIN)
|
||||
}
|
||||
|
||||
@@ -690,7 +613,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
pub fn writer<D: Document>(
|
||||
&self,
|
||||
memory_budget_in_bytes: usize,
|
||||
) -> crate::Result<IndexWriter<Codec, D>> {
|
||||
) -> crate::Result<IndexWriter<D>> {
|
||||
let mut num_threads = std::cmp::min(available_parallelism()?.get(), MAX_NUM_THREAD);
|
||||
let memory_budget_num_bytes_per_thread = memory_budget_in_bytes / num_threads;
|
||||
if memory_budget_num_bytes_per_thread < MEMORY_BUDGET_NUM_BYTES_MIN {
|
||||
@@ -717,7 +640,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
}
|
||||
|
||||
/// Returns the list of segments that are searchable
|
||||
pub fn searchable_segments(&self) -> crate::Result<Vec<Segment<Codec>>> {
|
||||
pub fn searchable_segments(&self) -> crate::Result<Vec<Segment>> {
|
||||
Ok(self
|
||||
.searchable_segment_metas()?
|
||||
.into_iter()
|
||||
@@ -726,12 +649,12 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn segment(&self, segment_meta: SegmentMeta) -> Segment<Codec> {
|
||||
pub fn segment(&self, segment_meta: SegmentMeta) -> Segment {
|
||||
Segment::for_index(self.clone(), segment_meta)
|
||||
}
|
||||
|
||||
/// Creates a new segment.
|
||||
pub fn new_segment(&self) -> Segment<Codec> {
|
||||
pub fn new_segment(&self) -> Segment {
|
||||
let segment_meta = self
|
||||
.inventory
|
||||
.new_segment_meta(SegmentId::generate_random(), 0);
|
||||
@@ -785,7 +708,7 @@ impl<Codec: crate::codec::Codec> Index<Codec> {
|
||||
}
|
||||
|
||||
impl fmt::Debug for Index {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "Index({:?})", self.directory)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,8 +5,7 @@ use std::path::PathBuf;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::SegmentComponent;
|
||||
use crate::codec::Codec;
|
||||
use crate::index::{CodecConfiguration, SegmentId};
|
||||
use crate::index::SegmentId;
|
||||
use crate::schema::Schema;
|
||||
use crate::store::Compressor;
|
||||
use crate::{Inventory, Opstamp, TrackedObject};
|
||||
@@ -287,10 +286,8 @@ pub struct IndexMeta {
|
||||
/// This payload is entirely unused by tantivy.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub payload: Option<String>,
|
||||
/// Codec configuration for the index.
|
||||
#[serde(skip_serializing_if = "CodecConfiguration::is_standard")]
|
||||
pub codec: CodecConfiguration,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
struct UntrackedIndexMeta {
|
||||
pub segments: Vec<InnerSegmentMeta>,
|
||||
@@ -300,8 +297,6 @@ struct UntrackedIndexMeta {
|
||||
pub opstamp: Opstamp,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub payload: Option<String>,
|
||||
#[serde(default)]
|
||||
pub codec: CodecConfiguration,
|
||||
}
|
||||
|
||||
impl UntrackedIndexMeta {
|
||||
@@ -316,7 +311,6 @@ impl UntrackedIndexMeta {
|
||||
schema: self.schema,
|
||||
opstamp: self.opstamp,
|
||||
payload: self.payload,
|
||||
codec: self.codec,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -327,14 +321,13 @@ impl IndexMeta {
|
||||
///
|
||||
/// This new index does not contains any segments.
|
||||
/// Opstamp will the value `0u64`.
|
||||
pub fn with_schema_and_codec<C: Codec>(schema: Schema, codec: &C) -> IndexMeta {
|
||||
pub fn with_schema(schema: Schema) -> IndexMeta {
|
||||
IndexMeta {
|
||||
index_settings: IndexSettings::default(),
|
||||
segments: vec![],
|
||||
schema,
|
||||
opstamp: 0u64,
|
||||
payload: None,
|
||||
codec: CodecConfiguration::from(codec),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -385,38 +378,14 @@ mod tests {
|
||||
schema,
|
||||
opstamp: 0u64,
|
||||
payload: None,
|
||||
codec: Default::default(),
|
||||
};
|
||||
let json_value: serde_json::Value =
|
||||
serde_json::to_value(&index_metas).expect("serialization failed");
|
||||
let json = serde_json::ser::to_string(&index_metas).expect("serialization failed");
|
||||
assert_eq!(
|
||||
&json_value,
|
||||
&serde_json::json!(
|
||||
{
|
||||
"index_settings": {
|
||||
"docstore_compression": "none",
|
||||
"docstore_blocksize": 16384
|
||||
},
|
||||
"segments": [],
|
||||
"schema": [
|
||||
{
|
||||
"name": "text",
|
||||
"type": "text",
|
||||
"options": {
|
||||
"indexing": {
|
||||
"record": "position",
|
||||
"fieldnorms": true,
|
||||
"tokenizer": "default"
|
||||
},
|
||||
"stored": false,
|
||||
"fast": false
|
||||
}
|
||||
}
|
||||
],
|
||||
"opstamp": 0
|
||||
})
|
||||
json,
|
||||
r#"{"index_settings":{"docstore_compression":"none","docstore_blocksize":16384},"segments":[],"schema":[{"name":"text","type":"text","options":{"indexing":{"record":"position","fieldnorms":true,"tokenizer":"default"},"stored":false,"fast":false}}],"opstamp":0}"#
|
||||
);
|
||||
let deser_meta: UntrackedIndexMeta = serde_json::from_value(json_value).unwrap();
|
||||
|
||||
let deser_meta: UntrackedIndexMeta = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(index_metas.index_settings, deser_meta.index_settings);
|
||||
assert_eq!(index_metas.schema, deser_meta.schema);
|
||||
assert_eq!(index_metas.opstamp, deser_meta.opstamp);
|
||||
@@ -442,39 +411,14 @@ mod tests {
|
||||
schema,
|
||||
opstamp: 0u64,
|
||||
payload: None,
|
||||
codec: Default::default(),
|
||||
};
|
||||
let json_value = serde_json::to_value(&index_metas).expect("serialization failed");
|
||||
let json = serde_json::ser::to_string(&index_metas).expect("serialization failed");
|
||||
assert_eq!(
|
||||
&json_value,
|
||||
&serde_json::json!(
|
||||
{
|
||||
"index_settings": {
|
||||
"docstore_compression": "zstd(compression_level=4)",
|
||||
"docstore_blocksize": 1000000
|
||||
},
|
||||
"segments": [],
|
||||
"schema": [
|
||||
{
|
||||
"name": "text",
|
||||
"type": "text",
|
||||
"options": {
|
||||
"indexing": {
|
||||
"record": "position",
|
||||
"fieldnorms": true,
|
||||
"tokenizer": "default"
|
||||
},
|
||||
"stored": false,
|
||||
"fast": false
|
||||
}
|
||||
}
|
||||
],
|
||||
"opstamp": 0
|
||||
}
|
||||
)
|
||||
json,
|
||||
r#"{"index_settings":{"docstore_compression":"zstd(compression_level=4)","docstore_blocksize":1000000},"segments":[],"schema":[{"name":"text","type":"text","options":{"indexing":{"record":"position","fieldnorms":true,"tokenizer":"default"},"stored":false,"fast":false}}],"opstamp":0}"#
|
||||
);
|
||||
|
||||
let deser_meta: UntrackedIndexMeta = serde_json::from_value(json_value).unwrap();
|
||||
let deser_meta: UntrackedIndexMeta = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(index_metas.index_settings, deser_meta.index_settings);
|
||||
assert_eq!(index_metas.schema, deser_meta.schema);
|
||||
assert_eq!(index_metas.opstamp, deser_meta.opstamp);
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::json_path_writer::JSON_END_OF_PATH;
|
||||
use common::{BinarySerializable, ByteCount, OwnedBytes};
|
||||
use common::{BinarySerializable, ByteCount};
|
||||
#[cfg(feature = "quickwit")]
|
||||
use futures_util::{FutureExt, StreamExt, TryStreamExt};
|
||||
#[cfg(feature = "quickwit")]
|
||||
@@ -10,13 +9,9 @@ use itertools::Itertools;
|
||||
#[cfg(feature = "quickwit")]
|
||||
use tantivy_fst::automaton::{AlwaysMatch, Automaton};
|
||||
|
||||
use crate::codec::postings::PostingsCodec;
|
||||
use crate::codec::{Codec, ObjectSafeCodec, StandardCodec};
|
||||
use crate::directory::FileSlice;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::{Postings, TermInfo};
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::{Bm25Weight, PhraseScorer, Scorer};
|
||||
use crate::positions::PositionReader;
|
||||
use crate::postings::{BlockSegmentPostings, SegmentPostings, TermInfo};
|
||||
use crate::schema::{IndexRecordOption, Term, Type};
|
||||
use crate::termdict::TermDictionary;
|
||||
|
||||
@@ -38,7 +33,6 @@ pub struct InvertedIndexReader {
|
||||
positions_file_slice: FileSlice,
|
||||
record_option: IndexRecordOption,
|
||||
total_num_tokens: u64,
|
||||
codec: Arc<dyn ObjectSafeCodec>,
|
||||
}
|
||||
|
||||
/// Object that records the amount of space used by a field in an inverted index.
|
||||
@@ -74,7 +68,6 @@ impl InvertedIndexReader {
|
||||
postings_file_slice: FileSlice,
|
||||
positions_file_slice: FileSlice,
|
||||
record_option: IndexRecordOption,
|
||||
codec: Arc<dyn ObjectSafeCodec>,
|
||||
) -> io::Result<InvertedIndexReader> {
|
||||
let (total_num_tokens_slice, postings_body) = postings_file_slice.split(8);
|
||||
let total_num_tokens = u64::deserialize(&mut total_num_tokens_slice.read_bytes()?)?;
|
||||
@@ -84,7 +77,6 @@ impl InvertedIndexReader {
|
||||
positions_file_slice,
|
||||
record_option,
|
||||
total_num_tokens,
|
||||
codec,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -97,7 +89,6 @@ impl InvertedIndexReader {
|
||||
positions_file_slice: FileSlice::empty(),
|
||||
record_option,
|
||||
total_num_tokens: 0u64,
|
||||
codec: Arc::new(StandardCodec),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -169,98 +160,61 @@ impl InvertedIndexReader {
|
||||
Ok(fields)
|
||||
}
|
||||
|
||||
pub(crate) fn new_term_scorer_specialized<C: Codec>(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
similarity_weight: Bm25Weight,
|
||||
codec: &C,
|
||||
) -> io::Result<TermScorer<<<C as Codec>::PostingsCodec as PostingsCodec>::Postings>> {
|
||||
let postings = self.read_postings_from_terminfo_specialized(term_info, option, codec)?;
|
||||
let term_scorer = TermScorer::new(postings, fieldnorm_reader, similarity_weight);
|
||||
Ok(term_scorer)
|
||||
}
|
||||
|
||||
pub(crate) fn new_phrase_scorer_type_specialized<C: Codec>(
|
||||
&self,
|
||||
term_infos: &[(usize, TermInfo)],
|
||||
similarity_weight_opt: Option<Bm25Weight>,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
slop: u32,
|
||||
codec: &C,
|
||||
) -> io::Result<PhraseScorer<<<C as Codec>::PostingsCodec as PostingsCodec>::Postings>> {
|
||||
let mut offset_and_term_postings: Vec<(
|
||||
usize,
|
||||
<<C as Codec>::PostingsCodec as PostingsCodec>::Postings,
|
||||
)> = Vec::with_capacity(term_infos.len());
|
||||
for (offset, term_info) in term_infos {
|
||||
let postings = self.read_postings_from_terminfo_specialized(
|
||||
term_info,
|
||||
IndexRecordOption::WithFreqsAndPositions,
|
||||
codec,
|
||||
)?;
|
||||
offset_and_term_postings.push((*offset, postings));
|
||||
}
|
||||
let phrase_scorer = PhraseScorer::new(
|
||||
offset_and_term_postings,
|
||||
similarity_weight_opt,
|
||||
fieldnorm_reader,
|
||||
slop,
|
||||
);
|
||||
Ok(phrase_scorer)
|
||||
}
|
||||
|
||||
/// Build a new term scorer.
|
||||
pub fn new_term_scorer(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
similarity_weight: Bm25Weight,
|
||||
) -> io::Result<Box<dyn Scorer>> {
|
||||
let term_scorer = self.codec.load_term_scorer_type_erased(
|
||||
term_info,
|
||||
option,
|
||||
self,
|
||||
fieldnorm_reader,
|
||||
similarity_weight,
|
||||
)?;
|
||||
Ok(term_scorer)
|
||||
}
|
||||
|
||||
/// Returns a postings object specific with a concrete type.
|
||||
/// Resets the block segment to another position of the postings
|
||||
/// file.
|
||||
///
|
||||
/// This requires you to provied the actual codec.
|
||||
pub fn read_postings_from_terminfo_specialized<C: Codec>(
|
||||
/// This is useful for enumerating through a list of terms,
|
||||
/// and consuming the associated posting lists while avoiding
|
||||
/// reallocating a [`BlockSegmentPostings`].
|
||||
///
|
||||
/// # Warning
|
||||
///
|
||||
/// This does not reset the positions list.
|
||||
pub fn reset_block_postings_from_terminfo(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
block_postings: &mut BlockSegmentPostings,
|
||||
) -> io::Result<()> {
|
||||
let postings_slice = self
|
||||
.postings_file_slice
|
||||
.slice(term_info.postings_range.clone());
|
||||
let postings_bytes = postings_slice.read_bytes()?;
|
||||
block_postings.reset(term_info.doc_freq, postings_bytes)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns a block postings given a `Term`.
|
||||
/// This method is for an advanced usage only.
|
||||
///
|
||||
/// Most users should prefer using [`Self::read_postings()`] instead.
|
||||
pub fn read_block_postings(
|
||||
&self,
|
||||
term: &Term,
|
||||
option: IndexRecordOption,
|
||||
codec: &C,
|
||||
) -> io::Result<<<C as Codec>::PostingsCodec as PostingsCodec>::Postings> {
|
||||
let option = option.downgrade(self.record_option);
|
||||
) -> io::Result<Option<BlockSegmentPostings>> {
|
||||
self.get_term_info(term)?
|
||||
.map(move |term_info| self.read_block_postings_from_terminfo(&term_info, option))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns a block postings given a `term_info`.
|
||||
/// This method is for an advanced usage only.
|
||||
///
|
||||
/// Most users should prefer using [`Self::read_postings()`] instead.
|
||||
pub fn read_block_postings_from_terminfo(
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
requested_option: IndexRecordOption,
|
||||
) -> io::Result<BlockSegmentPostings> {
|
||||
let postings_data = self
|
||||
.postings_file_slice
|
||||
.slice(term_info.postings_range.clone())
|
||||
.read_bytes()?;
|
||||
let positions_data: Option<OwnedBytes> = if option.has_positions() {
|
||||
let positions_data = self
|
||||
.positions_file_slice
|
||||
.slice(term_info.positions_range.clone())
|
||||
.read_bytes()?;
|
||||
Some(positions_data)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let postings: <<C as Codec>::PostingsCodec as PostingsCodec>::Postings =
|
||||
codec.postings_codec().load_postings(
|
||||
term_info.doc_freq,
|
||||
postings_data,
|
||||
self.record_option,
|
||||
option,
|
||||
positions_data,
|
||||
)?;
|
||||
Ok(postings)
|
||||
.slice(term_info.postings_range.clone());
|
||||
BlockSegmentPostings::open(
|
||||
term_info.doc_freq,
|
||||
postings_data,
|
||||
self.record_option,
|
||||
requested_option,
|
||||
)
|
||||
}
|
||||
|
||||
/// Returns a posting object given a `term_info`.
|
||||
@@ -271,9 +225,25 @@ impl InvertedIndexReader {
|
||||
&self,
|
||||
term_info: &TermInfo,
|
||||
option: IndexRecordOption,
|
||||
) -> io::Result<Box<dyn Postings>> {
|
||||
self.codec
|
||||
.load_postings_type_erased(term_info, option, self)
|
||||
) -> io::Result<SegmentPostings> {
|
||||
let option = option.downgrade(self.record_option);
|
||||
|
||||
let block_postings = self.read_block_postings_from_terminfo(term_info, option)?;
|
||||
let position_reader = {
|
||||
if option.has_positions() {
|
||||
let positions_data = self
|
||||
.positions_file_slice
|
||||
.read_bytes_slice(term_info.positions_range.clone())?;
|
||||
let position_reader = PositionReader::open(positions_data)?;
|
||||
Some(position_reader)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
Ok(SegmentPostings::from_block_postings(
|
||||
block_postings,
|
||||
position_reader,
|
||||
))
|
||||
}
|
||||
|
||||
/// Returns the total number of tokens recorded for all documents
|
||||
@@ -296,7 +266,7 @@ impl InvertedIndexReader {
|
||||
&self,
|
||||
term: &Term,
|
||||
option: IndexRecordOption,
|
||||
) -> io::Result<Option<Box<dyn Postings>>> {
|
||||
) -> io::Result<Option<SegmentPostings>> {
|
||||
self.get_term_info(term)?
|
||||
.map(move |term_info| self.read_postings_from_terminfo(&term_info, option))
|
||||
.transpose()
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
//!
|
||||
//! It contains `Index` and `Segment`, where a `Index` consists of one or more `Segment`s.
|
||||
|
||||
mod codec_configuration;
|
||||
mod index;
|
||||
mod index_meta;
|
||||
mod inverted_index_reader;
|
||||
@@ -11,7 +10,6 @@ mod segment_component;
|
||||
mod segment_id;
|
||||
mod segment_reader;
|
||||
|
||||
pub use self::codec_configuration::CodecConfiguration;
|
||||
pub use self::index::{Index, IndexBuilder};
|
||||
pub(crate) use self::index_meta::SegmentMetaInventory;
|
||||
pub use self::index_meta::{IndexMeta, IndexSettings, Order, SegmentMeta};
|
||||
|
||||
@@ -2,7 +2,6 @@ use std::fmt;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use super::SegmentComponent;
|
||||
use crate::codec::StandardCodec;
|
||||
use crate::directory::error::{OpenReadError, OpenWriteError};
|
||||
use crate::directory::{Directory, FileSlice, WritePtr};
|
||||
use crate::index::{Index, SegmentId, SegmentMeta};
|
||||
@@ -11,25 +10,25 @@ use crate::Opstamp;
|
||||
|
||||
/// A segment is a piece of the index.
|
||||
#[derive(Clone)]
|
||||
pub struct Segment<C: crate::codec::Codec = StandardCodec> {
|
||||
index: Index<C>,
|
||||
pub struct Segment {
|
||||
index: Index,
|
||||
meta: SegmentMeta,
|
||||
}
|
||||
|
||||
impl<C: crate::codec::Codec> fmt::Debug for Segment<C> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
impl fmt::Debug for Segment {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "Segment({:?})", self.id().uuid_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: crate::codec::Codec> Segment<C> {
|
||||
impl Segment {
|
||||
/// Creates a new segment given an `Index` and a `SegmentId`
|
||||
pub(crate) fn for_index(index: Index<C>, meta: SegmentMeta) -> Segment<C> {
|
||||
pub(crate) fn for_index(index: Index, meta: SegmentMeta) -> Segment {
|
||||
Segment { index, meta }
|
||||
}
|
||||
|
||||
/// Returns the index the segment belongs to.
|
||||
pub fn index(&self) -> &Index<C> {
|
||||
pub fn index(&self) -> &Index {
|
||||
&self.index
|
||||
}
|
||||
|
||||
@@ -47,7 +46,7 @@ impl<C: crate::codec::Codec> Segment<C> {
|
||||
///
|
||||
/// This method is only used when updating `max_doc` from 0
|
||||
/// as we finalize a fresh new segment.
|
||||
pub fn with_max_doc(self, max_doc: u32) -> Segment<C> {
|
||||
pub fn with_max_doc(self, max_doc: u32) -> Segment {
|
||||
Segment {
|
||||
index: self.index,
|
||||
meta: self.meta.with_max_doc(max_doc),
|
||||
@@ -56,7 +55,7 @@ impl<C: crate::codec::Codec> Segment<C> {
|
||||
|
||||
#[doc(hidden)]
|
||||
#[must_use]
|
||||
pub fn with_delete_meta(self, num_deleted_docs: u32, opstamp: Opstamp) -> Segment<C> {
|
||||
pub fn with_delete_meta(self, num_deleted_docs: u32, opstamp: Opstamp) -> Segment {
|
||||
Segment {
|
||||
index: self.index,
|
||||
meta: self.meta.with_delete_meta(num_deleted_docs, opstamp),
|
||||
|
||||
@@ -6,7 +6,6 @@ use common::{ByteCount, HasLen};
|
||||
use fnv::FnvHashMap;
|
||||
use itertools::Itertools;
|
||||
|
||||
use crate::codec::ObjectSafeCodec;
|
||||
use crate::directory::error::OpenReadError;
|
||||
use crate::directory::{CompositeFile, FileSlice};
|
||||
use crate::error::DataCorruption;
|
||||
@@ -49,7 +48,6 @@ pub struct SegmentReader {
|
||||
store_file: FileSlice,
|
||||
alive_bitset_opt: Option<AliveBitSet>,
|
||||
schema: Schema,
|
||||
codec: Arc<dyn ObjectSafeCodec>,
|
||||
}
|
||||
|
||||
impl SegmentReader {
|
||||
@@ -70,11 +68,6 @@ impl SegmentReader {
|
||||
&self.schema
|
||||
}
|
||||
|
||||
/// Returns the index codec.
|
||||
pub fn codec(&self) -> &dyn ObjectSafeCodec {
|
||||
&*self.codec
|
||||
}
|
||||
|
||||
/// Return the number of documents that have been
|
||||
/// deleted in the segment.
|
||||
pub fn num_deleted_docs(&self) -> DocId {
|
||||
@@ -148,16 +141,15 @@ impl SegmentReader {
|
||||
}
|
||||
|
||||
/// Open a new segment for reading.
|
||||
pub fn open<C: crate::codec::Codec>(segment: &Segment<C>) -> crate::Result<SegmentReader> {
|
||||
pub fn open(segment: &Segment) -> crate::Result<SegmentReader> {
|
||||
Self::open_with_custom_alive_set(segment, None)
|
||||
}
|
||||
|
||||
/// Open a new segment for reading.
|
||||
pub fn open_with_custom_alive_set<C: crate::codec::Codec>(
|
||||
segment: &Segment<C>,
|
||||
pub fn open_with_custom_alive_set(
|
||||
segment: &Segment,
|
||||
custom_bitset: Option<AliveBitSet>,
|
||||
) -> crate::Result<SegmentReader> {
|
||||
let codec: Arc<dyn ObjectSafeCodec> = Arc::new(segment.index().codec().clone());
|
||||
let termdict_file = segment.open_read(SegmentComponent::Terms)?;
|
||||
let termdict_composite = CompositeFile::open(&termdict_file)?;
|
||||
|
||||
@@ -211,7 +203,6 @@ impl SegmentReader {
|
||||
alive_bitset_opt,
|
||||
positions_composite,
|
||||
schema,
|
||||
codec,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -281,7 +272,6 @@ impl SegmentReader {
|
||||
postings_file,
|
||||
positions_file,
|
||||
record_option,
|
||||
self.codec.clone(),
|
||||
)?);
|
||||
|
||||
// by releasing the lock in between, we may end up opening the inverting index
|
||||
|
||||
@@ -9,7 +9,6 @@ use smallvec::smallvec;
|
||||
use super::operation::{AddOperation, UserOperation};
|
||||
use super::segment_updater::SegmentUpdater;
|
||||
use super::{AddBatch, AddBatchReceiver, AddBatchSender, PreparedCommit};
|
||||
use crate::codec::{Codec, StandardCodec};
|
||||
use crate::directory::{DirectoryLock, GarbageCollectionResult, TerminatingWrite};
|
||||
use crate::error::TantivyError;
|
||||
use crate::fastfield::write_alive_bitset;
|
||||
@@ -69,12 +68,12 @@ pub struct IndexWriterOptions {
|
||||
/// indexing queue.
|
||||
/// Each indexing thread builds its own independent [`Segment`], via
|
||||
/// a `SegmentWriter` object.
|
||||
pub struct IndexWriter<C: Codec = StandardCodec, D: Document = TantivyDocument> {
|
||||
pub struct IndexWriter<D: Document = TantivyDocument> {
|
||||
// the lock is just used to bind the
|
||||
// lifetime of the lock with that of the IndexWriter.
|
||||
_directory_lock: Option<DirectoryLock>,
|
||||
|
||||
index: Index<C>,
|
||||
index: Index,
|
||||
|
||||
options: IndexWriterOptions,
|
||||
|
||||
@@ -83,7 +82,7 @@ pub struct IndexWriter<C: Codec = StandardCodec, D: Document = TantivyDocument>
|
||||
index_writer_status: IndexWriterStatus<D>,
|
||||
operation_sender: AddBatchSender<D>,
|
||||
|
||||
segment_updater: SegmentUpdater<C>,
|
||||
segment_updater: SegmentUpdater,
|
||||
|
||||
worker_id: usize,
|
||||
|
||||
@@ -129,8 +128,8 @@ fn compute_deleted_bitset(
|
||||
/// is `==` target_opstamp.
|
||||
/// For instance, there was no delete operation between the state of the `segment_entry` and
|
||||
/// the `target_opstamp`, `segment_entry` is not updated.
|
||||
pub fn advance_deletes<C: Codec>(
|
||||
mut segment: Segment<C>,
|
||||
pub fn advance_deletes(
|
||||
mut segment: Segment,
|
||||
segment_entry: &mut SegmentEntry,
|
||||
target_opstamp: Opstamp,
|
||||
) -> crate::Result<()> {
|
||||
@@ -180,11 +179,11 @@ pub fn advance_deletes<C: Codec>(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn index_documents<C: crate::codec::Codec, D: Document>(
|
||||
fn index_documents<D: Document>(
|
||||
memory_budget: usize,
|
||||
segment: Segment<C>,
|
||||
segment: Segment,
|
||||
grouped_document_iterator: &mut dyn Iterator<Item = AddBatch<D>>,
|
||||
segment_updater: &SegmentUpdater<C>,
|
||||
segment_updater: &SegmentUpdater,
|
||||
mut delete_cursor: DeleteCursor,
|
||||
) -> crate::Result<()> {
|
||||
let mut segment_writer = SegmentWriter::for_segment(memory_budget, segment.clone())?;
|
||||
@@ -227,8 +226,8 @@ fn index_documents<C: crate::codec::Codec, D: Document>(
|
||||
}
|
||||
|
||||
/// `doc_opstamps` is required to be non-empty.
|
||||
fn apply_deletes<C: crate::codec::Codec>(
|
||||
segment: &Segment<C>,
|
||||
fn apply_deletes(
|
||||
segment: &Segment,
|
||||
delete_cursor: &mut DeleteCursor,
|
||||
doc_opstamps: &[Opstamp],
|
||||
) -> crate::Result<Option<BitSet>> {
|
||||
@@ -263,7 +262,7 @@ fn apply_deletes<C: crate::codec::Codec>(
|
||||
})
|
||||
}
|
||||
|
||||
impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
impl<D: Document> IndexWriter<D> {
|
||||
/// Create a new index writer. Attempts to acquire a lockfile.
|
||||
///
|
||||
/// The lockfile should be deleted on drop, but it is possible
|
||||
@@ -279,7 +278,7 @@ impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
/// If the memory arena per thread is too small or too big, returns
|
||||
/// `TantivyError::InvalidArgument`
|
||||
pub(crate) fn new(
|
||||
index: &Index<C>,
|
||||
index: &Index,
|
||||
options: IndexWriterOptions,
|
||||
directory_lock: DirectoryLock,
|
||||
) -> crate::Result<Self> {
|
||||
@@ -346,7 +345,7 @@ impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
}
|
||||
|
||||
/// Accessor to the index.
|
||||
pub fn index(&self) -> &Index<C> {
|
||||
pub fn index(&self) -> &Index {
|
||||
&self.index
|
||||
}
|
||||
|
||||
@@ -394,7 +393,7 @@ impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
/// It is safe to start writing file associated with the new `Segment`.
|
||||
/// These will not be garbage collected as long as an instance object of
|
||||
/// `SegmentMeta` object associated with the new `Segment` is "alive".
|
||||
pub fn new_segment(&self) -> Segment<C> {
|
||||
pub fn new_segment(&self) -> Segment {
|
||||
self.index.new_segment()
|
||||
}
|
||||
|
||||
@@ -616,7 +615,7 @@ impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
/// It is also possible to add a payload to the `commit`
|
||||
/// using this API.
|
||||
/// See [`PreparedCommit::set_payload()`].
|
||||
pub fn prepare_commit(&mut self) -> crate::Result<PreparedCommit<'_, C, D>> {
|
||||
pub fn prepare_commit(&mut self) -> crate::Result<PreparedCommit<'_, D>> {
|
||||
// Here, because we join all of the worker threads,
|
||||
// all of the segment update for this commit have been
|
||||
// sent.
|
||||
@@ -666,7 +665,7 @@ impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
self.prepare_commit()?.commit()
|
||||
}
|
||||
|
||||
pub(crate) fn segment_updater(&self) -> &SegmentUpdater<C> {
|
||||
pub(crate) fn segment_updater(&self) -> &SegmentUpdater {
|
||||
&self.segment_updater
|
||||
}
|
||||
|
||||
@@ -805,7 +804,7 @@ impl<C: Codec, D: Document> IndexWriter<C, D> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: Codec, D: Document> Drop for IndexWriter<C, D> {
|
||||
impl<D: Document> Drop for IndexWriter<D> {
|
||||
fn drop(&mut self) {
|
||||
self.segment_updater.kill();
|
||||
self.drop_sender();
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::codec::StandardCodec;
|
||||
use crate::collector::TopDocs;
|
||||
use crate::fastfield::AliveBitSet;
|
||||
use crate::index::Index;
|
||||
use crate::postings::{DocFreq, Postings};
|
||||
use crate::postings::Postings;
|
||||
use crate::query::QueryParser;
|
||||
use crate::schema::{
|
||||
self, BytesOptions, Facet, FacetOptions, IndexRecordOption, NumericOptions,
|
||||
@@ -122,26 +121,21 @@ mod tests {
|
||||
let my_text_field = index.schema().get_field("text_field").unwrap();
|
||||
let term_a = Term::from_field_text(my_text_field, "text");
|
||||
let inverted_index = segment_reader.inverted_index(my_text_field).unwrap();
|
||||
let term_info = inverted_index.get_term_info(&term_a).unwrap().unwrap();
|
||||
let mut postings = inverted_index
|
||||
.read_postings_from_terminfo_specialized(
|
||||
&term_info,
|
||||
IndexRecordOption::WithFreqsAndPositions,
|
||||
&StandardCodec,
|
||||
)
|
||||
.read_postings(&term_a, IndexRecordOption::WithFreqsAndPositions)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(postings.doc_freq(), DocFreq::Exact(2));
|
||||
assert_eq!(postings.doc_freq(), 2);
|
||||
let fallback_bitset = AliveBitSet::for_test_from_deleted_docs(&[0], 100);
|
||||
assert_eq!(
|
||||
crate::indexer::merger::doc_freq_given_deletes(
|
||||
&postings,
|
||||
postings.doc_freq_given_deletes(
|
||||
segment_reader.alive_bitset().unwrap_or(&fallback_bitset)
|
||||
),
|
||||
2
|
||||
);
|
||||
|
||||
assert_eq!(postings.term_freq(), 1);
|
||||
let mut output = Vec::new();
|
||||
let mut output = vec![];
|
||||
postings.positions(&mut output);
|
||||
assert_eq!(output, vec![1]);
|
||||
postings.advance();
|
||||
|
||||
@@ -7,8 +7,6 @@ use common::ReadOnlyBitSet;
|
||||
use itertools::Itertools;
|
||||
use measure_time::debug_time;
|
||||
|
||||
use crate::codec::postings::PostingsCodec;
|
||||
use crate::codec::{Codec, StandardCodec};
|
||||
use crate::directory::WritePtr;
|
||||
use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::error::DataCorruption;
|
||||
@@ -17,7 +15,7 @@ use crate::fieldnorm::{FieldNormReader, FieldNormReaders, FieldNormsSerializer,
|
||||
use crate::index::{Segment, SegmentComponent, SegmentReader};
|
||||
use crate::indexer::doc_id_mapping::{MappingType, SegmentDocIdMapping};
|
||||
use crate::indexer::SegmentSerializer;
|
||||
use crate::postings::{InvertedIndexSerializer, Postings};
|
||||
use crate::postings::{InvertedIndexSerializer, Postings, SegmentPostings};
|
||||
use crate::schema::{value_type_to_column_type, Field, FieldType, Schema};
|
||||
use crate::store::StoreWriter;
|
||||
use crate::termdict::{TermMerger, TermOrdinal};
|
||||
@@ -78,11 +76,10 @@ fn estimate_total_num_tokens(readers: &[SegmentReader], field: Field) -> crate::
|
||||
Ok(total_num_tokens)
|
||||
}
|
||||
|
||||
pub struct IndexMerger<C: Codec = StandardCodec> {
|
||||
pub struct IndexMerger {
|
||||
schema: Schema,
|
||||
pub(crate) readers: Vec<SegmentReader>,
|
||||
max_doc: u32,
|
||||
codec: C,
|
||||
}
|
||||
|
||||
struct DeltaComputer {
|
||||
@@ -147,8 +144,8 @@ fn extract_fast_field_required_columns(schema: &Schema) -> Vec<(String, ColumnTy
|
||||
.collect()
|
||||
}
|
||||
|
||||
impl<C: Codec> IndexMerger<C> {
|
||||
pub fn open(schema: Schema, segments: &[Segment<C>]) -> crate::Result<IndexMerger<C>> {
|
||||
impl IndexMerger {
|
||||
pub fn open(schema: Schema, segments: &[Segment]) -> crate::Result<IndexMerger> {
|
||||
let alive_bitset = segments.iter().map(|_| None).collect_vec();
|
||||
Self::open_with_custom_alive_set(schema, segments, alive_bitset)
|
||||
}
|
||||
@@ -165,15 +162,11 @@ impl<C: Codec> IndexMerger<C> {
|
||||
// This can be used to merge but also apply an additional filter.
|
||||
// One use case is demux, which is basically taking a list of
|
||||
// segments and partitions them e.g. by a value in a field.
|
||||
//
|
||||
// # Panics if segments is empty.
|
||||
pub fn open_with_custom_alive_set(
|
||||
schema: Schema,
|
||||
segments: &[Segment<C>],
|
||||
segments: &[Segment],
|
||||
alive_bitset_opt: Vec<Option<AliveBitSet>>,
|
||||
) -> crate::Result<IndexMerger<C>> {
|
||||
assert!(!segments.is_empty());
|
||||
let codec = segments[0].index().codec().clone();
|
||||
) -> crate::Result<IndexMerger> {
|
||||
let mut readers = vec![];
|
||||
for (segment, new_alive_bitset_opt) in segments.iter().zip(alive_bitset_opt) {
|
||||
if segment.meta().num_docs() > 0 {
|
||||
@@ -196,7 +189,6 @@ impl<C: Codec> IndexMerger<C> {
|
||||
schema,
|
||||
readers,
|
||||
max_doc,
|
||||
codec,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -295,7 +287,7 @@ impl<C: Codec> IndexMerger<C> {
|
||||
&self,
|
||||
indexed_field: Field,
|
||||
_field_type: &FieldType,
|
||||
serializer: &mut InvertedIndexSerializer<C>,
|
||||
serializer: &mut InvertedIndexSerializer,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
doc_id_mapping: &SegmentDocIdMapping,
|
||||
) -> crate::Result<()> {
|
||||
@@ -363,10 +355,7 @@ impl<C: Codec> IndexMerger<C> {
|
||||
indexed. Have you modified the schema?",
|
||||
);
|
||||
|
||||
let mut segment_postings_containing_the_term: Vec<(
|
||||
usize,
|
||||
<C::PostingsCodec as PostingsCodec>::Postings,
|
||||
)> = Vec::with_capacity(self.readers.len());
|
||||
let mut segment_postings_containing_the_term: Vec<(usize, SegmentPostings)> = vec![];
|
||||
|
||||
while merged_terms.advance() {
|
||||
segment_postings_containing_the_term.clear();
|
||||
@@ -378,24 +367,17 @@ impl<C: Codec> IndexMerger<C> {
|
||||
for (segment_ord, term_info) in merged_terms.current_segment_ords_and_term_infos() {
|
||||
let segment_reader = &self.readers[segment_ord];
|
||||
let inverted_index: &InvertedIndexReader = &field_readers[segment_ord];
|
||||
let postings = inverted_index.read_postings_from_terminfo_specialized(
|
||||
&term_info,
|
||||
segment_postings_option,
|
||||
&self.codec,
|
||||
)?;
|
||||
let segment_postings = inverted_index
|
||||
.read_postings_from_terminfo(&term_info, segment_postings_option)?;
|
||||
let alive_bitset_opt = segment_reader.alive_bitset();
|
||||
let doc_freq = if let Some(alive_bitset) = alive_bitset_opt {
|
||||
doc_freq_given_deletes(&postings, alive_bitset)
|
||||
segment_postings.doc_freq_given_deletes(alive_bitset)
|
||||
} else {
|
||||
// We do not an exact document frequency here.
|
||||
match postings.doc_freq() {
|
||||
crate::postings::DocFreq::Approximate(_) => exact_doc_freq(&postings),
|
||||
crate::postings::DocFreq::Exact(doc_freq) => doc_freq,
|
||||
}
|
||||
segment_postings.doc_freq()
|
||||
};
|
||||
if doc_freq > 0u32 {
|
||||
total_doc_freq += doc_freq;
|
||||
segment_postings_containing_the_term.push((segment_ord, postings));
|
||||
segment_postings_containing_the_term.push((segment_ord, segment_postings));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -413,7 +395,11 @@ impl<C: Codec> IndexMerger<C> {
|
||||
assert!(!segment_postings_containing_the_term.is_empty());
|
||||
|
||||
let has_term_freq = {
|
||||
let has_term_freq = segment_postings_containing_the_term[0].1.has_freq();
|
||||
let has_term_freq = !segment_postings_containing_the_term[0]
|
||||
.1
|
||||
.block_cursor
|
||||
.freqs()
|
||||
.is_empty();
|
||||
for (_, postings) in &segment_postings_containing_the_term[1..] {
|
||||
// This may look at a strange way to test whether we have term freq or not.
|
||||
// With JSON object, the schema is not sufficient to know whether a term
|
||||
@@ -429,7 +415,7 @@ impl<C: Codec> IndexMerger<C> {
|
||||
//
|
||||
// Overall the reliable way to know if we have actual frequencies loaded or not
|
||||
// is to check whether the actual decoded array is empty or not.
|
||||
if postings.has_freq() != has_term_freq {
|
||||
if has_term_freq == postings.block_cursor.freqs().is_empty() {
|
||||
return Err(DataCorruption::comment_only(
|
||||
"Term freqs are inconsistent across segments",
|
||||
)
|
||||
@@ -481,7 +467,7 @@ impl<C: Codec> IndexMerger<C> {
|
||||
|
||||
fn write_postings(
|
||||
&self,
|
||||
serializer: &mut InvertedIndexSerializer<C>,
|
||||
serializer: &mut InvertedIndexSerializer,
|
||||
fieldnorm_readers: FieldNormReaders,
|
||||
doc_id_mapping: &SegmentDocIdMapping,
|
||||
) -> crate::Result<()> {
|
||||
@@ -539,7 +525,7 @@ impl<C: Codec> IndexMerger<C> {
|
||||
///
|
||||
/// # Returns
|
||||
/// The number of documents in the resulting segment.
|
||||
pub fn write(&self, mut serializer: SegmentSerializer<C>) -> crate::Result<u32> {
|
||||
pub fn write(&self, mut serializer: SegmentSerializer) -> crate::Result<u32> {
|
||||
let doc_id_mapping = self.get_doc_id_from_concatenated_data()?;
|
||||
debug!("write-fieldnorms");
|
||||
if let Some(fieldnorms_serializer) = serializer.extract_fieldnorms_serializer() {
|
||||
@@ -567,43 +553,6 @@ impl<C: Codec> IndexMerger<C> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the number of non-deleted documents.
|
||||
///
|
||||
/// This method will clone and scan through the posting lists.
|
||||
/// (this is a rather expensive operation).
|
||||
pub(crate) fn doc_freq_given_deletes<P: Postings + Clone>(
|
||||
postings: &P,
|
||||
alive_bitset: &AliveBitSet,
|
||||
) -> u32 {
|
||||
let mut docset = postings.clone();
|
||||
let mut doc_freq = 0;
|
||||
loop {
|
||||
let doc = docset.doc();
|
||||
if doc == TERMINATED {
|
||||
return doc_freq;
|
||||
}
|
||||
if alive_bitset.is_alive(doc) {
|
||||
doc_freq += 1u32;
|
||||
}
|
||||
docset.advance();
|
||||
}
|
||||
}
|
||||
|
||||
/// If the postings is not able to inform us of the document frequency,
|
||||
/// we just scan through it.
|
||||
pub(crate) fn exact_doc_freq<P: Postings + Clone>(postings: &P) -> u32 {
|
||||
let mut docset = postings.clone();
|
||||
let mut doc_freq = 0;
|
||||
loop {
|
||||
let doc = docset.doc();
|
||||
if doc == TERMINATED {
|
||||
return doc_freq;
|
||||
}
|
||||
doc_freq += 1u32;
|
||||
docset.advance();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
@@ -612,16 +561,12 @@ mod tests {
|
||||
use proptest::strategy::Strategy;
|
||||
use schema::FAST;
|
||||
|
||||
use crate::codec::postings::PostingsCodec;
|
||||
use crate::codec::standard::postings::StandardPostingsCodec;
|
||||
use crate::collector::tests::{
|
||||
BytesFastFieldTestCollector, FastFieldTestCollector, TEST_COLLECTOR_WITH_SCORE,
|
||||
};
|
||||
use crate::collector::{Count, FacetCollector};
|
||||
use crate::fastfield::AliveBitSet;
|
||||
use crate::index::{Index, SegmentId};
|
||||
use crate::indexer::NoMergePolicy;
|
||||
use crate::postings::{DocFreq, Postings as _};
|
||||
use crate::query::{AllQuery, BooleanQuery, EnableScoring, Scorer, TermQuery};
|
||||
use crate::schema::{
|
||||
Facet, FacetOptions, IndexRecordOption, NumericOptions, TantivyDocument, Term,
|
||||
@@ -1573,10 +1518,10 @@ mod tests {
|
||||
let searcher = reader.searcher();
|
||||
let mut term_scorer = term_query
|
||||
.specialized_weight(EnableScoring::enabled_from_searcher(&searcher))?
|
||||
.term_scorer_for_test(searcher.segment_reader(0u32), 1.0)
|
||||
.term_scorer_for_test(searcher.segment_reader(0u32), 1.0)?
|
||||
.unwrap();
|
||||
assert_eq!(term_scorer.doc(), 0);
|
||||
assert_nearly_equals!(term_scorer.seek_block_max(0), 0.0079681855);
|
||||
assert_nearly_equals!(term_scorer.block_max_score(), 0.0079681855);
|
||||
assert_nearly_equals!(term_scorer.score(), 0.0079681855);
|
||||
for _ in 0..81 {
|
||||
writer.add_document(doc!(text=>"hello happy tax payer"))?;
|
||||
@@ -1589,13 +1534,13 @@ mod tests {
|
||||
for segment_reader in searcher.segment_readers() {
|
||||
let mut term_scorer = term_query
|
||||
.specialized_weight(EnableScoring::enabled_from_searcher(&searcher))?
|
||||
.term_scorer_for_test(segment_reader, 1.0)
|
||||
.term_scorer_for_test(segment_reader, 1.0)?
|
||||
.unwrap();
|
||||
// the difference compared to before is intrinsic to the bm25 formula. no worries
|
||||
// there.
|
||||
for doc in segment_reader.doc_ids_alive() {
|
||||
assert_eq!(term_scorer.doc(), doc);
|
||||
assert_nearly_equals!(term_scorer.seek_block_max(doc), 0.003478312);
|
||||
assert_nearly_equals!(term_scorer.block_max_score(), 0.003478312);
|
||||
assert_nearly_equals!(term_scorer.score(), 0.003478312);
|
||||
term_scorer.advance();
|
||||
}
|
||||
@@ -1615,12 +1560,12 @@ mod tests {
|
||||
let segment_reader = searcher.segment_reader(0u32);
|
||||
let mut term_scorer = term_query
|
||||
.specialized_weight(EnableScoring::enabled_from_searcher(&searcher))?
|
||||
.term_scorer_for_test(segment_reader, 1.0)
|
||||
.term_scorer_for_test(segment_reader, 1.0)?
|
||||
.unwrap();
|
||||
// the difference compared to before is intrinsic to the bm25 formula. no worries there.
|
||||
for doc in segment_reader.doc_ids_alive() {
|
||||
assert_eq!(term_scorer.doc(), doc);
|
||||
assert_nearly_equals!(term_scorer.seek_block_max(doc), 0.003478312);
|
||||
assert_nearly_equals!(term_scorer.block_max_score(), 0.003478312);
|
||||
assert_nearly_equals!(term_scorer.score(), 0.003478312);
|
||||
term_scorer.advance();
|
||||
}
|
||||
@@ -1634,16 +1579,4 @@ mod tests {
|
||||
assert!(((super::MAX_DOC_LIMIT - 1) as i32) >= 0);
|
||||
assert!((super::MAX_DOC_LIMIT as i32) < 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_doc_freq_given_delete() {
|
||||
let docs =
|
||||
<StandardPostingsCodec as PostingsCodec>::Postings::create_from_docs(&[0, 2, 10]);
|
||||
assert_eq!(docs.doc_freq(), DocFreq::Exact(3));
|
||||
let alive_bitset = AliveBitSet::for_test_from_deleted_docs(&[2], 12);
|
||||
assert_eq!(super::doc_freq_given_deletes(&docs, &alive_bitset), 2);
|
||||
let all_deleted =
|
||||
AliveBitSet::for_test_from_deleted_docs(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 12);
|
||||
assert_eq!(super::doc_freq_given_deletes(&docs, &all_deleted), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,17 +1,16 @@
|
||||
use super::IndexWriter;
|
||||
use crate::codec::Codec;
|
||||
use crate::schema::document::Document;
|
||||
use crate::{FutureResult, Opstamp, TantivyDocument};
|
||||
|
||||
/// A prepared commit
|
||||
pub struct PreparedCommit<'a, C: Codec, D: Document = TantivyDocument> {
|
||||
index_writer: &'a mut IndexWriter<C, D>,
|
||||
pub struct PreparedCommit<'a, D: Document = TantivyDocument> {
|
||||
index_writer: &'a mut IndexWriter<D>,
|
||||
payload: Option<String>,
|
||||
opstamp: Opstamp,
|
||||
}
|
||||
|
||||
impl<'a, C: Codec, D: Document> PreparedCommit<'a, C, D> {
|
||||
pub(crate) fn new(index_writer: &'a mut IndexWriter<C, D>, opstamp: Opstamp) -> Self {
|
||||
impl<'a, D: Document> PreparedCommit<'a, D> {
|
||||
pub(crate) fn new(index_writer: &'a mut IndexWriter<D>, opstamp: Opstamp) -> Self {
|
||||
Self {
|
||||
index_writer,
|
||||
payload: None,
|
||||
|
||||
@@ -8,17 +8,17 @@ use crate::store::StoreWriter;
|
||||
|
||||
/// Segment serializer is in charge of laying out on disk
|
||||
/// the data accumulated and sorted by the `SegmentWriter`.
|
||||
pub struct SegmentSerializer<C: crate::codec::Codec> {
|
||||
segment: Segment<C>,
|
||||
pub struct SegmentSerializer {
|
||||
segment: Segment,
|
||||
pub(crate) store_writer: StoreWriter,
|
||||
fast_field_write: WritePtr,
|
||||
fieldnorms_serializer: Option<FieldNormsSerializer>,
|
||||
postings_serializer: InvertedIndexSerializer<C>,
|
||||
postings_serializer: InvertedIndexSerializer,
|
||||
}
|
||||
|
||||
impl<C: crate::codec::Codec> SegmentSerializer<C> {
|
||||
impl SegmentSerializer {
|
||||
/// Creates a new `SegmentSerializer`.
|
||||
pub fn for_segment(mut segment: Segment<C>) -> crate::Result<SegmentSerializer<C>> {
|
||||
pub fn for_segment(mut segment: Segment) -> crate::Result<SegmentSerializer> {
|
||||
let settings = segment.index().settings().clone();
|
||||
let store_writer = {
|
||||
let store_write = segment.open_write(SegmentComponent::Store)?;
|
||||
@@ -50,12 +50,12 @@ impl<C: crate::codec::Codec> SegmentSerializer<C> {
|
||||
self.store_writer.mem_usage()
|
||||
}
|
||||
|
||||
pub fn segment(&self) -> &Segment<C> {
|
||||
pub fn segment(&self) -> &Segment {
|
||||
&self.segment
|
||||
}
|
||||
|
||||
/// Accessor to the `PostingsSerializer`.
|
||||
pub fn get_postings_serializer(&mut self) -> &mut InvertedIndexSerializer<C> {
|
||||
pub fn get_postings_serializer(&mut self) -> &mut InvertedIndexSerializer {
|
||||
&mut self.postings_serializer
|
||||
}
|
||||
|
||||
|
||||
@@ -10,13 +10,10 @@ use std::sync::{Arc, RwLock};
|
||||
use rayon::{ThreadPool, ThreadPoolBuilder};
|
||||
|
||||
use super::segment_manager::SegmentManager;
|
||||
use crate::codec::Codec;
|
||||
use crate::core::META_FILEPATH;
|
||||
use crate::directory::{Directory, DirectoryClone, GarbageCollectionResult};
|
||||
use crate::fastfield::AliveBitSet;
|
||||
use crate::index::{
|
||||
CodecConfiguration, Index, IndexMeta, IndexSettings, Segment, SegmentId, SegmentMeta,
|
||||
};
|
||||
use crate::index::{Index, IndexMeta, IndexSettings, Segment, SegmentId, SegmentMeta};
|
||||
use crate::indexer::delete_queue::DeleteCursor;
|
||||
use crate::indexer::index_writer::advance_deletes;
|
||||
use crate::indexer::merge_operation::MergeOperationInventory;
|
||||
@@ -64,10 +61,10 @@ pub(crate) fn save_metas(metas: &IndexMeta, directory: &dyn Directory) -> crate:
|
||||
// We voluntarily pass a merge_operation ref to guarantee that
|
||||
// the merge_operation is alive during the process
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SegmentUpdater<C: Codec>(Arc<InnerSegmentUpdater<C>>);
|
||||
pub(crate) struct SegmentUpdater(Arc<InnerSegmentUpdater>);
|
||||
|
||||
impl<C: Codec> Deref for SegmentUpdater<C> {
|
||||
type Target = InnerSegmentUpdater<C>;
|
||||
impl Deref for SegmentUpdater {
|
||||
type Target = InnerSegmentUpdater;
|
||||
|
||||
#[inline]
|
||||
fn deref(&self) -> &Self::Target {
|
||||
@@ -75,8 +72,8 @@ impl<C: Codec> Deref for SegmentUpdater<C> {
|
||||
}
|
||||
}
|
||||
|
||||
fn garbage_collect_files<C: Codec>(
|
||||
segment_updater: SegmentUpdater<C>,
|
||||
fn garbage_collect_files(
|
||||
segment_updater: SegmentUpdater,
|
||||
) -> crate::Result<GarbageCollectionResult> {
|
||||
info!("Running garbage collection");
|
||||
let mut index = segment_updater.index.clone();
|
||||
@@ -87,8 +84,8 @@ fn garbage_collect_files<C: Codec>(
|
||||
|
||||
/// Merges a list of segments the list of segment givens in the `segment_entries`.
|
||||
/// This function happens in the calling thread and is computationally expensive.
|
||||
fn merge<Codec: crate::codec::Codec>(
|
||||
index: &Index<Codec>,
|
||||
fn merge(
|
||||
index: &Index,
|
||||
mut segment_entries: Vec<SegmentEntry>,
|
||||
target_opstamp: Opstamp,
|
||||
) -> crate::Result<Option<SegmentEntry>> {
|
||||
@@ -111,13 +108,13 @@ fn merge<Codec: crate::codec::Codec>(
|
||||
|
||||
let delete_cursor = segment_entries[0].delete_cursor().clone();
|
||||
|
||||
let segments: Vec<Segment<Codec>> = segment_entries
|
||||
let segments: Vec<Segment> = segment_entries
|
||||
.iter()
|
||||
.map(|segment_entry| index.segment(segment_entry.meta().clone()))
|
||||
.collect();
|
||||
|
||||
// An IndexMerger is like a "view" of our merged segments.
|
||||
let merger: IndexMerger<Codec> = IndexMerger::open(index.schema(), &segments[..])?;
|
||||
let merger: IndexMerger = IndexMerger::open(index.schema(), &segments[..])?;
|
||||
|
||||
// ... we just serialize this index merger in our new segment to merge the segments.
|
||||
let segment_serializer = SegmentSerializer::for_segment(merged_segment.clone())?;
|
||||
@@ -142,10 +139,10 @@ fn merge<Codec: crate::codec::Codec>(
|
||||
/// meant to work if you have an `IndexWriter` running for the origin indices, or
|
||||
/// the destination `Index`.
|
||||
#[doc(hidden)]
|
||||
pub fn merge_indices<Codec: crate::codec::Codec>(
|
||||
indices: &[Index<Codec>],
|
||||
output_directory: Box<dyn Directory>,
|
||||
) -> crate::Result<Index<Codec>> {
|
||||
pub fn merge_indices<T: Into<Box<dyn Directory>>>(
|
||||
indices: &[Index],
|
||||
output_directory: T,
|
||||
) -> crate::Result<Index> {
|
||||
if indices.is_empty() {
|
||||
// If there are no indices to merge, there is no need to do anything.
|
||||
return Err(crate::TantivyError::InvalidArgument(
|
||||
@@ -166,7 +163,7 @@ pub fn merge_indices<Codec: crate::codec::Codec>(
|
||||
));
|
||||
}
|
||||
|
||||
let mut segments: Vec<Segment<Codec>> = Vec::new();
|
||||
let mut segments: Vec<Segment> = Vec::new();
|
||||
for index in indices {
|
||||
segments.extend(index.searchable_segments()?);
|
||||
}
|
||||
@@ -188,12 +185,12 @@ pub fn merge_indices<Codec: crate::codec::Codec>(
|
||||
/// meant to work if you have an `IndexWriter` running for the origin indices, or
|
||||
/// the destination `Index`.
|
||||
#[doc(hidden)]
|
||||
pub fn merge_filtered_segments<C: crate::codec::Codec, T: Into<Box<dyn Directory>>>(
|
||||
segments: &[Segment<C>],
|
||||
pub fn merge_filtered_segments<T: Into<Box<dyn Directory>>>(
|
||||
segments: &[Segment],
|
||||
target_settings: IndexSettings,
|
||||
filter_doc_ids: Vec<Option<AliveBitSet>>,
|
||||
output_directory: T,
|
||||
) -> crate::Result<Index<C>> {
|
||||
) -> crate::Result<Index> {
|
||||
if segments.is_empty() {
|
||||
// If there are no indices to merge, there is no need to do anything.
|
||||
return Err(crate::TantivyError::InvalidArgument(
|
||||
@@ -214,15 +211,14 @@ pub fn merge_filtered_segments<C: crate::codec::Codec, T: Into<Box<dyn Directory
|
||||
));
|
||||
}
|
||||
|
||||
let mut merged_index: Index<C> = Index::builder()
|
||||
.schema(target_schema.clone())
|
||||
.codec(segments[0].index().codec().clone())
|
||||
.settings(target_settings.clone())
|
||||
.create(output_directory.into())?;
|
||||
|
||||
let mut merged_index = Index::create(
|
||||
output_directory,
|
||||
target_schema.clone(),
|
||||
target_settings.clone(),
|
||||
)?;
|
||||
let merged_segment = merged_index.new_segment();
|
||||
let merged_segment_id = merged_segment.id();
|
||||
let merger: IndexMerger<C> =
|
||||
let merger: IndexMerger =
|
||||
IndexMerger::open_with_custom_alive_set(merged_index.schema(), segments, filter_doc_ids)?;
|
||||
let segment_serializer = SegmentSerializer::for_segment(merged_segment)?;
|
||||
let num_docs = merger.write(segment_serializer)?;
|
||||
@@ -239,7 +235,6 @@ pub fn merge_filtered_segments<C: crate::codec::Codec, T: Into<Box<dyn Directory
|
||||
))
|
||||
.trim_end()
|
||||
);
|
||||
let codec_configuration = CodecConfiguration::from(segments[0].index().codec());
|
||||
|
||||
let index_meta = IndexMeta {
|
||||
index_settings: target_settings, // index_settings of all segments should be the same
|
||||
@@ -247,7 +242,6 @@ pub fn merge_filtered_segments<C: crate::codec::Codec, T: Into<Box<dyn Directory
|
||||
schema: target_schema,
|
||||
opstamp: 0u64,
|
||||
payload: Some(stats),
|
||||
codec: codec_configuration,
|
||||
};
|
||||
|
||||
// save the meta.json
|
||||
@@ -256,7 +250,7 @@ pub fn merge_filtered_segments<C: crate::codec::Codec, T: Into<Box<dyn Directory
|
||||
Ok(merged_index)
|
||||
}
|
||||
|
||||
pub(crate) struct InnerSegmentUpdater<C: Codec> {
|
||||
pub(crate) struct InnerSegmentUpdater {
|
||||
// we keep a copy of the current active IndexMeta to
|
||||
// avoid loading the file every time we need it in the
|
||||
// `SegmentUpdater`.
|
||||
@@ -267,7 +261,7 @@ pub(crate) struct InnerSegmentUpdater<C: Codec> {
|
||||
pool: ThreadPool,
|
||||
merge_thread_pool: ThreadPool,
|
||||
|
||||
index: Index<C>,
|
||||
index: Index,
|
||||
segment_manager: SegmentManager,
|
||||
merge_policy: RwLock<Arc<dyn MergePolicy>>,
|
||||
killed: AtomicBool,
|
||||
@@ -275,13 +269,13 @@ pub(crate) struct InnerSegmentUpdater<C: Codec> {
|
||||
merge_operations: MergeOperationInventory,
|
||||
}
|
||||
|
||||
impl<Codec: crate::codec::Codec> SegmentUpdater<Codec> {
|
||||
impl SegmentUpdater {
|
||||
pub fn create(
|
||||
index: Index<Codec>,
|
||||
index: Index,
|
||||
stamper: Stamper,
|
||||
delete_cursor: &DeleteCursor,
|
||||
num_merge_threads: usize,
|
||||
) -> crate::Result<Self> {
|
||||
) -> crate::Result<SegmentUpdater> {
|
||||
let segments = index.searchable_segment_metas()?;
|
||||
let segment_manager = SegmentManager::from_segments(segments, delete_cursor);
|
||||
let pool = ThreadPoolBuilder::new()
|
||||
@@ -411,14 +405,12 @@ impl<Codec: crate::codec::Codec> SegmentUpdater<Codec> {
|
||||
// Segment 1 from disk 1, Segment 1 from disk 2, etc.
|
||||
committed_segment_metas
|
||||
.sort_by_key(|segment_meta| std::cmp::Reverse(segment_meta.max_doc()));
|
||||
let codec = CodecConfiguration::from(index.codec());
|
||||
let index_meta = IndexMeta {
|
||||
index_settings: index.settings().clone(),
|
||||
segments: committed_segment_metas,
|
||||
schema: index.schema(),
|
||||
opstamp,
|
||||
payload: commit_message,
|
||||
codec,
|
||||
};
|
||||
// TODO add context to the error.
|
||||
save_metas(&index_meta, directory.box_clone().borrow_mut())?;
|
||||
@@ -452,7 +444,7 @@ impl<Codec: crate::codec::Codec> SegmentUpdater<Codec> {
|
||||
opstamp: Opstamp,
|
||||
payload: Option<String>,
|
||||
) -> FutureResult<Opstamp> {
|
||||
let segment_updater: SegmentUpdater<Codec> = self.clone();
|
||||
let segment_updater: SegmentUpdater = self.clone();
|
||||
self.schedule_task(move || {
|
||||
let segment_entries = segment_updater.purge_deletes(opstamp)?;
|
||||
segment_updater.segment_manager.commit(segment_entries);
|
||||
@@ -708,7 +700,6 @@ impl<Codec: crate::codec::Codec> SegmentUpdater<Codec> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::merge_indices;
|
||||
use crate::codec::StandardCodec;
|
||||
use crate::collector::TopDocs;
|
||||
use crate::directory::RamDirectory;
|
||||
use crate::fastfield::AliveBitSet;
|
||||
@@ -939,7 +930,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_merge_empty_indices_array() {
|
||||
let merge_result = merge_indices::<StandardCodec>(&[], Box::new(RamDirectory::default()));
|
||||
let merge_result = merge_indices(&[], RamDirectory::default());
|
||||
assert!(merge_result.is_err());
|
||||
}
|
||||
|
||||
@@ -966,10 +957,7 @@ mod tests {
|
||||
};
|
||||
|
||||
// mismatched schema index list
|
||||
let result = merge_indices(
|
||||
&[first_index, second_index],
|
||||
Box::new(RamDirectory::default()),
|
||||
);
|
||||
let result = merge_indices(&[first_index, second_index], RamDirectory::default());
|
||||
assert!(result.is_err());
|
||||
|
||||
Ok(())
|
||||
|
||||
@@ -4,7 +4,6 @@ use itertools::Itertools;
|
||||
use tokenizer_api::BoxTokenStream;
|
||||
|
||||
use super::operation::AddOperation;
|
||||
use crate::codec::Codec;
|
||||
use crate::fastfield::FastFieldsWriter;
|
||||
use crate::fieldnorm::{FieldNormReaders, FieldNormsWriter};
|
||||
use crate::index::{Segment, SegmentComponent};
|
||||
@@ -13,7 +12,7 @@ use crate::indexer::segment_serializer::SegmentSerializer;
|
||||
use crate::json_utils::{index_json_value, IndexingPositionsPerPath};
|
||||
use crate::postings::{
|
||||
compute_table_memory_size, serialize_postings, IndexingContext, IndexingPosition,
|
||||
PerFieldPostingsWriter, PostingsWriter, PostingsWriterEnum,
|
||||
PerFieldPostingsWriter, PostingsWriter,
|
||||
};
|
||||
use crate::schema::document::{Document, Value};
|
||||
use crate::schema::{FieldEntry, FieldType, Schema, DATE_TIME_PRECISION_INDEXED};
|
||||
@@ -46,11 +45,11 @@ fn compute_initial_table_size(per_thread_memory_budget: usize) -> crate::Result<
|
||||
///
|
||||
/// They creates the postings list in anonymous memory.
|
||||
/// The segment is laid on disk when the segment gets `finalized`.
|
||||
pub struct SegmentWriter<Codec: crate::codec::Codec> {
|
||||
pub struct SegmentWriter {
|
||||
pub(crate) max_doc: DocId,
|
||||
pub(crate) ctx: IndexingContext,
|
||||
pub(crate) per_field_postings_writers: PerFieldPostingsWriter,
|
||||
pub(crate) segment_serializer: SegmentSerializer<Codec>,
|
||||
pub(crate) segment_serializer: SegmentSerializer,
|
||||
pub(crate) fast_field_writers: FastFieldsWriter,
|
||||
pub(crate) fieldnorms_writer: FieldNormsWriter,
|
||||
pub(crate) json_path_writer: JsonPathWriter,
|
||||
@@ -61,7 +60,7 @@ pub struct SegmentWriter<Codec: crate::codec::Codec> {
|
||||
schema: Schema,
|
||||
}
|
||||
|
||||
impl<Codec: crate::codec::Codec> SegmentWriter<Codec> {
|
||||
impl SegmentWriter {
|
||||
/// Creates a new `SegmentWriter`
|
||||
///
|
||||
/// The arguments are defined as follows
|
||||
@@ -71,10 +70,7 @@ impl<Codec: crate::codec::Codec> SegmentWriter<Codec> {
|
||||
/// behavior as a memory limit.
|
||||
/// - segment: The segment being written
|
||||
/// - schema
|
||||
pub fn for_segment(
|
||||
memory_budget_in_bytes: usize,
|
||||
segment: Segment<Codec>,
|
||||
) -> crate::Result<Self> {
|
||||
pub fn for_segment(memory_budget_in_bytes: usize, segment: Segment) -> crate::Result<Self> {
|
||||
let schema = segment.schema();
|
||||
let tokenizer_manager = segment.index().tokenizers().clone();
|
||||
let tokenizer_manager_fast_field = segment.index().fast_field_tokenizer().clone();
|
||||
@@ -173,7 +169,7 @@ impl<Codec: crate::codec::Codec> SegmentWriter<Codec> {
|
||||
}
|
||||
|
||||
let (term_buffer, ctx) = (&mut self.term_buffer, &mut self.ctx);
|
||||
let postings_writer: &mut PostingsWriterEnum =
|
||||
let postings_writer: &mut dyn PostingsWriter =
|
||||
self.per_field_postings_writers.get_for_field_mut(field);
|
||||
term_buffer.clear_with_field(field);
|
||||
|
||||
@@ -390,13 +386,13 @@ impl<Codec: crate::codec::Codec> SegmentWriter<Codec> {
|
||||
/// to the `SegmentSerializer`.
|
||||
///
|
||||
/// `doc_id_map` is used to map to the new doc_id order.
|
||||
fn remap_and_write<C: Codec>(
|
||||
fn remap_and_write(
|
||||
schema: Schema,
|
||||
per_field_postings_writers: &PerFieldPostingsWriter,
|
||||
ctx: IndexingContext,
|
||||
fast_field_writers: FastFieldsWriter,
|
||||
fieldnorms_writer: &FieldNormsWriter,
|
||||
mut serializer: SegmentSerializer<C>,
|
||||
mut serializer: SegmentSerializer,
|
||||
) -> crate::Result<()> {
|
||||
debug!("remap-and-write");
|
||||
if let Some(fieldnorms_serializer) = serializer.extract_fieldnorms_serializer() {
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use crate::codec::StandardCodec;
|
||||
use crate::index::CodecConfiguration;
|
||||
use crate::indexer::operation::AddOperation;
|
||||
use crate::indexer::segment_updater::save_metas;
|
||||
use crate::indexer::SegmentWriter;
|
||||
@@ -9,25 +7,22 @@ use crate::schema::document::Document;
|
||||
use crate::{Directory, Index, IndexMeta, Opstamp, Segment, TantivyDocument};
|
||||
|
||||
#[doc(hidden)]
|
||||
pub struct SingleSegmentIndexWriter<
|
||||
Codec: crate::codec::Codec = StandardCodec,
|
||||
D: Document = TantivyDocument,
|
||||
> {
|
||||
pub segment_writer: SegmentWriter<Codec>,
|
||||
segment: Segment<Codec>,
|
||||
pub struct SingleSegmentIndexWriter<D: Document = TantivyDocument> {
|
||||
segment_writer: SegmentWriter,
|
||||
segment: Segment,
|
||||
opstamp: Opstamp,
|
||||
_doc: PhantomData<D>,
|
||||
_phantom: PhantomData<D>,
|
||||
}
|
||||
|
||||
impl<Codec: crate::codec::Codec, D: Document> SingleSegmentIndexWriter<Codec, D> {
|
||||
pub fn new(index: Index<Codec>, mem_budget: usize) -> crate::Result<Self> {
|
||||
impl<D: Document> SingleSegmentIndexWriter<D> {
|
||||
pub fn new(index: Index, mem_budget: usize) -> crate::Result<Self> {
|
||||
let segment = index.new_segment();
|
||||
let segment_writer = SegmentWriter::for_segment(mem_budget, segment.clone())?;
|
||||
Ok(Self {
|
||||
segment_writer,
|
||||
segment,
|
||||
opstamp: 0,
|
||||
_doc: PhantomData,
|
||||
_phantom: PhantomData,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -42,10 +37,10 @@ impl<Codec: crate::codec::Codec, D: Document> SingleSegmentIndexWriter<Codec, D>
|
||||
.add_document(AddOperation { opstamp, document })
|
||||
}
|
||||
|
||||
pub fn finalize(self) -> crate::Result<Index<Codec>> {
|
||||
pub fn finalize(self) -> crate::Result<Index> {
|
||||
let max_doc = self.segment_writer.max_doc();
|
||||
self.segment_writer.finalize()?;
|
||||
let segment: Segment<Codec> = self.segment.with_max_doc(max_doc);
|
||||
let segment: Segment = self.segment.with_max_doc(max_doc);
|
||||
let index = segment.index();
|
||||
let index_meta = IndexMeta {
|
||||
index_settings: index.settings().clone(),
|
||||
@@ -53,7 +48,6 @@ impl<Codec: crate::codec::Codec, D: Document> SingleSegmentIndexWriter<Codec, D>
|
||||
schema: index.schema(),
|
||||
opstamp: 0,
|
||||
payload: None,
|
||||
codec: CodecConfiguration::from(index.codec()),
|
||||
};
|
||||
save_metas(&index_meta, index.directory())?;
|
||||
index.directory().sync_directory()?;
|
||||
|
||||
@@ -166,9 +166,6 @@ mod functional_test;
|
||||
|
||||
#[macro_use]
|
||||
mod macros;
|
||||
|
||||
/// Tantivy codecs describes how data is layed out on disk.
|
||||
pub mod codec;
|
||||
mod future_result;
|
||||
|
||||
// Re-exports
|
||||
|
||||
@@ -1,19 +1,28 @@
|
||||
use std::io;
|
||||
|
||||
use common::{OwnedBytes, VInt};
|
||||
use common::VInt;
|
||||
|
||||
use crate::codec::standard::postings::skip::{BlockInfo, SkipReader};
|
||||
use crate::codec::standard::postings::FreqReadingOption;
|
||||
use crate::directory::{FileSlice, OwnedBytes};
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::compression::{BlockDecoder, VIntDecoder as _, COMPRESSION_BLOCK_SIZE};
|
||||
use crate::postings::compression::{BlockDecoder, VIntDecoder, COMPRESSION_BLOCK_SIZE};
|
||||
use crate::postings::{BlockInfo, FreqReadingOption, SkipReader};
|
||||
use crate::query::Bm25Weight;
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::{DocId, Score, TERMINATED};
|
||||
|
||||
fn max_score<I: Iterator<Item = Score>>(mut it: I) -> Option<Score> {
|
||||
it.next().map(|first| it.fold(first, Score::max))
|
||||
}
|
||||
|
||||
/// `BlockSegmentPostings` is a cursor iterating over blocks
|
||||
/// of documents.
|
||||
///
|
||||
/// # Warning
|
||||
///
|
||||
/// While it is useful for some very specific high-performance
|
||||
/// use cases, you should prefer using `SegmentPostings` for most usage.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct BlockSegmentPostings {
|
||||
pub struct BlockSegmentPostings {
|
||||
pub(crate) doc_decoder: BlockDecoder,
|
||||
block_loaded: bool,
|
||||
freq_decoder: BlockDecoder,
|
||||
@@ -79,7 +88,7 @@ fn split_into_skips_and_postings(
|
||||
}
|
||||
|
||||
impl BlockSegmentPostings {
|
||||
/// Opens a `StandardPostingsReader`.
|
||||
/// Opens a `BlockSegmentPostings`.
|
||||
/// `doc_freq` is the number of documents in the posting list.
|
||||
/// `record_option` represents the amount of data available according to the schema.
|
||||
/// `requested_option` is the amount of data requested by the user.
|
||||
@@ -87,10 +96,11 @@ impl BlockSegmentPostings {
|
||||
/// term frequency blocks.
|
||||
pub(crate) fn open(
|
||||
doc_freq: u32,
|
||||
bytes: OwnedBytes,
|
||||
data: FileSlice,
|
||||
mut record_option: IndexRecordOption,
|
||||
requested_option: IndexRecordOption,
|
||||
) -> io::Result<BlockSegmentPostings> {
|
||||
let bytes = data.read_bytes()?;
|
||||
let (skip_data_opt, postings_data) = split_into_skips_and_postings(doc_freq, bytes)?;
|
||||
let skip_reader = match skip_data_opt {
|
||||
Some(skip_data) => {
|
||||
@@ -128,86 +138,6 @@ impl BlockSegmentPostings {
|
||||
block_segment_postings.load_block();
|
||||
Ok(block_segment_postings)
|
||||
}
|
||||
}
|
||||
|
||||
fn max_score<I: Iterator<Item = Score>>(mut it: I) -> Option<Score> {
|
||||
it.next().map(|first| it.fold(first, Score::max))
|
||||
}
|
||||
|
||||
impl BlockSegmentPostings {
|
||||
/// Returns the overall number of documents in the block postings.
|
||||
/// It does not take in account whether documents are deleted or not.
|
||||
///
|
||||
/// This `doc_freq` is simply the sum of the length of all of the blocks
|
||||
/// length, and it does not take in account deleted documents.
|
||||
pub fn doc_freq(&self) -> u32 {
|
||||
self.doc_freq
|
||||
}
|
||||
|
||||
/// Returns the array of docs in the current block.
|
||||
///
|
||||
/// Before the first call to `.advance()`, the block
|
||||
/// returned by `.docs()` is empty.
|
||||
#[inline]
|
||||
pub fn docs(&self) -> &[DocId] {
|
||||
debug_assert!(self.block_loaded);
|
||||
self.doc_decoder.output_array()
|
||||
}
|
||||
|
||||
/// Return the document at index `idx` of the block.
|
||||
#[inline]
|
||||
pub fn doc(&self, idx: usize) -> u32 {
|
||||
self.doc_decoder.output(idx)
|
||||
}
|
||||
|
||||
/// Return the array of `term freq` in the block.
|
||||
#[inline]
|
||||
pub fn freqs(&self) -> &[u32] {
|
||||
debug_assert!(self.block_loaded);
|
||||
self.freq_decoder.output_array()
|
||||
}
|
||||
|
||||
/// Return the frequency at index `idx` of the block.
|
||||
#[inline]
|
||||
pub fn freq(&self, idx: usize) -> u32 {
|
||||
debug_assert!(self.block_loaded);
|
||||
self.freq_decoder.output(idx)
|
||||
}
|
||||
|
||||
/// Position on a block that may contains `target_doc`.
|
||||
///
|
||||
/// If all docs are smaller than target, the block loaded may be empty,
|
||||
/// or be the last an incomplete VInt block.
|
||||
pub fn seek(&mut self, target_doc: DocId) -> usize {
|
||||
// Move to the block that might contain our document.
|
||||
self.seek_block_without_loading(target_doc);
|
||||
self.load_block();
|
||||
|
||||
// At this point we are on the block that might contain our document.
|
||||
let doc = self.doc_decoder.seek_within_block(target_doc);
|
||||
|
||||
// The last block is not full and padded with TERMINATED,
|
||||
// so we are guaranteed to have at least one value (real or padding)
|
||||
// that is >= target_doc.
|
||||
debug_assert!(doc < COMPRESSION_BLOCK_SIZE);
|
||||
|
||||
// `doc` is now the first element >= `target_doc`.
|
||||
// If all docs are smaller than target, the current block is incomplete and padded
|
||||
// with TERMINATED. After the search, the cursor points to the first TERMINATED.
|
||||
doc
|
||||
}
|
||||
|
||||
pub fn position_offset(&self) -> u64 {
|
||||
self.skip_reader.position_offset()
|
||||
}
|
||||
|
||||
/// Advance to the next block.
|
||||
pub fn advance(&mut self) {
|
||||
self.skip_reader.advance();
|
||||
self.block_loaded = false;
|
||||
self.block_max_score_cache = None;
|
||||
self.load_block();
|
||||
}
|
||||
|
||||
/// Returns the block_max_score for the current block.
|
||||
/// It does not require the block to be loaded. For instance, it is ok to call this method
|
||||
@@ -230,7 +160,7 @@ impl BlockSegmentPostings {
|
||||
}
|
||||
// this is the last block of the segment posting list.
|
||||
// If it is actually loaded, we can compute block max manually.
|
||||
if self.block_loaded {
|
||||
if self.block_is_loaded() {
|
||||
let docs = self.doc_decoder.output_array().iter().cloned();
|
||||
let freqs = self.freq_decoder.output_array().iter().cloned();
|
||||
let bm25_scores = docs.zip(freqs).map(|(doc, term_freq)| {
|
||||
@@ -247,25 +177,145 @@ impl BlockSegmentPostings {
|
||||
// We do not cache it however, so that it gets computed when once block is loaded.
|
||||
bm25_weight.max_score()
|
||||
}
|
||||
}
|
||||
|
||||
impl BlockSegmentPostings {
|
||||
/// Returns an empty segment postings object
|
||||
pub fn empty() -> BlockSegmentPostings {
|
||||
BlockSegmentPostings {
|
||||
doc_decoder: BlockDecoder::with_val(TERMINATED),
|
||||
block_loaded: true,
|
||||
freq_decoder: BlockDecoder::with_val(1),
|
||||
freq_reading_option: FreqReadingOption::NoFreq,
|
||||
block_max_score_cache: None,
|
||||
doc_freq: 0,
|
||||
data: OwnedBytes::empty(),
|
||||
skip_reader: SkipReader::new(OwnedBytes::empty(), 0, IndexRecordOption::Basic),
|
||||
}
|
||||
pub(crate) fn freq_reading_option(&self) -> FreqReadingOption {
|
||||
self.freq_reading_option
|
||||
}
|
||||
|
||||
pub(crate) fn skip_reader(&self) -> &SkipReader {
|
||||
&self.skip_reader
|
||||
// Resets the block segment postings on another position
|
||||
// in the postings file.
|
||||
//
|
||||
// This is useful for enumerating through a list of terms,
|
||||
// and consuming the associated posting lists while avoiding
|
||||
// reallocating a `BlockSegmentPostings`.
|
||||
//
|
||||
// # Warning
|
||||
//
|
||||
// This does not reset the positions list.
|
||||
pub(crate) fn reset(&mut self, doc_freq: u32, postings_data: OwnedBytes) -> io::Result<()> {
|
||||
let (skip_data_opt, postings_data) =
|
||||
split_into_skips_and_postings(doc_freq, postings_data)?;
|
||||
self.data = postings_data;
|
||||
self.block_max_score_cache = None;
|
||||
self.block_loaded = false;
|
||||
if let Some(skip_data) = skip_data_opt {
|
||||
self.skip_reader.reset(skip_data, doc_freq);
|
||||
} else {
|
||||
self.skip_reader.reset(OwnedBytes::empty(), doc_freq);
|
||||
}
|
||||
self.doc_freq = doc_freq;
|
||||
self.load_block();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the overall number of documents in the block postings.
|
||||
/// It does not take in account whether documents are deleted or not.
|
||||
///
|
||||
/// This `doc_freq` is simply the sum of the length of all of the blocks
|
||||
/// length, and it does not take in account deleted documents.
|
||||
pub fn doc_freq(&self) -> u32 {
|
||||
self.doc_freq
|
||||
}
|
||||
|
||||
/// Returns the array of docs in the current block.
|
||||
///
|
||||
/// Before the first call to `.advance()`, the block
|
||||
/// returned by `.docs()` is empty.
|
||||
#[inline]
|
||||
pub fn docs(&self) -> &[DocId] {
|
||||
debug_assert!(self.block_is_loaded());
|
||||
self.doc_decoder.output_array()
|
||||
}
|
||||
|
||||
/// Return the document at index `idx` of the block.
|
||||
#[inline]
|
||||
pub fn doc(&self, idx: usize) -> u32 {
|
||||
self.doc_decoder.output(idx)
|
||||
}
|
||||
|
||||
/// Return the array of `term freq` in the block.
|
||||
#[inline]
|
||||
pub fn freqs(&self) -> &[u32] {
|
||||
debug_assert!(self.block_is_loaded());
|
||||
self.freq_decoder.output_array()
|
||||
}
|
||||
|
||||
/// Return the frequency at index `idx` of the block.
|
||||
#[inline]
|
||||
pub fn freq(&self, idx: usize) -> u32 {
|
||||
debug_assert!(self.block_is_loaded());
|
||||
self.freq_decoder.output(idx)
|
||||
}
|
||||
|
||||
/// Returns the length of the current block.
|
||||
///
|
||||
/// Returns the decoded term-frequency buffer for the current block.
|
||||
#[inline]
|
||||
pub(crate) fn freq_output_array(&self) -> &[u32] {
|
||||
self.freq_decoder.output_array()
|
||||
}
|
||||
|
||||
/// All blocks have a length of `NUM_DOCS_PER_BLOCK`,
|
||||
/// except the last block that may have a length
|
||||
/// of any number between 1 and `NUM_DOCS_PER_BLOCK - 1`
|
||||
#[inline]
|
||||
pub fn block_len(&self) -> usize {
|
||||
debug_assert!(self.block_is_loaded());
|
||||
self.doc_decoder.output_len
|
||||
}
|
||||
|
||||
/// Position on a block that may contains `target_doc`.
|
||||
///
|
||||
/// If all docs are smaller than target, the block loaded may be empty,
|
||||
/// or be the last an incomplete VInt block.
|
||||
pub fn seek(&mut self, target_doc: DocId) -> usize {
|
||||
// Move to the block that might contain our document.
|
||||
self.seek_block(target_doc);
|
||||
self.load_block();
|
||||
|
||||
// At this point we are on the block that might contain our document.
|
||||
let doc = self.doc_decoder.seek_within_block(target_doc);
|
||||
|
||||
// The last block is not full and padded with TERMINATED,
|
||||
// so we are guaranteed to have at least one value (real or padding)
|
||||
// that is >= target_doc.
|
||||
debug_assert!(doc < COMPRESSION_BLOCK_SIZE);
|
||||
|
||||
// `doc` is now the first element >= `target_doc`.
|
||||
// If all docs are smaller than target, the current block is incomplete and padded
|
||||
// with TERMINATED. After the search, the cursor points to the first TERMINATED.
|
||||
doc
|
||||
}
|
||||
|
||||
/// Returns the number of documents with a doc id strictly smaller than `target`
|
||||
/// (i.e. the *rank* of `target` in this posting list).
|
||||
///
|
||||
/// This jumps to the block that may contain `target` through the skip list, so no
|
||||
/// skipped block is decoded; a single block is then decoded to locate `target`
|
||||
/// within it. The cost is therefore `O(number_of_skip_list_entries)` plus one block
|
||||
/// decode, rather than `O(doc_freq)`.
|
||||
///
|
||||
/// Like [`Self::seek`], the underlying cursor only ever moves forward. This method
|
||||
/// must be called with **non-decreasing** `target` values (galloping); calling it
|
||||
/// with a `target` smaller than a previous one yields an incorrect result. `target`
|
||||
/// must be a valid doc id (i.e. `target <= TERMINATED`), exactly as for `seek`.
|
||||
///
|
||||
/// Edge cases: returns `0` when `target` is smaller than every doc id, and
|
||||
/// `doc_freq()` when `target` is larger than every doc id.
|
||||
pub fn rank(&mut self, target: DocId) -> u32 {
|
||||
if self.doc_freq == 0 {
|
||||
return 0;
|
||||
}
|
||||
// `within` = number of docs in the landed block with a doc id < target.
|
||||
let within = self.seek(target);
|
||||
// `remaining_docs` counts the landed block and everything after it, so the
|
||||
// difference is the number of docs in all blocks strictly before it.
|
||||
let docs_before_block = self.doc_freq - self.skip_reader.remaining_docs();
|
||||
docs_before_block + within as u32
|
||||
}
|
||||
|
||||
pub(crate) fn position_offset(&self) -> u64 {
|
||||
self.skip_reader.position_offset()
|
||||
}
|
||||
|
||||
/// Dangerous API! This calls seeks the next block on the skip list,
|
||||
@@ -274,15 +324,24 @@ impl BlockSegmentPostings {
|
||||
/// `.load_block()` needs to be called manually afterwards.
|
||||
/// If all docs are smaller than target, the block loaded may be empty,
|
||||
/// or be the last an incomplete VInt block.
|
||||
pub(crate) fn seek_block_without_loading(&mut self, target_doc: DocId) {
|
||||
pub(crate) fn seek_block(&mut self, target_doc: DocId) {
|
||||
if self.skip_reader.seek(target_doc) {
|
||||
self.block_max_score_cache = None;
|
||||
self.block_loaded = false;
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn has_remaining_docs(&self) -> bool {
|
||||
self.skip_reader.has_remaining_docs()
|
||||
}
|
||||
|
||||
pub(crate) fn block_is_loaded(&self) -> bool {
|
||||
self.block_loaded
|
||||
}
|
||||
|
||||
pub(crate) fn load_block(&mut self) {
|
||||
if self.block_loaded {
|
||||
if self.block_is_loaded() {
|
||||
return;
|
||||
}
|
||||
let offset = self.skip_reader.byte_offset();
|
||||
@@ -330,40 +389,68 @@ impl BlockSegmentPostings {
|
||||
}
|
||||
self.block_loaded = true;
|
||||
}
|
||||
|
||||
/// Advance to the next block.
|
||||
pub fn advance(&mut self) {
|
||||
self.skip_reader.advance();
|
||||
self.block_loaded = false;
|
||||
self.block_max_score_cache = None;
|
||||
self.load_block();
|
||||
}
|
||||
|
||||
/// Returns an empty segment postings object
|
||||
pub fn empty() -> BlockSegmentPostings {
|
||||
BlockSegmentPostings {
|
||||
doc_decoder: BlockDecoder::with_val(TERMINATED),
|
||||
block_loaded: true,
|
||||
freq_decoder: BlockDecoder::with_val(1),
|
||||
freq_reading_option: FreqReadingOption::NoFreq,
|
||||
block_max_score_cache: None,
|
||||
doc_freq: 0,
|
||||
data: OwnedBytes::empty(),
|
||||
skip_reader: SkipReader::new(OwnedBytes::empty(), 0, IndexRecordOption::Basic),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn skip_reader(&self) -> &SkipReader {
|
||||
&self.skip_reader
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use common::OwnedBytes;
|
||||
use common::HasLen;
|
||||
|
||||
use super::BlockSegmentPostings;
|
||||
use crate::codec::postings::PostingsSerializer;
|
||||
use crate::codec::standard::postings::segment_postings::SegmentPostings;
|
||||
use crate::codec::standard::postings::StandardPostingsSerializer;
|
||||
use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::index::Index;
|
||||
use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::postings::postings::Postings;
|
||||
use crate::postings::SegmentPostings;
|
||||
use crate::schema::{IndexRecordOption, Schema, Term, INDEXED};
|
||||
use crate::DocId;
|
||||
|
||||
#[cfg(test)]
|
||||
fn build_block_postings(docs: &[u32]) -> BlockSegmentPostings {
|
||||
let doc_freq = docs.len() as u32;
|
||||
let mut postings_serializer =
|
||||
StandardPostingsSerializer::new(1.0f32, IndexRecordOption::Basic, None);
|
||||
postings_serializer.new_term(docs.len() as u32, false);
|
||||
for doc in docs {
|
||||
postings_serializer.write_doc(*doc, 1u32);
|
||||
}
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
postings_serializer
|
||||
.close_term(doc_freq, &mut buffer)
|
||||
.unwrap();
|
||||
BlockSegmentPostings::open(
|
||||
doc_freq,
|
||||
OwnedBytes::new(buffer),
|
||||
IndexRecordOption::Basic,
|
||||
IndexRecordOption::Basic,
|
||||
)
|
||||
.unwrap()
|
||||
#[test]
|
||||
fn test_empty_segment_postings() {
|
||||
let mut postings = SegmentPostings::empty();
|
||||
assert_eq!(postings.doc(), TERMINATED);
|
||||
assert_eq!(postings.advance(), TERMINATED);
|
||||
assert_eq!(postings.advance(), TERMINATED);
|
||||
assert_eq!(postings.doc_freq(), 0);
|
||||
assert_eq!(postings.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_postings_doc_returns_terminated() {
|
||||
let mut postings = SegmentPostings::empty();
|
||||
assert_eq!(postings.doc(), TERMINATED);
|
||||
assert_eq!(postings.advance(), TERMINATED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_postings_doc_term_freq_returns_0() {
|
||||
let postings = SegmentPostings::empty();
|
||||
assert_eq!(postings.term_freq(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -378,7 +465,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_block_segment_postings() -> crate::Result<()> {
|
||||
let mut block_segments = build_block_postings(&(0..100_000).collect::<Vec<u32>>());
|
||||
let mut block_segments = build_block_postings(&(0..100_000).collect::<Vec<u32>>())?;
|
||||
let mut offset: u32 = 0u32;
|
||||
// checking that the `doc_freq` is correct
|
||||
assert_eq!(block_segments.doc_freq(), 100_000);
|
||||
@@ -403,7 +490,7 @@ mod tests {
|
||||
doc_ids.push(129);
|
||||
doc_ids.push(130);
|
||||
{
|
||||
let block_segments = build_block_postings(&doc_ids);
|
||||
let block_segments = build_block_postings(&doc_ids)?;
|
||||
let mut docset = SegmentPostings::from_block_postings(block_segments, None);
|
||||
assert_eq!(docset.seek(128), 129);
|
||||
assert_eq!(docset.doc(), 129);
|
||||
@@ -412,7 +499,7 @@ mod tests {
|
||||
assert_eq!(docset.advance(), TERMINATED);
|
||||
}
|
||||
{
|
||||
let block_segments = build_block_postings(&doc_ids);
|
||||
let block_segments = build_block_postings(&doc_ids).unwrap();
|
||||
let mut docset = SegmentPostings::from_block_postings(block_segments, None);
|
||||
assert_eq!(docset.seek(129), 129);
|
||||
assert_eq!(docset.doc(), 129);
|
||||
@@ -421,7 +508,7 @@ mod tests {
|
||||
assert_eq!(docset.advance(), TERMINATED);
|
||||
}
|
||||
{
|
||||
let block_segments = build_block_postings(&doc_ids);
|
||||
let block_segments = build_block_postings(&doc_ids)?;
|
||||
let mut docset = SegmentPostings::from_block_postings(block_segments, None);
|
||||
assert_eq!(docset.doc(), 0);
|
||||
assert_eq!(docset.seek(131), TERMINATED);
|
||||
@@ -430,13 +517,38 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build_block_postings(docs: &[DocId]) -> crate::Result<BlockSegmentPostings> {
|
||||
let mut schema_builder = Schema::builder();
|
||||
let int_field = schema_builder.add_u64_field("id", INDEXED);
|
||||
let schema = schema_builder.build();
|
||||
let index = Index::create_in_ram(schema);
|
||||
let mut index_writer = index.writer_for_tests()?;
|
||||
let mut last_doc = 0u32;
|
||||
for &doc in docs {
|
||||
for _ in last_doc..doc {
|
||||
index_writer.add_document(doc!(int_field=>1u64))?;
|
||||
}
|
||||
index_writer.add_document(doc!(int_field=>0u64))?;
|
||||
last_doc = doc + 1;
|
||||
}
|
||||
index_writer.commit()?;
|
||||
let searcher = index.reader()?.searcher();
|
||||
let segment_reader = searcher.segment_reader(0);
|
||||
let inverted_index = segment_reader.inverted_index(int_field).unwrap();
|
||||
let term = Term::from_field_u64(int_field, 0u64);
|
||||
let term_info = inverted_index.get_term_info(&term)?.unwrap();
|
||||
let block_postings = inverted_index
|
||||
.read_block_postings_from_terminfo(&term_info, IndexRecordOption::Basic)?;
|
||||
Ok(block_postings)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_block_segment_postings_seek() -> crate::Result<()> {
|
||||
let mut docs = Vec::new();
|
||||
let mut docs = vec![0];
|
||||
for i in 0..1300 {
|
||||
docs.push((i * i / 100) + i);
|
||||
}
|
||||
let mut block_postings = build_block_postings(&docs[..]);
|
||||
let mut block_postings = build_block_postings(&docs[..])?;
|
||||
for i in &[0, 424, 10000] {
|
||||
block_postings.seek(*i);
|
||||
let docs = block_postings.docs();
|
||||
@@ -447,4 +559,74 @@ mod tests {
|
||||
assert_eq!(block_postings.doc(COMPRESSION_BLOCK_SIZE - 1), TERMINATED);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset_block_segment_postings() -> crate::Result<()> {
|
||||
let mut schema_builder = Schema::builder();
|
||||
let int_field = schema_builder.add_u64_field("id", INDEXED);
|
||||
let schema = schema_builder.build();
|
||||
let index = Index::create_in_ram(schema);
|
||||
let mut index_writer = index.writer_for_tests()?;
|
||||
// create two postings list, one containing even number,
|
||||
// the other containing odd numbers.
|
||||
for i in 0..6 {
|
||||
let doc = doc!(int_field=> (i % 2) as u64);
|
||||
index_writer.add_document(doc)?;
|
||||
}
|
||||
index_writer.commit()?;
|
||||
let searcher = index.reader()?.searcher();
|
||||
let segment_reader = searcher.segment_reader(0);
|
||||
|
||||
let mut block_segments;
|
||||
{
|
||||
let term = Term::from_field_u64(int_field, 0u64);
|
||||
let inverted_index = segment_reader.inverted_index(int_field)?;
|
||||
let term_info = inverted_index.get_term_info(&term)?.unwrap();
|
||||
block_segments = inverted_index
|
||||
.read_block_postings_from_terminfo(&term_info, IndexRecordOption::Basic)?;
|
||||
}
|
||||
assert_eq!(block_segments.docs(), &[0, 2, 4]);
|
||||
{
|
||||
let term = Term::from_field_u64(int_field, 1u64);
|
||||
let inverted_index = segment_reader.inverted_index(int_field)?;
|
||||
let term_info = inverted_index.get_term_info(&term)?.unwrap();
|
||||
inverted_index.reset_block_postings_from_terminfo(&term_info, &mut block_segments)?;
|
||||
}
|
||||
assert_eq!(block_segments.docs(), &[1, 3, 5]);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_block_segment_postings_rank() -> crate::Result<()> {
|
||||
// ~8 blocks worth of docs so the skip list is actually exercised.
|
||||
let docs: Vec<DocId> = (0..1000u32).map(|i| i * 3).collect();
|
||||
let mut block_postings = build_block_postings(&docs[..])?;
|
||||
let doc_freq = block_postings.doc_freq();
|
||||
|
||||
// rank(target) must equal the number of docs strictly below target.
|
||||
// Targets are queried in non-decreasing order, as the API requires.
|
||||
// `target` values must be a valid doc id (<= TERMINATED) and non-decreasing.
|
||||
let targets = [
|
||||
0u32, 1, 2, 3, 4, 299, 300, 301, 1500, 2996, 2997, 3000, 10_000,
|
||||
];
|
||||
for &target in &targets {
|
||||
let expected = docs.iter().filter(|&&d| d < target).count() as u32;
|
||||
assert_eq!(
|
||||
block_postings.rank(target),
|
||||
expected,
|
||||
"rank({target}) mismatch"
|
||||
);
|
||||
}
|
||||
|
||||
// Edge cases: below the first doc -> 0, above the last doc -> doc_freq.
|
||||
let mut fresh = build_block_postings(&docs[..])?;
|
||||
assert_eq!(fresh.rank(0), 0);
|
||||
let mut fresh = build_block_postings(&docs[..])?;
|
||||
assert_eq!(fresh.rank(1_000_000), doc_freq);
|
||||
|
||||
// Empty postings: rank is always 0.
|
||||
let mut empty = BlockSegmentPostings::empty();
|
||||
assert_eq!(empty.rank(42), 0);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,6 @@ use std::io;
|
||||
use common::json_path_writer::JSON_END_OF_PATH;
|
||||
use stacker::Addr;
|
||||
|
||||
use crate::codec::Codec;
|
||||
use crate::indexer::indexing_term::IndexingTerm;
|
||||
use crate::indexer::path_to_unordered_id::OrderedPathId;
|
||||
use crate::postings::postings_writer::SpecializedPostingsWriter;
|
||||
@@ -23,6 +22,12 @@ pub(crate) struct JsonPostingsWriter<Rec: Recorder> {
|
||||
non_str_posting_writer: SpecializedPostingsWriter<DocIdRecorder>,
|
||||
}
|
||||
|
||||
impl<Rec: Recorder> From<JsonPostingsWriter<Rec>> for Box<dyn PostingsWriter> {
|
||||
fn from(json_postings_writer: JsonPostingsWriter<Rec>) -> Box<dyn PostingsWriter> {
|
||||
Box::new(json_postings_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<Rec: Recorder> PostingsWriter for JsonPostingsWriter<Rec> {
|
||||
#[inline]
|
||||
fn subscribe(
|
||||
@@ -53,12 +58,12 @@ impl<Rec: Recorder> PostingsWriter for JsonPostingsWriter<Rec> {
|
||||
}
|
||||
|
||||
/// The actual serialization format is handled by the `PostingsSerializer`.
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
ordered_term_addrs: &[(Field, OrderedPathId, &[u8], Addr)],
|
||||
ordered_id_to_path: &[&str],
|
||||
ctx: &IndexingContext,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer,
|
||||
) -> io::Result<()> {
|
||||
let mut term_buffer = JsonTermSerializer(Vec::with_capacity(48));
|
||||
let mut buffer_lender = BufferLender::default();
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::postings::{DocFreq, Postings};
|
||||
use crate::postings::{Postings, SegmentPostings};
|
||||
use crate::DocId;
|
||||
|
||||
/// `LoadedPostings` is a `DocSet` and `Postings` implementation.
|
||||
@@ -25,16 +25,16 @@ impl LoadedPostings {
|
||||
/// Creates a new `LoadedPostings` from a `SegmentPostings`.
|
||||
///
|
||||
/// It will also preload positions, if positions are available in the SegmentPostings.
|
||||
pub fn load(postings: &mut Box<dyn Postings>) -> LoadedPostings {
|
||||
let num_docs: usize = u32::from(postings.doc_freq()) as usize;
|
||||
pub fn load(segment_postings: &mut SegmentPostings) -> LoadedPostings {
|
||||
let num_docs = segment_postings.doc_freq() as usize;
|
||||
let mut doc_ids = Vec::with_capacity(num_docs);
|
||||
let mut positions = Vec::with_capacity(num_docs);
|
||||
let mut position_offsets = Vec::with_capacity(num_docs);
|
||||
while postings.doc() != TERMINATED {
|
||||
while segment_postings.doc() != TERMINATED {
|
||||
position_offsets.push(positions.len() as u32);
|
||||
doc_ids.push(postings.doc());
|
||||
postings.append_positions_with_offset(0, &mut positions);
|
||||
postings.advance();
|
||||
doc_ids.push(segment_postings.doc());
|
||||
segment_postings.append_positions_with_offset(0, &mut positions);
|
||||
segment_postings.advance();
|
||||
}
|
||||
position_offsets.push(positions.len() as u32);
|
||||
LoadedPostings {
|
||||
@@ -101,14 +101,6 @@ impl Postings for LoadedPostings {
|
||||
output.push(*pos + offset);
|
||||
}
|
||||
}
|
||||
|
||||
fn has_freq(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn doc_freq(&self) -> DocFreq {
|
||||
DocFreq::Exact(self.doc_ids.len() as u32)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -4,6 +4,7 @@ mod block_search;
|
||||
|
||||
pub(crate) use self::block_search::branchless_binary_search;
|
||||
|
||||
mod block_segment_postings;
|
||||
pub(crate) mod compression;
|
||||
mod indexing_context;
|
||||
mod json_postings_writer;
|
||||
@@ -12,23 +13,33 @@ mod per_field_postings_writer;
|
||||
mod postings;
|
||||
mod postings_writer;
|
||||
mod recorder;
|
||||
mod segment_postings;
|
||||
/// Serializer module for the inverted index
|
||||
pub mod serializer;
|
||||
mod skip;
|
||||
mod term_info;
|
||||
|
||||
pub(crate) use loaded_postings::LoadedPostings;
|
||||
pub use postings::DocFreq;
|
||||
pub(crate) use stacker::compute_table_memory_size;
|
||||
|
||||
pub use self::block_segment_postings::BlockSegmentPostings;
|
||||
pub(crate) use self::indexing_context::IndexingContext;
|
||||
pub(crate) use self::per_field_postings_writer::PerFieldPostingsWriter;
|
||||
pub use self::postings::Postings;
|
||||
pub(crate) use self::postings_writer::{
|
||||
serialize_postings, IndexingPosition, PostingsWriter, PostingsWriterEnum,
|
||||
};
|
||||
pub(crate) use self::postings_writer::{serialize_postings, IndexingPosition, PostingsWriter};
|
||||
pub use self::segment_postings::SegmentPostings;
|
||||
pub use self::serializer::{FieldSerializer, InvertedIndexSerializer};
|
||||
pub(crate) use self::skip::{BlockInfo, SkipReader};
|
||||
pub use self::term_info::TermInfo;
|
||||
|
||||
#[expect(clippy::enum_variant_names)]
|
||||
#[derive(Debug, PartialEq, Clone, Copy, Eq)]
|
||||
pub(crate) enum FreqReadingOption {
|
||||
NoFreq,
|
||||
SkipFreq,
|
||||
ReadFreq,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod tests {
|
||||
use std::mem;
|
||||
@@ -39,7 +50,6 @@ pub(crate) mod tests {
|
||||
use crate::index::{Index, SegmentComponent, SegmentReader};
|
||||
use crate::indexer::operation::AddOperation;
|
||||
use crate::indexer::SegmentWriter;
|
||||
use crate::postings::DocFreq;
|
||||
use crate::query::Scorer;
|
||||
use crate::schema::{
|
||||
Field, IndexRecordOption, Schema, Term, TextFieldIndexing, TextOptions, INDEXED, TEXT,
|
||||
@@ -270,11 +280,11 @@ pub(crate) mod tests {
|
||||
}
|
||||
{
|
||||
let term_a = Term::from_field_text(text_field, "a");
|
||||
let mut postings_a: Box<dyn Postings> = segment_reader
|
||||
let mut postings_a = segment_reader
|
||||
.inverted_index(term_a.field())?
|
||||
.read_postings(&term_a, IndexRecordOption::WithFreqsAndPositions)?
|
||||
.unwrap();
|
||||
assert_eq!(postings_a.doc_freq(), DocFreq::Exact(1000));
|
||||
assert_eq!(postings_a.len(), 1000);
|
||||
assert_eq!(postings_a.doc(), 0);
|
||||
assert_eq!(postings_a.term_freq(), 6);
|
||||
postings_a.positions(&mut positions);
|
||||
@@ -297,7 +307,7 @@ pub(crate) mod tests {
|
||||
.inverted_index(term_e.field())?
|
||||
.read_postings(&term_e, IndexRecordOption::WithFreqsAndPositions)?
|
||||
.unwrap();
|
||||
assert_eq!(postings_e.doc_freq(), DocFreq::Exact(1000 - 2));
|
||||
assert_eq!(postings_e.len(), 1000 - 2);
|
||||
for i in 2u32..1000u32 {
|
||||
assert_eq!(postings_e.term_freq(), i);
|
||||
postings_e.positions(&mut positions);
|
||||
|
||||
@@ -1,15 +1,16 @@
|
||||
use crate::postings::json_postings_writer::JsonPostingsWriter;
|
||||
use crate::postings::postings_writer::{PostingsWriterEnum, SpecializedPostingsWriter};
|
||||
use crate::postings::postings_writer::SpecializedPostingsWriter;
|
||||
use crate::postings::recorder::{DocIdRecorder, TermFrequencyRecorder, TfAndPositionRecorder};
|
||||
use crate::postings::PostingsWriter;
|
||||
use crate::schema::{Field, FieldEntry, FieldType, IndexRecordOption, Schema};
|
||||
|
||||
pub(crate) struct PerFieldPostingsWriter {
|
||||
per_field_postings_writers: Vec<PostingsWriterEnum>,
|
||||
per_field_postings_writers: Vec<Box<dyn PostingsWriter>>,
|
||||
}
|
||||
|
||||
impl PerFieldPostingsWriter {
|
||||
pub fn for_schema(schema: &Schema) -> Self {
|
||||
let per_field_postings_writers: Vec<PostingsWriterEnum> = schema
|
||||
let per_field_postings_writers = schema
|
||||
.fields()
|
||||
.map(|(_, field_entry)| posting_writer_from_field_entry(field_entry))
|
||||
.collect();
|
||||
@@ -18,16 +19,16 @@ impl PerFieldPostingsWriter {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn get_for_field(&self, field: Field) -> &PostingsWriterEnum {
|
||||
&self.per_field_postings_writers[field.field_id() as usize]
|
||||
pub(crate) fn get_for_field(&self, field: Field) -> &dyn PostingsWriter {
|
||||
self.per_field_postings_writers[field.field_id() as usize].as_ref()
|
||||
}
|
||||
|
||||
pub(crate) fn get_for_field_mut(&mut self, field: Field) -> &mut PostingsWriterEnum {
|
||||
&mut self.per_field_postings_writers[field.field_id() as usize]
|
||||
pub(crate) fn get_for_field_mut(&mut self, field: Field) -> &mut dyn PostingsWriter {
|
||||
self.per_field_postings_writers[field.field_id() as usize].as_mut()
|
||||
}
|
||||
}
|
||||
|
||||
fn posting_writer_from_field_entry(field_entry: &FieldEntry) -> PostingsWriterEnum {
|
||||
fn posting_writer_from_field_entry(field_entry: &FieldEntry) -> Box<dyn PostingsWriter> {
|
||||
match *field_entry.field_type() {
|
||||
FieldType::Str(ref text_options) => text_options
|
||||
.get_indexing_options()
|
||||
@@ -50,7 +51,7 @@ fn posting_writer_from_field_entry(field_entry: &FieldEntry) -> PostingsWriterEn
|
||||
| FieldType::Date(_)
|
||||
| FieldType::Bytes(_)
|
||||
| FieldType::IpAddr(_)
|
||||
| FieldType::Facet(_) => <SpecializedPostingsWriter<DocIdRecorder>>::default().into(),
|
||||
| FieldType::Facet(_) => Box::<SpecializedPostingsWriter<DocIdRecorder>>::default(),
|
||||
FieldType::JsonObject(ref json_object_options) => {
|
||||
if let Some(text_indexing_option) = json_object_options.get_text_indexing_options() {
|
||||
match text_indexing_option.index_option() {
|
||||
|
||||
@@ -1,25 +1,5 @@
|
||||
use crate::docset::DocSet;
|
||||
|
||||
/// Result of the doc_freq method.
|
||||
///
|
||||
/// Postings can inform us that the document frequency is approximate.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum DocFreq {
|
||||
/// The document frequency is approximate.
|
||||
Approximate(u32),
|
||||
/// The document frequency is exact.
|
||||
Exact(u32),
|
||||
}
|
||||
|
||||
impl From<DocFreq> for u32 {
|
||||
fn from(doc_freq: DocFreq) -> Self {
|
||||
match doc_freq {
|
||||
DocFreq::Approximate(approximate_doc_freq) => approximate_doc_freq,
|
||||
DocFreq::Exact(doc_freq) => doc_freq,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Postings (also called inverted list)
|
||||
///
|
||||
/// For a given term, it is the list of doc ids of the doc
|
||||
@@ -34,9 +14,6 @@ pub trait Postings: DocSet + 'static {
|
||||
/// The number of times the term appears in the document.
|
||||
fn term_freq(&self) -> u32;
|
||||
|
||||
/// Returns the number of documents containing the term in the segment.
|
||||
fn doc_freq(&self) -> DocFreq;
|
||||
|
||||
/// Returns the positions offsetted with a given value.
|
||||
/// It is not necessary to clear the `output` before calling this method.
|
||||
/// The output vector will be resized to the `term_freq`.
|
||||
@@ -54,16 +31,6 @@ pub trait Postings: DocSet + 'static {
|
||||
fn positions(&mut self, output: &mut Vec<u32>) {
|
||||
self.positions_with_offset(0u32, output);
|
||||
}
|
||||
|
||||
/// Returns true if the term_frequency is available.
|
||||
///
|
||||
/// This is a tricky question, because on JSON fields, it is possible
|
||||
/// for a text term to have term freq, whereas a number term in the field has none.
|
||||
///
|
||||
/// This function returns whether the actual term has term frequencies or not.
|
||||
/// In this above JSON field example, `has_freq` should return true for the
|
||||
/// earlier and false for the latter.
|
||||
fn has_freq(&self) -> bool;
|
||||
}
|
||||
|
||||
impl Postings for Box<dyn Postings> {
|
||||
@@ -74,12 +41,4 @@ impl Postings for Box<dyn Postings> {
|
||||
fn append_positions_with_offset(&mut self, offset: u32, output: &mut Vec<u32>) {
|
||||
(**self).append_positions_with_offset(offset, output);
|
||||
}
|
||||
|
||||
fn has_freq(&self) -> bool {
|
||||
(**self).has_freq()
|
||||
}
|
||||
|
||||
fn doc_freq(&self) -> DocFreq {
|
||||
(**self).doc_freq()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,14 +4,10 @@ use std::ops::Range;
|
||||
|
||||
use stacker::Addr;
|
||||
|
||||
use crate::codec::Codec;
|
||||
use crate::fieldnorm::FieldNormReaders;
|
||||
use crate::indexer::indexing_term::IndexingTerm;
|
||||
use crate::indexer::path_to_unordered_id::OrderedPathId;
|
||||
use crate::postings::json_postings_writer::JsonPostingsWriter;
|
||||
use crate::postings::recorder::{
|
||||
BufferLender, DocIdRecorder, Recorder, TermFrequencyRecorder, TfAndPositionRecorder,
|
||||
};
|
||||
use crate::postings::recorder::{BufferLender, Recorder};
|
||||
use crate::postings::{
|
||||
FieldSerializer, IndexingContext, InvertedIndexSerializer, PerFieldPostingsWriter,
|
||||
};
|
||||
@@ -49,12 +45,12 @@ fn make_field_partition(
|
||||
/// Serialize the inverted index.
|
||||
/// It pushes all term, one field at a time, towards the
|
||||
/// postings serializer.
|
||||
pub(crate) fn serialize_postings<C: Codec>(
|
||||
pub(crate) fn serialize_postings(
|
||||
ctx: IndexingContext,
|
||||
schema: Schema,
|
||||
per_field_postings_writers: &PerFieldPostingsWriter,
|
||||
fieldnorm_readers: FieldNormReaders,
|
||||
serializer: &mut InvertedIndexSerializer<C>,
|
||||
serializer: &mut InvertedIndexSerializer,
|
||||
) -> crate::Result<()> {
|
||||
// Replace unordered ids by ordered ids to be able to sort
|
||||
let unordered_id_to_ordered_id: Vec<OrderedPathId> =
|
||||
@@ -104,141 +100,6 @@ pub(crate) struct IndexingPosition {
|
||||
pub end_position: u32,
|
||||
}
|
||||
|
||||
pub enum PostingsWriterEnum {
|
||||
DocId(SpecializedPostingsWriter<DocIdRecorder>),
|
||||
DocIdTf(SpecializedPostingsWriter<TermFrequencyRecorder>),
|
||||
DocTfAndPosition(SpecializedPostingsWriter<TfAndPositionRecorder>),
|
||||
JsonDocId(JsonPostingsWriter<DocIdRecorder>),
|
||||
JsonDocIdTf(JsonPostingsWriter<TermFrequencyRecorder>),
|
||||
JsonDocTfAndPosition(JsonPostingsWriter<TfAndPositionRecorder>),
|
||||
}
|
||||
|
||||
impl From<SpecializedPostingsWriter<DocIdRecorder>> for PostingsWriterEnum {
|
||||
fn from(doc_id_recorder_writer: SpecializedPostingsWriter<DocIdRecorder>) -> Self {
|
||||
PostingsWriterEnum::DocId(doc_id_recorder_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SpecializedPostingsWriter<TermFrequencyRecorder>> for PostingsWriterEnum {
|
||||
fn from(doc_id_tf_recorder_writer: SpecializedPostingsWriter<TermFrequencyRecorder>) -> Self {
|
||||
PostingsWriterEnum::DocIdTf(doc_id_tf_recorder_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SpecializedPostingsWriter<TfAndPositionRecorder>> for PostingsWriterEnum {
|
||||
fn from(
|
||||
doc_id_tf_and_positions_recorder_writer: SpecializedPostingsWriter<TfAndPositionRecorder>,
|
||||
) -> Self {
|
||||
PostingsWriterEnum::DocTfAndPosition(doc_id_tf_and_positions_recorder_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<JsonPostingsWriter<DocIdRecorder>> for PostingsWriterEnum {
|
||||
fn from(doc_id_recorder_writer: JsonPostingsWriter<DocIdRecorder>) -> Self {
|
||||
PostingsWriterEnum::JsonDocId(doc_id_recorder_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<JsonPostingsWriter<TermFrequencyRecorder>> for PostingsWriterEnum {
|
||||
fn from(doc_id_tf_recorder_writer: JsonPostingsWriter<TermFrequencyRecorder>) -> Self {
|
||||
PostingsWriterEnum::JsonDocIdTf(doc_id_tf_recorder_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<JsonPostingsWriter<TfAndPositionRecorder>> for PostingsWriterEnum {
|
||||
fn from(
|
||||
doc_id_tf_and_positions_recorder_writer: JsonPostingsWriter<TfAndPositionRecorder>,
|
||||
) -> Self {
|
||||
PostingsWriterEnum::JsonDocTfAndPosition(doc_id_tf_and_positions_recorder_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl PostingsWriter for PostingsWriterEnum {
|
||||
fn subscribe(&mut self, doc: DocId, pos: u32, term: &IndexingTerm, ctx: &mut IndexingContext) {
|
||||
match self {
|
||||
PostingsWriterEnum::DocId(writer) => writer.subscribe(doc, pos, term, ctx),
|
||||
PostingsWriterEnum::DocIdTf(writer) => writer.subscribe(doc, pos, term, ctx),
|
||||
PostingsWriterEnum::DocTfAndPosition(writer) => writer.subscribe(doc, pos, term, ctx),
|
||||
PostingsWriterEnum::JsonDocId(writer) => writer.subscribe(doc, pos, term, ctx),
|
||||
PostingsWriterEnum::JsonDocIdTf(writer) => writer.subscribe(doc, pos, term, ctx),
|
||||
PostingsWriterEnum::JsonDocTfAndPosition(writer) => {
|
||||
writer.subscribe(doc, pos, term, ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn serialize<C: Codec>(
|
||||
&self,
|
||||
term_addrs: &[(Field, OrderedPathId, &[u8], Addr)],
|
||||
ordered_id_to_path: &[&str],
|
||||
ctx: &IndexingContext,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
) -> io::Result<()> {
|
||||
match self {
|
||||
PostingsWriterEnum::DocId(writer) => {
|
||||
writer.serialize(term_addrs, ordered_id_to_path, ctx, serializer)
|
||||
}
|
||||
PostingsWriterEnum::DocIdTf(writer) => {
|
||||
writer.serialize(term_addrs, ordered_id_to_path, ctx, serializer)
|
||||
}
|
||||
PostingsWriterEnum::DocTfAndPosition(writer) => {
|
||||
writer.serialize(term_addrs, ordered_id_to_path, ctx, serializer)
|
||||
}
|
||||
PostingsWriterEnum::JsonDocId(writer) => {
|
||||
writer.serialize(term_addrs, ordered_id_to_path, ctx, serializer)
|
||||
}
|
||||
PostingsWriterEnum::JsonDocIdTf(writer) => {
|
||||
writer.serialize(term_addrs, ordered_id_to_path, ctx, serializer)
|
||||
}
|
||||
PostingsWriterEnum::JsonDocTfAndPosition(writer) => {
|
||||
writer.serialize(term_addrs, ordered_id_to_path, ctx, serializer)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tokenize a text and subscribe all of its token.
|
||||
fn index_text(
|
||||
&mut self,
|
||||
doc_id: DocId,
|
||||
token_stream: &mut dyn TokenStream,
|
||||
term_buffer: &mut IndexingTerm,
|
||||
ctx: &mut IndexingContext,
|
||||
indexing_position: &mut IndexingPosition,
|
||||
) {
|
||||
match self {
|
||||
PostingsWriterEnum::DocId(writer) => {
|
||||
writer.index_text(doc_id, token_stream, term_buffer, ctx, indexing_position)
|
||||
}
|
||||
PostingsWriterEnum::DocIdTf(writer) => {
|
||||
writer.index_text(doc_id, token_stream, term_buffer, ctx, indexing_position)
|
||||
}
|
||||
PostingsWriterEnum::DocTfAndPosition(writer) => {
|
||||
writer.index_text(doc_id, token_stream, term_buffer, ctx, indexing_position)
|
||||
}
|
||||
PostingsWriterEnum::JsonDocId(writer) => {
|
||||
writer.index_text(doc_id, token_stream, term_buffer, ctx, indexing_position)
|
||||
}
|
||||
PostingsWriterEnum::JsonDocIdTf(writer) => {
|
||||
writer.index_text(doc_id, token_stream, term_buffer, ctx, indexing_position)
|
||||
}
|
||||
PostingsWriterEnum::JsonDocTfAndPosition(writer) => {
|
||||
writer.index_text(doc_id, token_stream, term_buffer, ctx, indexing_position)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn total_num_tokens(&self) -> u64 {
|
||||
match self {
|
||||
PostingsWriterEnum::DocId(writer) => writer.total_num_tokens(),
|
||||
PostingsWriterEnum::DocIdTf(writer) => writer.total_num_tokens(),
|
||||
PostingsWriterEnum::DocTfAndPosition(writer) => writer.total_num_tokens(),
|
||||
PostingsWriterEnum::JsonDocId(writer) => writer.total_num_tokens(),
|
||||
PostingsWriterEnum::JsonDocIdTf(writer) => writer.total_num_tokens(),
|
||||
PostingsWriterEnum::JsonDocTfAndPosition(writer) => writer.total_num_tokens(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The `PostingsWriter` is in charge of receiving documenting
|
||||
/// and building a `Segment` in anonymous memory.
|
||||
///
|
||||
@@ -255,12 +116,12 @@ pub(crate) trait PostingsWriter: Send + Sync {
|
||||
|
||||
/// Serializes the postings on disk.
|
||||
/// The actual serialization format is handled by the `PostingsSerializer`.
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
term_addrs: &[(Field, OrderedPathId, &[u8], Addr)],
|
||||
ordered_id_to_path: &[&str],
|
||||
ctx: &IndexingContext,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer,
|
||||
) -> io::Result<()>;
|
||||
|
||||
/// Tokenize a text and subscribe all of its token.
|
||||
@@ -310,14 +171,22 @@ pub(crate) struct SpecializedPostingsWriter<Rec: Recorder> {
|
||||
_recorder_type: PhantomData<Rec>,
|
||||
}
|
||||
|
||||
impl<Rec: Recorder> From<SpecializedPostingsWriter<Rec>> for Box<dyn PostingsWriter> {
|
||||
fn from(
|
||||
specialized_postings_writer: SpecializedPostingsWriter<Rec>,
|
||||
) -> Box<dyn PostingsWriter> {
|
||||
Box::new(specialized_postings_writer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<Rec: Recorder> SpecializedPostingsWriter<Rec> {
|
||||
#[inline]
|
||||
pub(crate) fn serialize_one_term<C: Codec>(
|
||||
pub(crate) fn serialize_one_term(
|
||||
term: &[u8],
|
||||
addr: Addr,
|
||||
buffer_lender: &mut BufferLender,
|
||||
ctx: &IndexingContext,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer,
|
||||
) -> io::Result<()> {
|
||||
let recorder: Rec = ctx.term_index.read(addr);
|
||||
let term_doc_freq = recorder.term_doc_freq().unwrap_or(0u32);
|
||||
@@ -358,12 +227,12 @@ impl<Rec: Recorder> PostingsWriter for SpecializedPostingsWriter<Rec> {
|
||||
});
|
||||
}
|
||||
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
term_addrs: &[(Field, OrderedPathId, &[u8], Addr)],
|
||||
_ordered_id_to_path: &[&str],
|
||||
ctx: &IndexingContext,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer,
|
||||
) -> io::Result<()> {
|
||||
let mut buffer_lender = BufferLender::default();
|
||||
for (_field, _path_id, term, addr) in term_addrs {
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
use common::read_u32_vint;
|
||||
use stacker::{ExpUnrolledLinkedList, MemoryArena};
|
||||
|
||||
use crate::codec::Codec;
|
||||
use crate::postings::FieldSerializer;
|
||||
use crate::DocId;
|
||||
|
||||
@@ -68,10 +67,10 @@ pub(crate) trait Recorder: Copy + Default + Send + Sync + 'static {
|
||||
/// Close the document. It will help record the term frequency.
|
||||
fn close_doc(&mut self, arena: &mut MemoryArena);
|
||||
/// Pushes the postings information to the serializer.
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer<'_>,
|
||||
buffer_lender: &mut BufferLender,
|
||||
);
|
||||
/// Returns the number of document containing this term.
|
||||
@@ -111,10 +110,10 @@ impl Recorder for DocIdRecorder {
|
||||
#[inline]
|
||||
fn close_doc(&mut self, _arena: &mut MemoryArena) {}
|
||||
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer<'_>,
|
||||
buffer_lender: &mut BufferLender,
|
||||
) {
|
||||
let buffer = buffer_lender.lend_u8();
|
||||
@@ -179,10 +178,10 @@ impl Recorder for TermFrequencyRecorder {
|
||||
self.current_tf = 0;
|
||||
}
|
||||
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer<'_>,
|
||||
buffer_lender: &mut BufferLender,
|
||||
) {
|
||||
let buffer = buffer_lender.lend_u8();
|
||||
@@ -236,10 +235,10 @@ impl Recorder for TfAndPositionRecorder {
|
||||
self.stack.writer(arena).write_u32_vint(POSITION_END);
|
||||
}
|
||||
|
||||
fn serialize<C: Codec>(
|
||||
fn serialize(
|
||||
&self,
|
||||
arena: &MemoryArena,
|
||||
serializer: &mut FieldSerializer<C>,
|
||||
serializer: &mut FieldSerializer<'_>,
|
||||
buffer_lender: &mut BufferLender,
|
||||
) {
|
||||
let (buffer_u8, buffer_positions) = buffer_lender.lend_all();
|
||||
|
||||
@@ -1,14 +1,11 @@
|
||||
use common::BitSet;
|
||||
use common::HasLen;
|
||||
|
||||
use super::BlockSegmentPostings;
|
||||
use crate::codec::postings::PostingsWithBlockMax;
|
||||
use crate::docset::DocSet;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::fastfield::AliveBitSet;
|
||||
use crate::positions::PositionReader;
|
||||
use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
|
||||
use crate::postings::{DocFreq, Postings};
|
||||
use crate::query::Bm25Weight;
|
||||
use crate::{DocId, Score};
|
||||
use crate::postings::{BlockSegmentPostings, Postings};
|
||||
use crate::{DocId, TERMINATED};
|
||||
|
||||
/// `SegmentPostings` represents the inverted list or postings associated with
|
||||
/// a term in a `Segment`.
|
||||
@@ -32,6 +29,31 @@ impl SegmentPostings {
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the number of non-deleted documents.
|
||||
///
|
||||
/// This method will clone and scan through the posting lists.
|
||||
/// (this is a rather expensive operation).
|
||||
pub fn doc_freq_given_deletes(&self, alive_bitset: &AliveBitSet) -> u32 {
|
||||
let mut docset = self.clone();
|
||||
let mut doc_freq = 0;
|
||||
loop {
|
||||
let doc = docset.doc();
|
||||
if doc == TERMINATED {
|
||||
return doc_freq;
|
||||
}
|
||||
if alive_bitset.is_alive(doc) {
|
||||
doc_freq += 1u32;
|
||||
}
|
||||
docset.advance();
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the overall number of documents in the block postings.
|
||||
/// It does not take in account whether documents are deleted or not.
|
||||
pub fn doc_freq(&self) -> u32 {
|
||||
self.block_cursor.doc_freq()
|
||||
}
|
||||
|
||||
/// Creates a segment postings object with the given documents
|
||||
/// and no frequency encoded.
|
||||
///
|
||||
@@ -42,19 +64,13 @@ impl SegmentPostings {
|
||||
/// buffer with the serialized data.
|
||||
#[cfg(test)]
|
||||
pub fn create_from_docs(docs: &[u32]) -> SegmentPostings {
|
||||
use common::OwnedBytes;
|
||||
|
||||
use crate::directory::FileSlice;
|
||||
use crate::postings::serializer::PostingsSerializer;
|
||||
use crate::schema::IndexRecordOption;
|
||||
let mut buffer = Vec::new();
|
||||
{
|
||||
use crate::codec::postings::PostingsSerializer;
|
||||
|
||||
let mut postings_serializer =
|
||||
crate::codec::standard::postings::StandardPostingsSerializer::new(
|
||||
0.0,
|
||||
IndexRecordOption::Basic,
|
||||
None,
|
||||
);
|
||||
PostingsSerializer::new(0.0, IndexRecordOption::Basic, None);
|
||||
postings_serializer.new_term(docs.len() as u32, false);
|
||||
for &doc in docs {
|
||||
postings_serializer.write_doc(doc, 1u32);
|
||||
@@ -65,7 +81,7 @@ impl SegmentPostings {
|
||||
}
|
||||
let block_segment_postings = BlockSegmentPostings::open(
|
||||
docs.len() as u32,
|
||||
OwnedBytes::new(buffer),
|
||||
FileSlice::from(buffer),
|
||||
IndexRecordOption::Basic,
|
||||
IndexRecordOption::Basic,
|
||||
)
|
||||
@@ -79,11 +95,9 @@ impl SegmentPostings {
|
||||
doc_and_tfs: &[(u32, u32)],
|
||||
fieldnorms: Option<&[u32]>,
|
||||
) -> SegmentPostings {
|
||||
use common::OwnedBytes;
|
||||
|
||||
use crate::codec::postings::PostingsSerializer as _;
|
||||
use crate::codec::standard::postings::StandardPostingsSerializer;
|
||||
use crate::directory::FileSlice;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::serializer::PostingsSerializer;
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::Score;
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
@@ -100,7 +114,7 @@ impl SegmentPostings {
|
||||
total_num_tokens as Score / fieldnorms.len() as Score
|
||||
})
|
||||
.unwrap_or(0.0);
|
||||
let mut postings_serializer = StandardPostingsSerializer::new(
|
||||
let mut postings_serializer = PostingsSerializer::new(
|
||||
average_field_norm,
|
||||
IndexRecordOption::WithFreqs,
|
||||
fieldnorm_reader,
|
||||
@@ -114,7 +128,7 @@ impl SegmentPostings {
|
||||
.unwrap();
|
||||
let block_segment_postings = BlockSegmentPostings::open(
|
||||
doc_and_tfs.len() as u32,
|
||||
OwnedBytes::new(buffer),
|
||||
FileSlice::from(buffer),
|
||||
IndexRecordOption::WithFreqs,
|
||||
IndexRecordOption::WithFreqs,
|
||||
)
|
||||
@@ -144,6 +158,7 @@ impl DocSet for SegmentPostings {
|
||||
// next needs to be called a first time to point to the correct element.
|
||||
#[inline]
|
||||
fn advance(&mut self) -> DocId {
|
||||
debug_assert!(self.block_cursor.block_is_loaded());
|
||||
if self.cur == COMPRESSION_BLOCK_SIZE - 1 {
|
||||
self.cur = 0;
|
||||
self.block_cursor.advance();
|
||||
@@ -182,31 +197,13 @@ impl DocSet for SegmentPostings {
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> u32 {
|
||||
self.doc_freq().into()
|
||||
self.len() as u32
|
||||
}
|
||||
}
|
||||
|
||||
fn fill_bitset(&mut self, bitset: &mut BitSet) {
|
||||
let bitset_max_value: DocId = bitset.max_value();
|
||||
loop {
|
||||
let docs = self.block_cursor.docs();
|
||||
let Some(&last_doc) = docs.last() else {
|
||||
break;
|
||||
};
|
||||
if last_doc < bitset_max_value {
|
||||
// All docs are within the range of the bitset
|
||||
for &doc in docs {
|
||||
bitset.insert(doc);
|
||||
}
|
||||
} else {
|
||||
for &doc in docs {
|
||||
if doc < bitset_max_value {
|
||||
bitset.insert(doc);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
self.block_cursor.advance();
|
||||
}
|
||||
impl HasLen for SegmentPostings {
|
||||
fn len(&self) -> usize {
|
||||
self.block_cursor.doc_freq() as usize
|
||||
}
|
||||
}
|
||||
|
||||
@@ -232,13 +229,6 @@ impl Postings for SegmentPostings {
|
||||
self.block_cursor.freq(self.cur)
|
||||
}
|
||||
|
||||
/// Returns the overall number of documents in the block postings.
|
||||
/// It does not take in account whether documents are deleted or not.
|
||||
#[inline(always)]
|
||||
fn doc_freq(&self) -> DocFreq {
|
||||
DocFreq::Exact(self.block_cursor.doc_freq())
|
||||
}
|
||||
|
||||
fn append_positions_with_offset(&mut self, offset: u32, output: &mut Vec<u32>) {
|
||||
let term_freq = self.term_freq();
|
||||
let prev_len = output.len();
|
||||
@@ -262,42 +252,24 @@ impl Postings for SegmentPostings {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn has_freq(&self) -> bool {
|
||||
!self.block_cursor.freqs().is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl PostingsWithBlockMax for SegmentPostings {
|
||||
fn seek_block_max(
|
||||
&mut self,
|
||||
target_doc: crate::DocId,
|
||||
fieldnorm_reader: &FieldNormReader,
|
||||
similarity_weight: &Bm25Weight,
|
||||
) -> Score {
|
||||
self.block_cursor.seek_block_without_loading(target_doc);
|
||||
self.block_cursor
|
||||
.block_max_score(fieldnorm_reader, similarity_weight)
|
||||
}
|
||||
|
||||
fn last_doc_in_block(&self) -> crate::DocId {
|
||||
self.block_cursor.skip_reader().last_doc_in_block()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use common::HasLen;
|
||||
|
||||
use super::SegmentPostings;
|
||||
use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::postings::Postings;
|
||||
use crate::fastfield::AliveBitSet;
|
||||
use crate::postings::postings::Postings;
|
||||
|
||||
#[test]
|
||||
fn test_empty_segment_postings() {
|
||||
let mut postings = SegmentPostings::empty();
|
||||
assert_eq!(postings.doc(), TERMINATED);
|
||||
assert_eq!(postings.advance(), TERMINATED);
|
||||
assert_eq!(postings.advance(), TERMINATED);
|
||||
assert_eq!(postings.doc_freq(), crate::postings::DocFreq::Exact(0));
|
||||
assert_eq!(postings.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -312,4 +284,15 @@ mod tests {
|
||||
let postings = SegmentPostings::empty();
|
||||
assert_eq!(postings.term_freq(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_doc_freq() {
|
||||
let docs = SegmentPostings::create_from_docs(&[0, 2, 10]);
|
||||
assert_eq!(docs.doc_freq(), 3);
|
||||
let alive_bitset = AliveBitSet::for_test_from_deleted_docs(&[2], 12);
|
||||
assert_eq!(docs.doc_freq_given_deletes(&alive_bitset), 2);
|
||||
let all_deleted =
|
||||
AliveBitSet::for_test_from_deleted_docs(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 12);
|
||||
assert_eq!(docs.doc_freq_given_deletes(&all_deleted), 0);
|
||||
}
|
||||
}
|
||||
@@ -1,15 +1,17 @@
|
||||
use std::cmp::Ordering;
|
||||
use std::io::{self, Write};
|
||||
|
||||
use common::{BinarySerializable, CountingWriter};
|
||||
use common::{BinarySerializable, CountingWriter, VInt};
|
||||
|
||||
use super::TermInfo;
|
||||
use crate::codec::postings::PostingsSerializer;
|
||||
use crate::codec::Codec;
|
||||
use crate::directory::{CompositeWrite, WritePtr};
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::index::Segment;
|
||||
use crate::positions::PositionSerializer;
|
||||
use crate::schema::{Field, FieldEntry, FieldType, IndexRecordOption, Schema};
|
||||
use crate::postings::compression::{BlockEncoder, VIntEncoder, COMPRESSION_BLOCK_SIZE};
|
||||
use crate::postings::skip::SkipSerializer;
|
||||
use crate::query::Bm25Weight;
|
||||
use crate::schema::{Field, FieldEntry, IndexRecordOption, Schema};
|
||||
use crate::termdict::TermDictionaryBuilder;
|
||||
use crate::{DocId, Score};
|
||||
|
||||
@@ -44,27 +46,22 @@ use crate::{DocId, Score};
|
||||
///
|
||||
/// A description of the serialization format is
|
||||
/// [available here](https://fulmicoton.gitbooks.io/tantivy-doc/content/inverted-index.html).
|
||||
pub struct InvertedIndexSerializer<C: Codec> {
|
||||
pub struct InvertedIndexSerializer {
|
||||
terms_write: CompositeWrite<WritePtr>,
|
||||
postings_write: CompositeWrite<WritePtr>,
|
||||
positions_write: CompositeWrite<WritePtr>,
|
||||
schema: Schema,
|
||||
codec: C,
|
||||
}
|
||||
|
||||
use crate::codec::postings::PostingsCodec;
|
||||
|
||||
impl<C: Codec> InvertedIndexSerializer<C> {
|
||||
impl InvertedIndexSerializer {
|
||||
/// Open a new `InvertedIndexSerializer` for the given segment
|
||||
pub fn open(segment: &mut Segment<C>) -> crate::Result<InvertedIndexSerializer<C>> {
|
||||
pub fn open(segment: &mut Segment) -> crate::Result<InvertedIndexSerializer> {
|
||||
use crate::index::SegmentComponent::{Positions, Postings, Terms};
|
||||
let codec = segment.index().codec().clone();
|
||||
let inv_index_serializer = InvertedIndexSerializer {
|
||||
terms_write: CompositeWrite::wrap(segment.open_write(Terms)?),
|
||||
postings_write: CompositeWrite::wrap(segment.open_write(Postings)?),
|
||||
positions_write: CompositeWrite::wrap(segment.open_write(Positions)?),
|
||||
schema: segment.schema(),
|
||||
codec,
|
||||
};
|
||||
Ok(inv_index_serializer)
|
||||
}
|
||||
@@ -78,19 +75,22 @@ impl<C: Codec> InvertedIndexSerializer<C> {
|
||||
field: Field,
|
||||
total_num_tokens: u64,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
) -> io::Result<FieldSerializer<'_, C>> {
|
||||
) -> io::Result<FieldSerializer<'_>> {
|
||||
let field_entry: &FieldEntry = self.schema.get_field_entry(field);
|
||||
let term_dictionary_write = self.terms_write.for_field(field);
|
||||
let postings_write = self.postings_write.for_field(field);
|
||||
let positions_write = self.positions_write.for_field(field);
|
||||
let index_record_option = field_entry
|
||||
.field_type()
|
||||
.index_record_option()
|
||||
.unwrap_or(IndexRecordOption::Basic);
|
||||
FieldSerializer::create(
|
||||
field_entry.field_type(),
|
||||
index_record_option,
|
||||
total_num_tokens,
|
||||
term_dictionary_write,
|
||||
postings_write,
|
||||
positions_write,
|
||||
fieldnorm_reader,
|
||||
&self.codec,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -105,40 +105,34 @@ impl<C: Codec> InvertedIndexSerializer<C> {
|
||||
|
||||
/// The field serializer is in charge of
|
||||
/// the serialization of a specific field.
|
||||
pub struct FieldSerializer<'a, C: Codec> {
|
||||
term_dictionary_builder: TermDictionaryBuilder<&'a mut CountingWriter<WritePtr>>,
|
||||
postings_serializer: <C::PostingsCodec as PostingsCodec>::PostingsSerializer,
|
||||
positions_serializer_opt: Option<PositionSerializer<&'a mut CountingWriter<WritePtr>>>,
|
||||
pub struct FieldSerializer<'a, W: Write = WritePtr> {
|
||||
term_dictionary_builder: TermDictionaryBuilder<&'a mut CountingWriter<W>>,
|
||||
postings_serializer: PostingsSerializer,
|
||||
positions_serializer_opt: Option<PositionSerializer<&'a mut CountingWriter<W>>>,
|
||||
current_term_info: TermInfo,
|
||||
term_open: bool,
|
||||
postings_write: &'a mut CountingWriter<WritePtr>,
|
||||
postings_write: &'a mut CountingWriter<W>,
|
||||
postings_start_offset: u64,
|
||||
}
|
||||
|
||||
impl<'a, C: Codec> FieldSerializer<'a, C> {
|
||||
fn create(
|
||||
field_type: &FieldType,
|
||||
impl<'a, W: Write> FieldSerializer<'a, W> {
|
||||
/// Creates a new `FieldSerializer` for the given field type.
|
||||
pub fn create(
|
||||
index_record_option: IndexRecordOption,
|
||||
total_num_tokens: u64,
|
||||
term_dictionary_write: &'a mut CountingWriter<WritePtr>,
|
||||
postings_write: &'a mut CountingWriter<WritePtr>,
|
||||
positions_write: &'a mut CountingWriter<WritePtr>,
|
||||
term_dictionary_write: &'a mut CountingWriter<W>,
|
||||
postings_write: &'a mut CountingWriter<W>,
|
||||
positions_write: &'a mut CountingWriter<W>,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
codec: &C,
|
||||
) -> io::Result<FieldSerializer<'a, C>> {
|
||||
let index_record_option = field_type
|
||||
.index_record_option()
|
||||
.unwrap_or(IndexRecordOption::Basic);
|
||||
) -> io::Result<FieldSerializer<'a, W>> {
|
||||
total_num_tokens.serialize(postings_write)?;
|
||||
let term_dictionary_builder = TermDictionaryBuilder::create(term_dictionary_write)?;
|
||||
let average_fieldnorm = fieldnorm_reader
|
||||
.as_ref()
|
||||
.map(|ff_reader| total_num_tokens as Score / ff_reader.num_docs() as Score)
|
||||
.unwrap_or(0.0);
|
||||
let postings_serializer = codec.postings_codec().new_serializer(
|
||||
average_fieldnorm,
|
||||
index_record_option,
|
||||
fieldnorm_reader,
|
||||
);
|
||||
let postings_serializer =
|
||||
PostingsSerializer::new(average_fieldnorm, index_record_option, fieldnorm_reader);
|
||||
let positions_serializer_opt = if index_record_option.has_positions() {
|
||||
Some(PositionSerializer::new(positions_write))
|
||||
} else {
|
||||
@@ -191,6 +185,7 @@ impl<'a, C: Codec> FieldSerializer<'a, C> {
|
||||
"Called new_term, while the previous term was not closed."
|
||||
);
|
||||
self.term_open = true;
|
||||
self.postings_serializer.clear();
|
||||
self.current_term_info = self.current_term_info();
|
||||
self.term_dictionary_builder.insert_key(term)?;
|
||||
self.postings_serializer
|
||||
@@ -259,3 +254,234 @@ impl<'a, C: Codec> FieldSerializer<'a, C> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
struct Block {
|
||||
doc_ids: [DocId; COMPRESSION_BLOCK_SIZE],
|
||||
term_freqs: [u32; COMPRESSION_BLOCK_SIZE],
|
||||
len: usize,
|
||||
}
|
||||
|
||||
impl Block {
|
||||
fn new() -> Self {
|
||||
Block {
|
||||
doc_ids: [0u32; COMPRESSION_BLOCK_SIZE],
|
||||
term_freqs: [0u32; COMPRESSION_BLOCK_SIZE],
|
||||
len: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn doc_ids(&self) -> &[DocId] {
|
||||
&self.doc_ids[..self.len]
|
||||
}
|
||||
|
||||
fn term_freqs(&self) -> &[u32] {
|
||||
&self.term_freqs[..self.len]
|
||||
}
|
||||
|
||||
fn clear(&mut self) {
|
||||
self.len = 0;
|
||||
}
|
||||
|
||||
fn append_doc(&mut self, doc: DocId, term_freq: u32) {
|
||||
let len = self.len;
|
||||
self.doc_ids[len] = doc;
|
||||
self.term_freqs[len] = term_freq;
|
||||
self.len = len + 1;
|
||||
}
|
||||
|
||||
fn is_full(&self) -> bool {
|
||||
self.len == COMPRESSION_BLOCK_SIZE
|
||||
}
|
||||
|
||||
fn is_empty(&self) -> bool {
|
||||
self.len == 0
|
||||
}
|
||||
|
||||
fn last_doc(&self) -> DocId {
|
||||
assert_eq!(self.len, COMPRESSION_BLOCK_SIZE);
|
||||
self.doc_ids[COMPRESSION_BLOCK_SIZE - 1]
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializer for postings lists.
|
||||
pub struct PostingsSerializer {
|
||||
last_doc_id_encoded: u32,
|
||||
|
||||
block_encoder: BlockEncoder,
|
||||
block: Box<Block>,
|
||||
|
||||
postings_write: Vec<u8>,
|
||||
skip_write: SkipSerializer,
|
||||
|
||||
mode: IndexRecordOption,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
|
||||
bm25_weight: Option<Bm25Weight>,
|
||||
avg_fieldnorm: Score, /* Average number of term in the field for that segment.
|
||||
* this value is used to compute the block wand information. */
|
||||
term_has_freq: bool,
|
||||
}
|
||||
|
||||
impl PostingsSerializer {
|
||||
/// Creates a new `PostingsSerializer`.
|
||||
/// * avg_fieldnorm - average field norm for the field being serialized.
|
||||
/// * mode - indexing options for the field being serialized.
|
||||
pub fn new(
|
||||
avg_fieldnorm: Score,
|
||||
mode: IndexRecordOption,
|
||||
fieldnorm_reader: Option<FieldNormReader>,
|
||||
) -> PostingsSerializer {
|
||||
PostingsSerializer {
|
||||
block_encoder: BlockEncoder::new(),
|
||||
block: Box::new(Block::new()),
|
||||
|
||||
postings_write: Vec::new(),
|
||||
skip_write: SkipSerializer::new(),
|
||||
|
||||
last_doc_id_encoded: 0u32,
|
||||
mode,
|
||||
|
||||
fieldnorm_reader,
|
||||
bm25_weight: None,
|
||||
avg_fieldnorm,
|
||||
term_has_freq: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Starts the serialization for a new term.
|
||||
/// * term_doc_freq - the number of documents containing the term.
|
||||
pub fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool) {
|
||||
self.bm25_weight = None;
|
||||
|
||||
self.term_has_freq = self.mode.has_freq() && record_term_freq;
|
||||
if !self.term_has_freq {
|
||||
return;
|
||||
}
|
||||
|
||||
let num_docs_in_segment: u64 =
|
||||
if let Some(fieldnorm_reader) = self.fieldnorm_reader.as_ref() {
|
||||
fieldnorm_reader.num_docs() as u64
|
||||
} else {
|
||||
return;
|
||||
};
|
||||
|
||||
if num_docs_in_segment == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
self.bm25_weight = Some(Bm25Weight::for_one_term_without_explain(
|
||||
term_doc_freq as u64,
|
||||
num_docs_in_segment,
|
||||
self.avg_fieldnorm,
|
||||
));
|
||||
}
|
||||
|
||||
fn write_block(&mut self) {
|
||||
{
|
||||
// encode the doc ids
|
||||
let (num_bits, block_encoded): (u8, &[u8]) = self
|
||||
.block_encoder
|
||||
.compress_block_sorted(self.block.doc_ids(), self.last_doc_id_encoded);
|
||||
self.last_doc_id_encoded = self.block.last_doc();
|
||||
self.skip_write
|
||||
.write_doc(self.last_doc_id_encoded, num_bits);
|
||||
// last el block 0, offset block 1,
|
||||
self.postings_write.extend(block_encoded);
|
||||
}
|
||||
if self.term_has_freq {
|
||||
// encode the term frequencies
|
||||
let (num_bits, block_encoded): (u8, &[u8]) = self
|
||||
.block_encoder
|
||||
.compress_block_unsorted(self.block.term_freqs(), true);
|
||||
self.postings_write.extend(block_encoded);
|
||||
self.skip_write.write_term_freq(num_bits);
|
||||
if self.mode.has_positions() {
|
||||
// We serialize the sum of term freqs within the skip information
|
||||
// in order to navigate through positions.
|
||||
let sum_freq = self.block.term_freqs().iter().cloned().sum();
|
||||
self.skip_write.write_total_term_freq(sum_freq);
|
||||
}
|
||||
let mut blockwand_params = (0u8, 0u32);
|
||||
if let Some(bm25_weight) = self.bm25_weight.as_ref() {
|
||||
if let Some(fieldnorm_reader) = self.fieldnorm_reader.as_ref() {
|
||||
let docs = self.block.doc_ids().iter().cloned();
|
||||
let term_freqs = self.block.term_freqs().iter().cloned();
|
||||
let fieldnorms = docs.map(|doc| fieldnorm_reader.fieldnorm_id(doc));
|
||||
blockwand_params = fieldnorms
|
||||
.zip(term_freqs)
|
||||
.max_by(
|
||||
|(left_fieldnorm_id, left_term_freq),
|
||||
(right_fieldnorm_id, right_term_freq)| {
|
||||
let left_score =
|
||||
bm25_weight.tf_factor(*left_fieldnorm_id, *left_term_freq);
|
||||
let right_score =
|
||||
bm25_weight.tf_factor(*right_fieldnorm_id, *right_term_freq);
|
||||
left_score
|
||||
.partial_cmp(&right_score)
|
||||
.unwrap_or(Ordering::Equal)
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
let (fieldnorm_id, term_freq) = blockwand_params;
|
||||
self.skip_write.write_blockwand_max(fieldnorm_id, term_freq);
|
||||
}
|
||||
self.block.clear();
|
||||
}
|
||||
|
||||
/// Register that the given document contains the current term.
|
||||
/// * doc_id - the document id.
|
||||
/// * term_freq - the term frequency within the document.
|
||||
pub fn write_doc(&mut self, doc_id: DocId, term_freq: u32) {
|
||||
self.block.append_doc(doc_id, term_freq);
|
||||
if self.block.is_full() {
|
||||
self.write_block();
|
||||
}
|
||||
}
|
||||
|
||||
/// Finish the serialization for this term.
|
||||
pub fn close_term(
|
||||
&mut self,
|
||||
doc_freq: u32,
|
||||
output_write: &mut impl std::io::Write,
|
||||
) -> io::Result<()> {
|
||||
if !self.block.is_empty() {
|
||||
// we have doc ids waiting to be written
|
||||
// this happens when the number of doc ids is
|
||||
// not a perfect multiple of our block size.
|
||||
//
|
||||
// In that case, the remaining part is encoded
|
||||
// using variable int encoding.
|
||||
{
|
||||
let block_encoded = self
|
||||
.block_encoder
|
||||
.compress_vint_sorted(self.block.doc_ids(), self.last_doc_id_encoded);
|
||||
self.postings_write.write_all(block_encoded)?;
|
||||
}
|
||||
// ... Idem for term frequencies
|
||||
if self.term_has_freq {
|
||||
let block_encoded = self
|
||||
.block_encoder
|
||||
.compress_vint_unsorted(self.block.term_freqs());
|
||||
self.postings_write.write_all(block_encoded)?;
|
||||
}
|
||||
self.block.clear();
|
||||
}
|
||||
if doc_freq >= COMPRESSION_BLOCK_SIZE as u32 {
|
||||
let skip_data = self.skip_write.data();
|
||||
VInt(skip_data.len() as u64).serialize(output_write)?;
|
||||
output_write.write_all(skip_data)?;
|
||||
}
|
||||
output_write.write_all(&self.postings_write[..])?;
|
||||
self.skip_write.clear();
|
||||
self.postings_write.clear();
|
||||
self.bm25_weight = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn clear(&mut self) {
|
||||
self.block.clear();
|
||||
self.last_doc_id_encoded = 0;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -146,6 +146,28 @@ impl SkipReader {
|
||||
skip_reader
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn has_remaining_docs(&self) -> bool {
|
||||
self.remaining_docs != 0
|
||||
}
|
||||
|
||||
pub fn reset(&mut self, data: OwnedBytes, doc_freq: u32) {
|
||||
self.last_doc_in_block = if doc_freq >= COMPRESSION_BLOCK_SIZE as u32 {
|
||||
0
|
||||
} else {
|
||||
TERMINATED
|
||||
};
|
||||
self.last_doc_in_previous_block = 0u32;
|
||||
self.owned_read = data;
|
||||
self.block_info = BlockInfo::VInt { num_docs: doc_freq };
|
||||
self.byte_offset = 0;
|
||||
self.remaining_docs = doc_freq;
|
||||
self.position_offset = 0u64;
|
||||
if doc_freq >= COMPRESSION_BLOCK_SIZE as u32 {
|
||||
self.read_block_info();
|
||||
}
|
||||
}
|
||||
|
||||
// Returns the block max score for this block if available.
|
||||
//
|
||||
// The block max score is available for all full bitpacked block,
|
||||
@@ -165,6 +187,12 @@ impl SkipReader {
|
||||
self.last_doc_in_block
|
||||
}
|
||||
|
||||
/// Number of docs from the start of the current block to the end of the postings
|
||||
/// (i.e. the current block plus every block after it).
|
||||
pub(crate) fn remaining_docs(&self) -> u32 {
|
||||
self.remaining_docs
|
||||
}
|
||||
|
||||
pub fn position_offset(&self) -> u64 {
|
||||
self.position_offset
|
||||
}
|
||||
@@ -2,7 +2,7 @@ use crate::docset::{DocSet, COLLECT_BLOCK_BUFFER_LEN, TERMINATED};
|
||||
use crate::index::SegmentReader;
|
||||
use crate::query::boost_query::BoostScorer;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::{box_scorer, EnableScoring, Explanation, Query, Scorer, Weight};
|
||||
use crate::query::{EnableScoring, Explanation, Query, Scorer, Weight};
|
||||
use crate::{DocId, Score};
|
||||
|
||||
/// Query that matches all of the documents.
|
||||
@@ -24,9 +24,9 @@ impl Weight for AllWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
let all_scorer = AllScorer::new(reader.max_doc());
|
||||
if boost != 1.0 {
|
||||
Ok(box_scorer(BoostScorer::new(all_scorer, boost)))
|
||||
Ok(Box::new(BoostScorer::new(all_scorer, boost)))
|
||||
} else {
|
||||
Ok(box_scorer(all_scorer))
|
||||
Ok(Box::new(all_scorer))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ use crate::postings::TermInfo;
|
||||
use crate::query::{BitSetDocSet, ConstScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::{Field, IndexRecordOption};
|
||||
use crate::termdict::{TermDictionary, TermStreamer};
|
||||
use crate::{DocId, DocSet, Score, TantivyError};
|
||||
use crate::{DocId, Score, TantivyError};
|
||||
|
||||
/// A weight struct for Fuzzy Term and Regex Queries
|
||||
pub struct AutomatonWeight<A> {
|
||||
@@ -92,9 +92,18 @@ where
|
||||
let mut term_stream = self.automaton_stream(term_dict)?;
|
||||
while term_stream.advance() {
|
||||
let term_info = term_stream.value();
|
||||
let mut block_segment_postings =
|
||||
inverted_index.read_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
|
||||
block_segment_postings.fill_bitset(&mut doc_bitset);
|
||||
let mut block_segment_postings = inverted_index
|
||||
.read_block_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
|
||||
loop {
|
||||
let docs = block_segment_postings.docs();
|
||||
if docs.is_empty() {
|
||||
break;
|
||||
}
|
||||
for &doc in docs {
|
||||
doc_bitset.insert(doc);
|
||||
}
|
||||
block_segment_postings.advance();
|
||||
}
|
||||
}
|
||||
let doc_bitset = BitSetDocSet::from(doc_bitset);
|
||||
let const_scorer = ConstScorer::new(doc_bitset, boost);
|
||||
|
||||
@@ -24,13 +24,6 @@ impl BitSetDocSet {
|
||||
self.cursor_bucket = bucket_addr;
|
||||
self.cursor_tinybitset = self.docs.tinyset(bucket_addr);
|
||||
}
|
||||
|
||||
/// Returns the number of documents in the bitset.
|
||||
///
|
||||
/// This call is not free: it will bitcount the number of bits in the bitset.
|
||||
pub fn doc_freq(&self) -> u32 {
|
||||
self.docs.len() as u32
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BitSet> for BitSetDocSet {
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
use std::ops::{Deref, DerefMut};
|
||||
|
||||
use crate::codec::postings::PostingsWithBlockMax;
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::Scorer;
|
||||
use crate::{DocId, DocSet, Score, TERMINATED};
|
||||
@@ -14,8 +13,8 @@ use crate::{DocId, DocSet, Score, TERMINATED};
|
||||
/// We always have `before_pivot_len` < `pivot_len`.
|
||||
///
|
||||
/// `None` is returned if we establish that no document can exceed the threshold.
|
||||
fn find_pivot_doc<TPostings: PostingsWithBlockMax>(
|
||||
term_scorers: &[TermScorerWithMaxScore<TPostings>],
|
||||
fn find_pivot_doc(
|
||||
term_scorers: &[TermScorerWithMaxScore],
|
||||
threshold: Score,
|
||||
) -> Option<(usize, usize, DocId)> {
|
||||
let mut max_score = 0.0;
|
||||
@@ -47,8 +46,8 @@ fn find_pivot_doc<TPostings: PostingsWithBlockMax>(
|
||||
/// the next doc candidate defined by the min of `last_doc_in_block + 1` for
|
||||
/// scorer in scorers[..pivot_len] and `scorer.doc()` for scorer in scorers[pivot_len..].
|
||||
/// Note: before and after calling this method, scorers need to be sorted by their `.doc()`.
|
||||
fn block_max_was_too_low_advance_one_scorer<TPostings: PostingsWithBlockMax>(
|
||||
scorers: &mut [TermScorerWithMaxScore<TPostings>],
|
||||
fn block_max_was_too_low_advance_one_scorer(
|
||||
scorers: &mut [TermScorerWithMaxScore],
|
||||
pivot_len: usize,
|
||||
) {
|
||||
debug_assert!(scorers.iter().map(|scorer| scorer.doc()).is_sorted());
|
||||
@@ -83,10 +82,7 @@ fn block_max_was_too_low_advance_one_scorer<TPostings: PostingsWithBlockMax>(
|
||||
// Given a list of term_scorers and a `ord` and assuming that `term_scorers[ord]` is sorted
|
||||
// except term_scorers[ord] that might be in advance compared to its ranks,
|
||||
// bubble up term_scorers[ord] in order to restore the ordering.
|
||||
fn restore_ordering<TPostings: PostingsWithBlockMax>(
|
||||
term_scorers: &mut [TermScorerWithMaxScore<TPostings>],
|
||||
ord: usize,
|
||||
) {
|
||||
fn restore_ordering(term_scorers: &mut [TermScorerWithMaxScore], ord: usize) {
|
||||
let doc = term_scorers[ord].doc();
|
||||
for i in ord + 1..term_scorers.len() {
|
||||
if term_scorers[i].doc() >= doc {
|
||||
@@ -101,10 +97,9 @@ fn restore_ordering<TPostings: PostingsWithBlockMax>(
|
||||
// If this works, return true.
|
||||
// If this fails (ie: one of the term_scorer does not contain `pivot_doc` and seek goes past the
|
||||
// pivot), reorder the term_scorers to ensure the list is still sorted and returns `false`.
|
||||
// If a term_scorer reach TERMINATED in the process return false remove the term_scorer and
|
||||
// return.
|
||||
fn align_scorers<TPostings: PostingsWithBlockMax>(
|
||||
term_scorers: &mut Vec<TermScorerWithMaxScore<TPostings>>,
|
||||
// If a term_scorer reach TERMINATED in the process return false remove the term_scorer and return.
|
||||
fn align_scorers(
|
||||
term_scorers: &mut Vec<TermScorerWithMaxScore>,
|
||||
pivot_doc: DocId,
|
||||
before_pivot_len: usize,
|
||||
) -> bool {
|
||||
@@ -131,10 +126,7 @@ fn align_scorers<TPostings: PostingsWithBlockMax>(
|
||||
// Assumes terms_scorers[..pivot_len] are positioned on the same doc (pivot_doc).
|
||||
// Advance term_scorers[..pivot_len] and out of these removes the terminated scores.
|
||||
// Restores the ordering of term_scorers.
|
||||
fn advance_all_scorers_on_pivot<TPostings: PostingsWithBlockMax>(
|
||||
term_scorers: &mut Vec<TermScorerWithMaxScore<TPostings>>,
|
||||
pivot_len: usize,
|
||||
) {
|
||||
fn advance_all_scorers_on_pivot(term_scorers: &mut Vec<TermScorerWithMaxScore>, pivot_len: usize) {
|
||||
for term_scorer in &mut term_scorers[..pivot_len] {
|
||||
term_scorer.advance();
|
||||
}
|
||||
@@ -153,8 +145,8 @@ fn advance_all_scorers_on_pivot<TPostings: PostingsWithBlockMax>(
|
||||
/// Implements the WAND (Weak AND) algorithm for dynamic pruning
|
||||
/// described in the paper "Faster Top-k Document Retrieval Using Block-Max Indexes".
|
||||
/// Link: <http://engineering.nyu.edu/~suel/papers/bmw.pdf>
|
||||
pub fn block_wand<TPostings: PostingsWithBlockMax>(
|
||||
mut scorers: Vec<TermScorer<TPostings>>,
|
||||
pub fn block_wand(
|
||||
mut scorers: Vec<TermScorer>,
|
||||
mut threshold: Score,
|
||||
callback: &mut dyn FnMut(u32, Score) -> Score,
|
||||
) {
|
||||
@@ -163,7 +155,7 @@ pub fn block_wand<TPostings: PostingsWithBlockMax>(
|
||||
let scorer = scorers.pop().unwrap();
|
||||
return block_wand_single_scorer(scorer, threshold, callback);
|
||||
}
|
||||
let mut scorers: Vec<TermScorerWithMaxScore<TPostings>> = scorers
|
||||
let mut scorers: Vec<TermScorerWithMaxScore> = scorers
|
||||
.iter_mut()
|
||||
.map(TermScorerWithMaxScore::from)
|
||||
.collect();
|
||||
@@ -178,7 +170,10 @@ pub fn block_wand<TPostings: PostingsWithBlockMax>(
|
||||
|
||||
let block_max_score_upperbound: Score = scorers[..pivot_len]
|
||||
.iter_mut()
|
||||
.map(|scorer| scorer.seek_block_max(pivot_doc))
|
||||
.map(|scorer| {
|
||||
scorer.seek_block(pivot_doc);
|
||||
scorer.block_max_score()
|
||||
})
|
||||
.sum();
|
||||
|
||||
// Beware after shallow advance, skip readers can be in advance compared to
|
||||
@@ -229,22 +224,21 @@ pub fn block_wand<TPostings: PostingsWithBlockMax>(
|
||||
/// - On a block, advance until the end and execute `callback` when the doc score is greater or
|
||||
/// equal to the `threshold`.
|
||||
pub fn block_wand_single_scorer(
|
||||
mut scorer: TermScorer<impl PostingsWithBlockMax>,
|
||||
mut scorer: TermScorer,
|
||||
mut threshold: Score,
|
||||
callback: &mut dyn FnMut(u32, Score) -> Score,
|
||||
) {
|
||||
let mut doc = scorer.doc();
|
||||
let mut block_max_score = scorer.seek_block_max(doc);
|
||||
loop {
|
||||
// We position the scorer on a block that can reach
|
||||
// the threshold.
|
||||
while block_max_score < threshold {
|
||||
while scorer.block_max_score() <= threshold {
|
||||
let last_doc_in_block = scorer.last_doc_in_block();
|
||||
if last_doc_in_block == TERMINATED {
|
||||
return;
|
||||
}
|
||||
doc = last_doc_in_block + 1;
|
||||
block_max_score = scorer.seek_block_max(doc);
|
||||
scorer.seek_block(doc);
|
||||
}
|
||||
// Seek will effectively load that block.
|
||||
doc = scorer.seek(doc);
|
||||
@@ -266,33 +260,31 @@ pub fn block_wand_single_scorer(
|
||||
}
|
||||
}
|
||||
doc += 1;
|
||||
block_max_score = scorer.seek_block_max(doc);
|
||||
scorer.seek_block(doc);
|
||||
}
|
||||
}
|
||||
|
||||
struct TermScorerWithMaxScore<'a, TPostings: PostingsWithBlockMax> {
|
||||
scorer: &'a mut TermScorer<TPostings>,
|
||||
struct TermScorerWithMaxScore<'a> {
|
||||
scorer: &'a mut TermScorer,
|
||||
max_score: Score,
|
||||
}
|
||||
|
||||
impl<'a, TPostings: PostingsWithBlockMax> From<&'a mut TermScorer<TPostings>>
|
||||
for TermScorerWithMaxScore<'a, TPostings>
|
||||
{
|
||||
fn from(scorer: &'a mut TermScorer<TPostings>) -> Self {
|
||||
impl<'a> From<&'a mut TermScorer> for TermScorerWithMaxScore<'a> {
|
||||
fn from(scorer: &'a mut TermScorer) -> Self {
|
||||
let max_score = scorer.max_score();
|
||||
TermScorerWithMaxScore { scorer, max_score }
|
||||
}
|
||||
}
|
||||
|
||||
impl<TPostings: PostingsWithBlockMax> Deref for TermScorerWithMaxScore<'_, TPostings> {
|
||||
type Target = TermScorer<TPostings>;
|
||||
impl Deref for TermScorerWithMaxScore<'_> {
|
||||
type Target = TermScorer;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.scorer
|
||||
}
|
||||
}
|
||||
|
||||
impl<TPostings: PostingsWithBlockMax> DerefMut for TermScorerWithMaxScore<'_, TPostings> {
|
||||
impl DerefMut for TermScorerWithMaxScore<'_> {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
self.scorer
|
||||
}
|
||||
@@ -1,18 +1,25 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::codec::{ObjectSafeCodec, SumOrDoNothingCombiner};
|
||||
use crate::docset::COLLECT_BLOCK_BUFFER_LEN;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::postings::FreqReadingOption;
|
||||
use crate::query::disjunction::Disjunction;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::score_combiner::{DoNothingCombiner, ScoreCombiner};
|
||||
use crate::query::weight::for_each_docset_buffered;
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::weight::{for_each_docset_buffered, for_each_pruning_scorer, for_each_scorer};
|
||||
use crate::query::{
|
||||
box_scorer, intersect_scorers, AllScorer, BufferedUnionScorer, EmptyScorer, Exclude,
|
||||
Explanation, Occur, RequiredOptionalScorer, Scorer, SumCombiner, Weight,
|
||||
intersect_scorers, AllScorer, BufferedUnionScorer, EmptyScorer, Exclude, Explanation, Occur,
|
||||
RequiredOptionalScorer, Scorer, Weight,
|
||||
};
|
||||
use crate::{DocId, Score};
|
||||
|
||||
enum SpecializedScorer {
|
||||
TermUnion(Vec<TermScorer>),
|
||||
TermIntersection(Vec<TermScorer>),
|
||||
Other(Box<dyn Scorer>),
|
||||
}
|
||||
|
||||
fn scorer_disjunction<TScoreCombiner>(
|
||||
scorers: Vec<Box<dyn Scorer>>,
|
||||
score_combiner: TScoreCombiner,
|
||||
@@ -26,7 +33,7 @@ where
|
||||
if scorers.len() == 1 {
|
||||
return scorers.into_iter().next().unwrap(); // Safe unwrap.
|
||||
}
|
||||
box_scorer(Disjunction::new(
|
||||
Box::new(Disjunction::new(
|
||||
scorers,
|
||||
score_combiner,
|
||||
minimum_match_required,
|
||||
@@ -38,41 +45,66 @@ fn scorer_union<TScoreCombiner>(
|
||||
scorers: Vec<Box<dyn Scorer>>,
|
||||
score_combiner_fn: impl Fn() -> TScoreCombiner,
|
||||
num_docs: u32,
|
||||
codec: &dyn ObjectSafeCodec,
|
||||
) -> Box<dyn Scorer>
|
||||
) -> SpecializedScorer
|
||||
where
|
||||
TScoreCombiner: ScoreCombiner,
|
||||
{
|
||||
match scorers.len() {
|
||||
0 => box_scorer(EmptyScorer),
|
||||
1 => scorers.into_iter().next().unwrap(),
|
||||
_ => {
|
||||
let combiner_opt: Option<SumOrDoNothingCombiner> = if std::any::TypeId::of::<
|
||||
TScoreCombiner,
|
||||
>() == std::any::TypeId::of::<
|
||||
SumCombiner,
|
||||
>() {
|
||||
Some(SumOrDoNothingCombiner::Sum)
|
||||
} else if std::any::TypeId::of::<TScoreCombiner>()
|
||||
== std::any::TypeId::of::<DoNothingCombiner>()
|
||||
assert!(!scorers.is_empty());
|
||||
if scorers.len() == 1 && !scorers[0].is::<TermScorer>() {
|
||||
return SpecializedScorer::Other(scorers.into_iter().next().unwrap()); //< we checked the size beforehand
|
||||
}
|
||||
{
|
||||
let is_all_term_queries = scorers.iter().all(|scorer| scorer.is::<TermScorer>());
|
||||
if is_all_term_queries {
|
||||
let scorers: Vec<TermScorer> = scorers
|
||||
.into_iter()
|
||||
.map(|scorer| *(scorer.downcast::<TermScorer>().map_err(|_| ()).unwrap()))
|
||||
.collect();
|
||||
if scorers
|
||||
.iter()
|
||||
.all(|scorer| scorer.freq_reading_option() == FreqReadingOption::ReadFreq)
|
||||
{
|
||||
Some(SumOrDoNothingCombiner::DoNothing)
|
||||
// Block wand is only available if we read frequencies.
|
||||
return SpecializedScorer::TermUnion(scorers);
|
||||
} else if scorers.len() == 1 {
|
||||
// Single TermScorer without freq reading — unwrap directly.
|
||||
return SpecializedScorer::Other(Box::new(scorers.into_iter().next().unwrap()));
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(combiner) = combiner_opt {
|
||||
let scorer =
|
||||
codec.build_union_scorer_with_sum_combiner(scorers, num_docs, combiner);
|
||||
scorer
|
||||
} else {
|
||||
box_scorer(BufferedUnionScorer::build(
|
||||
return SpecializedScorer::Other(Box::new(BufferedUnionScorer::build(
|
||||
scorers,
|
||||
score_combiner_fn,
|
||||
num_docs,
|
||||
))
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
SpecializedScorer::Other(Box::new(BufferedUnionScorer::build(
|
||||
scorers,
|
||||
score_combiner_fn,
|
||||
num_docs,
|
||||
)))
|
||||
}
|
||||
|
||||
fn into_box_scorer<TScoreCombiner: ScoreCombiner>(
|
||||
scorer: SpecializedScorer,
|
||||
score_combiner_fn: impl Fn() -> TScoreCombiner,
|
||||
num_docs: u32,
|
||||
) -> Box<dyn Scorer> {
|
||||
match scorer {
|
||||
SpecializedScorer::TermUnion(term_scorers) => {
|
||||
let union_scorer =
|
||||
BufferedUnionScorer::build(term_scorers, score_combiner_fn, num_docs);
|
||||
Box::new(union_scorer)
|
||||
}
|
||||
SpecializedScorer::TermIntersection(term_scorers) => {
|
||||
let boxed_scorers: Vec<Box<dyn Scorer>> = term_scorers
|
||||
.into_iter()
|
||||
.map(|s| Box::new(s) as Box<dyn Scorer>)
|
||||
.collect();
|
||||
intersect_scorers(boxed_scorers, num_docs)
|
||||
}
|
||||
SpecializedScorer::Other(scorer) => scorer,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the effective MUST scorer, accounting for removed AllScorers.
|
||||
@@ -88,7 +120,7 @@ fn effective_must_scorer(
|
||||
if must_scorers.is_empty() {
|
||||
if removed_all_scorer_count > 0 {
|
||||
// Had AllScorer(s) only - all docs match
|
||||
Some(box_scorer(AllScorer::new(max_doc)))
|
||||
Some(Box::new(AllScorer::new(max_doc)))
|
||||
} else {
|
||||
// No MUST constraint at all
|
||||
None
|
||||
@@ -106,26 +138,28 @@ fn effective_must_scorer(
|
||||
/// When `scoring_enabled` is false, we can just return AllScorer alone since
|
||||
/// we don't need score contributions from the should_scorer.
|
||||
fn effective_should_scorer_for_union<TScoreCombiner: ScoreCombiner>(
|
||||
should_scorer: Box<dyn Scorer>,
|
||||
should_scorer: SpecializedScorer,
|
||||
removed_all_scorer_count: usize,
|
||||
max_doc: DocId,
|
||||
num_docs: u32,
|
||||
score_combiner_fn: impl Fn() -> TScoreCombiner,
|
||||
scoring_enabled: bool,
|
||||
) -> Box<dyn Scorer> {
|
||||
) -> SpecializedScorer {
|
||||
if removed_all_scorer_count > 0 {
|
||||
if scoring_enabled {
|
||||
// Need to union to get score contributions from both
|
||||
let all_scorers: Vec<Box<dyn Scorer>> =
|
||||
vec![should_scorer, box_scorer(AllScorer::new(max_doc))];
|
||||
box_scorer(BufferedUnionScorer::build(
|
||||
let all_scorers: Vec<Box<dyn Scorer>> = vec![
|
||||
into_box_scorer(should_scorer, &score_combiner_fn, num_docs),
|
||||
Box::new(AllScorer::new(max_doc)),
|
||||
];
|
||||
SpecializedScorer::Other(Box::new(BufferedUnionScorer::build(
|
||||
all_scorers,
|
||||
score_combiner_fn,
|
||||
num_docs,
|
||||
))
|
||||
)))
|
||||
} else {
|
||||
// Scoring disabled - AllScorer alone is sufficient
|
||||
box_scorer(AllScorer::new(max_doc))
|
||||
SpecializedScorer::Other(Box::new(AllScorer::new(max_doc)))
|
||||
}
|
||||
} else {
|
||||
should_scorer
|
||||
@@ -136,9 +170,9 @@ enum ShouldScorersCombinationMethod {
|
||||
// Should scorers are irrelevant.
|
||||
Ignored,
|
||||
// Only contributes to final score.
|
||||
Optional(Box<dyn Scorer>),
|
||||
Optional(SpecializedScorer),
|
||||
// Regardless of score, the should scorers may impact whether a document is matching or not.
|
||||
Required(Box<dyn Scorer>),
|
||||
Required(SpecializedScorer),
|
||||
}
|
||||
|
||||
/// Weight associated to the `BoolQuery`.
|
||||
@@ -200,7 +234,7 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
reader: &SegmentReader,
|
||||
boost: Score,
|
||||
score_combiner_fn: impl Fn() -> TComplexScoreCombiner,
|
||||
) -> crate::Result<Box<dyn Scorer>> {
|
||||
) -> crate::Result<SpecializedScorer> {
|
||||
let num_docs = reader.num_docs();
|
||||
let mut per_occur_scorers = self.per_occur_scorers(reader, boost)?;
|
||||
|
||||
@@ -210,7 +244,7 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
let must_special_scorer_counts = remove_and_count_all_and_empty_scorers(&mut must_scorers);
|
||||
|
||||
if must_special_scorer_counts.num_empty_scorers > 0 {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(SpecializedScorer::Other(Box::new(EmptyScorer)));
|
||||
}
|
||||
|
||||
let mut should_scorers = per_occur_scorers.remove(&Occur::Should).unwrap_or_default();
|
||||
@@ -225,7 +259,7 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
|
||||
if exclude_special_scorer_counts.num_all_scorers > 0 {
|
||||
// We exclude all documents at one point.
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(SpecializedScorer::Other(Box::new(EmptyScorer)));
|
||||
}
|
||||
|
||||
let effective_minimum_number_should_match = self
|
||||
@@ -237,7 +271,7 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
if effective_minimum_number_should_match > num_of_should_scorers {
|
||||
// We don't have enough scorers to satisfy the minimum number of should matches.
|
||||
// The request will match no documents.
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(SpecializedScorer::Other(Box::new(EmptyScorer)));
|
||||
}
|
||||
match effective_minimum_number_should_match {
|
||||
0 if num_of_should_scorers == 0 => ShouldScorersCombinationMethod::Ignored,
|
||||
@@ -245,13 +279,11 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
should_scorers,
|
||||
&score_combiner_fn,
|
||||
num_docs,
|
||||
reader.codec(),
|
||||
)),
|
||||
1 => ShouldScorersCombinationMethod::Required(scorer_union(
|
||||
should_scorers,
|
||||
&score_combiner_fn,
|
||||
num_docs,
|
||||
reader.codec(),
|
||||
)),
|
||||
n if num_of_should_scorers == n => {
|
||||
// When num_of_should_scorers equals the number of should clauses,
|
||||
@@ -259,41 +291,59 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
must_scorers.append(&mut should_scorers);
|
||||
ShouldScorersCombinationMethod::Ignored
|
||||
}
|
||||
_ => ShouldScorersCombinationMethod::Required(scorer_disjunction(
|
||||
should_scorers,
|
||||
score_combiner_fn(),
|
||||
effective_minimum_number_should_match,
|
||||
_ => ShouldScorersCombinationMethod::Required(SpecializedScorer::Other(
|
||||
scorer_disjunction(
|
||||
should_scorers,
|
||||
score_combiner_fn(),
|
||||
effective_minimum_number_should_match,
|
||||
),
|
||||
)),
|
||||
}
|
||||
};
|
||||
|
||||
let exclude_scorer_opt: Option<Box<dyn Scorer>> = if exclude_scorers.is_empty() {
|
||||
None
|
||||
} else {
|
||||
let exclude_scorers_union: Box<dyn Scorer> = scorer_union(
|
||||
exclude_scorers,
|
||||
DoNothingCombiner::default,
|
||||
num_docs,
|
||||
reader.codec(),
|
||||
);
|
||||
Some(exclude_scorers_union)
|
||||
};
|
||||
|
||||
|
||||
let include_scorer = match (should_scorers, must_scorers) {
|
||||
(ShouldScorersCombinationMethod::Ignored, must_scorers) => {
|
||||
// No SHOULD clauses (or they were absorbed into MUST).
|
||||
// Result depends entirely on MUST + any removed AllScorers.
|
||||
let combined_all_scorer_count = must_special_scorer_counts.num_all_scorers
|
||||
+ should_special_scorer_counts.num_all_scorers;
|
||||
let boxed_scorer: Box<dyn Scorer> = effective_must_scorer(
|
||||
must_scorers,
|
||||
combined_all_scorer_count,
|
||||
reader.max_doc(),
|
||||
num_docs,
|
||||
)
|
||||
.unwrap_or_else(|| box_scorer(EmptyScorer));
|
||||
boxed_scorer
|
||||
|
||||
// Try to detect a pure TermScorer intersection for block-max optimization.
|
||||
// Preconditions: no removed AllScorers, at least 2 scorers, all TermScorer
|
||||
// with frequency reading enabled.
|
||||
if combined_all_scorer_count == 0
|
||||
&& must_scorers.len() >= 2
|
||||
&& must_scorers.iter().all(|s| s.is::<TermScorer>())
|
||||
{
|
||||
let term_scorers: Vec<TermScorer> = must_scorers
|
||||
.into_iter()
|
||||
.map(|s| *(s.downcast::<TermScorer>().map_err(|_| ()).unwrap()))
|
||||
.collect();
|
||||
if term_scorers
|
||||
.iter()
|
||||
.all(|s| s.freq_reading_option() == FreqReadingOption::ReadFreq)
|
||||
{
|
||||
SpecializedScorer::TermIntersection(term_scorers)
|
||||
} else {
|
||||
let must_scorers: Vec<Box<dyn Scorer>> = term_scorers
|
||||
.into_iter()
|
||||
.map(|s| Box::new(s) as Box<dyn Scorer>)
|
||||
.collect();
|
||||
let boxed_scorer: Box<dyn Scorer> =
|
||||
effective_must_scorer(must_scorers, 0, reader.max_doc(), num_docs)
|
||||
.unwrap_or_else(|| Box::new(EmptyScorer));
|
||||
SpecializedScorer::Other(boxed_scorer)
|
||||
}
|
||||
} else {
|
||||
let boxed_scorer: Box<dyn Scorer> = effective_must_scorer(
|
||||
must_scorers,
|
||||
combined_all_scorer_count,
|
||||
reader.max_doc(),
|
||||
num_docs,
|
||||
)
|
||||
.unwrap_or_else(|| Box::new(EmptyScorer));
|
||||
SpecializedScorer::Other(boxed_scorer)
|
||||
}
|
||||
}
|
||||
(ShouldScorersCombinationMethod::Optional(should_scorer), must_scorers) => {
|
||||
// Optional SHOULD: contributes to scoring but not required for matching.
|
||||
@@ -318,12 +368,16 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
Some(must_scorer) => {
|
||||
// Has MUST constraint: SHOULD only affects scoring.
|
||||
if self.scoring_enabled {
|
||||
box_scorer(RequiredOptionalScorer::<_, _, TScoreCombiner>::new(
|
||||
SpecializedScorer::Other(Box::new(RequiredOptionalScorer::<
|
||||
_,
|
||||
_,
|
||||
TScoreCombiner,
|
||||
>::new(
|
||||
must_scorer,
|
||||
should_scorer,
|
||||
))
|
||||
into_box_scorer(should_scorer, &score_combiner_fn, num_docs),
|
||||
)))
|
||||
} else {
|
||||
must_scorer
|
||||
SpecializedScorer::Other(must_scorer)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -343,16 +397,33 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
|
||||
}
|
||||
Some(must_scorer) => {
|
||||
// Has MUST constraint: intersect MUST with SHOULD.
|
||||
intersect_scorers(vec![must_scorer, should_scorer], num_docs)
|
||||
let should_boxed =
|
||||
into_box_scorer(should_scorer, &score_combiner_fn, num_docs);
|
||||
SpecializedScorer::Other(intersect_scorers(
|
||||
vec![must_scorer, should_boxed],
|
||||
num_docs,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
if let Some(exclude_scorer) = exclude_scorer_opt {
|
||||
Ok(box_scorer(Exclude::new(include_scorer, exclude_scorer)))
|
||||
} else {
|
||||
Ok(include_scorer)
|
||||
if exclude_scorers.is_empty() {
|
||||
return Ok(include_scorer);
|
||||
}
|
||||
|
||||
let include_scorer_boxed = into_box_scorer(include_scorer, &score_combiner_fn, num_docs);
|
||||
let scorer: Box<dyn Scorer> = if exclude_scorers.len() == 1 {
|
||||
let exclude_scorer = exclude_scorers.pop().unwrap();
|
||||
match exclude_scorer.downcast::<TermScorer>() {
|
||||
// Cast to TermScorer succeeded
|
||||
Ok(exclude_scorer) => Box::new(Exclude::new(include_scorer_boxed, *exclude_scorer)),
|
||||
// We get back the original Box<dyn Scorer>
|
||||
Err(exclude_scorer) => Box::new(Exclude::new(include_scorer_boxed, exclude_scorer)),
|
||||
}
|
||||
} else {
|
||||
Box::new(Exclude::new(include_scorer_boxed, exclude_scorers))
|
||||
};
|
||||
Ok(SpecializedScorer::Other(scorer))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -382,6 +453,7 @@ fn remove_and_count_all_and_empty_scorers(
|
||||
|
||||
impl<TScoreCombiner: ScoreCombiner + Sync> Weight for BooleanWeight<TScoreCombiner> {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
let num_docs = reader.num_docs();
|
||||
if self.weights.is_empty() {
|
||||
Ok(Box::new(EmptyScorer))
|
||||
} else if self.weights.len() == 1 {
|
||||
@@ -393,8 +465,14 @@ impl<TScoreCombiner: ScoreCombiner + Sync> Weight for BooleanWeight<TScoreCombin
|
||||
}
|
||||
} else if self.scoring_enabled {
|
||||
self.complex_scorer(reader, boost, &self.score_combiner_fn)
|
||||
.map(|specialized_scorer| {
|
||||
into_box_scorer(specialized_scorer, &self.score_combiner_fn, num_docs)
|
||||
})
|
||||
} else {
|
||||
self.complex_scorer(reader, boost, DoNothingCombiner::default)
|
||||
.map(|specialized_scorer| {
|
||||
into_box_scorer(specialized_scorer, DoNothingCombiner::default, num_docs)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -423,8 +501,26 @@ impl<TScoreCombiner: ScoreCombiner + Sync> Weight for BooleanWeight<TScoreCombin
|
||||
reader: &SegmentReader,
|
||||
callback: &mut dyn FnMut(DocId, Score),
|
||||
) -> crate::Result<()> {
|
||||
let mut scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn)?;
|
||||
scorer.for_each(callback);
|
||||
let scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn)?;
|
||||
let num_docs = reader.num_docs();
|
||||
match scorer {
|
||||
SpecializedScorer::TermUnion(term_scorers) => {
|
||||
let mut union_scorer =
|
||||
BufferedUnionScorer::build(term_scorers, &self.score_combiner_fn, num_docs);
|
||||
for_each_scorer(&mut union_scorer, callback);
|
||||
}
|
||||
SpecializedScorer::TermIntersection(term_scorers) => {
|
||||
let boxed_scorers: Vec<Box<dyn Scorer>> = term_scorers
|
||||
.into_iter()
|
||||
.map(|term_scorer| Box::new(term_scorer) as Box<dyn Scorer>)
|
||||
.collect();
|
||||
let mut intersection = intersect_scorers(boxed_scorers, num_docs);
|
||||
for_each_scorer(intersection.as_mut(), callback);
|
||||
}
|
||||
SpecializedScorer::Other(mut scorer) => {
|
||||
for_each_scorer(scorer.as_mut(), callback);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -433,9 +529,28 @@ impl<TScoreCombiner: ScoreCombiner + Sync> Weight for BooleanWeight<TScoreCombin
|
||||
reader: &SegmentReader,
|
||||
callback: &mut dyn FnMut(&[DocId]),
|
||||
) -> crate::Result<()> {
|
||||
let mut scorer = self.complex_scorer(reader, 1.0, || DoNothingCombiner)?;
|
||||
let scorer = self.complex_scorer(reader, 1.0, || DoNothingCombiner)?;
|
||||
let num_docs = reader.num_docs();
|
||||
let mut buffer = [0u32; COLLECT_BLOCK_BUFFER_LEN];
|
||||
for_each_docset_buffered(scorer.as_mut(), &mut buffer, callback);
|
||||
|
||||
match scorer {
|
||||
SpecializedScorer::TermUnion(term_scorers) => {
|
||||
let mut union_scorer =
|
||||
BufferedUnionScorer::build(term_scorers, &self.score_combiner_fn, num_docs);
|
||||
for_each_docset_buffered(&mut union_scorer, &mut buffer, callback);
|
||||
}
|
||||
SpecializedScorer::TermIntersection(term_scorers) => {
|
||||
let boxed_scorers: Vec<Box<dyn Scorer>> = term_scorers
|
||||
.into_iter()
|
||||
.map(|term_scorer| Box::new(term_scorer) as Box<dyn Scorer>)
|
||||
.collect();
|
||||
let mut intersection = intersect_scorers(boxed_scorers, num_docs);
|
||||
for_each_docset_buffered(intersection.as_mut(), &mut buffer, callback);
|
||||
}
|
||||
SpecializedScorer::Other(mut scorer) => {
|
||||
for_each_docset_buffered(scorer.as_mut(), &mut buffer, callback);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -456,7 +571,17 @@ impl<TScoreCombiner: ScoreCombiner + Sync> Weight for BooleanWeight<TScoreCombin
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) -> crate::Result<()> {
|
||||
let scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn)?;
|
||||
reader.codec().for_each_pruning(threshold, scorer, callback);
|
||||
match scorer {
|
||||
SpecializedScorer::TermUnion(term_scorers) => {
|
||||
super::block_wand(term_scorers, threshold, callback);
|
||||
}
|
||||
SpecializedScorer::TermIntersection(term_scorers) => {
|
||||
super::block_wand_intersection(term_scorers, threshold, callback);
|
||||
}
|
||||
SpecializedScorer::Other(mut scorer) => {
|
||||
for_each_pruning_scorer(scorer.as_mut(), threshold, callback);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
mod block_wand_intersection;
|
||||
mod block_wand_union;
|
||||
mod boolean_query;
|
||||
mod boolean_weight;
|
||||
|
||||
|
||||
pub(crate) use self::block_wand_intersection::block_wand_intersection;
|
||||
pub(crate) use self::block_wand_union::{block_wand, block_wand_single_scorer};
|
||||
pub use self::boolean_query::BooleanQuery;
|
||||
pub use self::boolean_weight::BooleanWeight;
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use std::fmt;
|
||||
|
||||
use crate::docset::COLLECT_BLOCK_BUFFER_LEN;
|
||||
use crate::query::{box_scorer, EnableScoring, Explanation, Query, Scorer, Weight};
|
||||
use crate::docset::{SeekDangerResult, COLLECT_BLOCK_BUFFER_LEN};
|
||||
use crate::query::{EnableScoring, Explanation, Query, Scorer, Weight};
|
||||
use crate::{DocId, DocSet, Score, SegmentReader, TantivyError, Term};
|
||||
|
||||
/// `ConstScoreQuery` is a wrapper over a query to provide a constant score.
|
||||
@@ -65,10 +65,7 @@ impl ConstWeight {
|
||||
impl Weight for ConstWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
let inner_scorer = self.weight.scorer(reader, boost)?;
|
||||
Ok(box_scorer(ConstScorer::new(
|
||||
inner_scorer,
|
||||
boost * self.score,
|
||||
)))
|
||||
Ok(Box::new(ConstScorer::new(inner_scorer, boost * self.score)))
|
||||
}
|
||||
|
||||
fn explain(&self, reader: &SegmentReader, doc: u32) -> crate::Result<Explanation> {
|
||||
@@ -122,6 +119,10 @@ impl<TDocSet: DocSet> DocSet for ConstScorer<TDocSet> {
|
||||
self.docset.seek(target)
|
||||
}
|
||||
|
||||
fn seek_danger(&mut self, target: DocId) -> SeekDangerResult {
|
||||
self.docset.seek_danger(target)
|
||||
}
|
||||
|
||||
fn fill_buffer(&mut self, buffer: &mut [DocId; COLLECT_BLOCK_BUFFER_LEN]) -> usize {
|
||||
self.docset.fill_buffer(buffer)
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ use super::Scorer;
|
||||
use crate::docset::TERMINATED;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::{box_scorer, EnableScoring, Explanation, Query, Weight};
|
||||
use crate::query::{EnableScoring, Explanation, Query, Weight};
|
||||
use crate::{DocId, DocSet, Score, Searcher};
|
||||
|
||||
/// `EmptyQuery` is a dummy `Query` in which no document matches.
|
||||
@@ -27,7 +27,7 @@ impl Query for EmptyQuery {
|
||||
pub struct EmptyWeight;
|
||||
impl Weight for EmptyWeight {
|
||||
fn scorer(&self, _reader: &SegmentReader, _boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
Ok(box_scorer(EmptyScorer))
|
||||
Ok(Box::new(EmptyScorer))
|
||||
}
|
||||
|
||||
fn explain(&self, _reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
|
||||
|
||||
@@ -3,7 +3,7 @@ use core::fmt::Debug;
|
||||
use columnar::{ColumnIndex, DynamicColumn};
|
||||
use common::BitSet;
|
||||
|
||||
use super::{box_scorer, ConstScorer, EmptyScorer};
|
||||
use super::{ConstScorer, EmptyScorer};
|
||||
use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::index::SegmentReader;
|
||||
use crate::query::all_query::AllScorer;
|
||||
@@ -117,7 +117,7 @@ impl Weight for ExistsWeight {
|
||||
}
|
||||
}
|
||||
if non_empty_columns.is_empty() {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
}
|
||||
|
||||
// If any column is full, all docs match.
|
||||
@@ -128,9 +128,9 @@ impl Weight for ExistsWeight {
|
||||
{
|
||||
let all_scorer = AllScorer::new(max_doc);
|
||||
if boost != 1.0f32 {
|
||||
return Ok(box_scorer(BoostScorer::new(all_scorer, boost)));
|
||||
return Ok(Box::new(BoostScorer::new(all_scorer, boost)));
|
||||
} else {
|
||||
return Ok(box_scorer(all_scorer));
|
||||
return Ok(Box::new(all_scorer));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -138,7 +138,7 @@ impl Weight for ExistsWeight {
|
||||
// NOTE: A lower number may be better for very sparse columns
|
||||
if non_empty_columns.len() < 4 {
|
||||
let docset = ExistsDocSet::new(non_empty_columns, reader.max_doc());
|
||||
return Ok(box_scorer(ConstScorer::new(docset, boost)));
|
||||
return Ok(Box::new(ConstScorer::new(docset, boost)));
|
||||
}
|
||||
|
||||
// If we have many dynamic columns, precompute a bitset of matching docs
|
||||
@@ -162,7 +162,7 @@ impl Weight for ExistsWeight {
|
||||
}
|
||||
}
|
||||
let docset = BitSetDocSet::from(doc_bitset);
|
||||
Ok(box_scorer(ConstScorer::new(docset, boost)))
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
}
|
||||
|
||||
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
|
||||
|
||||
@@ -3,7 +3,7 @@ use common::TinySet;
|
||||
use super::size_hint::estimate_intersection;
|
||||
use crate::docset::{DocSet, SeekDangerResult, BLOCK_NUM_TINYBITSETS, TERMINATED};
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::{box_scorer, EmptyScorer, Scorer};
|
||||
use crate::query::{EmptyScorer, Scorer};
|
||||
use crate::{DocId, Score};
|
||||
|
||||
/// Returns the intersection scorer.
|
||||
@@ -22,7 +22,7 @@ pub fn intersect_scorers(
|
||||
segment_num_docs: u32,
|
||||
) -> Box<dyn Scorer> {
|
||||
if scorers.is_empty() {
|
||||
return box_scorer(EmptyScorer);
|
||||
return Box::new(EmptyScorer);
|
||||
}
|
||||
if scorers.len() == 1 {
|
||||
return scorers.pop().unwrap();
|
||||
@@ -31,7 +31,7 @@ pub fn intersect_scorers(
|
||||
scorers.sort_by_key(|scorer| scorer.cost());
|
||||
let doc = go_to_first_doc(&mut scorers[..]);
|
||||
if doc == TERMINATED {
|
||||
return box_scorer(EmptyScorer);
|
||||
return Box::new(EmptyScorer);
|
||||
}
|
||||
// We know that we have at least 2 elements.
|
||||
let left = scorers.remove(0);
|
||||
@@ -40,14 +40,14 @@ pub fn intersect_scorers(
|
||||
.iter()
|
||||
.all(|&scorer| scorer.is::<TermScorer>());
|
||||
if all_term_scorers {
|
||||
return box_scorer(Intersection {
|
||||
return Box::new(Intersection {
|
||||
left: *(left.downcast::<TermScorer>().map_err(|_| ()).unwrap()),
|
||||
right: *(right.downcast::<TermScorer>().map_err(|_| ()).unwrap()),
|
||||
others: scorers,
|
||||
segment_num_docs,
|
||||
});
|
||||
}
|
||||
box_scorer(Intersection {
|
||||
Box::new(Intersection {
|
||||
left,
|
||||
right,
|
||||
others: scorers,
|
||||
|
||||
@@ -24,7 +24,7 @@ mod reqopt_scorer;
|
||||
mod scorer;
|
||||
mod set_query;
|
||||
mod size_hint;
|
||||
pub(crate) mod term_query;
|
||||
mod term_query;
|
||||
mod union;
|
||||
mod weight;
|
||||
|
||||
@@ -54,14 +54,13 @@ pub use self::more_like_this::{MoreLikeThisQuery, MoreLikeThisQueryBuilder};
|
||||
pub use self::phrase_prefix_query::PhrasePrefixQuery;
|
||||
pub use self::phrase_query::regex_phrase_query::{wildcard_query_to_regex_str, RegexPhraseQuery};
|
||||
pub use self::phrase_query::PhraseQuery;
|
||||
pub(crate) use self::phrase_query::PhraseScorer;
|
||||
pub use self::query::{EnableScoring, Query, QueryClone};
|
||||
pub use self::query_parser::{QueryParser, QueryParserError};
|
||||
pub use self::range_query::*;
|
||||
pub use self::regex_query::RegexQuery;
|
||||
pub use self::reqopt_scorer::RequiredOptionalScorer;
|
||||
pub use self::score_combiner::{DisjunctionMaxCombiner, ScoreCombiner, SumCombiner};
|
||||
pub use self::scorer::{box_scorer, Scorer};
|
||||
pub use self::scorer::Scorer;
|
||||
pub use self::set_query::TermSetQuery;
|
||||
pub use self::term_query::TermQuery;
|
||||
pub use self::union::BufferedUnionScorer;
|
||||
|
||||
@@ -2,7 +2,7 @@ use crate::docset::{DocSet, SeekDangerResult, TERMINATED};
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::Postings;
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::phrase_query::{intersection_exists, PhraseScorer};
|
||||
use crate::query::phrase_query::{intersection_count, PhraseScorer};
|
||||
use crate::query::Scorer;
|
||||
use crate::{DocId, Score};
|
||||
|
||||
@@ -100,6 +100,7 @@ pub struct PhrasePrefixScorer<TPostings: Postings> {
|
||||
phrase_scorer: PhraseKind<TPostings>,
|
||||
suffixes: Vec<TPostings>,
|
||||
suffix_offset: u32,
|
||||
phrase_count: u32,
|
||||
suffix_position_buffer: Vec<u32>,
|
||||
}
|
||||
|
||||
@@ -143,6 +144,7 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
|
||||
phrase_scorer,
|
||||
suffixes,
|
||||
suffix_offset: (max_offset - suffix_pos) as u32,
|
||||
phrase_count: 0,
|
||||
suffix_position_buffer: Vec::with_capacity(100),
|
||||
};
|
||||
if phrase_prefix_scorer.doc() != TERMINATED && !phrase_prefix_scorer.matches_prefix() {
|
||||
@@ -151,7 +153,12 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
|
||||
phrase_prefix_scorer
|
||||
}
|
||||
|
||||
pub fn phrase_count(&self) -> u32 {
|
||||
self.phrase_count
|
||||
}
|
||||
|
||||
fn matches_prefix(&mut self) -> bool {
|
||||
let mut count = 0;
|
||||
let current_doc = self.doc();
|
||||
let pos_matching = self.phrase_scorer.get_intersection();
|
||||
for suffix in &mut self.suffixes {
|
||||
@@ -161,12 +168,11 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
|
||||
let doc = suffix.seek(current_doc);
|
||||
if doc == current_doc {
|
||||
suffix.positions_with_offset(self.suffix_offset, &mut self.suffix_position_buffer);
|
||||
if intersection_exists(pos_matching, &self.suffix_position_buffer) {
|
||||
return true;
|
||||
}
|
||||
count += intersection_count(pos_matching, &self.suffix_position_buffer);
|
||||
}
|
||||
}
|
||||
false
|
||||
self.phrase_count = count as u32;
|
||||
count != 0
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
use super::{prefix_end, PhrasePrefixScorer};
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::postings::Postings;
|
||||
use crate::postings::SegmentPostings;
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::{box_scorer, EmptyScorer, Scorer, Weight};
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::{EmptyScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::{IndexRecordOption, Term};
|
||||
use crate::Score;
|
||||
use crate::{DocId, DocSet, Score};
|
||||
|
||||
pub struct PhrasePrefixWeight {
|
||||
phrase_terms: Vec<(usize, Term)>,
|
||||
@@ -45,13 +46,13 @@ impl PhrasePrefixWeight {
|
||||
&self,
|
||||
reader: &SegmentReader,
|
||||
boost: Score,
|
||||
) -> crate::Result<Option<Box<dyn Scorer>>> {
|
||||
) -> crate::Result<Option<PhrasePrefixScorer<SegmentPostings>>> {
|
||||
let similarity_weight_opt = self
|
||||
.similarity_weight_opt
|
||||
.as_ref()
|
||||
.map(|similarity_weight| similarity_weight.boost_by(boost));
|
||||
let fieldnorm_reader = self.fieldnorm_reader(reader)?;
|
||||
let mut term_postings_list: Vec<(usize, Box<dyn Postings>)> = Vec::new();
|
||||
let mut term_postings_list = Vec::new();
|
||||
for &(offset, ref term) in &self.phrase_terms {
|
||||
if let Some(postings) = reader
|
||||
.inverted_index(term.field())?
|
||||
@@ -102,32 +103,49 @@ impl PhrasePrefixWeight {
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Some(box_scorer(PhrasePrefixScorer::new(
|
||||
Ok(Some(PhrasePrefixScorer::new(
|
||||
term_postings_list,
|
||||
similarity_weight_opt,
|
||||
fieldnorm_reader,
|
||||
suffixes,
|
||||
self.prefix.0,
|
||||
))))
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
impl Weight for PhrasePrefixWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
if let Some(scorer) = self.phrase_scorer(reader, boost)? {
|
||||
Ok(scorer)
|
||||
Ok(Box::new(scorer))
|
||||
} else {
|
||||
Ok(box_scorer(EmptyScorer))
|
||||
Ok(Box::new(EmptyScorer))
|
||||
}
|
||||
}
|
||||
|
||||
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
|
||||
let scorer_opt = self.phrase_scorer(reader, 1.0)?;
|
||||
if scorer_opt.is_none() {
|
||||
return Err(does_not_match(doc));
|
||||
}
|
||||
let mut scorer = scorer_opt.unwrap();
|
||||
if scorer.seek(doc) != doc {
|
||||
return Err(does_not_match(doc));
|
||||
}
|
||||
let fieldnorm_reader = self.fieldnorm_reader(reader)?;
|
||||
let fieldnorm_id = fieldnorm_reader.fieldnorm_id(doc);
|
||||
let phrase_count = scorer.phrase_count();
|
||||
let mut explanation = Explanation::new("Phrase Prefix Scorer", scorer.score());
|
||||
if let Some(similarity_weight) = self.similarity_weight_opt.as_ref() {
|
||||
explanation.add_detail(similarity_weight.explain(fieldnorm_id, phrase_count));
|
||||
}
|
||||
Ok(explanation)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::docset::TERMINATED;
|
||||
use crate::index::Index;
|
||||
use crate::postings::Postings;
|
||||
use crate::query::phrase_prefix_query::PhrasePrefixScorer;
|
||||
use crate::query::{EnableScoring, PhrasePrefixQuery, Query};
|
||||
use crate::schema::{Schema, TEXT};
|
||||
use crate::{DocSet, IndexWriter, Term};
|
||||
@@ -168,14 +186,14 @@ mod tests {
|
||||
.phrase_prefix_query_weight(enable_scoring)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut phrase_scorer_boxed = phrase_weight
|
||||
let mut phrase_scorer = phrase_weight
|
||||
.phrase_scorer(searcher.segment_reader(0u32), 1.0)?
|
||||
.unwrap();
|
||||
let phrase_scorer: &mut PhrasePrefixScorer<Box<dyn Postings>> =
|
||||
phrase_scorer_boxed.as_any_mut().downcast_mut().unwrap();
|
||||
assert_eq!(phrase_scorer.doc(), 1);
|
||||
assert_eq!(phrase_scorer.phrase_count(), 2);
|
||||
assert_eq!(phrase_scorer.advance(), 2);
|
||||
assert_eq!(phrase_scorer.doc(), 2);
|
||||
assert_eq!(phrase_scorer.phrase_count(), 1);
|
||||
assert_eq!(phrase_scorer.advance(), TERMINATED);
|
||||
Ok(())
|
||||
}
|
||||
@@ -195,15 +213,14 @@ mod tests {
|
||||
.phrase_prefix_query_weight(enable_scoring)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let mut phrase_scorer_boxed = phrase_weight
|
||||
let mut phrase_scorer = phrase_weight
|
||||
.phrase_scorer(searcher.segment_reader(0u32), 1.0)?
|
||||
.unwrap();
|
||||
let phrase_scorer = phrase_scorer_boxed
|
||||
.downcast_mut::<PhrasePrefixScorer<Box<dyn Postings>>>()
|
||||
.unwrap();
|
||||
assert_eq!(phrase_scorer.doc(), 1);
|
||||
assert_eq!(phrase_scorer.phrase_count(), 2);
|
||||
assert_eq!(phrase_scorer.advance(), 2);
|
||||
assert_eq!(phrase_scorer.doc(), 2);
|
||||
assert_eq!(phrase_scorer.phrase_count(), 1);
|
||||
assert_eq!(phrase_scorer.advance(), TERMINATED);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ pub mod regex_phrase_query;
|
||||
mod regex_phrase_weight;
|
||||
|
||||
pub use self::phrase_query::PhraseQuery;
|
||||
pub(crate) use self::phrase_scorer::intersection_exists;
|
||||
pub(crate) use self::phrase_scorer::intersection_count;
|
||||
pub use self::phrase_scorer::PhraseScorer;
|
||||
pub use self::phrase_weight::PhraseWeight;
|
||||
|
||||
|
||||
@@ -126,7 +126,7 @@ impl PhraseQuery {
|
||||
};
|
||||
let mut weight = PhraseWeight::new(self.phrase_terms.clone(), bm25_weight_opt);
|
||||
if self.slop > 0 {
|
||||
weight.set_slop(self.slop);
|
||||
weight.slop(self.slop);
|
||||
}
|
||||
Ok(weight)
|
||||
}
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
use std::cmp::Ordering;
|
||||
|
||||
use crate::codec::standard::postings::StandardPostings;
|
||||
use crate::docset::{DocSet, SeekDangerResult, TERMINATED};
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::Postings;
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::{Explanation, Intersection, Scorer};
|
||||
use crate::query::{Intersection, Scorer};
|
||||
use crate::{DocId, Score};
|
||||
|
||||
struct PostingsWithOffset<TPostings> {
|
||||
@@ -44,7 +43,7 @@ impl<TPostings: Postings> DocSet for PostingsWithOffset<TPostings> {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct PhraseScorer<TPostings: Postings = StandardPostings> {
|
||||
pub struct PhraseScorer<TPostings: Postings> {
|
||||
intersection_docset: Intersection<PostingsWithOffset<TPostings>, PostingsWithOffset<TPostings>>,
|
||||
num_terms: usize,
|
||||
left_positions: Vec<u32>,
|
||||
@@ -59,7 +58,7 @@ pub struct PhraseScorer<TPostings: Postings = StandardPostings> {
|
||||
}
|
||||
|
||||
/// Returns true if and only if the two sorted arrays contain a common element
|
||||
pub(crate) fn intersection_exists(left: &[u32], right: &[u32]) -> bool {
|
||||
fn intersection_exists(left: &[u32], right: &[u32]) -> bool {
|
||||
let mut left_index = 0;
|
||||
let mut right_index = 0;
|
||||
while left_index < left.len() && right_index < right.len() {
|
||||
@@ -80,7 +79,7 @@ pub(crate) fn intersection_exists(left: &[u32], right: &[u32]) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn intersection_count(left: &[u32], right: &[u32]) -> usize {
|
||||
pub(crate) fn intersection_count(left: &[u32], right: &[u32]) -> usize {
|
||||
let mut left_index = 0;
|
||||
let mut right_index = 0;
|
||||
let mut count = 0;
|
||||
@@ -403,7 +402,6 @@ impl<TPostings: Postings> PhraseScorer<TPostings> {
|
||||
scorer
|
||||
}
|
||||
|
||||
/// Returns the number of phrases identified in the current matching doc.
|
||||
pub fn phrase_count(&self) -> u32 {
|
||||
self.phrase_count
|
||||
}
|
||||
@@ -586,17 +584,6 @@ impl<TPostings: Postings> Scorer for PhraseScorer<TPostings> {
|
||||
1.0f32
|
||||
}
|
||||
}
|
||||
|
||||
fn explain(&mut self) -> Explanation {
|
||||
let doc = self.doc();
|
||||
let phrase_count = self.phrase_count();
|
||||
let fieldnorm_id = self.fieldnorm_reader.fieldnorm_id(doc);
|
||||
let mut explanation = Explanation::new("Phrase Scorer", self.score());
|
||||
if let Some(similarity_weight) = self.similarity_weight_opt.as_ref() {
|
||||
explanation.add_detail(similarity_weight.explain(fieldnorm_id, phrase_count));
|
||||
}
|
||||
explanation
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
use super::PhraseScorer;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::postings::TermInfo;
|
||||
use crate::postings::SegmentPostings;
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::{box_scorer, EmptyScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::Term;
|
||||
use crate::query::{EmptyScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::{IndexRecordOption, Term};
|
||||
use crate::{DocId, DocSet, Score};
|
||||
|
||||
pub struct PhraseWeight {
|
||||
@@ -20,10 +21,11 @@ impl PhraseWeight {
|
||||
phrase_terms: Vec<(usize, Term)>,
|
||||
similarity_weight_opt: Option<Bm25Weight>,
|
||||
) -> PhraseWeight {
|
||||
let slop = 0;
|
||||
PhraseWeight {
|
||||
phrase_terms,
|
||||
similarity_weight_opt,
|
||||
slop: 0,
|
||||
slop,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -41,52 +43,32 @@ impl PhraseWeight {
|
||||
&self,
|
||||
reader: &SegmentReader,
|
||||
boost: Score,
|
||||
) -> crate::Result<Option<Box<dyn Scorer>>> {
|
||||
) -> crate::Result<Option<PhraseScorer<SegmentPostings>>> {
|
||||
let similarity_weight_opt = self
|
||||
.similarity_weight_opt
|
||||
.as_ref()
|
||||
.map(|similarity_weight| similarity_weight.boost_by(boost));
|
||||
let fieldnorm_reader = self.fieldnorm_reader(reader)?;
|
||||
|
||||
if self.phrase_terms.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let field = self.phrase_terms[0].1.field();
|
||||
|
||||
if !self
|
||||
.phrase_terms
|
||||
.iter()
|
||||
.all(|(_offset, term)| term.field() == field)
|
||||
{
|
||||
return Err(crate::TantivyError::InvalidArgument(
|
||||
"All terms in a phrase query must belong to the same field".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let inverted_index_reader = reader.inverted_index(field)?;
|
||||
|
||||
let mut term_infos: Vec<(usize, TermInfo)> = Vec::with_capacity(self.phrase_terms.len());
|
||||
|
||||
let mut term_postings_list = Vec::new();
|
||||
for &(offset, ref term) in &self.phrase_terms {
|
||||
let Some(term_info) = inverted_index_reader.get_term_info(term)? else {
|
||||
if let Some(postings) = reader
|
||||
.inverted_index(term.field())?
|
||||
.read_postings(term, IndexRecordOption::WithFreqsAndPositions)?
|
||||
{
|
||||
term_postings_list.push((offset, postings));
|
||||
} else {
|
||||
return Ok(None);
|
||||
};
|
||||
term_infos.push((offset, term_info));
|
||||
}
|
||||
}
|
||||
|
||||
let scorer = reader.codec().new_phrase_scorer_type_erased(
|
||||
&term_infos[..],
|
||||
Ok(Some(PhraseScorer::new(
|
||||
term_postings_list,
|
||||
similarity_weight_opt,
|
||||
fieldnorm_reader,
|
||||
self.slop,
|
||||
&inverted_index_reader,
|
||||
)?;
|
||||
|
||||
Ok(Some(scorer))
|
||||
)))
|
||||
}
|
||||
|
||||
/// Sets the slop for the given PhraseWeight.
|
||||
pub fn set_slop(&mut self, slop: u32) {
|
||||
pub fn slop(&mut self, slop: u32) {
|
||||
self.slop = slop;
|
||||
}
|
||||
}
|
||||
@@ -94,9 +76,9 @@ impl PhraseWeight {
|
||||
impl Weight for PhraseWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
if let Some(scorer) = self.phrase_scorer(reader, boost)? {
|
||||
Ok(scorer)
|
||||
Ok(Box::new(scorer))
|
||||
} else {
|
||||
Ok(box_scorer(EmptyScorer))
|
||||
Ok(Box::new(EmptyScorer))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -109,7 +91,14 @@ impl Weight for PhraseWeight {
|
||||
if scorer.seek(doc) != doc {
|
||||
return Err(does_not_match(doc));
|
||||
}
|
||||
Ok(scorer.explain())
|
||||
let fieldnorm_reader = self.fieldnorm_reader(reader)?;
|
||||
let fieldnorm_id = fieldnorm_reader.fieldnorm_id(doc);
|
||||
let phrase_count = scorer.phrase_count();
|
||||
let mut explanation = Explanation::new("Phrase Scorer", scorer.score());
|
||||
if let Some(similarity_weight) = self.similarity_weight_opt.as_ref() {
|
||||
explanation.add_detail(similarity_weight.explain(fieldnorm_id, phrase_count));
|
||||
}
|
||||
Ok(explanation)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -117,8 +106,7 @@ impl Weight for PhraseWeight {
|
||||
mod tests {
|
||||
use super::super::tests::create_index;
|
||||
use crate::docset::TERMINATED;
|
||||
use crate::query::phrase_query::PhraseScorer;
|
||||
use crate::query::{EnableScoring, PhraseQuery, Scorer};
|
||||
use crate::query::{EnableScoring, PhraseQuery};
|
||||
use crate::{DocSet, Term};
|
||||
|
||||
#[test]
|
||||
@@ -133,11 +121,9 @@ mod tests {
|
||||
]);
|
||||
let enable_scoring = EnableScoring::enabled_from_searcher(&searcher);
|
||||
let phrase_weight = phrase_query.phrase_weight(enable_scoring).unwrap();
|
||||
let phrase_scorer_boxed: Box<dyn Scorer> = phrase_weight
|
||||
let mut phrase_scorer = phrase_weight
|
||||
.phrase_scorer(searcher.segment_reader(0u32), 1.0)?
|
||||
.unwrap();
|
||||
let mut phrase_scorer: Box<PhraseScorer> =
|
||||
phrase_scorer_boxed.downcast::<PhraseScorer>().ok().unwrap();
|
||||
assert_eq!(phrase_scorer.doc(), 1);
|
||||
assert_eq!(phrase_scorer.phrase_count(), 2);
|
||||
assert_eq!(phrase_scorer.advance(), 2);
|
||||
|
||||
@@ -6,13 +6,11 @@ use tantivy_fst::Regex;
|
||||
use super::PhraseScorer;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::postings::{LoadedPostings, Postings, TermInfo};
|
||||
use crate::postings::{LoadedPostings, Postings, SegmentPostings, TermInfo};
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::union::{BitSetPostingUnion, SimpleUnion};
|
||||
use crate::query::{
|
||||
box_scorer, AutomatonWeight, BitSetDocSet, EmptyScorer, Explanation, Scorer, Weight,
|
||||
};
|
||||
use crate::query::{AutomatonWeight, BitSetDocSet, EmptyScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::{Field, IndexRecordOption};
|
||||
use crate::{DocId, DocSet, InvertedIndexReader, Score};
|
||||
|
||||
@@ -105,9 +103,18 @@ impl RegexPhraseWeight {
|
||||
term_info: &TermInfo,
|
||||
doc_bitset: &mut BitSet,
|
||||
) -> crate::Result<()> {
|
||||
let mut segment_postings =
|
||||
inverted_index.read_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
|
||||
segment_postings.fill_bitset(doc_bitset);
|
||||
let mut block_segment_postings = inverted_index
|
||||
.read_block_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
|
||||
loop {
|
||||
let docs = block_segment_postings.docs();
|
||||
if docs.is_empty() {
|
||||
break;
|
||||
}
|
||||
for &doc in docs {
|
||||
doc_bitset.insert(doc);
|
||||
}
|
||||
block_segment_postings.advance();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -181,7 +188,7 @@ impl RegexPhraseWeight {
|
||||
// - Bucket 1: Terms appearing in 0.1% to 1% of documents
|
||||
// - Bucket 2: Terms appearing in 1% to 10% of documents
|
||||
// - Bucket 3: Terms appearing in more than 10% of documents
|
||||
let mut buckets: Vec<(BitSet, Vec<Box<dyn Postings>>)> = (0..4)
|
||||
let mut buckets: Vec<(BitSet, Vec<SegmentPostings>)> = (0..4)
|
||||
.map(|_| (BitSet::with_max_value(max_doc), Vec::new()))
|
||||
.collect();
|
||||
|
||||
@@ -190,7 +197,7 @@ impl RegexPhraseWeight {
|
||||
for term_info in term_infos {
|
||||
let mut term_posting = inverted_index
|
||||
.read_postings_from_terminfo(term_info, IndexRecordOption::WithFreqsAndPositions)?;
|
||||
let num_docs = u32::from(term_posting.doc_freq());
|
||||
let num_docs = term_posting.doc_freq();
|
||||
|
||||
if num_docs < SPARSE_TERM_DOC_THRESHOLD {
|
||||
let current_bucket = &mut sparse_buckets[0];
|
||||
@@ -264,9 +271,9 @@ impl RegexPhraseWeight {
|
||||
impl Weight for RegexPhraseWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
if let Some(scorer) = self.phrase_scorer(reader, boost)? {
|
||||
Ok(box_scorer(scorer))
|
||||
Ok(Box::new(scorer))
|
||||
} else {
|
||||
Ok(box_scorer(EmptyScorer))
|
||||
Ok(Box::new(EmptyScorer))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::ops::RangeInclusive;
|
||||
|
||||
use columnar::Column;
|
||||
|
||||
use crate::docset::SeekDangerResult;
|
||||
use crate::{DocId, DocSet, TERMINATED};
|
||||
|
||||
/// Helper to have a cursor over a vec of docids
|
||||
@@ -184,6 +185,37 @@ impl<T: Send + Sync + PartialOrd + Copy + Debug + 'static> DocSet for RangeDocSe
|
||||
doc
|
||||
}
|
||||
|
||||
/// `seek_danger` only needs to answer whether `target` itself matches, so it does a cheap
|
||||
/// point lookup on the column instead of scanning forward to materialize the next match (the
|
||||
/// expensive part of a regular `seek`).
|
||||
fn seek_danger(&mut self, target: DocId) -> SeekDangerResult {
|
||||
// Covers `target == TERMINATED` and any target past the last doc: no match is possible.
|
||||
if target >= self.column.num_docs() {
|
||||
return SeekDangerResult::SeekLowerBound(TERMINATED);
|
||||
}
|
||||
|
||||
if self.is_last_seek_distance_large(target) {
|
||||
self.reset_fetch_range();
|
||||
}
|
||||
self.last_seek_pos_opt = Some(target);
|
||||
|
||||
let is_match = self
|
||||
.column
|
||||
.values_for_doc(target)
|
||||
.any(|value| self.value_range.contains(&value));
|
||||
if is_match {
|
||||
// Leave the docset in a valid state positioned on `target`, so `doc()` returns it and a
|
||||
// following `advance()` resumes the scan right after it.
|
||||
self.loaded_docs.get_cleared_data().push(target);
|
||||
self.next_fetch_start = target + 1;
|
||||
SeekDangerResult::Found
|
||||
} else {
|
||||
// `target` is not in the docset. The next match is strictly greater than `target`, so
|
||||
// `target + 1` is a valid lower bound. We may leave the docset in an invalid state.
|
||||
SeekDangerResult::SeekLowerBound(target + 1)
|
||||
}
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> u32 {
|
||||
// TODO: Implement a better size hint
|
||||
self.column.num_docs() / 10
|
||||
@@ -209,12 +241,148 @@ impl<T: Send + Sync + PartialOrd + Copy + Debug + 'static> DocSet for RangeDocSe
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ops::Bound;
|
||||
use std::ops::{Bound, RangeInclusive};
|
||||
|
||||
use columnar::Column;
|
||||
|
||||
use super::RangeDocSet;
|
||||
use crate::collector::Count;
|
||||
use crate::directory::RamDirectory;
|
||||
use crate::docset::{SeekDangerResult, TERMINATED};
|
||||
use crate::query::RangeQuery;
|
||||
use crate::{schema, IndexBuilder, TantivyDocument, Term};
|
||||
use crate::{schema, DocSet, Index, IndexBuilder, TantivyDocument, Term};
|
||||
|
||||
/// Builds a single-segment index where doc `i` carries `values_for_doc(i)` in a u64 fast
|
||||
/// field, then returns its column so we can drive a `RangeDocSet` directly.
|
||||
fn build_u64_column(
|
||||
num_docs: usize,
|
||||
values_for_doc: impl Fn(usize) -> Vec<u64>,
|
||||
) -> Column<u64> {
|
||||
let mut schema_builder = schema::SchemaBuilder::new();
|
||||
let value_field = schema_builder.add_u64_field("value", schema::FAST);
|
||||
let index = Index::create_in_ram(schema_builder.build());
|
||||
{
|
||||
let mut writer = index.writer_for_tests().unwrap();
|
||||
for i in 0..num_docs {
|
||||
let mut doc = TantivyDocument::new();
|
||||
for v in values_for_doc(i) {
|
||||
doc.add_u64(value_field, v);
|
||||
}
|
||||
writer.add_document(doc).unwrap();
|
||||
}
|
||||
writer.commit().unwrap();
|
||||
}
|
||||
let searcher = index.reader().unwrap().searcher();
|
||||
assert_eq!(searcher.segment_readers().len(), 1);
|
||||
searcher
|
||||
.segment_reader(0)
|
||||
.fast_fields()
|
||||
.u64("value")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn range_docset(
|
||||
value_range: RangeInclusive<u64>,
|
||||
num_docs: usize,
|
||||
values_for_doc: impl Fn(usize) -> Vec<u64>,
|
||||
) -> RangeDocSet<u64> {
|
||||
RangeDocSet::new(value_range, build_u64_column(num_docs, values_for_doc))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_danger_found_leaves_valid_state() {
|
||||
// Even docs match the range, odd docs do not.
|
||||
let mut docset = range_docset(0..=0, 100, |i| vec![(i % 2) as u64]);
|
||||
|
||||
// Matching target: `Found`, and the docset is positioned exactly on it.
|
||||
assert_eq!(docset.seek_danger(10), SeekDangerResult::Found);
|
||||
assert_eq!(docset.doc(), 10);
|
||||
// A following advance resumes the scan right after the found doc.
|
||||
assert_eq!(docset.advance(), 12);
|
||||
assert_eq!(docset.doc(), 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_danger_miss_returns_lower_bound() {
|
||||
let mut docset = range_docset(0..=0, 100, |i| vec![(i % 2) as u64]);
|
||||
|
||||
// Odd target does not match: lower bound is strictly greater than the target and never
|
||||
// skips past the next real match (here doc 12, the first even doc after 11).
|
||||
match docset.seek_danger(11) {
|
||||
SeekDangerResult::SeekLowerBound(lower_bound) => {
|
||||
assert!(lower_bound > 11);
|
||||
assert!(lower_bound <= 12);
|
||||
}
|
||||
SeekDangerResult::Found => panic!("11 should not match"),
|
||||
}
|
||||
// After a miss we may be in an invalid state; another seek_danger recovers it.
|
||||
assert_eq!(docset.seek_danger(12), SeekDangerResult::Found);
|
||||
assert_eq!(docset.doc(), 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_danger_terminated_and_out_of_bounds() {
|
||||
let mut docset = range_docset(0..=0, 10, |i| vec![(i % 2) as u64]);
|
||||
assert_eq!(
|
||||
docset.seek_danger(TERMINATED),
|
||||
SeekDangerResult::SeekLowerBound(TERMINATED)
|
||||
);
|
||||
// A target past the last doc has no possible match either.
|
||||
assert_eq!(
|
||||
docset.seek_danger(10),
|
||||
SeekDangerResult::SeekLowerBound(TERMINATED)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_danger_multivalued() {
|
||||
// Doc `i` holds values [i, i+1]; the range {5} matches docs 4 and 5.
|
||||
let mut docset = range_docset(5..=5, 20, |i| vec![i as u64, i as u64 + 1]);
|
||||
|
||||
assert_eq!(docset.seek_danger(4), SeekDangerResult::Found);
|
||||
assert_eq!(docset.doc(), 4);
|
||||
assert_eq!(docset.advance(), 5);
|
||||
// No further match after doc 5.
|
||||
assert_eq!(docset.advance(), TERMINATED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_danger_matches_seek() {
|
||||
// Cross-check seek_danger against the true next match for every target, on a column with a
|
||||
// few sparse matches.
|
||||
let matches = [3u32, 7, 50, 51, 99];
|
||||
let num_docs = 100;
|
||||
let values_for_doc = |i: usize| {
|
||||
vec![if matches.contains(&(i as u32)) {
|
||||
1u64
|
||||
} else {
|
||||
0u64
|
||||
}]
|
||||
};
|
||||
|
||||
for target in 0..num_docs as u32 {
|
||||
// The first matching doc greater than or equal to `target`, i.e. what `seek` returns.
|
||||
let expected = matches
|
||||
.iter()
|
||||
.copied()
|
||||
.find(|&m| m >= target)
|
||||
.unwrap_or(TERMINATED);
|
||||
|
||||
let mut danger = range_docset(1..=1, num_docs, values_for_doc);
|
||||
match danger.seek_danger(target) {
|
||||
SeekDangerResult::Found => {
|
||||
assert_eq!(expected, target, "target {target} reported Found");
|
||||
assert_eq!(danger.doc(), target);
|
||||
}
|
||||
SeekDangerResult::SeekLowerBound(lower_bound) => {
|
||||
assert_ne!(expected, target, "target {target} should have been Found");
|
||||
assert!(lower_bound > target);
|
||||
// The lower bound must never skip past the true next match.
|
||||
assert!(lower_bound <= expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_query_fast_optional_field_minimum() {
|
||||
|
||||
@@ -8,12 +8,10 @@ use super::range_query_fastfield::FastFieldRangeWeight;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::range_query::is_type_valid_for_fastfield_range_query;
|
||||
use crate::query::{
|
||||
box_scorer, BitSetDocSet, ConstScorer, EnableScoring, Explanation, Query, Scorer, Weight,
|
||||
};
|
||||
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, DocSet, Score};
|
||||
use crate::{DocId, Score};
|
||||
|
||||
/// `RangeQuery` matches all documents that have at least one term within a defined range.
|
||||
///
|
||||
@@ -230,12 +228,21 @@ impl Weight for InvertedIndexRangeWeight {
|
||||
}
|
||||
processed_count += 1;
|
||||
let term_info = term_range.value();
|
||||
let mut postings =
|
||||
inverted_index.read_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
|
||||
postings.fill_bitset(&mut doc_bitset);
|
||||
let mut block_segment_postings = inverted_index
|
||||
.read_block_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
|
||||
loop {
|
||||
let docs = block_segment_postings.docs();
|
||||
if docs.is_empty() {
|
||||
break;
|
||||
}
|
||||
for &doc in block_segment_postings.docs() {
|
||||
doc_bitset.insert(doc);
|
||||
}
|
||||
block_segment_postings.advance();
|
||||
}
|
||||
}
|
||||
let doc_bitset = BitSetDocSet::from(doc_bitset);
|
||||
Ok(box_scorer(ConstScorer::new(doc_bitset, boost)))
|
||||
Ok(Box::new(ConstScorer::new(doc_bitset, boost)))
|
||||
}
|
||||
|
||||
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
|
||||
|
||||
@@ -13,8 +13,7 @@ use common::bounds::{BoundsRange, TransformBound};
|
||||
|
||||
use super::fast_field_range_doc_set::RangeDocSet;
|
||||
use crate::query::{
|
||||
box_scorer, AllScorer, ConstScorer, EmptyScorer, EnableScoring, Explanation, Query, Scorer,
|
||||
Weight,
|
||||
AllScorer, ConstScorer, EmptyScorer, EnableScoring, Explanation, Query, Scorer, Weight,
|
||||
};
|
||||
use crate::schema::{Type, ValueBytes};
|
||||
use crate::{DocId, DocSet, Score, SegmentReader, TantivyError, Term};
|
||||
@@ -56,7 +55,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
|
||||
// Check if both bounds are Bound::Unbounded
|
||||
if self.bounds.is_unbounded() {
|
||||
return Ok(box_scorer(AllScorer::new(reader.max_doc())));
|
||||
return Ok(Box::new(AllScorer::new(reader.max_doc())));
|
||||
}
|
||||
|
||||
let term = self
|
||||
@@ -96,7 +95,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
let Some(str_dict_column): Option<StrColumn> =
|
||||
reader.fast_fields().str(&field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
let dict = str_dict_column.dictionary();
|
||||
|
||||
@@ -108,7 +107,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
let Some((column, _col_type)) = fast_field_reader
|
||||
.u64_lenient_for_type(Some(&[ColumnType::Str]), &field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
search_on_u64_ff(column, boost, BoundsRange::new(lower_bound, upper_bound))
|
||||
}
|
||||
@@ -120,7 +119,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
let Some((column, _col_type)) = fast_field_reader
|
||||
.u64_lenient_for_type(Some(&[ColumnType::DateTime]), &field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
let bounds = bounds.map_bound(|term| term.as_date().unwrap().to_u64());
|
||||
search_on_u64_ff(
|
||||
@@ -147,7 +146,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
let Some(ip_addr_column): Option<Column<Ipv6Addr>> =
|
||||
reader.fast_fields().column_opt(&field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
let value_range = bound_range_inclusive_ip(
|
||||
&bounds.lower_bound,
|
||||
@@ -156,11 +155,11 @@ impl Weight for FastFieldRangeWeight {
|
||||
ip_addr_column.max_value(),
|
||||
);
|
||||
let docset = RangeDocSet::new(value_range, ip_addr_column);
|
||||
Ok(box_scorer(ConstScorer::new(docset, boost)))
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
} else if field_type.is_str() {
|
||||
let Some(str_dict_column): Option<StrColumn> = reader.fast_fields().str(&field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
let dict = str_dict_column.dictionary();
|
||||
|
||||
@@ -172,7 +171,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
let Some((column, _col_type)) =
|
||||
fast_field_reader.u64_lenient_for_type(None, &field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
search_on_u64_ff(column, boost, BoundsRange::new(lower_bound, upper_bound))
|
||||
} else if field_type.is_bytes() {
|
||||
@@ -229,7 +228,7 @@ impl Weight for FastFieldRangeWeight {
|
||||
&field_name,
|
||||
)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
search_on_u64_ff(
|
||||
column,
|
||||
@@ -270,7 +269,7 @@ fn search_on_json_numerical_field(
|
||||
let Some((column, col_type)) =
|
||||
fast_field_reader.u64_lenient_for_type(allowed_column_types, field_name)?
|
||||
else {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
};
|
||||
let actual_column_type: NumericalType = col_type
|
||||
.numerical_type()
|
||||
@@ -428,18 +427,18 @@ fn search_on_u64_ff(
|
||||
)
|
||||
.unwrap_or(1..=0); // empty range
|
||||
if value_range.is_empty() {
|
||||
return Ok(box_scorer(EmptyScorer));
|
||||
return Ok(Box::new(EmptyScorer));
|
||||
}
|
||||
if col_min_value >= *value_range.start() && col_max_value <= *value_range.end() {
|
||||
// all values in the column are within the range.
|
||||
if column.index.get_cardinality() == Cardinality::Full {
|
||||
if boost != 1.0f32 {
|
||||
return Ok(box_scorer(ConstScorer::new(
|
||||
return Ok(Box::new(ConstScorer::new(
|
||||
AllScorer::new(column.num_docs()),
|
||||
boost,
|
||||
)));
|
||||
} else {
|
||||
return Ok(box_scorer(AllScorer::new(column.num_docs())));
|
||||
return Ok(Box::new(AllScorer::new(column.num_docs())));
|
||||
}
|
||||
} else {
|
||||
// TODO Make it a field presence request for that specific column
|
||||
@@ -447,7 +446,7 @@ fn search_on_u64_ff(
|
||||
}
|
||||
|
||||
let docset = RangeDocSet::new(value_range, column);
|
||||
Ok(box_scorer(ConstScorer::new(docset, boost)))
|
||||
Ok(Box::new(ConstScorer::new(docset, boost)))
|
||||
}
|
||||
|
||||
/// Returns true if the type maps to a u64 fast field
|
||||
|
||||
@@ -1,11 +1,9 @@
|
||||
use std::mem::{transmute_copy, ManuallyDrop};
|
||||
use std::ops::DerefMut;
|
||||
|
||||
use downcast_rs::impl_downcast;
|
||||
|
||||
use crate::docset::DocSet;
|
||||
use crate::query::Explanation;
|
||||
use crate::{DocId, Score, TERMINATED};
|
||||
use crate::Score;
|
||||
|
||||
/// Scored set of documents matching a query within a specific segment.
|
||||
///
|
||||
@@ -15,53 +13,6 @@ pub trait Scorer: downcast_rs::Downcast + DocSet + 'static {
|
||||
///
|
||||
/// This method will perform a bit of computation and is not cached.
|
||||
fn score(&mut self) -> Score;
|
||||
|
||||
/// Calls `callback` with all of the `(doc, score)` for which score
|
||||
/// is exceeding a given threshold.
|
||||
///
|
||||
/// This method is useful for the TopDocs collector.
|
||||
/// For all docsets, the blanket implementation has the benefit
|
||||
/// of prefiltering (doc, score) pairs, avoiding the
|
||||
/// virtual dispatch cost.
|
||||
///
|
||||
/// More importantly, it makes it possible for scorers to implement
|
||||
/// important optimization (e.g. BlockWAND for union).
|
||||
fn for_each_pruning(
|
||||
&mut self,
|
||||
threshold: Score,
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) {
|
||||
for_each_pruning_scorer_default_impl(self, threshold, callback);
|
||||
}
|
||||
|
||||
/// Calls `callback` with all of the `(doc, score)` in the scorer.
|
||||
fn for_each(&mut self, callback: &mut dyn FnMut(DocId, Score)) {
|
||||
let mut doc = self.doc();
|
||||
while doc != TERMINATED {
|
||||
callback(doc, self.score());
|
||||
doc = self.advance();
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns an explanation for the score of the current document.
|
||||
fn explain(&mut self) -> Explanation {
|
||||
let score = self.score();
|
||||
let name = std::any::type_name_of_val(self);
|
||||
Explanation::new(name, score)
|
||||
}
|
||||
}
|
||||
|
||||
/// Boxes a scorer. Prefer this to Box::new as it avoids double boxing
|
||||
/// when TScorer is already a Box<dyn Scorer>.
|
||||
pub fn box_scorer<TScorer: Scorer>(scorer: TScorer) -> Box<dyn Scorer> {
|
||||
if std::any::TypeId::of::<TScorer>() == std::any::TypeId::of::<Box<dyn Scorer>>() {
|
||||
unsafe {
|
||||
let forget_me = ManuallyDrop::new(scorer);
|
||||
transmute_copy::<TScorer, Box<dyn Scorer>>(&forget_me)
|
||||
}
|
||||
} else {
|
||||
Box::new(scorer)
|
||||
}
|
||||
}
|
||||
|
||||
impl_downcast!(Scorer);
|
||||
@@ -71,41 +22,4 @@ impl Scorer for Box<dyn Scorer> {
|
||||
fn score(&mut self) -> Score {
|
||||
self.deref_mut().score()
|
||||
}
|
||||
|
||||
fn for_each_pruning(
|
||||
&mut self,
|
||||
threshold: Score,
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) {
|
||||
self.deref_mut().for_each_pruning(threshold, callback);
|
||||
}
|
||||
|
||||
fn for_each(&mut self, callback: &mut dyn FnMut(DocId, Score)) {
|
||||
self.deref_mut().for_each(callback);
|
||||
}
|
||||
}
|
||||
|
||||
/// Calls `callback` with all of the `(doc, score)` for which score
|
||||
/// is exceeding a given threshold.
|
||||
///
|
||||
/// This method is useful for the [`TopDocs`](crate::collector::TopDocs) collector.
|
||||
/// For all docsets, the blanket implementation has the benefit
|
||||
/// of prefiltering (doc, score) pairs, avoiding the
|
||||
/// virtual dispatch cost.
|
||||
///
|
||||
/// More importantly, it makes it possible for scorers to implement
|
||||
/// important optimization (e.g. BlockWAND for union).
|
||||
pub(crate) fn for_each_pruning_scorer_default_impl<TScorer: Scorer + ?Sized>(
|
||||
scorer: &mut TScorer,
|
||||
mut threshold: Score,
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) {
|
||||
let mut doc = scorer.doc();
|
||||
while doc != TERMINATED {
|
||||
let score = scorer.score();
|
||||
if score > threshold {
|
||||
threshold = callback(doc, score);
|
||||
}
|
||||
doc = scorer.advance();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,9 +4,9 @@ mod term_weight;
|
||||
|
||||
pub use self::term_query::TermQuery;
|
||||
pub use self::term_scorer::TermScorer;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use crate::collector::TopDocs;
|
||||
use crate::docset::DocSet;
|
||||
use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
|
||||
|
||||
@@ -1,27 +1,23 @@
|
||||
use crate::codec::postings::{PostingsCodec, PostingsWithBlockMax};
|
||||
use crate::codec::{Codec, StandardCodec};
|
||||
use crate::docset::DocSet;
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::postings::Postings;
|
||||
use crate::postings::{BlockSegmentPostings, FreqReadingOption, Postings, SegmentPostings};
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::{Explanation, Scorer};
|
||||
use crate::{DocId, Score};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct TermScorer<
|
||||
TPostings: Postings = <<StandardCodec as Codec>::PostingsCodec as PostingsCodec>::Postings,
|
||||
> {
|
||||
postings: TPostings,
|
||||
pub struct TermScorer {
|
||||
postings: SegmentPostings,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
similarity_weight: Bm25Weight,
|
||||
}
|
||||
|
||||
impl<TPostings: Postings> TermScorer<TPostings> {
|
||||
impl TermScorer {
|
||||
pub fn new(
|
||||
postings: TPostings,
|
||||
postings: SegmentPostings,
|
||||
fieldnorm_reader: FieldNormReader,
|
||||
similarity_weight: Bm25Weight,
|
||||
) -> TermScorer<TPostings> {
|
||||
) -> TermScorer {
|
||||
TermScorer {
|
||||
postings,
|
||||
fieldnorm_reader,
|
||||
@@ -29,35 +25,10 @@ impl<TPostings: Postings> TermScorer<TPostings> {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn term_freq(&self) -> u32 {
|
||||
self.postings.term_freq()
|
||||
pub(crate) fn seek_block(&mut self, target_doc: DocId) {
|
||||
self.postings.block_cursor.seek_block(target_doc);
|
||||
}
|
||||
|
||||
pub fn fieldnorm_id(&self) -> u8 {
|
||||
self.fieldnorm_reader.fieldnorm_id(self.doc())
|
||||
}
|
||||
|
||||
pub fn max_score(&self) -> Score {
|
||||
self.similarity_weight.max_score()
|
||||
}
|
||||
}
|
||||
|
||||
impl<TPostingsWithBlockMax: PostingsWithBlockMax> TermScorer<TPostingsWithBlockMax> {
|
||||
pub(crate) fn last_doc_in_block(&self) -> DocId {
|
||||
self.postings.last_doc_in_block()
|
||||
}
|
||||
|
||||
/// Advances the term scorer to the block containing target_doc and returns
|
||||
/// an upperbound for the score all of the documents in the block.
|
||||
/// (BlockMax). This score is not guaranteed to be the
|
||||
/// effective maximum score of the block.
|
||||
pub(crate) fn seek_block_max(&mut self, target_doc: DocId) -> Score {
|
||||
self.postings
|
||||
.seek_block_max(target_doc, &self.fieldnorm_reader, &self.similarity_weight)
|
||||
}
|
||||
}
|
||||
|
||||
impl TermScorer {
|
||||
#[cfg(test)]
|
||||
pub fn create_for_test(
|
||||
doc_and_tfs: &[(DocId, u32)],
|
||||
@@ -73,15 +44,75 @@ impl TermScorer {
|
||||
.unwrap_or(0u32)
|
||||
< fieldnorms.len() as u32
|
||||
);
|
||||
type SegmentPostings = <<StandardCodec as Codec>::PostingsCodec as PostingsCodec>::Postings;
|
||||
let segment_postings: SegmentPostings =
|
||||
let segment_postings =
|
||||
SegmentPostings::create_from_docs_and_tfs(doc_and_tfs, Some(fieldnorms));
|
||||
let fieldnorm_reader = FieldNormReader::for_test(fieldnorms);
|
||||
TermScorer::new(segment_postings, fieldnorm_reader, similarity_weight)
|
||||
}
|
||||
|
||||
/// See `FreqReadingOption`.
|
||||
pub(crate) fn freq_reading_option(&self) -> FreqReadingOption {
|
||||
self.postings.block_cursor.freq_reading_option()
|
||||
}
|
||||
|
||||
/// Returns the maximum score for the current block.
|
||||
///
|
||||
/// In some rare case, the result may not be exact. In this case a lower value is returned,
|
||||
/// (and may lead us to return a lesser document).
|
||||
///
|
||||
/// At index time, we store the (fieldnorm_id, term frequency) pair that maximizes the
|
||||
/// score assuming the average fieldnorm computed on this segment.
|
||||
///
|
||||
/// Though extremely rare, it is theoretically possible that the actual average fieldnorm
|
||||
/// is different enough from the current segment average fieldnorm that the maximum over a
|
||||
/// specific is achieved on a different document.
|
||||
///
|
||||
/// (The result is on the other hand guaranteed to be correct if there is only one segment).
|
||||
pub fn block_max_score(&mut self) -> Score {
|
||||
self.postings
|
||||
.block_cursor
|
||||
.block_max_score(&self.fieldnorm_reader, &self.similarity_weight)
|
||||
}
|
||||
|
||||
pub fn term_freq(&self) -> u32 {
|
||||
self.postings.term_freq()
|
||||
}
|
||||
|
||||
pub fn fieldnorm_id(&self) -> u8 {
|
||||
self.fieldnorm_reader.fieldnorm_id(self.doc())
|
||||
}
|
||||
|
||||
pub fn explain(&self) -> Explanation {
|
||||
let fieldnorm_id = self.fieldnorm_id();
|
||||
let term_freq = self.term_freq();
|
||||
self.similarity_weight.explain(fieldnorm_id, term_freq)
|
||||
}
|
||||
|
||||
pub fn max_score(&self) -> Score {
|
||||
self.similarity_weight.max_score()
|
||||
}
|
||||
|
||||
pub fn last_doc_in_block(&self) -> DocId {
|
||||
self.postings.block_cursor.skip_reader().last_doc_in_block()
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying block cursor.
|
||||
pub(crate) fn block_cursor(&mut self) -> &mut BlockSegmentPostings {
|
||||
&mut self.postings.block_cursor
|
||||
}
|
||||
|
||||
/// Returns a reference to the fieldnorm reader for batch lookups.
|
||||
pub(crate) fn fieldnorm_reader(&self) -> &FieldNormReader {
|
||||
&self.fieldnorm_reader
|
||||
}
|
||||
|
||||
/// Returns a reference to the BM25 weight for batch score computation.
|
||||
pub(crate) fn bm25_weight(&self) -> &Bm25Weight {
|
||||
&self.similarity_weight
|
||||
}
|
||||
}
|
||||
|
||||
impl<TPostings: Postings> DocSet for TermScorer<TPostings> {
|
||||
impl DocSet for TermScorer {
|
||||
#[inline]
|
||||
fn advance(&mut self) -> DocId {
|
||||
self.postings.advance()
|
||||
@@ -109,19 +140,13 @@ impl<TPostings: Postings> DocSet for TermScorer<TPostings> {
|
||||
// and do not have access to x86 to investigate.
|
||||
}
|
||||
|
||||
impl<TPostings: Postings> Scorer for TermScorer<TPostings> {
|
||||
impl Scorer for TermScorer {
|
||||
#[inline]
|
||||
fn score(&mut self) -> Score {
|
||||
let fieldnorm_id = self.fieldnorm_id();
|
||||
let term_freq = self.term_freq();
|
||||
self.similarity_weight.score(fieldnorm_id, term_freq)
|
||||
}
|
||||
|
||||
fn explain(&mut self) -> Explanation {
|
||||
let fieldnorm_id = self.fieldnorm_id();
|
||||
let term_freq = self.term_freq();
|
||||
self.similarity_weight.explain(fieldnorm_id, term_freq)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -130,7 +155,7 @@ mod tests {
|
||||
|
||||
use crate::index::SegmentId;
|
||||
use crate::indexer::index_writer::MEMORY_BUDGET_NUM_BYTES_MIN;
|
||||
use crate::indexer::NoMergePolicy;
|
||||
use crate::merge_policy::NoMergePolicy;
|
||||
use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::{Bm25Weight, EnableScoring, Scorer, TermQuery};
|
||||
@@ -151,7 +176,7 @@ mod tests {
|
||||
crate::assert_nearly_equals!(max_scorer, 1.3990127);
|
||||
assert_eq!(term_scorer.doc(), 2);
|
||||
assert_eq!(term_scorer.term_freq(), 3);
|
||||
assert_nearly_equals!(term_scorer.seek_block_max(2), 1.3676447);
|
||||
assert_nearly_equals!(term_scorer.block_max_score(), 1.3676447);
|
||||
assert_nearly_equals!(term_scorer.score(), 1.0892314);
|
||||
assert_eq!(term_scorer.advance(), 3);
|
||||
assert_eq!(term_scorer.doc(), 3);
|
||||
@@ -166,9 +191,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_term_scorer_shallow_advance() {
|
||||
fn test_term_scorer_shallow_advance() -> crate::Result<()> {
|
||||
let bm25_weight = Bm25Weight::for_one_term(300, 1024, 10.0);
|
||||
let mut doc_and_tfs = Vec::new();
|
||||
let mut doc_and_tfs = vec![];
|
||||
for i in 0u32..300u32 {
|
||||
let doc = i * 10;
|
||||
doc_and_tfs.push((doc, 1u32 + doc % 3u32));
|
||||
@@ -176,10 +201,11 @@ mod tests {
|
||||
let fieldnorms: Vec<u32> = std::iter::repeat_n(10u32, 3_000).collect();
|
||||
let mut term_scorer = TermScorer::create_for_test(&doc_and_tfs, &fieldnorms, bm25_weight);
|
||||
assert_eq!(term_scorer.doc(), 0u32);
|
||||
term_scorer.seek_block_max(1289);
|
||||
term_scorer.seek_block(1289);
|
||||
assert_eq!(term_scorer.doc(), 0u32);
|
||||
term_scorer.seek(1289);
|
||||
assert_eq!(term_scorer.doc(), 1290);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
proptest! {
|
||||
@@ -213,7 +239,7 @@ mod tests {
|
||||
|
||||
let docs: Vec<DocId> = (0..term_doc_freq).map(|doc| doc as DocId).collect();
|
||||
for block in docs.chunks(COMPRESSION_BLOCK_SIZE) {
|
||||
let block_max_score: Score = term_scorer.seek_block_max(0);
|
||||
let block_max_score: Score = term_scorer.block_max_score();
|
||||
let mut block_max_score_computed: Score = 0.0;
|
||||
for &doc in block {
|
||||
assert_eq!(term_scorer.doc(), doc);
|
||||
@@ -241,24 +267,25 @@ mod tests {
|
||||
let fieldnorms: Vec<u32> = std::iter::repeat_n(20u32, 300).collect();
|
||||
let bm25_weight = Bm25Weight::for_one_term(10, 129, 20.0);
|
||||
let mut docs = TermScorer::create_for_test(&doc_tfs[..], &fieldnorms[..], bm25_weight);
|
||||
assert_nearly_equals!(docs.seek_block_max(0), 2.5161593);
|
||||
assert_nearly_equals!(docs.seek_block_max(135), 3.4597192);
|
||||
assert_nearly_equals!(docs.block_max_score(), 2.5161593);
|
||||
docs.seek_block(135);
|
||||
assert_nearly_equals!(docs.block_max_score(), 3.4597192);
|
||||
docs.seek_block(256);
|
||||
// the block is not loaded yet.
|
||||
assert_nearly_equals!(docs.seek_block_max(256), 5.2971773);
|
||||
assert_nearly_equals!(docs.block_max_score(), 5.2971773);
|
||||
assert_eq!(256, docs.seek(256));
|
||||
assert_nearly_equals!(docs.seek_block_max(256), 3.9539647);
|
||||
assert_nearly_equals!(docs.block_max_score(), 3.9539647);
|
||||
}
|
||||
|
||||
fn test_block_wand_aux(term_query: &TermQuery, searcher: &Searcher) {
|
||||
let term_weight = term_query
|
||||
.specialized_weight(EnableScoring::enabled_from_searcher(searcher))
|
||||
.unwrap();
|
||||
fn test_block_wand_aux(term_query: &TermQuery, searcher: &Searcher) -> crate::Result<()> {
|
||||
let term_weight =
|
||||
term_query.specialized_weight(EnableScoring::enabled_from_searcher(searcher))?;
|
||||
for reader in searcher.segment_readers() {
|
||||
let mut block_max_scores = vec![];
|
||||
let mut block_max_scores_b = vec![];
|
||||
let mut docs = vec![];
|
||||
{
|
||||
let mut term_scorer = term_weight.term_scorer_for_test(reader, 1.0).unwrap();
|
||||
let mut term_scorer = term_weight.term_scorer_for_test(reader, 1.0)?.unwrap();
|
||||
while term_scorer.doc() != TERMINATED {
|
||||
let mut score = term_scorer.score();
|
||||
docs.push(term_scorer.doc());
|
||||
@@ -272,10 +299,10 @@ mod tests {
|
||||
}
|
||||
}
|
||||
{
|
||||
let mut term_scorer = term_weight.term_scorer_for_test(reader, 1.0).unwrap();
|
||||
let mut term_scorer = term_weight.term_scorer_for_test(reader, 1.0)?.unwrap();
|
||||
for d in docs {
|
||||
let block_max_score = term_scorer.seek_block_max(d);
|
||||
block_max_scores_b.push(block_max_score);
|
||||
term_scorer.seek_block(d);
|
||||
block_max_scores_b.push(term_scorer.block_max_score());
|
||||
}
|
||||
}
|
||||
for (l, r) in block_max_scores
|
||||
@@ -286,18 +313,18 @@ mod tests {
|
||||
assert_nearly_equals!(l, r);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[ignore]
|
||||
#[test]
|
||||
fn test_block_wand_long_test() {
|
||||
fn test_block_wand_long_test() -> crate::Result<()> {
|
||||
let mut schema_builder = Schema::builder();
|
||||
let text_field = schema_builder.add_text_field("text", TEXT);
|
||||
let schema = schema_builder.build();
|
||||
let index = Index::create_in_ram(schema);
|
||||
let mut writer: IndexWriter = index
|
||||
.writer_with_num_threads(3, 3 * MEMORY_BUDGET_NUM_BYTES_MIN)
|
||||
.unwrap();
|
||||
let mut writer: IndexWriter =
|
||||
index.writer_with_num_threads(3, 3 * MEMORY_BUDGET_NUM_BYTES_MIN)?;
|
||||
use rand::Rng;
|
||||
let mut rng = rand::rng();
|
||||
writer.set_merge_policy(Box::new(NoMergePolicy));
|
||||
@@ -305,15 +332,15 @@ mod tests {
|
||||
let term_freq = rng.random_range(1..10000);
|
||||
let words: Vec<&str> = std::iter::repeat_n("bbbb", term_freq).collect();
|
||||
let text = words.join(" ");
|
||||
writer.add_document(doc!(text_field=>text)).unwrap();
|
||||
writer.add_document(doc!(text_field=>text))?;
|
||||
}
|
||||
writer.commit().unwrap();
|
||||
writer.commit()?;
|
||||
let term_query = TermQuery::new(
|
||||
Term::from_field_text(text_field, "bbbb"),
|
||||
IndexRecordOption::WithFreqs,
|
||||
);
|
||||
let segment_ids: Vec<SegmentId>;
|
||||
let reader = index.reader().unwrap();
|
||||
let reader = index.reader()?;
|
||||
{
|
||||
let searcher = reader.searcher();
|
||||
segment_ids = searcher
|
||||
@@ -321,14 +348,15 @@ mod tests {
|
||||
.iter()
|
||||
.map(|segment| segment.segment_id())
|
||||
.collect();
|
||||
test_block_wand_aux(&term_query, &searcher);
|
||||
test_block_wand_aux(&term_query, &searcher)?;
|
||||
}
|
||||
writer.merge(&segment_ids[..]).wait().unwrap();
|
||||
{
|
||||
reader.reload().unwrap();
|
||||
reader.reload()?;
|
||||
let searcher = reader.searcher();
|
||||
assert_eq!(searcher.segment_readers().len(), 1);
|
||||
test_block_wand_aux(&term_query, &searcher);
|
||||
test_block_wand_aux(&term_query, &searcher)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
use super::term_scorer::TermScorer;
|
||||
use crate::docset::{DocSet, COLLECT_BLOCK_BUFFER_LEN};
|
||||
use crate::fieldnorm::FieldNormReader;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::postings::SegmentPostings;
|
||||
use crate::query::bm25::Bm25Weight;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::term_query::TermScorer;
|
||||
use crate::query::weight::for_each_docset_buffered;
|
||||
use crate::query::{box_scorer, AllScorer, AllWeight, EmptyScorer, Explanation, Scorer, Weight};
|
||||
use crate::query::weight::{for_each_docset_buffered, for_each_scorer};
|
||||
use crate::query::{AllScorer, AllWeight, EmptyScorer, Explanation, Scorer, Weight};
|
||||
use crate::schema::IndexRecordOption;
|
||||
use crate::{DocId, Score, TantivyError, Term};
|
||||
|
||||
@@ -17,7 +18,7 @@ pub struct TermWeight {
|
||||
}
|
||||
|
||||
enum TermOrEmptyOrAllScorer {
|
||||
TermScorer(Box<dyn Scorer>),
|
||||
TermScorer(Box<TermScorer>),
|
||||
Empty,
|
||||
AllMatch(AllScorer),
|
||||
}
|
||||
@@ -26,8 +27,8 @@ impl TermOrEmptyOrAllScorer {
|
||||
pub fn into_boxed_scorer(self) -> Box<dyn Scorer> {
|
||||
match self {
|
||||
TermOrEmptyOrAllScorer::TermScorer(scorer) => scorer,
|
||||
TermOrEmptyOrAllScorer::Empty => box_scorer(EmptyScorer),
|
||||
TermOrEmptyOrAllScorer::AllMatch(scorer) => box_scorer(scorer),
|
||||
TermOrEmptyOrAllScorer::Empty => Box::new(EmptyScorer),
|
||||
TermOrEmptyOrAllScorer::AllMatch(scorer) => Box::new(scorer),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -43,7 +44,6 @@ impl Weight for TermWeight {
|
||||
if term_scorer.doc() > doc || term_scorer.seek(doc) != doc {
|
||||
return Err(does_not_match(doc));
|
||||
}
|
||||
let mut term_scorer = term_scorer.downcast::<TermScorer>().ok().unwrap();
|
||||
let mut explanation = term_scorer.explain();
|
||||
explanation.add_context(format!("Term={:?}", self.term,));
|
||||
Ok(explanation)
|
||||
@@ -73,11 +73,11 @@ impl Weight for TermWeight {
|
||||
) -> crate::Result<()> {
|
||||
match self.specialized_scorer(reader, 1.0)? {
|
||||
TermOrEmptyOrAllScorer::TermScorer(mut term_scorer) => {
|
||||
term_scorer.for_each(callback);
|
||||
for_each_scorer(&mut *term_scorer, callback);
|
||||
}
|
||||
TermOrEmptyOrAllScorer::Empty => {}
|
||||
TermOrEmptyOrAllScorer::AllMatch(mut all_scorer) => {
|
||||
all_scorer.for_each(callback);
|
||||
for_each_scorer(&mut all_scorer, callback);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
@@ -124,9 +124,11 @@ impl Weight for TermWeight {
|
||||
let specialized_scorer = self.specialized_scorer(reader, 1.0)?;
|
||||
match specialized_scorer {
|
||||
TermOrEmptyOrAllScorer::TermScorer(term_scorer) => {
|
||||
reader
|
||||
.codec()
|
||||
.for_each_pruning(threshold, term_scorer, callback);
|
||||
crate::query::boolean_query::block_wand_single_scorer(
|
||||
*term_scorer,
|
||||
threshold,
|
||||
callback,
|
||||
);
|
||||
}
|
||||
TermOrEmptyOrAllScorer::Empty => {}
|
||||
TermOrEmptyOrAllScorer::AllMatch(_) => {
|
||||
@@ -166,15 +168,12 @@ impl TermWeight {
|
||||
&self,
|
||||
reader: &SegmentReader,
|
||||
boost: Score,
|
||||
) -> Option<super::TermScorer> {
|
||||
let scorer = self.specialized_scorer(reader, boost).unwrap();
|
||||
match scorer {
|
||||
TermOrEmptyOrAllScorer::TermScorer(scorer) => {
|
||||
let term_scorer = scorer.downcast::<super::TermScorer>().ok()?;
|
||||
Some(*term_scorer)
|
||||
}
|
||||
) -> crate::Result<Option<TermScorer>> {
|
||||
let scorer = self.specialized_scorer(reader, boost)?;
|
||||
Ok(match scorer {
|
||||
TermOrEmptyOrAllScorer::TermScorer(scorer) => Some(*scorer),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn specialized_scorer(
|
||||
@@ -197,16 +196,14 @@ impl TermWeight {
|
||||
)));
|
||||
}
|
||||
|
||||
let segment_postings: SegmentPostings =
|
||||
inverted_index.read_postings_from_terminfo(&term_info, self.index_record_option)?;
|
||||
|
||||
let fieldnorm_reader = self.fieldnorm_reader(reader)?;
|
||||
let similarity_weight = self.similarity_weight.boost_by(boost);
|
||||
let term_scorer = inverted_index.new_term_scorer(
|
||||
&term_info,
|
||||
self.index_record_option,
|
||||
fieldnorm_reader,
|
||||
similarity_weight,
|
||||
)?;
|
||||
|
||||
Ok(TermOrEmptyOrAllScorer::TermScorer(term_scorer))
|
||||
Ok(TermOrEmptyOrAllScorer::TermScorer(Box::new(
|
||||
TermScorer::new(segment_postings, fieldnorm_reader, similarity_weight),
|
||||
)))
|
||||
}
|
||||
|
||||
fn fieldnorm_reader(&self, segment_reader: &SegmentReader) -> crate::Result<FieldNormReader> {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use std::cell::RefCell;
|
||||
|
||||
use crate::docset::DocSet;
|
||||
use crate::postings::{DocFreq, Postings};
|
||||
use crate::postings::Postings;
|
||||
use crate::query::BitSetDocSet;
|
||||
use crate::DocId;
|
||||
|
||||
@@ -16,9 +16,6 @@ pub struct BitSetPostingUnion<TDocSet> {
|
||||
docsets: RefCell<Vec<TDocSet>>,
|
||||
/// The already unionized BitSet of the docsets
|
||||
bitset: BitSetDocSet,
|
||||
/// The total number of documents in the union (regardless of the position we are in the
|
||||
/// bitset).
|
||||
doc_freq: u32,
|
||||
}
|
||||
|
||||
impl<TDocSet: DocSet> BitSetPostingUnion<TDocSet> {
|
||||
@@ -26,11 +23,9 @@ impl<TDocSet: DocSet> BitSetPostingUnion<TDocSet> {
|
||||
docsets: Vec<TDocSet>,
|
||||
bitset: BitSetDocSet,
|
||||
) -> BitSetPostingUnion<TDocSet> {
|
||||
let doc_freq = bitset.doc_freq();
|
||||
BitSetPostingUnion {
|
||||
docsets: RefCell::new(docsets),
|
||||
bitset,
|
||||
doc_freq,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -51,10 +46,6 @@ impl<TDocSet: Postings> Postings for BitSetPostingUnion<TDocSet> {
|
||||
term_freq
|
||||
}
|
||||
|
||||
fn has_freq(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn append_positions_with_offset(&mut self, offset: u32, output: &mut Vec<u32>) {
|
||||
let curr_doc = self.bitset.doc();
|
||||
let mut docsets = self.docsets.borrow_mut();
|
||||
@@ -73,10 +64,6 @@ impl<TDocSet: Postings> Postings for BitSetPostingUnion<TDocSet> {
|
||||
output.sort_unstable();
|
||||
output.dedup();
|
||||
}
|
||||
|
||||
fn doc_freq(&self) -> DocFreq {
|
||||
DocFreq::Exact(self.doc_freq)
|
||||
}
|
||||
}
|
||||
|
||||
impl<TDocSet: DocSet> DocSet for BitSetPostingUnion<TDocSet> {
|
||||
|
||||
@@ -31,7 +31,7 @@ where P: FnMut(&mut T) -> bool {
|
||||
/// Creates a `DocSet` that iterate through the union of two or more `DocSet`s.
|
||||
pub struct BufferedUnionScorer<TScorer, TScoreCombiner = DoNothingCombiner> {
|
||||
/// Active scorers (already filtered of `TERMINATED`).
|
||||
scorers: Vec<TScorer>,
|
||||
docsets: Vec<TScorer>,
|
||||
/// Sliding window presence map for upcoming docs.
|
||||
///
|
||||
/// There are `HORIZON_NUM_TINYBITSETS` buckets, each covering
|
||||
@@ -46,8 +46,6 @@ pub struct BufferedUnionScorer<TScorer, TScoreCombiner = DoNothingCombiner> {
|
||||
/// hit the same doc within the buffered window.
|
||||
scores: Box<[TScoreCombiner; HORIZON as usize]>,
|
||||
/// Start doc ID (inclusive) of the current sliding window.
|
||||
/// None if the window is not loaded yet. This is true for a freshly created
|
||||
/// BufferedUnionScorer.
|
||||
window_start_doc: DocId,
|
||||
/// Current doc ID of the union.
|
||||
doc: DocId,
|
||||
@@ -83,81 +81,51 @@ fn refill<TScorer: Scorer, TScoreCombiner: ScoreCombiner>(
|
||||
}
|
||||
|
||||
impl<TScorer: Scorer, TScoreCombiner: ScoreCombiner> BufferedUnionScorer<TScorer, TScoreCombiner> {
|
||||
/// Returns the underlying scorers in the union.
|
||||
pub fn into_scorers(self) -> Vec<TScorer> {
|
||||
self.scorers
|
||||
}
|
||||
|
||||
/// Accessor for the underlying scorers in the union.
|
||||
pub fn scorers(&self) -> &[TScorer] {
|
||||
&self.scorers[..]
|
||||
}
|
||||
|
||||
/// num_docs is the number of documents in the segment.
|
||||
pub(crate) fn build(
|
||||
docsets: Vec<TScorer>,
|
||||
score_combiner_fn: impl FnOnce() -> TScoreCombiner,
|
||||
num_docs: u32,
|
||||
) -> BufferedUnionScorer<TScorer, TScoreCombiner> {
|
||||
let score_combiner = score_combiner_fn();
|
||||
let mut non_empty_docsets: Vec<TScorer> = docsets
|
||||
let non_empty_docsets: Vec<TScorer> = docsets
|
||||
.into_iter()
|
||||
.filter(|docset| docset.doc() != TERMINATED)
|
||||
.collect();
|
||||
|
||||
let first_doc: DocId = non_empty_docsets
|
||||
.iter()
|
||||
.map(|docset| docset.doc())
|
||||
.min()
|
||||
.unwrap_or(TERMINATED);
|
||||
let mut score_combiner_cloned = score_combiner;
|
||||
let mut i = 0;
|
||||
while i < non_empty_docsets.len() {
|
||||
let should_remove_docset: bool = {
|
||||
let non_empty_docset = &mut non_empty_docsets[i];
|
||||
if non_empty_docset.doc() != first_doc {
|
||||
false
|
||||
} else {
|
||||
score_combiner_cloned.update(non_empty_docset);
|
||||
non_empty_docsets[i].advance() == TERMINATED
|
||||
}
|
||||
};
|
||||
if should_remove_docset {
|
||||
non_empty_docsets.swap_remove(i);
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
let first_score: Score = score_combiner_cloned.score();
|
||||
BufferedUnionScorer {
|
||||
scorers: non_empty_docsets,
|
||||
let mut union = BufferedUnionScorer {
|
||||
docsets: non_empty_docsets,
|
||||
bitsets: Box::new([TinySet::empty(); HORIZON_NUM_TINYBITSETS]),
|
||||
scores: Box::new([score_combiner; HORIZON as usize]),
|
||||
scores: Box::new([score_combiner_fn(); HORIZON as usize]),
|
||||
bucket_idx: HORIZON_NUM_TINYBITSETS,
|
||||
// That way we will be detected as outside the window,
|
||||
window_start_doc: u32::MAX - HORIZON,
|
||||
doc: first_doc,
|
||||
score: first_score,
|
||||
window_start_doc: 0,
|
||||
doc: 0,
|
||||
score: 0.0,
|
||||
num_docs,
|
||||
};
|
||||
if union.refill() {
|
||||
union.advance();
|
||||
} else {
|
||||
union.doc = TERMINATED;
|
||||
}
|
||||
union
|
||||
}
|
||||
|
||||
fn refill(&mut self) -> bool {
|
||||
let Some(min_doc) = self.scorers.iter().map(DocSet::doc).min() else {
|
||||
return false;
|
||||
};
|
||||
// Reset the sliding window to start at the smallest doc
|
||||
// across all scorers and prebuffer within the horizon.
|
||||
self.window_start_doc = min_doc;
|
||||
self.bucket_idx = 0;
|
||||
self.doc = min_doc;
|
||||
refill(
|
||||
&mut self.scorers,
|
||||
&mut self.bitsets,
|
||||
&mut self.scores,
|
||||
min_doc,
|
||||
);
|
||||
true
|
||||
if let Some(min_doc) = self.docsets.iter().map(DocSet::doc).min() {
|
||||
// Reset the sliding window to start at the smallest doc
|
||||
// across all scorers and prebuffer within the horizon.
|
||||
self.window_start_doc = min_doc;
|
||||
self.bucket_idx = 0;
|
||||
self.doc = min_doc;
|
||||
refill(
|
||||
&mut self.docsets,
|
||||
&mut self.bitsets,
|
||||
&mut self.scores,
|
||||
min_doc,
|
||||
);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
@@ -179,7 +147,6 @@ impl<TScorer: Scorer, TScoreCombiner: ScoreCombiner> BufferedUnionScorer<TScorer
|
||||
|
||||
fn is_in_horizon(&self, target: DocId) -> bool {
|
||||
// wrapping_sub, because target may be < window_start_doc
|
||||
// in particular during initialization.
|
||||
let gap = target.wrapping_sub(self.window_start_doc);
|
||||
gap < HORIZON
|
||||
}
|
||||
@@ -249,10 +216,11 @@ where
|
||||
if self.doc >= target {
|
||||
return self.doc;
|
||||
}
|
||||
if self.is_in_horizon(target) {
|
||||
let gap = target - self.window_start_doc;
|
||||
if gap < HORIZON {
|
||||
// Our value is within the buffered horizon.
|
||||
|
||||
// Skipping to corresponding bucket.
|
||||
let gap = target.wrapping_sub(self.window_start_doc);
|
||||
let new_bucket_idx = gap as usize / 64;
|
||||
for obsolete_tinyset in &mut self.bitsets[self.bucket_idx..new_bucket_idx] {
|
||||
obsolete_tinyset.clear();
|
||||
@@ -271,14 +239,16 @@ where
|
||||
doc
|
||||
} else {
|
||||
// clear the buffered info.
|
||||
self.bitsets.fill(TinySet::empty());
|
||||
for obsolete_tinyset in self.bitsets.iter_mut() {
|
||||
*obsolete_tinyset = TinySet::empty();
|
||||
}
|
||||
for score_combiner in self.scores.iter_mut() {
|
||||
score_combiner.clear();
|
||||
}
|
||||
|
||||
// The target is outside of the buffered horizon.
|
||||
// advance all docsets to a doc >= to the target.
|
||||
unordered_drain_filter(&mut self.scorers, |docset| {
|
||||
unordered_drain_filter(&mut self.docsets, |docset| {
|
||||
if docset.doc() < target {
|
||||
docset.seek(target);
|
||||
}
|
||||
@@ -315,8 +285,8 @@ where
|
||||
let mut is_hit = false;
|
||||
let mut min_new_target = TERMINATED;
|
||||
|
||||
for scorer in self.scorers.iter_mut() {
|
||||
match scorer.seek_danger(target) {
|
||||
for docset in self.docsets.iter_mut() {
|
||||
match docset.seek_danger(target) {
|
||||
SeekDangerResult::Found => {
|
||||
is_hit = true;
|
||||
break;
|
||||
@@ -345,11 +315,11 @@ where
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> u32 {
|
||||
estimate_union(self.scorers.iter().map(DocSet::size_hint), self.num_docs)
|
||||
estimate_union(self.docsets.iter().map(DocSet::size_hint), self.num_docs)
|
||||
}
|
||||
|
||||
fn cost(&self) -> u64 {
|
||||
self.scorers.iter().map(|docset| docset.cost()).sum()
|
||||
self.docsets.iter().map(|docset| docset.cost()).sum()
|
||||
}
|
||||
|
||||
// TODO Also implement `count` with deletes efficiently.
|
||||
@@ -357,17 +327,21 @@ where
|
||||
if self.doc == TERMINATED {
|
||||
return 0;
|
||||
}
|
||||
let mut count = 1 + self.bitsets[self.bucket_idx..HORIZON_NUM_TINYBITSETS]
|
||||
let mut count = self.bitsets[self.bucket_idx..HORIZON_NUM_TINYBITSETS]
|
||||
.iter()
|
||||
.copied()
|
||||
.map(TinySet::len)
|
||||
.sum::<u32>();
|
||||
.map(|bitset| bitset.len())
|
||||
.sum::<u32>()
|
||||
+ 1;
|
||||
for bitset in self.bitsets.iter_mut() {
|
||||
bitset.clear();
|
||||
}
|
||||
while self.refill() {
|
||||
count += self.bitsets.iter().copied().map(TinySet::len).sum::<u32>();
|
||||
self.bitsets.fill(TinySet::empty());
|
||||
count += self.bitsets.iter().map(|bitset| bitset.len()).sum::<u32>();
|
||||
for bitset in self.bitsets.iter_mut() {
|
||||
bitset.clear();
|
||||
}
|
||||
}
|
||||
self.bucket_idx = HORIZON_NUM_TINYBITSETS;
|
||||
self.doc = TERMINATED;
|
||||
count
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use crate::docset::{DocSet, TERMINATED};
|
||||
use crate::postings::{DocFreq, Postings};
|
||||
use crate::postings::Postings;
|
||||
use crate::DocId;
|
||||
|
||||
/// A `SimpleUnion` is a `DocSet` that is the union of multiple `DocSet`.
|
||||
@@ -56,22 +56,6 @@ impl<TDocSet: Postings> Postings for SimpleUnion<TDocSet> {
|
||||
term_freq
|
||||
}
|
||||
|
||||
fn has_freq(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// We do not know the actual document frequency, so we return
|
||||
/// the maximum document frequency of the docsets.
|
||||
fn doc_freq(&self) -> DocFreq {
|
||||
let approximate_doc_freq = self
|
||||
.docsets
|
||||
.iter()
|
||||
.map(|docset| u32::from(docset.doc_freq()))
|
||||
.max()
|
||||
.unwrap_or(0u32);
|
||||
DocFreq::Approximate(approximate_doc_freq)
|
||||
}
|
||||
|
||||
fn append_positions_with_offset(&mut self, offset: u32, output: &mut Vec<u32>) {
|
||||
for docset in &mut self.docsets {
|
||||
let doc = docset.doc();
|
||||
|
||||
@@ -1,9 +1,21 @@
|
||||
use super::Scorer;
|
||||
use crate::docset::COLLECT_BLOCK_BUFFER_LEN;
|
||||
use crate::index::SegmentReader;
|
||||
use crate::query::explanation::does_not_match;
|
||||
use crate::query::Explanation;
|
||||
use crate::{DocId, DocSet, Score};
|
||||
use crate::{DocId, DocSet, Score, TERMINATED};
|
||||
|
||||
/// Iterates through all of the documents and scores matched by the DocSet
|
||||
/// `DocSet`.
|
||||
pub(crate) fn for_each_scorer<TScorer: Scorer + ?Sized>(
|
||||
scorer: &mut TScorer,
|
||||
callback: &mut dyn FnMut(DocId, Score),
|
||||
) {
|
||||
let mut doc = scorer.doc();
|
||||
while doc != TERMINATED {
|
||||
callback(doc, scorer.score());
|
||||
doc = scorer.advance();
|
||||
}
|
||||
}
|
||||
|
||||
/// Iterates through all of the documents matched by the DocSet
|
||||
/// `DocSet`.
|
||||
@@ -22,6 +34,31 @@ pub(crate) fn for_each_docset_buffered<T: DocSet + ?Sized>(
|
||||
}
|
||||
}
|
||||
|
||||
/// Calls `callback` with all of the `(doc, score)` for which score
|
||||
/// is exceeding a given threshold.
|
||||
///
|
||||
/// This method is useful for the [`TopDocs`](crate::collector::TopDocs) collector.
|
||||
/// For all docsets, the blanket implementation has the benefit
|
||||
/// of prefiltering (doc, score) pairs, avoiding the
|
||||
/// virtual dispatch cost.
|
||||
///
|
||||
/// More importantly, it makes it possible for scorers to implement
|
||||
/// important optimization (e.g. BlockWAND for union).
|
||||
pub(crate) fn for_each_pruning_scorer<TScorer: Scorer + ?Sized>(
|
||||
scorer: &mut TScorer,
|
||||
mut threshold: Score,
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) {
|
||||
let mut doc = scorer.doc();
|
||||
while doc != TERMINATED {
|
||||
let score = scorer.score();
|
||||
if score > threshold {
|
||||
threshold = callback(doc, score);
|
||||
}
|
||||
doc = scorer.advance();
|
||||
}
|
||||
}
|
||||
|
||||
/// A Weight is the specialization of a `Query`
|
||||
/// for a given set of segments.
|
||||
///
|
||||
@@ -35,13 +72,7 @@ pub trait Weight: Send + Sync + 'static {
|
||||
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>>;
|
||||
|
||||
/// Returns an [`Explanation`] for the given document.
|
||||
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
|
||||
let mut scorer = self.scorer(reader, 1.0)?;
|
||||
if scorer.doc() > doc || scorer.seek(doc) != doc {
|
||||
return Err(does_not_match(doc));
|
||||
}
|
||||
Ok(scorer.explain())
|
||||
}
|
||||
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation>;
|
||||
|
||||
/// Returns the number documents within the given [`SegmentReader`].
|
||||
fn count(&self, reader: &SegmentReader) -> crate::Result<u32> {
|
||||
@@ -61,7 +92,7 @@ pub trait Weight: Send + Sync + 'static {
|
||||
callback: &mut dyn FnMut(DocId, Score),
|
||||
) -> crate::Result<()> {
|
||||
let mut scorer = self.scorer(reader, 1.0)?;
|
||||
scorer.for_each(callback);
|
||||
for_each_scorer(scorer.as_mut(), callback);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -96,7 +127,7 @@ pub trait Weight: Send + Sync + 'static {
|
||||
callback: &mut dyn FnMut(DocId, Score) -> Score,
|
||||
) -> crate::Result<()> {
|
||||
let mut scorer = self.scorer(reader, 1.0)?;
|
||||
scorer.for_each_pruning(threshold, callback);
|
||||
for_each_pruning_scorer(scorer.as_mut(), threshold, callback);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,6 @@ use arc_swap::ArcSwap;
|
||||
pub use warming::Warmer;
|
||||
|
||||
use self::warming::WarmingState;
|
||||
use crate::codec::Codec;
|
||||
use crate::core::searcher::{SearcherGeneration, SearcherInner};
|
||||
use crate::directory::{Directory, WatchCallback, WatchHandle, META_LOCK};
|
||||
use crate::store::DOCSTORE_CACHE_CAPACITY;
|
||||
@@ -39,17 +38,17 @@ pub enum ReloadPolicy {
|
||||
/// - number of warming threads, for parallelizing warming work
|
||||
/// - The cache size of the underlying doc store readers.
|
||||
#[derive(Clone)]
|
||||
pub struct IndexReaderBuilder<C: Codec = crate::codec::StandardCodec> {
|
||||
pub struct IndexReaderBuilder {
|
||||
reload_policy: ReloadPolicy,
|
||||
index: Index<C>,
|
||||
index: Index,
|
||||
warmers: Vec<Weak<dyn Warmer>>,
|
||||
num_warming_threads: usize,
|
||||
doc_store_cache_num_blocks: usize,
|
||||
}
|
||||
|
||||
impl<C: Codec> IndexReaderBuilder<C> {
|
||||
impl IndexReaderBuilder {
|
||||
#[must_use]
|
||||
pub(crate) fn new(index: Index<C>) -> IndexReaderBuilder<C> {
|
||||
pub(crate) fn new(index: Index) -> IndexReaderBuilder {
|
||||
IndexReaderBuilder {
|
||||
reload_policy: ReloadPolicy::OnCommitWithDelay,
|
||||
index,
|
||||
@@ -64,7 +63,7 @@ impl<C: Codec> IndexReaderBuilder<C> {
|
||||
/// Building the reader is a non-trivial operation that requires
|
||||
/// to open different segment readers. It may take hundreds of milliseconds
|
||||
/// of time and it may return an error.
|
||||
pub fn try_into(self) -> crate::Result<IndexReader<C>> {
|
||||
pub fn try_into(self) -> crate::Result<IndexReader> {
|
||||
let searcher_generation_inventory = Inventory::default();
|
||||
let warming_state = WarmingState::new(
|
||||
self.num_warming_threads,
|
||||
@@ -107,7 +106,7 @@ impl<C: Codec> IndexReaderBuilder<C> {
|
||||
///
|
||||
/// See [`ReloadPolicy`] for more details.
|
||||
#[must_use]
|
||||
pub fn reload_policy(mut self, reload_policy: ReloadPolicy) -> IndexReaderBuilder<C> {
|
||||
pub fn reload_policy(mut self, reload_policy: ReloadPolicy) -> IndexReaderBuilder {
|
||||
self.reload_policy = reload_policy;
|
||||
self
|
||||
}
|
||||
@@ -119,14 +118,14 @@ impl<C: Codec> IndexReaderBuilder<C> {
|
||||
pub fn doc_store_cache_num_blocks(
|
||||
mut self,
|
||||
doc_store_cache_num_blocks: usize,
|
||||
) -> IndexReaderBuilder<C> {
|
||||
) -> IndexReaderBuilder {
|
||||
self.doc_store_cache_num_blocks = doc_store_cache_num_blocks;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the [`Warmer`]s that are invoked when reloading searchable segments.
|
||||
#[must_use]
|
||||
pub fn warmers(mut self, warmers: Vec<Weak<dyn Warmer>>) -> IndexReaderBuilder<C> {
|
||||
pub fn warmers(mut self, warmers: Vec<Weak<dyn Warmer>>) -> IndexReaderBuilder {
|
||||
self.warmers = warmers;
|
||||
self
|
||||
}
|
||||
@@ -136,33 +135,33 @@ impl<C: Codec> IndexReaderBuilder<C> {
|
||||
/// This allows parallelizing warming work when there are multiple [`Warmer`] registered with
|
||||
/// the [`IndexReader`].
|
||||
#[must_use]
|
||||
pub fn num_warming_threads(mut self, num_warming_threads: usize) -> IndexReaderBuilder<C> {
|
||||
pub fn num_warming_threads(mut self, num_warming_threads: usize) -> IndexReaderBuilder {
|
||||
self.num_warming_threads = num_warming_threads;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: Codec> TryInto<IndexReader<C>> for IndexReaderBuilder<C> {
|
||||
impl TryInto<IndexReader> for IndexReaderBuilder {
|
||||
type Error = crate::TantivyError;
|
||||
|
||||
fn try_into(self) -> crate::Result<IndexReader<C>> {
|
||||
fn try_into(self) -> crate::Result<IndexReader> {
|
||||
IndexReaderBuilder::try_into(self)
|
||||
}
|
||||
}
|
||||
|
||||
struct InnerIndexReader<C: Codec> {
|
||||
struct InnerIndexReader {
|
||||
doc_store_cache_num_blocks: usize,
|
||||
index: Index<C>,
|
||||
index: Index,
|
||||
warming_state: WarmingState,
|
||||
searcher: arc_swap::ArcSwap<SearcherInner>,
|
||||
searcher_generation_counter: Arc<AtomicU64>,
|
||||
searcher_generation_inventory: Inventory<SearcherGeneration>,
|
||||
}
|
||||
|
||||
impl<C: Codec> InnerIndexReader<C> {
|
||||
impl InnerIndexReader {
|
||||
fn new(
|
||||
doc_store_cache_num_blocks: usize,
|
||||
index: Index<C>,
|
||||
index: Index,
|
||||
warming_state: WarmingState,
|
||||
// The searcher_generation_inventory is not used as source, but as target to track the
|
||||
// loaded segments.
|
||||
@@ -190,7 +189,7 @@ impl<C: Codec> InnerIndexReader<C> {
|
||||
///
|
||||
/// This function acquires a lock to prevent GC from removing files
|
||||
/// as we are opening our index.
|
||||
fn open_segment_readers(index: &Index<C>) -> crate::Result<Vec<SegmentReader>> {
|
||||
fn open_segment_readers(index: &Index) -> crate::Result<Vec<SegmentReader>> {
|
||||
// Prevents segment files from getting deleted while we are in the process of opening them
|
||||
let _meta_lock = index.directory().acquire_lock(&META_LOCK)?;
|
||||
let searchable_segments = index.searchable_segments()?;
|
||||
@@ -213,7 +212,7 @@ impl<C: Codec> InnerIndexReader<C> {
|
||||
}
|
||||
|
||||
fn create_searcher(
|
||||
index: &Index<C>,
|
||||
index: &Index,
|
||||
doc_store_cache_num_blocks: usize,
|
||||
warming_state: &WarmingState,
|
||||
searcher_generation_counter: &Arc<AtomicU64>,
|
||||
@@ -227,10 +226,9 @@ impl<C: Codec> InnerIndexReader<C> {
|
||||
);
|
||||
|
||||
let schema = index.schema();
|
||||
// SearcherInner uses Index<StandardCodec> since the codec doesn't affect reading
|
||||
let searcher = Arc::new(SearcherInner::new(
|
||||
schema,
|
||||
index.with_standard_codec(),
|
||||
index.clone(),
|
||||
segment_readers,
|
||||
searcher_generation,
|
||||
doc_store_cache_num_blocks,
|
||||
@@ -266,14 +264,14 @@ impl<C: Codec> InnerIndexReader<C> {
|
||||
///
|
||||
/// `IndexReader` just wraps an `Arc`.
|
||||
#[derive(Clone)]
|
||||
pub struct IndexReader<C: Codec = crate::codec::StandardCodec> {
|
||||
inner: Arc<InnerIndexReader<C>>,
|
||||
pub struct IndexReader {
|
||||
inner: Arc<InnerIndexReader>,
|
||||
_watch_handle_opt: Option<WatchHandle>,
|
||||
}
|
||||
|
||||
impl<C: Codec> IndexReader<C> {
|
||||
impl IndexReader {
|
||||
#[cfg(test)]
|
||||
pub(crate) fn index(&self) -> Index<C> {
|
||||
pub(crate) fn index(&self) -> Index {
|
||||
self.inner.index.clone()
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user