22337622218

Committed 24 Feb 2026 05:13AM UTC coverage: 91.641% (-1.0%) from 92.675%

Build # 22337622218

Build Type

Pull #7

github

Committed by

web-flow

Commit Message

Merge ed641cc4f into 980559192

Pull Request Pull Request #7: Add HF source

Coverage Stats

3602 of 4094 new or added lines in 6 files covered. (87.98%)

99 existing lines in 3 files now uncovered.

13057 of 14248 relevant lines covered (91.64%)

2770.99 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.51

/src/source/mod.rs

//! Data source interfaces and paging helpers.
//!
//! Ownership model:
//! - `DataSource` is the sampler-facing interface that produces batches.
//! - `IndexableSource` exposes stable, index-based access into a corpus.
//! - `IndexablePager` owns the deterministic pseudo-random paging logic and
//!   can page any indexable source without retaining per-record state.

use chrono::{DateTime, Utc};
use std::hash::Hash;
use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::config::{SamplerConfig, TripletRecipe};
use crate::data::DataRecord;
use crate::errors::SamplerError;
use crate::hash::stable_hash_with;
use crate::types::SourceId;

/// Source implementation modules.
pub mod sources;
/// Utility helpers used by source implementations.
pub mod utilities;
pub use sources::file_source::{
    FileSource, FileSourceConfig, SectionBuilder, TaxonomyBuilder, anchor_context_sections,
    taxonomy_from_path,
};
#[cfg(feature = "huggingface")]
pub use sources::huggingface_source::{HuggingFaceRowSource, HuggingFaceRowsConfig};

/// Source-owned incremental refresh position.
///
/// The sampler stores and returns this value between refresh calls.
/// `revision` is opaque to the sampler and interpreted only by the source.
#[derive(Clone, Debug)]
pub struct SourceCursor {
    /// Most recent observation timestamp produced by the source.
    pub last_seen: DateTime<Utc>,
    /// Opaque paging position token used to continue incremental refresh.
    pub revision: u64,
}

/// Result of a single source refresh call.
///
/// Pass the returned `cursor` back into the next refresh to continue paging.
#[derive(Clone, Debug)]
pub struct SourceSnapshot {
    /// Records returned by the refresh operation.
    pub records: Vec<DataRecord>,
    /// Next cursor to pass into a future refresh call.
    pub cursor: SourceCursor,
}

/// Sampler-facing data source interface.
///
/// Implementations may be streaming or index-backed. For a fixed dataset state
/// and cursor, refresh output should be deterministic.
pub trait DataSource: Send + Sync {
    /// Stable source identifier used in records, metrics, and persistence state.
    fn id(&self) -> &str;
    /// Fetch up to `limit` records starting from `cursor` state.
    ///
    /// Return the next cursor position in `SourceSnapshot.cursor`.
    fn refresh(
        &self,
        cursor: Option<&SourceCursor>,
        limit: Option<usize>,
    ) -> Result<SourceSnapshot, SamplerError>;

    /// Exact metadata record count reported by the source.
    ///
    /// This is intended for estimators that must avoid iterating records.
    /// Implementations should return `Ok(count)` only when the count is
    /// exact for the source scope. Return `Err` when exact counting is not
    /// possible or the source is unavailable.
    ///
    /// Keep this consistent with `refresh` by using the same backend scope,
    /// filtering, and logical corpus definition.
    fn reported_record_count(&self) -> Result<u128, SamplerError>;

    /// Provide the active sampler configuration to this source.
    ///
    /// Called when the source is registered with a sampler. Sources can use
    /// this to align internal heuristics with runtime sampler settings.
    fn configure_sampler(&self, _config: &SamplerConfig) {}

    /// Optional source-provided default triplet recipes.
    ///
    /// Used when sampler config does not provide explicit recipes.
    fn default_triplet_recipes(&self) -> Vec<TripletRecipe> {
        Vec::new()
    }
}

/// Index-addressable source interface used by deterministic pagers.
///
/// `len_hint` must be stable within an epoch, and `record_at` must return the
/// record corresponding to the same index across runs.
///
/// Dense indexing is strongly recommended: implement indices as `0..len_hint`
/// with minimal gaps. Sparse indexes (returning `None` for many positions)
/// still work but waste paging capacity and reduce batch fill rates.
pub trait IndexableSource: Send + Sync {
    /// Stable source identifier.
    fn id(&self) -> &str;
    /// Current index domain size, typically `Some(total_records)`.
    fn len_hint(&self) -> Option<usize>;
    /// Return the record at index `idx`, or `None` for sparse/missing positions.
    fn record_at(&self, idx: usize) -> Result<Option<DataRecord>, SamplerError>;
}

/// Deterministic pager for `IndexableSource`.
///
/// Encapsulates shuffle seed and cursor math so callers can reuse a stable
/// paging algorithm without implementing permutation logic themselves.
pub struct IndexablePager {
    source_id: SourceId,
}

impl IndexablePager {
    /// Create a new deterministic pager for `source_id`.
    pub fn new(source_id: impl Into<SourceId>) -> Self {
        Self {
            source_id: source_id.into(),
        }
    }

    /// Page records from an `IndexableSource` using the provided cursor.
    pub fn refresh(
        &self,
        source: &dyn IndexableSource,
        cursor: Option<&SourceCursor>,
        limit: Option<usize>,
    ) -> Result<SourceSnapshot, SamplerError> {
        let total = source
            .len_hint()
            .ok_or_else(|| SamplerError::SourceInconsistent {
                source_id: source.id().to_string(),
                details: "indexable source did not provide len_hint".into(),
            })?;
        self.refresh_with(total, cursor, limit, |idx| source.record_at(idx))
    }

