densa_data/
snapshot.rs

1//! A table's live rows as plain record batches, for whatever else reads Arrow.
2
3use std::sync::Arc;
4
5use arrow_array::builder::BooleanBufferBuilder;
6use arrow_array::{BooleanArray, RecordBatch};
7use arrow_select::filter::filter_record_batch;
8
9use crate::{DataError, LiveTable};
10
11impl LiveTable {
12    /// The live rows (the latest values of each key) as record batches, all with the table's
13    /// [schema](Self::schema) (stored batches keep the schema they arrived with, whose
14    /// columns may not be nullable), to hand to something that reads Arrow: a SQL engine, a
15    /// file writer, another thread (the table itself is used from one thread at a time). The
16    /// columns of stored batches whose rows are all live are shared, not copied; the others
17    /// are filtered down to their live rows. Rows come in no particular order.
18    ///
19    /// A table without live rows gives no batches: take the schema from
20    /// [`schema`](Self::schema). A table fed many small batches gives as many small ones;
21    /// [`compact`](Self::compact) first to have them merged (and superseded rows dropped).
22    ///
23    /// The batches are a snapshot: later upserts and deletes do not change them. Taking one
24    /// allocates: a bit per stored row of partly live batches, and their live rows' copies.
25    ///
26    /// # Errors
27    /// [`DataError::Arrow`] if Arrow cannot filter a batch or give it the table's schema
28    /// (neither should happen to rows the table accepted).
29    pub fn snapshot(&self) -> Result<Vec<RecordBatch>, DataError> {
30        let batches = self.batches();
31        let mut live = vec![0_usize; batches.len()];
32        for row in self.live_rows() {
33            live[row.batch] += 1;
34        }
35        // A mask, a bit per row, only for the batches that are partly live.
36        let mut masks: Vec<Option<BooleanBufferBuilder>> = batches
37            .iter()
38            .zip(&live)
39            .map(|(batch, &count)| {
40                let rows = batch.num_rows();
41                (count > 0 && count < rows).then(|| {
42                    let mut mask = BooleanBufferBuilder::new(rows);
43                    mask.append_n(rows, false);
44                    mask
45                })
46            })
47            .collect();
48        for row in self.live_rows() {
49            if let Some(mask) = &mut masks[row.batch] {
50                mask.set_bit(row.row, true);
51            }
52        }
53        batches
54            .iter()
55            .zip(live)
56            .zip(masks)
57            .filter(|((_, count), _)| *count > 0)
58            .map(|((batch, _), mask)| {
59                let live = match mask {
60                    None => batch.clone(),
61                    Some(mut mask) => {
62                        let mask = BooleanArray::new(mask.finish(), None);
63                        filter_record_batch(batch, &mask).map_err(DataError::arrow)?
64                    }
65                };
66                RecordBatch::try_new(Arc::clone(self.schema()), live.columns().to_vec())
67                    .map_err(DataError::arrow)
68            })
69            .collect()
70    }
71}
densa_data/snapshot.rs

densa_data/
snapshot.rs
