Skip to main content

triplets_core/sampler/
mod.rs

1mod backends;
2
3use crate::chunking::{ChunkingAlgorithm, SlidingWindowChunker};
4use chrono::Duration;
5use indexmap::IndexMap;
6use rand::prelude::*;
7use rayon::prelude::*;
8use std::borrow::Cow;
9use std::collections::{HashMap, HashSet};
10use std::path::Path;
11use std::sync::atomic::{AtomicUsize, Ordering};
12use std::sync::mpsc;
13use std::sync::{Arc, Mutex};
14use std::thread;
15
16use crate::config::{
17    ChunkingStrategy, NegativeStrategy, SamplerConfig, Selector, TextRecipe, TripletRecipe,
18};
19use crate::constants::sampler::AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME;
20use crate::constants::sampler::{
21    ANCHOR_POSITIVE_SWAP_MASK, EPOCH_SEED_OFFSET, EXHAUSTION_RETRY_LIMIT, NEG_REASON_SAME_RECORD,
22    NEG_REASON_WRONG_ARTICLE, NEG_REASON_WRONG_DATE, NEG_REASON_WRONG_QA, PREFETCHER_SOURCE_ID,
23    PREFETCHER_STOPPED_REASON, RECIPE_LABEL_TEXT, RECIPE_LABEL_TRIPLETS,
24    RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER, ROLE_LABEL_ANCHOR, ROLE_LABEL_CONTEXT,
25    SAME_SELECTOR_PAIR_RETRY_LIMIT,
26};
27use crate::data::{
28    ChunkView, DataRecord, PairLabel, RecordChunk, RecordSection, SampleBatch, SamplePair,
29    SampleTriplet, SectionRole, TextBatch, TextSample, TripletBatch,
30};
31use crate::epoch::EpochTracker;
32use crate::errors::SamplerError;
33use crate::hash::{derive_epoch_seed, stable_hash_key, stable_hash_key_section, stable_hash_str};
34use crate::ingestion::IngestionManager;
35use crate::metadata::{META_FIELD_DATE, MetadataKey};
36use crate::metrics::{chunk_proximity_score, window_index_proximity};
37use crate::source::DataSource;
38use crate::splits::{
39    EpochStateStore, PersistedSamplerState, SamplerStateStore, SplitLabel, SplitStore,
40};
41use crate::tokenizer::{Tokenizer, WhitespaceTokenizer};
42use crate::types::{RecipeKey, RecordKey, SourceId};
43use crate::utils::platform_newline;
44
45// AUTO-RECIPE HANDLING OVERVIEW (end-to-end):
46// Stage A: Source-level injection eligibility ("should this source even get the recipe?")
47//   - `sync_records_from_cache` -> `record_has_long_anchor_or_context_section` marks
48//     `long_section_counts`.
49//   - `resolve_source_triplet_plan` -> `should_auto_inject_chunk_pair_recipe` appends
50//     `source_chunk_pair_recipe` for eligible sources.
51// Stage B: Record-level execution eligibility ("can this specific record run it now?")
52//   - `make_auto_chunk_pair_triplet_with_anchor` checks
53//     `record_has_at_least_two_window_chunks_for_selector` before sampling.
54// Stage C: Actual chunk window layout ("how are windows formed?")
55//   - `materialize_chunks` is the source of truth for token window construction.
56
57#[derive(Debug, Clone)]
58/// Small deterministic RNG used for reproducible sampler behavior.
59struct DeterministicRng {
60    state: u64,
61}
62
63impl DeterministicRng {
64    fn new(seed: u64) -> Self {
65        Self { state: seed }
66    }
67
68    fn from_state(state: u64) -> Self {
69        Self { state }
70    }
71
72    fn state(&self) -> u64 {
73        self.state
74    }
75
76    fn next_u64_internal(&mut self) -> u64 {
77        let mut z = self.state.wrapping_add(0x9E3779B97F4A7C15);
78        self.state = z;
79        z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
80        z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
81        z ^ (z >> 31)
82    }
83}
84
85impl rand::RngCore for DeterministicRng {
86    fn next_u32(&mut self) -> u32 {
87        self.next_u64_internal() as u32
88    }
89
90    fn next_u64(&mut self) -> u64 {
91        self.next_u64_internal()
92    }
93
94    fn fill_bytes(&mut self, dest: &mut [u8]) {
95        let mut offset = 0;
96        while offset < dest.len() {
97            let value = self.next_u64_internal();
98            let bytes = value.to_le_bytes();
99            let remaining = dest.len() - offset;
100            let copy_len = remaining.min(bytes.len());
101            dest[offset..offset + copy_len].copy_from_slice(&bytes[..copy_len]);
102            offset += copy_len;
103        }
104    }
105}
106
107/// Public helper so callers (and examples) can compute the same chunk weighting the sampler uses.
108///
109/// Definitions:
110/// - `trust` (range 0.0-1.0): comes from `RecordChunk.quality.trust`; scales the contribution of a chunk.
111/// - `base`: window chunks use proximity-to-head (`1 / (index + 1)`); summary chunks use their configured `summary_fallback_weight`.
112/// - `chunk_weight_floor`: minimum weight applied after scaling.
113///
114/// Formula: `max(chunk_weight_floor, base * trust)`.
115pub fn chunk_weight(strategy: &ChunkingStrategy, chunk: &RecordChunk) -> f32 {
116    let floor = strategy.chunk_weight_floor;
117    let trust = chunk.quality.trust.clamp(0.0, 1.0);
118    let base = match &chunk.view {
119        ChunkView::Window { index, .. } => window_index_proximity(*index),
120        ChunkView::SummaryFallback { weight, .. } => *weight,
121    };
122    (base * trust).max(floor)
123}
124
125/// Public sampling interface for pair, triplet, and text batch generation.
126pub trait Sampler {
127    /// Returns a batch of pairs. Consumes the shared epoch cursor for anchor selection.
128    ///
129    /// Samples all sources uniformly. Prefer [`Sampler::next_pair_batch_with_weights`]
130    /// with an explicit per-source weight map to honor a data mixture.
131    #[deprecated(
132        note = "samples all sources uniformly; call next_pair_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
133    )]
134    fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
135        self.next_pair_batch_with_weights(split, &HashMap::new())
136    }
137    /// Returns a batch of text samples. Consumes the shared epoch cursor for anchor selection.
138    ///
139    /// Samples all sources uniformly. Prefer [`Sampler::next_text_batch_with_weights`]
140    /// with an explicit per-source weight map to honor a data mixture.
141    #[deprecated(
142        note = "samples all sources uniformly; call next_text_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
143    )]
144    fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
145        self.next_text_batch_with_weights(split, &HashMap::new())
146    }
147    /// Returns a batch of triplets. Consumes the shared epoch cursor for anchor selection.
148    ///
149    /// Samples all sources uniformly. Prefer [`Sampler::next_triplet_batch_with_weights`]
150    /// with an explicit per-source weight map to honor a data mixture.
151    #[deprecated(
152        note = "samples all sources uniformly; call next_triplet_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
153    )]
154    fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
155        self.next_triplet_batch_with_weights(split, &HashMap::new())
156    }
157    /// Returns a batch of pairs with per-call source weights.
158    fn next_pair_batch_with_weights(
159        &self,
160        split: SplitLabel,
161        weights: &HashMap<SourceId, f32>,
162    ) -> Result<SampleBatch, SamplerError>;
163    /// Returns a batch of text samples with per-call source weights.
164    fn next_text_batch_with_weights(
165        &self,
166        split: SplitLabel,
167        weights: &HashMap<SourceId, f32>,
168    ) -> Result<TextBatch, SamplerError>;
169    /// Returns a batch of triplets with per-call source weights.
170    fn next_triplet_batch_with_weights(
171        &self,
172        split: SplitLabel,
173        weights: &HashMap<SourceId, f32>,
174    ) -> Result<TripletBatch, SamplerError>;
175    /// Persist sampler and split runtime state for restart-resume.
176    ///
177    /// Default is a no-op, allowing mock/test implementations to skip persistence.
178    fn save_sampler_state(&self, _save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
179        Ok(())
180    }
181}
182
183/// Background prefetcher that fills a bounded queue with sample batches.
184pub struct BatchPrefetcher<T> {
185    receiver: Option<mpsc::Receiver<Result<T, SamplerError>>>,
186    handle: Option<thread::JoinHandle<()>>,
187    stats: Arc<PrefetcherStats>,
188}
189
190#[derive(Default)]
191/// Prefetcher runtime counters.
192struct PrefetcherStats {
193    queued: AtomicUsize,
194    produced: AtomicUsize,
195    errors: AtomicUsize,
196}
197
198impl<T: Send + 'static> BatchPrefetcher<T> {
199    fn new<F>(capacity: usize, mut producer: F) -> Self
200    where
201        F: FnMut() -> Result<T, SamplerError> + Send + 'static,
202    {
203        let (sender, receiver) = mpsc::sync_channel(capacity.max(1));
204        let stats = Arc::new(PrefetcherStats::default());
205        let stats_thread = Arc::clone(&stats);
206        let handle = thread::spawn(move || {
207            loop {
208                let result = producer();
209                if result.is_err() {
210                    stats_thread.errors.fetch_add(1, Ordering::Relaxed);
211                }
212                if sender.send(result).is_err() {
213                    return;
214                }
215                stats_thread.queued.fetch_add(1, Ordering::Relaxed);
216                stats_thread.produced.fetch_add(1, Ordering::Relaxed);
217            }
218        });
219        Self {
220            receiver: Some(receiver),
221            handle: Some(handle),
222            stats,
223        }
224    }
225
226    /// Block until the next prefetched batch is available.
227    pub fn next(&self) -> Result<T, SamplerError> {
228        let receiver = self
229            .receiver
230            .as_ref()
231            .ok_or_else(|| SamplerError::SourceUnavailable {
232                source_id: PREFETCHER_SOURCE_ID.into(),
233                reason: PREFETCHER_STOPPED_REASON.into(),
234            })?;
235        let result = receiver.recv().unwrap_or_else(|_| {
236            Err(SamplerError::SourceUnavailable {
237                source_id: PREFETCHER_SOURCE_ID.into(),
238                reason: PREFETCHER_STOPPED_REASON.into(),
239            })
240        });
241        self.stats
242            .queued
243            .try_update(Ordering::Relaxed, Ordering::Relaxed, |value| {
244                Some(value.saturating_sub(1))
245            })
246            .ok();
247        result
248    }
249
250    /// Number of prefetched batches currently queued.
251    pub fn queue_len(&self) -> usize {
252        self.stats.queued.load(Ordering::Relaxed)
253    }
254
255    /// Total number of batches produced by the background worker.
256    pub fn produced_count(&self) -> usize {
257        self.stats.produced.load(Ordering::Relaxed)
258    }
259
260    /// Total number of errors produced by the background worker.
261    pub fn error_count(&self) -> usize {
262        self.stats.errors.load(Ordering::Relaxed)
263    }
264}
265
266impl<T> Drop for BatchPrefetcher<T> {
267    fn drop(&mut self) {
268        self.receiver.take();
269        if let Some(handle) = self.handle.take() {
270            let _ = handle.join();
271        }
272    }
273}
274
275/// Sampler that draws anchors from a single shared epoch cursor and then
276/// selects chunks from those records. Ingestion happens on demand when sampling.
277pub struct TripletSampler<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
278    inner: Mutex<TripletSamplerInner<S>>,
279}
280
281/// Internal sampler state implementation guarded by `TripletSampler`.
282struct TripletSamplerInner<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
283    /// Immutable sampler configuration (seed, batch size, recipes, splits, etc.).
284    config: SamplerConfig,
285    /// Active chunking implementation used to materialize section chunks.
286    chunker: Arc<dyn ChunkingAlgorithm>,
287    /// Split store backing train/val/test assignments and persisted sampler state.
288    split_store: Arc<S>,
289    /// On-demand ingestion manager that fills the batch-sized buffer.
290    ingestion: IngestionManager,
291    /// Current in-memory record pool keyed by composite record identity.
292    /// Composite keys (not bare ids) so overlapping ids from different
293    /// sources coexist instead of overwriting each other.
294    records: IndexMap<RecordKey, Arc<DataRecord>>,
295    /// Deficit counters for weighted source selection (smooth weighted
296    /// round-robin over `source_order`). Only engaged when the caller passes
297    /// explicit non-uniform per-source weights; uniform/empty maps take the
298    /// legacy uniform-cycling path bit-identically. New sources start at 0.0
299    /// via entry-or-insert (no sync hook needed); stale entries are pruned
300    /// with the source index. Intentionally NOT persisted: deficits
301    /// self-correct within a few draws, and no consumer persists sampler
302    /// state across processes with weights engaged.
303    source_deficits: HashMap<SourceId, f32>,
304    /// Deterministic RNG for per-batch shuffles and sampling.
305    rng: DeterministicRng,
306    /// Config-level triplet recipes used when sources do not supply their own.
307    triplet_recipes: Vec<TripletRecipe>,
308    /// Config-level text recipes used when sources do not supply their own.
309    text_recipes: Vec<TextRecipe>,
310    /// Per-source triplet recipes keyed by source id.
311    source_triplet_recipes: HashMap<SourceId, Vec<TripletRecipe>>,
312    /// Per-source count of pooled records containing at least one section
313    /// larger than the chunk window. A source is eligible for the
314    /// auto-injected chunk-pair recipe while its count is greater than zero.
315    /// Maintained incrementally on record insert/evict (strict O(Δ)) instead
316    /// of rescanning the pool on every sync.
317    long_section_counts: HashMap<SourceId, usize>,
318    /// Per-source text recipes keyed by source id.
319    source_text_recipes: HashMap<SourceId, Vec<TextRecipe>>,
320    /// True if triplet recipes came from config (no source defaults).
321    using_config_triplet_recipes: bool,
322    /// True if text recipes came from config (no derivation).
323    using_config_text_recipes: bool,
324    /// Last observed ingestion counter for cache updates.
325    last_observed_ingest: u64,
326    /// Epoch tracker for split-aware deterministic sampling.
327    epoch_tracker: EpochTracker,
328    /// Per-record, per-section chunk cursor to rotate through chunk windows.
329    chunk_cursors: HashMap<(RecordKey, usize), usize>,
330    /// Per-record, per-role cursor to rotate through role-specific sections.
331    role_cursors: HashMap<(RecordKey, String), usize>,
332    /// Pluggable negative-selection backend (uniform-random or BM25).
333    negative_backend: Box<dyn backends::NegativeBackend>,
334    /// Chunk key to parent record key lookup (used by epoch tracker).
335    chunk_index: HashMap<RecordKey, RecordKey>,
336    /// Round-robin order of source ids (deterministic).
337    source_order: Vec<SourceId>,
338    /// Current index into `source_order` for shuffled-cycle sampling.
339    source_cycle_idx: usize,
340    /// True once persisted source state is loaded.
341    source_state_loaded: bool,
342    /// True once persisted ingestion cursors are loaded.
343    ingestion_cursors_loaded: bool,
344    /// Marks whether source state needs persistence.
345    source_state_dirty: bool,
346    /// Record keys per source for round-robin within a source.
347    source_record_indices: HashMap<SourceId, Vec<RecordKey>>,
348    /// Per-source cursor into `source_record_indices`.
349    source_record_cursors: HashMap<SourceId, usize>,
350    /// Round-robin index for triplet recipe cycling.
351    triplet_recipe_rr_idx: usize,
352    /// Hashes of text already emitted by the text batch path while the record
353    /// pool is stable.  Pruned on sync_records_from_cache only when records
354    /// are evicted.  Prevents the same text from being sampled again within
355    /// the same ingestion window, regardless of source wrapping or epoch.
356    ///
357    /// Per-record tracking lets us surgically remove hashes only for evicted
358    /// records, preserving cross-batch dedup across small window advances
359    /// (e.g. window=512, batch_size=4 where only 4/512 records are swapped).
360    emitted_text_hashes: HashMap<RecordKey, HashSet<u64>>,
361    /// Round-robin index for text recipe cycling.
362    text_recipe_rr_idx: usize,
363    /// Epoch counter for per-source deterministic shuffling (seed ^ epoch).
364    epoch: u64,
365
366    /// Tracks whether each source has wrapped its cursor in the current epoch.
367    source_wrapped: HashMap<SourceId, bool>,
368
369    /// Last sync version watermark for incremental delta synchronization.
370    last_sync_version: u64,
371
372    /// Monotonic counter advanced only when record-pool membership actually
373    /// changes (genuinely new IDs added or IDs truly evicted). Pure version
374    /// refreshes of already-pooled IDs do not advance it, so the cached
375    /// `records_by_split` map below stays valid across steady-state advances.
376    pool_generation: u64,
377
378    /// Cached `records_by_split` output, valid while
379    /// `cached_records_by_split_gen == pool_generation`. Avoids rebuilding the
380    /// full per-split map (O(N) clones + label lookups) on every batch.
381    cached_records_by_split: Option<HashMap<SplitLabel, Vec<RecordKey>>>,
382    /// Pool generation the cached map was built from.
383    cached_records_by_split_gen: u64,
384
385    /// Last observed force-refresh generation. When the ingestion manager
386    /// reports a newer force-refresh generation, the next sync falls back to
387    /// a full pool resync (exact full-snapshot semantics) instead of the
388    /// O(Δ) delta path, because a force refresh replaces the cache contents
389    /// wholesale. Steady-state advances keep the delta fast path.
390    last_force_refresh_generation: u64,
391
392    /// Persistent cache of split labels keyed by composite record identity.
393    /// Eliminates redundant DataStore::read() calls for split label lookups.
394    split_labels: HashMap<RecordKey, SplitLabel>,
395}
396
397impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSamplerInner<S> {
398    fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
399        Self::new_with_chunker(config, split_store, Arc::new(SlidingWindowChunker))
400    }
401
402    fn new_with_chunker(
403        config: SamplerConfig,
404        split_store: Arc<S>,
405        chunker: Arc<dyn ChunkingAlgorithm>,
406    ) -> Self {
407        let buffer_size = config.ingestion_max_records.max(config.batch_size).max(2);
408        let using_config_triplet_recipes = !config.recipes.is_empty();
409        let using_config_text_recipes = !config.text_recipes.is_empty();
410        let triplet_recipes = if using_config_triplet_recipes {
411            config.recipes.clone()
412        } else {
413            Vec::new()
414        };
415        let text_recipes = if using_config_text_recipes {
416            config.text_recipes.clone()
417        } else if !triplet_recipes.is_empty() {
418            Self::build_derived_text_recipes(&triplet_recipes)
419        } else {
420            Vec::new()
421        };
422        let ingestion = IngestionManager::new(buffer_size, config.clone());
423        let epoch_backend = Some(Arc::clone(&split_store) as Arc<dyn EpochStateStore>);
424        let epoch_tracker = EpochTracker::new(
425            true,
426            epoch_backend,
427            derive_epoch_seed(config.seed, EPOCH_SEED_OFFSET),
428        );
429        let mut sampler = Self {
430            rng: DeterministicRng::new(config.seed),
431            config,
432            chunker,
433            split_store,
434            ingestion,
435            records: IndexMap::new(),
436            triplet_recipes,
437            text_recipes,
438            source_triplet_recipes: HashMap::new(),
439            long_section_counts: HashMap::new(),
440            source_text_recipes: HashMap::new(),
441            using_config_triplet_recipes,
442            using_config_text_recipes,
443            last_observed_ingest: 0,
444            epoch_tracker,
445            chunk_cursors: HashMap::new(),
446            role_cursors: HashMap::new(),
447            negative_backend: {
448                #[cfg(feature = "bm25-mining")]
449                {
450                    Box::new(backends::Bm25Backend::new())
451                }
452                #[cfg(not(feature = "bm25-mining"))]
453                {
454                    Box::new(backends::DefaultBackend)
455                }
456            },
457            chunk_index: HashMap::new(),
458            source_order: Vec::new(),
459            source_cycle_idx: 0,
460            source_deficits: HashMap::new(),
461            source_state_loaded: false,
462            ingestion_cursors_loaded: false,
463            source_state_dirty: false,
464            source_record_indices: HashMap::new(),
465            source_record_cursors: HashMap::new(),
466            emitted_text_hashes: HashMap::new(),
467            triplet_recipe_rr_idx: 0,
468            text_recipe_rr_idx: 0,
469            epoch: 0,
470            source_wrapped: HashMap::new(),
471            last_sync_version: 0,
472            last_force_refresh_generation: 0,
473            pool_generation: 0,
474            cached_records_by_split: None,
475            cached_records_by_split_gen: 0,
476            split_labels: HashMap::new(),
477        };
478        if !sampler.using_config_text_recipes {
479            sampler.rebuild_derived_text_recipes();
480        }
481        sampler
482    }
483
484    fn text_recipes(&self) -> &[TextRecipe] {
485        &self.text_recipes
486    }
487
488    /// Current epoch-adjusted seed: mixes `epoch` into `config.seed` so every epoch
489    /// produces a distinct permutation across all seed-dependent operations.
490    fn epoch_seed(&self) -> u64 {
491        derive_epoch_seed(self.config.seed, self.epoch)
492    }
493
494    /// Get or compute the split label for a record, using the persistent cache.
495    fn get_or_insert_split_label(&mut self, key: &RecordKey) -> Result<SplitLabel, SamplerError> {
496        if let Some(&label) = self.split_labels.get(key) {
497            return Ok(label);
498        }
499        let label = match self.split_store.label_for(key) {
500            Some(label) => label,
501            None => self.split_store.ensure(key.clone())?,
502        };
503        self.split_labels.insert(key.clone(), label);
504        Ok(label)
505    }
506
507    fn register_source(
508        &mut self,
509        source: Box<dyn DataSource + 'static>,
510    ) -> Result<(), SamplerError> {
511        let source_id = source.id().to_string();
512        if !self.using_config_triplet_recipes {
513            let triplets = source.default_triplet_recipes();
514            if !triplets.is_empty() {
515                self.source_triplet_recipes
516                    .insert(source_id.clone(), triplets.clone());
517                if !self.using_config_text_recipes {
518                    let derived = Self::build_derived_text_recipes(&triplets);
519                    self.source_text_recipes
520                        .insert(source_id.clone(), derived.clone());
521                    self.extend_text_recipes_unique(&derived);
522                }
523            }
524        }
525        self.ingestion.register_source(source)?;
526        Ok(())
527    }
528
529    fn set_epoch(&mut self, epoch: u64) -> Result<(), SamplerError> {
530        self.epoch_tracker.ensure_loaded()?;
531        self.epoch_tracker.force_epoch(epoch);
532        self.epoch = epoch;
533        self.ingestion.set_epoch(epoch);
534        self.ingestion.reset_stream_cursors();
535        // Reset epoch step at epoch boundary so each epoch
536        // starts with step=0, giving deterministic step sequences.
537        self.ingestion.reset_epoch_step();
538        self.source_record_cursors.clear();
539        self.source_cycle_idx = 0;
540        for source in &self.source_order {
541            self.source_wrapped.insert(source.clone(), false);
542        }
543        self.rebuild_source_index()?;
544        self.source_state_dirty = self.source_order.len() > 1;
545        Ok(())
546    }
547
548    fn next_chunk_from_pool(
549        &mut self,
550        key: &RecordKey,
551        section_idx: usize,
552        pool: Vec<RecordChunk>,
553    ) -> Option<RecordChunk> {
554        if pool.is_empty() {
555            return None;
556        }
557        let map_key = (key.clone(), section_idx);
558        if !self.chunk_cursors.contains_key(&map_key) {
559            // First touch should **not** always start at window 0.
560            // We derive a deterministic per-(record,section) starting offset so:
561            // - repeated runs with the same seed/epoch are reproducible,
562            // - first sampled window is spread across records/sections,
563            // - subsequent calls still rotate cyclically through the pool.
564            let cursor_key = stable_hash_key_section(self.epoch_seed(), key, section_idx);
565            let start = (cursor_key as usize) % pool.len();
566            self.chunk_cursors.insert(map_key.clone(), start);
567        }
568        let cursor = self.chunk_cursors.entry(map_key).or_insert(0);
569        if *cursor >= pool.len() {
570            *cursor = 0;
571        }
572        let chunk = pool.get(*cursor).cloned();
573        *cursor = (*cursor + 1) % pool.len();
574        chunk
575    }
576
577    fn prune_cursor_state(&mut self) {
578        if self.chunk_cursors.is_empty()
579            && self.role_cursors.is_empty()
580            && self.negative_backend.cursors_empty()
581        {
582            return;
583        }
584        let valid_keys: HashSet<RecordKey> = self.records.keys().cloned().collect();
585        self.chunk_cursors
586            .retain(|(record_key, _), _| valid_keys.contains(record_key));
587        self.role_cursors
588            .retain(|(record_key, _), _| valid_keys.contains(record_key));
589        self.negative_backend.prune_cursors(&valid_keys);
590    }
591
592    /// Ordered removal of a record id from one source's index vector.
593    /// Empty source entries are left in place; the caller prunes them after
594    /// all evictions and additions are processed (mirroring
595    /// `rebuild_source_index`, which only contains sources with at least one
596    /// allowed-split record).
597    fn remove_id_from_source_index(&mut self, source: &SourceId, record_key: &RecordKey) {
598        if let Some(ids) = self.source_record_indices.get_mut(source)
599            && let Some(pos) = ids.iter().position(|key| key == record_key)
600        {
601            ids.remove(pos);
602        }
603    }
604
605    /// Insert a record key into one source's index vector, preserving the
606    /// deterministic `stable_hash_key(epoch_seed, key)` sort order via binary
607    /// search. Records whose split is not allowed are not indexed, mirroring
608    /// `rebuild_source_index`. Creates the source entry when needed; the
609    /// caller keeps `source_order` sorted afterwards.
610    fn insert_id_into_source_index(
611        &mut self,
612        source: SourceId,
613        record_key: RecordKey,
614        label: SplitLabel,
615        allowed: &HashSet<SplitLabel>,
616    ) {
617        if !allowed.contains(&label) {
618            return;
619        }
620        let seed = self.epoch_seed();
621        let key = stable_hash_key(seed, &record_key);
622        // Upper-bound insertion: with distinct hashes this is the unique
623        // sorted position; on a hash collision it places the key after
624        // existing equals, matching stable-sort semantics.
625        let pos = match self.source_record_indices.get(&source) {
626            Some(ids) => ids.partition_point(|k| stable_hash_key(seed, k) <= key),
627            None => 0,
628        };
629        self.source_record_indices
630            .entry(source)
631            .or_default()
632            .insert(pos, record_key);
633    }
634
635    /// Increment the long-section record count for a source.
636    fn increment_long_section_count(&mut self, source: &SourceId) {
637        *self.long_section_counts.entry(source.clone()).or_default() += 1;
638    }
639
640    /// Decrement the long-section record count for a source, removing the
641    /// entry when no counted records remain.
642    fn decrement_long_section_count(&mut self, source: &SourceId) {
643        let remove = if let Some(count) = self.long_section_counts.get_mut(source) {
644            *count = count.saturating_sub(1);
645            *count == 0
646        } else {
647            false
648        };
649        if remove {
650            self.long_section_counts.remove(source);
651        }
652    }
653
654    fn rebuild_chunk_index(&mut self) {
655        self.chunk_index.clear();
656        for record in self.records.values() {
657            let key = RecordKey::from(record);
658            self.chunk_index.insert(key.clone(), key);
659        }
660    }
661
662    fn rebuild_source_index(&mut self) -> Result<(), SamplerError> {
663        self.source_record_indices.clear();
664        let allowed = self.allowed_target_splits();
665        let allowed_set: HashSet<SplitLabel> = allowed.into_iter().collect();
666        // Collect record keys first to avoid borrowing issues
667        let record_keys: Vec<RecordKey> = self.records.values().map(RecordKey::of_arc).collect();
668        for record_key in record_keys {
669            let label = self.get_or_insert_split_label(&record_key)?;
670            if !allowed_set.contains(&label) {
671                continue;
672            }
673            self.source_record_indices
674                .entry(record_key.source.clone())
675                .or_default()
676                .push(record_key);
677        }
678
679        let shuffle_seed = self.epoch_seed();
680        for record_keys in self.source_record_indices.values_mut() {
681            record_keys.sort_by_key(|key| stable_hash_key(shuffle_seed, key));
682        }
683
684        self.source_order = self.source_record_indices.keys().cloned().collect();
685        self.source_order.sort();
686        self.refresh_source_wrapped();
687        // Drop deficit counters for sources that left the pool.
688        self.source_deficits
689            .retain(|s, _| self.source_record_indices.contains_key(s));
690
691        self.source_record_cursors
692            .retain(|source, _| self.source_record_indices.contains_key(source));
693        if self.source_state_loaded {
694            if self.source_order.is_empty() {
695                self.source_cycle_idx = 0;
696            }
697            self.source_state_dirty = self.source_order.len() > 1;
698        }
699        Ok(())
700    }
701
702    fn refresh_source_wrapped(&mut self) {
703        self.source_wrapped.clear();
704        for source in &self.source_order {
705            let len = self
706                .source_record_indices
707                .get(source)
708                .map(|items| items.len())
709                .unwrap_or(0);
710            if len == 0 {
711                self.source_wrapped.insert(source.clone(), false);
712                continue;
713            }
714            let cursor = self.source_record_cursors.get(source).copied().unwrap_or(0);
715            let wrapped = cursor > 0 && cursor % len == 0;
716            self.source_wrapped.insert(source.clone(), wrapped);
717        }
718    }
719
720    fn shuffled_source_cycle(&self, cycle: u64) -> Vec<SourceId> {
721        let mut sources = self.source_order.clone();
722        let seed = self.epoch_seed() ^ cycle;
723        sources.sort_by_key(|source| stable_hash_str(seed, source));
724        sources
725    }
726
727    /// True when `weights` carries no selective signal: empty or all equal.
728    /// Uniform maps take the legacy uniform-cycling path bit-identically.
729    fn weights_are_selective(weights: &HashMap<SourceId, f32>) -> bool {
730        if weights.is_empty() {
731            return false;
732        }
733        let first = *weights.values().next().unwrap();
734        weights.values().any(|&w| w != first)
735    }
736
737    /// Effective quantum for `source`: its weight when finite and
738    /// non-negative, else 0. Unlisted sources default to 1.0 (mirrors the
739    /// ingestion drain).
740    fn weight_quantum(weights: &HashMap<SourceId, f32>, source: &str) -> f32 {
741        weights.get(source).copied().unwrap_or(1.0).max(0.0)
742    }
743
744    /// Weighted source choice: smooth weighted round-robin over `candidates`.
745    /// Credits each eligible source its weight quantum, picks the highest
746    /// deficit, deducts the round total. Deterministic given pool + deficits
747    /// (no RNG): same state, same choice. New sources start at 0.0 via
748    /// entry-or-insert. Returns `None` when no candidate has records.
749    fn select_weighted_source(
750        &mut self,
751        candidates: &[SourceId],
752        weights: &HashMap<SourceId, f32>,
753    ) -> Option<SourceId> {
754        let eligible: Vec<&SourceId> = candidates
755            .iter()
756            .filter(|s| {
757                self.source_record_indices
758                    .get(*s)
759                    .map(|v| !v.is_empty())
760                    .unwrap_or(false)
761            })
762            .collect();
763        if eligible.is_empty() {
764            return None;
765        }
766        let mut total = 0.0f32;
767        for source in &eligible {
768            let q = Self::weight_quantum(weights, source);
769            total += q;
770            *self.source_deficits.entry((*source).clone()).or_insert(0.0) += q;
771        }
772        if total <= 0.0 {
773            return None;
774        }
775        let mut best: Option<(f32, SourceId)> = None;
776        for source in &eligible {
777            let d = self.source_deficits.get(*source).copied().unwrap_or(0.0);
778            if best.as_ref().map(|(bd, _)| d > *bd).unwrap_or(true) {
779                best = Some((d, (*source).clone()));
780            }
781        }
782        let (_, source) = best?;
783        *self
784            .source_deficits
785            .get_mut(&source)
786            .expect("deficit credited above") -= total;
787        Some(source)
788    }
789
790    fn ensure_source_state(&mut self) -> Result<(), SamplerError> {
791        if self.source_state_loaded {
792            return Ok(());
793        }
794        let persisted = self.split_store.load_sampler_state()?;
795        self.source_cycle_idx = persisted
796            .as_ref()
797            .map(|state| state.source_cycle_idx as usize)
798            .unwrap_or(0);
799        if let Some(state) = persisted {
800            for (source, cursor) in state.source_record_cursors {
801                if self.source_record_indices.contains_key(&source) {
802                    self.source_record_cursors.insert(source, cursor as usize);
803                }
804            }
805            self.epoch = state.epoch;
806            self.ingestion.set_epoch(state.epoch);
807            self.rng = DeterministicRng::from_state(state.rng_state);
808            self.triplet_recipe_rr_idx = state.triplet_recipe_rr_idx as usize;
809            self.text_recipe_rr_idx = state.text_recipe_rr_idx as usize;
810        }
811        self.refresh_source_wrapped();
812        self.source_state_loaded = true;
813        self.source_state_dirty = true;
814        Ok(())
815    }
816
817    fn persist_source_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
818        if !self.source_state_loaded {
819            return Ok(());
820        }
821        let state = PersistedSamplerState {
822            source_cycle_idx: self.source_cycle_idx as u64,
823            source_record_cursors: self
824                .source_record_cursors
825                .iter()
826                .map(|(source, cursor)| (source.clone(), *cursor as u64))
827                .collect(),
828            epoch: self.epoch,
829            epoch_step: self.ingestion.epoch_step(),
830            rng_state: self.rng.state(),
831            triplet_recipe_rr_idx: self.triplet_recipe_rr_idx as u64,
832            text_recipe_rr_idx: self.text_recipe_rr_idx as u64,
833            source_stream_cursors: self.ingestion.snapshot_cursors(),
834        };
835        self.split_store.save_sampler_state(&state, save_to)?;
836        self.source_state_dirty = false;
837        Ok(())
838    }
839
840    fn rebuild_derived_text_recipes(&mut self) {
841        if self.using_config_text_recipes {
842            return;
843        }
844        if self.triplet_recipes.is_empty() {
845            self.text_recipes.clear();
846        } else {
847            self.text_recipes = Self::build_derived_text_recipes(&self.triplet_recipes);
848        }
849    }
850
851    fn extend_text_recipes_unique(&mut self, recipes: &[TextRecipe]) {
852        for recipe in recipes {
853            if self
854                .text_recipes
855                .iter()
856                .any(|existing| existing.name == recipe.name)
857            {
858                continue;
859            }
860            self.text_recipes.push(recipe.clone());
861        }
862    }
863
864    fn configured_triplet_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TripletRecipe] {
865        if self.using_config_triplet_recipes {
866            return &self.triplet_recipes;
867        }
868        self.source_triplet_recipes
869            .get(source)
870            .map(|recipes| recipes.as_slice())
871            .unwrap_or(&[])
872    }
873
874    /// Returns true when `recipes` already contains the long-section auto recipe.
875    fn contains_auto_chunk_pair_recipe(recipes: &[TripletRecipe]) -> bool {
876        recipes
877            .iter()
878            .any(|recipe| recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME)
879    }
880
881    fn source_supports_chunk_pair_recipe(&self, source: &str) -> bool {
882        if self.config.chunking.max_window_tokens == 0 {
883            return false;
884        }
885        self.long_section_counts
886            .get(source)
887            .is_some_and(|&count| count > 0)
888    }
889
890    /// Stage A (source-level): decide whether to append auto recipe for `source`.
891    ///
892    /// Decision criteria:
893    /// - source must be eligible (`source_supports_chunk_pair_recipe`), and
894    /// - configured pool must not already include the auto recipe name.
895    fn should_auto_inject_chunk_pair_recipe(
896        &self,
897        source: &str,
898        recipes: &[TripletRecipe],
899    ) -> bool {
900        self.source_supports_chunk_pair_recipe(source)
901            && !Self::contains_auto_chunk_pair_recipe(recipes)
902    }
903
904    /// Build the auto-injected long-section recipe.
905    ///
906    /// Semantics:
907    /// - Anchor selector: `Context`
908    /// - Positive selector: `Context`
909    /// - Negative selector: `Context` on a different record (`WrongArticle`)
910    ///
911    /// Anchor and positive are selected by two independent `select_chunk` calls
912    /// from the same context chunk candidate pool for the chosen record.
913    /// They are not concatenated and one is not derived from the other's text.
914    fn source_chunk_pair_recipe() -> TripletRecipe {
915        TripletRecipe {
916            name: Cow::Borrowed(AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME),
917            anchor: Selector::Role(SectionRole::Context),
918            positive_selector: Selector::Role(SectionRole::Context),
919            negative_selector: Selector::Role(SectionRole::Context),
920            negative_strategy: NegativeStrategy::WrongArticle,
921            weight: 1.0,
922            instruction: None,
923            allow_same_anchor_positive: false,
924        }
925    }
926
927    /// Resolve the source triplet plan.
928    ///
929    /// This is Stage A (source-level injection), not record-level execution.
930    ///
931    /// Algorithm for the auto-triplet injection path:
932    /// 1) Start from the configured source recipe pool (`configured_triplet_recipes_for_source`).
933    /// 2) Check auto-injection eligibility (`should_auto_inject_chunk_pair_recipe`), which is true only when:
934    ///    - the source has long sections discovered during ingest sync, and
935    ///    - the pool does not already include the auto recipe name.
936    /// 3) If eligible, append `source_chunk_pair_recipe`.
937    /// 4) Return both the effective recipe pool and whether step (3) happened.
938    fn resolve_source_triplet_plan(&self, source: &str) -> (Vec<TripletRecipe>, bool) {
939        let mut recipes = self.configured_triplet_recipes_for_source(source).to_vec();
940        let mut auto_injected = false;
941        if self.should_auto_inject_chunk_pair_recipe(source, &recipes) {
942            recipes.push(Self::source_chunk_pair_recipe());
943            auto_injected = true;
944        }
945        (recipes, auto_injected)
946    }
947
948    #[cfg(test)]
949    fn triplet_recipes_for_source(&self, source: &str) -> Vec<TripletRecipe> {
950        self.resolve_source_triplet_plan(source).0
951    }
952
953    fn triplet_recipe_count_for_source(&self, source: &str) -> usize {
954        let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
955        recipes.len()
956    }
957
958    fn text_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TextRecipe] {
959        if self.using_config_text_recipes || self.using_config_triplet_recipes {
960            return &self.text_recipes;
961        }
962        self.source_text_recipes
963            .get(source)
964            .map(|recipes| recipes.as_slice())
965            .unwrap_or(&[])
966    }
967
968    /// Build a weighted, shuffled selection order from a slice of recipe weights.
969    ///
970    /// Each recipe with `weight > 0.0` receives a number of slots proportional to its weight
971    /// relative to the smallest positive weight, capped at
972    /// [`RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER`].  The slot list is then shuffled so that
973    /// deterministic round-robin cycling naturally draws recipes at the requested frequency:
974    ///
975    /// * `weight = [1.0, 1.0, 1.0]` → one slot each — identical to uniform round-robin.
976    /// * `weight = [3.0, 1.0]`       → three slots for recipe 0, one for recipe 1.
977    /// * `weight = [1.0, 0.0]`       → recipe 1 is excluded entirely (zero slots).
978    fn recipe_order_weighted_shuffled(
979        &mut self,
980        weights: &[f32],
981        rng: &mut DeterministicRng,
982    ) -> Vec<usize> {
983        weighted_recipe_order(weights, rng)
984    }
985
986    /// Weighted shuffle rotated by `rr_idx` for deterministic round-robin cycling.
987    ///
988    /// Calls [`Self::recipe_order_weighted_shuffled`] and rotates the result by
989    /// `rr_idx % order.len()` so successive batches start from different recipes,
990    /// preserving both the proportional frequency and the cross-batch cycling.
991    fn recipe_order_weighted_cycled(
992        &mut self,
993        weights: &[f32],
994        rr_idx: usize,
995        rng: &mut DeterministicRng,
996    ) -> Vec<usize> {
997        let base = self.recipe_order_weighted_shuffled(weights, rng);
998        if base.is_empty() {
999            return base;
1000        }
1001        let start = rr_idx % base.len();
1002        let mut order = Vec::with_capacity(base.len());
1003        order.extend_from_slice(&base[start..]);
1004        order.extend_from_slice(&base[..start]);
1005        order
1006    }
1007
1008    /// Weighted shuffled selection order for text recipes — same algorithm as
1009    /// [`Self::recipe_order_weighted_shuffled`], kept separate so the text and triplet
1010    /// round-robin counters can advance independently.
1011    fn text_recipe_order_weighted_shuffled(
1012        &mut self,
1013        weights: &[f32],
1014        rng: &mut DeterministicRng,
1015    ) -> Vec<usize> {
1016        weighted_recipe_order(weights, rng)
1017    }
1018
1019    fn text_recipe_order_weighted_cycled(
1020        &mut self,
1021        weights: &[f32],
1022        rr_idx: usize,
1023        rng: &mut DeterministicRng,
1024    ) -> Vec<usize> {
1025        let base = self.text_recipe_order_weighted_shuffled(weights, rng);
1026        if base.is_empty() {
1027            return base;
1028        }
1029        let start = rr_idx % base.len();
1030        let mut order = Vec::with_capacity(base.len());
1031        order.extend_from_slice(&base[start..]);
1032        order.extend_from_slice(&base[..start]);
1033        order
1034    }
1035
1036    fn allowed_target_splits(&self) -> Vec<SplitLabel> {
1037        self.config.allowed_splits.clone()
1038    }
1039
1040    fn ensure_split_allowed(&self, split: SplitLabel) -> Result<(), SamplerError> {
1041        let allowed = self.allowed_target_splits();
1042        if allowed.contains(&split) {
1043            return Ok(());
1044        }
1045        Err(SamplerError::Configuration(format!(
1046            "requested split {:?} is not in allowed_splits {:?}",
1047            split, allowed
1048        )))
1049    }
1050
1051    fn ensure_split_has_records(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
1052        self.refresh_records_by_split_cache()?;
1053        let has_records = self
1054            .cached_records_by_split
1055            .as_ref()
1056            .and_then(|map| map.get(&target_split))
1057            .is_some_and(|records| !records.is_empty());
1058        if has_records {
1059            return Ok(());
1060        }
1061        Err(SamplerError::Exhausted(
1062            "no records available for target split".into(),
1063        ))
1064    }
1065
1066    /// Rebuild the cached `records_by_split` map when the record pool changed
1067    /// since the last build. Steady-state batches hit the cache (O(1)) instead
1068    /// of rebuilding the full map (O(N) clones + label lookups).
1069    fn refresh_records_by_split_cache(&mut self) -> Result<(), SamplerError> {
1070        if self.cached_records_by_split_gen == self.pool_generation
1071            && self.cached_records_by_split.is_some()
1072        {
1073            return Ok(());
1074        }
1075        let mut map: HashMap<SplitLabel, Vec<RecordKey>> = HashMap::new();
1076        // Collect chunk_index entries first to avoid borrowing issues
1077        let entries: Vec<(RecordKey, RecordKey)> = self
1078            .chunk_index
1079            .iter()
1080            .map(|(k, v)| (k.clone(), v.clone()))
1081            .collect();
1082        for (chunk_key, record_key) in entries {
1083            if !self.records.contains_key(&record_key) {
1084                continue;
1085            }
1086            let label = self.get_or_insert_split_label(&record_key)?;
1087            map.entry(label).or_default().push(chunk_key);
1088        }
1089        self.cached_records_by_split = Some(map);
1090        self.cached_records_by_split_gen = self.pool_generation;
1091        Ok(())
1092    }
1093
1094    /// Owned-map convenience used by tests that hand-mutate the pool.
1095    /// Production paths use [`Self::refresh_records_by_split_cache`] plus a
1096    /// direct reference to avoid the per-call clone.
1097    #[cfg(test)]
1098    fn records_by_split(&mut self) -> Result<HashMap<SplitLabel, Vec<RecordKey>>, SamplerError> {
1099        self.refresh_records_by_split_cache()?;
1100        Ok(self.cached_records_by_split.clone().unwrap_or_default())
1101    }
1102
1103    fn choose_anchor_record(
1104        &mut self,
1105        source: Option<&str>,
1106        split: SplitLabel,
1107    ) -> Option<Arc<DataRecord>> {
1108        if let Some(source) = source {
1109            let record_ids = self.source_record_indices.get(source)?;
1110            if record_ids.is_empty() {
1111                return None;
1112            }
1113            let len = record_ids.len();
1114            let mut cursor = *self.source_record_cursors.get(source).unwrap_or(&0);
1115            let cycle = cursor / len;
1116            let offset_seed = self.epoch_seed() ^ (cycle as u64);
1117            let offset = (stable_hash_str(offset_seed, source) as usize) % len;
1118            let mut wrapped = false;
1119            let mut selected_key: Option<RecordKey> = None;
1120
1121            for _ in 0..len {
1122                let pos = (cursor % len).saturating_add(offset) % len;
1123                let record_key = &record_ids[pos];
1124                cursor = cursor.saturating_add(1);
1125                if cursor.is_multiple_of(len) {
1126                    wrapped = true;
1127                }
1128                if self.split_labels.get(record_key) == Some(&split) {
1129                    selected_key = Some(record_key.clone());
1130                    break;
1131                }
1132            }
1133
1134            self.source_record_cursors
1135                .insert(source.to_string(), cursor);
1136            if wrapped {
1137                self.mark_source_wrapped(source);
1138            }
1139
1140            selected_key.and_then(|key| self.records.get(&key).map(Arc::clone))
1141        } else {
1142            while let Some(chunk_key) = self.epoch_tracker.next_record(split) {
1143                if let Some(record_key) = self.chunk_index.get(&chunk_key)
1144                    && let Some(record) = self.records.get(record_key)
1145                {
1146                    return Some(Arc::clone(record));
1147                }
1148            }
1149            None
1150        }
1151    }
1152
1153    fn save_sampler_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
1154        if self.epoch_tracker.is_enabled() {
1155            self.epoch_tracker.persist()?;
1156        }
1157        self.persist_source_state(save_to)?;
1158        Ok(())
1159    }
1160
1161    fn mark_source_wrapped(&mut self, source: &str) {
1162        self.source_wrapped.insert(source.to_string(), true);
1163        if self.source_order.is_empty() {
1164            return;
1165        }
1166        let all_wrapped = self
1167            .source_order
1168            .iter()
1169            .all(|name| self.source_wrapped.get(name).copied().unwrap_or(false));
1170        if all_wrapped {
1171            self.advance_epoch();
1172        }
1173    }
1174
1175    fn advance_epoch(&mut self) {
1176        // The epoch advanced, changing the sampling permutation.
1177        // emitted_text_hashes is NOT cleared: the record pool is unchanged,
1178        // so texts already emitted are still off-limits.  The epoch
1179        // changes the shuffling permutation, not the dedup window.
1180        // Advance the epoch counter so source refreshes receive a different
1181        // permutation seed (seed ^ epoch), producing a fresh traversal order.
1182        self.epoch = self.epoch.saturating_add(1);
1183        self.ingestion.set_epoch(self.epoch);
1184        // Do NOT reset the epoch step here: it tracks the number of
1185        // refresh batches, which is orthogonal to the sampling epoch.
1186        // Resetting it would corrupt the persisted cursor value on resume.
1187        // Reset per-source record cursors so the new epoch starts from the
1188        // beginning of each source's permuted index space (not from where the
1189        // previous epoch left off).
1190        self.source_record_cursors.clear();
1191        self.source_cycle_idx = 0;
1192        // Reset the wrapped-tracking so all sources are eligible for sampling
1193        // in the new epoch.
1194        for source in &self.source_order {
1195            self.source_wrapped.insert(source.clone(), false);
1196        }
1197        let _ = self.rebuild_source_index();
1198        self.source_state_dirty = self.source_order.len() > 1;
1199    }
1200
1201    fn select_temporal_neighbor(
1202        &'_ self,
1203        record: &DataRecord,
1204        offset_days: i32,
1205    ) -> Option<Arc<DataRecord>> {
1206        let target = record.created_at + Duration::days(offset_days.into());
1207        let key = record.taxonomy.first().cloned();
1208        let record_split = self.split_store.label_for(&RecordKey::from(record))?;
1209        self.records
1210            .values()
1211            .filter(|candidate| {
1212                RecordKey::from(*candidate) != RecordKey::from(record)
1213                    && self
1214                        .split_store
1215                        .label_for(&RecordKey::from(*candidate))
1216                        .map(|label| label == record_split)
1217                        .unwrap_or(false)
1218                    && (candidate.source == record.source
1219                        || key
1220                            .as_ref()
1221                            .zip(candidate.taxonomy.first())
1222                            .map(|(a, b)| a == b)
1223                            .unwrap_or(false))
1224            })
1225            .min_by_key(|candidate| (candidate.created_at - target).num_seconds().abs())
1226            .cloned()
1227    }
1228
1229    fn select_negative_record(
1230        &self,
1231        anchor_record: &DataRecord,
1232        strategy: &NegativeStrategy,
1233        anchor_query_text: Option<&str>,
1234        rng: &mut dyn rand::RngCore,
1235    ) -> Option<(Arc<DataRecord>, bool)> {
1236        let anchor_split = self
1237            .split_store
1238            .label_for(&RecordKey::from(anchor_record))?;
1239
1240        let in_anchor_split = |candidate: &DataRecord| {
1241            self.split_store
1242                .label_for(&RecordKey::from(candidate))
1243                .map(|label| label == anchor_split)
1244                .unwrap_or(false)
1245        };
1246
1247        match strategy {
1248            NegativeStrategy::WrongArticle => {
1249                let anchor_date =
1250                    taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1251                let mut same_date: Vec<Arc<DataRecord>> = self
1252                    .records
1253                    .values()
1254                    .filter(|candidate| {
1255                        candidate.source == anchor_record.source
1256                            && candidate.id != anchor_record.id
1257                            && in_anchor_split(candidate)
1258                    })
1259                    .filter(|candidate| {
1260                        anchor_date
1261                            .as_deref()
1262                            .zip(taxonomy_value(candidate, META_FIELD_DATE))
1263                            .map(|(a, b)| a == b)
1264                            .unwrap_or(false)
1265                    })
1266                    .cloned()
1267                    .collect();
1268                if same_date.is_empty() {
1269                    same_date = self
1270                        .records
1271                        .values()
1272                        .filter(|candidate| {
1273                            candidate.source == anchor_record.source
1274                                && candidate.id != anchor_record.id
1275                                && in_anchor_split(candidate)
1276                        })
1277                        .cloned()
1278                        .collect();
1279                }
1280                if !same_date.is_empty() {
1281                    return self.negative_backend.choose_negative(
1282                        anchor_record,
1283                        anchor_split,
1284                        same_date,
1285                        false,
1286                        anchor_query_text,
1287                        rng,
1288                    );
1289                }
1290                let pool = self
1291                    .records
1292                    .values()
1293                    .filter(|candidate| {
1294                        candidate.id != anchor_record.id && in_anchor_split(candidate)
1295                    })
1296                    .cloned()
1297                    .collect::<Vec<_>>();
1298                self.negative_backend.choose_negative(
1299                    anchor_record,
1300                    anchor_split,
1301                    pool,
1302                    true,
1303                    anchor_query_text,
1304                    rng,
1305                )
1306            }
1307            NegativeStrategy::WrongPublicationDate => {
1308                let anchor_date =
1309                    taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1310                let pool: Vec<Arc<DataRecord>> = self
1311                    .records
1312                    .values()
1313                    .filter(|candidate| {
1314                        candidate.source == anchor_record.source
1315                            && candidate.id != anchor_record.id
1316                            && in_anchor_split(candidate)
1317                    })
1318                    .filter(|candidate| {
1319                        match (
1320                            anchor_date.as_deref(),
1321                            taxonomy_value(candidate, META_FIELD_DATE),
1322                        ) {
1323                            (Some(anchor), Some(candidate_date)) => anchor != candidate_date,
1324                            (Some(_), None) => true,
1325                            (None, Some(_)) => true,
1326                            (None, None) => false,
1327                        }
1328                    })
1329                    .cloned()
1330                    .collect();
1331                if pool.is_empty() {
1332                    // Fallback to any other record in the same split so split boundaries
1333                    // remain strictly isolated.
1334                    let fallback_pool = self
1335                        .records
1336                        .values()
1337                        .filter(|candidate| {
1338                            candidate.id != anchor_record.id && in_anchor_split(candidate)
1339                        })
1340                        .cloned()
1341                        .collect::<Vec<_>>();
1342
1343                    return self.negative_backend.choose_negative(
1344                        anchor_record,
1345                        anchor_split,
1346                        fallback_pool,
1347                        true,
1348                        anchor_query_text,
1349                        rng,
1350                    );
1351                }
1352
1353                self.negative_backend.choose_negative(
1354                    anchor_record,
1355                    anchor_split,
1356                    pool,
1357                    false,
1358                    anchor_query_text,
1359                    rng,
1360                )
1361            }
1362            NegativeStrategy::QuestionAnswerMismatch => {
1363                let pool: Vec<Arc<DataRecord>> = self
1364                    .records
1365                    .values()
1366                    .filter(|candidate| {
1367                        candidate.source == anchor_record.source
1368                            && candidate.id != anchor_record.id
1369                            && in_anchor_split(candidate)
1370                    })
1371                    .cloned()
1372                    .collect();
1373                if pool.is_empty() {
1374                    // Fallback to any other record in the same split so split boundaries
1375                    // remain strictly isolated.
1376                    let fallback_pool = self
1377                        .records
1378                        .values()
1379                        .filter(|candidate| {
1380                            candidate.id != anchor_record.id && in_anchor_split(candidate)
1381                        })
1382                        .cloned()
1383                        .collect::<Vec<_>>();
1384
1385                    return self.negative_backend.choose_negative(
1386                        anchor_record,
1387                        anchor_split,
1388                        fallback_pool,
1389                        true,
1390                        anchor_query_text,
1391                        rng,
1392                    );
1393                }
1394
1395                self.negative_backend.choose_negative(
1396                    anchor_record,
1397                    anchor_split,
1398                    pool,
1399                    false,
1400                    anchor_query_text,
1401                    rng,
1402                )
1403            }
1404            NegativeStrategy::SameRecord => {
1405                // Return the anchor record itself.  The caller
1406                // (finalize_triplet_with_negative) will call select_chunk on
1407                // this record using the recipe's negative_selector, which
1408                // picks from Context sections.  The dedup guard rejects
1409                // triplets where any two slots share identical text.
1410                Some((Arc::new(anchor_record.clone()), false))
1411            }
1412        }
1413    }
1414
1415    /// True when the recipe is the special auto-injected long-section chunk-pair recipe.
1416    fn is_auto_chunk_pair_recipe(recipe: &TripletRecipe) -> bool {
1417        recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME
1418    }
1419
1420    /// Select anchor/positive chunk pair from selectors for a single record.
1421    ///
1422    /// When selectors are the same, retries until a valid pair is found or retry limit is hit.
1423    fn select_anchor_positive_pair(
1424        &mut self,
1425        record: &DataRecord,
1426        anchor_selector: &Selector,
1427        positive_selector: &Selector,
1428        enforce_window_pair: bool,
1429    ) -> Option<(RecordChunk, RecordChunk)> {
1430        let mut anchor_chunk = self.select_chunk(record, anchor_selector)?;
1431        let mut positive_chunk = self.select_chunk(record, positive_selector)?;
1432        if anchor_selector == positive_selector {
1433            let mut retries = 0usize;
1434            while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1435                && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1436            {
1437                let Some(redraw_anchor) = self.select_chunk(record, anchor_selector) else {
1438                    break;
1439                };
1440                let Some(redraw_positive) = self.select_chunk(record, positive_selector) else {
1441                    break;
1442                };
1443                anchor_chunk = redraw_anchor;
1444                positive_chunk = redraw_positive;
1445                retries += 1;
1446            }
1447            if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1448                return None;
1449            }
1450        }
1451        Some((anchor_chunk, positive_chunk))
1452    }
1453
1454    /// Select two distinct *window* chunks for the auto-injected recipe using recipe selectors.
1455    ///
1456    /// This method intentionally uses the recipe's selectors (no hardcoded selector values).
1457    /// For the current auto recipe those selectors are both `Role(Context)`, but the behavior
1458    /// remains tied to recipe configuration, not this helper body.
1459    ///
1460    /// NOTE: this does *not* decide whether the auto recipe is enabled for a source.
1461    /// Source-level enablement happens in `resolve_source_triplet_plan` via
1462    /// `should_auto_inject_chunk_pair_recipe`.
1463    fn select_distinct_window_pair_for_auto_recipe(
1464        &mut self,
1465        recipe: &TripletRecipe,
1466        record: &DataRecord,
1467    ) -> Option<(RecordChunk, RecordChunk)> {
1468        if recipe.anchor != recipe.positive_selector {
1469            return None;
1470        }
1471        self.select_anchor_positive_pair(record, &recipe.anchor, &recipe.positive_selector, true)
1472    }
1473
1474    /// Stage B (record-level): guard for the auto-injected recipe.
1475    ///
1476    /// Requires at least two `ChunkView::Window` candidates across sections addressed by
1477    /// `selector`. This makes the auto recipe eligibility check explicit at execution time.
1478    fn record_has_at_least_two_window_chunks_for_selector(
1479        &self,
1480        record: &DataRecord,
1481        selector: &Selector,
1482    ) -> bool {
1483        let section_indices: Vec<usize> = match selector {
1484            Selector::Role(role) => record
1485                .sections
1486                .iter()
1487                .enumerate()
1488                .filter(|(_, section)| roles_match(role, &section.role))
1489                .map(|(idx, _)| idx)
1490                .collect(),
1491            Selector::Paragraph(idx) => {
1492                if *idx < record.sections.len() {
1493                    vec![*idx]
1494                } else {
1495                    Vec::new()
1496                }
1497            }
1498            Selector::Random => (0..record.sections.len()).collect(),
1499            Selector::TemporalOffset(_) => return false,
1500        };
1501
1502        let mut window_count = 0usize;
1503        for section_idx in section_indices {
1504            let Some(section) = record.sections.get(section_idx) else {
1505                continue;
1506            };
1507            let chunks = self.materialize_chunks(record, section_idx, section);
1508            window_count += chunks
1509                .iter()
1510                .filter(|chunk| matches!(chunk.view, ChunkView::Window { .. }))
1511                .count();
1512            if window_count >= 2 {
1513                return true;
1514            }
1515        }
1516        false
1517    }
1518
1519    /// Shared selector-pair stage for both standard and auto-injected triplet recipes.
1520    fn build_triplet_with_selector_pair_policy(
1521        &mut self,
1522        recipe: &TripletRecipe,
1523        record: &DataRecord,
1524        enforce_window_pair: bool,
1525        rng: &mut DeterministicRng,
1526    ) -> Option<SampleTriplet> {
1527        let (mut anchor_chunk, mut positive_chunk) = self.select_anchor_positive_pair(
1528            record,
1529            &recipe.anchor,
1530            &recipe.positive_selector,
1531            enforce_window_pair,
1532        )?;
1533        let anchor_raw_text = anchor_chunk.text.clone();
1534        self.decorate_chunk(record, &mut anchor_chunk, rng);
1535        self.decorate_chunk(record, &mut positive_chunk, rng);
1536        // Snapshot the raw anchor text for BM25 querying before decoration
1537        // added the metadata prefix — prefix tokens are absent from the index.
1538        self.finalize_triplet_with_negative(
1539            recipe,
1540            record,
1541            anchor_chunk,
1542            positive_chunk,
1543            &anchor_raw_text,
1544            rng,
1545        )
1546    }
1547
1548    /// Execute the special auto-injected long-section chunk-pair recipe.
1549    ///
1550    /// Algorithm:
1551    /// 0) Stage B check: current record must expose at least two window chunks for selector.
1552    /// 1) Draw anchor and positive from the same record using recipe selectors.
1553    /// 2) Require the two chunks to be distinct.
1554    /// 3) Additionally require both chunks to be `ChunkView::Window` variants.
1555    /// 4) Apply metadata decoration to both chunks.
1556    /// 5) Draw the negative according to recipe negative strategy.
1557    fn make_auto_chunk_pair_triplet_with_anchor(
1558        &mut self,
1559        recipe: &TripletRecipe,
1560        record: &DataRecord,
1561        rng: &mut DeterministicRng,
1562    ) -> Option<SampleTriplet> {
1563        if !self.record_has_at_least_two_window_chunks_for_selector(record, &recipe.anchor) {
1564            return None;
1565        }
1566        let (mut anchor_chunk, mut positive_chunk) =
1567            self.select_distinct_window_pair_for_auto_recipe(recipe, record)?;
1568        let anchor_raw_text = anchor_chunk.text.clone();
1569        self.decorate_chunk(record, &mut anchor_chunk, rng);
1570        self.decorate_chunk(record, &mut positive_chunk, rng);
1571        self.finalize_triplet_with_negative(
1572            recipe,
1573            record,
1574            anchor_chunk,
1575            positive_chunk,
1576            &anchor_raw_text,
1577            rng,
1578        )
1579    }
1580
1581    fn make_standard_triplet_with_anchor(
1582        &mut self,
1583        recipe: &TripletRecipe,
1584        record: &DataRecord,
1585        rng: &mut DeterministicRng,
1586    ) -> Option<SampleTriplet> {
1587        self.build_triplet_with_selector_pair_policy(recipe, record, false, rng)
1588    }
1589
1590    /// Finalize a triplet by selecting a negative and applying a deterministic 50 % coin-flip
1591    /// that swaps anchor and positive.
1592    ///
1593    /// # Anchor/positive swap
1594    ///
1595    /// Before constructing the [`SampleTriplet`], the sampler tests the least-significant bit of
1596    /// the next RNG word. When the bit is zero (≈ 50 % of the time) the anchor and positive
1597    /// slots are exchanged so that what was originally selected as the positive becomes the
1598    /// anchor, and vice-versa.
1599    ///
1600    /// **Why this matters:** contrastive objectives such as InfoNCE treat the two non-negative
1601    /// slots asymmetrically only in so far as downstream code does. If the model can learn a
1602    /// positional shortcut — for example "the anchor is always the shorter text" — it can
1603    /// achieve low loss without learning the intended similarity structure. The swap eliminates
1604    /// that opportunity by ensuring both orderings appear at equal frequency, forcing the model
1605    /// to treat the two positive-pair slots symmetrically.
1606    ///
1607    /// The negative chunk is unaffected by the swap.
1608    fn finalize_triplet_with_negative(
1609        &mut self,
1610        recipe: &TripletRecipe,
1611        record: &DataRecord,
1612        anchor_chunk: RecordChunk,
1613        positive_chunk: RecordChunk,
1614        anchor_raw_text: &str,
1615        rng: &mut DeterministicRng,
1616    ) -> Option<SampleTriplet> {
1617        let (negative_record, fallback_used) = self.select_negative_record(
1618            record,
1619            &recipe.negative_strategy,
1620            Some(anchor_raw_text),
1621            rng,
1622        )?;
1623        let mut negative_chunk = self.select_chunk(&negative_record, &recipe.negative_selector)?;
1624        self.decorate_chunk(&negative_record, &mut negative_chunk, rng);
1625
1626        // 50 % coin-flip: swap anchor and positive to prevent positional shortcuts.
1627        let (anchor_chunk, positive_chunk) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0 {
1628            (positive_chunk, anchor_chunk)
1629        } else {
1630            (anchor_chunk, positive_chunk)
1631        };
1632
1633        // Reject if any two slots share identical rendered text.  This catches sources
1634        // that produce multiple records with the same string content: a negative whose
1635        // text matches the positive (or the anchor) would produce a trivially-invalid
1636        // training example.  String equality short-circuits on the first differing byte,
1637        // so the guard is effectively free when texts differ (the common case).
1638        //
1639        // When `allow_same_anchor_positive` is set the anchor==positive check is skipped
1640        // to support SimCSE-style training where both slots carry identical text and the
1641        // model's dropout layers produce the required embedding variation at train time.
1642        // The negative must still differ from both anchor and positive.
1643        if (!recipe.allow_same_anchor_positive && anchor_chunk.text == positive_chunk.text)
1644            || negative_chunk.text == positive_chunk.text
1645            || negative_chunk.text == anchor_chunk.text
1646        {
1647            return None;
1648        }
1649
1650        let chunk_weight =
1651            self.triplet_chunk_weight(recipe, &anchor_chunk, &positive_chunk, &negative_chunk);
1652        let weight = recipe.weight * chunk_weight;
1653        let recipe_name = if fallback_used {
1654            format!("{}_fallback_same_split", recipe.name)
1655        } else {
1656            recipe.name.to_string()
1657        };
1658        Some(SampleTriplet {
1659            recipe: recipe_name,
1660            anchor: anchor_chunk,
1661            positive: positive_chunk,
1662            negative: negative_chunk,
1663            weight,
1664            instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1665        })
1666    }
1667
1668    fn make_triplet_with_anchor(
1669        &mut self,
1670        recipe: &TripletRecipe,
1671        record: &DataRecord,
1672        rng: &mut DeterministicRng,
1673    ) -> Option<SampleTriplet> {
1674        if Self::is_auto_chunk_pair_recipe(recipe) {
1675            return self.make_auto_chunk_pair_triplet_with_anchor(recipe, record, rng);
1676        }
1677        self.make_standard_triplet_with_anchor(recipe, record, rng)
1678    }
1679
1680    fn make_text_sample_for_split(
1681        &mut self,
1682        recipe: &TextRecipe,
1683        source: Option<&str>,
1684        split: SplitLabel,
1685        rng: &mut DeterministicRng,
1686    ) -> Option<TextSample> {
1687        let record = self.choose_anchor_record(source, split)?;
1688        let mut chunk = self.select_chunk(&record, &recipe.selector)?;
1689        self.decorate_chunk(&record, &mut chunk, rng);
1690        let weight = recipe.weight * self.chunk_weight(&chunk);
1691        Some(TextSample {
1692            recipe: recipe.name.to_string(),
1693            chunk,
1694            weight,
1695            instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1696        })
1697    }
1698
1699    fn chunk_weight(&self, chunk: &RecordChunk) -> f32 {
1700        chunk_weight(&self.config.chunking, chunk)
1701    }
1702
1703    fn triplet_chunk_weight(
1704        &self,
1705        recipe: &TripletRecipe,
1706        anchor: &RecordChunk,
1707        positive: &RecordChunk,
1708        negative: &RecordChunk,
1709    ) -> f32 {
1710        let floor = self.config.chunking.chunk_weight_floor;
1711        let negative_weight = negative.quality.trust.clamp(0.0, 1.0).max(floor);
1712        if Self::is_auto_chunk_pair_recipe(recipe) {
1713            // For the auto long-section recipe, use one coherence signal (proximity)
1714            // for the anchor/positive pair.
1715            let pair_trust = ((anchor.quality.trust.clamp(0.0, 1.0)
1716                + positive.quality.trust.clamp(0.0, 1.0))
1717                / 2.0)
1718                .clamp(0.0, 1.0);
1719            let pair_weight = (chunk_proximity_score(anchor, positive) * pair_trust).max(floor);
1720            // Keep negative weighting simple in this recipe: trust + floor only.
1721            return (pair_weight + pair_weight + negative_weight) / 3.0;
1722        }
1723        // Non-auto recipes also apply anchor-positive proximity, while keeping
1724        // negative weighting trust-only.
1725        let pair_proximity = chunk_proximity_score(anchor, positive);
1726        let anchor_weight = (self.chunk_weight(anchor) * pair_proximity).max(floor);
1727        let positive_weight = (self.chunk_weight(positive) * pair_proximity).max(floor);
1728        (anchor_weight + positive_weight + negative_weight) / 3.0
1729    }
1730
1731    fn decorate_chunk(
1732        &mut self,
1733        record: &DataRecord,
1734        chunk: &mut RecordChunk,
1735        rng: &mut DeterministicRng,
1736    ) {
1737        chunk.kvp_meta = record
1738            .meta_prefix
1739            .as_ref()
1740            .map(|s| s.all_metadata())
1741            .unwrap_or_default();
1742        if let Some(spec) = record.meta_prefix.as_ref()
1743            && let Some(prefix) = spec.sample(rng)
1744        {
1745            let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1746            let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1747            let total_tokens = prefix_tokens.len() + body_tokens.len();
1748            let max_window = self.config.chunking.max_window_tokens;
1749            if max_window > 0 && total_tokens > max_window {
1750                if prefix_tokens.len() >= max_window {
1751                    chunk.text = prefix_tokens
1752                        .into_iter()
1753                        .take(max_window)
1754                        .collect::<Vec<_>>()
1755                        .join(" ");
1756                    chunk.tokens_estimate = max_window;
1757                } else {
1758                    let remaining = max_window - prefix_tokens.len();
1759                    let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1760                    chunk.text =
1761                        format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1762                    chunk.tokens_estimate = max_window;
1763                }
1764            } else {
1765                chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1766                chunk.tokens_estimate = total_tokens;
1767            }
1768        }
1769    }
1770
1771    // ── parallel-batch helpers (&self, rng-based, no cursor maps) ────────────
1772
1773    /// Select a chunk using rng instead of cursor maps (for parallel execution).
1774    fn select_chunk_parallel(
1775        &self,
1776        record: &DataRecord,
1777        selector: &Selector,
1778        rng: &mut DeterministicRng,
1779    ) -> Option<RecordChunk> {
1780        match selector {
1781            Selector::Role(role) => self.select_role_parallel(record, role, rng),
1782            Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1783                let pool = self.materialize_chunks(record, *idx, section);
1784                if pool.is_empty() {
1785                    return None;
1786                }
1787                let i = rng.random_range(0..pool.len());
1788                pool.into_iter().nth(i)
1789            }),
1790            Selector::TemporalOffset(offset) => self
1791                .select_temporal_neighbor(record, *offset)
1792                .and_then(|neighbor| {
1793                    self.select_role_parallel(&neighbor, &SectionRole::Context, rng)
1794                }),
1795            Selector::Random => {
1796                if record.sections.is_empty() {
1797                    return None;
1798                }
1799                let idx = rng.random_range(0..record.sections.len());
1800                record.sections.get(idx).and_then(|section| {
1801                    let pool = self.materialize_chunks(record, idx, section);
1802                    if pool.is_empty() {
1803                        return None;
1804                    }
1805                    let i = rng.random_range(0..pool.len());
1806                    pool.into_iter().nth(i)
1807                })
1808            }
1809        }
1810    }
1811
1812    /// Select a chunk by role using rng (parallel path — no role_cursors written).
1813    fn select_role_parallel(
1814        &self,
1815        record: &DataRecord,
1816        role: &SectionRole,
1817        rng: &mut DeterministicRng,
1818    ) -> Option<RecordChunk> {
1819        let indices: Vec<usize> = record
1820            .sections
1821            .iter()
1822            .enumerate()
1823            .filter(|(_, s)| roles_match(role, &s.role))
1824            .map(|(i, _)| i)
1825            .collect();
1826        if indices.is_empty() {
1827            return None;
1828        }
1829        let start = rng.random_range(0..indices.len());
1830        for offset in 0..indices.len() {
1831            let section_idx = indices[(start + offset) % indices.len()];
1832            let section = &record.sections[section_idx];
1833            let pool = self.materialize_chunks(record, section_idx, section);
1834            if !pool.is_empty() {
1835                let i = rng.random_range(0..pool.len());
1836                return pool.into_iter().nth(i);
1837            }
1838        }
1839        None
1840    }
1841
1842    /// Decorate a chunk using a provided rng (parallel-safe, &self).
1843    fn decorate_chunk_parallel(
1844        &self,
1845        record: &DataRecord,
1846        chunk: &mut RecordChunk,
1847        rng: &mut DeterministicRng,
1848    ) {
1849        chunk.kvp_meta = record
1850            .meta_prefix
1851            .as_ref()
1852            .map(|s| s.all_metadata())
1853            .unwrap_or_default();
1854        if let Some(spec) = record.meta_prefix.as_ref()
1855            && let Some(prefix) = spec.sample(rng)
1856        {
1857            let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1858            let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1859            let total_tokens = prefix_tokens.len() + body_tokens.len();
1860            let max_window = self.config.chunking.max_window_tokens;
1861            if max_window > 0 && total_tokens > max_window {
1862                if prefix_tokens.len() >= max_window {
1863                    chunk.text = prefix_tokens
1864                        .into_iter()
1865                        .take(max_window)
1866                        .collect::<Vec<_>>()
1867                        .join(" ");
1868                    chunk.tokens_estimate = max_window;
1869                } else {
1870                    let remaining = max_window - prefix_tokens.len();
1871                    let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1872                    chunk.text =
1873                        format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1874                    chunk.tokens_estimate = max_window;
1875                }
1876            } else {
1877                chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1878                chunk.tokens_estimate = total_tokens;
1879            }
1880        }
1881    }
1882
1883    /// Select an anchor/positive pair for the parallel path (rng-based).
1884    fn select_anchor_positive_parallel(
1885        &self,
1886        record: &DataRecord,
1887        anchor_selector: &Selector,
1888        positive_selector: &Selector,
1889        enforce_window_pair: bool,
1890        rng: &mut DeterministicRng,
1891    ) -> Option<(RecordChunk, RecordChunk)> {
1892        let anchor_chunk = self.select_chunk_parallel(record, anchor_selector, rng)?;
1893        let mut positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1894        if anchor_selector == positive_selector {
1895            let mut retries = 0usize;
1896            while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1897                && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1898            {
1899                positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1900                retries += 1;
1901            }
1902            if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1903                return None;
1904            }
1905        }
1906        // No swap here — finalize_triplet_parallel applies the coin-flip swap.
1907        Some((anchor_chunk, positive_chunk))
1908    }
1909
1910    /// Select (anchor_chunk, positive_chunk, anchor_raw_text) for a recipe — parallel-safe.
1911    /// Does NOT select a negative; that happens sequentially in Phase 3.
1912    fn select_anchor_positive_for_recipe(
1913        &self,
1914        recipe: &TripletRecipe,
1915        anchor_record: &DataRecord,
1916        rng: &mut DeterministicRng,
1917    ) -> Option<(RecordChunk, RecordChunk, String)> {
1918        if Self::is_auto_chunk_pair_recipe(recipe) {
1919            if !self
1920                .record_has_at_least_two_window_chunks_for_selector(anchor_record, &recipe.anchor)
1921            {
1922                return None;
1923            }
1924            let mut anchor_chunk =
1925                self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1926            let mut positive_chunk =
1927                self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1928            let mut tries = 0usize;
1929            while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, true) {
1930                tries += 1;
1931                if tries >= SAME_SELECTOR_PAIR_RETRY_LIMIT {
1932                    return None;
1933                }
1934                anchor_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1935                positive_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1936            }
1937            let anchor_raw_text = anchor_chunk.text.clone();
1938            self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1939            self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1940            return Some((anchor_chunk, positive_chunk, anchor_raw_text));
1941        }
1942        let (mut anchor_chunk, mut positive_chunk) = self.select_anchor_positive_parallel(
1943            anchor_record,
1944            &recipe.anchor,
1945            &recipe.positive_selector,
1946            false,
1947            rng,
1948        )?;
1949        let anchor_raw_text = anchor_chunk.text.clone();
1950        self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1951        self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1952        Some((anchor_chunk, positive_chunk, anchor_raw_text))
1953    }
1954
1955    fn select_chunk(&mut self, record: &DataRecord, selector: &Selector) -> Option<RecordChunk> {
1956        match selector {
1957            Selector::Role(role) => self.select_by_role(record, role),
1958            Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1959                let pool = self.materialize_chunks(record, *idx, section);
1960                self.next_chunk_from_pool(&RecordKey::from(record), *idx, pool)
1961            }),
1962            Selector::TemporalOffset(offset) => self
1963                .select_temporal_neighbor(record, *offset)
1964                .and_then(|neighbor| self.select_by_role(&neighbor, &SectionRole::Context)),
1965            Selector::Random => {
1966                if record.sections.is_empty() {
1967                    return None;
1968                }
1969                let idx = self.rng.random_range(0..record.sections.len());
1970                record.sections.get(idx).and_then(|section| {
1971                    let pool = self.materialize_chunks(record, idx, section);
1972                    self.next_chunk_from_pool(&RecordKey::from(record), idx, pool)
1973                })
1974            }
1975        }
1976    }
1977
1978    fn select_by_role(&mut self, record: &DataRecord, role: &SectionRole) -> Option<RecordChunk> {
1979        let indices: Vec<usize> = record
1980            .sections
1981            .iter()
1982            .enumerate()
1983            .filter(|(_, section)| roles_match(role, &section.role))
1984            .map(|(idx, _)| idx)
1985            .collect();
1986        if indices.is_empty() {
1987            return None;
1988        }
1989        let key = role_cursor_key(&RecordKey::from(record), role);
1990        let start_offset = self
1991            .role_cursors
1992            .get(&key)
1993            .and_then(|last_idx| indices.iter().position(|idx| idx == last_idx))
1994            .map(|pos| (pos + 1) % indices.len())
1995            .unwrap_or_else(|| {
1996                // On first use for this (record,role), choose a deterministic hashed
1997                // section offset instead of always starting at the first matching section.
1998                // This avoids systematic head-bias while preserving reproducibility.
1999                // Preimage replicates the historical "{id}::{role}" form so
2000                // sequences are unchanged for existing corpora.
2001                let record_key = RecordKey::from(record);
2002                let seed_key = format!("{}::{}", record_key.id, role_label(role));
2003                (stable_hash_str(self.epoch_seed(), &seed_key) as usize) % indices.len()
2004            });
2005        for offset in 0..indices.len() {
2006            let section_idx = indices[(start_offset + offset) % indices.len()];
2007            let section = &record.sections[section_idx];
2008            let pool = self.materialize_chunks(record, section_idx, section);
2009            if let Some(chunk) =
2010                self.next_chunk_from_pool(&RecordKey::from(record), section_idx, pool)
2011            {
2012                self.role_cursors.insert(key.clone(), section_idx);
2013                return Some(chunk);
2014            }
2015        }
2016        None
2017    }
2018
2019    /// Materialize chunk windows for one section according to `ChunkingStrategy`.
2020    ///
2021    /// Window layout algorithm:
2022    /// 1) Tokenize section text with `split_whitespace`.
2023    /// 2) Set `span = min(max_window_tokens, total_tokens)`.
2024    /// 3) If `span == total_tokens`, emit one full-section `ChunkView::Window`.
2025    /// 4) Otherwise, for each configured overlap:
2026    ///    - `stride = max(1, span - overlap)`
2027    ///    - emit windows `[start, min(start + span, total_tokens))` while advancing by `stride`.
2028    /// 5) Optionally append one `SummaryFallback` chunk when section is longer than `max_window_tokens`.
2029    ///
2030    /// This is the source of truth for how chunk windows are "laid" in a record.
2031    fn materialize_chunks(
2032        &self,
2033        record: &DataRecord,
2034        section_idx: usize,
2035        section: &RecordSection,
2036    ) -> Vec<RecordChunk> {
2037        let mut chunks =
2038            self.chunker
2039                .materialize(&self.config.chunking, record, section_idx, section);
2040        // Propagate the section's precomputed embedding (if any) into every
2041        // chunk materialized from it. The chunker stays agnostic; Arc makes
2042        // this a pointer copy, never a vector copy.
2043        for chunk in &mut chunks {
2044            chunk.embedding = section.embedding.clone();
2045            chunk.source = record.source.clone();
2046        }
2047        chunks
2048    }
2049
2050    fn build_derived_text_recipes(recipes: &[TripletRecipe]) -> Vec<TextRecipe> {
2051        let mut derived = Vec::new();
2052        for recipe in recipes {
2053            let base = recipe.name.as_ref();
2054            derived.push(TextRecipe {
2055                name: Cow::Owned(format!("{base}_anchor")),
2056                selector: recipe.anchor.clone(),
2057                weight: recipe.weight,
2058                instruction: None,
2059            });
2060            derived.push(TextRecipe {
2061                name: Cow::Owned(format!("{base}_positive")),
2062                selector: recipe.positive_selector.clone(),
2063                weight: recipe.weight,
2064                instruction: None,
2065            });
2066            derived.push(TextRecipe {
2067                name: Cow::Owned(format!("{base}_negative")),
2068                selector: recipe.negative_selector.clone(),
2069                weight: recipe.weight,
2070                instruction: None,
2071            });
2072        }
2073        derived
2074    }
2075
2076    /// Stage A helper: true when record contains an Anchor/Context section whose token
2077    /// count exceeds `chunking.max_window_tokens`.
2078    fn record_has_long_anchor_or_context_section(&self, record: &DataRecord) -> bool {
2079        let window = self.config.chunking.max_window_tokens;
2080        if window == 0 {
2081            return false;
2082        }
2083        record.sections.iter().any(|s| {
2084            matches!(s.role, SectionRole::Anchor | SectionRole::Context) && s.token_count > window
2085        })
2086    }
2087    fn sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
2088        if self.last_sync_version == 0 {
2089            self.full_sync_records_from_cache()?;
2090            return Ok(());
2091        }
2092
2093        // Force refreshes replace cache contents wholesale; fall back to a
2094        // full pool resync with exact full-snapshot semantics (evicted =
2095        // pool minus snapshot, dedup pruned only for records that left).
2096        // Steady-state advances keep the O(Δ) delta fast path below.
2097        if self.ingestion.force_refresh_generation() != self.last_force_refresh_generation {
2098            self.last_force_refresh_generation = self.ingestion.force_refresh_generation();
2099            self.full_sync_records_from_cache()?;
2100            return Ok(());
2101        }
2102
2103        // Use sync_delta to get only added and evicted records.
2104        // No full snapshot, no sort, no HashSet diff — O(ΔN) path.
2105        let (new_version, added, evicted) = self.ingestion.sync_delta(self.last_sync_version);
2106
2107        if added.is_empty() && evicted.is_empty() {
2108            self.negative_backend.on_sync_start();
2109            self.prune_cursor_state();
2110            self.last_sync_version = new_version;
2111            return Ok(());
2112        }
2113
2114        // --- Process evictions (strict O(Δ)): prune pool state and the
2115        // per-source index only for IDs that truly left. IDs cleared and
2116        // immediately re-added in this same sync (e.g. force-refresh
2117        // returning the same records) keep their cross-batch dedup and
2118        // split-label state. This preserves the
2119        // `emitted_text_hashes_blocks_repeats` contract while still freeing
2120        // state for evicted records.
2121        let mut truly_evicted = false;
2122        if !evicted.is_empty() {
2123            let added_keys: HashSet<RecordKey> = added.iter().map(RecordKey::of_arc).collect();
2124            for evicted_key in &evicted {
2125                if added_keys.contains(evicted_key) {
2126                    continue;
2127                }
2128                let Some(record) = self.records.get(evicted_key) else {
2129                    continue;
2130                };
2131                let source = record.source.clone();
2132                let is_long = self.record_has_long_anchor_or_context_section(record);
2133                self.records.shift_remove(evicted_key);
2134                self.emitted_text_hashes.remove(evicted_key);
2135                self.split_labels.remove(evicted_key);
2136                self.chunk_index.remove(evicted_key);
2137                self.remove_id_from_source_index(&source, evicted_key);
2138                if is_long {
2139                    self.decrement_long_section_count(&source);
2140                }
2141                truly_evicted = true;
2142            }
2143        }
2144
2145        self.negative_backend.on_sync_start();
2146
2147        // --- Process additions (strict O(Δ log N) via binary search) ---
2148        // `pool_changed` tracks actual membership change: genuinely new IDs
2149        // or true evictions. Pure version refreshes of already-pooled IDs
2150        // leave the pool (and the cached records_by_split map) intact.
2151        let mut pool_changed = truly_evicted;
2152        if !added.is_empty() {
2153            let allowed: HashSet<SplitLabel> = self.allowed_target_splits().into_iter().collect();
2154            for record in &added {
2155                // Re-ingested keys (version bump for records already pooled,
2156                // e.g. a wrapped source re-draining known records) must not
2157                // create duplicate index entries or double-count long
2158                // sections. A full rebuild derives each key exactly once from
2159                // the pool, so skip index/count work here and only refresh
2160                // the stored Arc.
2161                let key = RecordKey::from(record);
2162                if self.records.contains_key(&key) {
2163                    self.records.insert(key, Arc::clone(record));
2164                    continue;
2165                }
2166                pool_changed = true;
2167                let label = self.get_or_insert_split_label(&key)?;
2168                if self.record_has_long_anchor_or_context_section(record) {
2169                    self.increment_long_section_count(&record.source);
2170                }
2171                self.records.insert(key.clone(), Arc::clone(record));
2172                self.chunk_index.insert(key.clone(), key.clone());
2173                self.insert_id_into_source_index(record.source.clone(), key, label, &allowed);
2174            }
2175        }
2176
2177        // --- Refresh derived source state (mirrors `rebuild_source_index` tail) ---
2178        // Drop index entries for sources left with no allowed-split records,
2179        // then refresh ordering, wrap flags, and cursors exactly as a full
2180        // rebuild would.
2181        let empty_sources: Vec<SourceId> = self
2182            .source_record_indices
2183            .iter()
2184            .filter(|(_, ids)| ids.is_empty())
2185            .map(|(source, _)| source.clone())
2186            .collect();
2187        let mut order_changed = !empty_sources.is_empty();
2188        for source in &empty_sources {
2189            self.source_record_indices.remove(source);
2190        }
2191        // New sources may have appeared via the insertions above.
2192        if !order_changed {
2193            for source in self.source_record_indices.keys() {
2194                if !self.source_order.contains(source) {
2195                    order_changed = true;
2196                    break;
2197                }
2198            }
2199        }
2200        if order_changed {
2201            self.source_order = self.source_record_indices.keys().cloned().collect();
2202            self.source_order.sort();
2203        }
2204        // Drop deficit counters for sources that left the pool.
2205        self.source_deficits
2206            .retain(|s, _| self.source_record_indices.contains_key(s));
2207        self.refresh_source_wrapped();
2208        self.source_record_cursors
2209            .retain(|source, _| self.source_record_indices.contains_key(source));
2210        if self.source_state_loaded {
2211            if self.source_order.is_empty() {
2212                self.source_cycle_idx = 0;
2213            }
2214            self.source_state_dirty = self.source_order.len() > 1;
2215        }
2216        // Cursor pruning is a no-op unless the valid pool shrank.
2217        if truly_evicted {
2218            self.prune_cursor_state();
2219        }
2220
2221        if pool_changed {
2222            self.pool_generation = self.pool_generation.saturating_add(1);
2223        }
2224        self.last_sync_version = new_version;
2225        Ok(())
2226    }
2227
2228    /// Full snapshot sync: used only for the initial bootstrap (last_sync_version == 0).
2229    /// Inserts records in ID-sorted order so subsequent delta syncs can append
2230    /// new records in sorted order without breaking the IndexMap ordering.
2231    fn full_sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
2232        let mut snapshot = self.ingestion.all_records_snapshot();
2233        snapshot.sort_by_key(RecordKey::of_arc);
2234
2235        // Determine which records are being evicted BEFORE clearing self.records.
2236        let old_keys: HashSet<RecordKey> = self.records.keys().cloned().collect();
2237        let incoming_keys: HashSet<RecordKey> = snapshot.iter().map(RecordKey::of_arc).collect();
2238        let evicted: HashSet<&RecordKey> = old_keys.difference(&incoming_keys).collect();
2239
2240        self.records.clear();
2241        self.long_section_counts.clear();
2242        for evicted_key in &evicted {
2243            self.emitted_text_hashes.remove(*evicted_key);
2244            self.split_labels.remove(*evicted_key);
2245        }
2246        self.negative_backend.on_sync_start();
2247        for record in snapshot {
2248            let key = RecordKey::from(&record);
2249            let _ = self.get_or_insert_split_label(&key)?;
2250            if self.record_has_long_anchor_or_context_section(&record) {
2251                self.increment_long_section_count(&record.source);
2252            }
2253            self.records.insert(key, Arc::clone(&record));
2254        }
2255        self.prune_cursor_state();
2256        self.rebuild_chunk_index();
2257        self.rebuild_source_index()?;
2258        self.last_sync_version = self.ingestion.global_version();
2259        // The pool was fully replaced; invalidate the cached split map.
2260        self.pool_generation = self.pool_generation.saturating_add(1);
2261        Ok(())
2262    }
2263
2264    fn ingest_internal_for_split(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
2265        if !self.ingestion.has_sources() {
2266            return Ok(());
2267        }
2268        if !self.ingestion_cursors_loaded {
2269            if let Some(state) = self.split_store.load_sampler_state()? {
2270                self.ingestion.load_cursors(&state.source_stream_cursors);
2271                self.ingestion.set_epoch(state.epoch);
2272                self.ingestion.set_epoch_step(state.epoch_step);
2273            }
2274            self.ingestion_cursors_loaded = true;
2275        }
2276        if self.ingestion.all_caches_empty() {
2277            self.ingestion.refresh_all();
2278        } else {
2279            self.ingestion.advance(self.config.batch_size);
2280        }
2281        let observed = self.ingestion.total_ingest_count();
2282        if observed == self.last_observed_ingest && !self.records.is_empty() {
2283            return Ok(());
2284        }
2285        self.last_observed_ingest = observed;
2286        self.sync_records_from_cache()?;
2287        let max_window_tokens = self.config.chunking.max_window_tokens;
2288        self.negative_backend.on_records_refreshed(
2289            &self.records,
2290            max_window_tokens,
2291            &|id| self.split_labels.get(id).copied(),
2292            self.ingestion.last_refreshed_sources(),
2293        );
2294        // Epoch tracking and source-state management must run every batch regardless
2295        // of whether the record pool changed — reconcile advances the sampling cursor.
2296        self.epoch_tracker.ensure_loaded()?;
2297        self.refresh_records_by_split_cache()?;
2298        let cached_records_by_split = self
2299            .cached_records_by_split
2300            .as_ref()
2301            .expect("records_by_split cache populated by refresh");
2302        self.epoch_tracker
2303            .reconcile(target_split, cached_records_by_split);
2304        self.ensure_source_state()?;
2305        Ok(())
2306    }
2307
2308    /// Ensure ingestion cursors (including epoch step) are loaded from the
2309    /// persisted state so the step counter reflects the saved value before
2310    /// the wrapper increments it for the current call.
2311    ///
2312    /// Only loads ingestion cursors and source epoch.  The sampling state
2313    /// (RNG, source_record_cursors, …) is restored later by
2314    /// `ingest_internal_for_split` → `ensure_source_state`, after
2315    /// `sync_records_from_cache` has populated `source_record_indices`.
2316    fn ensure_ingestion_cursors_loaded(&mut self) -> Result<(), SamplerError> {
2317        if !self.ingestion_cursors_loaded {
2318            if let Some(state) = self.split_store.load_sampler_state()? {
2319                self.ingestion.load_cursors(&state.source_stream_cursors);
2320                self.ingestion.set_epoch(state.epoch);
2321                self.ingestion.set_epoch_step(state.epoch_step);
2322            }
2323            self.ingestion_cursors_loaded = true;
2324        }
2325        Ok(())
2326    }
2327
2328    #[cfg(test)]
2329    fn ingest_internal(&mut self, split: SplitLabel) -> Result<(), SamplerError> {
2330        self.ingest_internal_for_split(split)
2331    }
2332
2333    fn ingest_internal_with_weights_for_split(
2334        &mut self,
2335        target_split: SplitLabel,
2336        weights: &HashMap<SourceId, f32>,
2337    ) -> Result<(), SamplerError> {
2338        if !self.ingestion.has_sources() {
2339            return Ok(());
2340        }
2341        if !self.ingestion_cursors_loaded {
2342            if let Some(state) = self.split_store.load_sampler_state()? {
2343                self.ingestion.load_cursors(&state.source_stream_cursors);
2344                self.ingestion.set_epoch(state.epoch);
2345                self.ingestion.set_epoch_step(state.epoch_step);
2346            }
2347            self.ingestion_cursors_loaded = true;
2348        }
2349        if self.ingestion.all_caches_empty() {
2350            self.ingestion.refresh_all_with_weights(weights)?;
2351        } else {
2352            self.ingestion
2353                .advance_with_weights(self.config.batch_size, weights)?;
2354        }
2355        let observed = self.ingestion.total_ingest_count();
2356        if observed == self.last_observed_ingest && !self.records.is_empty() {
2357            return Ok(());
2358        }
2359        self.last_observed_ingest = observed;
2360        self.sync_records_from_cache()?;
2361        let max_window_tokens = self.config.chunking.max_window_tokens;
2362        self.negative_backend.on_records_refreshed(
2363            &self.records,
2364            max_window_tokens,
2365            &|id| self.split_labels.get(id).copied(),
2366            self.ingestion.last_refreshed_sources(),
2367        );
2368        self.epoch_tracker.ensure_loaded()?;
2369        self.refresh_records_by_split_cache()?;
2370        let cached_records_by_split = self
2371            .cached_records_by_split
2372            .as_ref()
2373            .expect("records_by_split cache populated by refresh");
2374        self.epoch_tracker
2375            .reconcile(target_split, cached_records_by_split);
2376        self.ensure_source_state()?;
2377        Ok(())
2378    }
2379
2380    /// Helper that centralises the weight-fallback logic.
2381    ///
2382    /// If `weights` is `Some` and non-uniform or if `weights` is `None`,
2383    /// delegates to the appropriate ingest method.
2384    fn ingest_with_weights_fallback(
2385        &mut self,
2386        target_split: SplitLabel,
2387        weights: Option<&HashMap<SourceId, f32>>,
2388    ) -> Result<(), SamplerError> {
2389        match weights {
2390            Some(weights)
2391                if !weights.is_empty()
2392                    && !weights
2393                        .values()
2394                        .all(|&w| w == *weights.values().next().unwrap()) =>
2395            {
2396                self.ingest_internal_with_weights_for_split(target_split, weights)?
2397            }
2398            _ => self.ingest_internal_for_split(target_split)?,
2399        }
2400        Ok(())
2401    }
2402
2403    fn force_ingest_refresh_with_weights_for_split(
2404        &mut self,
2405        target_split: SplitLabel,
2406        weights: &HashMap<SourceId, f32>,
2407    ) -> Result<(), SamplerError> {
2408        if !self.ingestion.has_sources() {
2409            return Ok(());
2410        }
2411        if !self.ingestion_cursors_loaded {
2412            if let Some(state) = self.split_store.load_sampler_state()? {
2413                self.ingestion.load_cursors(&state.source_stream_cursors);
2414                self.ingestion.set_epoch(state.epoch);
2415                self.ingestion.set_epoch_step(state.epoch_step);
2416            }
2417            self.ingestion_cursors_loaded = true;
2418        }
2419        self.ingestion.force_refresh_all_with_weights(weights)?;
2420        self.last_observed_ingest = self.ingestion.total_ingest_count();
2421        self.sync_records_from_cache()?;
2422        let max_window_tokens = self.config.chunking.max_window_tokens;
2423        self.negative_backend.on_records_refreshed(
2424            &self.records,
2425            max_window_tokens,
2426            &|id| self.split_labels.get(id).copied(),
2427            self.ingestion.last_refreshed_sources(),
2428        );
2429        self.epoch_tracker.ensure_loaded()?;
2430        self.refresh_records_by_split_cache()?;
2431        let cached_records_by_split = self
2432            .cached_records_by_split
2433            .as_ref()
2434            .expect("records_by_split cache populated by refresh");
2435        self.epoch_tracker
2436            .reconcile(target_split, cached_records_by_split);
2437        self.ensure_source_state()?;
2438        Ok(())
2439    }
2440
2441    /// Select one triplet candidate for a specific source and split.
2442    ///
2443    /// This helper is the source-level recipe execution path used by both pair and
2444    /// triplet batch builders. It always starts from `triplet_recipes_for_source`,
2445    /// which already includes any auto-injected long-section recipe when eligible.
2446    ///
2447    /// Returns:
2448    /// - candidate triplet (with recipe) if one was sampled,
2449    /// - number of recipe attempts consumed (for round-robin bookkeeping).
2450    fn sample_source_triplet_candidate(
2451        &mut self,
2452        source: &str,
2453        target_split: SplitLabel,
2454        recipe_orders: &mut HashMap<RecipeKey, Vec<usize>>,
2455        recipe_positions: &mut HashMap<RecipeKey, usize>,
2456        rng: &mut DeterministicRng,
2457    ) -> (Option<(TripletRecipe, SampleTriplet)>, usize) {
2458        // Stage A (source-level injection): resolve effective recipe pool,
2459        // including auto-injected long-section recipe when source is eligible.
2460        let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
2461        if recipes.is_empty() {
2462            return (None, 0);
2463        }
2464        if !recipe_orders.contains_key(source) {
2465            let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2466            let order =
2467                self.recipe_order_weighted_cycled(&recipe_weights, self.triplet_recipe_rr_idx, rng);
2468            recipe_orders.insert(source.to_string(), order);
2469        }
2470        let order = recipe_orders
2471            .get(source)
2472            .expect("recipe order missing for source");
2473        let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2474        let Some(anchor) = self.choose_anchor_record(Some(source), target_split) else {
2475            return (None, 0);
2476        };
2477
2478        let mut attempts = 0usize;
2479        for offset in 0..order.len() {
2480            let idx = order[(*pos + offset) % order.len()];
2481            attempts = attempts.saturating_add(1);
2482            let recipe = recipes[idx].clone();
2483            // Stage B/C happen inside `make_triplet_with_anchor`:
2484            // - Stage B: record-level auto-recipe eligibility gate,
2485            // - Stage C: chunk window materialization/selection.
2486            if let Some(sample) = self.make_triplet_with_anchor(&recipe, &anchor, rng) {
2487                *pos = (*pos + offset + 1) % order.len();
2488                return (Some((recipe, sample)), attempts);
2489            }
2490        }
2491
2492        (None, attempts)
2493    }
2494
2495    fn next_pair_batch_inner_with_weights(
2496        &mut self,
2497        target_split: SplitLabel,
2498        weights: Option<&HashMap<SourceId, f32>>,
2499    ) -> Result<SampleBatch, SamplerError> {
2500        self.ingest_with_weights_fallback(target_split, weights)?;
2501        self.ensure_split_has_records(target_split)?;
2502        let sources = self.source_order.clone();
2503        if sources.is_empty() {
2504            if self.triplet_recipes.is_empty() {
2505                return Err(SamplerError::Configuration(
2506                    "no triplet recipes available".into(),
2507                ));
2508            }
2509            let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2510            let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2511            let recipe_order = self.recipe_order_weighted_cycled(
2512                &recipe_weights,
2513                self.triplet_recipe_rr_idx,
2514                &mut rng,
2515            );
2516            let mut pairs = Vec::new();
2517            let mut seen = HashSet::new();
2518            let mut last_recipe_name = None;
2519            let mut recipe_pos = 0usize;
2520            let mut recipe_steps = 0usize;
2521            let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2522            for _ in 0..attempts {
2523                if pairs.len() >= self.config.batch_size {
2524                    break;
2525                }
2526                let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2527                    break;
2528                };
2529                let mut triplet = None;
2530                for offset in 0..recipe_order.len() {
2531                    let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2532                    recipe_steps = recipe_steps.saturating_add(1);
2533                    let recipe = self.triplet_recipes[idx].clone();
2534                    last_recipe_name = Some(recipe.name.clone());
2535                    if let Some(sample) = self.make_triplet_with_anchor(&recipe, &anchor, &mut rng)
2536                    {
2537                        triplet = Some((recipe, sample));
2538                        recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2539                        break;
2540                    }
2541                }
2542                if let Some((recipe, triplet)) = triplet {
2543                    let key = (
2544                        RecordKey::of_chunk(&triplet.anchor),
2545                        RecordKey::of_chunk(&triplet.positive),
2546                        RecordKey::of_chunk(&triplet.negative),
2547                    );
2548                    if seen.insert(key) {
2549                        let SampleTriplet {
2550                            recipe: triplet_recipe_name,
2551                            anchor,
2552                            positive,
2553                            negative,
2554                            weight,
2555                            instruction,
2556                        } = triplet;
2557
2558                        // Look up the anchor record's label to override structural defaults.
2559                        // For SRD pair-mode data, the label is stored in DataRecord.
2560                        let record_label = self
2561                            .records
2562                            .get(&RecordKey::of_chunk(&anchor))
2563                            .and_then(|r| r.label.clone());
2564
2565                        if pairs.len() < self.config.batch_size {
2566                            pairs.push(SamplePair {
2567                                recipe: triplet_recipe_name.clone(),
2568                                anchor: anchor.clone(),
2569                                positive: positive.clone(),
2570                                weight,
2571                                instruction: instruction.clone(),
2572                                label: record_label.clone().unwrap_or(PairLabel::Positive),
2573                                reason: None,
2574                            });
2575                        }
2576                        if pairs.len() < self.config.batch_size {
2577                            pairs.push(SamplePair {
2578                                recipe: triplet_recipe_name,
2579                                anchor,
2580                                positive: negative,
2581                                weight,
2582                                instruction,
2583                                label: record_label.unwrap_or(PairLabel::Negative),
2584                                reason: Some(
2585                                    strategy_reason(&recipe.negative_strategy).to_string(),
2586                                ),
2587                            });
2588                        }
2589                    }
2590                }
2591            }
2592            if recipe_steps > 0 {
2593                self.triplet_recipe_rr_idx =
2594                    self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2595            }
2596            self.rng = rng;
2597            pad_with_reuse(&mut pairs, self.config.batch_size);
2598            if pairs.len() == self.config.batch_size {
2599                return Ok(SampleBatch { pairs });
2600            }
2601            return Err(SamplerError::Exhausted(
2602                last_recipe_name
2603                    .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2604                    .to_string(),
2605            ));
2606        }
2607
2608        let mut pairs = Vec::new();
2609        let mut seen = HashSet::new();
2610        let mut source_steps = 0usize;
2611        let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2612        let mut source_idx = self.source_cycle_idx % sources.len();
2613        let mut cycle_sources = self.shuffled_source_cycle(cycle);
2614        // Explicit non-uniform weights engage deficit round-robin source
2615        // choice; uniform/empty maps keep legacy uniform cycling bit-identical.
2616        let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2617            Some(w) if Self::weights_are_selective(w) => Some(w),
2618            _ => None,
2619        };
2620        let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2621        let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2622        let mut recipe_steps = 0usize;
2623        let max_recipe_len = sources
2624            .iter()
2625            .map(|source| self.triplet_recipe_count_for_source(source))
2626            .max()
2627            .unwrap_or(1)
2628            .max(1);
2629        let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2630        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2631        for _ in 0..attempts {
2632            if pairs.len() >= self.config.batch_size {
2633                break;
2634            }
2635            let weighted_source;
2636            let source = if let Some(w) = weighted {
2637                match self.select_weighted_source(&sources, w) {
2638                    Some(s) => {
2639                        weighted_source = s;
2640                        weighted_source.as_str()
2641                    }
2642                    None => break,
2643                }
2644            } else {
2645                cycle_sources[source_idx].as_str()
2646            };
2647            let (triplet, attempts_used) = self.sample_source_triplet_candidate(
2648                source,
2649                target_split,
2650                &mut recipe_orders,
2651                &mut recipe_positions,
2652                &mut rng,
2653            );
2654            recipe_steps = recipe_steps.saturating_add(attempts_used);
2655            if let Some((recipe, triplet)) = triplet {
2656                let key = (
2657                    RecordKey::of_chunk(&triplet.anchor),
2658                    RecordKey::of_chunk(&triplet.positive),
2659                    RecordKey::of_chunk(&triplet.negative),
2660                );
2661                if seen.insert(key) {
2662                    let SampleTriplet {
2663                        recipe: triplet_recipe_name,
2664                        anchor,
2665                        positive,
2666                        negative,
2667                        weight,
2668                        instruction,
2669                    } = triplet;
2670
2671                    // Look up the anchor record's label to override structural defaults.
2672                    let record_label = self
2673                        .records
2674                        .get(&RecordKey::of_chunk(&anchor))
2675                        .and_then(|r| r.label.clone());
2676
2677                    if pairs.len() < self.config.batch_size {
2678                        pairs.push(SamplePair {
2679                            recipe: triplet_recipe_name.clone(),
2680                            anchor: anchor.clone(),
2681                            positive: positive.clone(),
2682                            weight,
2683                            instruction: instruction.clone(),
2684                            label: record_label.clone().unwrap_or(PairLabel::Positive),
2685                            reason: None,
2686                        });
2687                    }
2688                    if pairs.len() < self.config.batch_size {
2689                        pairs.push(SamplePair {
2690                            recipe: triplet_recipe_name,
2691                            anchor,
2692                            positive: negative,
2693                            weight,
2694                            instruction,
2695                            label: record_label.unwrap_or(PairLabel::Negative),
2696                            reason: Some(strategy_reason(&recipe.negative_strategy).to_string()),
2697                        });
2698                    }
2699                }
2700            }
2701            source_idx += 1;
2702            source_steps += 1;
2703            if source_idx >= cycle_sources.len() {
2704                source_idx = 0;
2705                cycle = cycle.saturating_add(1);
2706                cycle_sources = self.shuffled_source_cycle(cycle);
2707            }
2708        }
2709        if recipe_steps > 0 {
2710            self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2711        }
2712        self.rng = rng;
2713        pad_with_reuse(&mut pairs, self.config.batch_size);
2714        if pairs.len() == self.config.batch_size {
2715            self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2716            self.source_state_dirty = sources.len() > 1;
2717            return Ok(SampleBatch { pairs });
2718        }
2719        Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
2720    }
2721
2722    /// Check whether `sample` passes both per-batch and per-window text dedup.
2723    /// The per-batch set is never cleared mid-batch (survives epoch advances),
2724    /// while `emitted_text_hashes` is cleared on epoch advance to allow fresh
2725    /// permutations.  Both must pass for the sample to be emitted.
2726    fn try_emit_text_sample(
2727        &mut self,
2728        sample: &TextSample,
2729        batch_texts: &mut HashSet<String>,
2730    ) -> bool {
2731        batch_texts.insert(sample.chunk.text.clone())
2732            && self
2733                .emitted_text_hashes
2734                .entry(RecordKey::of_chunk(&sample.chunk))
2735                .or_default()
2736                .insert(stable_hash_str(0, &sample.chunk.text))
2737    }
2738
2739    fn next_text_batch_inner_with_weights(
2740        &mut self,
2741        target_split: SplitLabel,
2742        weights: Option<&HashMap<SourceId, f32>>,
2743    ) -> Result<TextBatch, SamplerError> {
2744        self.ingest_with_weights_fallback(target_split, weights)?;
2745        self.ensure_split_has_records(target_split)?;
2746        let sources = self.source_order.clone();
2747        if sources.is_empty() {
2748            if self.text_recipes.is_empty() {
2749                return Err(SamplerError::Configuration(
2750                    "no text recipes configured".into(),
2751                ));
2752            }
2753            let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2754            let recipe_weights: Vec<f32> = self.text_recipes.iter().map(|r| r.weight).collect();
2755            let recipe_order = self.text_recipe_order_weighted_cycled(
2756                &recipe_weights,
2757                self.text_recipe_rr_idx,
2758                &mut rng,
2759            );
2760            let mut samples = Vec::new();
2761            let mut batch_texts = HashSet::new();
2762            let mut last_recipe_name = None;
2763            let mut recipe_pos = 0usize;
2764            let mut recipe_steps = 0usize;
2765            let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2766            for _ in 0..attempts {
2767                if samples.len() >= self.config.batch_size {
2768                    break;
2769                }
2770                let recipe_idx = recipe_order[recipe_pos];
2771                recipe_pos = (recipe_pos + 1) % recipe_order.len();
2772                recipe_steps = recipe_steps.saturating_add(1);
2773                let recipe = self.text_recipes[recipe_idx].clone();
2774                last_recipe_name = Some(recipe.name.clone());
2775                if let Some(sample) =
2776                    self.make_text_sample_for_split(&recipe, None, target_split, &mut rng)
2777                    && self.try_emit_text_sample(&sample, &mut batch_texts)
2778                {
2779                    samples.push(sample);
2780                }
2781            }
2782            if recipe_steps > 0 {
2783                self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2784            }
2785            self.rng = rng;
2786            pad_with_reuse(&mut samples, self.config.batch_size);
2787            if samples.len() == self.config.batch_size {
2788                return Ok(TextBatch { samples });
2789            }
2790            return Err(SamplerError::Exhausted(
2791                last_recipe_name
2792                    .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TEXT))
2793                    .to_string(),
2794            ));
2795        }
2796
2797        let mut samples = Vec::new();
2798        let mut batch_texts = HashSet::new();
2799        let mut source_steps = 0usize;
2800        let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2801        let mut idx = self.source_cycle_idx % sources.len();
2802        let mut cycle_sources = self.shuffled_source_cycle(cycle);
2803        let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2804        let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2805        let mut recipe_steps = 0usize;
2806        let max_recipe_len = sources
2807            .iter()
2808            .map(|source| self.text_recipes_for_source(source).len())
2809            .max()
2810            .unwrap_or(1)
2811            .max(1);
2812        let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2813        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2814        // Explicit non-uniform weights engage deficit round-robin source
2815        // choice; uniform/empty maps keep legacy uniform cycling bit-identical.
2816        let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2817            Some(w) if Self::weights_are_selective(w) => Some(w),
2818            _ => None,
2819        };
2820        for _ in 0..attempts {
2821            if samples.len() >= self.config.batch_size {
2822                break;
2823            }
2824            let weighted_source;
2825            let source = if let Some(w) = weighted {
2826                match self.select_weighted_source(&sources, w) {
2827                    Some(s) => {
2828                        weighted_source = s;
2829                        weighted_source.as_str()
2830                    }
2831                    None => break,
2832                }
2833            } else {
2834                cycle_sources[idx].as_str()
2835            };
2836            let recipes = self.text_recipes_for_source(source).to_vec();
2837            if recipes.is_empty() {
2838                idx += 1;
2839                source_steps += 1;
2840                if idx >= cycle_sources.len() {
2841                    idx = 0;
2842                    cycle = cycle.saturating_add(1);
2843                    cycle_sources = self.shuffled_source_cycle(cycle);
2844                }
2845                continue;
2846            }
2847            if !recipe_orders.contains_key(source) {
2848                let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2849                let order = self.text_recipe_order_weighted_cycled(
2850                    &recipe_weights,
2851                    self.text_recipe_rr_idx,
2852                    &mut rng,
2853                );
2854                recipe_orders.insert(source.to_string(), order);
2855            }
2856            let order = recipe_orders
2857                .get(source)
2858                .expect("recipe order missing for source");
2859            let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2860            let mut sample: Option<(TextRecipe, TextSample)> = None;
2861            for offset in 0..order.len() {
2862                let recipe_idx = order[(*pos + offset) % order.len()];
2863                let recipe = recipes[recipe_idx].clone();
2864                if let Some(item) =
2865                    self.make_text_sample_for_split(&recipe, Some(source), target_split, &mut rng)
2866                    && self.try_emit_text_sample(&item, &mut batch_texts)
2867                {
2868                    recipe_steps = recipe_steps.saturating_add(offset + 1);
2869                    *pos = (*pos + offset + 1) % order.len();
2870                    sample = Some((recipe, item));
2871                    break;
2872                }
2873            }
2874            if sample.is_none() {
2875                recipe_steps = recipe_steps.saturating_add(order.len());
2876            }
2877            if let Some((_recipe, sample)) = sample {
2878                samples.push(sample);
2879            }
2880            idx += 1;
2881            source_steps += 1;
2882            if idx >= cycle_sources.len() {
2883                idx = 0;
2884                cycle = cycle.saturating_add(1);
2885                cycle_sources = self.shuffled_source_cycle(cycle);
2886            }
2887        }
2888        if samples.len() != self.config.batch_size {
2889            pad_with_reuse(&mut samples, self.config.batch_size);
2890        }
2891        if samples.len() != self.config.batch_size {
2892            self.rng = rng;
2893            return Err(SamplerError::Exhausted(RECIPE_LABEL_TEXT.into()));
2894        }
2895        self.rng = rng;
2896        self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2897        self.source_state_dirty = sources.len() > 1;
2898        if recipe_steps > 0 {
2899            self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2900        }
2901        Ok(TextBatch { samples })
2902    }
2903
2904    fn next_triplet_batch_inner_with_weights(
2905        &mut self,
2906        target_split: SplitLabel,
2907        weights: Option<&HashMap<SourceId, f32>>,
2908    ) -> Result<TripletBatch, SamplerError> {
2909        self.ingest_with_weights_fallback(target_split, weights)?;
2910        self.ensure_split_has_records(target_split)?;
2911        let sources = self.source_order.clone();
2912        if sources.is_empty() {
2913            if self.triplet_recipes.is_empty() {
2914                return Err(SamplerError::Configuration(
2915                    "no triplet recipes configured".into(),
2916                ));
2917            }
2918            let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2919            let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2920            let recipe_order = self.recipe_order_weighted_cycled(
2921                &recipe_weights,
2922                self.triplet_recipe_rr_idx,
2923                &mut rng,
2924            );
2925            let mut triplets = Vec::new();
2926            let mut seen = HashSet::new();
2927            let mut last_recipe_name = None;
2928            let mut recipe_pos = 0usize;
2929            let mut recipe_steps = 0usize;
2930            let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2931            for _ in 0..attempts {
2932                if triplets.len() >= self.config.batch_size {
2933                    break;
2934                }
2935                let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2936                    break;
2937                };
2938                let mut triplet = None;
2939                for offset in 0..recipe_order.len() {
2940                    let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2941                    recipe_steps = recipe_steps.saturating_add(1);
2942                    let recipe = self.triplet_recipes[idx].clone();
2943                    last_recipe_name = Some(recipe.name.clone());
2944                    if let Some(sample) = self.make_triplet_with_anchor(&recipe, &anchor, &mut rng)
2945                    {
2946                        triplet = Some(sample);
2947                        recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2948                        break;
2949                    }
2950                }
2951                if let Some(triplet) = triplet {
2952                    let key = (
2953                        RecordKey::of_chunk(&triplet.anchor),
2954                        RecordKey::of_chunk(&triplet.positive),
2955                        RecordKey::of_chunk(&triplet.negative),
2956                    );
2957                    if seen.insert(key) {
2958                        triplets.push(triplet);
2959                    }
2960                }
2961            }
2962            if recipe_steps > 0 {
2963                self.triplet_recipe_rr_idx =
2964                    self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2965            }
2966            self.rng = rng;
2967            pad_with_reuse(&mut triplets, self.config.batch_size);
2968            if triplets.len() == self.config.batch_size {
2969                return Ok(TripletBatch { triplets });
2970            }
2971            return Err(SamplerError::Exhausted(
2972                last_recipe_name
2973                    .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2974                    .to_string(),
2975            ));
2976        }
2977
2978        // ── Declarations for multi-source parallel path ────────────────────────
2979        let mut triplets: Vec<SampleTriplet> = Vec::new();
2980        let mut seen: HashSet<(RecordKey, RecordKey, RecordKey)> = HashSet::new();
2981        let mut source_steps = 0usize;
2982        let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2983        let mut source_idx = self.source_cycle_idx % sources.len();
2984        let mut cycle_sources = self.shuffled_source_cycle(cycle);
2985        // Explicit non-uniform weights engage deficit round-robin source
2986        // choice; uniform/empty maps keep legacy uniform cycling bit-identical.
2987        let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2988            Some(w) if Self::weights_are_selective(w) => Some(w),
2989            _ => None,
2990        };
2991        let mut recipe_steps = 0usize;
2992        let max_recipe_len = sources
2993            .iter()
2994            .map(|source| self.triplet_recipe_count_for_source(source.as_str()))
2995            .max()
2996            .unwrap_or(1)
2997            .max(1);
2998        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2999
3000        // ── Phase 1: sequential anchor pre-assignment ──────────────────────────
3001        // Advance cursors and pick anchors deterministically. Capture a fork seed
3002        // for each slot so that Phase 2 can reconstruct independent rngs in
3003        // deterministic order.
3004        struct SlotPlan {
3005            anchor: Arc<DataRecord>,
3006            recipes: Vec<TripletRecipe>,
3007            fork_seed: u64,
3008        }
3009        let target_slots = self.config.batch_size * 4 * max_recipe_len;
3010        let mut slot_plans: Vec<SlotPlan> = Vec::with_capacity(target_slots);
3011
3012        for _ in 0..target_slots {
3013            if slot_plans.len() >= target_slots {
3014                break;
3015            }
3016            let weighted_source;
3017            let source = if let Some(w) = weighted {
3018                match self.select_weighted_source(&sources, w) {
3019                    Some(s) => {
3020                        weighted_source = s;
3021                        weighted_source.as_str()
3022                    }
3023                    None => break,
3024                }
3025            } else {
3026                cycle_sources[source_idx].as_str()
3027            };
3028            let (recipes, _) = self.resolve_source_triplet_plan(source);
3029            if !recipes.is_empty() {
3030                let fork_seed = rng.next_u64();
3031                if let Some(anchor) = self.choose_anchor_record(Some(source), target_split) {
3032                    slot_plans.push(SlotPlan {
3033                        anchor,
3034                        recipes,
3035                        fork_seed,
3036                    });
3037                    recipe_steps = recipe_steps.saturating_add(1);
3038                }
3039            }
3040            source_idx += 1;
3041            source_steps += 1;
3042            if source_idx >= cycle_sources.len() {
3043                source_idx = 0;
3044                cycle = cycle.saturating_add(1);
3045                cycle_sources = self.shuffled_source_cycle(cycle);
3046            }
3047        }
3048        self.rng = rng;
3049
3050        // ── Phase 2: parallel anchor+positive selection ────────────────────────
3051        // Only chunk selection runs in parallel; negative selection and the
3052        // coin-flip swap remain sequential in Phase 3 so that the BM25 cursor
3053        // rotation (shared RwLock state) is deterministic across runs.
3054        struct SlotCandidate {
3055            recipe: TripletRecipe,
3056            anchor: Arc<DataRecord>,
3057            anchor_chunk: RecordChunk,
3058            positive_chunk: RecordChunk,
3059            anchor_raw_text: String,
3060        }
3061        let mut raw_candidates: Vec<(usize, Option<SlotCandidate>)> = slot_plans
3062            .par_iter()
3063            .enumerate()
3064            .map(|(slot_idx, plan)| {
3065                let mut fork_rng = DeterministicRng::new(plan.fork_seed);
3066                // Use weighted shuffle so recipes with higher weight are tried more
3067                // often and zero-weight recipes are excluded entirely.
3068                let weights: Vec<f32> = plan.recipes.iter().map(|r| r.weight).collect();
3069                let order = weighted_recipe_order(&weights, &mut fork_rng);
3070                let mut candidate = None;
3071                for &idx in &order {
3072                    let recipe = &plan.recipes[idx];
3073                    if let Some((ac, pc, raw)) =
3074                        self.select_anchor_positive_for_recipe(recipe, &plan.anchor, &mut fork_rng)
3075                    {
3076                        candidate = Some(SlotCandidate {
3077                            recipe: recipe.clone(),
3078                            anchor: Arc::clone(&plan.anchor),
3079                            anchor_chunk: ac,
3080                            positive_chunk: pc,
3081                            anchor_raw_text: raw,
3082                        });
3083                        break;
3084                    }
3085                }
3086                (slot_idx, candidate)
3087            })
3088            .collect();
3089
3090        // ── Phase 3: sequential finalization (negative selection + coin-flip) ───
3091        // Sort by slot_idx to restore deterministic ordering, then finalize each
3092        // candidate sequentially so that BM25 cursor rotation advances in a fixed
3093        // order regardless of rayon's thread scheduling.
3094        raw_candidates.sort_unstable_by_key(|(i, _)| *i);
3095        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3096        for (_, candidate) in raw_candidates {
3097            if triplets.len() >= self.config.batch_size {
3098                break;
3099            }
3100            let Some(sc) = candidate else { continue };
3101            let (negative_record, fallback_used) = match self.select_negative_record(
3102                &sc.anchor,
3103                &sc.recipe.negative_strategy,
3104                Some(&sc.anchor_raw_text),
3105                &mut rng,
3106            ) {
3107                Some(pair) => pair,
3108                None => continue,
3109            };
3110            let mut negative_chunk = match self.select_chunk_parallel(
3111                &negative_record,
3112                &sc.recipe.negative_selector,
3113                &mut rng,
3114            ) {
3115                Some(c) => c,
3116                None => continue,
3117            };
3118            self.decorate_chunk_parallel(&negative_record, &mut negative_chunk, &mut rng);
3119
3120            // 50 % coin-flip swap using the pre-sampled per-slot seed + main rng.
3121            let (anchor_chunk, positive_chunk) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0
3122            {
3123                (sc.positive_chunk, sc.anchor_chunk)
3124            } else {
3125                (sc.anchor_chunk, sc.positive_chunk)
3126            };
3127
3128            if (!sc.recipe.allow_same_anchor_positive && anchor_chunk.text == positive_chunk.text)
3129                || negative_chunk.text == positive_chunk.text
3130                || negative_chunk.text == anchor_chunk.text
3131            {
3132                continue;
3133            }
3134
3135            let chunk_weight = self.triplet_chunk_weight(
3136                &sc.recipe,
3137                &anchor_chunk,
3138                &positive_chunk,
3139                &negative_chunk,
3140            );
3141            let weight = sc.recipe.weight * chunk_weight;
3142            let recipe_name = if fallback_used {
3143                format!("{}_fallback_same_split", sc.recipe.name)
3144            } else {
3145                sc.recipe.name.to_string()
3146            };
3147            let triplet = SampleTriplet {
3148                recipe: recipe_name,
3149                anchor: anchor_chunk,
3150                positive: positive_chunk,
3151                negative: negative_chunk,
3152                weight,
3153                instruction: sc.recipe.instruction.as_ref().map(|s| s.to_string()),
3154            };
3155            let key = (
3156                RecordKey::of_chunk(&triplet.anchor),
3157                RecordKey::of_chunk(&triplet.positive),
3158                RecordKey::of_chunk(&triplet.negative),
3159            );
3160            if seen.insert(key) && triplets.len() < self.config.batch_size {
3161                triplets.push(triplet);
3162            }
3163        }
3164        self.rng = rng;
3165
3166        if recipe_steps > 0 {
3167            self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
3168        }
3169        pad_with_reuse(&mut triplets, self.config.batch_size);
3170        if triplets.len() == self.config.batch_size {
3171            self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
3172            self.source_state_dirty = sources.len() > 1;
3173            let batch = TripletBatch { triplets };
3174            return Ok(batch);
3175        }
3176        Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
3177    }
3178
3179    // ── test-only helpers ─────────────────────────────────────────────────────
3180
3181    /// Return a mutable reference to the concrete [`backends::Bm25Backend`], panicking
3182    /// when the feature is disabled (which won't happen — call sites are
3183    /// equally gated on `#[cfg(feature = "bm25-mining")]`).
3184    #[cfg(test)]
3185    #[cfg(feature = "bm25-mining")]
3186    fn bm25_backend_mut(&mut self) -> &mut backends::Bm25Backend {
3187        self.negative_backend
3188            .as_any_mut()
3189            .downcast_mut::<backends::Bm25Backend>()
3190            .expect("bm25_backend_mut: negative_backend is Bm25Backend when bm25-mining feature is active")
3191    }
3192
3193    /// Test shim: `recipe_order_weighted_shuffled` using the inner rng.
3194    #[cfg(test)]
3195    fn recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
3196        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3197        let result = self.recipe_order_weighted_shuffled(weights, &mut rng);
3198        self.rng = rng;
3199        result
3200    }
3201
3202    /// Test shim: `recipe_order_weighted_cycled` using the inner rng.
3203    #[cfg(test)]
3204    fn recipe_order_weighted_cycled_seeded(
3205        &mut self,
3206        weights: &[f32],
3207        rr_idx: usize,
3208    ) -> Vec<usize> {
3209        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3210        let result = self.recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
3211        self.rng = rng;
3212        result
3213    }
3214
3215    /// Test shim: `text_recipe_order_weighted_shuffled` using the inner rng.
3216    #[cfg(test)]
3217    fn text_recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
3218        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3219        let result = self.text_recipe_order_weighted_shuffled(weights, &mut rng);
3220        self.rng = rng;
3221        result
3222    }
3223
3224    /// Test shim: `text_recipe_order_weighted_cycled` using the inner rng.
3225    #[cfg(test)]
3226    fn text_recipe_order_weighted_cycled_seeded(
3227        &mut self,
3228        weights: &[f32],
3229        rr_idx: usize,
3230    ) -> Vec<usize> {
3231        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3232        let result = self.text_recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
3233        self.rng = rng;
3234        result
3235    }
3236
3237    /// Test shim: `select_negative_record` using the inner rng.
3238    #[cfg(test)]
3239    fn select_negative_record_seeded(
3240        &mut self,
3241        anchor_record: &DataRecord,
3242        strategy: &NegativeStrategy,
3243        anchor_query_text: Option<&str>,
3244    ) -> Option<(Arc<DataRecord>, bool)> {
3245        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3246        let result =
3247            self.select_negative_record(anchor_record, strategy, anchor_query_text, &mut rng);
3248        self.rng = rng;
3249        result
3250    }
3251
3252    /// Test shim: `make_triplet_with_anchor` using the inner rng.
3253    #[cfg(test)]
3254    fn make_triplet_with_anchor_seeded(
3255        &mut self,
3256        recipe: &TripletRecipe,
3257        anchor: &DataRecord,
3258    ) -> Option<SampleTriplet> {
3259        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3260        let result = self.make_triplet_with_anchor(recipe, anchor, &mut rng);
3261        self.rng = rng;
3262        result
3263    }
3264
3265    /// Test shim: `decorate_chunk` using the inner rng.
3266    #[cfg(test)]
3267    fn decorate_chunk_seeded(&mut self, record: &DataRecord, chunk: &mut RecordChunk) {
3268        let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3269        self.decorate_chunk(record, chunk, &mut rng);
3270        self.rng = rng;
3271    }
3272
3273    // ── extended-metrics helpers ───────────────────────────────────────────────
3274
3275    /// Return cumulative `(fallback_count, selection_count)` from the BM25
3276    /// backend since it was created.  Only available when both `bm25-mining`
3277    /// and `extended-metrics` features are active.
3278    #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3279    fn bm25_fallback_stats(&self) -> (u64, u64) {
3280        self.negative_backend.bm25_fallback_stats()
3281    }
3282
3283    /// Look up `anchor`'s split label and return BM25-ranked candidate IDs.
3284    ///
3285    /// Mirrors the old `bm25_ranked_candidates` method that tests called on
3286    /// `TripletSamplerInner` before BM25 state moved into `Bm25Backend`.
3287    #[cfg(test)]
3288    #[cfg(feature = "bm25-mining")]
3289    fn bm25_ranked_candidates(&mut self, anchor: &crate::data::DataRecord) -> Vec<RecordKey> {
3290        let split = self
3291            .split_store
3292            .label_for(&RecordKey::from(anchor))
3293            .unwrap_or(SplitLabel::Train);
3294        self.negative_backend
3295            .as_any_mut()
3296            .downcast_mut::<backends::Bm25Backend>()
3297            .expect("bm25_ranked_candidates: Bm25Backend")
3298            .ranked_candidates_pub(anchor, split)
3299    }
3300}
3301
3302/// Build a shuffled selection order from a slice of recipe weights.
3303///
3304/// Each recipe with `weight > 0.0` receives a number of slots proportional to its weight
3305/// relative to the smallest positive weight, capped at
3306/// [`RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER`].  The resulting expanded slot list is then
3307/// shuffled, giving each recipe approximately the requested frequency in round-robin
3308/// cycling.
3309///
3310/// This is a free function (no `self`) so it can be called from rayon parallel closures as
3311/// well as from within the `TripletSamplerInner` instance methods.
3312fn weighted_recipe_order(weights: &[f32], rng: &mut DeterministicRng) -> Vec<usize> {
3313    let nonzero: Vec<(usize, f32)> = weights
3314        .iter()
3315        .enumerate()
3316        .filter(|(_, w)| **w > 0.0)
3317        .map(|(i, &w)| (i, w))
3318        .collect();
3319    if nonzero.is_empty() {
3320        return Vec::new();
3321    }
3322    let w_min = nonzero
3323        .iter()
3324        .map(|(_, w)| *w)
3325        .fold(f32::INFINITY, f32::min);
3326    let mut order: Vec<usize> = Vec::new();
3327    for (recipe_idx, w) in &nonzero {
3328        let tickets = ((w / w_min).round() as usize).clamp(1, RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER);
3329        for _ in 0..tickets {
3330            order.push(*recipe_idx);
3331        }
3332    }
3333    order.shuffle(rng);
3334    order
3335}
3336
3337fn same_selector_pair_is_valid(
3338    anchor_chunk: &RecordChunk,
3339    positive_chunk: &RecordChunk,
3340    enforce_window_pair: bool,
3341) -> bool {
3342    if triplet_chunk_key(anchor_chunk) == triplet_chunk_key(positive_chunk) {
3343        return false;
3344    }
3345    if !enforce_window_pair {
3346        return true;
3347    }
3348    matches!(
3349        (&anchor_chunk.view, &positive_chunk.view),
3350        (ChunkView::Window { .. }, ChunkView::Window { .. })
3351    )
3352}
3353
3354impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSampler<S> {
3355    /// Create a sampler from config and a split-state backend.
3356    pub fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
3357        let inner = TripletSamplerInner::new(config, split_store);
3358        Self {
3359            inner: Mutex::new(inner),
3360        }
3361    }
3362
3363    /// Create a sampler from config with a custom chunking implementation.
3364    pub fn new_with_chunker(
3365        config: SamplerConfig,
3366        split_store: Arc<S>,
3367        chunker: Arc<dyn ChunkingAlgorithm>,
3368    ) -> Self {
3369        let inner = TripletSamplerInner::new_with_chunker(config, split_store, chunker);
3370        Self {
3371            inner: Mutex::new(inner),
3372        }
3373    }
3374
3375    /// Return an unweighted pair batch for `split`.
3376    ///
3377    /// Samples all sources uniformly. Prefer
3378    /// [`Self::next_pair_batch_with_weights_for_split`] with an explicit
3379    /// per-source weight map to honor a data mixture.
3380    #[deprecated(
3381        note = "samples all sources uniformly; call next_pair_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3382    )]
3383    pub fn next_pair_batch_for_split(
3384        &self,
3385        split: SplitLabel,
3386    ) -> Result<SampleBatch, SamplerError> {
3387        self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3388    }
3389
3390    /// Return an unweighted text batch for `split`.
3391    ///
3392    /// Samples all sources uniformly. Prefer
3393    /// [`Self::next_text_batch_with_weights_for_split`] with an explicit
3394    /// per-source weight map to honor a data mixture.
3395    #[deprecated(
3396        note = "samples all sources uniformly; call next_text_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3397    )]
3398    pub fn next_text_batch_for_split(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3399        self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3400    }
3401
3402    /// Return an unweighted triplet batch for `split`.
3403    ///
3404    /// Samples all sources uniformly. Prefer
3405    /// [`Self::next_triplet_batch_with_weights_for_split`] with an explicit
3406    /// per-source weight map to honor a data mixture.
3407    #[deprecated(
3408        note = "samples all sources uniformly; call next_triplet_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3409    )]
3410    pub fn next_triplet_batch_for_split(
3411        &self,
3412        split: SplitLabel,
3413    ) -> Result<TripletBatch, SamplerError> {
3414        self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3415    }
3416
3417    /// Shared retry loop for all three batch types: load state, bump step,
3418    /// retry on exhaustion up to EXHAUSTION_RETRY_LIMIT times, and return
3419    /// the batch or propagate the final error.
3420    fn next_batch_with_retry<B>(
3421        &self,
3422        split: SplitLabel,
3423        weights: &HashMap<SourceId, f32>,
3424        inner_fn: impl Fn(
3425            &mut TripletSamplerInner<S>,
3426            SplitLabel,
3427            Option<&HashMap<SourceId, f32>>,
3428        ) -> Result<B, SamplerError>,
3429    ) -> Result<B, SamplerError> {
3430        let mut inner = self.inner.lock().unwrap();
3431        inner.ensure_split_allowed(split)?;
3432        // Validate weight keys up front, on every path: a typo'd source id
3433        // must fail loudly, never silently sample as if unweighted.
3434        if !weights.is_empty() {
3435            inner.ingestion.validate_weights(weights)?;
3436        }
3437        // Load persisted state BEFORE bumping step so the restored step
3438        // counter is reflected before we increment it for this call.
3439        inner.ensure_ingestion_cursors_loaded()?;
3440        // Bump epoch step once per public call, regardless of success/exhaustion.
3441        inner.ingestion.increment_epoch_step();
3442        for attempt in 0..=EXHAUSTION_RETRY_LIMIT {
3443            match inner_fn(&mut inner, split, Some(weights)) {
3444                Ok(batch) => return Ok(batch),
3445                Err(SamplerError::Exhausted(_)) if attempt < EXHAUSTION_RETRY_LIMIT => {
3446                    inner.force_ingest_refresh_with_weights_for_split(split, weights)?;
3447                }
3448                Err(err) => return Err(err),
3449            }
3450        }
3451        // The for-loop always returns via the match arms above.
3452        unreachable!()
3453    }
3454
3455    /// Return a weighted pair batch for `split` using per-source weights.
3456    ///
3457    /// Source choice is deficit round-robin over the map: deterministic and
3458    /// exact over each weight cycle (0.75/0.25 yields exactly 3:1 anchors).
3459    /// Unlisted sources get quantum 1.0; non-finite/negative quanta count as
3460    /// 0. Empty or all-equal maps keep legacy uniform cycling bit-identically.
3461    /// Weights also steer ingestion; see `ingest_with_weights_fallback`.
3462    pub fn next_pair_batch_with_weights_for_split(
3463        &self,
3464        split: SplitLabel,
3465        weights: &HashMap<SourceId, f32>,
3466    ) -> Result<SampleBatch, SamplerError> {
3467        self.next_batch_with_retry(split, weights, |inner, split, weights| {
3468            inner.next_pair_batch_inner_with_weights(split, weights)
3469        })
3470    }
3471
3472    /// Return a weighted text batch for `split` using per-source weights.
3473    ///
3474    /// Source choice is deficit round-robin over the map: deterministic and
3475    /// exact over each weight cycle. Unlisted sources get quantum 1.0;
3476    /// non-finite/negative quanta count as 0. Empty or all-equal maps keep
3477    /// legacy uniform cycling bit-identically.
3478    pub fn next_text_batch_with_weights_for_split(
3479        &self,
3480        split: SplitLabel,
3481        weights: &HashMap<SourceId, f32>,
3482    ) -> Result<TextBatch, SamplerError> {
3483        self.next_batch_with_retry(split, weights, |inner, split, weights| {
3484            inner.next_text_batch_inner_with_weights(split, weights)
3485        })
3486    }
3487
3488    /// Return a weighted triplet batch for `split` using per-source weights.
3489    ///
3490    /// Source choice is deficit round-robin over the map: deterministic and
3491    /// exact over each weight cycle. Unlisted sources get quantum 1.0;
3492    /// non-finite/negative quanta count as 0. Empty or all-equal maps keep
3493    /// legacy uniform cycling bit-identically.
3494    pub fn next_triplet_batch_with_weights_for_split(
3495        &self,
3496        split: SplitLabel,
3497        weights: &HashMap<SourceId, f32>,
3498    ) -> Result<TripletBatch, SamplerError> {
3499        self.next_batch_with_retry(split, weights, |inner, split, weights| {
3500            inner.next_triplet_batch_inner_with_weights(split, weights)
3501        })
3502    }
3503
3504    /// Spawn a background prefetcher for triplet batches.
3505    ///
3506    /// Samples all sources uniformly. Prefer
3507    /// [`Self::prefetch_triplet_batches_with_weights`] with an explicit
3508    /// per-source weight map to honor a data mixture.
3509    #[deprecated(
3510        note = "samples all sources uniformly; call prefetch_triplet_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3511    )]
3512    pub fn prefetch_triplet_batches(
3513        self: Arc<Self>,
3514        split: SplitLabel,
3515        capacity: usize,
3516    ) -> BatchPrefetcher<TripletBatch> {
3517        BatchPrefetcher::new(capacity, move || {
3518            self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3519        })
3520    }
3521
3522    /// Spawn a background prefetcher for weighted triplet batches.
3523    pub fn prefetch_triplet_batches_with_weights(
3524        self: Arc<Self>,
3525        split: SplitLabel,
3526        capacity: usize,
3527        weights: HashMap<SourceId, f32>,
3528    ) -> BatchPrefetcher<TripletBatch> {
3529        BatchPrefetcher::new(capacity, move || {
3530            self.next_triplet_batch_with_weights_for_split(split, &weights)
3531        })
3532    }
3533
3534    /// Spawn a background prefetcher for pair batches.
3535    ///
3536    /// Samples all sources uniformly. Prefer
3537    /// [`Self::prefetch_pair_batches_with_weights`] with an explicit
3538    /// per-source weight map to honor a data mixture.
3539    #[deprecated(
3540        note = "samples all sources uniformly; call prefetch_pair_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3541    )]
3542    pub fn prefetch_pair_batches(
3543        self: Arc<Self>,
3544        split: SplitLabel,
3545        capacity: usize,
3546    ) -> BatchPrefetcher<SampleBatch> {
3547        BatchPrefetcher::new(capacity, move || {
3548            self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3549        })
3550    }
3551
3552    /// Spawn a background prefetcher for weighted pair batches.
3553    pub fn prefetch_pair_batches_with_weights(
3554        self: Arc<Self>,
3555        split: SplitLabel,
3556        capacity: usize,
3557        weights: HashMap<SourceId, f32>,
3558    ) -> BatchPrefetcher<SampleBatch> {
3559        BatchPrefetcher::new(capacity, move || {
3560            self.next_pair_batch_with_weights_for_split(split, &weights)
3561        })
3562    }
3563
3564    /// Spawn a background prefetcher for text batches.
3565    ///
3566    /// Samples all sources uniformly. Prefer
3567    /// [`Self::prefetch_text_batches_with_weights`] with an explicit
3568    /// per-source weight map to honor a data mixture.
3569    #[deprecated(
3570        note = "samples all sources uniformly; call prefetch_text_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3571    )]
3572    pub fn prefetch_text_batches(
3573        self: Arc<Self>,
3574        split: SplitLabel,
3575        capacity: usize,
3576    ) -> BatchPrefetcher<TextBatch> {
3577        BatchPrefetcher::new(capacity, move || {
3578            self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3579        })
3580    }
3581
3582    /// Spawn a background prefetcher for weighted text batches.
3583    pub fn prefetch_text_batches_with_weights(
3584        self: Arc<Self>,
3585        split: SplitLabel,
3586        capacity: usize,
3587        weights: HashMap<SourceId, f32>,
3588    ) -> BatchPrefetcher<TextBatch> {
3589        BatchPrefetcher::new(capacity, move || {
3590            self.next_text_batch_with_weights_for_split(split, &weights)
3591        })
3592    }
3593
3594    /// Return the currently active text recipes.
3595    pub fn text_recipes(&self) -> Vec<TextRecipe> {
3596        let inner = self.inner.lock().unwrap();
3597        inner.text_recipes().to_vec()
3598    }
3599
3600    /// Register a data source for ingestion and sampling.
3601    ///
3602    /// Returns an error if the source's `id()` matches the reserved `__*__`
3603    /// pattern used for internal synthetic/metadata source identifiers.
3604    pub fn register_source(
3605        &self,
3606        source: Box<dyn DataSource + 'static>,
3607    ) -> Result<(), SamplerError> {
3608        let mut inner = self.inner.lock().unwrap();
3609        inner.register_source(source)
3610    }
3611
3612    /// Force sampler epoch to `epoch` (advanced deterministic replay control).
3613    pub fn set_epoch(&self, epoch: u64) -> Result<(), SamplerError> {
3614        let mut inner = self.inner.lock().unwrap();
3615        inner.set_epoch(epoch)
3616    }
3617
3618    /// Persist sampler and split runtime state for restart-resume.
3619    ///
3620    /// When `save_to` is `Some(path)`, current persisted runtime state is also
3621    /// mirrored to `path` when supported by the split-store backend.
3622    pub fn save_sampler_state(&self, save_to: Option<&Path>) -> Result<(), SamplerError> {
3623        let mut inner = self.inner.lock().unwrap();
3624        inner.save_sampler_state(save_to)
3625    }
3626
3627    /// Return cumulative `(fallback_count, selection_count)` from the BM25
3628    /// backend since it was created.  Only available when both `bm25-mining`
3629    /// and `extended-metrics` features are active.
3630    ///
3631    /// `selection_count` counts calls where the pool was non-empty.
3632    /// `fallback_count` counts the subset where BM25 produced no candidates
3633    /// intersecting the pool and random selection was used instead.
3634    #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3635    pub fn bm25_fallback_stats(&self) -> (u64, u64) {
3636        let inner = self.inner.lock().unwrap();
3637        inner.bm25_fallback_stats()
3638    }
3639}
3640
3641impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> Sampler for TripletSampler<S> {
3642    fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
3643        self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3644    }
3645
3646    fn next_pair_batch_with_weights(
3647        &self,
3648        split: SplitLabel,
3649        weights: &HashMap<SourceId, f32>,
3650    ) -> Result<SampleBatch, SamplerError> {
3651        self.next_pair_batch_with_weights_for_split(split, weights)
3652    }
3653
3654    fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3655        self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3656    }
3657
3658    fn next_text_batch_with_weights(
3659        &self,
3660        split: SplitLabel,
3661        weights: &HashMap<SourceId, f32>,
3662    ) -> Result<TextBatch, SamplerError> {
3663        self.next_text_batch_with_weights_for_split(split, weights)
3664    }
3665
3666    fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
3667        self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3668    }
3669
3670    fn next_triplet_batch_with_weights(
3671        &self,
3672        split: SplitLabel,
3673        weights: &HashMap<SourceId, f32>,
3674    ) -> Result<TripletBatch, SamplerError> {
3675        self.next_triplet_batch_with_weights_for_split(split, weights)
3676    }
3677
3678    fn save_sampler_state(&self, save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
3679        TripletSampler::save_sampler_state(self, save_to)
3680    }
3681}
3682
3683fn roles_match(target: &SectionRole, candidate: &SectionRole) -> bool {
3684    target == candidate
3685}
3686
3687fn role_cursor_key(record_key: &RecordKey, role: &SectionRole) -> (RecordKey, String) {
3688    (record_key.clone(), role_label(role))
3689}
3690
3691fn role_label(role: &SectionRole) -> String {
3692    match role {
3693        SectionRole::Anchor => ROLE_LABEL_ANCHOR.into(),
3694        SectionRole::Context => ROLE_LABEL_CONTEXT.into(),
3695    }
3696}
3697
3698fn taxonomy_value(record: &DataRecord, field: MetadataKey) -> Option<&str> {
3699    record.taxonomy.iter().find_map(|entry| field.strip(entry))
3700}
3701
3702fn strategy_reason(strategy: &NegativeStrategy) -> &'static str {
3703    match strategy {
3704        NegativeStrategy::WrongPublicationDate => NEG_REASON_WRONG_DATE,
3705        NegativeStrategy::WrongArticle => NEG_REASON_WRONG_ARTICLE,
3706        NegativeStrategy::QuestionAnswerMismatch => NEG_REASON_WRONG_QA,
3707        NegativeStrategy::SameRecord => NEG_REASON_SAME_RECORD,
3708    }
3709}
3710
3711/// Dedup key for text batch generation: (record_id, text).
3712///
3713/// Unlike `triplet_chunk_key` which compares structural identity
3714/// (record_id, section_idx, window_index), this key compares actual
3715/// text content. This prevents text-columns mode sources (where Anchor
3716/// and Context sections carry identical text) from producing duplicate
3717/// samples from the same record.
3718#[cfg(test)]
3719fn text_dedup_key(chunk: &RecordChunk) -> (String, String) {
3720    (chunk.record_id.clone(), chunk.text.clone())
3721}
3722
3723fn triplet_chunk_key(chunk: &RecordChunk) -> String {
3724    match &chunk.view {
3725        ChunkView::Window { index, .. } => {
3726            format!("{}|{}|w|{}", chunk.record_id, chunk.section_idx, index)
3727        }
3728        ChunkView::SummaryFallback { strategy, .. } => {
3729            format!("{}|{}|s|{}", chunk.record_id, chunk.section_idx, strategy)
3730        }
3731    }
3732}
3733
3734fn pad_with_reuse<T: Clone>(items: &mut Vec<T>, target: usize) {
3735    if items.is_empty() || items.len() >= target {
3736        return;
3737    }
3738    let seed = items.clone();
3739    let base_len = seed.len();
3740    for idx in 0..(target - items.len()) {
3741        items.push(seed[idx % base_len].clone());
3742    }
3743}
3744
3745#[cfg(test)]
3746mod tests;