    /// Page records using a custom index fetcher.
    ///
    /// Useful when records are indexable but not exposed through `IndexableSource`
    /// (for example, temporary index stores or precomputed path lists).
    pub fn refresh_with(
        &self,
        total: usize,
        cursor: Option<&SourceCursor>,
        limit: Option<usize>,
        mut fetch: impl FnMut(usize) -> Result<Option<DataRecord>, SamplerError>,
    ) -> Result<SourceSnapshot, SamplerError> {
        if total == 0 {
            return Ok(SourceSnapshot {
                records: Vec::new(),
                cursor: SourceCursor {
                    last_seen: Utc::now(),
                    revision: 0,
                },
            });
        }
        let mut start = cursor.map(|cursor| cursor.revision as usize).unwrap_or(0);
        if start >= total {
            start = 0;
        }
        let max = limit.unwrap_or(total);
        let mut records = Vec::new();
        let seed = Self::seed_for(&self.source_id, total);
        let mut permutation = IndexPermutation::new(total, seed, start as u64);
        let report_every = Duration::from_millis(750);
        let refresh_start = Instant::now();
        let mut last_report = refresh_start;
        let mut attempts = 0usize;
        let should_report = total >= 10_000 || max >= 1_024;
        if should_report {
            eprintln!(
                "[triplets:source] refresh start source='{}' total={} target={}",
                self.source_id, total, max
            );
        }
        for _ in 0..total {
            attempts += 1;
            if records.len() >= max {
                break;
            }
            let idx = permutation.next();
            if let Some(record) = fetch(idx)? {
                records.push(record);
            }
            if should_report && last_report.elapsed() >= report_every {
                eprintln!(
                    "[triplets:source] refresh progress source='{}' attempted={}/{} fetched={}/{} elapsed={:.1}s",
                    self.source_id,
                    attempts,
                    total,
                    records.len(),
                    max,
                    refresh_start.elapsed().as_secs_f64()
                );
                last_report = Instant::now();
            }
        }
        if should_report {
            eprintln!(
                "[triplets:source] refresh done source='{}' attempted={} fetched={} elapsed={:.2}s",
                self.source_id,
                attempts,
                records.len(),
                refresh_start.elapsed().as_secs_f64()
            );
        }
        let last_seen = records
            .iter()
            .map(|record| record.updated_at)
            .max()
            .unwrap_or_else(Utc::now);
        let next_start = permutation.cursor();
        Ok(SourceSnapshot {
            records,
            cursor: SourceCursor {
                last_seen,
                revision: next_start as u64,
            },
        })
    }

    /// Build a deterministic seed for a source and total size.
    pub(crate) fn seed_for(source_id: &SourceId, total: usize) -> u64 {
        Self::stable_index_shuffle_key(source_id, 0)
            ^ Self::stable_index_shuffle_key(source_id, total)
    }

    fn stable_index_shuffle_key(source_id: &SourceId, idx: usize) -> u64 {
        stable_hash_with(|hasher| {
            source_id.hash(hasher);
            idx.hash(hasher);
        })
    }
}

/// DataSource adapter that pages an `IndexableSource` via `IndexablePager`.
pub struct IndexableAdapter<T: IndexableSource> {
    inner: T,
}

impl<T: IndexableSource> IndexableAdapter<T> {
    /// Wrap an `IndexableSource` so it can be registered as a `DataSource`.
    pub fn new(inner: T) -> Self {
        Self { inner }
    }
}

impl<T: IndexableSource> DataSource for IndexableAdapter<T> {
    fn id(&self) -> &str {
        self.inner.id()
    }

    fn refresh(
        &self,
        cursor: Option<&SourceCursor>,
        limit: Option<usize>,
    ) -> Result<SourceSnapshot, SamplerError> {
        let pager = IndexablePager::new(self.inner.id());
        pager.refresh(&self.inner, cursor, limit)
    }

    fn reported_record_count(&self) -> Result<u128, SamplerError> {
        self.inner
            .len_hint()
            .map(|value| value as u128)
            .ok_or_else(|| SamplerError::SourceInconsistent {
                source_id: self.inner.id().to_string(),
                details: "indexable source did not provide len_hint".into(),
            })
    }
}

/// Internal permutation used by `IndexablePager`.
pub(crate) struct IndexPermutation {
    total: u64,
    domain_bits: u32,
    domain_size: u64,
    seed: u64,
    counter: u64,
}

impl IndexPermutation {
    fn new(total: usize, seed: u64, counter: u64) -> Self {
        let total_u64 = total as u64;
        let domain_bits = (64 - (total_u64 - 1).leading_zeros()).max(1);
        let domain_size = 1u64 << domain_bits;
        Self {
            total: total_u64,
            domain_bits,
            domain_size,
            seed,
            counter,
        }
    }

    fn next(&mut self) -> usize {
        loop {
            let v =
                Self::permute_bits(self.counter % self.domain_size, self.domain_bits, self.seed);
            self.counter = self.counter.wrapping_add(1);
            if v < self.total {
                return v as usize;
            }
        }
    }

    fn cursor(&self) -> usize {
        (self.counter as usize) % (self.total as usize)
    }
    fn permute_bits(value: u64, bits: u32, seed: u64) -> u64 {
        if bits == 0 {
            return 0;
        }
        let mask = if bits == 64 {
            u64::MAX
        } else {
            (1u64 << bits) - 1
        };
        let mut a = (seed | 1) & mask;
        if a == 0 {
            a = 1;
        }
        let b = (seed >> 1) & mask;
        a.wrapping_mul(value).wrapping_add(b) & mask
    }
}

