@@ -23,7 +23,7 @@ use crate::{
2323} ;
2424
2525/// A built-in helper for benchmarking the K-nearest neighbors method
26- /// [`graph::DiskANNIndex::search`].
26+ /// [`graph::DiskANNIndex::search`] with optional post-processing support .
2727///
2828/// This is intended to be used in conjunction with [`search::search`] or
2929/// [`search::search_all`] and provides some basic additional metrics for
@@ -32,21 +32,31 @@ use crate::{
3232///
3333/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
3434/// and returns [`Metrics`] as additional output.
35+ ///
36+ /// # Type Parameters
37+ ///
38+ /// - `DP`: The data provider type
39+ /// - `T`: The query element type
40+ /// - `S`: The search strategy type
41+ /// - `PP`: Post-processor selector. Defaults to [`Defaulted`], which uses the
42+ /// strategy's default post-processor. Use [`KNN::with_postprocessor`] to
43+ /// supply an explicit post-processor.
3544#[ derive( Debug ) ]
36- pub struct KNN < DP , T , S >
45+ pub struct KNN < DP , T , S , PP = Defaulted >
3746where
3847 DP : provider:: DataProvider ,
3948{
4049 index : Arc < graph:: DiskANNIndex < DP > > ,
4150 queries : Arc < Matrix < T > > ,
4251 strategy : Strategy < S > ,
52+ post_processor : PP ,
4353}
4454
45- impl < DP , T , S > KNN < DP , T , S >
55+ impl < DP , T , S > KNN < DP , T , S , Defaulted >
4656where
4757 DP : provider:: DataProvider ,
4858{
49- /// Construct a new [`KNN`] searcher.
59+ /// Construct a new [`KNN`] searcher using the strategy's default post-processor .
5060///
5161/// If `strategy` is one of the container variants of [`Strategy`], its length
5262/// must match the number of rows in `queries`. If this is the case, then the
@@ -68,10 +78,98 @@ where
6878 index,
6979 queries,
7080 strategy,
81+ post_processor : Defaulted ,
82+ } ) )
83+ }
84+ }
85+
86+ impl < DP , T , S , PP > KNN < DP , T , S , Forwarded < PP > >
87+ where
88+ DP : provider:: DataProvider ,
89+ {
90+ /// Construct a new [`KNN`] searcher with an explicit post-processor.
91+ ///
92+ /// # Errors
93+ ///
94+ /// Returns an error if the number of elements in `strategy` is not compatible with
95+ /// the number of rows in `queries`.
96+ pub fn with_postprocessor (
97+ index : Arc < graph:: DiskANNIndex < DP > > ,
98+ queries : Arc < Matrix < T > > ,
99+ strategy : Strategy < S > ,
100+ post_processor : PP ,
101+ ) -> anyhow:: Result < Arc < Self > > {
102+ strategy. length_compatible ( queries. nrows ( ) ) ?;
103+
104+ Ok ( Arc :: new ( Self {
105+ index,
106+ queries,
107+ strategy,
108+ post_processor : Forwarded ( post_processor) ,
71109 } ) )
72110 }
73111}
74112
113+ impl < DP , T , S , PP > KNN < DP , T , S , PP >
114+ where
115+ DP : provider:: DataProvider ,
116+ {
117+ /// Access the index.
118+ pub fn index ( & self ) -> & Arc < graph:: DiskANNIndex < DP > > {
119+ & self . index
120+ }
121+ }
122+
123+ /// Resolves a post-processor for [`KNN`] given a search strategy.
124+ ///
125+ /// This trait lets [`KNN`] support both "use the strategy's default post-processor"
126+ /// ([`Defaulted`]) and "use this explicit post-processor" ([`Forwarded`]) without
127+ /// duplicating the search loop.
128+ pub trait AsPostProcessor < ' a , S , DP , T >
129+ where
130+ DP : provider:: DataProvider ,
131+ S : glue:: SearchStrategy < ' a , DP , T > ,
132+ {
133+ /// The concrete post-processor used for a single search.
134+ type Processor : glue:: SearchPostProcess < S :: SearchAccessor , T , DP :: ExternalId > + Send + Sync ;
135+
136+ /// Construct the post-processor to use for a single search.
137+ fn as_post_processor ( & ' a self , strategy : & ' a S ) -> Self :: Processor ;
138+ }
139+
140+ /// Marker indicating that [`KNN`] should use the strategy's default post-processor.
141+ #[ derive( Debug , Clone , Copy ) ]
142+ pub struct Defaulted ;
143+
144+ impl < ' a , S , DP , T > AsPostProcessor < ' a , S , DP , T > for Defaulted
145+ where
146+ DP : provider:: DataProvider ,
147+ S : glue:: DefaultPostProcessor < ' a , DP , T , DP :: ExternalId > ,
148+ {
149+ type Processor = S :: Processor ;
150+
151+ fn as_post_processor ( & ' a self , strategy : & ' a S ) -> Self :: Processor {
152+ strategy. default_post_processor ( )
153+ }
154+ }
155+
156+ /// Wraps an explicit post-processor for use with [`KNN::with_postprocessor`].
157+ #[ derive( Debug , Clone , Copy ) ]
158+ pub struct Forwarded < PP > ( PP ) ;
159+
160+ impl < ' a , S , DP , T , PP > AsPostProcessor < ' a , S , DP , T > for Forwarded < PP >
161+ where
162+ DP : provider:: DataProvider ,
163+ S : glue:: SearchStrategy < ' a , DP , T > ,
164+ PP : glue:: SearchPostProcess < S :: SearchAccessor , T , DP :: ExternalId > + Clone + AsyncFriendly ,
165+ {
166+ type Processor = PP ;
167+
168+ fn as_post_processor ( & ' a self , _strategy : & ' a S ) -> Self :: Processor {
169+ self . 0 . clone ( )
170+ }
171+ }
172+
75173/// Additional metrics collected during [`KNN`] search.
76174///
77175/// # Note
@@ -86,10 +184,13 @@ pub struct Metrics {
86184pub hops : u32 ,
87185}
88186
89- impl < DP , T , S > Search for KNN < DP , T , S >
187+ impl < DP , T , S , PP > Search for KNN < DP , T , S , PP >
90188where
91189 DP : provider:: DataProvider < Context : Default , ExternalId : search:: Id > ,
92- S : for < ' a > glue:: DefaultSearchStrategy < ' a , DP , & ' a [ T ] , DP :: ExternalId > + Clone + AsyncFriendly ,
190+ S : for < ' a > glue:: SearchStrategy < ' a , DP , & ' a [ T ] > + Clone + AsyncFriendly ,
191+ PP : for < ' a > AsPostProcessor < ' a , S , DP , & ' a [ T ] > + AsyncFriendly ,
192+ graph:: search:: Knn :
193+ for < ' a > graph:: Search < ' a , DP , S , & ' a [ T ] , Output = graph:: index:: SearchStats > ,
93194 T : AsyncFriendly + Clone ,
94195{
95196 type Id = DP :: ExternalId ;
@@ -115,11 +216,15 @@ where
115216 {
116217 let context = DP :: Context :: default ( ) ;
117218 let knn_search = * parameters;
219+ let strategy = self . strategy . get ( index) ?;
220+ let processor = self . post_processor . as_post_processor ( strategy) ;
221+
118222 let stats = self
119223 . index
120- . search (
224+ . search_with (
121225 knn_search,
122- self . strategy . get ( index) ?,
226+ strategy,
227+ processor,
123228 & context,
124229 self . queries . row ( index) ,
125230 buffer,
0 commit comments