ruda_dataset/dataset/
indexed_jsonl.rs

1//! Random-access JSONL without retaining deserialized examples in memory.
2use super::Dataset;
3use serde::{Deserialize, Serialize, de::DeserializeOwned};
4use std::{
5    fs::File,
6    io::{self, Read, Seek, SeekFrom},
7    marker::PhantomData,
8    path::Path,
9    sync::{Arc, Mutex},
10    time::UNIX_EPOCH,
11};
12
13fn invalid(message: impl Into<String>) -> io::Error {
14    io::Error::new(io::ErrorKind::InvalidData, message.into())
15}
16
17/// Caller-selected handling of physically empty or JSON-whitespace-only rows.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
19pub enum JsonlBlankLines {
20    /// Return an error without assigning that row a dataset index.
21    Reject,
22    /// Exclude blank rows, preserving the order of all other rows.
23    Skip,
24}
25
26/// Immutable-source identity and explicit JSONL indexing policies.
27#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
28pub struct JsonlIndexOptions {
29    /// Caller-supplied identity of the immutable source/version, not its path.
30    pub source_id: String,
31    /// Whether physically blank rows belong to an invalid source or are skipped.
32    pub blank_lines: JsonlBlankLines,
33    /// Optional maximum raw row length, excluding LF but including CR.
34    /// None imposes no row-size limit.
35    pub max_record_bytes: Option<u64>,
36}
37
38#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
39struct FileIdentity {
40    length: u64,
41    modified: Option<(u64, u32)>,
42}
43
44impl FileIdentity {
45    fn read(file: &File) -> io::Result<Self> {
46        let metadata = file.metadata()?;
47        if !metadata.is_file() {
48            return Err(invalid("indexed JSONL requires a regular immutable file"));
49        }
50        let modified = metadata.modified().ok()
51            .and_then(|time| time.duration_since(UNIX_EPOCH).ok())
52            .map(|duration| (duration.as_secs(), duration.subsec_nanos()));
53        Ok(Self { length: metadata.len(), modified })
54    }
55
56    fn check(&self, file: &File) -> io::Result<()> {
57        if self != &Self::read(file)? {
58            return Err(invalid("JSONL source length or modification time changed"));
59        }
60        Ok(())
61    }
62}
63
64/// Actual byte range of one retained physical row, without its LF delimiter.
65#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
66pub struct JsonlRecordLocation {
67    /// Absolute byte position in the original source.
68    pub offset: u64,
69    /// Raw row length; CR in a CRLF row is retained as JSON whitespace.
70    pub length: u64,
71    /// One-based physical row number, including any skipped blank rows.
72    pub line_number: u64,
73}
74
75/// Serializable indexing checkpoint; no source data or filesystem sidecar.
76///
77/// Byte ranges retain physical source order. The caller owns the source's
78/// immutability and its version identity: length/mtime are change detection,
79/// not a content hash. The index records rows, not their deserialization result.
80#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
81pub struct JsonlIndexState {
82    version: u32,
83    options: JsonlIndexOptions,
84    identity: FileIdentity,
85    processed_bytes: u64,
86    pending_start: u64,
87    pending_nonblank: bool,
88    next_line: u64,
89    skipped_lines: u64,
90    records: Vec<JsonlRecordLocation>,
91    finished: bool,
92}
93
94/// Work actually completed by a bounded indexing call.
95#[derive(Clone, Copy, Debug, PartialEq, Eq)]
96pub struct JsonlIndexProgress {
97    /// Bytes whose indexing state has been committed, not buffered read-ahead.
98    pub processed_bytes: u64,
99    /// Immutable file length, including delimiters and skipped rows.
100    pub total_bytes: u64,
101    /// Completed retained records; a partial final record is not counted.
102    pub records: usize,
103    /// Physically blank rows explicitly excluded by the selected policy.
104    pub skipped_lines: u64,
105    /// True only after the final physical row has been finalized.
106    pub finished: bool,
107}
108
109impl JsonlIndexState {
110    /// Source identity and row policies recorded by this checkpoint.
111    pub fn options(&self) -> &JsonlIndexOptions { &self.options }
112
113    /// The completed rows in their original source order.
114    pub fn records(&self) -> &[JsonlRecordLocation] { &self.records }
115
116    /// Actual bytes/records processed so far, including an unfinished row.
117    pub fn progress(&self) -> JsonlIndexProgress {
118        JsonlIndexProgress {
119            processed_bytes: self.processed_bytes,
120            total_bytes: self.identity.length,
121            records: self.records.len(),
122            skipped_lines: self.skipped_lines,
123            finished: self.finished,
124        }
125    }
126
127    /// Validate checkpoint accounting before restoring it against a source.
128    pub fn validate(&self) -> io::Result<()> {
129        if self.version != 1 || self.options.source_id.is_empty() || self.next_line == 0 {
130            return Err(invalid("invalid JSONL checkpoint version or source identity"));
131        }
132        if self.pending_start > self.processed_bytes || self.processed_bytes > self.identity.length
133            || (self.pending_start == self.processed_bytes && self.pending_nonblank)
134            || (self.finished && (self.processed_bytes != self.identity.length
135                || self.pending_start != self.processed_bytes)) {
136            return Err(invalid("JSONL checkpoint byte accounting is inconsistent"));
137        }
138        let completed_lines = u64::try_from(self.records.len())
139            .ok().and_then(|count| count.checked_add(self.skipped_lines));
140        if completed_lines.and_then(|count| count.checked_add(1)) != Some(self.next_line) {
141            return Err(invalid("JSONL checkpoint physical row accounting is inconsistent"));
142        }
143        let pending_length = self.processed_bytes - self.pending_start;
144        if self.options.max_record_bytes.is_some_and(|limit| pending_length > limit) {
145            return Err(invalid("JSONL checkpoint contains an oversized pending row"));
146        }
147        let mut previous_end = 0;
148        let mut previous_line = 0;
149        for row in &self.records {
150            let end = row.offset.checked_add(row.length)
151                .ok_or_else(|| invalid("JSONL row byte range overflow"))?;
152            if row.length == 0 || row.offset < previous_end || end > self.pending_start
153                || row.line_number <= previous_line || row.line_number >= self.next_line
154                || self.options.max_record_bytes.is_some_and(|limit| row.length > limit) {
155                return Err(invalid("JSONL checkpoint row locations are inconsistent"));
156            }
157            previous_end = end.checked_add(1).unwrap_or(end);
158            previous_line = row.line_number;
159        }
160        Ok(())
161    }
162}
163
164/// Incremental byte-range indexer with serializable, caller-persisted state.
165///
166/// `scan_step` reads at most its byte budget. Index memory is proportional to
167/// completed rows; scan memory is bounded independently of row/file length.
168/// Invalid rows are never silently truncated, replaced, or indexed as blanks.
169pub struct JsonlIndexBuilder {
170    file: File,
171    state: JsonlIndexState,
172}
173
174impl JsonlIndexBuilder {
175    /// Open a source using explicit blank-line and row-size policies.
176    pub fn open(path: impl AsRef<Path>, options: JsonlIndexOptions) -> io::Result<Self> {
177        if options.source_id.is_empty() {
178            return Err(invalid("supply an immutable JSONL source identity"));
179        }
180        let file = File::open(path)?;
181        let identity = FileIdentity::read(&file)?;
182        let state = JsonlIndexState {
183            version: 1, options, identity, processed_bytes: 0, pending_start: 0,
184            pending_nonblank: false, next_line: 1, skipped_lines: 0,
185            records: Vec::new(), finished: false,
186        };
187        Ok(Self { file, state })
188    }
189
190    /// Resume exact byte/row accounting, including a partially scanned row.
191    /// A matching caller identity is required even if the path has changed.
192    pub fn resume(path: impl AsRef<Path>, source_id: &str, state: JsonlIndexState) -> io::Result<Self> {
193        state.validate()?;
194        if source_id != state.options.source_id {
195            return Err(invalid("JSONL continuation belongs to a different source version"));
196        }
197        let file = File::open(path)?;
198        state.identity.check(&file)?;
199        Ok(Self { file, state })
200    }
201
202    /// Current state, suitable for caller-controlled checkpoint serialization.
203    pub fn state(&self) -> &JsonlIndexState { &self.state }
204
205    /// Current actual byte/record work counters.
206    pub fn progress(&self) -> JsonlIndexProgress { self.state.progress() }
207
208    fn finish_line(&mut self) -> io::Result<()> {
209        let length = self.state.processed_bytes - self.state.pending_start;
210        let next_line = self.state.next_line.checked_add(1)
211            .ok_or_else(|| invalid("JSONL physical line count overflow"))?;
212        if !self.state.pending_nonblank {
213            match self.state.options.blank_lines {
214                JsonlBlankLines::Reject => return Err(invalid(format!(
215                    "blank JSONL row at physical line {}", self.state.next_line))),
216                JsonlBlankLines::Skip => self.state.skipped_lines += 1,
217            }
218        } else {
219            self.state.records.push(JsonlRecordLocation {
220                offset: self.state.pending_start, length, line_number: self.state.next_line,
221            });
222        }
223        self.state.next_line = next_line;
224        self.state.pending_nonblank = false;
225        Ok(())
226    }
227
228    /// Scan at most `byte_budget` source bytes, retaining exact restart state.
229    /// Zero returns current counters without performing work. EOF finalizes a
230    /// nonempty unterminated final row without adding a phantom trailing row.
231    pub fn scan_step(&mut self, byte_budget: u64) -> io::Result<JsonlIndexProgress> {
232        self.state.identity.check(&self.file)?;
233        if self.state.finished || byte_budget == 0 { return Ok(self.progress()); }
234        self.file.seek(SeekFrom::Start(self.state.processed_bytes))?;
235        let mut remaining = byte_budget.min(self.state.identity.length - self.state.processed_bytes);
236        let mut buffer = [0u8; 16 * 1024];
237        while remaining != 0 {
238            let length = remaining.min(buffer.len() as u64) as usize;
239            self.file.read_exact(&mut buffer[..length])?;
240            for &byte in &buffer[..length] {
241                if byte == b'\n' {
242                    self.finish_line()?;
243                    self.state.processed_bytes += 1;
244                    self.state.pending_start = self.state.processed_bytes;
245                } else {
246                    let row_length = self.state.processed_bytes - self.state.pending_start;
247                    if self.state.options.max_record_bytes.is_some_and(|limit| row_length >= limit) {
248                        return Err(invalid(format!("JSONL row {} exceeds the selected byte limit", self.state.next_line)));
249                    }
250                    self.state.pending_nonblank |= !matches!(byte, b' ' | b'\t' | b'\r');
251                    self.state.processed_bytes += 1;
252                }
253            }
254            remaining -= length as u64;
255        }
256        if self.state.processed_bytes == self.state.identity.length {
257            if self.state.pending_start != self.state.processed_bytes {
258                self.finish_line()?;
259                self.state.pending_start = self.state.processed_bytes;
260            }
261            self.state.finished = true;
262        }
263        self.state.identity.check(&self.file)?;
264        Ok(self.progress())
265    }
266
267    /// Consume a completed index and retain the already-open source handle.
268    pub fn into_dataset<I: DeserializeOwned>(self) -> io::Result<IndexedJsonlDataset<I>> {
269        self.state.validate()?;
270        if !self.state.finished { return Err(invalid("JSONL indexing has not reached EOF")); }
271        self.state.identity.check(&self.file)?;
272        Ok(IndexedJsonlDataset {
273            file: Arc::new(Mutex::new(self.file)), state: Arc::new(self.state), item: PhantomData,
274        })
275    }
276}
277
278/// Immutable indexed JSONL, with deserialization performed only on requested rows.
279///
280/// Clones share the source/index, not materialized examples. Reads serialize on
281/// the source handle; deserialization occurs outside that lock. Storage remains
282/// proportional to the offset index plus concurrently requested row buffers.
283pub struct IndexedJsonlDataset<I> {
284    file: Arc<Mutex<File>>,
285    state: Arc<JsonlIndexState>,
286    item: PhantomData<fn() -> I>,
287}
288
289impl<I> Clone for IndexedJsonlDataset<I> {
290    fn clone(&self) -> Self {
291        Self { file: self.file.clone(), state: self.state.clone(), item: PhantomData }
292    }
293}
294
295impl<I: DeserializeOwned> IndexedJsonlDataset<I> {
296    /// Reopen a fully indexed source without rebuilding or copying its contents.
297    pub fn from_index(path: impl AsRef<Path>, source_id: &str, state: JsonlIndexState) -> io::Result<Self> {
298        JsonlIndexBuilder::resume(path, source_id, state)?.into_dataset()
299    }
300
301    /// Completed immutable byte-range index, suitable for caller persistence.
302    pub fn index(&self) -> &JsonlIndexState { &self.state }
303
304    /// Return the actual row bytes, retaining CR and excluding LF.
305    /// None means only that the dataset index is out of bounds.
306    pub fn get_raw(&self, index: usize) -> io::Result<Option<Vec<u8>>> {
307        let Some(row) = self.state.records.get(index) else { return Ok(None); };
308        let mut file = self.file.lock().map_err(|_| io::Error::other("JSONL source lock poisoned"))?;
309        self.state.identity.check(&file)?;
310        let bytes = Self::read_row(&mut file, row)?;
311        self.state.identity.check(&file)?;
312        Ok(Some(bytes))
313    }
314
315    fn read_row(file: &mut File, row: &JsonlRecordLocation) -> io::Result<Vec<u8>> {
316        let length = usize::try_from(row.length)
317            .map_err(|_| invalid("JSONL row is larger than the process address space"))?;
318        let mut bytes = Vec::new();
319        bytes.try_reserve_exact(length).map_err(|error| io::Error::other(error.to_string()))?;
320        bytes.resize(length, 0);
321        file.seek(SeekFrom::Start(row.offset))?;
322        file.read_exact(&mut bytes)?;
323        Ok(bytes)
324    }
325
326    /// Deserialize one actual row, distinguishing missing index from parse/I/O failure.
327    pub fn get_result(&self, index: usize) -> io::Result<Option<I>> {
328        let Some(bytes) = self.get_raw(index)? else { return Ok(None); };
329        serde_json::from_slice(&bytes).map(Some).map_err(|error| invalid(format!(
330            "JSONL dataset index {index}, physical line {}: {error}",
331            self.state.records[index].line_number)))
332    }
333
334    /// Fetch requested actual rows under one source lock and metadata check pair.
335    /// Requested order and duplicate indices are preserved; any missing index
336    /// returns None before reading. An empty request returns an empty collection.
337    pub fn get_many_raw(&self, indices: &[usize]) -> io::Result<Option<Vec<Vec<u8>>>> {
338        if indices.iter().any(|&index| index >= self.state.records.len()) { return Ok(None); }
339        if indices.is_empty() { return Ok(Some(Vec::new())); }
340        let mut file = self.file.lock().map_err(|_| io::Error::other("JSONL source lock poisoned"))?;
341        self.state.identity.check(&file)?;
342        let rows = indices.iter().map(|&index| Self::read_row(&mut file, &self.state.records[index]))
343            .collect::<io::Result<Vec<_>>>()?;
344        self.state.identity.check(&file)?;
345        Ok(Some(rows))
346    }
347
348    /// Deserialize a requested batch after releasing the shared source lock.
349    /// Parse/I/O failure is distinct from an out-of-bounds request, and no
350    /// partially loaded collection is returned as a successful batch.
351    pub fn get_many_result(&self, indices: &[usize]) -> io::Result<Option<Vec<I>>> {
352        let Some(rows) = self.get_many_raw(indices)? else { return Ok(None); };
353        indices.iter().zip(rows).map(|(&index, bytes)| {
354            serde_json::from_slice(&bytes).map_err(|error| invalid(format!(
355                "JSONL dataset index {index}, physical line {}: {error}",
356                self.state.records[index].line_number)))
357        }).collect::<io::Result<Vec<_>>>().map(Some)
358    }
359}
360
361impl<I: DeserializeOwned> Dataset<I> for IndexedJsonlDataset<I> {
362    /// The collection interface cannot carry errors: malformed/I/O-failed rows
363    /// panic rather than masquerading as absent data. Use get_result to handle them.
364    fn get(&self, index: usize) -> Option<I> {
365        self.get_result(index).unwrap_or_else(|error| panic!("indexed JSONL read failed: {error}"))
366    }
367
368    fn get_many(&self, indices: &[usize]) -> Option<Vec<I>> {
369        self.get_many_result(indices).unwrap_or_else(|error| panic!("indexed JSONL batch read failed: {error}"))
370    }
371
372    fn len(&self) -> usize { self.state.records.len() }
373}
ruda_dataset/dataset/indexed_jsonl.rs

ruda_dataset/dataset/
indexed_jsonl.rs
