From eb2c2815b63c049d5e2aeabe6b67283bb2759ca6 Mon Sep 17 00:00:00 2001
From: Mubelotix <simon@meilisearch.com>
Date: Tue, 1 Jul 2025 10:00:10 +0200
Subject: [PATCH] Fix panic

---
 .../milli/src/facet/facet_sort_recursive.rs   | 65 +++++++++++++------
 1 file changed, 44 insertions(+), 21 deletions(-)

diff --git a/crates/milli/src/facet/facet_sort_recursive.rs b/crates/milli/src/facet/facet_sort_recursive.rs
index 6f26ad16f..213d18624 100644
--- a/crates/milli/src/facet/facet_sort_recursive.rs
+++ b/crates/milli/src/facet/facet_sort_recursive.rs
@@ -128,37 +128,60 @@ impl<'ctx> SortedDocumentsIteratorBuilder<'ctx> {
         let mut cache = VecDeque::new();
         let mut rtree = None;
         let size = candidates.len() as usize;
+        let not_geo_candidates = candidates.clone() - geo_candidates;
+        let mut geo_remaining = size - not_geo_candidates.len() as usize;
+        let mut not_geo_candidates = Some(not_geo_candidates);
 
         let next_children = std::iter::from_fn(move || {
-            match next_bucket(
-                index,
-                rtxn,
-                &candidates,
-                ascending,
-                target_point,
-                &Some(field_ids),
-                &mut rtree,
-                &mut cache,
-                geo_candidates,
-                GeoSortParameter::default(),
-            ) {
-                Ok(Some((docids, _point))) => Some(Ok(SortedDocumentsIteratorBuilder {
+            // Find the next bucket of geo-sorted documents.
+            // next_bucket loops and will go back to the beginning so we use a variable to track how many are left.
+            if geo_remaining > 0 {
+                if let Ok(Some((docids, _point))) = next_bucket(
                     index,
                     rtxn,
-                    number_db,
-                    string_db,
-                    fields: &fields[1..],
-                    candidates: docids,
+                    &candidates,
+                    ascending,
+                    target_point,
+                    &Some(field_ids),
+                    &mut rtree,
+                    &mut cache,
                     geo_candidates,
-                })),
-                Ok(None) => None,
-                Err(e) => Some(Err(e)),
+                    GeoSortParameter::default(),
+                ) {
+                    geo_remaining -= docids.len() as usize;
+                    return Some(Ok(SortedDocumentsIteratorBuilder {
+                        index,
+                        rtxn,
+                        number_db,
+                        string_db,
+                        fields: &fields[1..],
+                        candidates: docids,
+                        geo_candidates,
+                    }));
+                }
             }
+            
+            // Once all geo candidates have been processed, we can return the others
+            if let Some(not_geo_candidates) = not_geo_candidates.take() {
+                if !not_geo_candidates.is_empty() {
+                    return Some(Ok(SortedDocumentsIteratorBuilder {
+                        index,
+                        rtxn,
+                        number_db,
+                        string_db,
+                        fields: &fields[1..],
+                        candidates: not_geo_candidates,
+                        geo_candidates,
+                    }));
+                }
+            }
+
+            None
         });
 
         Ok(SortedDocumentsIterator::Branch {
             current_child: None,
-            next_children_size: size, // TODO: confirm all candidates will be included
+            next_children_size: size,
             next_children: Box::new(next_children),
         })
     }