1#[cfg(feature = "native")]
13use rayon::prelude::*;
14const TERM_DEGREE_VALUE_BYTES: usize = std::mem::size_of::<[u32; 2]>();
15const CANDIDATE_ENTRY_BYTES: usize = std::mem::size_of::<(usize, u32)>();
16const PARALLEL_BP_MIN_ENTITIES: usize = 1_048_576;
20const MIN_RELATIVE_OBJECTIVE_IMPROVEMENT: f64 = 1e-6;
21const MIN_OBJECTIVE_ITERATIONS: usize = 4;
22const OBJECTIVE_STALL_ITERATIONS: usize = 2;
23
24fn term_degree_bytes(num_terms: usize) -> usize {
25 let bitmap_words = num_terms.div_ceil(64);
26 num_terms
27 .saturating_mul(TERM_DEGREE_VALUE_BYTES)
28 .saturating_add(bitmap_words.saturating_mul(std::mem::size_of::<u64>()))
29 .saturating_add(bitmap_words.saturating_mul(std::mem::size_of::<u32>()))
33}
34
35#[cfg(feature = "native")]
40trait FrequencyParallelSafe: Sync {}
41#[cfg(feature = "native")]
42impl<T: Sync + ?Sized> FrequencyParallelSafe for T {}
43
44#[cfg(not(feature = "native"))]
45trait FrequencyParallelSafe {}
46#[cfg(not(feature = "native"))]
47impl<T: ?Sized> FrequencyParallelSafe for T {}
48
49fn count_frequencies_bounded<T: FrequencyParallelSafe>(
50 items: &[T],
51 num_terms: usize,
52 available_bytes: usize,
53 count_item: impl Fn(&T, &mut [u32]) + FrequencyParallelSafe,
54) -> Option<Vec<u32>> {
55 if num_terms == 0 {
56 return Some(Vec::new());
57 }
58 let table_bytes = num_terms
59 .checked_mul(std::mem::size_of::<u32>())?
60 .checked_add(std::mem::size_of::<Vec<u32>>())?;
61 let affordable_tables = available_bytes.checked_div(table_bytes)?;
62 if affordable_tables == 0 {
63 return None;
64 }
65
66 #[cfg(feature = "native")]
67 {
68 let lanes = affordable_tables
69 .min(rayon::current_num_threads().max(1))
70 .min(items.len().max(1));
71 if lanes > 1 {
72 let chunk_len = items.len().div_ceil(lanes);
73 return items
74 .par_chunks(chunk_len)
75 .map(|chunk| {
76 let mut counts = vec![0u32; num_terms];
77 for item in chunk {
78 count_item(item, &mut counts);
79 }
80 counts
81 })
82 .reduce_with(|mut left, right| {
83 for (total, count) in left.iter_mut().zip(right) {
84 *total = total.saturating_add(count);
85 }
86 left
87 });
88 }
89 }
90
91 let mut counts = vec![0u32; num_terms];
92 for item in items {
93 count_item(item, &mut counts);
94 }
95 Some(counts)
96}
97
98fn select_frequency_candidates(
101 frequencies: &[u32],
102 min_frequency: usize,
103 max_frequency: usize,
104 candidate_budget_bytes: usize,
105) -> (Vec<(u32, usize)>, bool) {
106 let eligible_count = frequencies
107 .iter()
108 .filter(|&&frequency| {
109 let frequency = frequency as usize;
110 frequency >= min_frequency && frequency <= max_frequency
111 })
112 .count();
113 let capacity = candidate_budget_bytes
114 .checked_div(CANDIDATE_ENTRY_BYTES)
115 .unwrap_or(0)
116 .min(eligible_count);
117 let mut candidates = std::collections::BinaryHeap::with_capacity(capacity);
118 for (term_id, &frequency) in frequencies.iter().enumerate() {
119 let frequency = frequency as usize;
120 if frequency < min_frequency || frequency > max_frequency {
121 continue;
122 }
123 let candidate = (frequency, term_id as u32);
124 if candidates.len() < capacity {
125 candidates.push(candidate);
126 } else if capacity > 0 && candidate < *candidates.peek().unwrap() {
127 candidates.pop();
128 candidates.push(candidate);
129 }
130 }
131 let selected = candidates
132 .into_vec()
133 .into_iter()
134 .map(|(frequency, term_id)| (term_id, frequency))
135 .collect();
136 (selected, capacity < eligible_count)
137}
138
139struct CandidateFit {
140 estimated_bytes: usize,
141 retained_postings: usize,
142 dropped: usize,
143}
144
145fn fit_candidates_to_budget(
149 candidates: &mut Vec<(u32, usize)>,
150 fixed_bytes: usize,
151 memory_budget_bytes: usize,
152) -> CandidateFit {
153 let total_postings = candidates.iter().fold(0usize, |total, (_, frequency)| {
154 total.saturating_add(*frequency)
155 });
156 let estimated_bytes = total_postings
157 .saturating_mul(std::mem::size_of::<u32>())
158 .saturating_add(fixed_bytes)
159 .saturating_add(candidates.len().saturating_mul(CANDIDATE_ENTRY_BYTES))
160 .saturating_add(term_degree_bytes(candidates.len()));
161 if estimated_bytes <= memory_budget_bytes || candidates.is_empty() {
162 return CandidateFit {
163 estimated_bytes,
164 retained_postings: total_postings,
165 dropped: 0,
166 };
167 }
168
169 candidates.sort_by_key(|&(_, frequency)| frequency);
170 let mut used_bytes = fixed_bytes;
171 let mut retained_postings = 0usize;
172 let mut keep = 0usize;
173 for &(_, frequency) in candidates.iter() {
174 let term_bytes = frequency
175 .saturating_mul(std::mem::size_of::<u32>())
176 .saturating_add(TERM_DEGREE_VALUE_BYTES + 1)
177 .saturating_add(CANDIDATE_ENTRY_BYTES);
178 if term_bytes > memory_budget_bytes.saturating_sub(used_bytes) {
179 break;
180 }
181 used_bytes = used_bytes.saturating_add(term_bytes);
182 retained_postings = retained_postings.saturating_add(frequency);
183 keep += 1;
184 }
185 let dropped = candidates.len() - keep;
186 candidates.truncate(keep);
187 candidates.shrink_to_fit();
190 CandidateFit {
191 estimated_bytes,
192 retained_postings,
193 dropped,
194 }
195}
196
197fn parallel_bisect_lanes(
198 memory_budget_bytes: usize,
199 non_degree_bytes: usize,
200 num_terms: usize,
201) -> usize {
202 let per_node = term_degree_bytes(num_terms).max(1);
203 let affordable_nodes = memory_budget_bytes
204 .saturating_sub(non_degree_bytes)
205 .checked_div(per_node)
206 .unwrap_or(0)
207 .max(1);
208 #[cfg(feature = "native")]
209 let worker_limit = rayon::current_num_threads().max(1);
210 #[cfg(not(feature = "native"))]
211 let worker_limit = 1usize;
212 affordable_nodes.min(worker_limit)
213}
214
215struct TermDegrees {
223 values: Vec<std::mem::MaybeUninit<[u32; 2]>>,
224 initialized: Vec<u64>,
225 touched_words: Vec<u32>,
226}
227
228impl TermDegrees {
229 fn new(num_terms: usize) -> Self {
230 let bitmap_words = num_terms.div_ceil(64);
231 let mut values = Vec::with_capacity(num_terms);
232 values.resize_with(num_terms, std::mem::MaybeUninit::uninit);
233 Self {
234 values,
235 initialized: vec![0; bitmap_words],
236 touched_words: Vec::with_capacity(bitmap_words),
237 }
238 }
239
240 fn reset(&mut self) {
243 for word in self.touched_words.drain(..) {
244 self.initialized[word as usize] = 0;
245 }
246 }
247
248 fn sort_touched_words(&mut self) {
249 self.touched_words.sort_unstable();
250 }
251
252 #[inline]
253 fn entry_mut(&mut self, term: usize) -> &mut [u32; 2] {
254 let word = term / 64;
255 let mask = 1u64 << (term % 64);
256 if self.initialized[word] & mask == 0 {
257 if self.initialized[word] == 0 {
258 self.touched_words.push(word as u32);
259 }
260 self.values[term].write([0, 0]);
261 self.initialized[word] |= mask;
262 }
263 unsafe { self.values[term].assume_init_mut() }
265 }
266
267 #[inline]
268 fn get(&self, term: usize) -> [u32; 2] {
269 let word = term / 64;
270 let mask = 1u64 << (term % 64);
271 if self.initialized[word] & mask == 0 {
272 return [0, 0];
273 }
274 unsafe { *self.values[term].assume_init_ref() }
277 }
278
279 fn merge_from(&mut self, other: &Self) {
280 for &word_idx in &other.touched_words {
281 let word_idx = word_idx as usize;
282 let mut pending = other.initialized[word_idx];
283 while pending != 0 {
284 let bit = pending.trailing_zeros() as usize;
285 let term = word_idx * 64 + bit;
286 let [left, right] = unsafe { *other.values[term].assume_init_ref() };
288 let entry = self.entry_mut(term);
289 entry[0] += left;
290 entry[1] += right;
291 pending &= pending - 1;
292 }
293 }
294 }
295
296 fn apply_moves_to(&self, degrees: &mut Self) {
302 for &word_idx in &self.touched_words {
303 let word_idx = word_idx as usize;
304 let mut pending = self.initialized[word_idx];
305 while pending != 0 {
306 let bit = pending.trailing_zeros() as usize;
307 let term = word_idx * 64 + bit;
308 let [right_to_left, left_to_right] =
310 unsafe { *self.values[term].assume_init_ref() };
311 debug_assert!(
312 degrees.initialized[word_idx] & (1u64 << bit) != 0,
313 "a moved term must already exist in the partition degrees"
314 );
315 let degree = degrees.entry_mut(term);
316 let new_left =
317 i64::from(degree[0]) + i64::from(right_to_left) - i64::from(left_to_right);
318 let new_right =
319 i64::from(degree[1]) + i64::from(left_to_right) - i64::from(right_to_left);
320 debug_assert!(new_left >= 0 && new_right >= 0);
321 debug_assert!(new_left <= i64::from(u32::MAX));
322 debug_assert!(new_right <= i64::from(u32::MAX));
323 degree[0] = new_left as u32;
324 degree[1] = new_right as u32;
325 pending &= pending - 1;
326 }
327 }
328 }
329
330 fn bisection_objective(&self, left_size: usize, right_size: usize, log_table: &[f32]) -> f64 {
337 let mut objective = 0.0f64;
338 let side_log = [
339 fast_log2_lookup(left_size, log_table) as f64,
340 fast_log2_lookup(right_size, log_table) as f64,
341 ];
342 for &word_idx in &self.touched_words {
343 let word_idx = word_idx as usize;
344 let mut pending = self.initialized[word_idx];
345 while pending != 0 {
346 let bit = pending.trailing_zeros() as usize;
347 let term = word_idx * 64 + bit;
348 let [left, right] = unsafe { *self.values[term].assume_init_ref() };
350 for (side, count) in [left, right].into_iter().enumerate() {
351 if count > 0 {
352 objective += count as f64
353 * (fast_log2_lookup(count as usize + 1, log_table) as f64
354 - side_log[side]);
355 }
356 }
357 pending &= pending - 1;
358 }
359 }
360 objective
361 }
362}
363
364pub(crate) struct ForwardIndex {
370 terms: Vec<u32>,
371 offsets: Vec<u64>,
376 pub num_terms: usize,
377 parallel_bisect_lanes: usize,
381 budget_limited: bool,
385}
386
387fn build_csr_offsets(counts: &[u32]) -> Vec<u64> {
390 let mut offsets = Vec::with_capacity(counts.len() + 1);
391 offsets.push(0u64);
392 for &c in counts {
393 offsets.push(offsets.last().unwrap() + c as u64);
394 }
395 offsets
396}
397
398impl ForwardIndex {
399 #[inline]
400 pub fn num_docs(&self) -> usize {
401 if self.offsets.is_empty() {
402 0
403 } else {
404 self.offsets.len() - 1
405 }
406 }
407
408 #[inline]
409 fn doc_terms(&self, doc: usize) -> &[u32] {
410 let start = self.offsets[doc] as usize;
411 let end = self.offsets[doc + 1] as usize;
412 &self.terms[start..end]
413 }
414
415 pub fn total_postings(&self) -> u64 {
417 self.offsets.last().copied().unwrap_or(0)
418 }
419
420 #[inline]
421 pub fn budget_limited(&self) -> bool {
422 self.budget_limited
423 }
424}
425
426pub(crate) fn build_vid_maps(
436 bmp: &crate::segment::reader::bmp::BmpIndex,
437) -> crate::Result<(Vec<u32>, Vec<u32>)> {
438 let ids = bmp.doc_map_ids_slice();
439 let num_virtual = bmp.num_virtual_docs as usize;
440 let expected_real = bmp.num_real_docs() as usize;
441 let mut virtual_to_real = vec![u32::MAX; num_virtual];
442 let mut real_to_virtual = Vec::with_capacity(expected_real);
443 debug_assert_eq!(ids.len(), virtual_to_real.len() * 4);
444 bmp.visit_real_slots_for_rewrite(|vid| {
445 virtual_to_real[vid] = real_to_virtual.len() as u32;
446 real_to_virtual.push(vid as u32);
447 })?;
448 Ok((virtual_to_real, real_to_virtual))
449}
450
451struct BlockJob {
457 src: u32,
458 block_id: u32,
459 real_start: u32,
461 real_len: u32,
463}
464
465fn build_block_jobs(
468 bmps: &[&crate::segment::reader::bmp::BmpIndex],
469 vid_maps: &[(Vec<u32>, Vec<u32>)],
470) -> Vec<BlockJob> {
471 let total_blocks: usize = bmps.iter().map(|b| b.num_blocks as usize).sum();
472 let mut jobs = Vec::with_capacity(total_blocks);
473 for (src, (bmp, (v2r, _))) in bmps.iter().zip(vid_maps).enumerate() {
474 let block_size = bmp.bmp_block_size as usize;
475 let mut real_cursor = 0u32;
476 for block_id in 0..bmp.num_blocks as usize {
477 let vid_start = block_id * block_size;
478 let vid_end = ((block_id + 1) * block_size).min(v2r.len());
479 let real_len = v2r[vid_start..vid_end]
480 .iter()
481 .filter(|&&r| r != u32::MAX)
482 .count() as u32;
483 jobs.push(BlockJob {
484 src: src as u32,
485 block_id: block_id as u32,
486 real_start: real_cursor,
487 real_len,
488 });
489 real_cursor += real_len;
490 }
491 }
492 jobs
493}
494
495#[cfg(test)]
509pub(crate) fn build_forward_index_from_bmps(
510 bmps: &[&crate::segment::reader::bmp::BmpIndex],
511 min_doc_freq: usize,
512 max_doc_freq: usize,
513 memory_budget_bytes: usize,
514) -> crate::Result<ForwardIndex> {
515 let vid_maps: Vec<(Vec<u32>, Vec<u32>)> = bmps
516 .iter()
517 .map(|bmp| build_vid_maps(bmp))
518 .collect::<crate::Result<_>>()?;
519 Ok(build_forward_index_from_bmps_with_maps(
520 bmps,
521 &vid_maps,
522 min_doc_freq,
523 max_doc_freq,
524 memory_budget_bytes,
525 ))
526}
527
528pub(crate) fn build_forward_index_from_bmps_with_maps(
532 bmps: &[&crate::segment::reader::bmp::BmpIndex],
533 vid_maps: &[(Vec<u32>, Vec<u32>)],
534 min_doc_freq: usize,
535 max_doc_freq: usize,
536 memory_budget_bytes: usize,
537) -> ForwardIndex {
538 debug_assert_eq!(bmps.len(), vid_maps.len());
539 let total_docs: usize = vid_maps.iter().map(|(_, r2v)| r2v.len()).sum();
540
541 if total_docs == 0 {
542 return ForwardIndex {
543 terms: Vec::new(),
544 offsets: Vec::new(),
545 num_terms: 0,
546 parallel_bisect_lanes: 1,
547 budget_limited: false,
548 };
549 }
550
551 let jobs = build_block_jobs(bmps, vid_maps);
556
557 let max_dims = bmps
561 .iter()
562 .map(|bmp| bmp.dims() as usize)
563 .max()
564 .unwrap_or(0);
565 let jobs_bytes = jobs
566 .len()
567 .saturating_mul(std::mem::size_of::<BlockJob>().saturating_add(40));
568 let frequency_bytes = max_dims.saturating_mul(std::mem::size_of::<u32>());
569 if frequency_bytes > memory_budget_bytes.saturating_sub(jobs_bytes) {
570 log::warn!(
571 "[reorder] memory budget {} cannot hold the {} dimension-frequency table; using identity order",
572 crate::format_bytes(memory_budget_bytes as u64),
573 crate::format_bytes(frequency_bytes as u64),
574 );
575 return ForwardIndex {
576 terms: Vec::new(),
577 offsets: Vec::new(),
578 num_terms: 0,
579 parallel_bisect_lanes: 1,
580 budget_limited: true,
581 };
582 }
583 let Some(dim_df) = count_frequencies_bounded(
584 &jobs,
585 max_dims,
586 memory_budget_bytes.saturating_sub(jobs_bytes),
587 |job, counts| {
588 let bmp = bmps[job.src as usize];
589 let (v2r, _) = &vid_maps[job.src as usize];
590 let block_size = bmp.bmp_block_size as usize;
591 for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
592 let mut frequency = 0u32;
593 for posting in postings {
594 let vid = job.block_id as usize * block_size + posting.local_slot as usize;
595 if v2r.get(vid).is_some_and(|&real| real != u32::MAX) && posting.impact > 0 {
596 frequency = frequency.saturating_add(1);
597 }
598 }
599 if frequency > 0
600 && let Some(total) = counts.get_mut(dim_id as usize)
601 {
602 *total = total.saturating_add(frequency);
603 }
604 }
605 },
606 ) else {
607 log::warn!(
608 "[reorder] memory budget {} cannot hold a bounded dimension-frequency table; using identity order",
609 crate::format_bytes(memory_budget_bytes as u64),
610 );
611 return ForwardIndex {
612 terms: Vec::new(),
613 offsets: Vec::new(),
614 num_terms: 0,
615 parallel_bisect_lanes: 1,
616 budget_limited: true,
617 };
618 };
619
620 let (mut eligible, mut budget_limited) = select_frequency_candidates(
624 &dim_df,
625 min_doc_freq,
626 max_doc_freq,
627 memory_budget_bytes
628 .saturating_sub(jobs_bytes)
629 .saturating_sub(frequency_bytes),
630 );
631 drop(dim_df);
632
633 let entity_scratch_bytes = total_docs.saturating_mul(32);
637 let remap_bytes = max_dims.saturating_mul(4);
638 let fixed_bytes = entity_scratch_bytes
639 .saturating_add(remap_bytes)
640 .saturating_add(jobs_bytes);
641 let fit = fit_candidates_to_budget(&mut eligible, fixed_bytes, memory_budget_bytes);
642 if fit.dropped > 0 {
643 budget_limited = true;
644 log::warn!(
645 "[reorder] memory budget {}: estimated {}, dropped {} highest-df dims, keeping {} ({} postings)",
646 crate::format_bytes(memory_budget_bytes as u64),
647 crate::format_bytes(fit.estimated_bytes as u64),
648 fit.dropped,
649 eligible.len(),
650 fit.retained_postings,
651 );
652 }
653
654 if eligible.is_empty() {
655 return ForwardIndex {
660 terms: Vec::new(),
661 offsets: Vec::new(),
662 num_terms: 0,
663 parallel_bisect_lanes: 1,
664 budget_limited,
665 };
666 }
667
668 let mut term_remap = vec![u32::MAX; max_dims];
669 for (compact_id, &(dim_id, _)) in eligible.iter().enumerate() {
670 term_remap[dim_id as usize] = compact_id as u32;
671 }
672 let num_active_terms = eligible.len();
673 let non_degree_bytes = entity_scratch_bytes.saturating_add(
678 fit.retained_postings
679 .saturating_mul(std::mem::size_of::<u32>()),
680 );
681 let parallel_bisect_lanes =
682 parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_active_terms);
683 drop(eligible);
684
685 let mut counts = vec![0u32; total_docs];
687 let fill_block_counts = |job: &BlockJob, out: &mut [u32]| {
688 let bmp = bmps[job.src as usize];
689 let (v2r, _) = &vid_maps[job.src as usize];
690 let block_size = bmp.bmp_block_size as usize;
691 for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
692 if term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX) == u32::MAX {
693 continue;
694 }
695 for p in postings {
696 let vid = job.block_id as usize * block_size + p.local_slot as usize;
697 let Some(&real) = v2r.get(vid) else {
698 continue;
699 };
700 if real != u32::MAX && p.impact > 0 {
701 out[(real - job.real_start) as usize] += 1;
702 }
703 }
704 }
705 };
706 {
707 let mut slices: Vec<(&BlockJob, &mut [u32])> = Vec::with_capacity(jobs.len());
708 let mut rest: &mut [u32] = &mut counts;
709 for job in &jobs {
710 let (head, tail) = rest.split_at_mut(job.real_len as usize);
711 slices.push((job, head));
712 rest = tail;
713 }
714 #[cfg(feature = "native")]
715 slices
716 .into_par_iter()
717 .for_each(|(job, out)| fill_block_counts(job, out));
718 #[cfg(not(feature = "native"))]
719 for (job, out) in slices {
720 fill_block_counts(job, out);
721 }
722 }
723
724 let offsets = build_csr_offsets(&counts);
726 let total = *offsets.last().unwrap() as usize;
727 drop(counts);
728
729 let mut terms = vec![0u32; total];
732 let fill_block_terms = |job: &BlockJob, global_real_start: usize, out: &mut [u32]| {
733 let bmp = bmps[job.src as usize];
734 let (v2r, _) = &vid_maps[job.src as usize];
735 let block_size = bmp.bmp_block_size as usize;
736 assert!(job.real_len as usize <= 256, "BMP block exceeds 256 docs");
738 let mut cursor = [0u32; 256];
739 let base = offsets[global_real_start] as usize;
740 for (dim_id, _, postings) in bmp.iter_block_terms(job.block_id) {
741 let compact = term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX);
742 if compact == u32::MAX {
743 continue;
744 }
745 for p in postings {
746 let vid = job.block_id as usize * block_size + p.local_slot as usize;
747 let Some(&real) = v2r.get(vid) else {
748 continue;
749 };
750 if real != u32::MAX && p.impact > 0 {
751 let local = (real - job.real_start) as usize;
752 let pos =
753 offsets[global_real_start + local] as usize - base + cursor[local] as usize;
754 out[pos] = compact;
755 cursor[local] += 1;
756 }
757 }
758 }
759 };
760 {
761 let mut slices: Vec<(&BlockJob, usize, &mut [u32])> = Vec::with_capacity(jobs.len());
762 let mut rest: &mut [u32] = &mut terms;
763 let mut global_real = 0usize;
764 for job in &jobs {
765 let len =
766 (offsets[global_real + job.real_len as usize] - offsets[global_real]) as usize;
767 let (head, tail) = rest.split_at_mut(len);
768 slices.push((job, global_real, head));
769 rest = tail;
770 global_real += job.real_len as usize;
771 }
772 #[cfg(feature = "native")]
773 slices
774 .into_par_iter()
775 .for_each(|(job, g, out)| fill_block_terms(job, g, out));
776 #[cfg(not(feature = "native"))]
777 for (job, g, out) in slices {
778 fill_block_terms(job, g, out);
779 }
780 }
781
782 ForwardIndex {
783 terms,
784 offsets,
785 num_terms: num_active_terms,
786 parallel_bisect_lanes,
787 budget_limited,
788 }
789}
790
791pub(crate) fn build_forward_index_from_blocks(
799 bmps: &[&crate::segment::reader::bmp::BmpIndex],
800 memory_budget_bytes: usize,
801) -> ForwardIndex {
802 let total_blocks: usize = bmps.iter().map(|b| b.num_blocks as usize).sum();
803 if total_blocks == 0 {
804 return ForwardIndex {
805 terms: Vec::new(),
806 offsets: Vec::new(),
807 num_terms: 0,
808 parallel_bisect_lanes: 1,
809 budget_limited: false,
810 };
811 }
812
813 let blocks: Vec<(u32, u32)> = bmps
815 .iter()
816 .enumerate()
817 .flat_map(|(src, bmp)| (0..bmp.num_blocks).map(move |b| (src as u32, b)))
818 .collect();
819
820 let max_dims = bmps
823 .iter()
824 .map(|bmp| bmp.dims() as usize)
825 .max()
826 .unwrap_or(0);
827 let blocks_bytes = blocks
828 .len()
829 .saturating_mul(std::mem::size_of::<(u32, u32)>().saturating_add(32));
830 let frequency_bytes = max_dims.saturating_mul(std::mem::size_of::<u32>());
831 if frequency_bytes > memory_budget_bytes.saturating_sub(blocks_bytes) {
832 log::warn!(
833 "[reorder] block-level frequency table exceeds memory budget; using identity order"
834 );
835 return ForwardIndex {
836 terms: Vec::new(),
837 offsets: Vec::new(),
838 num_terms: 0,
839 parallel_bisect_lanes: 1,
840 budget_limited: true,
841 };
842 }
843 let Some(dim_bf) = count_frequencies_bounded(
844 &blocks,
845 max_dims,
846 memory_budget_bytes.saturating_sub(blocks_bytes),
847 |&(src, block_id), counts| {
848 for (dim_id, _, _) in bmps[src as usize].iter_block_terms(block_id) {
849 if let Some(count) = counts.get_mut(dim_id as usize) {
850 *count = count.saturating_add(1);
851 }
852 }
853 },
854 ) else {
855 log::warn!(
856 "[reorder] block-level frequency table cannot fit its bounded allocation; using identity order"
857 );
858 return ForwardIndex {
859 terms: Vec::new(),
860 offsets: Vec::new(),
861 num_terms: 0,
862 parallel_bisect_lanes: 1,
863 budget_limited: true,
864 };
865 };
866
867 let max_bf = (total_blocks as f64 * 0.9) as usize;
868 let (mut eligible, mut budget_limited) = select_frequency_candidates(
869 &dim_bf,
870 2,
871 max_bf.max(2),
872 memory_budget_bytes
873 .saturating_sub(blocks_bytes)
874 .saturating_sub(frequency_bytes),
875 );
876 drop(dim_bf);
877
878 let entity_scratch_bytes = total_blocks.saturating_mul(32);
879 let remap_bytes = max_dims.saturating_mul(4);
880 let fixed_bytes = entity_scratch_bytes
881 .saturating_add(remap_bytes)
882 .saturating_add(blocks_bytes);
883 let fit = fit_candidates_to_budget(&mut eligible, fixed_bytes, memory_budget_bytes);
884 if fit.dropped > 0 {
885 budget_limited = true;
886 log::warn!(
887 "[reorder] block-level fwd index over budget — dropped {} highest-bf dims",
888 fit.dropped,
889 );
890 }
891
892 if eligible.is_empty() {
893 return ForwardIndex {
894 terms: Vec::new(),
895 offsets: Vec::new(),
896 num_terms: 0,
897 parallel_bisect_lanes: 1,
898 budget_limited,
899 };
900 }
901
902 let mut term_remap = vec![u32::MAX; max_dims];
903 for (compact, &(dim_id, _)) in eligible.iter().enumerate() {
904 term_remap[dim_id as usize] = compact as u32;
905 }
906 let num_terms = eligible.len();
907 let non_degree_bytes = entity_scratch_bytes.saturating_add(
908 fit.retained_postings
909 .saturating_mul(std::mem::size_of::<u32>()),
910 );
911 let parallel_bisect_lanes =
912 parallel_bisect_lanes(memory_budget_bytes, non_degree_bytes, num_terms);
913 drop(eligible);
914
915 let count_remapped = |&(src, block_id): &(u32, u32)| -> u32 {
918 bmps[src as usize]
919 .iter_block_terms(block_id)
920 .filter(|(dim_id, _, _)| {
921 term_remap
922 .get(*dim_id as usize)
923 .copied()
924 .unwrap_or(u32::MAX)
925 != u32::MAX
926 })
927 .count() as u32
928 };
929 #[cfg(feature = "native")]
930 let counts: Vec<u32> = blocks.par_iter().map(count_remapped).collect();
931 #[cfg(not(feature = "native"))]
932 let counts: Vec<u32> = blocks.iter().map(count_remapped).collect();
933
934 let offsets = build_csr_offsets(&counts);
935 let total = *offsets.last().unwrap() as usize;
936 drop(counts);
937
938 let mut terms = vec![0u32; total];
939 let fill_block = |&(src, block_id): &(u32, u32), out: &mut [u32]| {
940 let mut n = 0usize;
941 for (dim_id, _, _) in bmps[src as usize].iter_block_terms(block_id) {
942 let compact = term_remap.get(dim_id as usize).copied().unwrap_or(u32::MAX);
943 if compact != u32::MAX {
944 out[n] = compact;
945 n += 1;
946 }
947 }
948 };
949 {
950 let mut slices: Vec<(&(u32, u32), &mut [u32])> = Vec::with_capacity(blocks.len());
951 let mut rest: &mut [u32] = &mut terms;
952 for (gb, b) in blocks.iter().enumerate() {
953 let len = (offsets[gb + 1] - offsets[gb]) as usize;
954 let (head, tail) = rest.split_at_mut(len);
955 slices.push((b, head));
956 rest = tail;
957 }
958 #[cfg(feature = "native")]
959 slices
960 .into_par_iter()
961 .for_each(|(b, out)| fill_block(b, out));
962 #[cfg(not(feature = "native"))]
963 for (b, out) in slices {
964 fill_block(b, out);
965 }
966 }
967
968 ForwardIndex {
969 terms,
970 offsets,
971 num_terms,
972 parallel_bisect_lanes,
973 budget_limited,
974 }
975}
976
977#[derive(Clone, Copy, Debug, Default)]
985pub struct BpBudget {
986 pub min_partition_docs: Option<usize>,
991 pub time_budget: Option<std::time::Duration>,
995}
996
997impl BpBudget {
998 pub fn full() -> Self {
1000 Self::default()
1001 }
1002}
1003
1004fn build_term_degrees(
1011 docs: &[u32],
1012 mid: usize,
1013 fwd: &ForwardIndex,
1014 workspaces: &mut [TermDegrees],
1015 cancellation: Option<&std::sync::atomic::AtomicBool>,
1016) {
1017 debug_assert!(!workspaces.is_empty());
1018 let build_range = |degrees: &mut TermDegrees, start: usize, chunk: &[u32]| {
1019 degrees.reset();
1020 for (offset, &doc) in chunk.iter().enumerate() {
1021 if offset.is_multiple_of(1024)
1022 && cancellation
1023 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
1024 {
1025 break;
1026 }
1027 let side = usize::from(start + offset >= mid);
1028 for &term in fwd.doc_terms(doc as usize) {
1029 degrees.entry_mut(term as usize)[side] += 1;
1030 }
1031 }
1032 };
1033
1034 #[cfg(feature = "native")]
1035 {
1036 let workers = workspaces
1037 .len()
1038 .min(docs.len().div_ceil(PARALLEL_BP_MIN_ENTITIES).max(1));
1039 if workers > 1 {
1040 let chunk_len = docs.len().div_ceil(workers);
1041 workspaces[..workers]
1042 .par_iter_mut()
1043 .enumerate()
1044 .for_each(|(worker, degrees)| {
1045 let start = worker * chunk_len;
1046 let end = (start + chunk_len).min(docs.len());
1047 build_range(degrees, start, &docs[start..end]);
1048 });
1049 let (degrees, partials) = workspaces[..workers]
1050 .split_first_mut()
1051 .expect("at least one BP degree workspace");
1052 for partial in partials {
1053 degrees.merge_from(partial);
1054 }
1055 return;
1056 }
1057 build_range(&mut workspaces[0], 0, docs);
1058 }
1059
1060 #[cfg(not(feature = "native"))]
1061 build_range(&mut workspaces[0], 0, docs);
1062}
1063
1064#[inline]
1070fn gain_order_key(gain: f32) -> u32 {
1071 let bits = gain.to_bits();
1072 if bits & 0x8000_0000 != 0 {
1073 !bits
1074 } else {
1075 bits ^ 0x8000_0000
1076 }
1077}
1078
1079fn select_gain_threshold(gains: &[f32], left_count: usize) -> (u32, usize) {
1083 debug_assert!(left_count > 0 && left_count <= gains.len());
1084 let mut rank_within_prefix = left_count - 1;
1085 let mut strictly_lower = 0usize;
1086 let mut prefix = 0u32;
1087 let mut prefix_mask = 0u32;
1088
1089 for shift in [24u32, 16, 8, 0] {
1090 let histogram = {
1091 #[cfg(feature = "native")]
1092 {
1093 if gains.len() >= PARALLEL_BP_MIN_ENTITIES {
1094 gains
1095 .par_iter()
1096 .fold(
1097 || Box::new([0usize; 256]),
1098 |mut counts, &gain| {
1099 let key = gain_order_key(gain);
1100 if key & prefix_mask == prefix {
1101 counts[((key >> shift) & 0xff) as usize] += 1;
1102 }
1103 counts
1104 },
1105 )
1106 .reduce(
1107 || Box::new([0usize; 256]),
1108 |mut left, right| {
1109 for (dst, &count) in left.iter_mut().zip(right.iter()) {
1110 *dst += count;
1111 }
1112 left
1113 },
1114 )
1115 } else {
1116 let mut counts = [0usize; 256];
1117 for &gain in gains {
1118 let key = gain_order_key(gain);
1119 if key & prefix_mask == prefix {
1120 counts[((key >> shift) & 0xff) as usize] += 1;
1121 }
1122 }
1123 Box::new(counts)
1124 }
1125 }
1126 #[cfg(not(feature = "native"))]
1127 {
1128 let mut counts = [0usize; 256];
1129 for &gain in gains {
1130 let key = gain_order_key(gain);
1131 if key & prefix_mask == prefix {
1132 counts[((key >> shift) & 0xff) as usize] += 1;
1133 }
1134 }
1135 counts
1136 }
1137 };
1138
1139 let mut before_bucket = 0usize;
1140 let mut selected_bucket = None;
1141 for (bucket, count) in histogram.iter().copied().enumerate() {
1142 if rank_within_prefix < before_bucket + count {
1143 selected_bucket = Some(bucket as u32);
1144 rank_within_prefix -= before_bucket;
1145 strictly_lower += before_bucket;
1146 break;
1147 }
1148 before_bucket += count;
1149 }
1150 let selected_bucket = selected_bucket.expect("BP radix selection lost the requested rank");
1151 prefix |= selected_bucket << shift;
1152 prefix_mask |= 0xffu32 << shift;
1153 }
1154
1155 (prefix, strictly_lower)
1156}
1157
1158enum PartitionDegreeUpdate {
1159 Moves,
1162 Ranked,
1166 Threshold {
1169 threshold_key: u32,
1170 ties_left: usize,
1171 },
1172}
1173
1174struct PartitionOutcome {
1175 swap_count: usize,
1176 degree_update: PartitionDegreeUpdate,
1177}
1178
1179#[derive(Clone, Copy)]
1180struct PartitionChunk {
1181 start: usize,
1182 end: usize,
1183 strictly_lower: usize,
1184 equal: usize,
1185 ties_left: usize,
1186}
1187
1188#[inline]
1189fn select_left(key: u32, threshold_key: u32, equal_seen: &mut usize, ties_left: usize) -> bool {
1190 if key < threshold_key {
1191 true
1192 } else if key == threshold_key {
1193 let selected = *equal_seen < ties_left;
1194 *equal_seen += 1;
1195 selected
1196 } else {
1197 false
1198 }
1199}
1200
1201fn partition_by_gain(
1210 docs: &[u32],
1211 gains: &[f32],
1212 mid: usize,
1213 fwd: &ForwardIndex,
1214 movement_workspaces: &mut [TermDegrees],
1215 output: &mut [u32],
1216 ranked_scratch: &mut [usize],
1217) -> PartitionOutcome {
1218 #[cfg(not(feature = "native"))]
1219 let _ = (fwd, &movement_workspaces);
1220
1221 #[cfg(feature = "native")]
1222 if !movement_workspaces.is_empty() && docs.len() >= PARALLEL_BP_MIN_ENTITIES {
1223 let (threshold_key, strictly_lower) = select_gain_threshold(gains, mid);
1224 let ties_left = mid - strictly_lower;
1225
1226 let chunk_count = movement_workspaces
1229 .len()
1230 .min(docs.len().div_ceil(PARALLEL_BP_MIN_ENTITIES).max(1));
1231 let chunk_len = docs.len().div_ceil(chunk_count);
1232 let mut chunks: Vec<PartitionChunk> = gains
1233 .par_chunks(chunk_len)
1234 .enumerate()
1235 .map(|(chunk_id, chunk)| {
1236 let mut lower = 0usize;
1237 let mut equal = 0usize;
1238 for &gain in chunk {
1239 match gain_order_key(gain).cmp(&threshold_key) {
1240 std::cmp::Ordering::Less => lower += 1,
1241 std::cmp::Ordering::Equal => equal += 1,
1242 std::cmp::Ordering::Greater => {}
1243 }
1244 }
1245 let start = chunk_id * chunk_len;
1246 PartitionChunk {
1247 start,
1248 end: start + chunk.len(),
1249 strictly_lower: lower,
1250 equal,
1251 ties_left: 0,
1252 }
1253 })
1254 .collect();
1255
1256 let mut remaining_ties = ties_left;
1257 for chunk in &mut chunks {
1258 chunk.ties_left = remaining_ties.min(chunk.equal);
1259 remaining_ties -= chunk.ties_left;
1260 }
1261 debug_assert_eq!(remaining_ties, 0);
1262
1263 let (mut left_rest, mut right_rest) = output.split_at_mut(mid);
1264 let mut jobs = Vec::with_capacity(chunks.len());
1265 for chunk in chunks {
1266 let left_len = chunk.strictly_lower + chunk.ties_left;
1267 let right_len = chunk.end - chunk.start - left_len;
1268 let (left_out, next_left) = left_rest.split_at_mut(left_len);
1269 let (right_out, next_right) = right_rest.split_at_mut(right_len);
1270 jobs.push((
1271 chunk.start,
1272 &docs[chunk.start..chunk.end],
1273 &gains[chunk.start..chunk.end],
1274 chunk.ties_left,
1275 left_out,
1276 right_out,
1277 ));
1278 left_rest = next_left;
1279 right_rest = next_right;
1280 }
1281 debug_assert!(left_rest.is_empty() && right_rest.is_empty());
1282
1283 let swap_count = movement_workspaces[..chunk_count]
1284 .par_iter_mut()
1285 .zip(jobs.into_par_iter())
1286 .map(
1287 |(moves, (start, docs, gains, ties_for_chunk, left_out, right_out))| {
1288 moves.reset();
1289 let mut equal_seen = 0usize;
1290 let mut left_cursor = 0usize;
1291 let mut right_cursor = 0usize;
1292 let mut swaps = 0usize;
1293
1294 for (offset, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1295 let key = gain_order_key(gain);
1296 let now_left =
1297 select_left(key, threshold_key, &mut equal_seen, ties_for_chunk);
1298 if now_left {
1299 left_out[left_cursor] = doc;
1300 left_cursor += 1;
1301 } else {
1302 right_out[right_cursor] = doc;
1303 right_cursor += 1;
1304 }
1305
1306 let was_left = start + offset < mid;
1307 if was_left != now_left {
1308 swaps += 1;
1309 let direction = usize::from(was_left);
1311 for &term in fwd.doc_terms(doc as usize) {
1312 moves.entry_mut(term as usize)[direction] += 1;
1313 }
1314 }
1315 }
1316 debug_assert_eq!(left_cursor, left_out.len());
1317 debug_assert_eq!(right_cursor, right_out.len());
1318 swaps
1319 },
1320 )
1321 .sum();
1322
1323 let (moves, partials) = movement_workspaces[..chunk_count]
1324 .split_first_mut()
1325 .expect("parallel BP partition must use at least one movement workspace");
1326 for partial in partials {
1327 moves.merge_from(partial);
1328 }
1329
1330 return PartitionOutcome {
1331 swap_count,
1332 degree_update: PartitionDegreeUpdate::Moves,
1333 };
1334 }
1335
1336 if docs.len() < PARALLEL_BP_MIN_ENTITIES {
1337 debug_assert_eq!(ranked_scratch.len(), docs.len());
1338 for (index, rank) in ranked_scratch.iter_mut().enumerate() {
1339 *rank = index;
1340 }
1341 ranked_scratch.select_nth_unstable_by(mid, |&left, &right| {
1342 gains[left]
1343 .total_cmp(&gains[right])
1344 .then_with(|| left.cmp(&right))
1345 });
1346
1347 let mut swaps = 0usize;
1348 for (rank, &old_index) in ranked_scratch.iter().enumerate() {
1349 output[rank] = docs[old_index];
1350 swaps += usize::from((old_index < mid) != (rank < mid));
1351 }
1352 return PartitionOutcome {
1353 swap_count: swaps,
1354 degree_update: PartitionDegreeUpdate::Ranked,
1355 };
1356 }
1357
1358 let (threshold_key, strictly_lower) = select_gain_threshold(gains, mid);
1359 let ties_left = mid - strictly_lower;
1360 let mut equal_seen = 0usize;
1361 let mut left_cursor = 0usize;
1362 let mut right_cursor = mid;
1363 let mut swaps = 0usize;
1364 for (idx, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1365 let key = gain_order_key(gain);
1366 let now_left = select_left(key, threshold_key, &mut equal_seen, ties_left);
1367 if now_left {
1368 output[left_cursor] = doc;
1369 left_cursor += 1;
1370 } else {
1371 output[right_cursor] = doc;
1372 right_cursor += 1;
1373 }
1374 swaps += usize::from((idx < mid) != now_left);
1375 }
1376 debug_assert_eq!(left_cursor, mid);
1377 debug_assert_eq!(right_cursor, docs.len());
1378
1379 PartitionOutcome {
1380 swap_count: swaps,
1381 degree_update: PartitionDegreeUpdate::Threshold {
1382 threshold_key,
1383 ties_left,
1384 },
1385 }
1386}
1387
1388fn update_degrees_for_threshold_partition(
1390 docs: &[u32],
1391 gains: &[f32],
1392 mid: usize,
1393 threshold_key: u32,
1394 ties_left: usize,
1395 fwd: &ForwardIndex,
1396 degrees: &mut TermDegrees,
1397) {
1398 let mut equal_seen = 0usize;
1399 for (idx, (&doc, &gain)) in docs.iter().zip(gains).enumerate() {
1400 let key = gain_order_key(gain);
1401 let now_left = select_left(key, threshold_key, &mut equal_seen, ties_left);
1402 let was_left = idx < mid;
1403 if was_left == now_left {
1404 continue;
1405 }
1406 let left_delta = if was_left { -1i64 } else { 1i64 };
1407 for &term in fwd.doc_terms(doc as usize) {
1408 let degree = degrees.entry_mut(term as usize);
1409 let new_left = degree[0] as i64 + left_delta;
1410 let new_right = degree[1] as i64 - left_delta;
1411 debug_assert!(new_left >= 0 && new_right >= 0);
1412 degree[0] = new_left as u32;
1413 degree[1] = new_right as u32;
1414 }
1415 }
1416}
1417
1418fn update_degrees_for_ranked_partition(
1421 docs: &[u32],
1422 ranked: &[usize],
1423 mid: usize,
1424 fwd: &ForwardIndex,
1425 degrees: &mut TermDegrees,
1426) {
1427 for (rank, &old_index) in ranked.iter().enumerate() {
1428 let was_left = old_index < mid;
1429 let now_left = rank < mid;
1430 if was_left == now_left {
1431 continue;
1432 }
1433 let left_delta = if was_left { -1i64 } else { 1i64 };
1434 for &term in fwd.doc_terms(docs[old_index] as usize) {
1435 let degree = degrees.entry_mut(term as usize);
1436 let new_left = degree[0] as i64 + left_delta;
1437 let new_right = degree[1] as i64 - left_delta;
1438 debug_assert!(new_left >= 0 && new_right >= 0);
1439 degree[0] = new_left as u32;
1440 degree[1] = new_right as u32;
1441 }
1442 }
1443}
1444
1445#[derive(Clone, Copy)]
1446pub(crate) struct BpProgressLabel<'a> {
1447 pub index: &'a str,
1448 pub field: &'a str,
1449 pub entity_kind: &'static str,
1450}
1451
1452#[cfg(test)]
1453impl BpProgressLabel<'static> {
1454 fn anonymous() -> Self {
1455 Self {
1456 index: "unknown",
1457 field: "unknown",
1458 entity_kind: "entities",
1459 }
1460 }
1461}
1462
1463#[cfg(feature = "native")]
1464struct BpProgress<'a> {
1465 label: BpProgressLabel<'a>,
1466 start: std::time::Instant,
1467 total_entities: usize,
1468 total_postings: u64,
1469 expected_depth: usize,
1470 next_log_ms: std::sync::atomic::AtomicU64,
1471 active_partitions: std::sync::atomic::AtomicU64,
1472 partitions_started: std::sync::atomic::AtomicU64,
1473 partitions_completed: std::sync::atomic::AtomicU64,
1474 iterations: std::sync::atomic::AtomicU64,
1475 entity_passes: std::sync::atomic::AtomicU64,
1476 swaps: std::sync::atomic::AtomicU64,
1477 deepest_level: std::sync::atomic::AtomicU64,
1478 objective_stops: std::sync::atomic::AtomicU64,
1479 last_objective_delta_bits: std::sync::atomic::AtomicU64,
1480 last_relative_delta_bits: std::sync::atomic::AtomicU64,
1481 active_metric_released: std::sync::atomic::AtomicBool,
1482}
1483
1484#[cfg(feature = "native")]
1485impl<'a> BpProgress<'a> {
1486 fn new(
1487 label: BpProgressLabel<'a>,
1488 total_entities: usize,
1489 total_postings: u64,
1490 expected_depth: usize,
1491 ) -> Self {
1492 log::info!(
1493 "[reorder][bp] started: index={} field={} entity_kind={} entities={} postings={} expected_depth={} objective_stall_threshold={:.1e}x{} min_objective_iterations={}",
1494 label.index,
1495 label.field,
1496 label.entity_kind,
1497 total_entities,
1498 total_postings,
1499 expected_depth,
1500 MIN_RELATIVE_OBJECTIVE_IMPROVEMENT,
1501 OBJECTIVE_STALL_ITERATIONS,
1502 MIN_OBJECTIVE_ITERATIONS,
1503 );
1504 crate::observe::reorder_bp_started(label.index, label.field, label.entity_kind);
1505 Self {
1506 label,
1507 start: std::time::Instant::now(),
1508 total_entities,
1509 total_postings,
1510 expected_depth,
1511 next_log_ms: std::sync::atomic::AtomicU64::new(30_000),
1512 active_partitions: std::sync::atomic::AtomicU64::new(0),
1513 partitions_started: std::sync::atomic::AtomicU64::new(0),
1514 partitions_completed: std::sync::atomic::AtomicU64::new(0),
1515 iterations: std::sync::atomic::AtomicU64::new(0),
1516 entity_passes: std::sync::atomic::AtomicU64::new(0),
1517 swaps: std::sync::atomic::AtomicU64::new(0),
1518 deepest_level: std::sync::atomic::AtomicU64::new(0),
1519 objective_stops: std::sync::atomic::AtomicU64::new(0),
1520 last_objective_delta_bits: std::sync::atomic::AtomicU64::new(0f64.to_bits()),
1521 last_relative_delta_bits: std::sync::atomic::AtomicU64::new(0f64.to_bits()),
1522 active_metric_released: std::sync::atomic::AtomicBool::new(false),
1523 }
1524 }
1525
1526 fn partition_started(&self, level: usize) {
1527 self.active_partitions
1528 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1529 self.partitions_started
1530 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1531 self.deepest_level
1532 .fetch_max(level as u64, std::sync::atomic::Ordering::Relaxed);
1533 }
1534
1535 fn partition_finished(&self) {
1536 self.partitions_completed
1537 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1538 self.active_partitions
1539 .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
1540 }
1541
1542 fn iteration(&self, entities: usize, swaps: usize, objective_delta: f64, relative_delta: f64) {
1543 self.iterations
1544 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1545 self.entity_passes
1546 .fetch_add(entities as u64, std::sync::atomic::Ordering::Relaxed);
1547 self.swaps
1548 .fetch_add(swaps as u64, std::sync::atomic::Ordering::Relaxed);
1549 self.last_objective_delta_bits.store(
1550 objective_delta.to_bits(),
1551 std::sync::atomic::Ordering::Relaxed,
1552 );
1553 self.last_relative_delta_bits.store(
1554 relative_delta.to_bits(),
1555 std::sync::atomic::Ordering::Relaxed,
1556 );
1557 self.maybe_log();
1558 }
1559
1560 fn objective_stop(&self) {
1561 self.objective_stops
1562 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1563 }
1564
1565 fn maybe_log(&self) {
1566 let elapsed_ms = self.start.elapsed().as_millis().min(u64::MAX as u128) as u64;
1567 let next = self.next_log_ms.load(std::sync::atomic::Ordering::Relaxed);
1568 if elapsed_ms < next
1569 || self
1570 .next_log_ms
1571 .compare_exchange(
1572 next,
1573 elapsed_ms.saturating_add(30_000),
1574 std::sync::atomic::Ordering::Relaxed,
1575 std::sync::atomic::Ordering::Relaxed,
1576 )
1577 .is_err()
1578 {
1579 return;
1580 }
1581
1582 let active = self
1583 .active_partitions
1584 .load(std::sync::atomic::Ordering::Relaxed);
1585 let started = self
1586 .partitions_started
1587 .load(std::sync::atomic::Ordering::Relaxed);
1588 let completed = self
1589 .partitions_completed
1590 .load(std::sync::atomic::Ordering::Relaxed);
1591 let iterations = self.iterations.load(std::sync::atomic::Ordering::Relaxed);
1592 let entity_passes = self
1593 .entity_passes
1594 .load(std::sync::atomic::Ordering::Relaxed);
1595 let swaps = self.swaps.load(std::sync::atomic::Ordering::Relaxed);
1596 let deepest = self
1597 .deepest_level
1598 .load(std::sync::atomic::Ordering::Relaxed);
1599 let objective_delta = f64::from_bits(
1600 self.last_objective_delta_bits
1601 .load(std::sync::atomic::Ordering::Relaxed),
1602 );
1603 let relative_delta = f64::from_bits(
1604 self.last_relative_delta_bits
1605 .load(std::sync::atomic::Ordering::Relaxed),
1606 );
1607 log::info!(
1608 "[reorder][bp] progress: index={} field={} entity_kind={} elapsed={:.1}s depth={}/{} partitions={}/{} active={} iterations={} entity_passes={} swaps={} last_objective_delta={:.3} relative={:.3e}",
1609 self.label.index,
1610 self.label.field,
1611 self.label.entity_kind,
1612 self.start.elapsed().as_secs_f64(),
1613 deepest,
1614 self.expected_depth,
1615 completed,
1616 started,
1617 active,
1618 iterations,
1619 entity_passes,
1620 swaps,
1621 objective_delta,
1622 relative_delta,
1623 );
1624 }
1625
1626 fn finish(
1627 &self,
1628 converged: bool,
1629 memory_limited: bool,
1630 deadline_exhausted: bool,
1631 cancelled: bool,
1632 ) {
1633 let elapsed = self.start.elapsed().as_secs_f64();
1634 let partitions = self
1635 .partitions_completed
1636 .load(std::sync::atomic::Ordering::Relaxed);
1637 let iterations = self.iterations.load(std::sync::atomic::Ordering::Relaxed);
1638 let entity_passes = self
1639 .entity_passes
1640 .load(std::sync::atomic::Ordering::Relaxed);
1641 let swaps = self.swaps.load(std::sync::atomic::Ordering::Relaxed);
1642 let deepest = self
1643 .deepest_level
1644 .load(std::sync::atomic::Ordering::Relaxed);
1645 let objective_stops = self
1646 .objective_stops
1647 .load(std::sync::atomic::Ordering::Relaxed);
1648 let stop_reason = if cancelled {
1649 "shutdown"
1650 } else if memory_limited {
1651 "memory_budget"
1652 } else if deadline_exhausted {
1653 "time_budget"
1654 } else if objective_stops > 0 {
1655 "objective"
1656 } else {
1657 "complete"
1658 };
1659 log::info!(
1660 "[reorder][bp] completed: index={} field={} entity_kind={} entities={} postings={} elapsed={:.1}s depth={}/{} partitions={} iterations={} entity_passes={} swaps={} objective_stops={} converged={} stop_reason={}",
1661 self.label.index,
1662 self.label.field,
1663 self.label.entity_kind,
1664 self.total_entities,
1665 self.total_postings,
1666 elapsed,
1667 deepest,
1668 self.expected_depth,
1669 partitions,
1670 iterations,
1671 entity_passes,
1672 swaps,
1673 objective_stops,
1674 converged,
1675 stop_reason,
1676 );
1677 crate::observe::reorder_bp_pass(
1678 self.label.index,
1679 self.label.field,
1680 self.label.entity_kind,
1681 stop_reason,
1682 elapsed,
1683 self.total_entities,
1684 self.total_postings,
1685 partitions,
1686 iterations,
1687 entity_passes,
1688 swaps,
1689 converged,
1690 );
1691 self.release_active_metric();
1692 }
1693
1694 fn release_active_metric(&self) {
1695 if !self
1696 .active_metric_released
1697 .swap(true, std::sync::atomic::Ordering::AcqRel)
1698 {
1699 crate::observe::reorder_bp_finished(
1700 self.label.index,
1701 self.label.field,
1702 self.label.entity_kind,
1703 );
1704 }
1705 }
1706}
1707
1708#[cfg(feature = "native")]
1709impl Drop for BpProgress<'_> {
1710 fn drop(&mut self) {
1711 self.release_active_metric();
1712 }
1713}
1714
1715#[cfg(not(feature = "native"))]
1716struct BpProgress<'a>(std::marker::PhantomData<&'a ()>);
1717
1718#[cfg(not(feature = "native"))]
1719impl BpProgress<'_> {
1720 fn new(_: BpProgressLabel<'_>, _: usize, _: u64, _: usize) -> Self {
1721 Self(std::marker::PhantomData)
1722 }
1723 fn partition_started(&self, _: usize) {}
1724 fn partition_finished(&self) {}
1725 fn iteration(&self, _: usize, _: usize, _: f64, _: f64) {}
1726 fn objective_stop(&self) {}
1727 fn finish(&self, _: bool, _: bool, _: bool, _: bool) {}
1728}
1729
1730#[cfg(test)]
1741pub(crate) fn graph_bisection(
1742 fwd: &ForwardIndex,
1743 min_partition_size: usize,
1744 max_iters: usize,
1745 budget: BpBudget,
1746) -> (Vec<u32>, bool) {
1747 graph_bisection_with_progress(
1748 fwd,
1749 min_partition_size,
1750 max_iters,
1751 budget,
1752 None,
1753 BpProgressLabel::anonymous(),
1754 )
1755}
1756
1757pub(crate) fn graph_bisection_with_progress(
1758 fwd: &ForwardIndex,
1759 min_partition_size: usize,
1760 max_iters: usize,
1761 budget: BpBudget,
1762 cancellation: Option<&std::sync::atomic::AtomicBool>,
1763 progress_label: BpProgressLabel<'_>,
1764) -> (Vec<u32>, bool) {
1765 let n = fwd.num_docs();
1766 if n == 0 {
1767 return (Vec::new(), !fwd.budget_limited);
1768 }
1769
1770 let effective_min_partition = budget
1771 .min_partition_docs
1772 .unwrap_or(0)
1773 .max(min_partition_size);
1774
1775 let mut docs: Vec<u32> = (0..n as u32).collect();
1776 let depth = if effective_min_partition > 0 {
1777 ((n as f64) / (effective_min_partition as f64))
1778 .log2()
1779 .ceil() as usize
1780 } else {
1781 0
1782 };
1783 let log_table = build_log_table(4096);
1784 let progress = BpProgress::new(progress_label, n, fwd.total_postings(), depth);
1785
1786 log::debug!(
1787 "BP graph_bisection: n={}, min_partition={}, max_iters={}, depth=~{}, time_budget={:?}",
1788 n,
1789 effective_min_partition,
1790 max_iters,
1791 depth,
1792 budget.time_budget,
1793 );
1794
1795 #[cfg(feature = "native")]
1796 let deadline = budget.time_budget.map(|duration| {
1797 let now = std::time::Instant::now();
1798 now.checked_add(duration).unwrap_or(now)
1799 });
1800 #[cfg(not(feature = "native"))]
1801 let deadline: Option<()> = None;
1802
1803 let exhausted = std::sync::atomic::AtomicBool::new(false);
1804 let context = BisectContext {
1805 fwd,
1806 min_partition_size: effective_min_partition,
1807 max_iters,
1808 log_table: &log_table,
1809 #[cfg(feature = "native")]
1810 deadline,
1811 #[cfg(not(feature = "native"))]
1812 deadline,
1813 exhausted: &exhausted,
1814 cancellation,
1815 progress: &progress,
1816 };
1817 let immediately_exhausted = {
1818 #[cfg(feature = "native")]
1819 {
1820 deadline.is_some_and(|deadline| std::time::Instant::now() >= deadline)
1821 }
1822 #[cfg(not(feature = "native"))]
1823 {
1824 false
1825 }
1826 };
1827 let initially_cancelled =
1828 cancellation.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire));
1829 if immediately_exhausted || initially_cancelled {
1830 exhausted.store(true, std::sync::atomic::Ordering::Relaxed);
1831 } else if n > effective_min_partition {
1832 let mut gains = vec![0.0f32; n];
1836 let mut partitioned = vec![0u32; n];
1837 let mut ranked = vec![0usize; n];
1838 #[cfg(feature = "native")]
1839 let degree_lanes = fwd.parallel_bisect_lanes.max(1);
1840 #[cfg(not(feature = "native"))]
1841 let degree_lanes = 1usize;
1842 let mut degree_workspaces: Vec<TermDegrees> = (0..degree_lanes)
1843 .map(|_| TermDegrees::new(fwd.num_terms))
1844 .collect();
1845 bisect(
1846 &mut docs,
1847 &mut gains,
1848 &mut partitioned,
1849 &mut ranked,
1850 &mut degree_workspaces,
1851 0,
1852 &context,
1853 );
1854 }
1855
1856 let stopped_early = exhausted.load(std::sync::atomic::Ordering::Relaxed);
1857 let cancelled =
1858 cancellation.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire));
1859 let deadline_exhausted = stopped_early && !cancelled;
1860 let converged = !fwd.budget_limited && !stopped_early && !cancelled;
1861 progress.finish(converged, fwd.budget_limited, deadline_exhausted, cancelled);
1862 if !converged {
1863 log::info!(
1864 "BP graph_bisection: pass incomplete at n={} (time={:?}, memory_limited={}, cancelled={}) — emitting partial (still valid) permutation",
1865 n,
1866 budget.time_budget,
1867 fwd.budget_limited,
1868 cancelled,
1869 );
1870 }
1871 (docs, converged)
1872}
1873
1874struct BisectContext<'a> {
1883 fwd: &'a ForwardIndex,
1884 min_partition_size: usize,
1885 max_iters: usize,
1886 log_table: &'a [f32],
1887 #[cfg(feature = "native")]
1888 deadline: Option<std::time::Instant>,
1889 #[cfg(not(feature = "native"))]
1890 deadline: Option<()>,
1891 exhausted: &'a std::sync::atomic::AtomicBool,
1892 cancellation: Option<&'a std::sync::atomic::AtomicBool>,
1893 progress: &'a BpProgress<'a>,
1894}
1895
1896#[allow(clippy::too_many_arguments)]
1897fn bisect(
1898 docs: &mut [u32],
1899 gains: &mut [f32],
1900 partitioned: &mut [u32],
1901 ranked_scratch: &mut [usize],
1902 degree_workspaces: &mut [TermDegrees],
1903 level: usize,
1904 context: &BisectContext<'_>,
1905) {
1906 let n = docs.len();
1907 debug_assert_eq!(gains.len(), n);
1908 debug_assert_eq!(partitioned.len(), n);
1909 debug_assert_eq!(ranked_scratch.len(), n);
1910 debug_assert!(!degree_workspaces.is_empty());
1911 if n <= context.min_partition_size {
1912 return;
1913 }
1914 if context.exhausted.load(std::sync::atomic::Ordering::Relaxed)
1916 || context
1917 .cancellation
1918 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
1919 {
1920 context
1921 .exhausted
1922 .store(true, std::sync::atomic::Ordering::Relaxed);
1923 return;
1924 }
1925 #[cfg(feature = "native")]
1926 if let Some(dl) = context.deadline
1927 && std::time::Instant::now() >= dl
1928 {
1929 context
1930 .exhausted
1931 .store(true, std::sync::atomic::Ordering::Relaxed);
1932 return;
1933 }
1934 #[cfg(not(feature = "native"))]
1935 let _ = context.deadline;
1936
1937 context.progress.partition_started(level);
1938 let mid = n / 2;
1939
1940 let effective_iters = if n > 100_000 {
1944 context.max_iters.min(12)
1945 } else {
1946 context.max_iters
1947 };
1948
1949 build_term_degrees(
1954 docs,
1955 mid,
1956 context.fwd,
1957 degree_workspaces,
1958 context.cancellation,
1959 );
1960 if context
1961 .cancellation
1962 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
1963 {
1964 context
1965 .exhausted
1966 .store(true, std::sync::atomic::Ordering::Relaxed);
1967 return;
1968 }
1969 let track_objective = n >= PARALLEL_BP_MIN_ENTITIES;
1974 if track_objective {
1975 degree_workspaces[0].sort_touched_words();
1979 }
1980 let mut previous_objective = if track_objective {
1981 degree_workspaces[0].bisection_objective(mid, n - mid, context.log_table)
1982 } else {
1983 0.0
1984 };
1985 let mut best_objective = previous_objective;
1986 let mut objective_stalls = 0usize;
1987
1988 for iter in 0..effective_iters {
1989 if context
1991 .cancellation
1992 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
1993 {
1994 context
1995 .exhausted
1996 .store(true, std::sync::atomic::Ordering::Relaxed);
1997 break;
1998 }
1999 #[cfg(feature = "native")]
2000 if let Some(dl) = context.deadline
2001 && std::time::Instant::now() >= dl
2002 {
2003 context
2004 .exhausted
2005 .store(true, std::sync::atomic::Ordering::Relaxed);
2006 break;
2007 }
2008 compute_gains(
2011 docs,
2012 context.fwd,
2013 mid,
2014 °ree_workspaces[0],
2015 context.log_table,
2016 gains,
2017 );
2018 if context
2019 .cancellation
2020 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Acquire))
2021 {
2022 context
2023 .exhausted
2024 .store(true, std::sync::atomic::Ordering::Relaxed);
2025 break;
2026 }
2027
2028 let partition = partition_by_gain(
2033 docs,
2034 gains,
2035 mid,
2036 context.fwd,
2037 &mut degree_workspaces[1..],
2038 partitioned,
2039 ranked_scratch,
2040 );
2041
2042 if partition.swap_count == 0 {
2043 context.progress.iteration(n, 0, 0.0, 0.0);
2044 docs.copy_from_slice(partitioned);
2048 break;
2049 }
2050
2051 match &partition.degree_update {
2052 PartitionDegreeUpdate::Moves => {
2053 let (degrees, movement_workspaces) = degree_workspaces
2054 .split_first_mut()
2055 .expect("BP always has one degree workspace");
2056 movement_workspaces
2057 .first()
2058 .expect("parallel BP movement update requires a spare workspace")
2059 .apply_moves_to(degrees);
2060 }
2061 PartitionDegreeUpdate::Ranked => update_degrees_for_ranked_partition(
2062 docs,
2063 ranked_scratch,
2064 mid,
2065 context.fwd,
2066 &mut degree_workspaces[0],
2067 ),
2068 PartitionDegreeUpdate::Threshold {
2069 threshold_key,
2070 ties_left,
2071 } => update_degrees_for_threshold_partition(
2072 docs,
2073 gains,
2074 mid,
2075 *threshold_key,
2076 *ties_left,
2077 context.fwd,
2078 &mut degree_workspaces[0],
2079 ),
2080 }
2081
2082 let (new_objective, objective_improvement, relative_improvement) = if track_objective {
2083 let new_objective =
2084 degree_workspaces[0].bisection_objective(mid, n - mid, context.log_table);
2085 let objective_improvement = new_objective - previous_objective;
2086 let relative_improvement = objective_improvement / previous_objective.abs().max(1.0);
2087 (new_objective, objective_improvement, relative_improvement)
2088 } else {
2089 (0.0, 0.0, 0.0)
2090 };
2091 context.progress.iteration(
2092 n,
2093 partition.swap_count,
2094 objective_improvement,
2095 relative_improvement,
2096 );
2097
2098 docs.copy_from_slice(partitioned);
2105 if track_objective {
2106 previous_objective = new_objective;
2107 let relative_best_improvement =
2108 (new_objective - best_objective) / best_objective.abs().max(1.0);
2109 if relative_best_improvement >= MIN_RELATIVE_OBJECTIVE_IMPROVEMENT {
2110 best_objective = new_objective;
2111 objective_stalls = 0;
2112 } else if iter + 1 >= MIN_OBJECTIVE_ITERATIONS {
2113 objective_stalls += 1;
2114 }
2115 if objective_stalls >= OBJECTIVE_STALL_ITERATIONS {
2116 context.progress.objective_stop();
2117 break;
2118 }
2119 }
2120
2121 if iter > 2 && partition.swap_count < n / 200 {
2123 break;
2124 }
2125
2126 if !track_objective && iter > 5 {
2130 let max_abs_gain = gains
2131 .iter()
2132 .copied()
2133 .fold(0.0f32, |max_gain, gain| max_gain.max(gain.abs()));
2134 if max_abs_gain < 0.001 {
2135 break;
2136 }
2137 }
2138 }
2139
2140 context.progress.partition_finished();
2141
2142 let (left, right) = docs.split_at_mut(mid);
2143 let (left_gains, right_gains) = gains.split_at_mut(mid);
2144 let (left_partitioned, right_partitioned) = partitioned.split_at_mut(mid);
2145 let (left_ranked, right_ranked) = ranked_scratch.split_at_mut(mid);
2146 #[cfg(feature = "native")]
2147 if degree_workspaces.len() > 1 {
2148 let left_lanes = degree_workspaces.len() / 2;
2149 let (left_workspaces, right_workspaces) = degree_workspaces.split_at_mut(left_lanes);
2150 rayon::join(
2151 || {
2152 bisect(
2153 left,
2154 left_gains,
2155 left_partitioned,
2156 left_ranked,
2157 left_workspaces,
2158 level + 1,
2159 context,
2160 )
2161 },
2162 || {
2163 bisect(
2164 right,
2165 right_gains,
2166 right_partitioned,
2167 right_ranked,
2168 right_workspaces,
2169 level + 1,
2170 context,
2171 )
2172 },
2173 );
2174 } else {
2175 bisect(
2179 left,
2180 left_gains,
2181 left_partitioned,
2182 left_ranked,
2183 degree_workspaces,
2184 level + 1,
2185 context,
2186 );
2187 bisect(
2188 right,
2189 right_gains,
2190 right_partitioned,
2191 right_ranked,
2192 degree_workspaces,
2193 level + 1,
2194 context,
2195 );
2196 }
2197 #[cfg(not(feature = "native"))]
2198 {
2199 bisect(
2200 left,
2201 left_gains,
2202 left_partitioned,
2203 left_ranked,
2204 degree_workspaces,
2205 level + 1,
2206 context,
2207 );
2208 bisect(
2209 right,
2210 right_gains,
2211 right_partitioned,
2212 right_ranked,
2213 degree_workspaces,
2214 level + 1,
2215 context,
2216 );
2217 }
2218}
2219
2220#[inline(never)]
2226fn compute_gains(
2227 docs: &[u32],
2228 fwd: &ForwardIndex,
2229 mid: usize,
2230 degrees: &TermDegrees,
2231 log_table: &[f32],
2232 gains: &mut [f32],
2233) {
2234 let gain_for_doc = |i: usize| -> f32 {
2243 let doc = docs[i] as usize;
2244 let in_left = i < mid;
2245 let mut g = 0.0f32;
2246 for &term in fwd.doc_terms(doc) {
2247 let [left, right] = degrees.get(term as usize);
2248 let (from, to) = if in_left {
2249 (left, right)
2250 } else {
2251 (right, left)
2252 };
2253 let move_gain = fast_log2_lookup(to as usize + 2, log_table)
2254 - fast_log2_lookup(from as usize, log_table)
2255 - std::f32::consts::LOG2_E / (1.0 + to as f32);
2256 g += if in_left { move_gain } else { -move_gain };
2257 }
2258 g
2259 };
2260
2261 #[cfg(feature = "native")]
2262 {
2263 if docs.len() > 4096 {
2264 gains
2265 .par_iter_mut()
2266 .enumerate()
2267 .for_each(|(i, gain)| *gain = gain_for_doc(i));
2268 } else {
2269 for (i, gain) in gains.iter_mut().enumerate().take(docs.len()) {
2270 *gain = gain_for_doc(i);
2271 }
2272 }
2273 }
2274 #[cfg(not(feature = "native"))]
2275 {
2276 for (i, gain) in gains.iter_mut().enumerate().take(docs.len()) {
2277 *gain = gain_for_doc(i);
2278 }
2279 }
2280}
2281
2282fn build_log_table(size: usize) -> Vec<f32> {
2286 let mut table = vec![0.0f32; size];
2287 table[0] = -10.0;
2289 for (i, entry) in table.iter_mut().enumerate().skip(1) {
2290 *entry = (i as f32).log2();
2291 }
2292 table
2293}
2294
2295#[inline]
2297fn fast_log2_lookup(val: usize, table: &[f32]) -> f32 {
2298 if val < table.len() {
2299 table[val]
2300 } else {
2301 (val as f32).log2()
2302 }
2303}
2304
2305#[cfg(test)]
2306mod tests {
2307 use super::*;
2308
2309 #[test]
2310 fn bounded_frequency_count_and_candidate_selection_are_exact() {
2311 let items: Vec<u32> = (0..10_000).collect();
2312 let frequencies = count_frequencies_bounded(&items, 7, 16 * 1024, |item, counts| {
2313 let term = *item as usize % counts.len();
2314 counts[term] += 1;
2315 })
2316 .unwrap();
2317 assert_eq!(frequencies.iter().sum::<u32>(), items.len() as u32);
2318 for (term, &frequency) in frequencies.iter().enumerate() {
2319 let expected = (term..items.len()).step_by(frequencies.len()).count() as u32;
2320 assert_eq!(frequency, expected);
2321 }
2322
2323 let (selected, limited) =
2324 select_frequency_candidates(&[8, 3, 0, 5, 3], 1, 10, 2 * CANDIDATE_ENTRY_BYTES);
2325 assert!(limited);
2326 let mut selected = selected;
2327 selected.sort_unstable();
2328 assert_eq!(selected, vec![(1, 3), (4, 3)]);
2329 }
2330
2331 #[test]
2332 fn bounded_frequency_count_rejects_an_undersized_budget() {
2333 assert!(
2334 count_frequencies_bounded(&[0u32], 32, 32, |_, _| {}).is_none(),
2335 "one complete dense frequency table must fit before counting"
2336 );
2337 }
2338
2339 #[test]
2340 fn lazy_term_degrees_initialize_only_on_first_write() {
2341 let mut degrees = TermDegrees::new(130);
2342 let values_ptr = degrees.values.as_ptr();
2343 let bitmap_ptr = degrees.initialized.as_ptr();
2344 assert_eq!(degrees.get(65), [0, 0]);
2345 degrees.entry_mut(65)[0] += 3;
2346 degrees.entry_mut(65)[1] += 2;
2347 assert_eq!(degrees.get(65), [3, 2]);
2348 assert_eq!(degrees.get(64), [0, 0]);
2349 assert_eq!(
2350 degrees
2351 .initialized
2352 .iter()
2353 .map(|w| w.count_ones())
2354 .sum::<u32>(),
2355 1
2356 );
2357
2358 degrees.reset();
2359 assert_eq!(degrees.values.as_ptr(), values_ptr);
2360 assert_eq!(degrees.initialized.as_ptr(), bitmap_ptr);
2361 assert_eq!(degrees.get(65), [0, 0]);
2362 assert!(degrees.touched_words.is_empty());
2363 degrees.entry_mut(129)[1] = 7;
2364 assert_eq!(degrees.get(129), [0, 7]);
2365 assert_eq!(degrees.get(65), [0, 0]);
2366 }
2367
2368 #[test]
2369 fn sparse_objective_keeps_the_original_ascending_term_order() {
2370 let mut degrees = TermDegrees::new(130);
2371 *degrees.entry_mut(129) = [5, 2];
2372 *degrees.entry_mut(1) = [3, 7];
2373 *degrees.entry_mut(65) = [11, 13];
2374 assert_eq!(degrees.touched_words, vec![2, 0, 1]);
2375 degrees.sort_touched_words();
2376 assert_eq!(degrees.touched_words, vec![0, 1, 2]);
2377
2378 let log_table = build_log_table(4096);
2379 let actual = degrees.bisection_objective(31, 32, &log_table);
2380 let side_log = [
2381 fast_log2_lookup(31, &log_table) as f64,
2382 fast_log2_lookup(32, &log_table) as f64,
2383 ];
2384 let mut original_bitmap_scan = 0.0f64;
2385 for (word_idx, &initialized) in degrees.initialized.iter().enumerate() {
2386 let mut pending = initialized;
2387 while pending != 0 {
2388 let bit = pending.trailing_zeros() as usize;
2389 let term = word_idx * 64 + bit;
2390 let [left, right] = degrees.get(term);
2391 for (side, count) in [left, right].into_iter().enumerate() {
2392 if count > 0 {
2393 original_bitmap_scan += count as f64
2394 * (fast_log2_lookup(count as usize + 1, &log_table) as f64
2395 - side_log[side]);
2396 }
2397 }
2398 pending &= pending - 1;
2399 }
2400 }
2401 assert_eq!(actual.to_bits(), original_bitmap_scan.to_bits());
2402 }
2403
2404 #[test]
2405 fn gain_radix_key_matches_total_cmp() {
2406 let values = [
2407 f32::from_bits(0xffc0_0001),
2408 f32::NEG_INFINITY,
2409 -42.0,
2410 -0.0,
2411 0.0,
2412 42.0,
2413 f32::INFINITY,
2414 f32::from_bits(0x7fc0_0001),
2415 ];
2416 let mut by_cmp = values;
2417 by_cmp.sort_by(f32::total_cmp);
2418 let mut by_key = values;
2419 by_key.sort_by_key(|value| gain_order_key(*value));
2420 assert_eq!(
2421 by_cmp.map(f32::to_bits),
2422 by_key.map(f32::to_bits),
2423 "radix selection must preserve the former total_cmp order"
2424 );
2425 }
2426
2427 #[test]
2428 fn radix_threshold_matches_exact_rank_with_ties() {
2429 let gains = [3.0, -1.0, 7.0, -1.0, 0.0, -0.0, 3.0, 9.0, 3.0, 2.0, 2.0];
2430 let mut sorted: Vec<(u32, usize)> = gains
2431 .iter()
2432 .enumerate()
2433 .map(|(idx, &gain)| (gain_order_key(gain), idx))
2434 .collect();
2435 sorted.sort_unstable();
2436
2437 for left_count in 1..=gains.len() {
2438 let (threshold, lower) = select_gain_threshold(&gains, left_count);
2439 assert_eq!(threshold, sorted[left_count - 1].0);
2440 assert_eq!(lower, sorted.partition_point(|&(key, _)| key < threshold),);
2441 }
2442 }
2443
2444 #[cfg(feature = "native")]
2445 #[test]
2446 fn parallel_partition_matches_exact_selection_and_degree_rebuild() {
2447 const N: usize = PARALLEL_BP_MIN_ENTITIES + 1;
2448 const TERMS: usize = 101;
2449 let mut terms = Vec::with_capacity(N * 3);
2450 let mut offsets = Vec::with_capacity(N + 1);
2451 offsets.push(0);
2452 for doc in 0..N {
2453 terms.extend_from_slice(&[
2454 (doc % TERMS) as u32,
2455 ((doc / 7) % TERMS) as u32,
2456 ((doc * 13) % TERMS) as u32,
2457 ]);
2458 offsets.push(terms.len() as u64);
2459 }
2460 let fwd = ForwardIndex {
2461 terms,
2462 offsets,
2463 num_terms: TERMS,
2464 parallel_bisect_lanes: 4,
2465 budget_limited: false,
2466 };
2467 let docs: Vec<u32> = (0..N as u32)
2468 .map(|idx| ((idx as usize * 7_919) % N) as u32)
2469 .collect();
2470 let gains: Vec<f32> = docs
2471 .iter()
2472 .map(|&doc| ((doc as usize * 37) % 257) as f32 - 128.0)
2473 .collect();
2474 let mid = N / 2;
2475
2476 let mut ranked: Vec<usize> = (0..N).collect();
2477 ranked.sort_unstable_by(|&left, &right| {
2478 gains[left]
2479 .total_cmp(&gains[right])
2480 .then_with(|| left.cmp(&right))
2481 });
2482 let mut selected_left = vec![false; N];
2483 for &idx in &ranked[..mid] {
2484 selected_left[idx] = true;
2485 }
2486 let expected: Vec<u32> = docs
2487 .iter()
2488 .enumerate()
2489 .filter(|(idx, _)| selected_left[*idx])
2490 .chain(
2491 docs.iter()
2492 .enumerate()
2493 .filter(|(idx, _)| !selected_left[*idx]),
2494 )
2495 .map(|(_, &doc)| doc)
2496 .collect();
2497
2498 let mut output = vec![0; N];
2499 let mut ranked_scratch = Vec::new();
2500 let mut movement_workspaces: Vec<_> = (0..3).map(|_| TermDegrees::new(TERMS)).collect();
2501 let outcome = partition_by_gain(
2502 &docs,
2503 &gains,
2504 mid,
2505 &fwd,
2506 &mut movement_workspaces,
2507 &mut output,
2508 &mut ranked_scratch,
2509 );
2510 assert_eq!(output, expected);
2511 assert!(
2512 matches!(&outcome.degree_update, PartitionDegreeUpdate::Moves),
2513 "test must exercise parallel movement counts"
2514 );
2515
2516 let mut updated_workspaces: Vec<_> = (0..4).map(|_| TermDegrees::new(TERMS)).collect();
2517 build_term_degrees(&docs, mid, &fwd, &mut updated_workspaces, None);
2518 movement_workspaces[0].apply_moves_to(&mut updated_workspaces[0]);
2519
2520 let mut rebuilt_workspaces: Vec<_> = (0..4).map(|_| TermDegrees::new(TERMS)).collect();
2521 build_term_degrees(&output, mid, &fwd, &mut rebuilt_workspaces, None);
2522 for term in 0..TERMS {
2523 assert_eq!(
2524 updated_workspaces[0].get(term),
2525 rebuilt_workspaces[0].get(term),
2526 "parallel movement-count mismatch for term {term}"
2527 );
2528 }
2529 }
2530
2531 #[test]
2537 fn test_csr_offsets_do_not_wrap_past_u32() {
2538 let counts = [1_500_000_000u32; 3]; let offsets = build_csr_offsets(&counts);
2540 assert_eq!(
2541 offsets,
2542 vec![0, 1_500_000_000, 3_000_000_000, 4_500_000_000]
2543 );
2544 assert!(*offsets.last().unwrap() > u32::MAX as u64);
2545 }
2546
2547 fn make_fwd(docs: &[&[u32]], num_terms: usize) -> ForwardIndex {
2549 let mut terms = Vec::new();
2550 let mut offsets = vec![0u64];
2551 for doc_terms in docs {
2552 terms.extend_from_slice(doc_terms);
2553 offsets.push(terms.len() as u64);
2554 }
2555 ForwardIndex {
2556 terms,
2557 offsets,
2558 num_terms,
2559 parallel_bisect_lanes: 1,
2560 budget_limited: false,
2561 }
2562 }
2563
2564 #[test]
2565 fn test_bp_empty() {
2566 let fwd = ForwardIndex {
2567 terms: Vec::new(),
2568 offsets: Vec::new(),
2569 num_terms: 0,
2570 parallel_bisect_lanes: 1,
2571 budget_limited: false,
2572 };
2573 let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2574 assert!(perm.is_empty());
2575 }
2576
2577 #[test]
2578 fn test_bp_small() {
2579 let fwd = make_fwd(&[&[0, 1], &[0, 2], &[1, 3], &[2, 3]], 4);
2581 let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2582 assert_eq!(perm.len(), 4);
2583 let mut sorted = perm.clone();
2585 sorted.sort();
2586 assert_eq!(sorted, vec![0, 1, 2, 3]);
2587 }
2588
2589 #[test]
2590 fn test_bp_clusters() {
2591 let fwd = make_fwd(
2595 &[
2596 &[0, 1],
2597 &[0, 1],
2598 &[0, 1],
2599 &[0, 1],
2600 &[2, 3],
2601 &[2, 3],
2602 &[2, 3],
2603 &[2, 3],
2604 ],
2605 4,
2606 );
2607 let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2608 assert_eq!(perm.len(), 8);
2609
2610 let left: Vec<u32> = perm[..4].to_vec();
2612
2613 let a_in_left = left.iter().filter(|&&d| d < 4).count();
2615 let b_in_left = left.iter().filter(|&&d| d >= 4).count();
2616 assert!(
2617 (a_in_left == 4 && b_in_left == 0) || (a_in_left == 0 && b_in_left == 4),
2618 "Clusters should be separated: a_left={}, b_left={}",
2619 a_in_left,
2620 b_in_left,
2621 );
2622 }
2623
2624 #[test]
2625 fn test_bp_permutation_valid() {
2626 let docs: Vec<Vec<u32>> = (0..16).map(|i| vec![i / 4, 10 + i / 2]).collect();
2628 let doc_refs: Vec<&[u32]> = docs.iter().map(|v| v.as_slice()).collect();
2629 let fwd = make_fwd(&doc_refs, 18); let (perm, _) = graph_bisection(&fwd, 4, 20, BpBudget::full());
2631
2632 assert_eq!(perm.len(), 16);
2633 let mut sorted = perm.clone();
2635 sorted.sort();
2636 let expected: Vec<u32> = (0..16).collect();
2637 assert_eq!(sorted, expected);
2638 }
2639
2640 #[test]
2645 fn test_bp_depth_cap_separates_clusters_and_converges() {
2646 let fwd = make_fwd(
2649 &[
2650 &[0, 1],
2651 &[0, 1],
2652 &[0, 1],
2653 &[2, 3],
2654 &[0, 1],
2655 &[2, 3],
2656 &[2, 3],
2657 &[2, 3],
2658 ],
2659 4,
2660 );
2661 let budget = BpBudget {
2662 min_partition_docs: Some(4),
2663 time_budget: None,
2664 };
2665 let (perm, converged) = graph_bisection(&fwd, 2, 20, budget);
2666 assert!(converged, "depth cap must report converged");
2667 assert_eq!(perm.len(), 8);
2668 let mut sorted = perm.clone();
2669 sorted.sort();
2670 assert_eq!(
2671 sorted,
2672 (0..8).collect::<Vec<u32>>(),
2673 "must stay a valid permutation"
2674 );
2675 let cluster_a = [0u32, 1, 2, 4];
2678 let a_in_left = perm[..4].iter().filter(|d| cluster_a.contains(d)).count();
2679 assert!(
2680 a_in_left == 4 || a_in_left == 0,
2681 "clusters should separate at the top level: {:?}",
2682 perm
2683 );
2684 }
2685
2686 #[test]
2689 fn test_bp_zero_time_budget_emits_valid_partial_permutation() {
2690 let docs: Vec<Vec<u32>> = (0..64).map(|i| vec![i % 4]).collect();
2691 let doc_refs: Vec<&[u32]> = docs.iter().map(|v| v.as_slice()).collect();
2692 let fwd = make_fwd(&doc_refs, 4);
2693 let budget = BpBudget {
2694 min_partition_docs: None,
2695 time_budget: Some(std::time::Duration::ZERO),
2696 };
2697 let (perm, converged) = graph_bisection(&fwd, 4, 20, budget);
2698 assert!(!converged, "zero budget must report unconverged");
2699 assert_eq!(perm.len(), 64);
2700 let mut sorted = perm.clone();
2701 sorted.sort();
2702 assert_eq!(sorted, (0..64).collect::<Vec<u32>>());
2703 }
2704
2705 #[test]
2706 fn test_bp_shutdown_cancellation_emits_valid_partial_permutation() {
2707 let docs: Vec<Vec<u32>> = (0..64).map(|i| vec![i % 4]).collect();
2708 let doc_refs: Vec<&[u32]> = docs.iter().map(Vec::as_slice).collect();
2709 let fwd = make_fwd(&doc_refs, 4);
2710 let cancellation = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
2711 let budget = BpBudget {
2712 min_partition_docs: None,
2713 time_budget: None,
2714 };
2715
2716 let (perm, converged) = graph_bisection_with_progress(
2717 &fwd,
2718 4,
2719 20,
2720 budget,
2721 Some(cancellation.as_ref()),
2722 BpProgressLabel::anonymous(),
2723 );
2724
2725 assert!(!converged, "cancelled BP must report unconverged");
2726 assert_eq!(perm, (0..64).collect::<Vec<u32>>());
2727 }
2728
2729 #[test]
2730 fn test_memory_limited_graph_never_reports_converged() {
2731 let mut fwd = make_fwd(&[&[0], &[0], &[1], &[1]], 2);
2732 fwd.budget_limited = true;
2733
2734 let (perm, converged) = graph_bisection(&fwd, 2, 20, BpBudget::full());
2735
2736 assert!(!converged);
2737 let mut sorted = perm;
2738 sorted.sort_unstable();
2739 assert_eq!(sorted, vec![0, 1, 2, 3]);
2740 }
2741
2742 #[test]
2743 fn test_fast_log2() {
2744 let table = build_log_table(4096);
2745 assert!((table[1] - 0.0).abs() < 0.001);
2746 assert!((table[2] - 1.0).abs() < 0.001);
2747 assert!((table[4] - 2.0).abs() < 0.001);
2748 assert!((table[1024] - 10.0).abs() < 0.001);
2749 let val = fast_log2_lookup(8192, &table);
2751 assert!((val - 13.0).abs() < 0.001);
2752 }
2753}