diff --git a/src/aggregation/bucket/multi_terms/mod.rs b/src/aggregation/bucket/multi_terms/mod.rs index 5fa866906..6c37d0f15 100644 --- a/src/aggregation/bucket/multi_terms/mod.rs +++ b/src/aggregation/bucket/multi_terms/mod.rs @@ -390,10 +390,10 @@ where /// Partial keys for the current block, grouped by document. keys_buf: Vec, /// Distinct documents represented in `keys_buf`. - active_docs: Vec, - /// One document ID per key after Cartesian expansion. Empty while keys and active documents - /// are one-to-one. - expanded_key_docs: Vec, + unique_docs: Vec, + /// 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, } /// 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: &Packing, + field_idx: usize, + field: &MultiTermsFieldAccessor, + missing: Option<&MultiTermsMissingAccessor>, + block_accessor: &ColumnBlockAccessor, + keys_buf: &mut Vec, + unique_docs: &mut Vec, + doc_ids_per_partial_combination: &mut Vec, +) { + // 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(), })) }