1#[cfg(feature = "native")]
16use rayon::prelude::*;
17const TERM_DEGREE_VALUE_BYTES: usize = std::mem::size_of::<[u32; 2]>();
18const CANDIDATE_ENTRY_BYTES: usize = std::mem::size_of::<(usize, u32)>();
19const PARALLEL_BP_MIN_ENTITIES: usize = 1_048_576;
23const MIN_RELATIVE_OBJECTIVE_IMPROVEMENT: f64 = 1e-6;
24const MIN_OBJECTIVE_ITERATIONS: usize = 4;
25const OBJECTIVE_STALL_ITERATIONS: usize = 2;
26const LOG_TABLE_SIZE: usize = 4096;
27
28fn term_degree_bytes(num_terms: usize) -> usize {
29 let bitmap_words = num_terms.div_ceil(64);
30 num_terms
31 .saturating_mul(TERM_DEGREE_VALUE_BYTES)
32 .saturating_add(bitmap_words.saturating_mul(std::mem::size_of::<u64>()))
33 .saturating_add(bitmap_words.saturating_mul(std::mem::size_of::<u32>()))
37}
38
39#[cfg(feature = "native")]
44trait FrequencyParallelSafe: Sync {}
45#[cfg(feature = "native")]
46impl<T: Sync + ?Sized> FrequencyParallelSafe for T {}
47
48#[cfg(not(feature = "native"))]
49trait FrequencyParallelSafe {}
50#[cfg(not(feature = "native"))]
51impl<T: ?Sized> FrequencyParallelSafe for T {}
52
53fn count_frequencies_bounded<T: FrequencyParallelSafe>(
54 items: &[T],
55 num_terms: usize,
56 available_bytes: usize,
57 count_item: impl Fn(&T, &mut [u32]) + FrequencyParallelSafe,
58) -> Option<Vec<u32>> {
59 if num_terms == 0 {
60 return Some(Vec::new());
61 }
62 let table_bytes = num_terms
63 .checked_mul(std::mem::size_of::<u32>())?
64 .checked_add(std::mem::size_of::<Vec<u32>>())?;
65 let affordable_tables = available_bytes.checked_div(table_bytes)?;
66 if affordable_tables == 0 {
67 return None;
68 }
69
70 #[cfg(feature = "native")]
71 {
72 let lanes = affordable_tables
73 .min(rayon::current_num_threads().max(1))
74 .min(items.len().max(1));
75 if lanes > 1 {
76 let chunk_len = items.len().div_ceil(lanes);
77 return items
78 .par_chunks(chunk_len)
79 .map(|chunk| {
80 let mut counts = vec![0u32; num_terms];
81 for item in chunk {
82 count_item(item, &mut counts);
83 }
84 counts
85 })
86 .reduce_with(|mut left, right| {
87 for (total, count) in left.iter_mut().zip(right) {
88 *total = total.saturating_add(count);
89 }
90 left
91 });
92 }
93 }
94
95 let mut counts = vec![0u32; num_terms];
96 for item in items {
97 count_item(item, &mut counts);
98 }
99 Some(counts)
100}
101
102fn select_frequency_candidates(
105 frequencies: &[u32],
106 min_frequency: usize,
107 max_frequency: usize,
108 candidate_budget_bytes: usize,
109) -> (Vec<(u32, usize)>, bool) {
110 let eligible_count = frequencies
111 .iter()
112 .filter(|&&frequency| {
113 let frequency = frequency as usize;
114 frequency >= min_frequency && frequency <= max_frequency
115 })
116 .count();
117 let capacity = candidate_budget_bytes
118 .checked_div(CANDIDATE_ENTRY_BYTES)
119 .unwrap_or(0)
120 .min(eligible_count);
121 let mut candidates = std::collections::BinaryHeap::with_capacity(capacity);
122 for (term_id, &frequency) in frequencies.iter().enumerate() {
123 let frequency = frequency as usize;
124 if frequency < min_frequency || frequency > max_frequency {
125 continue;
126 }
127 let candidate = (frequency, term_id as u32);
128 if candidates.len() < capacity {
129 candidates.push(candidate);
130 } else if capacity > 0 && candidate < *candidates.peek().unwrap() {
131 candidates.pop();
132 candidates.push(candidate);
133 }
134 }
135 let selected = candidates
136 .into_vec()
137 .into_iter()
138 .map(|(frequency, term_id)| (term_id, frequency))
139 .collect();
140 (selected, capacity < eligible_count)
141}
142
143struct CandidateFit {
144 estimated_bytes: usize,
145 retained_postings: usize,
146 dropped: usize,
147}
148
149fn fit_candidates_to_budget(
153 candidates: &mut Vec<(u32, usize)>,
154 fixed_bytes: usize,
155 memory_budget_bytes: usize,
156) -> CandidateFit {
157 let total_postings = candidates.iter().fold(0usize, |total, (_, frequency)| {
158 total.saturating_add(*frequency)
159 });
160 let estimated_bytes = total_postings
161 .saturating_mul(std::mem::size_of::<u32>())
162 .saturating_add(fixed_bytes)
163 .saturating_add(candidates.len().saturating_mul(CANDIDATE_ENTRY_BYTES))
164 .saturating_add(term_degree_bytes(candidates.len()));
165 if estimated_bytes <= memory_budget_bytes || candidates.is_empty() {
166 return CandidateFit {
167 estimated_bytes,
168 retained_postings: total_postings,
169 dropped: 0,
170 };
171 }
172
173 candidates.sort_by_key(|&(_, frequency)| frequency);
174 let mut used_bytes = fixed_bytes;
175 let mut retained_postings = 0usize;
176 let mut keep = 0usize;
177 for &(_, frequency) in candidates.iter() {
178 let term_bytes = frequency
179 .saturating_mul(std::mem::size_of::<u32>())
180 .saturating_add(TERM_DEGREE_VALUE_BYTES + 1)
181 .saturating_add(CANDIDATE_ENTRY_BYTES);
182 if term_bytes > memory_budget_bytes.saturating_sub(used_bytes) {
183 break;
184 }
185 used_bytes = used_bytes.saturating_add(term_bytes);
186 retained_postings = retained_postings.saturating_add(frequency);
187 keep += 1;
188 }
189 let dropped = candidates.len() - keep;
190 candidates.truncate(keep);
191 candidates.shrink_to_fit();
194 CandidateFit {
195 estimated_bytes,
196 retained_postings,
197 dropped,
198 }
199}
200
201fn parallel_bisect_lanes(
202 memory_budget_bytes: usize,
203 non_degree_bytes: usize,
204 num_terms: usize,
205) -> usize {
206 let per_node = term_degree_bytes(num_terms).max(1);
207 let affordable_nodes = memory_budget_bytes
208 .saturating_sub(non_degree_bytes)
209 .checked_div(per_node)
210 .unwrap_or(0)
211 .max(1);
212 #[cfg(feature = "native")]
213 let worker_limit = rayon::current_num_threads().max(1);
214 #[cfg(not(feature = "native"))]
215 let worker_limit = 1usize;
216 affordable_nodes.min(worker_limit)
217}
218
219fn gain_cache_fits(budget: usize, non_degree_bytes: usize, terms: usize, lanes: usize) -> bool {
223 let required = terms
224 .checked_mul(std::mem::size_of::<[f32; 2]>())
225 .and_then(|cache| cache.checked_add(term_degree_bytes(terms)))
226 .and_then(|lane| lane.checked_add(std::mem::size_of::<TermDegrees>()))
227 .and_then(|lane| lane.checked_mul(lanes))
228 .and_then(|all_lanes| all_lanes.checked_add(non_degree_bytes))
229 .and_then(|graph| graph.checked_add(LOG_TABLE_SIZE * std::mem::size_of::<f32>()));
230 terms > 0 && required.is_some_and(|bytes| bytes <= budget)
231}
232
233struct TermDegrees {
241 values: Vec<std::mem::MaybeUninit<[u32; 2]>>,
242 initialized: Vec<u64>,
243 touched_words: Vec<u32>,
244 gain_cache: Vec<[f32; 2]>,
247}
248
249impl TermDegrees {
250 fn new(num_terms: usize) -> Self {
251 let bitmap_words = num_terms.div_ceil(64);
252 let mut values = Vec::with_capacity(num_terms);
253 values.resize_with(num_terms, std::mem::MaybeUninit::uninit);
254 Self {
255 values,
256 initialized: vec![0; bitmap_words],
257 touched_words: Vec::with_capacity(bitmap_words),
258 gain_cache: Vec::new(),
259 }
260 }
261
262 fn reset(&mut self) {
265 for word in self.touched_words.drain(..) {
266 self.initialized[word as usize] = 0;
267 }
268 }
269
270 fn sort_touched_words(&mut self) {
271 self.touched_words.sort_unstable();
272 }
273
274 #[inline]
275 fn entry_mut(&mut self, term: usize) -> &mut [u32; 2] {
276 let word = term / 64;
277 let mask = 1u64 << (term % 64);
278 if self.initialized[word] & mask == 0 {
279 if self.initialized[word] == 0 {
280 self.touched_words.push(word as u32);
281 }
282 self.values[term].write([0, 0]);
283 self.initialized[word] |= mask;
284 }
285 unsafe { self.values[term].assume_init_mut() }
287 }
288
289 #[inline]
290 fn get(&self, term: usize) -> [u32; 2] {
291 let word = term / 64;
292 let mask = 1u64 << (term % 64);
293 if self.initialized[word] & mask == 0 {
294 return [0, 0];
295 }
296 unsafe { *self.values[term].assume_init_ref() }
299 }
300
301 fn refresh_gain_cache(&mut self, log_table: &[f32]) {
302 if self.gain_cache.is_empty() {
303 return;
304 }
305 for &word_idx in &self.touched_words {
306 let word_idx = word_idx as usize;
307 let mut pending = self.initialized[word_idx];
308 while pending != 0 {
309 let bit = pending.trailing_zeros() as usize;
310 let term = word_idx * 64 + bit;
311 let [left, right] = unsafe { *self.values[term].assume_init_ref() };
313 self.gain_cache[term] = [
317 if left == 0 {
318 0.0
319 } else {
320 fast_log2_lookup(right as usize + 2, log_table)
321 - fast_log2_lookup(left as usize, log_table)
322 - std::f32::consts::LOG2_E / (1.0 + right as f32)
323 },
324 if right == 0 {
325 0.0
326 } else {
327 -(fast_log2_lookup(left as usize + 2, log_table)
328 - fast_log2_lookup(right as usize, log_table)
329 - std::f32::consts::LOG2_E / (1.0 + left as f32))
330 },
331 ];
332 pending &= pending - 1;
333 }
334 }
335 }
336
337 fn merge_from(&mut self, other: &Self) {
338 for &word_idx in &other.touched_words {
339 let word_idx = word_idx as usize;
340 let mut pending = other.initialized[word_idx];
341 while pending != 0 {
342 let bit = pending.trailing_zeros() as usize;
343 let term = word_idx * 64 + bit;
344 let [left, right] = unsafe { *other.values[term].assume_init_ref() };
346 let entry = self.entry_mut(term);
347 entry[0] += left;
348 entry[1] += right;
349 pending &= pending - 1;
350 }
351 }
352 }
353
354 fn apply_moves_to(&self, degrees: &mut Self) {
360 for &word_idx in &self.touched_words {
361 let word_idx = word_idx as usize;
362 let mut pending = self.initialized[word_idx];
363 while pending != 0 {
364 let bit = pending.trailing_zeros() as usize;
365 let term = word_idx * 64 + bit;
366 let [right_to_left, left_to_right] =
368 unsafe { *self.values[term].assume_init_ref() };
369 debug_assert!(
370 degrees.initialized[word_idx] & (1u64 << bit) != 0,
371 "a moved term must already exist in the partition degrees"
372 );
373 let degree = degrees.entry_mut(term);
374 let new_left =
375 i64::from(degree[0]) + i64::from(right_to_left) - i64::from(left_to_right);
376 let new_right =
377 i64::from(degree[1]) + i64::from(left_to_right) - i64::from(right_to_left);
378 debug_assert!(new_left >= 0 && new_right >= 0);
379 debug_assert!(new_left <= i64::from(u32::MAX));
380 debug_assert!(new_right <= i64::from(u32::MAX));
381 degree[0] = new_left as u32;
382 degree[1] = new_right as u32;
383 pending &= pending - 1;
384 }
385 }
386 }
387
388 fn bisection_objective(&self, left_size: usize, right_size: usize, log_table: &[f32]) -> f64 {
395 let mut objective = 0.0f64;
396 let side_log = [
397 fast_log2_lookup(left_size, log_table) as f64,
398 fast_log2_lookup(right_size, log_table) as f64,
399 ];
400 for &word_idx in &self.touched_words {
401 let word_idx = word_idx as usize;
402 let mut pending = self.initialized[word_idx];
403 while pending != 0 {
404 let bit = pending.trailing_zeros() as usize;
405 let term = word_idx * 64 + bit;
406 let [left, right] = unsafe { *self.values[term].assume_init_ref() };
408 for (side, count) in [left, right].into_iter().enumerate() {
409 if count > 0 {
410 objective += count as f64
411 * (fast_log2_lookup(count as usize + 1, log_table) as f64
412 - side_log[side]);
413 }
414 }
415 pending &= pending - 1;
416 }
417 }
418 objective
419 }
420}
421
422pub(crate) struct ForwardIndex {
428 terms: Vec<u32>,
429 offsets: Vec<u64>,
434 pub num_terms: usize,
435 parallel_bisect_lanes: usize,
439 cache_gains: bool,
440 budget_limited: bool,
444}
445
446fn build_csr_offsets(counts: &[u32]) -> Vec<u64> {
449 let mut offsets = Vec::with_capacity(counts.len() + 1);
450 offsets.push(0u64);
451 for &c in counts {
452 offsets.push(offsets.last().unwrap() + c as u64);
453 }
454 offsets
455}
456
457impl ForwardIndex {
458 pub(crate) fn from_csr(
462 terms: Vec<u32>,
463 offsets: Vec<u64>,
464 num_terms: usize,
465 memory_budget_bytes: usize,
466 budget_limited: bool,
467 ) -> Self {
468 let non_degree_bytes =
469 terms.len() * std::mem::size_of::<u32>() + offsets.len() * std::mem::size_of::<u64>();
470 let lanes = parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_terms);
471 Self {
472 terms,
473 offsets,
474 num_terms,
475 parallel_bisect_lanes: lanes.max(1),
476 cache_gains: false,
480 budget_limited,
481 }
482 }
483
484 #[inline]
485 pub fn num_docs(&self) -> usize {
486 if self.offsets.is_empty() {
487 0
488 } else {
489 self.offsets.len() - 1
490 }
491 }
492
493 #[inline]
494 fn doc_terms(&self, doc: usize) -> &[u32] {
495 let start = self.offsets[doc] as usize;
496 let end = self.offsets[doc + 1] as usize;
497 &self.terms[start..end]
498 }
499
500 pub fn total_postings(&self) -> u64 {
502 self.offsets.last().copied().unwrap_or(0)
503 }
504
505 #[inline]
506 pub fn budget_limited(&self) -> bool {
507 self.budget_limited
508 }
509}
510
511pub(crate) fn build_vid_maps(
521 bmp: &crate::segment::reader::bmp::BmpIndex,
522) -> crate::Result<(Vec<u32>, Vec<u32>)> {
523 let ids = bmp.doc_map_ids_slice();
524 let num_virtual = bmp.num_virtual_docs as usize;
525 let expected_real = bmp.num_real_docs() as usize;
526 let mut virtual_to_real = vec![u32::MAX; num_virtual];
527 let mut real_to_virtual = Vec::with_capacity(expected_real);
528 debug_assert_eq!(ids.len(), virtual_to_real.len() * 4);
529 bmp.visit_real_slots_for_rewrite(|vid| {
530 virtual_to_real[vid] = real_to_virtual.len() as u32;
531 real_to_virtual.push(vid as u32);
532 })?;
533 Ok((virtual_to_real, real_to_virtual))
534}
535
536struct BlockJob {
542 src: u32,
543 block_id: u32,
544 real_start: u32,
546 real_len: u32,
548}
549
550fn build_block_jobs(
553 bmps: &[&crate::segment::reader::bmp::BmpIndex],
554 vid_maps: &[(Vec<u32>, Vec<u32>)],
555) -> Vec<BlockJob> {
556 let total_blocks: usize = bmps.iter().map(|b| b.num_blocks as usize).sum();
557 let mut jobs = Vec::with_capacity(total_blocks);
558 for (src, (bmp, (v2r, _))) in bmps.iter().zip(vid_maps).enumerate() {
559 let block_size = bmp.bmp_block_size as usize;
560 let mut real_cursor = 0u32;
561 for block_id in 0..bmp.num_blocks as usize {
562 let vid_start = block_id * block_size;
563 let vid_end = ((block_id + 1) * block_size).min(v2r.len());
564 let real_len = v2r[vid_start..vid_end]
565 .iter()
566 .filter(|&&r| r != u32::MAX)
567 .count() as u32;
568 jobs.push(BlockJob {
569 src: src as u32,
570 block_id: block_id as u32,
571 real_start: real_cursor,
572 real_len,
573 });
574 real_cursor += real_len;
575 }
576 }
577 jobs
578}
579
580#[cfg(test)]
594pub(crate) fn build_forward_index_from_bmps(
595 bmps: &[&crate::segment::reader::bmp::BmpIndex],
596 min_doc_freq: usize,
597 max_doc_freq: usize,
598 memory_budget_bytes: usize,
599) -> crate::Result<ForwardIndex> {
600 let vid_maps: Vec<(Vec<u32>, Vec<u32>)> = bmps
601 .iter()
602 .map(|bmp| build_vid_maps(bmp))
603 .collect::<crate::Result<_>>()?;
604 Ok(build_forward_index_from_bmps_with_maps(
605 bmps,
606 &vid_maps,
607 min_doc_freq,
608 max_doc_freq,
609 memory_budget_bytes,
610 ))
611}
612
613pub(crate) fn build_forward_index_from_bmps_with_maps(
617 bmps: &[&crate::segment::reader::bmp::BmpIndex],
618 vid_maps: &[(Vec<u32>, Vec<u32>)],
619 min_doc_freq: usize,
620 max_doc_freq: usize,
621 memory_budget_bytes: usize,
622) -> ForwardIndex {
623 debug_assert_eq!(bmps.len(), vid_maps.len());
624 let total_docs: usize = vid_maps.iter().map(|(_, r2v)| r2v.len()).sum();
625
626 if total_docs == 0 {
627 return ForwardIndex {
628 terms: Vec::new(),
629 offsets: Vec::new(),
630 num_terms: 0,
631 parallel_bisect_lanes: 1,
632 cache_gains: false,
633 budget_limited: false,
634 };
635 }
636
637 let jobs = build_block_jobs(bmps, vid_maps);
642
643 let max_dims = bmps
647 .iter()
648 .map(|bmp| bmp.dims() as usize)
649 .max()
650 .unwrap_or(0);
651 let jobs_bytes = jobs
652 .len()
653 .saturating_mul(std::mem::size_of::<BlockJob>().saturating_add(40));
654 let frequency_bytes = max_dims.saturating_mul(std::mem::size_of::<u32>());
655 if frequency_bytes > memory_budget_bytes.saturating_sub(jobs_bytes) {
656 log::warn!(
657 "[reorder] memory budget {} cannot hold the {} dimension-frequency table; using identity order",
658 crate::format_bytes(memory_budget_bytes as u64),
659 crate::format_bytes(frequency_bytes as u64),
660 );
661 return ForwardIndex {
662 terms: Vec::new(),
663 offsets: Vec::new(),
664 num_terms: 0,
665 parallel_bisect_lanes: 1,
666 cache_gains: false,
667 budget_limited: true,
668 };
669 }
670 let Some(dim_df) = count_frequencies_bounded(
671 &jobs,
672 max_dims,
673 memory_budget_bytes.saturating_sub(jobs_bytes),
674 |job, counts| {
675 let bmp = bmps[job.src as usize];
676 let (v2r, _) = &vid_maps[job.src as usize];
677 let block_size = bmp.bmp_block_size as usize;
678 for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
679 let mut frequency = 0u32;
680 for posting in postings {
681 let vid = job.block_id as usize * block_size + posting.local_slot as usize;
682 if v2r.get(vid).is_some_and(|&real| real != u32::MAX) && posting.impact > 0 {
683 frequency = frequency.saturating_add(1);
684 }
685 }
686 if frequency > 0
687 && let Some(total) = counts.get_mut(dim_id as usize)
688 {
689 *total = total.saturating_add(frequency);
690 }
691 }
692 },
693 ) else {
694 log::warn!(
695 "[reorder] memory budget {} cannot hold a bounded dimension-frequency table; using identity order",
696 crate::format_bytes(memory_budget_bytes as u64),
697 );
698 return ForwardIndex {
699 terms: Vec::new(),
700 offsets: Vec::new(),
701 num_terms: 0,
702 parallel_bisect_lanes: 1,
703 cache_gains: false,
704 budget_limited: true,
705 };
706 };
707
708 let (mut eligible, mut budget_limited) = select_frequency_candidates(
712 &dim_df,
713 min_doc_freq,
714 max_doc_freq,
715 memory_budget_bytes
716 .saturating_sub(jobs_bytes)
717 .saturating_sub(frequency_bytes),
718 );
719 drop(dim_df);
720
721 let entity_scratch_bytes = total_docs.saturating_mul(32);
725 let remap_bytes = max_dims.saturating_mul(4);
726 let fixed_bytes = entity_scratch_bytes
727 .saturating_add(remap_bytes)
728 .saturating_add(jobs_bytes);
729 let fit = fit_candidates_to_budget(&mut eligible, fixed_bytes, memory_budget_bytes);
730 if fit.dropped > 0 {
731 budget_limited = true;
732 log::warn!(
733 "[reorder] memory budget {}: estimated {}, dropped {} highest-df dims, keeping {} ({} postings)",
734 crate::format_bytes(memory_budget_bytes as u64),
735 crate::format_bytes(fit.estimated_bytes as u64),
736 fit.dropped,
737 eligible.len(),
738 fit.retained_postings,
739 );
740 }
741
742 if eligible.is_empty() {
743 return ForwardIndex {
748 terms: Vec::new(),
749 offsets: Vec::new(),
750 num_terms: 0,
751 parallel_bisect_lanes: 1,
752 cache_gains: false,
753 budget_limited,
754 };
755 }
756
757 let mut term_remap = vec![u32::MAX; max_dims];
758 for (compact_id, &(dim_id, _)) in eligible.iter().enumerate() {
759 term_remap[dim_id as usize] = compact_id as u32;
760 }
761 let num_active_terms = eligible.len();
762 let non_degree_bytes = entity_scratch_bytes.saturating_add(
767 fit.retained_postings
768 .saturating_mul(std::mem::size_of::<u32>()),
769 );
770 let parallel_bisect_lanes =
771 parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_active_terms);
772 drop(eligible);
773
774 let mut counts = vec![0u32; total_docs];
776 let fill_block_counts = |job: &BlockJob, out: &mut [u32]| {
777 let bmp = bmps[job.src as usize];
778 let (v2r, _) = &vid_maps[job.src as usize];
779 let block_size = bmp.bmp_block_size as usize;
780 for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
781 if term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX) == u32::MAX {
782 continue;
783 }
784 for p in postings {
785 let vid = job.block_id as usize * block_size + p.local_slot as usize;
786 let Some(&real) = v2r.get(vid) else {
787 continue;
788 };
789 if real != u32::MAX && p.impact > 0 {
790 out[(real - job.real_start) as usize] += 1;
791 }
792 }
793 }
794 };
795 {
796 let mut slices: Vec<(&BlockJob, &mut [u32])> = Vec::with_capacity(jobs.len());
797 let mut rest: &mut [u32] = &mut counts;
798 for job in &jobs {
799 let (head, tail) = rest.split_at_mut(job.real_len as usize);
800 slices.push((job, head));
801 rest = tail;
802 }
803 #[cfg(feature = "native")]
804 slices
805 .into_par_iter()
806 .for_each(|(job, out)| fill_block_counts(job, out));
807 #[cfg(not(feature = "native"))]
808 for (job, out) in slices {
809 fill_block_counts(job, out);
810 }
811 }
812
813 let offsets = build_csr_offsets(&counts);
815 let total = *offsets.last().unwrap() as usize;
816 drop(counts);
817
818 let mut terms = vec![0u32; total];
821 let fill_block_terms = |job: &BlockJob, global_real_start: usize, out: &mut [u32]| {
822 let bmp = bmps[job.src as usize];
823 let (v2r, _) = &vid_maps[job.src as usize];
824 let block_size = bmp.bmp_block_size as usize;
825 assert!(job.real_len as usize <= 256, "BMP block exceeds 256 docs");
827 let mut cursor = [0u32; 256];
828 let base = offsets[global_real_start] as usize;
829 for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
830 let compact = term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX);
831 if compact == u32::MAX {
832 continue;
833 }
834 for p in postings {
835 let vid = job.block_id as usize * block_size + p.local_slot as usize;
836 let Some(&real) = v2r.get(vid) else {
837 continue;
838 };
839 if real != u32::MAX && p.impact > 0 {
840 let local = (real - job.real_start) as usize;
841 let pos =
842 offsets[global_real_start + local] as usize - base + cursor[local] as usize;
843 out[pos] = compact;
844 cursor[local] += 1;
845 }
846 }
847 }
848 };
849 {
850 let mut slices: Vec<(&BlockJob, usize, &mut [u32])> = Vec::with_capacity(jobs.len());
851 let mut rest: &mut [u32] = &mut terms;
852 let mut global_real = 0usize;
853 for job in &jobs {
854 let len =
855 (offsets[global_real + job.real_len as usize] - offsets[global_real]) as usize;
856 let (head, tail) = rest.split_at_mut(len);
857 slices.push((job, global_real, head));
858 rest = tail;
859 global_real += job.real_len as usize;
860 }
861 #[cfg(feature = "native")]
862 slices
863 .into_par_iter()
864 .for_each(|(job, g, out)| fill_block_terms(job, g, out));
865 #[cfg(not(feature = "native"))]
866 for (job, g, out) in slices {
867 fill_block_terms(job, g, out);
868 }
869 }
870
871 ForwardIndex {
872 terms,
873 offsets,
874 num_terms: num_active_terms,
875 parallel_bisect_lanes,
876 cache_gains: gain_cache_fits(
877 memory_budget_bytes,
878 non_degree_bytes,
879 num_active_terms,
880 parallel_bisect_lanes,
881 ),
882 budget_limited,
883 }
884}
885
886pub(crate) fn build_forward_index_from_blocks(
894 bmps: &[&crate::segment::reader::bmp::BmpIndex],
895 memory_budget_bytes: usize,
896) -> ForwardIndex {
897 let total_blocks: usize = bmps.iter().map(|b| b.num_blocks as usize).sum();
898 if total_blocks == 0 {
899 return ForwardIndex {
900 terms: Vec::new(),
901 offsets: Vec::new(),
902 num_terms: 0,
903 parallel_bisect_lanes: 1,
904 cache_gains: false,
905 budget_limited: false,
906 };
907 }
908
909 let blocks: Vec<(u32, u32)> = bmps
911 .iter()
912 .enumerate()
913 .flat_map(|(src, bmp)| (0..bmp.num_blocks).map(move |b| (src as u32, b)))
914 .collect();
915
916 let max_dims = bmps
919 .iter()
920 .map(|bmp| bmp.dims() as usize)
921 .max()
922 .unwrap_or(0);
923 let blocks_bytes = blocks
924 .len()
925 .saturating_mul(std::mem::size_of::<(u32, u32)>().saturating_add(32));
926 let frequency_bytes = max_dims.saturating_mul(std::mem::size_of::<u32>());
927 if frequency_bytes > memory_budget_bytes.saturating_sub(blocks_bytes) {
928 log::warn!(
929 "[reorder] block-level frequency table exceeds memory budget; using identity order"
930 );
931 return ForwardIndex {
932 terms: Vec::new(),
933 offsets: Vec::new(),
934 num_terms: 0,
935 parallel_bisect_lanes: 1,
936 cache_gains: false,
937 budget_limited: true,
938 };
939 }
940 let Some(dim_bf) = count_frequencies_bounded(
941 &blocks,
942 max_dims,
943 memory_budget_bytes.saturating_sub(blocks_bytes),
944 |&(src, block_id), counts| {
945 for (dim_id, _, _) in bmps[src as usize].iter_block_terms(block_id) {
946 if let Some(count) = counts.get_mut(dim_id as usize) {
947 *count = count.saturating_add(1);
948 }
949 }
950 },
951 ) else {
952 log::warn!(
953 "[reorder] block-level frequency table cannot fit its bounded allocation; using identity order"
954 );
955 return ForwardIndex {
956 terms: Vec::new(),
957 offsets: Vec::new(),
958 num_terms: 0,
959 parallel_bisect_lanes: 1,
960 cache_gains: false,
961 budget_limited: true,
962 };
963 };
964
965 let max_bf = (total_blocks as f64 * 0.9) as usize;
966 let (mut eligible, mut budget_limited) = select_frequency_candidates(
967 &dim_bf,
968 2,
969 max_bf.max(2),
970 memory_budget_bytes
971 .saturating_sub(blocks_bytes)
972 .saturating_sub(frequency_bytes),
973 );
974 drop(dim_bf);
975
976 let entity_scratch_bytes = total_blocks.saturating_mul(32);
977 let remap_bytes = max_dims.saturating_mul(4);
978 let fixed_bytes = entity_scratch_bytes
979 .saturating_add(remap_bytes)
980 .saturating_add(blocks_bytes);
981 let fit = fit_candidates_to_budget(&mut eligible, fixed_bytes, memory_budget_bytes);
982 if fit.dropped > 0 {
983 budget_limited = true;
984 log::warn!(
985 "[reorder] block-level fwd index over budget — dropped {} highest-bf dims",
986 fit.dropped,
987 );
988 }
989
990 if eligible.is_empty() {
991 return ForwardIndex {
992 terms: Vec::new(),
993 offsets: Vec::new(),
994 num_terms: 0,
995 parallel_bisect_lanes: 1,
996 cache_gains: false,
997 budget_limited,
998 };
999 }
1000
1001 let mut term_remap = vec![u32::MAX; max_dims];
1002 for (compact, &(dim_id, _)) in eligible.iter().enumerate() {
1003 term_remap[dim_id as usize] = compact as u32;
1004 }
1005 let num_terms = eligible.len();
1006 let non_degree_bytes = entity_scratch_bytes.saturating_add(
1007 fit.retained_postings
1008 .saturating_mul(std::mem::size_of::<u32>()),
1009 );
1010 let parallel_bisect_lanes =
1011 parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_terms);
1012 drop(eligible);
1013
1014 let count_remapped = |&(src, block_id): &(u32, u32)| -> u32 {
1017 bmps[src as usize]
1018 .iter_block_terms(block_id)
1019 .filter(|(dim_id, _, _)| {
1020 term_remap
1021 .get(*dim_id as usize)
1022 .copied()
1023 .unwrap_or(u32::MAX)
1024 != u32::MAX
1025 })
1026 .count() as u32
1027 };
1028 #[cfg(feature = "native")]
1029 let counts: Vec<u32> = blocks.par_iter().map(count_remapped).collect();
1030 #[cfg(not(feature = "native"))]
1031 let counts: Vec<u32> = blocks.iter().map(count_remapped).collect();
1032
1033 let offsets = build_csr_offsets(&counts);
1034 let total = *offsets.last().unwrap() as usize;
1035 drop(counts);
1036
1037 let mut terms = vec![0u32; total];
1038 let fill_block = |&(src, block_id): &(u32, u32), out: &mut [u32]| {
1039 let mut n = 0usize;
1040 for (dim_id, _, _) in bmps[src as usize].iter_block_terms(block_id) {
1041 let compact = term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX);
1042 if compact != u32::MAX {
1043 out[n] = compact;
1044 n += 1;
1045 }
1046 }
1047 };
1048 {
1049 let mut slices: Vec<(&(u32, u32), &mut [u32])> = Vec::with_capacity(blocks.len());
1050 let mut rest: &mut [u32] = &mut terms;
1051 for (gb, b) in blocks.iter().enumerate() {
1052 let len = (offsets[gb + 1] - offsets[gb]) as usize;
1053 let (head, tail) = rest.split_at_mut(len);
1054 slices.push((b, head));
1055 rest = tail;
1056 }
1057 #[cfg(feature = "native")]
1058 slices
1059 .into_par_iter()
1060 .for_each(|(b, out)| fill_block(b, out));
1061 #[cfg(not(feature = "native"))]
1062 for (b, out) in slices {
1063 fill_block(b, out);
1064 }
1065 }
1066
1067 ForwardIndex {
1068 terms,
1069 offsets,
1070 num_terms,
1071 parallel_bisect_lanes,
1072 cache_gains: gain_cache_fits(
1073 memory_budget_bytes,
1074 non_degree_bytes,
1075 num_terms,
1076 parallel_bisect_lanes,
1077 ),
1078 budget_limited,
1079 }
1080}
1081
1082#[derive(Clone, Copy, Debug, Default)]
1090pub struct BpBudget {
1091 pub min_partition_docs: Option<usize>,
1096 pub time_budget: Option<std::time::Duration>,
1100}
1101
1102impl BpBudget {
1103 pub fn full() -> Self {
1105 Self::default()
1106 }
1107}
1108
1109fn build_term_degrees(
1116 docs: &[u32],
1117 mid: usize,
1118 fwd: &ForwardIndex,
1119 workspaces: &mut [TermDegrees],
1120 cancellation: Option<&std::sync::atomic::AtomicBool>,
1121) {
1122 debug_assert!(!workspaces.is_empty());
1123 let build_range = |degrees: &mut TermDegrees, start: usize, chunk: &[u32]| {
1124 degrees.reset();
1125 for (offset, &doc) in chunk.iter().enumerate() {
1126 if offset.is_multiple_of(1024)
1127 && cancellation
1128 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
1129 {
1130 break;
1131 }
1132 let side = usize::from(start + offset >= mid);
1133 for &term in fwd.doc_terms(doc as usize) {
1134 degrees.entry_mut(term as usize)[side] += 1;
1135 }
1136 }
1137 };
1138
1139 #[cfg(feature = "native")]
1140 {
1141 let workers = workspaces
1142 .len()
1143 .min(docs.len().div_ceil(PARALLEL_BP_MIN_ENTITIES).max(1));
1144 if workers > 1 {
1145 let chunk_len = docs.len().div_ceil(workers);
1146 workspaces[..workers]
1147 .par_iter_mut()
1148 .enumerate()
1149 .for_each(|(worker, degrees)| {
1150 let start = worker * chunk_len;
1151 let end = (start + chunk_len).min(docs.len());
1152 build_range(degrees, start, &docs[start..end]);
1153 });
1154 let (degrees, partials) = workspaces[..workers]
1155 .split_first_mut()
1156 .expect("at least one BP degree workspace");
1157 for partial in partials {
1158 degrees.merge_from(partial);
1159 }
1160 return;
1161 }
1162 build_range(&mut workspaces[0], 0, docs);
1163 }
1164
1165 #[cfg(not(feature = "native"))]
1166 build_range(&mut workspaces[0], 0, docs);
1167}
1168
1169#[inline]
1175fn gain_order_key(gain: f32) -> u32 {
1176 let bits = gain.to_bits();
1177 if bits & 0x8000_0000 != 0 {
1178 !bits
1179 } else {
1180 bits ^ 0x8000_0000
1181 }
1182}
1183
1184fn select_gain_threshold(
1188 gains: &[f32],
1189 left_count: usize,
1190 allow_inner_parallelism: bool,
1191) -> (u32, usize) {
1192 #[cfg(not(feature = "native"))]
1193 let _ = allow_inner_parallelism;
1194
1195 debug_assert!(left_count > 0 && left_count <= gains.len());
1196 let mut rank_within_prefix = left_count - 1;
1197 let mut strictly_lower = 0usize;
1198 let mut prefix = 0u32;
1199 let mut prefix_mask = 0u32;
1200
1201 for shift in [24u32, 16, 8, 0] {
1202 let histogram = {
1203 #[cfg(feature = "native")]
1204 {
1205 if allow_inner_parallelism && gains.len() >= PARALLEL_BP_MIN_ENTITIES {
1206 gains
1207 .par_iter()
1208 .fold(
1209 || Box::new([0usize; 256]),
1210 |mut counts, &gain| {
1211 let key = gain_order_key(gain);
1212 if key & prefix_mask == prefix {
1213 counts[((key >> shift) & 0xff) as usize] += 1;
1214 }
1215 counts
1216 },
1217 )
1218 .reduce(
1219 || Box::new([0usize; 256]),
1220 |mut left, right| {
1221 for (dst, &count) in left.iter_mut().zip(right.iter()) {
1222 *dst += count;
1223 }
1224 left
1225 },
1226 )
1227 } else {
1228 let mut counts = [0usize; 256];
1229 for &gain in gains {
1230 let key = gain_order_key(gain);
1231 if key & prefix_mask == prefix {
1232 counts[((key >> shift) & 0xff) as usize] += 1;
1233 }
1234 }
1235 Box::new(counts)
1236 }
1237 }
1238 #[cfg(not(feature = "native"))]
1239 {
1240 let mut counts = [0usize; 256];
1241 for &gain in gains {
1242 let key = gain_order_key(gain);
1243 if key & prefix_mask == prefix {
1244 counts[((key >> shift) & 0xff) as usize] += 1;
1245 }
1246 }
1247 counts
1248 }
1249 };
1250
1251 let mut before_bucket = 0usize;
1252 let mut selected_bucket = None;
1253 for (bucket, count) in histogram.iter().copied().enumerate() {
1254 if rank_within_prefix < before_bucket + count {
1255 selected_bucket = Some(bucket as u32);
1256 rank_within_prefix -= before_bucket;
1257 strictly_lower += before_bucket;
1258 break;
1259 }
1260 before_bucket += count;
1261 }
1262 let selected_bucket = selected_bucket.expect("BP radix selection lost the requested rank");
1263 prefix |= selected_bucket << shift;
1264 prefix_mask |= 0xffu32 << shift;
1265 }
1266
1267 (prefix, strictly_lower)
1268}
1269
1270enum PartitionDegreeUpdate {
1271 Moves,
1274 Ranked,
1278 Threshold {
1281 threshold_key: u32,
1282 ties_left: usize,
1283 },
1284}
1285
1286struct PartitionOutcome {
1287 swap_count: usize,
1288 degree_update: PartitionDegreeUpdate,
1289}
1290
1291#[derive(Clone, Copy)]
1292struct PartitionChunk {
1293 start: usize,
1294 end: usize,
1295 strictly_lower: usize,
1296 equal: usize,
1297 ties_left: usize,
1298}
1299
1300#[inline]
1301fn select_left(key: u32, threshold_key: u32, equal_seen: &mut usize, ties_left: usize) -> bool {
1302 if key < threshold_key {
1303 true
1304 } else if key == threshold_key {
1305 let selected = *equal_seen < ties_left;
1306 *equal_seen += 1;
1307 selected
1308 } else {
1309 false
1310 }
1311}
1312
1313#[allow(clippy::too_many_arguments)]
1322fn partition_by_gain(
1323 docs: &[u32],
1324 gains: &[f32],
1325 mid: usize,
1326 fwd: &ForwardIndex,
1327 movement_workspaces: &mut [TermDegrees],
1328 output: &mut [u32],
1329 ranked_scratch: &mut [usize],
1330 allow_inner_parallelism: bool,
1331) -> PartitionOutcome {
1332 #[cfg(not(feature = "native"))]
1333 let _ = (fwd, &movement_workspaces);
1334
1335 #[cfg(feature = "native")]
1336 if allow_inner_parallelism
1337 && !movement_workspaces.is_empty()
1338 && docs.len() >= PARALLEL_BP_MIN_ENTITIES
1339 {
1340 let (threshold_key, strictly_lower) =
1341 select_gain_threshold(gains, mid, allow_inner_parallelism);
1342 let ties_left = mid - strictly_lower;
1343
1344 let chunk_count = movement_workspaces
1347 .len()
1348 .min(docs.len().div_ceil(PARALLEL_BP_MIN_ENTITIES).max(1));
1349 let chunk_len = docs.len().div_ceil(chunk_count);
1350 let mut chunks: Vec<PartitionChunk> = gains
1351 .par_chunks(chunk_len)
1352 .enumerate()
1353 .map(|(chunk_id, chunk)| {
1354 let mut lower = 0usize;
1355 let mut equal = 0usize;
1356 for &gain in chunk {
1357 match gain_order_key(gain).cmp(&threshold_key) {
1358 std::cmp::Ordering::Less => lower += 1,
1359 std::cmp::Ordering::Equal => equal += 1,
1360 std::cmp::Ordering::Greater => {}
1361 }
1362 }
1363 let start = chunk_id * chunk_len;
1364 PartitionChunk {
1365 start,
1366 end: start + chunk.len(),
1367 strictly_lower: lower,
1368 equal,
1369 ties_left: 0,
1370 }
1371 })
1372 .collect();
1373
1374 let mut remaining_ties = ties_left;
1375 for chunk in &mut chunks {
1376 chunk.ties_left = remaining_ties.min(chunk.equal);
1377 remaining_ties -= chunk.ties_left;
1378 }
1379 debug_assert_eq!(remaining_ties, 0);
1380
1381 let (mut left_rest, mut right_rest) = output.split_at_mut(mid);
1382 let mut jobs = Vec::with_capacity(chunks.len());
1383 for chunk in chunks {
1384 let left_len = chunk.strictly_lower + chunk.ties_left;
1385 let right_len = chunk.end - chunk.start - left_len;
1386 let (left_out, next_left) = left_rest.split_at_mut(left_len);
1387 let (right_out, next_right) = right_rest.split_at_mut(right_len);
1388 jobs.push((
1389 chunk.start,
1390 &docs[chunk.start..chunk.end],
1391 &gains[chunk.start..chunk.end],
1392 chunk.ties_left,
1393 left_out,
1394 right_out,
1395 ));
1396 left_rest = next_left;
1397 right_rest = next_right;
1398 }
1399 debug_assert!(left_rest.is_empty() && right_rest.is_empty());
1400
1401 let swap_count = movement_workspaces[..chunk_count]
1402 .par_iter_mut()
1403 .zip(jobs.into_par_iter())
1404 .map(
1405 |(moves, (start, docs, gains, ties_for_chunk, left_out, right_out))| {
1406 moves.reset();
1407 let mut equal_seen = 0usize;
1408 let mut left_cursor = 0usize;
1409 let mut right_cursor = 0usize;
1410 let mut swaps = 0usize;
1411
1412 for (offset, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1413 let key = gain_order_key(gain);
1414 let now_left =
1415 select_left(key, threshold_key, &mut equal_seen, ties_for_chunk);
1416 if now_left {
1417 left_out[left_cursor] = doc;
1418 left_cursor += 1;
1419 } else {
1420 right_out[right_cursor] = doc;
1421 right_cursor += 1;
1422 }
1423
1424 let was_left = start + offset < mid;
1425 if was_left != now_left {
1426 swaps += 1;
1427 let direction = usize::from(was_left);
1429 for &term in fwd.doc_terms(doc as usize) {
1430 moves.entry_mut(term as usize)[direction] += 1;
1431 }
1432 }
1433 }
1434 debug_assert_eq!(left_cursor, left_out.len());
1435 debug_assert_eq!(right_cursor, right_out.len());
1436 swaps
1437 },
1438 )
1439 .sum();
1440
1441 let (moves, partials) = movement_workspaces[..chunk_count]
1442 .split_first_mut()
1443 .expect("parallel BP partition must use at least one movement workspace");
1444 for partial in partials {
1445 moves.merge_from(partial);
1446 }
1447
1448 return PartitionOutcome {
1449 swap_count,
1450 degree_update: PartitionDegreeUpdate::Moves,
1451 };
1452 }
1453
1454 if docs.len() < PARALLEL_BP_MIN_ENTITIES {
1455 debug_assert_eq!(ranked_scratch.len(), docs.len());
1456 for (index, rank) in ranked_scratch.iter_mut().enumerate() {
1457 *rank = index;
1458 }
1459 ranked_scratch.select_nth_unstable_by(mid, |&left, &right| {
1460 gains[left]
1461 .total_cmp(&gains[right])
1462 .then_with(|| left.cmp(&right))
1463 });
1464
1465 let mut swaps = 0usize;
1466 for (rank, &old_index) in ranked_scratch.iter().enumerate() {
1467 output[rank] = docs[old_index];
1468 swaps += usize::from((old_index < mid) != (rank < mid));
1469 }
1470 return PartitionOutcome {
1471 swap_count: swaps,
1472 degree_update: PartitionDegreeUpdate::Ranked,
1473 };
1474 }
1475
1476 let (threshold_key, strictly_lower) =
1477 select_gain_threshold(gains, mid, allow_inner_parallelism);
1478 let ties_left = mid - strictly_lower;
1479 let mut equal_seen = 0usize;
1480 let mut left_cursor = 0usize;
1481 let mut right_cursor = mid;
1482 let mut swaps = 0usize;
1483 for (idx, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1484 let key = gain_order_key(gain);
1485 let now_left = select_left(key, threshold_key, &mut equal_seen, ties_left);
1486 if now_left {
1487 output[left_cursor] = doc;
1488 left_cursor += 1;
1489 } else {
1490 output[right_cursor] = doc;
1491 right_cursor += 1;
1492 }
1493 swaps += usize::from((idx < mid) != now_left);
1494 }
1495 debug_assert_eq!(left_cursor, mid);
1496 debug_assert_eq!(right_cursor, docs.len());
1497
1498 PartitionOutcome {
1499 swap_count: swaps,
1500 degree_update: PartitionDegreeUpdate::Threshold {
1501 threshold_key,
1502 ties_left,
1503 },
1504 }
1505}
1506
1507fn update_degrees_for_threshold_partition(
1509 docs: &[u32],
1510 gains: &[f32],
1511 mid: usize,
1512 threshold_key: u32,
1513 ties_left: usize,
1514 fwd: &ForwardIndex,
1515 degrees: &mut TermDegrees,
1516) {
1517 let mut equal_seen = 0usize;
1518 for (idx, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1519 let key = gain_order_key(gain);
1520 let now_left = select_left(key, threshold_key, &mut equal_seen, ties_left);
1521 let was_left = idx < mid;
1522 if was_left == now_left {
1523 continue;
1524 }
1525 let left_delta = if was_left { -1i64 } else { 1i64 };
1526 for &term in fwd.doc_terms(doc as usize) {
1527 let degree = degrees.entry_mut(term as usize);
1528 let new_left = degree[0] as i64 + left_delta;
1529 let new_right = degree[1] as i64 - left_delta;
1530 debug_assert!(new_left >= 0 && new_right >= 0);
1531 degree[0] = new_left as u32;
1532 degree[1] = new_right as u32;
1533 }
1534 }
1535}
1536
1537fn update_degrees_for_ranked_partition(
1540 docs: &[u32],
1541 ranked: &[usize],
1542 mid: usize,
1543 fwd: &ForwardIndex,
1544 degrees: &mut TermDegrees,
1545) {
1546 for (rank, &old_index) in ranked.iter().enumerate() {
1547 let was_left = old_index < mid;
1548 let now_left = rank < mid;
1549 if was_left == now_left {
1550 continue;
1551 }
1552 let left_delta = if was_left { -1i64 } else { 1i64 };
1553 for &term in fwd.doc_terms(docs[old_index] as usize) {
1554 let degree = degrees.entry_mut(term as usize);
1555 let new_left = degree[0] as i64 + left_delta;
1556 let new_right = degree[1] as i64 - left_delta;
1557 debug_assert!(new_left >= 0 && new_right >= 0);
1558 degree[0] = new_left as u32;
1559 degree[1] = new_right as u32;
1560 }
1561 }
1562}
1563
1564#[derive(Clone, Copy)]
1565pub(crate) struct BpProgressLabel<'a> {
1566 pub index: &'a str,
1567 pub field: &'a str,
1568 pub entity_kind: &'static str,
1569}
1570
1571#[cfg(test)]
1572impl BpProgressLabel<'static> {
1573 fn anonymous() -> Self {
1574 Self {
1575 index: "unknown",
1576 field: "unknown",
1577 entity_kind: "entities",
1578 }
1579 }
1580}
1581
1582#[cfg(feature = "native")]
1583struct BpProgress<'a> {
1584 label: BpProgressLabel<'a>,
1585 start: std::time::Instant,
1586 total_entities: usize,
1587 total_postings: u64,
1588 expected_depth: usize,
1589 next_log_ms: std::sync::atomic::AtomicU64,
1590 active_partitions: std::sync::atomic::AtomicU64,
1591 partitions_started: std::sync::atomic::AtomicU64,
1592 partitions_completed: std::sync::atomic::AtomicU64,
1593 iterations: std::sync::atomic::AtomicU64,
1594 entity_passes: std::sync::atomic::AtomicU64,
1595 swaps: std::sync::atomic::AtomicU64,
1596 deepest_level: std::sync::atomic::AtomicU64,
1597 objective_stops: std::sync::atomic::AtomicU64,
1598 last_objective_delta_bits: std::sync::atomic::AtomicU64,
1599 last_relative_delta_bits: std::sync::atomic::AtomicU64,
1600 active_metric_released: std::sync::atomic::AtomicBool,
1601}
1602
1603#[cfg(feature = "native")]
1604impl<'a> BpProgress<'a> {
1605 fn new(
1606 label: BpProgressLabel<'a>,
1607 total_entities: usize,
1608 total_postings: u64,
1609 expected_depth: usize,
1610 degree_lanes: usize,
1611 ) -> Self {
1612 log::info!(
1613 "[reorder][bp] started: index={} field={} entity_kind={} scheduler=level_synchronized degree_lanes={} entities={} postings={} expected_depth={} objective_stall_threshold={:.1e}x{} min_objective_iterations={}",
1614 label.index,
1615 label.field,
1616 label.entity_kind,
1617 degree_lanes,
1618 total_entities,
1619 total_postings,
1620 expected_depth,
1621 MIN_RELATIVE_OBJECTIVE_IMPROVEMENT,
1622 OBJECTIVE_STALL_ITERATIONS,
1623 MIN_OBJECTIVE_ITERATIONS,
1624 );
1625 crate::observe::reorder_bp_started(label.index, label.field, label.entity_kind);
1626 Self {
1627 label,
1628 start: std::time::Instant::now(),
1629 total_entities,
1630 total_postings,
1631 expected_depth,
1632 next_log_ms: std::sync::atomic::AtomicU64::new(30_000),
1633 active_partitions: std::sync::atomic::AtomicU64::new(0),
1634 partitions_started: std::sync::atomic::AtomicU64::new(0),
1635 partitions_completed: std::sync::atomic::AtomicU64::new(0),
1636 iterations: std::sync::atomic::AtomicU64::new(0),
1637 entity_passes: std::sync::atomic::AtomicU64::new(0),
1638 swaps: std::sync::atomic::AtomicU64::new(0),
1639 deepest_level: std::sync::atomic::AtomicU64::new(0),
1640 objective_stops: std::sync::atomic::AtomicU64::new(0),
1641 last_objective_delta_bits: std::sync::atomic::AtomicU64::new(0f64.to_bits()),
1642 last_relative_delta_bits: std::sync::atomic::AtomicU64::new(0f64.to_bits()),
1643 active_metric_released: std::sync::atomic::AtomicBool::new(false),
1644 }
1645 }
1646
1647 fn partition_started(&self, level: usize) {
1648 self.active_partitions
1649 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1650 self.partitions_started
1651 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1652 self.deepest_level
1653 .fetch_max(level as u64, std::sync::atomic::Ordering::Relaxed);
1654 }
1655
1656 fn partition_finished(&self) {
1657 self.partitions_completed
1658 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1659 self.active_partitions
1660 .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
1661 }
1662
1663 fn iteration(&self, entities: usize, swaps: usize, objective_delta: f64, relative_delta: f64) {
1664 self.iterations
1665 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1666 self.entity_passes
1667 .fetch_add(entities as u64, std::sync::atomic::Ordering::Relaxed);
1668 self.swaps
1669 .fetch_add(swaps as u64, std::sync::atomic::Ordering::Relaxed);
1670 self.last_objective_delta_bits.store(
1671 objective_delta.to_bits(),
1672 std::sync::atomic::Ordering::Relaxed,
1673 );
1674 self.last_relative_delta_bits.store(
1675 relative_delta.to_bits(),
1676 std::sync::atomic::Ordering::Relaxed,
1677 );
1678 self.maybe_log();
1679 }
1680
1681 fn objective_stop(&self) {
1682 self.objective_stops
1683 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1684 }
1685
1686 fn maybe_log(&self) {
1687 let elapsed_ms = self.start.elapsed().as_millis().min(u64::MAX as u128) as u64;
1688 let next = self.next_log_ms.load(std::sync::atomic::Ordering::Relaxed);
1689 if elapsed_ms < next
1690 || self
1691 .next_log_ms
1692 .compare_exchange(
1693 next,
1694 elapsed_ms.saturating_add(30_000),
1695 std::sync::atomic::Ordering::Relaxed,
1696 std::sync::atomic::Ordering::Relaxed,
1697 )
1698 .is_err()
1699 {
1700 return;
1701 }
1702
1703 let active = self
1704 .active_partitions
1705 .load(std::sync::atomic::Ordering::Relaxed);
1706 let started = self
1707 .partitions_started
1708 .load(std::sync::atomic::Ordering::Relaxed);
1709 let completed = self
1710 .partitions_completed
1711 .load(std::sync::atomic::Ordering::Relaxed);
1712 let iterations = self.iterations.load(std::sync::atomic::Ordering::Relaxed);
1713 let entity_passes = self
1714 .entity_passes
1715 .load(std::sync::atomic::Ordering::Relaxed);
1716 let swaps = self.swaps.load(std::sync::atomic::Ordering::Relaxed);
1717 let deepest = self
1718 .deepest_level
1719 .load(std::sync::atomic::Ordering::Relaxed);
1720 let objective_delta = f64::from_bits(
1721 self.last_objective_delta_bits
1722 .load(std::sync::atomic::Ordering::Relaxed),
1723 );
1724 let relative_delta = f64::from_bits(
1725 self.last_relative_delta_bits
1726 .load(std::sync::atomic::Ordering::Relaxed),
1727 );
1728 log::info!(
1729 "[reorder][bp] progress: index={} field={} entity_kind={} elapsed={:.1}s depth={}/{} partitions={}/{} active={} iterations={} entity_passes={} swaps={} last_objective_delta={:.3} relative={:.3e}",
1730 self.label.index,
1731 self.label.field,
1732 self.label.entity_kind,
1733 self.start.elapsed().as_secs_f64(),
1734 deepest,
1735 self.expected_depth,
1736 completed,
1737 started,
1738 active,
1739 iterations,
1740 entity_passes,
1741 swaps,
1742 objective_delta,
1743 relative_delta,
1744 );
1745 }
1746
1747 fn finish(
1748 &self,
1749 converged: bool,
1750 memory_limited: bool,
1751 deadline_exhausted: bool,
1752 cancelled: bool,
1753 ) {
1754 let elapsed = self.start.elapsed().as_secs_f64();
1755 let partitions = self
1756 .partitions_completed
1757 .load(std::sync::atomic::Ordering::Relaxed);
1758 let iterations = self.iterations.load(std::sync::atomic::Ordering::Relaxed);
1759 let entity_passes = self
1760 .entity_passes
1761 .load(std::sync::atomic::Ordering::Relaxed);
1762 let swaps = self.swaps.load(std::sync::atomic::Ordering::Relaxed);
1763 let deepest = self
1764 .deepest_level
1765 .load(std::sync::atomic::Ordering::Relaxed);
1766 let objective_stops = self
1767 .objective_stops
1768 .load(std::sync::atomic::Ordering::Relaxed);
1769 let stop_reason = if cancelled {
1770 "shutdown"
1771 } else if memory_limited {
1772 "memory_budget"
1773 } else if deadline_exhausted {
1774 "time_budget"
1775 } else if objective_stops > 0 {
1776 "objective"
1777 } else {
1778 "complete"
1779 };
1780 log::info!(
1781 "[reorder][bp] completed: index={} field={} entity_kind={} entities={} postings={} elapsed={:.1}s depth={}/{} partitions={} iterations={} entity_passes={} swaps={} objective_stops={} converged={} stop_reason={}",
1782 self.label.index,
1783 self.label.field,
1784 self.label.entity_kind,
1785 self.total_entities,
1786 self.total_postings,
1787 elapsed,
1788 deepest,
1789 self.expected_depth,
1790 partitions,
1791 iterations,
1792 entity_passes,
1793 swaps,
1794 objective_stops,
1795 converged,
1796 stop_reason,
1797 );
1798 crate::observe::reorder_bp_pass(
1799 self.label.index,
1800 self.label.field,
1801 self.label.entity_kind,
1802 stop_reason,
1803 elapsed,
1804 self.total_entities,
1805 self.total_postings,
1806 partitions,
1807 iterations,
1808 entity_passes,
1809 swaps,
1810 converged,
1811 );
1812 self.release_active_metric();
1813 }
1814
1815 fn release_active_metric(&self) {
1816 if !self
1817 .active_metric_released
1818 .swap(true, std::sync::atomic::Ordering::AcqRel)
1819 {
1820 crate::observe::reorder_bp_finished(
1821 self.label.index,
1822 self.label.field,
1823 self.label.entity_kind,
1824 );
1825 }
1826 }
1827}
1828
1829#[cfg(feature = "native")]
1830impl Drop for BpProgress<'_> {
1831 fn drop(&mut self) {
1832 self.release_active_metric();
1833 }
1834}
1835
1836#[cfg(not(feature = "native"))]
1837struct BpProgress<'a>(std::marker::PhantomData<&'a ()>);
1838
1839#[cfg(not(feature = "native"))]
1840impl BpProgress<'_> {
1841 fn new(_: BpProgressLabel<'_>, _: usize, _: u64, _: usize, _: usize) -> Self {
1842 Self(std::marker::PhantomData)
1843 }
1844 fn partition_started(&self, _: usize) {}
1845 fn partition_finished(&self) {}
1846 fn iteration(&self, _: usize, _: usize, _: f64, _: f64) {}
1847 fn objective_stop(&self) {}
1848 fn finish(&self, _: bool, _: bool, _: bool, _: bool) {}
1849}
1850
1851#[cfg(test)]
1862pub(crate) fn graph_bisection(
1863 fwd: &ForwardIndex,
1864 min_partition_size: usize,
1865 max_iters: usize,
1866 budget: BpBudget,
1867) -> (Vec<u32>, bool) {
1868 graph_bisection_with_progress(
1869 fwd,
1870 min_partition_size,
1871 max_iters,
1872 budget,
1873 None,
1874 BpProgressLabel::anonymous(),
1875 )
1876}
1877
1878pub(crate) fn graph_bisection_with_progress(
1879 fwd: &ForwardIndex,
1880 min_partition_size: usize,
1881 max_iters: usize,
1882 budget: BpBudget,
1883 cancellation: Option<&std::sync::atomic::AtomicBool>,
1884 progress_label: BpProgressLabel<'_>,
1885) -> (Vec<u32>, bool) {
1886 let n = fwd.num_docs();
1887 if n == 0 {
1888 return (Vec::new(), !fwd.budget_limited);
1889 }
1890
1891 let effective_min_partition = budget
1892 .min_partition_docs
1893 .unwrap_or(0)
1894 .max(min_partition_size)
1895 .max(1);
1900
1901 let mut docs: Vec<u32> = (0..n as u32).collect();
1902 let depth = if effective_min_partition > 0 {
1903 ((n as f64) / (effective_min_partition as f64))
1904 .log2()
1905 .ceil() as usize
1906 } else {
1907 0
1908 };
1909 #[cfg(feature = "native")]
1910 let degree_lanes = fwd.parallel_bisect_lanes.max(1);
1911 #[cfg(not(feature = "native"))]
1912 let degree_lanes = 1usize;
1913 let log_table = build_log_table(LOG_TABLE_SIZE);
1914 let progress = BpProgress::new(progress_label, n, fwd.total_postings(), depth, degree_lanes);
1915
1916 log::debug!(
1917 "BP graph_bisection: n={}, min_partition={}, max_iters={}, depth=~{}, time_budget={:?}, gain_cache={}",
1918 n,
1919 effective_min_partition,
1920 max_iters,
1921 depth,
1922 budget.time_budget,
1923 fwd.cache_gains,
1924 );
1925
1926 #[cfg(feature = "native")]
1927 let deadline = budget.time_budget.map(|duration| {
1928 let now = std::time::Instant::now();
1929 now.checked_add(duration).unwrap_or(now)
1930 });
1931 #[cfg(not(feature = "native"))]
1932 let deadline: Option<()> = None;
1933
1934 let exhausted = std::sync::atomic::AtomicBool::new(false);
1935 let context = BisectContext {
1936 fwd,
1937 min_partition_size: effective_min_partition,
1938 max_iters,
1939 log_table: &log_table,
1940 #[cfg(feature = "native")]
1941 deadline,
1942 #[cfg(not(feature = "native"))]
1943 deadline,
1944 exhausted: &exhausted,
1945 cancellation,
1946 progress: &progress,
1947 };
1948 let immediately_exhausted = {
1949 #[cfg(feature = "native")]
1950 {
1951 deadline.is_some_and(|deadline| std::time::Instant::now() >= deadline)
1952 }
1953 #[cfg(not(feature = "native"))]
1954 {
1955 false
1956 }
1957 };
1958 let initially_cancelled =
1959 cancellation.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire));
1960 if immediately_exhausted || initially_cancelled {
1961 exhausted.store(true, std::sync::atomic::Ordering::Relaxed);
1962 } else if n > effective_min_partition {
1963 let mut gains = vec![0.0f32; n];
1967 let mut partitioned = vec![0u32; n];
1968 let mut ranked = vec![0usize; n];
1969 let mut degree_workspaces: Vec<TermDegrees> = (0..degree_lanes)
1970 .map(|_| {
1971 let mut lane = TermDegrees::new(fwd.num_terms);
1972 if fwd.cache_gains {
1973 lane.gain_cache = vec![[0.0; 2]; fwd.num_terms];
1974 }
1975 lane
1976 })
1977 .collect();
1978 bisect_level_synchronized(
1979 &mut docs,
1980 &mut gains,
1981 &mut partitioned,
1982 &mut ranked,
1983 &mut degree_workspaces,
1984 &context,
1985 );
1986 }
1987
1988 let stopped_early = exhausted.load(std::sync::atomic::Ordering::Relaxed);
1989 let cancelled =
1990 cancellation.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire));
1991 let deadline_exhausted = stopped_early && !cancelled;
1992 let converged = !fwd.budget_limited && !stopped_early && !cancelled;
1993 progress.finish(converged, fwd.budget_limited, deadline_exhausted, cancelled);
1994 if !converged {
1995 log::info!(
1996 "BP graph_bisection: pass incomplete at n={} (time={:?}, memory_limited={}, cancelled={}) — emitting partial (still valid) permutation",
1997 n,
1998 budget.time_budget,
1999 fwd.budget_limited,
2000 cancelled,
2001 );
2002 }
2003 (docs, converged)
2004}
2005
2006struct BisectContext<'a> {
2012 fwd: &'a ForwardIndex,
2013 min_partition_size: usize,
2014 max_iters: usize,
2015 log_table: &'a [f32],
2016 #[cfg(feature = "native")]
2017 deadline: Option<std::time::Instant>,
2018 #[cfg(not(feature = "native"))]
2019 deadline: Option<()>,
2020 exhausted: &'a std::sync::atomic::AtomicBool,
2021 cancellation: Option<&'a std::sync::atomic::AtomicBool>,
2022 progress: &'a BpProgress<'a>,
2023}
2024
2025fn partition_range(total: usize, level: usize, partition_id: usize) -> std::ops::Range<usize> {
2031 debug_assert!(level < usize::BITS as usize);
2032 debug_assert!(partition_id < (1usize << level));
2033 let mut start = 0usize;
2034 let mut len = total;
2035 for bit in (0..level).rev() {
2036 let left_len = len / 2;
2037 if partition_id & (1usize << bit) == 0 {
2038 len = left_len;
2039 } else {
2040 start += left_len;
2041 len -= left_len;
2042 }
2043 }
2044 start..start + len
2045}
2046
2047fn small_active_partition_set(
2051 total: usize,
2052 level: usize,
2053 min_partition_size: usize,
2054 limit: usize,
2055) -> Option<Vec<usize>> {
2056 let partition_count = 1usize.checked_shl(level as u32)?;
2057 let mut active = Vec::with_capacity(partition_count.min(limit));
2058 for partition_id in 0..partition_count {
2059 if partition_range(total, level, partition_id).len() <= min_partition_size {
2060 continue;
2061 }
2062 active.push(partition_id);
2063 if active.len() > limit {
2064 return None;
2065 }
2066 }
2067 Some(active)
2068}
2069
2070#[cfg(feature = "native")]
2077#[derive(Clone, Copy)]
2078struct LevelBuffers {
2079 docs: *mut u32,
2080 gains: *mut f32,
2081 partitioned: *mut u32,
2082 ranked: *mut usize,
2083 len: usize,
2084}
2085
2086#[cfg(feature = "native")]
2087unsafe impl Send for LevelBuffers {}
2088#[cfg(feature = "native")]
2089unsafe impl Sync for LevelBuffers {}
2090
2091#[cfg(feature = "native")]
2092impl LevelBuffers {
2093 fn new(
2094 docs: &mut [u32],
2095 gains: &mut [f32],
2096 partitioned: &mut [u32],
2097 ranked: &mut [usize],
2098 ) -> Self {
2099 debug_assert_eq!(docs.len(), gains.len());
2100 debug_assert_eq!(docs.len(), partitioned.len());
2101 debug_assert_eq!(docs.len(), ranked.len());
2102 Self {
2103 docs: docs.as_mut_ptr(),
2104 gains: gains.as_mut_ptr(),
2105 partitioned: partitioned.as_mut_ptr(),
2106 ranked: ranked.as_mut_ptr(),
2107 len: docs.len(),
2108 }
2109 }
2110
2111 unsafe fn with_slices<R>(
2116 &self,
2117 range: std::ops::Range<usize>,
2118 use_slices: impl for<'a> FnOnce(
2119 &'a mut [u32],
2120 &'a mut [f32],
2121 &'a mut [u32],
2122 &'a mut [usize],
2123 ) -> R,
2124 ) -> R {
2125 debug_assert!(range.start <= range.end && range.end <= self.len);
2126 let len = range.len();
2127 unsafe {
2130 use_slices(
2131 std::slice::from_raw_parts_mut(self.docs.add(range.start), len),
2132 std::slice::from_raw_parts_mut(self.gains.add(range.start), len),
2133 std::slice::from_raw_parts_mut(self.partitioned.add(range.start), len),
2134 std::slice::from_raw_parts_mut(self.ranked.add(range.start), len),
2135 )
2136 }
2137 }
2138}
2139
2140fn bisect_level_synchronized(
2145 docs: &mut [u32],
2146 gains: &mut [f32],
2147 partitioned: &mut [u32],
2148 ranked_scratch: &mut [usize],
2149 degree_workspaces: &mut [TermDegrees],
2150 context: &BisectContext<'_>,
2151) {
2152 debug_assert_eq!(docs.len(), gains.len());
2153 debug_assert_eq!(docs.len(), partitioned.len());
2154 debug_assert_eq!(docs.len(), ranked_scratch.len());
2155 debug_assert!(!degree_workspaces.is_empty());
2156
2157 #[cfg(feature = "native")]
2158 {
2159 let buffers = LevelBuffers::new(docs, gains, partitioned, ranked_scratch);
2160 let mut level = 0usize;
2161 while let Some(partition_count) = 1usize.checked_shl(level as u32) {
2162 let small_set = small_active_partition_set(
2163 docs.len(),
2164 level,
2165 context.min_partition_size,
2166 degree_workspaces.len(),
2167 );
2168
2169 match small_set {
2170 Some(active) if active.is_empty() => break,
2171 Some(active) => {
2172 let active_count = active.len();
2176 let allow_inner_parallelism =
2177 active_count < rayon::current_num_threads().max(1);
2178 let mut rest: &mut [TermDegrees] = &mut *degree_workspaces;
2179 let mut groups = Vec::with_capacity(active_count);
2180 for remaining_groups in (1..=active_count).rev() {
2181 let group_len = rest.len().div_ceil(remaining_groups);
2182 let (group, tail) = rest.split_at_mut(group_len);
2183 groups.push(group);
2184 rest = tail;
2185 }
2186 debug_assert!(rest.is_empty());
2187 groups.into_par_iter().zip(active.into_par_iter()).for_each(
2188 |(workspaces, partition_id)| {
2189 let range = partition_range(docs.len(), level, partition_id);
2190 unsafe {
2194 buffers.with_slices(
2195 range,
2196 |part_docs, part_gains, part_output, part_ranked| {
2197 bisect_partition(
2198 part_docs,
2199 part_gains,
2200 part_output,
2201 part_ranked,
2202 workspaces,
2203 level,
2204 allow_inner_parallelism,
2205 context,
2206 );
2207 },
2208 );
2209 }
2210 },
2211 );
2212 }
2213 None => {
2214 let next = std::sync::atomic::AtomicUsize::new(0);
2220 let allow_inner_parallelism =
2221 degree_workspaces.len() < rayon::current_num_threads().max(1);
2222 degree_workspaces.par_iter_mut().for_each(|workspace| {
2223 loop {
2224 let partition_id =
2225 next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2226 if partition_id >= partition_count {
2227 break;
2228 }
2229 let range = partition_range(docs.len(), level, partition_id);
2230 if range.len() <= context.min_partition_size {
2231 continue;
2232 }
2233 unsafe {
2237 buffers.with_slices(
2238 range,
2239 |part_docs, part_gains, part_output, part_ranked| {
2240 bisect_partition(
2241 part_docs,
2242 part_gains,
2243 part_output,
2244 part_ranked,
2245 std::slice::from_mut(workspace),
2246 level,
2247 allow_inner_parallelism,
2248 context,
2249 );
2250 },
2251 );
2252 }
2253 }
2254 });
2255 }
2256 }
2257
2258 if context.exhausted.load(std::sync::atomic::Ordering::Relaxed) {
2259 break;
2260 }
2261 level += 1;
2262 }
2263 }
2264
2265 #[cfg(not(feature = "native"))]
2266 {
2267 let mut level = 0usize;
2268 while let Some(partition_count) = 1usize.checked_shl(level as u32) {
2269 let mut active = false;
2270 for partition_id in 0..partition_count {
2271 let range = partition_range(docs.len(), level, partition_id);
2272 if range.len() <= context.min_partition_size {
2273 continue;
2274 }
2275 active = true;
2276 bisect_partition(
2277 &mut docs[range.clone()],
2278 &mut gains[range.clone()],
2279 &mut partitioned[range.clone()],
2280 &mut ranked_scratch[range],
2281 degree_workspaces,
2282 level,
2283 false,
2284 context,
2285 );
2286 }
2287 if !active || context.exhausted.load(std::sync::atomic::Ordering::Relaxed) {
2288 break;
2289 }
2290 level += 1;
2291 }
2292 }
2293}
2294
2295#[allow(clippy::too_many_arguments)]
2301fn bisect_partition(
2302 docs: &mut [u32],
2303 gains: &mut [f32],
2304 partitioned: &mut [u32],
2305 ranked_scratch: &mut [usize],
2306 degree_workspaces: &mut [TermDegrees],
2307 level: usize,
2308 allow_inner_parallelism: bool,
2309 context: &BisectContext<'_>,
2310) {
2311 let n = docs.len();
2312 debug_assert_eq!(gains.len(), n);
2313 debug_assert_eq!(partitioned.len(), n);
2314 debug_assert_eq!(ranked_scratch.len(), n);
2315 debug_assert!(!degree_workspaces.is_empty());
2316 if n <= context.min_partition_size {
2317 return;
2318 }
2319 if context.exhausted.load(std::sync::atomic::Ordering::Relaxed)
2321 || context
2322 .cancellation
2323 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2324 {
2325 context
2326 .exhausted
2327 .store(true, std::sync::atomic::Ordering::Relaxed);
2328 return;
2329 }
2330 #[cfg(feature = "native")]
2331 if let Some(dl) = context.deadline
2332 && std::time::Instant::now() >= dl
2333 {
2334 context
2335 .exhausted
2336 .store(true, std::sync::atomic::Ordering::Relaxed);
2337 return;
2338 }
2339 #[cfg(not(feature = "native"))]
2340 let _ = context.deadline;
2341
2342 context.progress.partition_started(level);
2343 let mid = n / 2;
2344
2345 let effective_iters = if n > 100_000 {
2349 context.max_iters.min(12)
2350 } else {
2351 context.max_iters
2352 };
2353
2354 build_term_degrees(
2359 docs,
2360 mid,
2361 context.fwd,
2362 degree_workspaces,
2363 context.cancellation,
2364 );
2365 if context
2366 .cancellation
2367 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2368 {
2369 context
2370 .exhausted
2371 .store(true, std::sync::atomic::Ordering::Relaxed);
2372 context.progress.partition_finished();
2373 return;
2374 }
2375 let track_objective = n >= PARALLEL_BP_MIN_ENTITIES;
2380 if track_objective {
2381 degree_workspaces[0].sort_touched_words();
2385 }
2386 let mut previous_objective = if track_objective {
2387 degree_workspaces[0].bisection_objective(mid, n - mid, context.log_table)
2388 } else {
2389 0.0
2390 };
2391 let mut best_objective = previous_objective;
2392 let mut objective_stalls = 0usize;
2393
2394 for iter in 0..effective_iters {
2395 if context
2397 .cancellation
2398 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2399 {
2400 context
2401 .exhausted
2402 .store(true, std::sync::atomic::Ordering::Relaxed);
2403 break;
2404 }
2405 #[cfg(feature = "native")]
2406 if let Some(dl) = context.deadline
2407 && std::time::Instant::now() >= dl
2408 {
2409 context
2410 .exhausted
2411 .store(true, std::sync::atomic::Ordering::Relaxed);
2412 break;
2413 }
2414 compute_gains(
2417 docs,
2418 context.fwd,
2419 mid,
2420 &mut degree_workspaces[0],
2421 context.log_table,
2422 gains,
2423 allow_inner_parallelism,
2424 );
2425 if context
2426 .cancellation
2427 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2428 {
2429 context
2430 .exhausted
2431 .store(true, std::sync::atomic::Ordering::Relaxed);
2432 break;
2433 }
2434
2435 let partition = partition_by_gain(
2440 docs,
2441 gains,
2442 mid,
2443 context.fwd,
2444 &mut degree_workspaces[1..],
2445 partitioned,
2446 ranked_scratch,
2447 allow_inner_parallelism,
2448 );
2449
2450 if partition.swap_count == 0 {
2451 context.progress.iteration(n, 0, 0.0, 0.0);
2452 docs.copy_from_slice(partitioned);
2456 break;
2457 }
2458
2459 match &partition.degree_update {
2460 PartitionDegreeUpdate::Moves => {
2461 let (degrees, movement_workspaces) = degree_workspaces
2462 .split_first_mut()
2463 .expect("BP always has one degree workspace");
2464 movement_workspaces
2465 .first()
2466 .expect("parallel BP movement update requires a spare workspace")
2467 .apply_moves_to(degrees);
2468 }
2469 PartitionDegreeUpdate::Ranked => update_degrees_for_ranked_partition(
2470 docs,
2471 ranked_scratch,
2472 mid,
2473 context.fwd,
2474 &mut degree_workspaces[0],
2475 ),
2476 PartitionDegreeUpdate::Threshold {
2477 threshold_key,
2478 ties_left,
2479 } => update_degrees_for_threshold_partition(
2480 docs,
2481 gains,
2482 mid,
2483 *threshold_key,
2484 *ties_left,
2485 context.fwd,
2486 &mut degree_workspaces[0],
2487 ),
2488 }
2489
2490 let (new_objective, objective_improvement, relative_improvement) = if track_objective {
2491 let new_objective =
2492 degree_workspaces[0].bisection_objective(mid, n - mid, context.log_table);
2493 let objective_improvement = new_objective - previous_objective;
2494 let relative_improvement = objective_improvement / previous_objective.abs().max(1.0);
2495 (new_objective, objective_improvement, relative_improvement)
2496 } else {
2497 (0.0, 0.0, 0.0)
2498 };
2499 context.progress.iteration(
2500 n,
2501 partition.swap_count,
2502 objective_improvement,
2503 relative_improvement,
2504 );
2505
2506 docs.copy_from_slice(partitioned);
2513 if track_objective {
2514 previous_objective = new_objective;
2515 let relative_best_improvement =
2516 (new_objective - best_objective) / best_objective.abs().max(1.0);
2517 if relative_best_improvement >= MIN_RELATIVE_OBJECTIVE_IMPROVEMENT {
2518 best_objective = new_objective;
2519 objective_stalls = 0;
2520 } else if iter + 1 >= MIN_OBJECTIVE_ITERATIONS {
2521 objective_stalls += 1;
2522 }
2523 if objective_stalls >= OBJECTIVE_STALL_ITERATIONS {
2524 context.progress.objective_stop();
2525 break;
2526 }
2527 }
2528
2529 if iter > 2 && partition.swap_count < n / 200 {
2531 break;
2532 }
2533
2534 if !track_objective && iter > 5 {
2538 let max_abs_gain = gains
2539 .iter()
2540 .copied()
2541 .fold(0.0f32, |max_gain, gain| max_gain.max(gain.abs()));
2542 if max_abs_gain < 0.001 {
2543 break;
2544 }
2545 }
2546 }
2547
2548 context.progress.partition_finished();
2549}
2550
2551#[inline(never)]
2557fn compute_gains(
2558 docs: &[u32],
2559 fwd: &ForwardIndex,
2560 mid: usize,
2561 degrees: &mut TermDegrees,
2562 log_table: &[f32],
2563 gains: &mut [f32],
2564 allow_inner_parallelism: bool,
2565) {
2566 degrees.refresh_gain_cache(log_table);
2567 debug_assert_eq!(docs.len(), gains.len());
2568 if !degrees.gain_cache.is_empty() {
2569 let cache = degrees.gain_cache.as_slice();
2573 fill_gains(gains, allow_inner_parallelism, |i| {
2574 let side = usize::from(i >= mid);
2575 let mut gain = 0.0f32;
2576 for &term in fwd.doc_terms(docs[i] as usize) {
2577 gain += cache[term as usize][side];
2578 }
2579 gain
2580 });
2581 return;
2582 }
2583
2584 let gain_for_doc = |i: usize| -> f32 {
2593 let doc = docs[i] as usize;
2594 let in_left = i < mid;
2595 let mut g = 0.0f32;
2596 for &term in fwd.doc_terms(doc) {
2597 let [left, right] = degrees.get(term as usize);
2598 let (from, to) = if in_left {
2599 (left, right)
2600 } else {
2601 (right, left)
2602 };
2603 let move_gain = fast_log2_lookup(to as usize + 2, log_table)
2604 - fast_log2_lookup(from as usize, log_table)
2605 - std::f32::consts::LOG2_E / (1.0 + to as f32);
2606 g += if in_left { move_gain } else { -move_gain };
2607 }
2608 g
2609 };
2610
2611 fill_gains(gains, allow_inner_parallelism, gain_for_doc);
2612}
2613
2614fn fill_gains(
2615 gains: &mut [f32],
2616 allow_inner_parallelism: bool,
2617 gain_for_doc: impl Fn(usize) -> f32 + FrequencyParallelSafe,
2618) {
2619 #[cfg(not(feature = "native"))]
2620 let _ = allow_inner_parallelism;
2621
2622 #[cfg(feature = "native")]
2623 {
2624 if allow_inner_parallelism && gains.len() > 4096 {
2625 gains
2626 .par_iter_mut()
2627 .enumerate()
2628 .for_each(|(i, gain)| *gain = gain_for_doc(i));
2629 } else {
2630 for (i, gain) in gains.iter_mut().enumerate() {
2631 *gain = gain_for_doc(i);
2632 }
2633 }
2634 }
2635 #[cfg(not(feature = "native"))]
2636 {
2637 for (i, gain) in gains.iter_mut().enumerate() {
2638 *gain = gain_for_doc(i);
2639 }
2640 }
2641}
2642
2643fn build_log_table(size: usize) -> Vec<f32> {
2647 let mut table = vec![0.0f32; size];
2648 table[0] = -10.0;
2650 for (i, entry) in table.iter_mut().enumerate().skip(1) {
2651 *entry = (i as f32).log2();
2652 }
2653 table
2654}
2655
2656#[inline]
2658fn fast_log2_lookup(val: usize, table: &[f32]) -> f32 {
2659 if val < table.len() {
2660 table[val]
2661 } else {
2662 (val as f32).log2()
2663 }
2664}
2665
2666#[cfg(test)]
2667mod gain_tests;
2668
2669#[cfg(test)]
2670mod tests {
2671 use super::*;
2672
2673 #[test]
2674 fn bounded_frequency_count_and_candidate_selection_are_exact() {
2675 let items: Vec<u32> = (0..10_000).collect();
2676 let frequencies = count_frequencies_bounded(&items, 7, 16 * 1024, |item, counts| {
2677 let term = *item as usize % counts.len();
2678 counts[term] += 1;
2679 })
2680 .unwrap();
2681 assert_eq!(frequencies.iter().sum::<u32>(), items.len() as u32);
2682 for (term, &frequency) in frequencies.iter().enumerate() {
2683 let expected = (term..items.len()).step_by(frequencies.len()).count() as u32;
2684 assert_eq!(frequency, expected);
2685 }
2686
2687 let (selected, limited) =
2688 select_frequency_candidates(&[8, 3, 0, 5, 3], 1, 10, 2 * CANDIDATE_ENTRY_BYTES);
2689 assert!(limited);
2690 let mut selected = selected;
2691 selected.sort_unstable();
2692 assert_eq!(selected, vec![(1, 3), (4, 3)]);
2693 }
2694
2695 #[test]
2696 fn bounded_frequency_count_rejects_an_undersized_budget() {
2697 assert!(
2698 count_frequencies_bounded(&[0u32], 32, 32, |_, _| {}).is_none(),
2699 "one complete dense frequency table must fit before counting"
2700 );
2701 }
2702
2703 #[test]
2704 fn lazy_term_degrees_initialize_only_on_first_write() {
2705 let mut degrees = TermDegrees::new(130);
2706 let values_ptr = degrees.values.as_ptr();
2707 let bitmap_ptr = degrees.initialized.as_ptr();
2708 assert_eq!(degrees.get(65), [0, 0]);
2709 degrees.entry_mut(65)[0] += 3;
2710 degrees.entry_mut(65)[1] += 2;
2711 assert_eq!(degrees.get(65), [3, 2]);
2712 assert_eq!(degrees.get(64), [0, 0]);
2713 assert_eq!(
2714 degrees
2715 .initialized
2716 .iter()
2717 .map(|w| w.count_ones())
2718 .sum::<u32>(),
2719 1
2720 );
2721
2722 degrees.reset();
2723 assert_eq!(degrees.values.as_ptr(), values_ptr);
2724 assert_eq!(degrees.initialized.as_ptr(), bitmap_ptr);
2725 assert_eq!(degrees.get(65), [0, 0]);
2726 assert!(degrees.touched_words.is_empty());
2727 degrees.entry_mut(129)[1] = 7;
2728 assert_eq!(degrees.get(129), [0, 7]);
2729 assert_eq!(degrees.get(65), [0, 0]);
2730 }
2731
2732 #[test]
2733 fn sparse_objective_keeps_the_original_ascending_term_order() {
2734 let mut degrees = TermDegrees::new(130);
2735 *degrees.entry_mut(129) = [5, 2];
2736 *degrees.entry_mut(1) = [3, 7];
2737 *degrees.entry_mut(65) = [11, 13];
2738 assert_eq!(degrees.touched_words, vec![2, 0, 1]);
2739 degrees.sort_touched_words();
2740 assert_eq!(degrees.touched_words, vec![0, 1, 2]);
2741
2742 let log_table = build_log_table(4096);
2743 let actual = degrees.bisection_objective(31, 32, &log_table);
2744 let side_log = [
2745 fast_log2_lookup(31, &log_table) as f64,
2746 fast_log2_lookup(32, &log_table) as f64,
2747 ];
2748 let mut original_bitmap_scan = 0.0f64;
2749 for (word_idx, &initialized) in degrees.initialized.iter().enumerate() {
2750 let mut pending = initialized;
2751 while pending != 0 {
2752 let bit = pending.trailing_zeros() as usize;
2753 let term = word_idx * 64 + bit;
2754 let [left, right] = degrees.get(term);
2755 for (side, count) in [left, right].into_iter().enumerate() {
2756 if count > 0 {
2757 original_bitmap_scan += count as f64
2758 * (fast_log2_lookup(count as usize + 1, &log_table) as f64
2759 - side_log[side]);
2760 }
2761 }
2762 pending &= pending - 1;
2763 }
2764 }
2765 assert_eq!(actual.to_bits(), original_bitmap_scan.to_bits());
2766 }
2767
2768 #[test]
2769 fn gain_radix_key_matches_total_cmp() {
2770 let values = [
2771 f32::from_bits(0xffc0_0001),
2772 f32::NEG_INFINITY,
2773 -42.0,
2774 -0.0,
2775 0.0,
2776 42.0,
2777 f32::INFINITY,
2778 f32::from_bits(0x7fc0_0001),
2779 ];
2780 let mut by_cmp = values;
2781 by_cmp.sort_by(f32::total_cmp);
2782 let mut by_key = values;
2783 by_key.sort_by_key(|value| gain_order_key(*value));
2784 assert_eq!(
2785 by_cmp.map(f32::to_bits),
2786 by_key.map(f32::to_bits),
2787 "radix selection must preserve the former total_cmp order"
2788 );
2789 }
2790
2791 #[test]
2792 fn radix_threshold_matches_exact_rank_with_ties() {
2793 let gains = [3.0, -1.0, 7.0, -1.0, 0.0, -0.0, 3.0, 9.0, 3.0, 2.0, 2.0];
2794 let mut sorted: Vec<(u32, usize)> = gains
2795 .iter()
2796 .enumerate()
2797 .map(|(idx, &gain)| (gain_order_key(gain), idx))
2798 .collect();
2799 sorted.sort_unstable();
2800
2801 for left_count in 1..=gains.len() {
2802 let (threshold, lower) = select_gain_threshold(&gains, left_count, true);
2803 assert_eq!(threshold, sorted[left_count - 1].0);
2804 assert_eq!(lower, sorted.partition_point(|&(key, _)| key < threshold),);
2805 }
2806 }
2807
2808 #[cfg(feature = "native")]
2809 #[test]
2810 fn parallel_partition_matches_exact_selection_and_degree_rebuild() {
2811 const N: usize = PARALLEL_BP_MIN_ENTITIES + 1;
2812 const TERMS: usize = 101;
2813 let mut terms = Vec::with_capacity(N * 3);
2814 let mut offsets = Vec::with_capacity(N + 1);
2815 offsets.push(0);
2816 for doc in 0..N {
2817 terms.extend_from_slice(&[
2818 (doc % TERMS) as u32,
2819 ((doc / 7) % TERMS) as u32,
2820 ((doc * 13) % TERMS) as u32,
2821 ]);
2822 offsets.push(terms.len() as u64);
2823 }
2824 let fwd = ForwardIndex {
2825 terms,
2826 offsets,
2827 num_terms: TERMS,
2828 parallel_bisect_lanes: 4,
2829 cache_gains: false,
2830 budget_limited: false,
2831 };
2832 let docs: Vec<u32> = (0..N as u32)
2833 .map(|idx| ((idx as usize * 7_919) % N) as u32)
2834 .collect();
2835 let gains: Vec<f32> = docs
2836 .iter()
2837 .map(|&doc| ((doc as usize * 37) % 257) as f32 - 128.0)
2838 .collect();
2839 let mid = N / 2;
2840
2841 let mut ranked: Vec<usize> = (0..N).collect();
2842 ranked.sort_unstable_by(|&left, &right| {
2843 gains[left]
2844 .total_cmp(&gains[right])
2845 .then_with(|| left.cmp(&right))
2846 });
2847 let mut selected_left = vec![false; N];
2848 for &idx in &ranked[..mid] {
2849 selected_left[idx] = true;
2850 }
2851 let expected: Vec<u32> = docs
2852 .iter()
2853 .enumerate()
2854 .filter(|(idx, _)| selected_left[*idx])
2855 .chain(
2856 docs.iter()
2857 .enumerate()
2858 .filter(|(idx, _)| !selected_left[*idx]),
2859 )
2860 .map(|(_, &doc)| doc)
2861 .collect();
2862
2863 let mut output = vec![0; N];
2864 let mut ranked_scratch = Vec::new();
2865 let mut movement_workspaces: Vec<_> = (0..3).map(|_| TermDegrees::new(TERMS)).collect();
2866 let outcome = partition_by_gain(
2867 &docs,
2868 &gains,
2869 mid,
2870 &fwd,
2871 &mut movement_workspaces,
2872 &mut output,
2873 &mut ranked_scratch,
2874 true,
2875 );
2876 assert_eq!(output, expected);
2877 assert!(
2878 matches!(&outcome.degree_update, PartitionDegreeUpdate::Moves),
2879 "test must exercise parallel movement counts"
2880 );
2881
2882 let mut updated_workspaces: Vec<_> = (0..4).map(|_| TermDegrees::new(TERMS)).collect();
2883 build_term_degrees(&docs, mid, &fwd, &mut updated_workspaces, None);
2884 movement_workspaces[0].apply_moves_to(&mut updated_workspaces[0]);
2885
2886 let mut rebuilt_workspaces: Vec<_> = (0..4).map(|_| TermDegrees::new(TERMS)).collect();
2887 build_term_degrees(&output, mid, &fwd, &mut rebuilt_workspaces, None);
2888 for term in 0..TERMS {
2889 assert_eq!(
2890 updated_workspaces[0].get(term),
2891 rebuilt_workspaces[0].get(term),
2892 "parallel movement-count mismatch for term {term}"
2893 );
2894 }
2895 }
2896
2897 #[test]
2903 fn test_csr_offsets_do_not_wrap_past_u32() {
2904 let counts = [1_500_000_000u32; 3]; let offsets = build_csr_offsets(&counts);
2906 assert_eq!(
2907 offsets,
2908 vec![0, 1_500_000_000, 3_000_000_000, 4_500_000_000]
2909 );
2910 assert!(*offsets.last().unwrap() > u32::MAX as u64);
2911 }
2912
2913 fn make_fwd(docs: &[&[u32]], num_terms: usize) -> ForwardIndex {
2915 let mut terms = Vec::new();
2916 let mut offsets = vec![0u64];
2917 for doc_terms in docs {
2918 terms.extend_from_slice(doc_terms);
2919 offsets.push(terms.len() as u64);
2920 }
2921 ForwardIndex {
2922 terms,
2923 offsets,
2924 num_terms,
2925 parallel_bisect_lanes: 1,
2926 cache_gains: false,
2927 budget_limited: false,
2928 }
2929 }
2930
2931 #[test]
2932 fn level_partition_ranges_match_recursive_halving() {
2933 for total in 1..=129 {
2934 let mut expected: Vec<std::ops::Range<usize>> = std::iter::once(0..total).collect();
2935 for level in 0..8 {
2936 let actual: Vec<_> = (0..(1usize << level))
2937 .map(|partition_id| partition_range(total, level, partition_id))
2938 .collect();
2939 assert_eq!(actual, expected, "total={total}, level={level}");
2940
2941 expected = expected
2942 .into_iter()
2943 .flat_map(|range| {
2944 let mid = range.start + range.len() / 2;
2945 [range.start..mid, mid..range.end]
2946 })
2947 .collect();
2948 }
2949 }
2950 }
2951
2952 #[cfg(feature = "native")]
2953 #[test]
2954 fn level_synchronized_scheduler_is_thread_count_deterministic() {
2955 let docs: Vec<Vec<u32>> = (0..1_027)
2956 .map(|doc| {
2957 vec![
2958 (doc % 31) as u32,
2959 ((doc / 5) % 53) as u32,
2960 ((doc * 29 + 3) % 97) as u32,
2961 ]
2962 })
2963 .collect();
2964 let doc_refs: Vec<_> = docs.iter().map(Vec::as_slice).collect();
2965 let mut fwd = make_fwd(&doc_refs, 97);
2966 fwd.parallel_bisect_lanes = 8;
2967 let one_thread = rayon::ThreadPoolBuilder::new()
2968 .num_threads(1)
2969 .build()
2970 .unwrap();
2971 let eight_threads = rayon::ThreadPoolBuilder::new()
2972 .num_threads(8)
2973 .build()
2974 .unwrap();
2975
2976 let one = one_thread.install(|| graph_bisection(&fwd, 8, 12, BpBudget::full()).0);
2977 let eight = eight_threads.install(|| graph_bisection(&fwd, 8, 12, BpBudget::full()).0);
2978
2979 assert_eq!(eight, one);
2980 }
2981
2982 #[cfg(feature = "native")]
2986 #[test]
2987 #[ignore]
2988 fn bench_level_synchronized_scheduler() {
2989 const DOCS: usize = 500_000;
2990 const TERMS: usize = 32_768;
2991 const ROUNDS: usize = 3;
2992
2993 let mut terms = Vec::with_capacity(DOCS * 12);
2994 let mut offsets = Vec::with_capacity(DOCS + 1);
2995 offsets.push(0);
2996 for doc in 0..DOCS {
2997 let term_count = 4 + ((doc.wrapping_mul(2_654_435_761) >> 12) % 17);
3001 for term in 0..term_count {
3002 terms.push(
3003 ((doc / 64)
3004 .wrapping_mul(131)
3005 .wrapping_add(term.wrapping_mul(7_919))
3006 % TERMS) as u32,
3007 );
3008 }
3009 offsets.push(terms.len() as u64);
3010 }
3011 let fwd = ForwardIndex {
3012 terms,
3013 offsets,
3014 num_terms: TERMS,
3015 parallel_bisect_lanes: 8,
3016 cache_gains: false,
3017 budget_limited: false,
3018 };
3019 let pool = rayon::ThreadPoolBuilder::new()
3020 .num_threads(8)
3021 .build()
3022 .unwrap();
3023 let mut level_times = Vec::with_capacity(ROUNDS);
3024 let mut expected_output = None;
3025
3026 pool.install(|| {
3027 for _ in 0..ROUNDS {
3028 let started = std::time::Instant::now();
3029 let output = graph_bisection(&fwd, 32, 12, BpBudget::full()).0;
3030 level_times.push(started.elapsed());
3031 if let Some(expected) = &expected_output {
3032 assert_eq!(&output, expected);
3033 } else {
3034 expected_output = Some(output);
3035 }
3036 }
3037 });
3038
3039 level_times.sort_unstable();
3040 let level = level_times[ROUNDS / 2];
3041 println!(
3042 "BP level-synchronized scheduler median: {:.3}s",
3043 level.as_secs_f64(),
3044 );
3045 }
3046
3047 #[test]
3048 fn test_bp_empty() {
3049 let fwd = ForwardIndex {
3050 terms: Vec::new(),
3051 offsets: Vec::new(),
3052 num_terms: 0,
3053 parallel_bisect_lanes: 1,
3054 cache_gains: false,
3055 budget_limited: false,
3056 };
3057 let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
3058 assert!(perm.is_empty());
3059 }
3060
3061 #[test]
3062 fn test_bp_small() {
3063 let fwd = make_fwd(&[&[0, 1], &[0, 2], &[1, 3], &[2, 3]], 4);
3065 let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
3066 assert_eq!(perm.len(), 4);
3067 let mut sorted = perm.clone();
3069 sorted.sort();
3070 assert_eq!(sorted, vec![0, 1, 2, 3]);
3071 }
3072
3073 #[test]
3074 fn test_bp_clusters() {
3075 let fwd = make_fwd(
3079 &[
3080 &[0, 1],
3081 &[0, 1],
3082 &[0, 1],
3083 &[0, 1],
3084 &[2, 3],
3085 &[2, 3],
3086 &[2, 3],
3087 &[2, 3],
3088 ],
3089 4,
3090 );
3091 let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
3092 assert_eq!(perm.len(), 8);
3093
3094 let left: Vec<u32> = perm[..4].to_vec();
3096
3097 let a_in_left = left.iter().filter(|&&d| d < 4).count();
3099 let b_in_left = left.iter().filter(|&&d| d >= 4).count();
3100 assert!(
3101 (a_in_left == 4 && b_in_left == 0) || (a_in_left == 0 && b_in_left == 4),
3102 "Clusters should be separated: a_left={}, b_left={}",
3103 a_in_left,
3104 b_in_left,
3105 );
3106 }
3107
3108 #[test]
3109 fn test_bp_permutation_valid() {
3110 let docs: Vec<Vec<u32>> = (0..16).map(|i| vec![i / 4, 10 + i / 2]).collect();
3112 let doc_refs: Vec<&[u32]> = docs.iter().map(|v| v.as_slice()).collect();
3113 let fwd = make_fwd(&doc_refs, 18); let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
3115
3116 assert_eq!(perm.len(), 16);
3117 let mut sorted = perm.clone();
3119 sorted.sort();
3120 let expected: Vec<u32> = (0..16).collect();
3121 assert_eq!(sorted, expected);
3122 }
3123
3124 #[test]
3129 fn test_bp_depth_cap_separates_clusters_and_converges() {
3130 let fwd = make_fwd(
3133 &[
3134 &[0, 1],
3135 &[0, 1],
3136 &[0, 1],
3137 &[2, 3],
3138 &[0, 1],
3139 &[2, 3],
3140 &[2, 3],
3141 &[2, 3],
3142 ],
3143 4,
3144 );
3145 let budget = BpBudget {
3146 min_partition_docs: Some(4),
3147 time_budget: None,
3148 };
3149 let (perm, converged) = graph_bisection(&fwd, 2, 20, budget);
3150 assert!(converged, "depth cap must report converged");
3151 assert_eq!(perm.len(), 8);
3152 let mut sorted = perm.clone();
3153 sorted.sort();
3154 assert_eq!(
3155 sorted,
3156 (0..8).collect::<Vec<u32>>(),
3157 "must stay a valid permutation"
3158 );
3159 let cluster_a = [0u32, 1, 2, 4];
3162 let a_in_left = perm[..4].iter().filter(|d| cluster_a.contains(d)).count();
3163 assert!(
3164 a_in_left == 4 || a_in_left == 0,
3165 "clusters should separate at the top level: {:?}",
3166 perm
3167 );
3168 }
3169
3170 #[test]
3173 fn test_bp_zero_time_budget_emits_valid_partial_permutation() {
3174 let docs: Vec<Vec<u32>> = (0..64).map(|i| vec![i % 4]).collect();
3175 let doc_refs: Vec<&[u32]> = docs.iter().map(|v| v.as_slice()).collect();
3176 let fwd = make_fwd(&doc_refs, 4);
3177 let budget = BpBudget {
3178 min_partition_docs: None,
3179 time_budget: Some(std::time::Duration::ZERO),
3180 };
3181 let (perm, converged) = graph_bisection(&fwd, 4, 20, budget);
3182 assert!(!converged, "zero budget must report unconverged");
3183 assert_eq!(perm.len(), 64);
3184 let mut sorted = perm.clone();
3185 sorted.sort();
3186 assert_eq!(sorted, (0..64).collect::<Vec<u32>>());
3187 }
3188
3189 #[test]
3190 fn test_bp_shutdown_cancellation_emits_valid_partial_permutation() {
3191 let docs: Vec<Vec<u32>> = (0..64).map(|i| vec![i % 4]).collect();
3192 let doc_refs: Vec<&[u32]> = docs.iter().map(Vec::as_slice).collect();
3193 let fwd = make_fwd(&doc_refs, 4);
3194 let cancellation = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
3195 let budget = BpBudget {
3196 min_partition_docs: None,
3197 time_budget: None,
3198 };
3199
3200 let (perm, converged) = graph_bisection_with_progress(
3201 &fwd,
3202 4,
3203 20,
3204 budget,
3205 Some(cancellation.as_ref()),
3206 BpProgressLabel::anonymous(),
3207 );
3208
3209 assert!(!converged, "cancelled BP must report unconverged");
3210 assert_eq!(perm, (0..64).collect::<Vec<u32>>());
3211 }
3212
3213 #[test]
3214 fn test_memory_limited_graph_never_reports_converged() {
3215 let mut fwd = make_fwd(&[&[0], &[0], &[1], &[1]], 2);
3216 fwd.budget_limited = true;
3217
3218 let (perm, converged) = graph_bisection(&fwd, 2, 20, BpBudget::full());
3219
3220 assert!(!converged);
3221 let mut sorted = perm;
3222 sorted.sort_unstable();
3223 assert_eq!(sorted, vec![0, 1, 2, 3]);
3224 }
3225
3226 #[test]
3227 fn test_fast_log2() {
3228 let table = build_log_table(4096);
3229 assert!((table[1] - 0.0).abs() < 0.001);
3230 assert!((table[2] - 1.0).abs() < 0.001);
3231 assert!((table[4] - 2.0).abs() < 0.001);
3232 assert!((table[1024] - 10.0).abs() < 0.001);
3233 let val = fast_log2_lookup(8192, &table);
3235 assert!((val - 13.0).abs() < 0.001);
3236 }
3237}