dataloader_rs/loader/
iter.rs

1//! [`DataLoaderIter`]: one epoch of a loader, whatever its source.
2
3#[cfg(feature = "python")]
4use std::time::{Duration, Instant};
5
6use crate::{
7    collator::Collator,
8    error::{Error, Result},
9    loader::source::LoaderSource,
10    transform::{Identity, Transform},
11};
12
13/// The error of an epoch used in a process `fork()`ed after it started: the
14/// pool's threads do not exist there, so its batches would never arrive.
15pub(super) fn forked() -> Error {
16    "this epoch was started before fork(); start a new one in this process (iter())".into()
17}
18
19/// What an epoch of a loader yields its batches through: [`SyncEpoch`] for a
20/// [`Dataset`], `AsyncEpoch` for an `AsyncDataset`.
21pub trait BatchStream<B> {
22    /// The next batch in epoch order, `None` at the end of the epoch.
23    fn next(&mut self) -> Option<Result<B>>;
24
25    /// Batches left in the epoch.
26    fn len(&self) -> usize;
27
28    /// The next batch if it arrives by `deadline`, or (`None`) if it has
29    /// already arrived; `None` if not, so that the caller can wait again.
30    #[cfg(feature = "python")]
31    fn next_until(&mut self, deadline: Option<Instant>) -> Option<Option<Result<B>>>;
32
33    /// End the epoch and wait up to `timeout` for its background work to
34    /// exit; `true` if it did, so that nothing it held is left.
35    #[cfg(feature = "python")]
36    fn stop(&mut self, timeout: Duration) -> bool;
37}
38
39// ── DataLoaderIter ────────────────────────────────────────────────────────────
40
41/// Iterator over one epoch of a [`DataLoader`](crate::DataLoader).
42///
43/// Owns shared handles to the dataset, transform and collator, so it does not
44/// borrow the loader and can be stored, sent to another thread, or outlive it.
45pub struct DataLoaderIter<Src, C, T = Identity>
46where
47    Src: LoaderSource,
48    T: Transform<Src::Item>,
49    C: Collator<T::Output>,
50{
51    epoch: Src::Epoch<T, C, T::Output>,
52}
53
54impl<Src, C, T> DataLoaderIter<Src, C, T>
55where
56    Src: LoaderSource,
57    T: Transform<Src::Item>,
58    C: Collator<T::Output>,
59{
60    pub(super) fn new(epoch: Src::Epoch<T, C, T::Output>) -> Self {
61        Self { epoch }
62    }
63
64    /// The next item if it is ready by `deadline`, or (`None`) right away;
65    /// `None` if not. Lets the Python binding keep the interpreter when a
66    /// batch is already there, and wait in slices that signals can end.
67    #[cfg(feature = "python")]
68    pub(crate) fn next_until(
69        &mut self,
70        deadline: Option<Instant>,
71    ) -> Option<Option<Result<C::Batch>>> {
72        self.epoch.next_until(deadline)
73    }
74
75    /// End the epoch and wait up to `timeout` for the tasks in flight, so that
76    /// the Python objects they hold are released on the caller's thread.
77    /// Called with the interpreter released: the tasks may need it to finish.
78    #[cfg(feature = "python")]
79    pub(crate) fn stop(&mut self, timeout: Duration) -> bool {
80        self.epoch.stop(timeout)
81    }
82}
83
84impl<Src, C, T> Iterator for DataLoaderIter<Src, C, T>
85where
86    Src: LoaderSource,
87    T: Transform<Src::Item>,
88    C: Collator<T::Output>,
89{
90    type Item = Result<C::Batch>;
91
92    fn next(&mut self) -> Option<Self::Item> {
93        self.epoch.next()
94    }
95
96    fn size_hint(&self) -> (usize, Option<usize>) {
97        let n = self.epoch.len();
98        (n, Some(n))
99    }
100}
101
102impl<Src, C, T> ExactSizeIterator for DataLoaderIter<Src, C, T>
103where
104    Src: LoaderSource,
105    T: Transform<Src::Item>,
106    C: Collator<T::Output>,
107{
108}
109
110// ── IndexedIter ───────────────────────────────────────────────────────────────
111
112/// Iterator over one epoch of a [`DataLoader`](crate::DataLoader) that pairs
113/// each batch with the dataset indices of its items; see
114/// [`DataLoader::iter_with_indices`](crate::DataLoader::iter_with_indices).
115pub struct IndexedIter<Src, C, T = Identity>
116where
117    Src: LoaderSource,
118    T: Transform<Src::Item>,
119    C: Collator<T::Output>,
120{
121    indices: std::vec::IntoIter<Vec<usize>>,
122    batches: DataLoaderIter<Src, C, T>,
123}
124
125impl<Src, C, T> IndexedIter<Src, C, T>
126where
127    Src: LoaderSource,
128    T: Transform<Src::Item>,
129    C: Collator<T::Output>,
130{
131    /// `batches`, whose epoch runs over `indices`.
132    pub(super) fn new(indices: Vec<Vec<usize>>, batches: DataLoaderIter<Src, C, T>) -> Self {
133        Self {
134            indices: indices.into_iter(),
135            batches,
136        }
137    }
138}
139
140impl<Src, C, T> Iterator for IndexedIter<Src, C, T>
141where
142    Src: LoaderSource,
143    T: Transform<Src::Item>,
144    C: Collator<T::Output>,
145{
146    type Item = Result<(Vec<usize>, C::Batch)>;
147
148    fn next(&mut self) -> Option<Self::Item> {
149        // Every path yields its batches in epoch order, a failed one included, so the
150        // i-th batch is the i-th chunk of indices.
151        let batch = self.batches.next()?;
152        let indices = self.indices.next()?;
153        Some(batch.map(|batch| (indices, batch)))
154    }
155
156    fn size_hint(&self) -> (usize, Option<usize>) {
157        self.batches.size_hint()
158    }
159}
160
161impl<Src, C, T> ExactSizeIterator for IndexedIter<Src, C, T>
162where
163    Src: LoaderSource,
164    T: Transform<Src::Item>,
165    C: Collator<T::Output>,
166{
167}
dataloader_rs/loader/iter.rs

dataloader_rs/loader/
iter.rs
