burn_dataset/dataset/
base.rs

1use std::error::Error;
2use std::sync::Arc;
3
4use crate::{DatasetError, DatasetIterator};
5
6/// The dataset trait defines a basic collection of items with a predefined size.
7///
8/// # Panics
9///
10/// `get` panics if `index >= len()`, matching slice/`Vec` indexing conventions. `Err(_)` is
11/// reserved for genuine retrieval failures on an in-bounds index (e.g. an I/O or deserialization
12/// failure).
13pub trait Dataset<I, E = DatasetError>: Send + Sync
14where
15    E: Error + Send + Sync + 'static,
16{
17    /// Gets the item at the given index.
18    ///
19    /// # Panics
20    ///
21    /// Panics if `index >= len()`.
22    fn get(&self, index: usize) -> Result<I, E>;
23
24    /// Gets the items at given indexes
25    ///
26    /// # Panics
27    ///
28    /// panics if `indexes[i] >= len()`
29    fn get_many(&self, indexes: Vec<usize>) -> Result<Vec<I>, E> {
30        let len = self.len();
31        let mut items = Vec::new();
32
33        for i in indexes {
34            assert!(i < len, "Index out of bounds for dataset: {i} >= {len}");
35            let item = self.get(i)?;
36            items.push(item);
37        }
38
39        Ok(items)
40    }
41
42    /// Gets the number of items in the dataset.
43    fn len(&self) -> usize;
44
45    /// Checks if the dataset is empty.
46    fn is_empty(&self) -> bool {
47        self.len() == 0
48    }
49
50    /// Returns an iterator over the dataset.
51    fn iter(&self) -> DatasetIterator<'_, I, E>
52    where
53        Self: Sized,
54    {
55        DatasetIterator::new(self)
56    }
57}
58
59impl<D, I, E> Dataset<I, E> for Arc<D>
60where
61    D: Dataset<I, E>,
62    E: Error + Send + Sync + 'static,
63{
64    fn get(&self, index: usize) -> Result<I, E> {
65        self.as_ref().get(index)
66    }
67
68    fn get_many(&self, indexes: Vec<usize>) -> Result<Vec<I>, E> {
69        self.as_ref().get_many(indexes)
70    }
71
72    fn len(&self) -> usize {
73        self.as_ref().len()
74    }
75}
76
77impl<I, E> Dataset<I, E> for Arc<dyn Dataset<I, E>>
78where
79    E: Error + Send + Sync + 'static,
80{
81    fn get(&self, index: usize) -> Result<I, E> {
82        self.as_ref().get(index)
83    }
84
85    fn get_many(&self, indexes: Vec<usize>) -> Result<Vec<I>, E> {
86        self.as_ref().get_many(indexes)
87    }
88
89    fn len(&self) -> usize {
90        self.as_ref().len()
91    }
92}
93
94impl<D, I, E> Dataset<I, E> for Box<D>
95where
96    D: Dataset<I, E>,
97    E: Error + Send + Sync + 'static,
98{
99    fn get(&self, index: usize) -> Result<I, E> {
100        self.as_ref().get(index)
101    }
102
103    fn get_many(&self, indexes: Vec<usize>) -> Result<Vec<I>, E> {
104        self.as_ref().get_many(indexes)
105    }
106
107    fn len(&self) -> usize {
108        self.as_ref().len()
109    }
110}
111
112impl<I, E> Dataset<I, E> for Box<dyn Dataset<I, E>>
113where
114    E: Error + Send + Sync + 'static,
115{
116    fn get(&self, index: usize) -> Result<I, E> {
117        self.as_ref().get(index)
118    }
119
120    fn get_many(&self, indexes: Vec<usize>) -> Result<Vec<I>, E> {
121        self.as_ref().get_many(indexes)
122    }
123
124    fn len(&self) -> usize {
125        self.as_ref().len()
126    }
127}
burn_dataset/dataset/base.rs

burn_dataset/dataset/
base.rs
