rust/tantivy
1
0
Fork 0
mirror of https://github.com/quickwit-oss/tantivy.git synced 2026-11-14 14:47:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
Pascal Seitz
2026-07-31 17:07:51 +02:00
parent e81daa93da
commit cc0fad0214
1 changed files with 118 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aggregation/bucket/multi_terms/mod.rs
+118 -77
View File
@@ -390,10 +390,10 @@ where
/// Partial keys for the current block, grouped by document.
keys_buf: Vec<Packing::PackingType>,
/// Distinct documents represented in `keys_buf`.
active_docs: Vec<DocId>,
/// One document ID per key after Cartesian expansion. Empty while keys and active documents
/// are one-to-one.
expanded_key_docs: Vec<DocId>,
unique_docs: Vec<DocId>,
/// One document ID per partial Cartesian-product combination in `keys_buf`. Empty while
/// combinations and unique documents are one-to-one.
doc_ids_per_partial_combination: Vec<DocId>,
}
/// Validates the field configuration and, when ordering by a sub-aggregation, that the target
@@ -453,6 +453,91 @@ fn missing_value_for_doc(
(!field_has_value).then_some(missing.missing_value)
}
/// Expands the existing partial combinations with every value of one field.
fn expand_partial_combinations_for_field<Packing: MultiTermsPacking>(
packing: &Packing,
field_idx: usize,
field: &MultiTermsFieldAccessor,
missing: Option<&MultiTermsMissingAccessor>,
block_accessor: &ColumnBlockAccessor<u64>,
keys_buf: &mut Vec<Packing::PackingType>,
unique_docs: &mut Vec<DocId>,
doc_ids_per_partial_combination: &mut Vec<DocId>,
) {
// Keep the current partial combinations as a prefix while appending their expansion.
let partial_combination_count = keys_buf.len();
// In one-to-one mode the input document IDs are implicit in `unique_docs`; otherwise they
// occupy the prefix of `doc_ids_per_partial_combination`.
let doc_ids_have_input_prefix = !doc_ids_per_partial_combination.is_empty();
debug_assert_eq!(
partial_combination_count,
if doc_ids_have_input_prefix {
doc_ids_per_partial_combination.len()
} else {
unique_docs.len()
}
);
{
let mut field_values = block_accessor
.iter_docid_vals(unique_docs, &field.column)
.peekable();
let mut doc_values = SmallVec::<[u64; 2]>::new();
let mut combination_start = 0;
for &doc_id in unique_docs.iter() {
let num_combinations_for_doc = if doc_ids_have_input_prefix {
doc_ids_per_partial_combination[combination_start..partial_combination_count]
.iter()
.take_while(|&&combination_doc_id| combination_doc_id == doc_id)
.count()
} else {
1
};
debug_assert_ne!(num_combinations_for_doc, 0);
let combination_end = combination_start + num_combinations_for_doc;
debug_assert!(field_values
.peek()
.is_none_or(|(value_doc, _)| *value_doc >= doc_id));
doc_values.clear();
while let Some((_, value)) = field_values.next_if(|(value_doc, _)| *value_doc == doc_id)
{
doc_values.push(value);
}
if doc_values.is_empty() {
if let Some(value) = missing_value_for_doc(missing, doc_id) {
doc_values.push(value);
}
}
for combination_idx in combination_start..combination_end {
for &value in &doc_values {
let mut key = keys_buf[combination_idx].clone();
packing.push(&mut key, field_idx, value);
keys_buf.push(key);
doc_ids_per_partial_combination.push(doc_id);
}
}
combination_start = combination_end;
}
debug_assert_eq!(combination_start, partial_combination_count);
debug_assert!(field_values.next().is_none());
}
keys_buf.drain(..partial_combination_count);
if doc_ids_have_input_prefix {
doc_ids_per_partial_combination.drain(..partial_combination_count);
}
unique_docs.clear();
unique_docs.extend_from_slice(doc_ids_per_partial_combination);
unique_docs.dedup();
if doc_ids_per_partial_combination.len() == unique_docs.len() {
doc_ids_per_partial_combination.clear();
}
}
/// Builds the segment collector for a multi_terms aggregation, validating the request.
///
/// Every value, including missing, is represented as a plain `u64` and collected through block
@@ -645,9 +730,9 @@ where
for key in &mut self.keys_buf {
packing.clear_packing(key);
}
self.active_docs.clear();
self.active_docs.extend_from_slice(docs);
self.expanded_key_docs.clear();
self.unique_docs.clear();
self.unique_docs.extend_from_slice(docs);
self.doc_ids_per_partial_combination.clear();
let block_accessor = &mut agg_data.column_block_accessor;
if self.all_fields_full {
@@ -664,9 +749,9 @@ where
for (field_idx, field) in self.req_data.fields.iter().enumerate() {
let missing = self.req_data.missing_accessors[field_idx].as_ref();
let has_one_value_per_doc =
fetch_field_block(&self.active_docs, field, missing, block_accessor);
fetch_field_block(&self.unique_docs, field, missing, block_accessor);
if has_one_value_per_doc {
if self.expanded_key_docs.is_empty() {
if self.doc_ids_per_partial_combination.is_empty() {
packing.push_full_values(
&mut self.keys_buf,
field_idx,
@@ -675,8 +760,9 @@ where
} else {
let mut values = block_accessor.iter_vals();
let mut keys = self.keys_buf.iter_mut();
for key_doc_group in
self.expanded_key_docs.chunk_by(|left, right| left == right)
for key_doc_group in self
.doc_ids_per_partial_combination
.chunk_by(|left, right| left == right)
{
let value = values.next().unwrap();
for key in keys.by_ref().take(key_doc_group.len()) {
@@ -690,7 +776,7 @@ where
}
// Until expansion, sparse single-value fields can filter keys in place.
if self.expanded_key_docs.is_empty()
if self.doc_ids_per_partial_combination.is_empty()
&& !field.column.get_cardinality().is_multivalue()
&& missing.is_none()
{
@@ -699,88 +785,43 @@ where
for (&doc_id, &value) in
block_accessor.docids().iter().zip(block_accessor.values())
{
while self.active_docs[source_idx] < doc_id {
while self.unique_docs[source_idx] < doc_id {
source_idx += 1;
}
debug_assert_eq!(self.active_docs[source_idx], doc_id);
debug_assert_eq!(self.unique_docs[source_idx], doc_id);
if target_idx != source_idx {
self.active_docs[target_idx] = doc_id;
self.unique_docs[target_idx] = doc_id;
self.keys_buf.swap(target_idx, source_idx);
}
packing.push(&mut self.keys_buf[target_idx], field_idx, value);
source_idx += 1;
target_idx += 1;
}
self.active_docs.truncate(target_idx);
self.unique_docs.truncate(target_idx);
self.keys_buf.truncate(target_idx);
continue;
}
// Preserve source keys and IDs as prefixes while appending their expansions.
let source_key_count = self.keys_buf.len();
let source_doc_count = self.expanded_key_docs.len();
debug_assert!(source_doc_count == 0 || source_doc_count == source_key_count);
{
let mut field_values = block_accessor
.iter_docid_vals(&self.active_docs, &field.column)
.peekable();
let mut current_doc = None;
let mut doc_values = SmallVec::<[u64; 2]>::new();
for source_idx in 0..source_key_count {
let doc_id = if source_doc_count == 0 {
self.active_docs[source_idx]
} else {
self.expanded_key_docs[source_idx]
};
if current_doc != Some(doc_id) {
current_doc = Some(doc_id);
debug_assert!(field_values
.peek()
.is_none_or(|(value_doc, _)| *value_doc >= doc_id));
doc_values.clear();
while let Some((_, value)) =
field_values.next_if(|(value_doc, _)| *value_doc == doc_id)
{
doc_values.push(value);
}
if doc_values.is_empty() {
if let Some(value) = missing_value_for_doc(missing, doc_id) {
doc_values.push(value);
}
}
}
for &value in &doc_values {
let mut key = self.keys_buf[source_idx].clone();
packing.push(&mut key, field_idx, value);
self.keys_buf.push(key);
self.expanded_key_docs.push(doc_id);
}
}
debug_assert!(field_values.next().is_none());
}
self.keys_buf.drain(..source_key_count);
if source_doc_count != 0 {
self.expanded_key_docs.drain(..source_doc_count);
}
self.active_docs.clear();
self.active_docs.extend_from_slice(&self.expanded_key_docs);
self.active_docs.dedup();
if self.expanded_key_docs.len() == self.active_docs.len() {
self.expanded_key_docs.clear();
}
expand_partial_combinations_for_field(
packing,
field_idx,
field,
missing,
block_accessor,
&mut self.keys_buf,
&mut self.unique_docs,
&mut self.doc_ids_per_partial_combination,
);
}
}
if let Some(sub_agg) = sub_agg {
let key_docs = if self.expanded_key_docs.is_empty() {
self.active_docs.as_slice()
let doc_ids = if self.doc_ids_per_partial_combination.is_empty() {
self.unique_docs.as_slice()
} else {
self.expanded_key_docs.as_slice()
self.doc_ids_per_partial_combination.as_slice()
};
for (&doc_id, key) in key_docs.iter().zip(&self.keys_buf) {
for (&doc_id, key) in doc_ids.iter().zip(&self.keys_buf) {
let bucket_id = buckets.term_entry(key.clone(), bucket_id_provider);
sub_agg.push(bucket_id.to_bucket_id(), doc_id);
}
@@ -1032,8 +1073,8 @@ where
map_init_value,
all_fields_full,
keys_buf: Vec::new(),
active_docs: Vec::new(),
expanded_key_docs: Vec::new(),
unique_docs: Vec::new(),
doc_ids_per_partial_combination: Vec::new(),
}))
}