Optimize detect when primary condition is sufficient review (#6794)

* Optimize payload filtering to avoid redundant matching

* No boxing

* review suggestions

---------

Co-authored-by: Arnaud Gourlay <arnaud.gourlay@gmail.com>

lib/common/common/src/either_variant.rs

@@ -0,0 +1,114 @@
+use std::iter;
+
+/// Variant works similar as Either, but for 4 variants.
+pub enum EitherVariant<A, B, C, D> {
+    A(A),
+    B(B),
+    C(C),
+    D(D),
+}
+
+macro_rules! for_all {
+    ($value:expr, $pattern:pat => $result:expr) => {
+        match $value {
+            $crate::either_variant::EitherVariant::A($pattern) => $result,
+            $crate::either_variant::EitherVariant::B($pattern) => $result,
+            $crate::either_variant::EitherVariant::C($pattern) => $result,
+            $crate::either_variant::EitherVariant::D($pattern) => $result,
+        }
+    };
+}
+
+impl<A, B, C, D> Iterator for EitherVariant<A, B, C, D>
+where
+    A: Iterator,
+    B: Iterator<Item = A::Item>,
+    C: Iterator<Item = A::Item>,
+    D: Iterator<Item = A::Item>,
+{
+    type Item = A::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        for_all!(*self, ref mut inner => inner.next())
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        for_all!(*self, ref inner => inner.size_hint())
+    }
+
+    fn count(self) -> usize {
+        for_all!(self, inner => inner.count())
+    }
+
+    fn last(self) -> Option<Self::Item> {
+        for_all!(self, inner => inner.last())
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        for_all!(*self, ref mut inner => inner.nth(n))
+    }
+
+    fn for_each<F>(self, f: F)
+    where
+        F: FnMut(Self::Item),
+    {
+        for_all!(self, inner => inner.for_each(f))
+    }
+
+    fn collect<X>(self) -> X
+    where
+        X: iter::FromIterator<Self::Item>,
+    {
+        for_all!(self, inner => inner.collect())
+    }
+
+    fn partition<X, F>(self, f: F) -> (X, X)
+    where
+        X: Default + Extend<Self::Item>,
+        F: FnMut(&Self::Item) -> bool,
+    {
+        for_all!(self, inner => inner.partition(f))
+    }
+
+    fn fold<Acc, G>(self, init: Acc, f: G) -> Acc
+    where
+        G: FnMut(Acc, Self::Item) -> Acc,
+    {
+        for_all!(self, inner => inner.fold(init, f))
+    }
+
+    fn all<F>(&mut self, f: F) -> bool
+    where
+        F: FnMut(Self::Item) -> bool,
+    {
+        for_all!(*self, ref mut inner => inner.all(f))
+    }
+
+    fn any<F>(&mut self, f: F) -> bool
+    where
+        F: FnMut(Self::Item) -> bool,
+    {
+        for_all!(*self, ref mut inner => inner.any(f))
+    }
+
+    fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
+    where
+        P: FnMut(&Self::Item) -> bool,
+    {
+        for_all!(*self, ref mut inner => inner.find(predicate))
+    }
+
+    fn find_map<X, F>(&mut self, f: F) -> Option<X>
+    where
+        F: FnMut(Self::Item) -> Option<X>,
+    {
+        for_all!(*self, ref mut inner => inner.find_map(f))
+    }
+
+    fn position<P>(&mut self, predicate: P) -> Option<usize>
+    where
+        P: FnMut(Self::Item) -> bool,
+    {
+        for_all!(*self, ref mut inner => inner.position(predicate))
+    }
+}

lib/common/common/src/lib.rs

@@ -8,6 +8,7 @@ pub mod cpu;
 pub mod defaults;
 pub mod delta_pack;
 pub mod disk;
+pub mod either_variant;
 pub mod ext;
 pub mod fixed_length_priority_queue;
 pub mod flags;

lib/segment/src/index/struct_payload_index.rs

@@ -7,8 +7,8 @@ use ahash::AHashSet;
 use atomic_refcell::AtomicRefCell;
 use common::counter::hardware_counter::HardwareCounterCell;
 use common::counter::iterator_hw_measurement::HwMeasurementIteratorExt;
+use common::either_variant::EitherVariant;
 use common::types::PointOffsetType;
-use itertools::Either;
 use log::debug;
 use schemars::_serde_json::Value;
 
@@ -423,38 +423,61 @@ impl StructPayloadIndex {
         query_cardinality: &'a CardinalityEstimation,
         hw_counter: &'a HardwareCounterCell,
     ) -> impl Iterator<Item = PointOffsetType> + 'a {
-        let struct_filtered_context = self.struct_filtered_context(filter, hw_counter);
-
         if query_cardinality.primary_clauses.is_empty() {
             let full_scan_iterator = id_tracker.iter_ids();
-
+            let struct_filtered_context = self.struct_filtered_context(filter, hw_counter);
             // Worst case: query expected to return few matches, but index can't be used
             let matched_points =
                 full_scan_iterator.filter(move |i| struct_filtered_context.check(*i));
 
-            Either::Left(matched_points)
+            EitherVariant::A(matched_points)
         } else {
             // CPU-optimized strategy here: points are made unique before applying other filters.
             let mut visited_list = self.visited_pool.get(id_tracker.total_point_count());
 
-            let iter = query_cardinality
+            // If even one iterator is None, we should replace the whole thing with
+            // an iterator over all ids.
+            let primary_clause_iterators: Option<Vec<_>> = query_cardinality
                 .primary_clauses
                 .iter()
-                .flat_map(move |clause| {
-                    self.query_field(clause, hw_counter).unwrap_or_else(|| {
-                        // index is not built
-                        Box::new(id_tracker.iter_ids().measure_hw_with_cell(
-                            hw_counter,
-                            size_of::<PointOffsetType>(),
-                            |i| i.cpu_counter(),
-                        ))
-                    })
+                .map(move |clause| self.query_field(clause, hw_counter))
+                .collect();
+
+            if let Some(primary_iterators) = primary_clause_iterators {
+                let num_primary_iterators = primary_iterators.len();
+                let joined_primary_iterator = primary_iterators.into_iter().flatten();
+
+                return if num_primary_iterators == filter.total_conditions_count() {
+                    // All conditions are primary clauses,
+                    // We can avoid post-filtering
+                    let iter = joined_primary_iterator
+                        .filter(move |&id| !visited_list.check_and_update_visited(id));
+                    EitherVariant::B(iter)
+                } else {
+                    // Some conditions are primary clauses, some are not
+                    let struct_filtered_context = self.struct_filtered_context(filter, hw_counter);
+                    let iter = joined_primary_iterator.filter(move |&id| {
+                        !visited_list.check_and_update_visited(id)
+                            && struct_filtered_context.check(id)
+                    });
+                    EitherVariant::C(iter)
+                };
+            }
+
+            // We can't use primary conditions, so we fall back to iterating over all ids
+            // and applying full filter.
+            let struct_filtered_context = self.struct_filtered_context(filter, hw_counter);
+
+            let iter = id_tracker
+                .iter_ids()
+                .measure_hw_with_cell(hw_counter, size_of::<PointOffsetType>(), |i| {
+                    i.cpu_counter()
                 })
                 .filter(move |&id| {
                     !visited_list.check_and_update_visited(id) && struct_filtered_context.check(id)
                 });
 
-            Either::Right(iter)
+            EitherVariant::D(iter)
         }
     }