/// In-memory data source for tests and small datasets.
pub struct InMemorySource {
    id: SourceId,
    records: Arc<Vec<DataRecord>>,
}

impl InMemorySource {
    /// Create an in-memory source from prebuilt records.
    pub fn new(id: impl Into<SourceId>, records: Vec<DataRecord>) -> Self {
        Self {
            id: id.into(),
            records: Arc::new(records),
        }
    }
}

impl DataSource for InMemorySource {
    fn id(&self) -> &str {
        &self.id
    }

    fn refresh(
        &self,
        cursor: Option<&SourceCursor>,
        limit: Option<usize>,
    ) -> Result<SourceSnapshot, SamplerError> {
        let records = &*self.records;
        let total = records.len();
        let mut start = cursor.map(|cursor| cursor.revision as usize).unwrap_or(0);
        if total > 0 && start >= total {
            start = 0;
        }
        let max = limit.unwrap_or(total);
        let mut filtered = Vec::new();
        for idx in 0..total {
            if filtered.len() >= max {
                break;
            }
            let pos = (start + idx) % total;
            filtered.push(records[pos].clone());
        }
        let last_seen = filtered
            .iter()
            .map(|record| record.updated_at)
            .max()
            .unwrap_or_else(Utc::now);
        let next_start = if total == 0 {
            0
        } else {
            (start + filtered.len()) % total
        };
        Ok(SourceSnapshot {
            records: filtered,
            cursor: SourceCursor {
                last_seen,
                revision: next_start as u64,
            },
        })
    }

