Skip to main content

scirs2_io/pipeline/
sources.rs

1//! Data source abstractions for the streaming pipeline.
2//!
3//! Every source implements the [`DataSource`] trait which yields batches of
4//! `serde_json::Value` records.  Concrete implementations:
5//!
6//! | Source            | Description                                     |
7//! |-------------------|-------------------------------------------------|
8//! | [`FileSource`]    | Read CSV / JSON / JSONL files in batches        |
9//! | [`DatabaseSource`]| SQLite query-based loading                      |
10//! | [`StreamSource`]  | Kafka-like channel-based message consumption    |
11//! | [`GeneratorSource`]| Produce data from a user closure               |
12//!
13//! `DatabaseSource` uses `oxisql-sqlite-compat` (Pure Rust SQLite via Limbo engine)
14//! with a sync/async bridge via `tokio::task::block_in_place` / a fresh `Runtime`.
15
16#![allow(missing_docs)]
17
18use crate::error::{IoError, Result};
19use serde_json::Value;
20use std::fs::File;
21use std::io::{BufRead, BufReader};
22use std::path::{Path, PathBuf};
23use std::sync::{Arc, Mutex};
24
25// ─────────────────────────────────────────────────────────────────────────────
26// Core trait
27// ─────────────────────────────────────────────────────────────────────────────
28
29/// A pull-based data source that yields record batches.
30pub trait DataSource: Send {
31    /// Return the next batch of records, or `Ok(None)` when exhausted.
32    fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>>;
33
34    /// Reset the source to the beginning, if supported.
35    fn reset(&mut self) -> Result<()> {
36        Err(IoError::Other(
37            "reset not supported by this source".to_string(),
38        ))
39    }
40
41    /// Human-readable name for diagnostics.
42    fn name(&self) -> &str {
43        "unknown"
44    }
45
46    /// Estimated total number of records, if knowable.
47    fn estimated_len(&self) -> Option<usize> {
48        None
49    }
50}
51
52// ─────────────────────────────────────────────────────────────────────────────
53// FileSource
54// ─────────────────────────────────────────────────────────────────────────────
55
56/// Supported file formats for [`FileSource`].
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum FileSourceFormat {
59    /// Comma-separated values
60    Csv,
61    /// Newline-delimited JSON (one JSON object per line)
62    JsonLines,
63    /// A single JSON array `[…]` or object
64    Json,
65    /// Auto-detect from file extension
66    Auto,
67}
68
69/// Read structured data from a CSV / JSON / JSONL file in batches.
70pub struct FileSource {
71    path: PathBuf,
72    format: FileSourceFormat,
73    reader: Option<FileSourceReader>,
74    exhausted: bool,
75    /// CSV header row (populated lazily on first read).
76    csv_headers: Option<Vec<String>>,
77    /// Pre-loaded JSON records (for non-streaming JSON).
78    json_records: Option<std::collections::VecDeque<Value>>,
79}
80
81enum FileSourceReader {
82    Lines(BufReader<File>),
83}
84
85impl FileSource {
86    /// Create a new `FileSource` with automatic format detection.
87    pub fn new(path: impl AsRef<Path>) -> Self {
88        Self::with_format(path, FileSourceFormat::Auto)
89    }
90
91    /// Create a new `FileSource` with an explicit format.
92    pub fn with_format(path: impl AsRef<Path>, format: FileSourceFormat) -> Self {
93        Self {
94            path: path.as_ref().to_path_buf(),
95            format,
96            reader: None,
97            exhausted: false,
98            csv_headers: None,
99            json_records: None,
100        }
101    }
102
103    fn resolve_format(&self) -> FileSourceFormat {
104        if self.format != FileSourceFormat::Auto {
105            return self.format;
106        }
107        match self
108            .path
109            .extension()
110            .and_then(|e| e.to_str())
111            .map(|e| e.to_ascii_lowercase())
112            .as_deref()
113        {
114            Some("csv") => FileSourceFormat::Csv,
115            Some("jsonl") | Some("ndjson") => FileSourceFormat::JsonLines,
116            Some("json") => FileSourceFormat::Json,
117            _ => FileSourceFormat::JsonLines,
118        }
119    }
120
121    fn open(&mut self) -> Result<()> {
122        let fmt = self.resolve_format();
123        match fmt {
124            FileSourceFormat::Json => {
125                // Load the entire JSON file up front.
126                let file = File::open(&self.path).map_err(IoError::Io)?;
127                let value: Value = serde_json::from_reader(BufReader::new(file))
128                    .map_err(|e| IoError::ParseError(e.to_string()))?;
129                let records: Vec<Value> = match value {
130                    Value::Array(arr) => arr,
131                    other => vec![other],
132                };
133                self.json_records = Some(records.into());
134            }
135            _ => {
136                // Line-oriented formats (CSV, JSONL).
137                let file = File::open(&self.path).map_err(IoError::Io)?;
138                let reader = BufReader::new(file);
139                self.reader = Some(FileSourceReader::Lines(reader));
140            }
141        }
142        Ok(())
143    }
144
145    fn next_batch_csv(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
146        let reader = match &mut self.reader {
147            Some(FileSourceReader::Lines(r)) => r,
148            None => return Ok(None),
149        };
150
151        // Read header on first call.
152        if self.csv_headers.is_none() {
153            let mut line = String::new();
154            if reader.read_line(&mut line).map_err(IoError::Io)? == 0 {
155                return Ok(None);
156            }
157            let headers: Vec<String> = line
158                .trim_end_matches(['\n', '\r'])
159                .split(',')
160                .map(|h| h.trim().trim_matches('"').to_string())
161                .collect();
162            self.csv_headers = Some(headers);
163        }
164
165        let headers = self.csv_headers.clone().unwrap_or_default();
166        let mut records = Vec::with_capacity(batch_size);
167        let mut line = String::new();
168
169        while records.len() < batch_size {
170            line.clear();
171            let n = reader.read_line(&mut line).map_err(IoError::Io)?;
172            if n == 0 {
173                break;
174            }
175            let row = line.trim_end_matches(['\n', '\r']);
176            if row.is_empty() {
177                continue;
178            }
179            let fields: Vec<&str> = row.split(',').collect();
180            let obj: serde_json::Map<String, Value> = headers
181                .iter()
182                .enumerate()
183                .map(|(i, h)| {
184                    let raw = fields.get(i).copied().unwrap_or("").trim();
185                    let val = raw.trim_matches('"');
186                    // Try to parse as number, then bool, then leave as string.
187                    let json_val: Value = if let Ok(n) = val.parse::<i64>() {
188                        Value::Number(n.into())
189                    } else if let Ok(f) = val.parse::<f64>() {
190                        Value::Number(
191                            serde_json::Number::from_f64(f).unwrap_or(serde_json::Number::from(0)),
192                        )
193                    } else if val == "true" {
194                        Value::Bool(true)
195                    } else if val == "false" {
196                        Value::Bool(false)
197                    } else {
198                        Value::String(val.to_string())
199                    };
200                    (h.clone(), json_val)
201                })
202                .collect();
203            records.push(Value::Object(obj));
204        }
205
206        if records.is_empty() {
207            Ok(None)
208        } else {
209            Ok(Some(records))
210        }
211    }
212
213    fn next_batch_jsonl(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
214        let reader = match &mut self.reader {
215            Some(FileSourceReader::Lines(r)) => r,
216            None => return Ok(None),
217        };
218
219        let mut records = Vec::with_capacity(batch_size);
220        let mut line = String::new();
221
222        while records.len() < batch_size {
223            line.clear();
224            let n = reader.read_line(&mut line).map_err(IoError::Io)?;
225            if n == 0 {
226                break;
227            }
228            let trimmed = line.trim();
229            if trimmed.is_empty() {
230                continue;
231            }
232            let val: Value = serde_json::from_str(trimmed)
233                .map_err(|e| IoError::ParseError(format!("JSONL parse: {e}")))?;
234            records.push(val);
235        }
236
237        if records.is_empty() {
238            Ok(None)
239        } else {
240            Ok(Some(records))
241        }
242    }
243}
244
245impl DataSource for FileSource {
246    fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
247        if self.exhausted {
248            return Ok(None);
249        }
250
251        if self.reader.is_none() && self.json_records.is_none() {
252            self.open()?;
253        }
254
255        let result = if let Some(queue) = &mut self.json_records {
256            // Pre-loaded JSON array.
257            if queue.is_empty() {
258                Ok(None)
259            } else {
260                let batch: Vec<Value> = queue.drain(..batch_size.min(queue.len())).collect();
261                Ok(Some(batch))
262            }
263        } else {
264            match self.resolve_format() {
265                FileSourceFormat::Csv => self.next_batch_csv(batch_size),
266                _ => self.next_batch_jsonl(batch_size),
267            }
268        };
269
270        if matches!(result, Ok(None)) {
271            self.exhausted = true;
272        }
273        result
274    }
275
276    fn reset(&mut self) -> Result<()> {
277        self.reader = None;
278        self.json_records = None;
279        self.exhausted = false;
280        self.csv_headers = None;
281        Ok(())
282    }
283
284    fn name(&self) -> &str {
285        "file_source"
286    }
287}
288
289// ─────────────────────────────────────────────────────────────────────────────
290// DatabaseSource  (SQLite-backed)
291// ─────────────────────────────────────────────────────────────────────────────
292
293/// Query-based data source backed by an in-process SQLite database.
294///
295/// Requires the `sqlite` feature.  When the feature is absent the struct still
296/// compiles but `next_batch` always returns an "unsupported" error.
297pub struct DatabaseSource {
298    /// SQLite connection string (path or `:memory:`).
299    connection_string: String,
300    /// SQL query to execute.
301    query: String,
302    /// Pre-fetched rows (populated lazily on first read).
303    rows: Option<std::collections::VecDeque<Value>>,
304    exhausted: bool,
305}
306
307impl DatabaseSource {
308    /// Create a new `DatabaseSource`.
309    ///
310    /// - `connection_string`: path to the `.sqlite` / `.db` file, or `:memory:`.
311    /// - `query`: full SQL `SELECT` statement.
312    pub fn new(connection_string: impl Into<String>, query: impl Into<String>) -> Self {
313        Self {
314            connection_string: connection_string.into(),
315            query: query.into(),
316            rows: None,
317            exhausted: false,
318        }
319    }
320
321    #[cfg(feature = "sqlite")]
322    fn load_rows(&mut self) -> Result<()> {
323        use oxisql_core::{Connection as OxiConnection, Value as OxiValue};
324        use oxisql_sqlite_compat::SqliteConnection;
325        use std::future::Future;
326
327        /// Run an async future to completion, bridging from sync context.
328        fn run_sync_local<F, T, E>(fut: F) -> std::result::Result<T, E>
329        where
330            F: Future<Output = std::result::Result<T, E>>,
331        {
332            match tokio::runtime::Handle::try_current() {
333                Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fut)),
334                Err(_) => tokio::runtime::Builder::new_current_thread()
335                    .enable_all()
336                    .build()
337                    .expect("tokio runtime for sqlite source")
338                    .block_on(fut),
339            }
340        }
341
342        let conn = run_sync_local(SqliteConnection::open(&self.connection_string))
343            .map_err(|e| IoError::DatabaseError(e.to_string()))?;
344
345        let oxi_rows = run_sync_local(conn.query(&self.query, &[]))
346            .map_err(|e| IoError::DatabaseError(e.to_string()))?;
347
348        let mut rows_vec = Vec::with_capacity(oxi_rows.len());
349
350        for row in &oxi_rows {
351            let mut obj = serde_json::Map::new();
352            let col_count = row.column_count();
353            let col_names = row.columns().to_vec();
354
355            for i in 0..col_count {
356                let col_name = col_names
357                    .get(i)
358                    .cloned()
359                    .unwrap_or_else(|| format!("column_{i}"));
360                let json_val = match row.get_by_index(i) {
361                    Some(OxiValue::Null) | None => Value::Null,
362                    Some(OxiValue::Bool(b)) => Value::Bool(*b),
363                    Some(OxiValue::I64(n)) => Value::Number((*n).into()),
364                    Some(OxiValue::F64(f)) => Value::Number(
365                        serde_json::Number::from_f64(*f).unwrap_or(serde_json::Number::from(0)),
366                    ),
367                    Some(OxiValue::Text(t)) => Value::String(t.clone()),
368                    Some(OxiValue::Blob(b)) => Value::String(format!("<blob {} bytes>", b.len())),
369                    Some(OxiValue::Json(j)) => {
370                        serde_json::from_str(j).unwrap_or_else(|_| Value::String(j.clone()))
371                    }
372                    Some(OxiValue::Decimal(d)) => Value::String(d.clone()),
373                    Some(OxiValue::Timestamp(ts)) => Value::Number((*ts).into()),
374                    Some(OxiValue::Date(d)) => Value::Number((*d).into()),
375                    Some(OxiValue::Time(t)) => Value::Number((*t).into()),
376                    Some(OxiValue::Uuid(_)) => {
377                        // Format UUID using Display impl on OxiValue.
378                        Value::String(format!(
379                            "{}",
380                            row.get_by_index(i)
381                                .expect("index checked above in outer match")
382                        ))
383                    }
384                    Some(OxiValue::Array(_)) => Value::String(format!("{:?}", row.get_by_index(i))),
385                    // Typed arrays carry the same element data as a plain array
386                    // plus a nominal element type; represent them the same way.
387                    Some(OxiValue::TypedArray { .. }) => {
388                        Value::String(format!("{:?}", row.get_by_index(i)))
389                    }
390                };
391                obj.insert(col_name, json_val);
392            }
393            rows_vec.push(Value::Object(obj));
394        }
395
396        self.rows = Some(rows_vec.into());
397        Ok(())
398    }
399
400    #[cfg(not(feature = "sqlite"))]
401    fn load_rows(&mut self) -> Result<()> {
402        Err(IoError::Other(
403            "DatabaseSource requires the `sqlite` feature".to_string(),
404        ))
405    }
406}
407
408impl DataSource for DatabaseSource {
409    fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
410        if self.exhausted {
411            return Ok(None);
412        }
413        if self.rows.is_none() {
414            self.load_rows()?;
415        }
416        let queue = self
417            .rows
418            .as_mut()
419            .expect("rows must be populated after load_rows");
420        if queue.is_empty() {
421            self.exhausted = true;
422            return Ok(None);
423        }
424        let batch: Vec<Value> = queue.drain(..batch_size.min(queue.len())).collect();
425        Ok(Some(batch))
426    }
427
428    fn reset(&mut self) -> Result<()> {
429        self.rows = None;
430        self.exhausted = false;
431        Ok(())
432    }
433
434    fn name(&self) -> &str {
435        "database_source"
436    }
437}
438
439// ─────────────────────────────────────────────────────────────────────────────
440// StreamSource  (Kafka-like channel-based source)
441// ─────────────────────────────────────────────────────────────────────────────
442
443/// A streaming, channel-based message source that mimics a Kafka consumer.
444///
445/// Messages are enqueued via `StreamSource::sender` from any thread and
446/// consumed in batches via [`DataSource::next_batch`].
447pub struct StreamSource {
448    receiver: crossbeam_channel::Receiver<Value>,
449    timeout: std::time::Duration,
450    label: String,
451}
452
453/// Handle for sending messages into a [`StreamSource`].
454pub type StreamSender = crossbeam_channel::Sender<Value>;
455
456impl StreamSource {
457    /// Create an unbounded channel-backed source.
458    ///
459    /// Returns `(source, sender)`.  Messages sent via `sender` become
460    /// available through `source.next_batch`.
461    pub fn new_unbounded() -> (Self, StreamSender) {
462        let (tx, rx) = crossbeam_channel::unbounded();
463        let source = Self {
464            receiver: rx,
465            timeout: std::time::Duration::from_millis(50),
466            label: "stream_source".to_string(),
467        };
468        (source, tx)
469    }
470
471    /// Create a bounded channel-backed source.
472    ///
473    /// Senders will block once the channel is full (backpressure).
474    pub fn new_bounded(capacity: usize) -> (Self, StreamSender) {
475        let (tx, rx) = crossbeam_channel::bounded(capacity);
476        let source = Self {
477            receiver: rx,
478            timeout: std::time::Duration::from_millis(50),
479            label: "stream_source_bounded".to_string(),
480        };
481        (source, tx)
482    }
483
484    /// Set the receive timeout used when the channel is empty.
485    pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
486        self.timeout = timeout;
487        self
488    }
489
490    /// Set a human-readable label.
491    pub fn with_name(mut self, name: impl Into<String>) -> Self {
492        self.label = name.into();
493        self
494    }
495}
496
497impl DataSource for StreamSource {
498    /// Collect up to `batch_size` messages.  Returns `None` only when the
499    /// channel is both empty **and** all senders have been dropped.
500    fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
501        use crossbeam_channel::RecvTimeoutError;
502
503        let mut batch = Vec::with_capacity(batch_size);
504
505        // Block for the first message (or detect channel closure).
506        match self.receiver.recv_timeout(self.timeout) {
507            Ok(msg) => batch.push(msg),
508            Err(RecvTimeoutError::Disconnected) => return Ok(None),
509            Err(RecvTimeoutError::Timeout) => return Ok(Some(batch)), // empty batch OK
510        }
511
512        // Drain non-blockingly up to batch_size - 1 more items.
513        while batch.len() < batch_size {
514            match self.receiver.try_recv() {
515                Ok(msg) => batch.push(msg),
516                Err(_) => break,
517            }
518        }
519
520        Ok(Some(batch))
521    }
522
523    fn name(&self) -> &str {
524        &self.label
525    }
526}
527
528// ─────────────────────────────────────────────────────────────────────────────
529// GeneratorSource
530// ─────────────────────────────────────────────────────────────────────────────
531
532/// Produce data from a user-supplied closure.
533///
534/// The generator closure is called repeatedly; when it returns `None` the
535/// source is exhausted.
536pub struct GeneratorSource<F>
537where
538    F: FnMut() -> Option<Value> + Send,
539{
540    generator: F,
541    label: String,
542    exhausted: bool,
543}
544
545impl<F> GeneratorSource<F>
546where
547    F: FnMut() -> Option<Value> + Send,
548{
549    /// Create a new `GeneratorSource` with the given closure.
550    pub fn new(generator: F) -> Self {
551        Self {
552            generator,
553            label: "generator_source".to_string(),
554            exhausted: false,
555        }
556    }
557
558    /// Attach a human-readable label.
559    pub fn with_name(mut self, name: impl Into<String>) -> Self {
560        self.label = name.into();
561        self
562    }
563}
564
565impl<F> DataSource for GeneratorSource<F>
566where
567    F: FnMut() -> Option<Value> + Send,
568{
569    fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
570        if self.exhausted {
571            return Ok(None);
572        }
573        let mut batch = Vec::with_capacity(batch_size);
574        while batch.len() < batch_size {
575            match (self.generator)() {
576                Some(v) => batch.push(v),
577                None => {
578                    self.exhausted = true;
579                    break;
580                }
581            }
582        }
583        if batch.is_empty() {
584            Ok(None)
585        } else {
586            Ok(Some(batch))
587        }
588    }
589
590    fn reset(&mut self) -> Result<()> {
591        // Cannot reset a closure-based generator by default.
592        Err(IoError::Other(
593            "GeneratorSource does not support reset".to_string(),
594        ))
595    }
596
597    fn name(&self) -> &str {
598        &self.label
599    }
600}
601
602// ─────────────────────────────────────────────────────────────────────────────
603// SourceChain – iterate a source to completion
604// ─────────────────────────────────────────────────────────────────────────────
605
606/// Convenience: drain an entire [`DataSource`] into a `Vec<Value>`.
607pub fn drain_source(source: &mut dyn DataSource, batch_size: usize) -> Result<Vec<Value>> {
608    let mut all = Vec::new();
609    while let Some(batch) = source.next_batch(batch_size)? {
610        all.extend(batch);
611    }
612    Ok(all)
613}
614
615// ─────────────────────────────────────────────────────────────────────────────
616// Tests
617// ─────────────────────────────────────────────────────────────────────────────
618
619#[cfg(test)]
620mod tests {
621    use super::*;
622    use serde_json::json;
623    use std::io::Write;
624
625    fn temp_path(suffix: &str) -> PathBuf {
626        let mut p = std::env::temp_dir();
627        p.push(format!("scirs2_io_test_{suffix}_{}", std::process::id()));
628        p
629    }
630
631    #[test]
632    fn test_generator_source() {
633        let mut counter = 0u64;
634        let mut src = GeneratorSource::new(move || {
635            if counter < 5 {
636                let v = json!(counter);
637                counter += 1;
638                Some(v)
639            } else {
640                None
641            }
642        });
643
644        let batch = src.next_batch(3).unwrap().unwrap();
645        assert_eq!(batch.len(), 3);
646        let rest = drain_source(&mut src, 10).unwrap();
647        assert_eq!(rest.len(), 2);
648        assert!(src.next_batch(10).unwrap().is_none());
649    }
650
651    #[test]
652    fn test_stream_source_basic() {
653        let (mut src, tx) = StreamSource::new_unbounded();
654        tx.send(json!({"id": 1})).unwrap();
655        tx.send(json!({"id": 2})).unwrap();
656        drop(tx); // signal EOF
657
658        let all = drain_source(&mut src, 10).unwrap();
659        // We expect at least 2 records (timing-dependent for the empty-batch path)
660        let total: usize = all.len() + drain_source(&mut src, 10).unwrap_or_default().len();
661        assert!(total >= 2 || all.len() >= 2);
662    }
663
664    #[test]
665    fn test_file_source_jsonl() {
666        let path = temp_path("jsonl");
667        {
668            let mut f = std::fs::File::create(&path).unwrap();
669            writeln!(f, r#"{{"a":1,"b":"x"}}"#).unwrap();
670            writeln!(f, r#"{{"a":2,"b":"y"}}"#).unwrap();
671            writeln!(f, r#"{{"a":3,"b":"z"}}"#).unwrap();
672        }
673        let mut src = FileSource::with_format(&path, FileSourceFormat::JsonLines);
674        let all = drain_source(&mut src, 10).unwrap();
675        assert_eq!(all.len(), 3);
676        assert_eq!(all[0]["a"], json!(1));
677        std::fs::remove_file(&path).ok();
678    }
679
680    #[test]
681    fn test_file_source_csv() {
682        let path = temp_path("csv");
683        {
684            let mut f = std::fs::File::create(&path).unwrap();
685            writeln!(f, "name,score").unwrap();
686            writeln!(f, "alice,90").unwrap();
687            writeln!(f, "bob,85").unwrap();
688        }
689        let mut src = FileSource::with_format(&path, FileSourceFormat::Csv);
690        let all = drain_source(&mut src, 10).unwrap();
691        assert_eq!(all.len(), 2);
692        assert_eq!(all[0]["name"], json!("alice"));
693        assert_eq!(all[0]["score"], json!(90));
694        std::fs::remove_file(&path).ok();
695    }
696
697    #[test]
698    fn test_file_source_json_array() {
699        let path = temp_path("json");
700        {
701            let mut f = std::fs::File::create(&path).unwrap();
702            write!(f, r#"[{{"x":1}},{{"x":2}},{{"x":3}}]"#).unwrap();
703        }
704        let mut src = FileSource::with_format(&path, FileSourceFormat::Json);
705        // Read in two batches of 2.
706        let b1 = src.next_batch(2).unwrap().unwrap();
707        assert_eq!(b1.len(), 2);
708        let b2 = src.next_batch(2).unwrap().unwrap();
709        assert_eq!(b2.len(), 1);
710        assert!(src.next_batch(2).unwrap().is_none());
711        std::fs::remove_file(&path).ok();
712    }
713
714    #[test]
715    fn test_file_source_batching() {
716        let path = temp_path("jsonl2");
717        {
718            let mut f = std::fs::File::create(&path).unwrap();
719            for i in 0..10 {
720                writeln!(f, r#"{{"i":{i}}}"#).unwrap();
721            }
722        }
723        let mut src = FileSource::with_format(&path, FileSourceFormat::JsonLines);
724        let b1 = src.next_batch(4).unwrap().unwrap();
725        assert_eq!(b1.len(), 4);
726        let b2 = src.next_batch(4).unwrap().unwrap();
727        assert_eq!(b2.len(), 4);
728        let b3 = src.next_batch(4).unwrap().unwrap();
729        assert_eq!(b3.len(), 2);
730        assert!(src.next_batch(4).unwrap().is_none());
731        std::fs::remove_file(&path).ok();
732    }
733}