aggregation: make block value sources mutable

src/aggregation/block_accessor.rs

@@ -12,7 +12,7 @@ use crate::DocId;
 /// `docids` with one document id per value.
 pub(crate) trait BlockValueSource {
     fn load_block(
-        &self,
+        &mut self,
         docs: &[DocId],
         values: &mut Vec<u64>,
         docids: &mut Vec<DocId>,
@@ -46,7 +46,7 @@ pub(crate) struct ColumnBlockAccessor {
 impl BlockValueSource for Column<u64> {
     #[inline]
     fn load_block(
-        &self,
+        &mut self,
         docs: &[DocId],
         values: &mut Vec<u64>,
         docids: &mut Vec<DocId>,
@@ -68,7 +68,7 @@ impl BlockValueSource for Column<u64> {
 
 impl ColumnBlockAccessor {
     #[inline]
-    pub(crate) fn fetch_block(&mut self, docs: &[DocId], source: &impl BlockValueSource) {
+    pub(crate) fn fetch_block(&mut self, docs: &[DocId], source: &mut impl BlockValueSource) {
         self.cardinality = source.load_block(
             docs,
             &mut self.val_cache,
@@ -93,7 +93,7 @@ impl ColumnBlockAccessor {
     pub(crate) fn fetch_block_with_missing(
         &mut self,
         docs: &[DocId],
-        source: &impl BlockValueSource,
+        source: &mut impl BlockValueSource,
         missing_opt: Option<u64>,
     ) {
         self.fetch_block_with_missing_ordered(docs, source, missing_opt, false)
@@ -106,7 +106,7 @@ impl ColumnBlockAccessor {
     pub(crate) fn fetch_block_with_missing_ordered(
         &mut self,
         docs: &[DocId],
-        source: &impl BlockValueSource,
+        source: &mut impl BlockValueSource,
         missing_opt: Option<u64>,
         ordered: bool,
     ) {
@@ -176,7 +176,7 @@ impl ColumnBlockAccessor {
     pub(crate) fn fetch_block_with_missing_unique_per_doc(
         &mut self,
         docs: &[DocId],
-        source: &impl BlockValueSource,
+        source: &mut impl BlockValueSource,
         missing: Option<u64>,
         ordered: bool,
     ) {
@@ -365,7 +365,7 @@ mod tests {
 
     impl BlockValueSource for TestValueSource {
         fn load_block(
-            &self,
+            &mut self,
             docs: &[DocId],
             values: &mut Vec<u64>,
             docids: &mut Vec<DocId>,
@@ -421,13 +421,13 @@ mod tests {
     #[test]
     fn test_source_neutral_full_block_alignment() {
         let docs = [2, 4, 8];
-        let source = TestValueSource {
+        let mut source = TestValueSource {
             cardinality: Cardinality::Full,
             entries: vec![(2, 20), (4, 40), (8, 80)],
         };
         let mut accessor = ColumnBlockAccessor::default();
 
-        accessor.fetch_block(&docs, &source);
+        accessor.fetch_block(&docs, &mut source);
 
         assert!(accessor.has_one_value_per_doc(&docs));
         assert_eq!(
@@ -439,13 +439,13 @@ mod tests {
     #[test]
     fn test_source_neutral_optional_block_with_missing() {
         let docs = [0, 1, 2, 4];
-        let source = TestValueSource {
+        let mut source = TestValueSource {
             cardinality: Cardinality::Optional,
             entries: vec![(1, 10), (4, 40)],
         };
         let mut accessor = ColumnBlockAccessor::default();
 
-        accessor.fetch_block_with_missing_ordered(&docs, &source, Some(99), true);
+        accessor.fetch_block_with_missing_ordered(&docs, &mut source, Some(99), true);
 
         assert!(accessor.has_one_value_per_doc(&docs));
         assert_eq!(
@@ -457,13 +457,13 @@ mod tests {
     #[test]
     fn test_source_neutral_multivalue_block_deduplication() {
         let docs = [0, 1];
-        let source = TestValueSource {
+        let mut source = TestValueSource {
             cardinality: Cardinality::Multivalued,
             entries: vec![(0, 3), (0, 1), (0, 3), (1, 5), (1, 5)],
         };
         let mut accessor = ColumnBlockAccessor::default();
 
-        accessor.fetch_block_with_missing_unique_per_doc(&docs, &source, None, false);
+        accessor.fetch_block_with_missing_unique_per_doc(&docs, &mut source, None, false);
 
         assert!(!accessor.has_one_value_per_doc(&docs));
         assert_eq!(
@@ -482,14 +482,14 @@ mod tests {
         let vals = [10u64, 40, 70];
         let values =
             serialize_and_load_u64_based_column_values::<u64>(&&vals[..], &ALL_U64_CODEC_TYPES);
-        let column = Column {
+        let mut column = Column {
             index: ColumnIndex::Optional(OptionalIndex::for_test(9, &[1, 4, 7])),
             values,
         };
         let docs = [0, 1, 2, 4, 7, 8];
         let mut accessor = ColumnBlockAccessor::default();
 
-        accessor.fetch_block_with_missing_ordered(&docs, &column, Some(99), true);
+        accessor.fetch_block_with_missing_ordered(&docs, &mut column, Some(99), true);
 
         assert_eq!(
             accessor.iter_vals().collect::<Vec<_>>(),
@@ -573,13 +573,13 @@ mod tests {
         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 {
+        let mut column = Column {
             index: ColumnIndex::Full,
             values,
         };
 
-        let check = |accessor: &mut ColumnBlockAccessor, docs: &[u32]| {
-            accessor.fetch_block(docs, &column);
+        let check = |accessor: &mut ColumnBlockAccessor, docs: &[u32], column: &mut Column<u64>| {
+            accessor.fetch_block(docs, column);
             let got: Vec<(u32, u64)> = accessor.iter_docid_vals(docs).collect();
             let expected: Vec<(u32, u64)> = docs.iter().map(|&d| (d, vals[d as usize])).collect();
             assert_eq!(got, expected);
@@ -587,11 +587,11 @@ mod tests {
 
         let mut accessor = ColumnBlockAccessor::default();
         // Contiguous block -> get_range fast path.
-        check(&mut accessor, &(10..74).collect::<Vec<u32>>());
+        check(&mut accessor, &(10..74).collect::<Vec<u32>>(), &mut column);
         // Non-contiguous block -> get_vals gather path.
-        check(&mut accessor, &[0, 5, 9, 100, 199]);
+        check(&mut accessor, &[0, 5, 9, 100, 199], &mut column);
         // Single doc and full span.
-        check(&mut accessor, &[42]);
-        check(&mut accessor, &(0..200).collect::<Vec<u32>>());
+        check(&mut accessor, &[42], &mut column);
+        check(&mut accessor, &(0..200).collect::<Vec<u32>>(), &mut column);
     }
 }

src/aggregation/bucket/histogram/histogram.rs

@@ -483,7 +483,7 @@ impl<B: BucketIdSlot> SegmentAggregationCollector for SegmentHistogramCollector<
         // 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 req = &mut self.req_data;
         let bounds = req.bounds;
         let interval = req.req.interval;
         let offset = req.offset;
@@ -491,7 +491,7 @@ impl<B: BucketIdSlot> SegmentAggregationCollector for SegmentHistogramCollector<
 
         agg_data
             .column_block_accessor
-            .fetch_block(docs, &req.accessor);
+            .fetch_block(docs, &mut req.accessor);
         // 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) {

src/aggregation/bucket/multi_terms/mod.rs

@@ -223,14 +223,14 @@ fn block_missing_value(missing: Option<&MultiTermsMissingAccessor>) -> Option<u6
 #[inline]
 fn fetch_field_block(
     docs: &[crate::DocId],
-    field: &MultiTermsFieldAccessor,
+    field: &mut MultiTermsFieldAccessor,
     missing: Option<&MultiTermsMissingAccessor>,
     block_accessor: &mut ColumnBlockAccessor,
 ) -> bool {
     let missing_value = block_missing_value(missing);
     block_accessor.fetch_block_with_missing_unique_per_doc(
         docs,
-        &field.column,
+        &mut field.column,
         missing_value,
         true,
     );
@@ -571,7 +571,7 @@ where
         docs: &[DocId],
         block_accessor: &mut ColumnBlockAccessor,
     ) {
-        for (field_idx, field) in self.req_data.fields.iter().enumerate() {
+        for (field_idx, field) in self.req_data.fields.iter_mut().enumerate() {
             fetch_field_block(
                 docs,
                 field,
@@ -595,7 +595,7 @@ where
         self.alive_docs.extend_from_slice(docs);
         self.doc_ids_per_partial_combination.clear();
 
-        for (field_idx, field) in self.req_data.fields.iter().enumerate() {
+        for (field_idx, field) in self.req_data.fields.iter_mut().enumerate() {
             let missing = self.req_data.missing_accessors[field_idx].as_ref();
             let has_one_value_per_doc =
                 fetch_field_block(&self.alive_docs, field, missing, block_accessor);

src/aggregation/bucket/range.rs

@@ -281,7 +281,7 @@ impl<B: SubAggBuffer> SegmentAggregationCollector for SegmentRangeCollector<B> {
     ) -> crate::Result<()> {
         agg_data
             .column_block_accessor
-            .fetch_block(docs, &self.req_data.accessor);
+            .fetch_block(docs, &mut self.req_data.accessor);
 
         let buckets = &mut self.parent_buckets[parent_bucket_id as usize];
 

src/aggregation/bucket/term_agg/mod.rs

@@ -1064,7 +1064,7 @@ impl<TermMap: TermAggregationMap, B: SubAggBuffer> SegmentAggregationCollector
             .column_block_accessor
             .fetch_block_with_missing_unique_per_doc(
                 docs,
-                &req_data.accessor,
+                &mut req_data.accessor,
                 req_data.missing_value_for_accessor,
                 false,
             );

src/aggregation/metric/cardinality.rs

@@ -638,7 +638,7 @@ impl<S: TermOrdAccumulator> SegmentCardinalityCollector<S> {
     ) {
         agg_data.column_block_accessor.fetch_block_with_missing(
             docs,
-            &self.accessor,
+            &mut self.accessor,
             self.missing_value_for_accessor,
         );
     }

src/aggregation/metric/extended_stats.rs

@@ -373,9 +373,11 @@ impl SegmentAggregationCollector for SegmentExtendedStatsCollector {
     ) -> crate::Result<()> {
         let mut extended_stats = self.buckets[parent_bucket_id as usize].clone();
 
-        agg_data
-            .column_block_accessor
-            .fetch_block_with_missing(docs, &self.accessor, self.missing);
+        agg_data.column_block_accessor.fetch_block_with_missing(
+            docs,
+            &mut self.accessor,
+            self.missing,
+        );
         for val in agg_data.column_block_accessor.iter_vals() {
             let val1 = f64_from_fastfield_u64(val, self.field_type);
             extended_stats.collect(val1);

src/aggregation/metric/percentiles.rs

@@ -299,7 +299,7 @@ impl SegmentAggregationCollector for SegmentPercentilesCollector {
         let percentiles = &mut self.buckets[parent_bucket_id as usize];
         agg_data.column_block_accessor.fetch_block_with_missing(
             docs,
-            &self.accessor,
+            &mut self.accessor,
             self.missing_u64,
         );
 

src/aggregation/metric/stats.rs

@@ -302,7 +302,7 @@ impl<const COLUMN_TYPE_ID: u8> SegmentAggregationCollector
         }
         agg_data.column_block_accessor.fetch_block_with_missing(
             docs,
-            &self.accessor,
+            &mut self.accessor,
             self.missing_u64,
         );
         collect_stats::<COLUMN_TYPE_ID>(