diff --git a/src/collector/sort_key/order.rs b/src/collector/sort_key/order.rs index 923d5cb8e..5e7b80669 100644 --- a/src/collector/sort_key/order.rs +++ b/src/collector/sort_key/order.rs @@ -322,11 +322,12 @@ impl SegmentSortKeyComput for SegmentSortKeyComputerWithComparator where TSegmentSortKeyComputer: SegmentSortKeyComputer, - TSegmentSortKey: PartialOrd + Clone + 'static + Sync + Send, + TSegmentSortKey: Clone + 'static + Sync + Send, TComparator: Comparator + 'static + Sync + Send, { type SortKey = TSegmentSortKeyComputer::SortKey; type SegmentSortKey = TSegmentSortKey; + type SegmentComparator = TComparator; fn segment_sort_key(&mut self, doc: DocId, score: Score) -> Self::SegmentSortKey { self.segment_sort_key_computer.segment_sort_key(doc, score) diff --git a/src/collector/sort_key/sort_by_score.rs b/src/collector/sort_key/sort_by_score.rs index df8b0dd75..a23660e56 100644 --- a/src/collector/sort_key/sort_by_score.rs +++ b/src/collector/sort_key/sort_by_score.rs @@ -63,8 +63,8 @@ impl SortKeyComputer for SortBySimilarityScore { impl SegmentSortKeyComputer for SortBySimilarityScore { type SortKey = Score; - type SegmentSortKey = Score; + type SegmentComparator = NaturalComparator; #[inline(always)] fn segment_sort_key(&mut self, _doc: DocId, score: Score) -> Score { diff --git a/src/collector/sort_key/sort_by_static_fast_value.rs b/src/collector/sort_key/sort_by_static_fast_value.rs index b38b8b034..44a4e1d8d 100644 --- a/src/collector/sort_key/sort_by_static_fast_value.rs +++ b/src/collector/sort_key/sort_by_static_fast_value.rs @@ -34,9 +34,7 @@ impl SortByStaticFastValue { impl SortKeyComputer for SortByStaticFastValue { type Child = SortByFastValueSegmentSortKeyComputer; - type SortKey = Option; - type Comparator = NaturalComparator; fn check_schema(&self, schema: &crate::schema::Schema) -> crate::Result<()> { @@ -84,8 +82,8 @@ pub struct SortByFastValueSegmentSortKeyComputer { impl SegmentSortKeyComputer for SortByFastValueSegmentSortKeyComputer { type SortKey = Option; - type SegmentSortKey = Option; + type SegmentComparator = NaturalComparator; #[inline(always)] fn segment_sort_key(&mut self, doc: DocId, _score: Score) -> Self::SegmentSortKey { diff --git a/src/collector/sort_key/sort_by_string.rs b/src/collector/sort_key/sort_by_string.rs index 41ef22e9b..c7f449619 100644 --- a/src/collector/sort_key/sort_by_string.rs +++ b/src/collector/sort_key/sort_by_string.rs @@ -30,9 +30,7 @@ impl SortByString { impl SortKeyComputer for SortByString { type SortKey = Option; - type Child = ByStringColumnSegmentSortKeyComputer; - type Comparator = NaturalComparator; fn segment_sort_key_computer( @@ -50,8 +48,8 @@ pub struct ByStringColumnSegmentSortKeyComputer { impl SegmentSortKeyComputer for ByStringColumnSegmentSortKeyComputer { type SortKey = Option; - type SegmentSortKey = Option; + type SegmentComparator = NaturalComparator; #[inline(always)] fn segment_sort_key(&mut self, doc: DocId, _score: Score) -> Option { diff --git a/src/collector/sort_key/sort_key_computer.rs b/src/collector/sort_key/sort_key_computer.rs index d56fa7cd0..ccfe44e05 100644 --- a/src/collector/sort_key/sort_key_computer.rs +++ b/src/collector/sort_key/sort_key_computer.rs @@ -12,13 +12,21 @@ use crate::{DocAddress, DocId, Result, Score, SegmentReader}; /// It is the segment local version of the [`SortKeyComputer`]. pub trait SegmentSortKeyComputer: 'static { /// The final score being emitted. - type SortKey: 'static + PartialOrd + Send + Sync + Clone; + type SortKey: 'static + Send + Sync + Clone; /// Sort key used by at the segment level by the `SegmentSortKeyComputer`. /// /// It is typically small like a `u64`, and is meant to be converted /// to the final score at the end of the collection of the segment. - type SegmentSortKey: 'static + PartialOrd + Clone + Send + Sync + Clone; + type SegmentSortKey: 'static + Clone + Send + Sync + Clone; + + /// Comparator type. + type SegmentComparator: Comparator + 'static; + + /// Returns the segment sort key comparator. + fn segment_comparator(&self) -> Self::SegmentComparator { + Self::SegmentComparator::default() + } /// Computes the sort key for the given document and score. fn segment_sort_key(&mut self, doc: DocId, score: Score) -> Self::SegmentSortKey; @@ -47,7 +55,7 @@ pub trait SegmentSortKeyComputer: 'static { left: &Self::SegmentSortKey, right: &Self::SegmentSortKey, ) -> Ordering { - NaturalComparator.compare(left, right) + self.segment_comparator().compare(left, right) } /// Implementing this method makes it possible to avoid computing @@ -81,7 +89,7 @@ pub trait SegmentSortKeyComputer: 'static { /// the sort key at a segment scale. pub trait SortKeyComputer: Sync { /// The sort key type. - type SortKey: 'static + Send + Sync + PartialOrd + Clone + std::fmt::Debug; + type SortKey: 'static + Send + Sync + Clone + std::fmt::Debug; /// Type of the associated [`SegmentSortKeyComputer`]. type Child: SegmentSortKeyComputer; /// Comparator type. @@ -188,6 +196,11 @@ where TailSegmentSortKeyComputer::SegmentSortKey, ); + type SegmentComparator = ( + HeadSegmentSortKeyComputer::SegmentComparator, + TailSegmentSortKeyComputer::SegmentComparator, + ); + /// A SegmentSortKeyComputer maps to a SegmentSortKey, but it can also decide on /// its ordering. /// @@ -269,11 +282,12 @@ impl SegmentSortKeyComputer for MappedSegmentSortKeyComputer where T: SegmentSortKeyComputer, - PreviousScore: 'static + Clone + Send + Sync + PartialOrd, - NewScore: 'static + Clone + Send + Sync + PartialOrd, + PreviousScore: 'static + Clone + Send + Sync, + NewScore: 'static + Clone + Send + Sync, { type SortKey = NewScore; type SegmentSortKey = T::SegmentSortKey; + type SegmentComparator = T::SegmentComparator; fn segment_sort_key(&mut self, doc: DocId, score: Score) -> Self::SegmentSortKey { self.sort_key_computer.segment_sort_key(doc, score) @@ -463,6 +477,7 @@ where { type SortKey = TSortKey; type SegmentSortKey = TSortKey; + type SegmentComparator = NaturalComparator; fn segment_sort_key(&mut self, doc: DocId, _score: Score) -> TSortKey { (self)(doc) diff --git a/src/collector/top_score_collector.rs b/src/collector/top_score_collector.rs index 78c344dbe..631169bc0 100644 --- a/src/collector/top_score_collector.rs +++ b/src/collector/top_score_collector.rs @@ -325,7 +325,7 @@ impl TopDocs { sort_key_computer: impl SortKeyComputer + Send + 'static, ) -> impl Collector> where - TSortKey: 'static + Clone + Send + Sync + PartialOrd + std::fmt::Debug, + TSortKey: 'static + Clone + Send + Sync + std::fmt::Debug, { TopBySortKeyCollector::new(sort_key_computer, self.doc_range()) } @@ -446,7 +446,7 @@ where F: 'static + Send + Sync + Fn(&SegmentReader) -> TTweakScoreSortKeyFn, TTweakScoreSortKeyFn: 'static + Fn(DocId, Score) -> TSortKey, TweakScoreSegmentSortKeyComputer: - SegmentSortKeyComputer, + SegmentSortKeyComputer, TSortKey: 'static + PartialOrd + Clone + Send + Sync + std::fmt::Debug, { type SortKey = TSortKey; @@ -481,6 +481,7 @@ where { type SortKey = TSortKey; type SegmentSortKey = TSortKey; + type SegmentComparator = NaturalComparator; fn segment_sort_key(&mut self, doc: DocId, score: Score) -> TSortKey { (self.sort_key_fn)(doc, score)