rskit_dataset/record/
source.rs

1//! Record source adapter — stream a [`DatasetReader`] as a [`Source`] of records.
2
3use serde_json::{Value, json};
4use tokio_util::sync::CancellationToken;
5
6use crate::source::{BoxItemStream, Source};
7
8use super::model::DatasetRecord;
9use super::reader::DatasetReader;
10
11/// Adapts any [`DatasetReader`] into a [`Source`] the generic engine can collect.
12pub struct RecordSource {
13    name: String,
14    reader: Box<dyn DatasetReader>,
15    cache_key: Value,
16    max_items: Option<usize>,
17}
18
19impl RecordSource {
20    /// Create a record source with a stable name over the given reader.
21    #[must_use]
22    pub fn new(name: impl Into<String>, reader: Box<dyn DatasetReader>) -> Self {
23        let name = name.into();
24        let cache_key = json!({ "record-source": name });
25        Self {
26            name,
27            reader,
28            cache_key,
29            max_items: None,
30        }
31    }
32
33    /// Override the cache key describing this source's configured inputs.
34    #[must_use]
35    pub fn with_cache_key(mut self, cache_key: Value) -> Self {
36        self.cache_key = cache_key;
37        self
38    }
39
40    /// Declare the maximum number of records this source will emit.
41    #[must_use]
42    pub fn with_max_items(mut self, max_items: usize) -> Self {
43        self.max_items = Some(max_items);
44        self
45    }
46}
47
48impl Source<DatasetRecord> for RecordSource {
49    fn name(&self) -> &str {
50        &self.name
51    }
52
53    fn stream(self: Box<Self>, _cancel: CancellationToken) -> BoxItemStream<DatasetRecord> {
54        self.reader.stream()
55    }
56
57    fn cache_key(&self) -> Value {
58        self.cache_key.clone()
59    }
60
61    fn max_items(&self) -> Option<usize> {
62        self.max_items
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use futures_util::StreamExt as _;
69
70    use super::*;
71    use crate::record::JsonLinesReader;
72
73    #[tokio::test]
74    async fn record_source_streams_reader_output_and_exposes_metadata() {
75        let dir = tempfile::tempdir().unwrap();
76        let path = dir.path().join("data.jsonl");
77        std::fs::write(&path, "{\"id\":1}\n{\"id\":2}\n").unwrap();
78
79        let source = RecordSource::new("records", Box::new(JsonLinesReader::new(&path)))
80            .with_max_items(2)
81            .with_cache_key(json!({"path": "data.jsonl"}));
82        assert_eq!(source.name(), "records");
83        assert_eq!(source.max_items(), Some(2));
84        assert_eq!(source.cache_key(), json!({"path": "data.jsonl"}));
85
86        let collected = Box::new(source)
87            .stream(CancellationToken::new())
88            .collect::<Vec<_>>()
89            .await;
90        assert_eq!(collected.len(), 2);
91        assert!(collected.iter().all(Result::is_ok));
92    }
93}
rskit_dataset/record/source.rs

rskit_dataset/record/
source.rs
