spire_core/dataset/
data.rs

1use std::fmt;
2use std::future::Future;
3use std::sync::Arc;
4
5use futures::future::BoxFuture;
6use futures::{FutureExt, Stream, stream};
7
8use crate::Result;
9use crate::dataset::Dataset;
10
11/// How many items [`Data::read_all`] takes at once.
12pub const READ_BATCH: usize = 1024;
13
14/// A cloneable, type-erased handle to a [`Dataset`].
15pub struct Data<T> {
16    inner: Arc<dyn DynDataset<T>>,
17}
18
19impl<T: Send + 'static> Data<T> {
20    /// Wraps `dataset`.
21    #[must_use]
22    pub fn new(dataset: impl Dataset<T>) -> Self {
23        Self {
24            inner: Arc::new(dataset),
25        }
26    }
27
28    /// Stores an item.
29    ///
30    /// # Errors
31    ///
32    /// Returns an error if the underlying dataset fails.
33    pub async fn write(&self, item: T) -> Result<()> {
34        self.inner.write(item).await
35    }
36
37    /// Stores every item of `items` at once, see [`Dataset::write_batch`].
38    ///
39    /// # Errors
40    ///
41    /// Returns an error if the underlying dataset fails.
42    pub async fn write_all(&self, items: impl IntoIterator<Item = T>) -> Result<()> {
43        let items: Vec<T> = items.into_iter().collect();
44        if items.is_empty() {
45            return Ok(());
46        }
47        self.inner.write_batch(items).await
48    }
49
50    /// Takes the next item, or `None` if the dataset is empty.
51    ///
52    /// # Errors
53    ///
54    /// Returns an error if the underlying dataset fails.
55    pub async fn read(&self) -> Result<Option<T>> {
56        self.inner.read().await
57    }
58
59    /// Takes up to `max` items at once, see [`Dataset::read_batch`].
60    ///
61    /// # Errors
62    ///
63    /// Returns an error if the underlying dataset fails.
64    pub async fn read_batch(&self, max: usize) -> Result<Vec<T>> {
65        self.inner.read_batch(max).await
66    }
67
68    /// Takes every item until the dataset is empty, in batches of
69    /// [`READ_BATCH`] items.
70    ///
71    /// # Errors
72    ///
73    /// Returns an error if the underlying dataset fails.
74    pub async fn read_all(&self) -> Result<Vec<T>> {
75        let mut items = Vec::with_capacity(self.size_hint().unwrap_or_default());
76        loop {
77            let batch = self.inner.read_batch(READ_BATCH).await?;
78            let last = batch.len() < READ_BATCH;
79            items.extend(batch);
80            if last {
81                return Ok(items);
82            }
83        }
84    }
85
86    /// Returns a stream that takes items one at a time until the dataset is
87    /// empty, so it never takes an item it does not yield. To drain a large
88    /// dataset, [`into_batches`](Self::into_batches) takes fewer trips to
89    /// its storage.
90    pub fn into_stream(self) -> impl Stream<Item = Result<T>> + Send {
91        stream::try_unfold(self, |data| async move {
92            let item = data.read().await?;
93            Ok(item.map(|item| (item, data)))
94        })
95    }
96
97    /// Returns a stream that takes batches of up to `max` items, see
98    /// [`Dataset::read_batch`], until the dataset is empty.
99    pub fn into_batches(self, max: usize) -> impl Stream<Item = Result<Vec<T>>> + Send {
100        let max = max.max(1);
101        stream::try_unfold((self, false), move |(data, drained)| async move {
102            if drained {
103                return Ok(None);
104            }
105            let batch = data.read_batch(max).await?;
106            if batch.is_empty() {
107                return Ok(None);
108            }
109            let drained = batch.len() < max;
110            Ok(Some((batch, (data, drained))))
111        })
112    }
113
114    /// Returns the number of stored items, if known.
115    #[must_use]
116    pub fn size_hint(&self) -> Option<usize> {
117        self.inner.size_hint()
118    }
119}
120
121impl<T> Clone for Data<T> {
122    fn clone(&self) -> Self {
123        Self {
124            inner: self.inner.clone(),
125        }
126    }
127}
128
129impl<T> fmt::Debug for Data<T> {
130    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
131        f.debug_struct("Data")
132            .field("type", &std::any::type_name::<T>())
133            .finish_non_exhaustive()
134    }
135}
136
137impl<T: Send + 'static> Dataset<T> for Data<T> {
138    fn write(&self, item: T) -> impl Future<Output = Result<()>> + Send {
139        self.inner.write(item)
140    }
141
142    fn write_batch(&self, items: Vec<T>) -> impl Future<Output = Result<()>> + Send {
143        self.inner.write_batch(items)
144    }
145
146    fn read(&self) -> impl Future<Output = Result<Option<T>>> + Send {
147        self.inner.read()
148    }
149
150    fn read_batch(&self, max: usize) -> impl Future<Output = Result<Vec<T>>> + Send {
151        self.inner.read_batch(max)
152    }
153
154    fn size_hint(&self) -> Option<usize> {
155        self.inner.size_hint()
156    }
157}
158
159/// An object-safe [`Dataset`].
160trait DynDataset<T>: Send + Sync + 'static {
161    fn write(&self, item: T) -> BoxFuture<'_, Result<()>>;
162    fn write_batch(&self, items: Vec<T>) -> BoxFuture<'_, Result<()>>;
163    fn read(&self) -> BoxFuture<'_, Result<Option<T>>>;
164    fn read_batch(&self, max: usize) -> BoxFuture<'_, Result<Vec<T>>>;
165    fn size_hint(&self) -> Option<usize>;
166}
167
168impl<T, D> DynDataset<T> for D
169where
170    T: Send + 'static,
171    D: Dataset<T>,
172{
173    fn write(&self, item: T) -> BoxFuture<'_, Result<()>> {
174        Dataset::write(self, item).boxed()
175    }
176
177    fn write_batch(&self, items: Vec<T>) -> BoxFuture<'_, Result<()>> {
178        Dataset::write_batch(self, items).boxed()
179    }
180
181    fn read(&self) -> BoxFuture<'_, Result<Option<T>>> {
182        Dataset::read(self).boxed()
183    }
184
185    fn read_batch(&self, max: usize) -> BoxFuture<'_, Result<Vec<T>>> {
186        Dataset::read_batch(self, max).boxed()
187    }
188
189    fn size_hint(&self) -> Option<usize> {
190        Dataset::size_hint(self)
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use futures::TryStreamExt;
197
198    use super::*;
199    use crate::dataset::InMemDataset;
200
201    #[tokio::test]
202    async fn batch_and_stream() {
203        let data = Data::new(InMemDataset::new());
204        data.write_all([1, 2, 3]).await.unwrap();
205        assert_eq!(data.size_hint(), Some(3));
206        assert_eq!(data.read().await.unwrap(), Some(1));
207        assert_eq!(
208            data.clone()
209                .into_stream()
210                .try_collect::<Vec<_>>()
211                .await
212                .unwrap(),
213            [2, 3]
214        );
215        assert_eq!(data.read_all().await.unwrap(), Vec::<i32>::new());
216    }
217
218    #[tokio::test]
219    async fn batches_stream_until_the_dataset_is_empty() {
220        let data = Data::new((1..=5).collect::<InMemDataset<u32>>());
221        let batches: Vec<Vec<u32>> = data.clone().into_batches(2).try_collect().await.unwrap();
222        assert_eq!(batches, [vec![1, 2], vec![3, 4], vec![5]]);
223
224        // A last full batch ends on the next, empty read.
225        data.write_all([6, 7]).await.unwrap();
226        let batches: Vec<Vec<u32>> = data.clone().into_batches(2).try_collect().await.unwrap();
227        assert_eq!(batches, [vec![6, 7]]);
228        let batches: Vec<Vec<u32>> = data.into_batches(2).try_collect().await.unwrap();
229        assert_eq!(batches, Vec::<Vec<u32>>::new());
230    }
231}
spire_core/dataset/data.rs

spire_core/dataset/
data.rs
