dbsp/trace/spine_async/
snapshot.rs

1//! A snapshot of a spine, which can be used to read from the spine without
2//! holding a reference to the spine itself.
3
4use std::fmt::Debug;
5use std::sync::Arc;
6
7use futures::{StreamExt, stream::FuturesUnordered};
8use rand::Rng;
9use rkyv::ser::Serializer;
10use rkyv::{Archive, Archived, Deserialize, Fallible, Serialize};
11use size_of::SizeOf;
12
13use super::SpineCursor;
14use crate::NumEntries;
15use crate::dynamic::{DynVec, Factory};
16use crate::storage::file::FilterStats;
17use crate::trace::cursor::{CursorFactory, CursorList};
18use crate::trace::{
19    Batch, BatchReader, BatchReaderFactories, Cursor, Spine, merge_batches,
20    sample_keys_from_batches,
21};
22
23pub trait WithSnapshot: Sized {
24    type Batch: Batch;
25
26    fn into_ro_snapshot(self) -> SpineSnapshot<Self::Batch> {
27        self.ro_snapshot()
28    }
29
30    /// Returns a read-only, non-merging snapshot of the current trace
31    /// state.
32    fn ro_snapshot(&self) -> SpineSnapshot<Self::Batch>;
33}
34
35pub trait BatchReaderWithSnapshot:
36    BatchReader<
37        Key = <Self::Batch as BatchReader>::Key,
38        Val = <Self::Batch as BatchReader>::Val,
39        Time = <Self::Batch as BatchReader>::Time,
40        R = <Self::Batch as BatchReader>::R,
41    > + WithSnapshot
42{
43}
44
45impl<B> BatchReaderWithSnapshot for B where
46    B: BatchReader<
47            Key = <Self::Batch as BatchReader>::Key,
48            Val = <Self::Batch as BatchReader>::Val,
49            Time = <Self::Batch as BatchReader>::Time,
50            R = <Self::Batch as BatchReader>::R,
51        > + WithSnapshot
52{
53}
54
55#[derive(Clone, SizeOf)]
56pub struct SpineSnapshot<B>
57where
58    B: Batch + Send + Sync,
59{
60    batches: Vec<Arc<B>>,
61    #[size_of(skip)]
62    factories: B::Factories,
63}
64
65impl<B> WithSnapshot for SpineSnapshot<B>
66where
67    B: Batch + Send + Sync,
68{
69    type Batch = B;
70
71    fn into_ro_snapshot(self) -> SpineSnapshot<Self::Batch> {
72        self
73    }
74
75    fn ro_snapshot(&self) -> SpineSnapshot<B> {
76        self.clone()
77    }
78}
79
80impl<B> WithSnapshot for B
81where
82    B: Batch,
83{
84    type Batch = B;
85    fn into_ro_snapshot(self) -> SpineSnapshot<B> {
86        let factories = self.factories();
87
88        SpineSnapshot {
89            batches: vec![Arc::new(self)],
90            factories,
91        }
92    }
93
94    fn ro_snapshot(&self) -> SpineSnapshot<Self::Batch> {
95        SpineSnapshot {
96            batches: vec![Arc::new(self.clone())],
97            factories: self.factories(),
98        }
99    }
100}
101
102impl<B: Batch + Send + Sync> Debug for SpineSnapshot<B> {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        f.debug_struct("SpineSnapshot")
105            .field("batches", &self.batches)
106            .finish()
107    }
108}
109
110impl<B> SpineSnapshot<B>
111where
112    B: Batch + Send + Sync,
113{
114    pub fn new(factories: B::Factories) -> Self {
115        Self {
116            batches: Vec::new(),
117            factories,
118        }
119    }
120
121    pub fn with_batches(factories: &B::Factories, batches: Vec<Arc<B>>) -> Self {
122        Self {
123            batches,
124            factories: factories.clone(),
125        }
126    }
127
128    pub fn extend(&mut self, other: Self) {
129        self.batches.extend(other.batches.iter().cloned())
130    }
131
132    pub fn extend_with_batches<I>(&mut self, batches: I)
133    where
134        I: IntoIterator<Item = Arc<B>>,
135    {
136        self.batches.extend(batches);
137    }
138
139    pub fn concat<'a, I>(factories: B::Factories, snapshots: I) -> Self
140    where
141        I: IntoIterator<Item = &'a Self>,
142    {
143        Self {
144            batches: snapshots
145                .into_iter()
146                .flat_map(|snapshot| snapshot.batches.iter().cloned())
147                .collect::<Vec<_>>(),
148            factories,
149        }
150    }
151
152    pub fn batches(&self) -> &[Arc<B>] {
153        &self.batches
154    }
155
156    pub fn consolidate(&self) -> B {
157        merge_batches(
158            &self.factories,
159            self.batches().iter().map(|b| b.as_ref().clone()),
160            &None,
161            &None,
162        )
163    }
164
165    pub fn into_batches(self) -> Vec<Arc<B>> {
166        self.batches
167    }
168}
169
170impl<B> From<&Spine<B>> for SpineSnapshot<B>
171where
172    B: Batch + Send + Sync,
173{
174    fn from(spine: &Spine<B>) -> Self {
175        Self {
176            batches: spine.merger.get_batches(),
177            factories: spine.factories.clone(),
178        }
179    }
180}
181
182impl<B> NumEntries for SpineSnapshot<B>
183where
184    B: Batch + Send + Sync,
185{
186    const CONST_NUM_ENTRIES: Option<usize> = None;
187
188    fn num_entries_shallow(&self) -> usize {
189        self.batches.iter().fold(0, |acc, batch| acc + batch.len())
190    }
191
192    fn num_entries_deep(&self) -> usize {
193        self.num_entries_shallow()
194    }
195}
196
197impl<B> BatchReader for SpineSnapshot<B>
198where
199    B: Batch + Send + Sync,
200{
201    type Factories = B::Factories;
202    type Key = B::Key;
203    type Val = B::Val;
204    type Time = B::Time;
205    type R = B::R;
206
207    type Cursor<'s> = SpineCursor<B>;
208
209    fn factories(&self) -> Self::Factories {
210        self.factories.clone()
211    }
212
213    fn cursor(&self) -> Self::Cursor<'_> {
214        SpineCursor::new_cursor(&self.factories, self.batches.clone())
215    }
216
217    fn key_count(&self) -> usize {
218        self.batches
219            .iter()
220            .fold(0, |acc, batch| acc + batch.key_count())
221    }
222
223    fn len(&self) -> usize {
224        self.batches.iter().fold(0, |acc, batch| acc + batch.len())
225    }
226
227    fn approximate_byte_size(&self) -> usize {
228        self.batches
229            .iter()
230            .fold(0, |acc, batch| acc + batch.approximate_byte_size())
231    }
232
233    fn membership_filter_stats(&self) -> FilterStats {
234        self.batches
235            .iter()
236            .map(|b| b.membership_filter_stats())
237            .sum()
238    }
239
240    fn range_filter_stats(&self) -> FilterStats {
241        self.batches.iter().map(|b| b.range_filter_stats()).sum()
242    }
243
244    fn sample_keys<RG>(&self, rng: &mut RG, sample_size: usize, sample: &mut DynVec<Self::Key>)
245    where
246        RG: Rng,
247    {
248        let total_keys = self
249            .batches
250            .iter()
251            .map(|batch| batch.key_count())
252            .sum::<usize>();
253        let batch_refs: Vec<_> = self.batches.iter().map(Arc::as_ref).collect();
254        sample_keys_from_batches(
255            &self.factories,
256            &batch_refs,
257            rng,
258            |batch| {
259                if sample_size == 0 || total_keys == 0 {
260                    0
261                } else {
262                    ((batch.key_count() as u128) * (sample_size as u128) / (total_keys as u128))
263                        as usize
264                }
265            },
266            sample,
267        );
268    }
269
270    async fn fetch<K>(
271        &self,
272        keys: &K,
273    ) -> Option<Box<dyn CursorFactory<Self::Key, Self::Val, Self::Time, Self::R>>>
274    where
275        K: BatchReader<Key = Self::Key, Time = ()>,
276    {
277        Some(Box::new(
278            FetchList::new(self.batches.clone(), keys, self.factories.weight_factory()).await,
279        ))
280    }
281}
282
283pub struct FetchList<B>
284where
285    B: BatchReader,
286{
287    weight_factory: &'static dyn Factory<B::R>,
288    batches: Vec<Arc<B>>,
289    fetched: Vec<Box<dyn CursorFactory<B::Key, B::Val, B::Time, B::R>>>,
290}
291
292impl<B> FetchList<B>
293where
294    B: BatchReader,
295{
296    pub async fn new<K>(
297        inputs: Vec<Arc<B>>,
298        keys: &K,
299        weight_factory: &'static dyn Factory<B::R>,
300    ) -> Self
301    where
302        K: BatchReader<Key = B::Key, Time = ()>,
303    {
304        let mut batches = Vec::new();
305        let mut fetched = Vec::new();
306        let mut futures = inputs
307            .into_iter()
308            .map(|b| async move { (b.clone(), b.fetch(keys).await) })
309            .collect::<FuturesUnordered<_>>();
310        while let Some((batch, fetch)) = futures.next().await {
311            if let Some(fetch) = fetch {
312                fetched.push(fetch);
313            } else {
314                batches.push(batch);
315            }
316        }
317
318        Self {
319            weight_factory,
320            batches,
321            fetched,
322        }
323    }
324}
325
326impl<B> CursorFactory<B::Key, B::Val, B::Time, B::R> for FetchList<B>
327where
328    B: Batch,
329{
330    fn get_cursor<'a>(&'a self) -> Box<dyn Cursor<B::Key, B::Val, B::Time, B::R> + 'a> {
331        let cursors =
332            self.fetched
333                .iter()
334                .map(|hc| hc.get_cursor())
335                .chain(self.batches.iter().map(|b| {
336                    Box::new(b.cursor()) as Box<dyn Cursor<B::Key, B::Val, B::Time, B::R>>
337                }))
338                .collect::<Vec<_>>();
339        Box::new(CursorList::new(self.weight_factory, cursors))
340    }
341}
342
343impl<B> Archive for SpineSnapshot<B>
344where
345    B: Batch + Send + Sync,
346{
347    type Archived = ();
348    type Resolver = ();
349
350    unsafe fn resolve(&self, _pos: usize, _resolver: Self::Resolver, _out: *mut Self::Archived) {
351        unimplemented!();
352    }
353}
354
355impl<B, S: Serializer + ?Sized> Serialize<S> for SpineSnapshot<B>
356where
357    B: Batch + Send + Sync,
358{
359    fn serialize(&self, _serializer: &mut S) -> Result<Self::Resolver, S::Error> {
360        unimplemented!();
361    }
362}
363
364impl<B, D: Fallible> Deserialize<SpineSnapshot<B>, D> for Archived<SpineSnapshot<B>>
365where
366    B: Batch + Send + Sync,
367{
368    fn deserialize(&self, _deserializer: &mut D) -> Result<SpineSnapshot<B>, D::Error> {
369        unimplemented!();
370    }
371}
dbsp/trace/spine_async/snapshot.rs

dbsp/trace/spine_async/
snapshot.rs