    fn reported_record_count(&self) -> Result<u128, SamplerError> {
        Ok(self.records.len() as u128)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::data::{QualityScore, RecordSection, SectionRole};
    use crate::types::RecordId;
    use chrono::Duration;

    /// Minimal `IndexableSource` test fixture.
    struct IndexableStub {
        id: SourceId,
        count: usize,
    }

    struct NoLenHintStub {
        id: SourceId,
    }

    impl IndexableStub {
        fn new(id: &str, count: usize) -> Self {
            Self {
                id: id.to_string(),
                count,
            }
        }
    }

    impl NoLenHintStub {
        fn new(id: &str) -> Self {
            Self { id: id.to_string() }
        }
    }

    impl IndexableSource for IndexableStub {
        fn id(&self) -> &str {
            &self.id
        }

        fn len_hint(&self) -> Option<usize> {
            Some(self.count)
        }

        fn record_at(&self, idx: usize) -> Result<Option<DataRecord>, SamplerError> {
            if idx >= self.count {
                return Ok(None);
            }
            let now = Utc::now();
            Ok(Some(DataRecord {
                id: format!("record_{idx}"),
                source: self.id.clone(),
                created_at: now,
                updated_at: now,
                quality: QualityScore { trust: 1.0 },
                taxonomy: Vec::new(),
                sections: vec![RecordSection {
                    role: SectionRole::Anchor,
                    heading: None,
                    text: "stub".into(),
                    sentences: vec!["stub".into()],
                }],
                meta_prefix: None,
            }))
        }
    }

    impl IndexableSource for NoLenHintStub {
        fn id(&self) -> &str {
            &self.id
        }

        fn len_hint(&self) -> Option<usize> {
            None
        }

        fn record_at(&self, _idx: usize) -> Result<Option<DataRecord>, SamplerError> {
            Ok(None)
        }
    }

    #[test]
    fn indexable_adapter_pages_in_stable_order() {
        let adapter = IndexableAdapter::new(IndexableStub::new("stub", 6));
        let full = adapter.refresh(None, None).unwrap();
        let full_ids: Vec<RecordId> = full.records.into_iter().map(|r| r.id).collect();

        let mut cursor = None;
        let mut paged = Vec::new();
        for _ in 0..3 {
            let snapshot = adapter.refresh(cursor.as_ref(), Some(2)).unwrap();
            cursor = Some(snapshot.cursor);
            paged.extend(snapshot.records.into_iter().map(|r| r.id));
        }
        assert_eq!(paged, full_ids);
    }

    #[test]
    fn indexable_paging_spans_multiple_regimes() {
        // Use a source id whose permutation step is not 1 or -1 mod 2^k,
        // otherwise the sequence would be a simple rotation/reversal.
        let total = 256usize;
        let mask = (1u64 << (64 - (total as u64 - 1).leading_zeros())) - 1;
        let source_id = (0..512)
            .map(|idx| format!("regime_test_{idx}"))
            .find(|id| {
                let seed = IndexablePager::seed_for(id, total);
                let a = (seed | 1) & mask;
                a != 1 && a != mask
            })
            .unwrap();

        // Pull a single page and ensure the indices are spread across the space,
        // which indicates the permutation isn't stuck in a narrow regime.
        let adapter = IndexableAdapter::new(IndexableStub::new(&source_id, total));
        let snapshot = adapter.refresh(None, Some(64)).unwrap();
        let indices: Vec<usize> = snapshot
            .records
            .into_iter()
            .map(|r| {
                r.id.strip_prefix("record_")
                    .unwrap()
                    .parse::<usize>()
                    .unwrap()
            })
            .collect();
        let min_idx = *indices.iter().min().unwrap();
        let max_idx = *indices.iter().max().unwrap();
        assert!(
            max_idx - min_idx >= total / 2,
            "expected spread across the index space, got min={min_idx} max={max_idx}"
        );
    }

    #[test]
    fn indexable_pager_errors_when_len_hint_missing() {
        let pager = IndexablePager::new("no_len_hint");
        let source = NoLenHintStub::new("no_len_hint");
        let result = pager.refresh(&source, None, Some(3));
        assert!(result.is_err());
    }

    #[test]
    fn indexable_adapter_reported_count_errors_when_len_hint_missing() {
        let adapter = IndexableAdapter::new(NoLenHintStub::new("no_len_hint"));
        let result = adapter.reported_record_count();
        assert!(result.is_err());
    }

    #[test]
    fn indexable_pager_refresh_with_zero_total_returns_empty_snapshot() {
        let pager = IndexablePager::new("empty");
        let snapshot = pager
            .refresh_with(0, None, Some(4), |_idx| Ok(None))
            .unwrap();
        assert!(snapshot.records.is_empty());
        assert_eq!(snapshot.cursor.revision, 0);
    }

    #[test]
    fn in_memory_source_refresh_wraps_cursor_and_uses_latest_timestamp() {
        let now = Utc::now();
        let older = now - Duration::seconds(5);
        let newer = now + Duration::seconds(5);
        let mk = |id: &str, ts: chrono::DateTime<Utc>| DataRecord {
            id: id.to_string(),
            source: "mem".to_string(),
            created_at: ts,
            updated_at: ts,
            quality: QualityScore { trust: 1.0 },
            taxonomy: Vec::new(),
            sections: vec![RecordSection {
                role: SectionRole::Anchor,
                heading: None,
                text: id.to_string(),
                sentences: vec![id.to_string()],
            }],
            meta_prefix: None,
        };

        let source = InMemorySource::new("mem", vec![mk("a", older), mk("b", newer)]);
        let cursor = SourceCursor {
            last_seen: now,
            revision: 7,
        };

        let snapshot = source.refresh(Some(&cursor), Some(1)).unwrap();
        assert_eq!(snapshot.records.len(), 1);
        assert_eq!(snapshot.records[0].id, "a");
        assert_eq!(snapshot.cursor.revision, 1);
        assert_eq!(snapshot.cursor.last_seen, older);
    }

    #[test]
    fn index_permutation_permute_bits_handles_zero_bits_and_zero_seed_path() {
        assert_eq!(IndexPermutation::permute_bits(123, 0, 99), 0);

        let bits = 1;
        let value = 1;
        let out = IndexPermutation::permute_bits(value, bits, 0);
        assert!(out <= 1);
    }

    #[test]
    fn index_permutation_next_stays_within_total_and_cursor_advances() {
        let mut perm = IndexPermutation::new(3, 7, 0);
        let mut seen = Vec::new();
        for _ in 0..8 {
            seen.push(perm.next());
        }
        assert!(seen.iter().all(|idx| *idx < 3));
        assert!(perm.cursor() < 3);
    }
}

1	//! Data source interfaces and paging helpers.
2	//!
3	//! Ownership model:
4	//! - `DataSource` is the sampler-facing interface that produces batches.
5	//! - `IndexableSource` exposes stable, index-based access into a corpus.
6	//! - `IndexablePager` owns the deterministic pseudo-random paging logic and
7	//! can page any indexable source without retaining per-record state.
8
9	use chrono::{DateTime, Utc};
10	use std::hash::Hash;
11	use std::sync::Arc;
12	use std::time::{Duration, Instant};
13
14	use crate::config::{SamplerConfig, TripletRecipe};
15	use crate::data::DataRecord;
16	use crate::errors::SamplerError;
17	use crate::hash::stable_hash_with;
18	use crate::types::SourceId;
19
20	/// Source implementation modules.
21	pub mod sources;
22	/// Utility helpers used by source implementations.
23	pub mod utilities;
24	pub use sources::file_source::{
25	FileSource, FileSourceConfig, SectionBuilder, TaxonomyBuilder, anchor_context_sections,
26	taxonomy_from_path,
27	};
28	#[cfg(feature = "huggingface")]
29	pub use sources::huggingface_source::{HuggingFaceRowSource, HuggingFaceRowsConfig};
30
31	/// Source-owned incremental refresh position.
32	///
33	/// The sampler stores and returns this value between refresh calls.
34	/// `revision` is opaque to the sampler and interpreted only by the source.
35	#[derive(Clone, Debug)]
36	pub struct SourceCursor {
37	/// Most recent observation timestamp produced by the source.
38	pub last_seen: DateTime<Utc>,
39	/// Opaque paging position token used to continue incremental refresh.
40	pub revision: u64,
41	}
42
43	/// Result of a single source refresh call.
44	///
45	/// Pass the returned `cursor` back into the next refresh to continue paging.
46	#[derive(Clone, Debug)]
47	pub struct SourceSnapshot {
48	/// Records returned by the refresh operation.
49	pub records: Vec<DataRecord>,
50	/// Next cursor to pass into a future refresh call.
51	pub cursor: SourceCursor,
52	}
53
54	/// Sampler-facing data source interface.
55	///
56	/// Implementations may be streaming or index-backed. For a fixed dataset state
57	/// and cursor, refresh output should be deterministic.
58	pub trait DataSource: Send + Sync {
59	/// Stable source identifier used in records, metrics, and persistence state.
60	fn id(&self) -> &str;
61	/// Fetch up to `limit` records starting from `cursor` state.
62	///
63	/// Return the next cursor position in `SourceSnapshot.cursor`.
64	fn refresh(
65	&self,
66	cursor: Option<&SourceCursor>,
67	limit: Option<usize>,
68	) -> Result<SourceSnapshot, SamplerError>;
69
70	/// Exact metadata record count reported by the source.
71	///
72	/// This is intended for estimators that must avoid iterating records.
73	/// Implementations should return `Ok(count)` only when the count is
74	/// exact for the source scope. Return `Err` when exact counting is not
75	/// possible or the source is unavailable.
76	///
77	/// Keep this consistent with `refresh` by using the same backend scope,
78	/// filtering, and logical corpus definition.
79	fn reported_record_count(&self) -> Result<u128, SamplerError>;
80
81	/// Provide the active sampler configuration to this source.
82	///
83	/// Called when the source is registered with a sampler. Sources can use
84	/// this to align internal heuristics with runtime sampler settings.
85	fn configure_sampler(&self, _config: &SamplerConfig) {}	186✔
86
87	/// Optional source-provided default triplet recipes.
88	///
89	/// Used when sampler config does not provide explicit recipes.
90	fn default_triplet_recipes(&self) -> Vec<TripletRecipe> {	25✔
91	Vec::new()	25✔
92	}	25✔
93	}
94
95	/// Index-addressable source interface used by deterministic pagers.
96	///
97	/// `len_hint` must be stable within an epoch, and `record_at` must return the
98	/// record corresponding to the same index across runs.
99	///
100	/// Dense indexing is strongly recommended: implement indices as `0..len_hint`
101	/// with minimal gaps. Sparse indexes (returning `None` for many positions)
102	/// still work but waste paging capacity and reduce batch fill rates.
103	pub trait IndexableSource: Send + Sync {
104	/// Stable source identifier.
105	fn id(&self) -> &str;
106	/// Current index domain size, typically `Some(total_records)`.
107	fn len_hint(&self) -> Option<usize>;
108	/// Return the record at index `idx`, or `None` for sparse/missing positions.
109	fn record_at(&self, idx: usize) -> Result<Option<DataRecord>, SamplerError>;
110	}
111
112	/// Deterministic pager for `IndexableSource`.
113	///
114	/// Encapsulates shuffle seed and cursor math so callers can reuse a stable
115	/// paging algorithm without implementing permutation logic themselves.
116	pub struct IndexablePager {
117	source_id: SourceId,
118	}
119
120	impl IndexablePager {
121	/// Create a new deterministic pager for `source_id`.
122	pub fn new(source_id: impl Into<SourceId>) -> Self {	7✔
123	Self {	7✔
124	source_id: source_id.into(),	7✔
125	}	7✔
126	}	7✔
127
128	/// Page records from an `IndexableSource` using the provided cursor.
129	pub fn refresh(	6✔
130	&self,	6✔
131	source: &dyn IndexableSource,	6✔
132	cursor: Option<&SourceCursor>,	6✔
133	limit: Option<usize>,	6✔
134	) -> Result<SourceSnapshot, SamplerError> {	6✔
135	let total = source	6✔
136	.len_hint()	6✔
137	.ok_or_else(\|\| SamplerError::SourceInconsistent {	6✔
138	source_id: source.id().to_string(),	1✔
139	details: "indexable source did not provide len_hint".into(),	1✔
140	})?;	1✔
141	self.refresh_with(total, cursor, limit, \|idx\| source.record_at(idx))	76✔
142	}	6✔
143
144	/// Page records using a custom index fetcher.
145	///
146	/// Useful when records are indexable but not exposed through `IndexableSource`
147	/// (for example, temporary index stores or precomputed path lists).
148	pub fn refresh_with(	6✔
149	&self,	6✔
150	total: usize,	6✔
151	cursor: Option<&SourceCursor>,	6✔
152	limit: Option<usize>,	6✔
153	mut fetch: impl FnMut(usize) -> Result<Option<DataRecord>, SamplerError>,	6✔
154	) -> Result<SourceSnapshot, SamplerError> {	6✔
155	if total == 0 {	6✔
156	return Ok(SourceSnapshot {	1✔
157	records: Vec::new(),	1✔
158	cursor: SourceCursor {	1✔
159	last_seen: Utc::now(),	1✔
160	revision: 0,	1✔
161	},	1✔
162	});	1✔
163	}	5✔
164	let mut start = cursor.map(\|cursor\| cursor.revision as usize).unwrap_or(0);	5✔
165	if start >= total {	5✔
166	start = 0;	×
167	}	5✔
168	let max = limit.unwrap_or(total);	5✔
169	let mut records = Vec::new();	5✔
170	let seed = Self::seed_for(&self.source_id, total);	5✔
171	let mut permutation = IndexPermutation::new(total, seed, start as u64);	5✔
172	let report_every = Duration::from_millis(750);	5✔
173	let refresh_start = Instant::now();	5✔
174	let mut last_report = refresh_start;	5✔
175	let mut attempts = 0usize;	5✔
176	let should_report = total >= 10_000 \|\| max >= 1_024;	5✔
177	if should_report {	5✔
NEW 178	eprintln!(	×
NEW 179	"[triplets:source] refresh start source='{}' total={} target={}",	×
NEW 180	self.source_id, total, max	×
NEW 181	);	×
182	}	5✔
183	for _ in 0..total {	5✔
184	attempts += 1;	80✔
185	if records.len() >= max {	80✔
186	break;	4✔
187	}	76✔
188	let idx = permutation.next();	76✔
189	if let Some(record) = fetch(idx)? {	76✔
190	records.push(record);	76✔
191	}	76✔
192	if should_report && last_report.elapsed() >= report_every {	76✔
NEW 193	eprintln!(	×
NEW 194	"[triplets:source] refresh progress source='{}' attempted={}/{} fetched={}/{} elapsed={:.1}s",	×
NEW 195	self.source_id,	×
NEW 196	attempts,	×
NEW 197	total,	×
NEW 198	records.len(),	×
NEW 199	max,	×
NEW 200	refresh_start.elapsed().as_secs_f64()	×
NEW 201	);	×
NEW 202	last_report = Instant::now();	×
203	}	76✔
204	}
205	if should_report {	5✔
NEW 206	eprintln!(	×
NEW 207	"[triplets:source] refresh done source='{}' attempted={} fetched={} elapsed={:.2}s",	×
NEW 208	self.source_id,	×
NEW 209	attempts,	×
NEW 210	records.len(),	×
NEW 211	refresh_start.elapsed().as_secs_f64()	×
NEW 212	);	×
213	}	5✔
214	let last_seen = records	5✔
215	.iter()	5✔
216	.map(\|record\| record.updated_at)	5✔
217	.max()	5✔
218	.unwrap_or_else(Utc::now);	5✔
219	let next_start = permutation.cursor();	5✔
220	Ok(SourceSnapshot {	5✔
221	records,	5✔
222	cursor: SourceCursor {	5✔
223	last_seen,	5✔
224	revision: next_start as u64,	5✔
225	},	5✔
226	})	5✔
227	}	6✔
228
229	/// Build a deterministic seed for a source and total size.
230	pub(crate) fn seed_for(source_id: &SourceId, total: usize) -> u64 {	14✔
231	Self::stable_index_shuffle_key(source_id, 0)	14✔
232	^ Self::stable_index_shuffle_key(source_id, total)	14✔
233	}	14✔
234
235	fn stable_index_shuffle_key(source_id: &SourceId, idx: usize) -> u64 {	28✔
236	stable_hash_with(\|hasher\| {	28✔
237	source_id.hash(hasher);	28✔
238	idx.hash(hasher);	28✔
239	})	28✔
240	}	28✔
241	}
242
243	/// DataSource adapter that pages an `IndexableSource` via `IndexablePager`.
244	pub struct IndexableAdapter<T: IndexableSource> {
245	inner: T,
246	}
247
248	impl<T: IndexableSource> IndexableAdapter<T> {
249	/// Wrap an `IndexableSource` so it can be registered as a `DataSource`.
250	pub fn new(inner: T) -> Self {	3✔
251	Self { inner }	3✔
252	}	3✔
253	}
254
255	impl<T: IndexableSource> DataSource for IndexableAdapter<T> {
256	fn id(&self) -> &str {	×
257	self.inner.id()	×
258	}	×
259
260	fn refresh(	5✔
261	&self,	5✔
262	cursor: Option<&SourceCursor>,	5✔
263	limit: Option<usize>,	5✔
264	) -> Result<SourceSnapshot, SamplerError> {	5✔
265	let pager = IndexablePager::new(self.inner.id());	5✔
266	pager.refresh(&self.inner, cursor, limit)	5✔
267	}	5✔
268
269	fn reported_record_count(&self) -> Result<u128, SamplerError> {	1✔
270	self.inner	1✔
271	.len_hint()	1✔
272	.map(\|value\| value as u128)	1✔
273	.ok_or_else(\|\| SamplerError::SourceInconsistent {	1✔
274	source_id: self.inner.id().to_string(),	1✔
275	details: "indexable source did not provide len_hint".into(),	1✔
276	})	1✔
277	}	1✔
278	}
279
280	/// Internal permutation used by `IndexablePager`.
281	pub(crate) struct IndexPermutation {
282	total: u64,
283	domain_bits: u32,
284	domain_size: u64,
285	seed: u64,
286	counter: u64,
287	}
288
289	impl IndexPermutation {
290	fn new(total: usize, seed: u64, counter: u64) -> Self {	40✔
291	let total_u64 = total as u64;	40✔
292	let domain_bits = (64 - (total_u64 - 1).leading_zeros()).max(1);	40✔
293	let domain_size = 1u64 << domain_bits;	40✔
294	Self {	40✔
295	total: total_u64,	40✔
296	domain_bits,	40✔
297	domain_size,	40✔
298	seed,	40✔
299	counter,	40✔
300	}	40✔
301	}	40✔
302
303	fn next(&mut self) -> usize {	574✔
304	loop {
305	let v =	623✔
306	Self::permute_bits(self.counter % self.domain_size, self.domain_bits, self.seed);	623✔
307	self.counter = self.counter.wrapping_add(1);	623✔
308	if v < self.total {	623✔
309	return v as usize;	574✔
310	}	49✔
311	}
312	}	574✔
313
314	fn cursor(&self) -> usize {	20✔
315	(self.counter as usize) % (self.total as usize)	20✔
316	}	20✔
317	fn permute_bits(value: u64, bits: u32, seed: u64) -> u64 {	625✔
318	if bits == 0 {	625✔
319	return 0;	1✔
320	}	624✔
321	let mask = if bits == 64 {	624✔
322	u64::MAX	×
323	} else {
324	(1u64 << bits) - 1	624✔
325	};
326	let mut a = (seed \| 1) & mask;	624✔
327	if a == 0 {	624✔
328	a = 1;	×
329	}	624✔
330	let b = (seed >> 1) & mask;	624✔
331	a.wrapping_mul(value).wrapping_add(b) & mask	624✔
332	}	625✔
333	}
334
335	/// In-memory data source for tests and small datasets.
336	pub struct InMemorySource {
337	id: SourceId,
338	records: Arc<Vec<DataRecord>>,
339	}
340
341	impl InMemorySource {
342	/// Create an in-memory source from prebuilt records.
343	pub fn new(id: impl Into<SourceId>, records: Vec<DataRecord>) -> Self {	191✔
344	Self {	191✔
345	id: id.into(),	191✔
346	records: Arc::new(records),	191✔
347	}	191✔
348	}	191✔
349	}
350
351	impl DataSource for InMemorySource {
352	fn id(&self) -> &str {	4,873✔
353	&self.id	4,873✔
354	}	4,873✔
355
356	fn refresh(	1,350✔
357	&self,	1,350✔
358	cursor: Option<&SourceCursor>,	1,350✔
359	limit: Option<usize>,	1,350✔
360	) -> Result<SourceSnapshot, SamplerError> {	1,350✔
361	let records = &*self.records;	1,350✔
362	let total = records.len();	1,350✔
363	let mut start = cursor.map(\|cursor\| cursor.revision as usize).unwrap_or(0);	1,350✔
364	if total > 0 && start >= total {	1,350✔
365	start = 0;	1✔
366	}	1,349✔
367	let max = limit.unwrap_or(total);	1,350✔
368	let mut filtered = Vec::new();	1,350✔
369	for idx in 0..total {	9,434✔
370	if filtered.len() >= max {	9,434✔
371	break;	1,138✔
372	}	8,296✔
373	let pos = (start + idx) % total;	8,296✔
374	filtered.push(records[pos].clone());	8,296✔
375	}
376	let last_seen = filtered	1,350✔
377	.iter()	1,350✔
378	.map(\|record\| record.updated_at)	1,350✔
379	.max()	1,350✔
380	.unwrap_or_else(Utc::now);	1,350✔
381	let next_start = if total == 0 {	1,350✔
382	0	×
383	} else {
384	(start + filtered.len()) % total	1,350✔
385	};
386	Ok(SourceSnapshot {	1,350✔
387	records: filtered,	1,350✔
388	cursor: SourceCursor {	1,350✔
389	last_seen,	1,350✔
390	revision: next_start as u64,	1,350✔
391	},	1,350✔
392	})	1,350✔
393	}	1,350✔
394
NEW 395	fn reported_record_count(&self) -> Result<u128, SamplerError> {	×
NEW 396	Ok(self.records.len() as u128)	×
NEW 397	}	×
398	}
399
400	#[cfg(test)]
401	mod tests {
402	use super::*;
403	use crate::data::{QualityScore, RecordSection, SectionRole};
404	use crate::types::RecordId;
405	use chrono::Duration;
406
407	/// Minimal `IndexableSource` test fixture.
408	struct IndexableStub {
409	id: SourceId,
410	count: usize,
411	}
412
413	struct NoLenHintStub {
414	id: SourceId,
415	}
416
417	impl IndexableStub {
418	fn new(id: &str, count: usize) -> Self {	2✔
419	Self {	2✔
420	id: id.to_string(),	2✔
421	count,	2✔
422	}	2✔
423	}	2✔
424	}
425
426	impl NoLenHintStub {
427	fn new(id: &str) -> Self {	2✔
428	Self { id: id.to_string() }	2✔
429	}	2✔
430	}
431
432	impl IndexableSource for IndexableStub {
433	fn id(&self) -> &str {	5✔
434	&self.id	5✔
435	}	5✔
436
437	fn len_hint(&self) -> Option<usize> {	5✔
438	Some(self.count)	5✔
439	}	5✔
440
441	fn record_at(&self, idx: usize) -> Result<Option<DataRecord>, SamplerError> {	76✔
442	if idx >= self.count {	76✔
443	return Ok(None);	×
444	}	76✔
445	let now = Utc::now();	76✔
446	Ok(Some(DataRecord {	76✔
447	id: format!("record_{idx}"),	76✔
448	source: self.id.clone(),	76✔
449	created_at: now,	76✔
450	updated_at: now,	76✔
451	quality: QualityScore { trust: 1.0 },	76✔
452	taxonomy: Vec::new(),	76✔
453	sections: vec![RecordSection {	76✔
454	role: SectionRole::Anchor,	76✔
455	heading: None,	76✔
456	text: "stub".into(),	76✔
457	sentences: vec!["stub".into()],	76✔
458	}],	76✔
459	meta_prefix: None,	76✔
460	}))	76✔
461	}	76✔
462	}
463
464	impl IndexableSource for NoLenHintStub {
465	fn id(&self) -> &str {	2✔
466	&self.id	2✔
467	}	2✔
468
469	fn len_hint(&self) -> Option<usize> {	2✔
470	None	2✔
471	}	2✔
472
NEW 473	fn record_at(&self, _idx: usize) -> Result<Option<DataRecord>, SamplerError> {	×
NEW 474	Ok(None)	×
NEW 475	}	×
476	}
477
478	#[test]
479	fn indexable_adapter_pages_in_stable_order() {	1✔
480	let adapter = IndexableAdapter::new(IndexableStub::new("stub", 6));	1✔
481	let full = adapter.refresh(None, None).unwrap();	1✔
482	let full_ids: Vec<RecordId> = full.records.into_iter().map(\|r\| r.id).collect();	1✔
483
484	let mut cursor = None;	1✔
485	let mut paged = Vec::new();	1✔
486	for _ in 0..3 {	1✔
487	let snapshot = adapter.refresh(cursor.as_ref(), Some(2)).unwrap();	3✔
488	cursor = Some(snapshot.cursor);	3✔
489	paged.extend(snapshot.records.into_iter().map(\|r\| r.id));	3✔
490	}
491	assert_eq!(paged, full_ids);	1✔
492	}	1✔
493
494	#[test]
495	fn indexable_paging_spans_multiple_regimes() {	1✔
496	// Use a source id whose permutation step is not 1 or -1 mod 2^k,
497	// otherwise the sequence would be a simple rotation/reversal.
498	let total = 256usize;	1✔
499	let mask = (1u64 << (64 - (total as u64 - 1).leading_zeros())) - 1;	1✔
500	let source_id = (0..512)	1✔
501	.map(\|idx\| format!("regime_test_{idx}"))	1✔
502	.find(\|id\| {	1✔
503	let seed = IndexablePager::seed_for(id, total);	1✔
504	let a = (seed \| 1) & mask;	1✔
505	a != 1 && a != mask	1✔
506	})	1✔
507	.unwrap();	1✔
508
509	// Pull a single page and ensure the indices are spread across the space,
510	// which indicates the permutation isn't stuck in a narrow regime.
511	let adapter = IndexableAdapter::new(IndexableStub::new(&source_id, total));	1✔
512	let snapshot = adapter.refresh(None, Some(64)).unwrap();	1✔
513	let indices: Vec<usize> = snapshot	1✔
514	.records	1✔
515	.into_iter()	1✔
516	.map(\|r\| {	64✔
517	r.id.strip_prefix("record_")	64✔
518	.unwrap()	64✔
519	.parse::<usize>()	64✔
520	.unwrap()	64✔
521	})	64✔
522	.collect();	1✔
523	let min_idx = *indices.iter().min().unwrap();	1✔
524	let max_idx = *indices.iter().max().unwrap();	1✔
525	assert!(	1✔
526	max_idx - min_idx >= total / 2,	1✔
527	"expected spread across the index space, got min={min_idx} max={max_idx}"
528	);
529	}	1✔
530
531	#[test]
532	fn indexable_pager_errors_when_len_hint_missing() {	1✔
533	let pager = IndexablePager::new("no_len_hint");	1✔
534	let source = NoLenHintStub::new("no_len_hint");	1✔
535	let result = pager.refresh(&source, None, Some(3));	1✔
536	assert!(result.is_err());	1✔
537	}	1✔
538
539	#[test]
540	fn indexable_adapter_reported_count_errors_when_len_hint_missing() {	1✔
541	let adapter = IndexableAdapter::new(NoLenHintStub::new("no_len_hint"));	1✔
542	let result = adapter.reported_record_count();	1✔
543	assert!(result.is_err());	1✔
544	}	1✔
545
546	#[test]
547	fn indexable_pager_refresh_with_zero_total_returns_empty_snapshot() {	1✔
548	let pager = IndexablePager::new("empty");	1✔
549	let snapshot = pager	1✔
550	.refresh_with(0, None, Some(4), \|_idx\| Ok(None))	1✔
551	.unwrap();	1✔
552	assert!(snapshot.records.is_empty());	1✔
553	assert_eq!(snapshot.cursor.revision, 0);	1✔
554	}	1✔
555
556	#[test]
557	fn in_memory_source_refresh_wraps_cursor_and_uses_latest_timestamp() {	1✔
558	let now = Utc::now();	1✔
559	let older = now - Duration::seconds(5);	1✔
560	let newer = now + Duration::seconds(5);	1✔
561	let mk = \|id: &str, ts: chrono::DateTime<Utc>\| DataRecord {	1✔
562	id: id.to_string(),	2✔
563	source: "mem".to_string(),	2✔
564	created_at: ts,	2✔
565	updated_at: ts,	2✔
566	quality: QualityScore { trust: 1.0 },	2✔
567	taxonomy: Vec::new(),	2✔
568	sections: vec![RecordSection {	2✔
569	role: SectionRole::Anchor,	2✔
570	heading: None,	2✔
571	text: id.to_string(),	2✔
572	sentences: vec![id.to_string()],	2✔
573	}],	2✔
574	meta_prefix: None,	2✔
575	};	2✔
576
577	let source = InMemorySource::new("mem", vec![mk("a", older), mk("b", newer)]);	1✔
578	let cursor = SourceCursor {	1✔
579	last_seen: now,	1✔
580	revision: 7,	1✔
581	};	1✔
582
583	let snapshot = source.refresh(Some(&cursor), Some(1)).unwrap();	1✔
584	assert_eq!(snapshot.records.len(), 1);	1✔
585	assert_eq!(snapshot.records[0].id, "a");	1✔
586	assert_eq!(snapshot.cursor.revision, 1);	1✔
587	assert_eq!(snapshot.cursor.last_seen, older);	1✔
588	}	1✔
589
590	#[test]
591	fn index_permutation_permute_bits_handles_zero_bits_and_zero_seed_path() {	1✔
592	assert_eq!(IndexPermutation::permute_bits(123, 0, 99), 0);	1✔
593
594	let bits = 1;	1✔
595	let value = 1;	1✔
596	let out = IndexPermutation::permute_bits(value, bits, 0);	1✔
597	assert!(out <= 1);	1✔
598	}	1✔
599
600	#[test]
601	fn index_permutation_next_stays_within_total_and_cursor_advances() {	1✔
602	let mut perm = IndexPermutation::new(3, 7, 0);	1✔
603	let mut seen = Vec::new();	1✔
604	for _ in 0..8 {	8✔
605	seen.push(perm.next());	8✔
606	}	8✔
607	assert!(seen.iter().all(\|idx\| *idx < 3));	8✔
608	assert!(perm.cursor() < 3);	1✔
609	}	1✔
610	}

jzombie / rust-triplets / 22337622218

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous