Skip to main content

densa_data/
table.rs

1use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
2use std::sync::Arc;
3
4use arrow_array::{Array, ArrayRef, RecordBatch};
5use arrow_schema::{Field, Schema, SchemaRef};
6use rustc_hash::FxHashMap;
7
8use crate::dictionary::interleave_column;
9use crate::group::{GroupState, Groups};
10use crate::key::{key_reader, KeyFn, KeyIndex, KeyRef};
11use crate::view::{Rows, ViewState};
12use crate::{DataError, GroupSpec, RowId, Scalar, ViewSpec};
13
14/// Changing more than one row in this many rebuilds each view instead of patching it.
15const REBUILD_RATIO: usize = 8;
16/// Superseded rows tolerated before an automatic compaction (and only once they outnumber
17/// the live rows).
18const MIN_DEAD_ROWS: usize = 4_096;
19/// Stored batches tolerated before an automatic compaction merges them.
20const MAX_BATCHES: usize = 4_096;
21/// Rows per batch written by compaction.
22const COMPACT_ROWS: usize = 65_536;
23/// Compaction splits a chunk Arrow cannot build into halves, down to this many rows. Any
24/// dictionary key type indexes this many distinct strings, so a smaller chunk would not help.
25const MIN_SPLIT_ROWS: usize = 128;
26/// Row ids are `u32`: at most this many live keys.
27const MAX_ROWS: usize = 1 << 32;
28
29/// Hands out identities of tables, views and groups, unique within the process.
30static NEXT_SERIAL: AtomicU64 = AtomicU64::new(1);
31
32fn next_serial() -> u64 {
33    // ORDERING: Relaxed. Only uniqueness matters; no other memory is published with it.
34    NEXT_SERIAL.fetch_add(1, AtomicOrdering::Relaxed)
35}
36
37/// Where a row's values are: `table.batches()[batch]`, row `row`. Valid until the next
38/// compaction, which [`LiveTable::upsert`] and [`LiveTable::delete`] may run: compare
39/// [`LiveTable::generation`] before reusing one.
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
41pub struct RowRef {
42    /// Index into [`LiveTable::batches`].
43    pub batch: usize,
44    /// Row within that batch.
45    pub row: usize,
46}
47
48/// Where a live row id's values are, packed to 8 bytes.
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub(crate) struct Slot {
51    batch: u32,
52    row: u32,
53}
54
55impl Slot {
56    pub(crate) fn row_ref(self) -> RowRef {
57        RowRef {
58            batch: self.batch as usize,
59            row: self.row as usize,
60        }
61    }
62}
63
64/// What an upsert or delete did.
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
66#[non_exhaustive]
67pub struct Changes {
68    /// Keys that were new.
69    pub inserted: usize,
70    /// Keys whose row was replaced.
71    pub updated: usize,
72    /// Keys removed.
73    pub deleted: usize,
74}
75
76/// How a [`LiveTable`] manages its storage.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78#[non_exhaustive]
79pub struct TableOptions {
80    /// Whether [`LiveTable::upsert`] and [`LiveTable::delete`] compact the table when
81    /// superseded rows outnumber live ones (at least 4,096 of them) or more than 4,096 batches
82    /// were stored since the last compaction. Compaction copies every live row, a pause
83    /// proportional to the table;
84    /// without it, call [`LiveTable::compact`] when convenient. Default: on.
85    pub auto_compact: bool,
86}
87
88impl Default for TableOptions {
89    fn default() -> Self {
90        Self { auto_compact: true }
91    }
92}
93
94impl TableOptions {
95    /// The options with automatic compaction on or off.
96    #[must_use]
97    pub fn with_auto_compact(mut self, auto_compact: bool) -> Self {
98        self.auto_compact = auto_compact;
99        self
100    }
101}
102
103/// Identifies a view of one [`LiveTable`]. A clone of the table copies its views, so the id
104/// of a view added before cloning works on both; one added afterwards, only on its table.
105#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
106pub struct ViewId {
107    /// Unique to the view's addition.
108    serial: u64,
109    index: usize,
110}
111
112/// Identifies grouped aggregates of one [`LiveTable`] (and of clones made after they were
113/// added, as for [`ViewId`]).
114#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
115pub struct GroupsId {
116    serial: u64,
117    index: usize,
118}
119
120/// A record: one key's row for as long as the key is live, through updates, compactions and
121/// view changes. Deleting the key ends the record; inserting the key again starts a new one.
122/// Get one with [`LiveTable::record_of_key`] or [`View::record`], find its row with
123/// [`LiveTable::row_of`] and its place in a view with [`View::position`].
124///
125/// It is a plain number ([`to_bits`](Self::to_bits)), so a widget can keep it as a selection
126/// without borrowing the table. It names a record of the table it came from and of clones
127/// made after; in any other table it may name an unrelated row.
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
129pub struct RecordId(u64);
130
131impl RecordId {
132    fn new(id: RowId, reuse: u32) -> Self {
133        Self(u64::from(reuse) << 32 | u64::from(id))
134    }
135
136    /// The record as a number, for storing outside densa-data.
137    #[must_use]
138    pub const fn to_bits(self) -> u64 {
139        self.0
140    }
141
142    /// The record that [`to_bits`](Self::to_bits) gave `bits` for. A number it did not give
143    /// names no record, or an arbitrary one.
144    #[must_use]
145    pub const fn from_bits(bits: u64) -> Self {
146        Self(bits)
147    }
148
149    /// The row id, and how many times it had been freed when the record began.
150    fn parts(self) -> (RowId, u32) {
151        // Lossless: the low and high halves.
152        (self.0 as u32, (self.0 >> 32) as u32)
153    }
154}
155
156/// One planned row change: its id, where its values were (if it existed) and are now, and
157/// its key if new.
158type Change<'a> = (RowId, Option<Slot>, Slot, Option<KeyRef<'a>>);
159
160/// A table of rows identified by a key column, updated by Arrow record batches, with views
161/// (filtered, sorted) that follow every change.
162///
163/// Storage is log-structured: an upserted batch is kept as it is (its arrays shared, not
164/// copied) and each key points at its latest row. Superseded rows are dropped by compaction,
165/// automatically (see [`TableOptions`]) or with [`compact`](Self::compact). Every key has a
166/// stable internal id, so compaction never reorders a view.
167///
168/// Views and grouped aggregates ([`GroupSpec`]) are maintained incrementally: changing `k`
169/// rows costs each view about `k` sorted insertions and removals (a few binary searches over
170/// contiguous memory plus moving one block of at most 1,024 entries), and each group spec as
171/// much per min or max column plus `O(log groups)` per row, not a re-sort or a recount; a
172/// change touching more than an eighth of the rows rebuilds them instead. Reading a view
173/// ([`View::row`], `O(log n)`) never allocates; [`Groups::to_batch`] does.
174///
175/// Every column of the table is nullable, whatever the schema given says; batches must have
176/// the table's column names and data types, in order.
177#[derive(Debug)]
178pub struct LiveTable {
179    /// Unique to this table object; a clone gets its own.
180    instance: u64,
181    options: TableOptions,
182    schema: SchemaRef,
183    key_column: usize,
184    key_read: KeyFn,
185    keys: KeyIndex,
186    batches: Vec<RecordBatch>,
187    /// Where each row id's values are; `None` for a free id.
188    slots: Vec<Option<Slot>>,
189    free: Vec<RowId>,
190    /// How many times each row id was freed (missing: never), so a record of a deleted key
191    /// never names the key that reuses its id.
192    reuses: Vec<u32>,
193    /// Stored rows that are no longer live.
194    dead: usize,
195    /// Each view with the serial of its id.
196    views: Vec<Option<(u64, ViewState)>>,
197    groups: Vec<Option<(u64, GroupState)>>,
198    generation: u64,
199    /// The last automatic compaction's failure, until a compaction succeeds.
200    compaction_error: Option<DataError>,
201    /// Dead rows at which automatic compaction tries again after a failure.
202    retry_at_dead: usize,
203    /// Batches the last compaction wrote: automatic compaction counts the ones stored since.
204    compacted_batches: usize,
205}
206
207impl Clone for LiveTable {
208    fn clone(&self) -> Self {
209        Self {
210            instance: next_serial(),
211            options: self.options,
212            schema: Arc::clone(&self.schema),
213            key_column: self.key_column,
214            key_read: self.key_read,
215            keys: self.keys.clone(),
216            batches: self.batches.clone(),
217            slots: self.slots.clone(),
218            free: self.free.clone(),
219            reuses: self.reuses.clone(),
220            dead: self.dead,
221            views: self.views.clone(),
222            groups: self.groups.clone(),
223            generation: self.generation,
224            compaction_error: self.compaction_error.clone(),
225            retry_at_dead: self.retry_at_dead,
226            compacted_batches: self.compacted_batches,
227        }
228    }
229}
230
231impl LiveTable {
232    /// An empty table of `schema`, keyed by the column named `key` (integers or strings),
233    /// with default [`TableOptions`].
234    ///
235    /// # Errors
236    /// [`DataError::UnknownColumn`] or [`DataError::UnsupportedKey`].
237    pub fn new(schema: SchemaRef, key: &str) -> Result<Self, DataError> {
238        Self::with_options(schema, key, TableOptions::default())
239    }
240
241    /// An empty table of `schema`, keyed by the column named `key`, managed with `options`.
242    ///
243    /// # Errors
244    /// [`DataError::UnknownColumn`] or [`DataError::UnsupportedKey`].
245    pub fn with_options(
246        schema: SchemaRef,
247        key: &str,
248        options: TableOptions,
249    ) -> Result<Self, DataError> {
250        let (key_column, field) = schema
251            .column_with_name(key)
252            .ok_or_else(|| DataError::UnknownColumn(key.to_owned()))?;
253        let unsupported = || DataError::UnsupportedKey(key.to_owned());
254        let key_read = key_reader(field.data_type()).ok_or_else(unsupported)?;
255        let keys = KeyIndex::for_type(field.data_type()).ok_or_else(unsupported)?;
256        Ok(Self {
257            instance: next_serial(),
258            options,
259            schema: nullable(&schema),
260            key_column,
261            key_read,
262            keys,
263            batches: Vec::new(),
264            slots: Vec::new(),
265            free: Vec::new(),
266            reuses: Vec::new(),
267            dead: 0,
268            views: Vec::new(),
269            groups: Vec::new(),
270            generation: 0,
271            compaction_error: None,
272            retry_at_dead: MIN_DEAD_ROWS,
273            compacted_batches: 0,
274        })
275    }
276
277    /// The table's schema (every column nullable).
278    pub fn schema(&self) -> &SchemaRef {
279        &self.schema
280    }
281
282    /// Distinguishes this table object from every other one in the process, clones
283    /// included: something caching per table (a grid's formatters) compares it with
284    /// [`generation`](Self::generation).
285    pub fn instance(&self) -> u64 {
286        self.instance
287    }
288
289    /// Live rows: one per key.
290    pub fn len(&self) -> usize {
291        self.keys.len()
292    }
293
294    /// Whether there are no live rows.
295    pub fn is_empty(&self) -> bool {
296        self.len() == 0
297    }
298
299    /// The stored batches [`RowRef`]s point into. They also hold superseded rows until the
300    /// next compaction.
301    pub fn batches(&self) -> &[RecordBatch] {
302        &self.batches
303    }
304
305    /// Counts compactions: when it changes, batches were rewritten and earlier [`RowRef`]s
306    /// are invalid (views stay valid).
307    pub fn generation(&self) -> u64 {
308        self.generation
309    }
310
311    /// Why the last automatic compaction failed, if it did and no compaction has succeeded
312    /// since. The table stays correct, only larger; automatic compaction tries again once
313    /// twice as many rows are superseded.
314    pub fn compaction_error(&self) -> Option<&DataError> {
315        self.compaction_error.as_ref()
316    }
317
318    /// Where the live row with key `key` is.
319    pub fn find(&self, key: &Scalar) -> Option<RowRef> {
320        let id = self.key_id(key)?;
321        self.slots[id as usize].map(Slot::row_ref)
322    }
323
324    /// The record of key `key`, if it is live.
325    pub fn record_of_key(&self, key: &Scalar) -> Option<RecordId> {
326        let id = self.key_id(key)?;
327        Some(RecordId::new(id, reuse(&self.reuses, id)))
328    }
329
330    /// Where the row of `record` is, or `None` if its key was deleted.
331    pub fn row_of(&self, record: RecordId) -> Option<RowRef> {
332        live_slot(&self.slots, &self.reuses, record).map(|(_, slot)| slot.row_ref())
333    }
334
335    /// The row id of key `key`, if it is live.
336    fn key_id(&self, key: &Scalar) -> Option<RowId> {
337        let key = match (key, &self.keys) {
338            (Scalar::Int(k), KeyIndex::Int(_)) => KeyRef::Int(*k),
339            (Scalar::Int(k), KeyIndex::UInt(_)) => KeyRef::UInt(u64::try_from(*k).ok()?),
340            (Scalar::UInt(k), KeyIndex::UInt(_)) => KeyRef::UInt(*k),
341            (Scalar::UInt(k), KeyIndex::Int(_)) => KeyRef::Int(i64::try_from(*k).ok()?),
342            (Scalar::Str(k), KeyIndex::Str(_)) => KeyRef::Str(k),
343            _ => return None,
344        };
345        self.keys.get(key)
346    }
347
348    /// Inserts `batch`'s rows, replacing the rows of keys already present; within the batch
349    /// the last row for a key wins. Views follow.
350    ///
351    /// # Errors
352    /// A column error if `batch` has other column names or data types,
353    /// [`DataError::NullKey`] if a row has no key, [`DataError::TooManyRows`] past 2^32 live
354    /// keys or rows in a batch. Nothing changes on error.
355    pub fn upsert(&mut self, batch: RecordBatch) -> Result<Changes, DataError> {
356        self.check(&batch)?;
357        let latest = self.latest_rows(&batch)?;
358        if latest.is_empty() {
359            return Ok(Changes::default());
360        }
361        let plan = self.plan(&batch, &latest)?;
362        let updated = plan.iter().filter(|c| c.1.is_some()).count();
363        let changes = Changes {
364            inserted: plan.len() - updated,
365            updated,
366            deleted: 0,
367        };
368        let rebuild = plan.len().saturating_mul(REBUILD_RATIO) > self.len() + changes.inserted;
369
370        self.batches.push(batch.clone());
371        if !rebuild {
372            self.drop_from_views(plan.iter().filter_map(|c| Some((c.0, c.1?))));
373        }
374        for &(id, old, at, key) in &plan {
375            self.slots[id as usize] = Some(at);
376            self.dead += usize::from(old.is_some());
377            if let Some(key) = key {
378                self.keys.insert(key, id);
379            }
380        }
381        // Earlier rows of keys repeated within the batch were never live.
382        self.dead += batch.num_rows() - plan.len();
383        self.add_to_views(rebuild, plan.iter().map(|c| (c.0, c.2)));
384        self.compact_if_due();
385        Ok(changes)
386    }
387
388    /// Deletes the rows of `keys` (an array of the key column's kind; nulls and unknown keys
389    /// are skipped). Views follow.
390    ///
391    /// # Errors
392    /// [`DataError::KeyType`] if `keys` is not of the key column's kind: nothing changes.
393    pub fn delete(&mut self, keys: &dyn Array) -> Result<Changes, DataError> {
394        let read = key_reader(keys.data_type())
395            .filter(|_| self.keys.same_kind(keys.data_type()))
396            .ok_or(DataError::KeyType)?;
397        let mut deleted = 0;
398        for row in 0..keys.len() {
399            let Some(id) = read(keys, row).and_then(|key| self.keys.remove(key)) else {
400                continue;
401            };
402            if let Some(slot) = self.slots[id as usize] {
403                self.drop_from_views(std::iter::once((id, slot)));
404            }
405            self.slots[id as usize] = None;
406            self.free.push(id);
407            if self.reuses.len() <= id as usize {
408                self.reuses.resize(id as usize + 1, 0);
409            }
410            // Wrapping: an id freed 2^32 times could give an old record back; harmless odds.
411            self.reuses[id as usize] = self.reuses[id as usize].wrapping_add(1);
412            self.dead += 1;
413            deleted += 1;
414        }
415        self.compact_if_due();
416        Ok(Changes {
417            deleted,
418            ..Changes::default()
419        })
420    }
421
422    /// Rewrites the stored batches with only the live rows, in batches of up to 65,536 rows
423    /// (fewer where one batch cannot hold them: more distinct strings than a dictionary
424    /// column's key type indexes, over 2 GiB of strings). Dictionary columns of strings are
425    /// re-encoded with the strings their rows use. Views keep their order;
426    /// [`generation`](Self::generation) increases.
427    ///
428    /// # Errors
429    /// [`DataError::Arrow`] if Arrow cannot build a batch of even 128 of the rows; nothing
430    /// changes.
431    pub fn compact(&mut self) -> Result<(), DataError> {
432        let live: Vec<(RowId, Slot)> = self
433            .slots
434            .iter()
435            .enumerate()
436            .filter_map(|(id, slot)| Some((RowId::try_from(id).ok()?, (*slot)?)))
437            .collect();
438        let mut compacted = Vec::new();
439        for chunk in live.chunks(COMPACT_ROWS) {
440            self.gather_into(chunk, &mut compacted)?;
441        }
442        // The live rows, in order, fill the new batches in order.
443        let places = compacted.iter().enumerate().flat_map(|(batch, rows)| {
444            // Lossless: at most 2^32 live rows, so as many batches or rows in one.
445            (0..rows.num_rows()).map(move |row| Slot {
446                batch: batch as u32,
447                row: row as u32,
448            })
449        });
450        for (&(id, _), slot) in live.iter().zip(places) {
451            self.slots[id as usize] = Some(slot);
452        }
453        self.compacted_batches = compacted.len();
454        self.batches = compacted;
455        self.dead = 0;
456        self.generation += 1;
457        self.compaction_error = None;
458        self.retry_at_dead = MIN_DEAD_ROWS;
459        Ok(())
460    }
461
462    /// Adds a view of the rows matching `spec`, filled from the current rows.
463    ///
464    /// # Errors
465    /// [`DataError::UnknownColumn`], [`DataError::UnsupportedColumn`] or
466    /// [`DataError::LiteralMismatch`] if the spec does not fit the schema.
467    pub fn add_view(&mut self, spec: ViewSpec) -> Result<ViewId, DataError> {
468        let mut view = ViewState::compile(&spec, &self.schema)?;
469        view.rebuild(Rows {
470            batches: &self.batches,
471            slots: &self.slots,
472        });
473        let (serial, index) = store(&mut self.views, view);
474        Ok(ViewId { serial, index })
475    }
476
477    /// Removes a view; `false` if it was already gone or belongs to another table. Its id
478    /// never names another view (its storage may be reused).
479    pub fn remove_view(&mut self, id: ViewId) -> bool {
480        take(&mut self.views, id.index, id.serial)
481    }
482
483    /// The rows of view `id`, in order; `None` if it was removed or belongs to another
484    /// table.
485    pub fn view(&self, id: ViewId) -> Option<View<'_>> {
486        let state = find(&self.views, id.index, id.serial)?;
487        Some(View {
488            state,
489            batches: &self.batches,
490            slots: &self.slots,
491            reuses: &self.reuses,
492        })
493    }
494
495    /// Adds grouped aggregates of the rows matching `spec`, filled from the current rows.
496    ///
497    /// # Errors
498    /// [`DataError::UnknownColumn`] or [`DataError::UnsupportedColumn`] for a key or an
499    /// aggregated column, [`DataError::NotNumeric`] for a sum or mean of a column that is
500    /// not numeric, [`DataError::DuplicateName`] if two result columns share a name,
501    /// [`DataError::LiteralMismatch`] for a filter that does not fit.
502    pub fn add_groups(&mut self, spec: GroupSpec) -> Result<GroupsId, DataError> {
503        let mut groups = GroupState::compile(&spec, &self.schema)?;
504        groups.rebuild(Rows {
505            batches: &self.batches,
506            slots: &self.slots,
507        });
508        let (serial, index) = store(&mut self.groups, groups);
509        Ok(GroupsId { serial, index })
510    }
511
512    /// Removes grouped aggregates; `false` if already gone or another table's. Their id
513    /// never names other aggregates (their storage may be reused).
514    pub fn remove_groups(&mut self, id: GroupsId) -> bool {
515        take(&mut self.groups, id.index, id.serial)
516    }
517
518    /// The groups of `id`, in key order; `None` if removed or another table's.
519    pub fn groups(&self, id: GroupsId) -> Option<Groups<'_>> {
520        let state = find(&self.groups, id.index, id.serial)?;
521        Some(Groups {
522            state,
523            rows: Rows {
524                batches: &self.batches,
525                slots: &self.slots,
526            },
527        })
528    }
529
530    fn check(&self, batch: &RecordBatch) -> Result<(), DataError> {
531        let ours = self.schema.fields();
532        let theirs = batch.schema_ref().fields();
533        if ours.len() != theirs.len() {
534            return Err(DataError::ColumnCount {
535                expected: ours.len(),
536                found: theirs.len(),
537            });
538        }
539        match ours
540            .iter()
541            .zip(theirs.iter())
542            .position(|(a, b)| a.name() != b.name() || a.data_type() != b.data_type())
543        {
544            Some(column) => Err(DataError::ColumnMismatch { column }),
545            None => Ok(()),
546        }
547    }
548
549    /// The last row of each distinct key in `batch`, in row order.
550    fn latest_rows<'b>(
551        &self,
552        batch: &'b RecordBatch,
553    ) -> Result<Vec<(usize, KeyRef<'b>)>, DataError> {
554        let keys = batch.column(self.key_column).as_ref();
555        let mut last: FxHashMap<KeyRef<'b>, usize> =
556            FxHashMap::with_capacity_and_hasher(batch.num_rows(), Default::default());
557        for row in 0..batch.num_rows() {
558            let key = (self.key_read)(keys, row).ok_or(DataError::NullKey { row })?;
559            last.insert(key, row);
560        }
561        let mut rows: Vec<(usize, KeyRef<'b>)> = last.into_iter().map(|(k, r)| (r, k)).collect();
562        rows.sort_unstable_by_key(|&(row, _)| row);
563        Ok(rows)
564    }
565
566    /// Each changed row's id, old and new place and new key. Checks every limit before
567    /// taking any id, so an error leaves the table unchanged.
568    fn plan<'b>(
569        &mut self,
570        batch: &RecordBatch,
571        latest: &[(usize, KeyRef<'b>)],
572    ) -> Result<Vec<Change<'b>>, DataError> {
573        let segment = u32::try_from(self.batches.len()).map_err(|_| DataError::TooManyRows)?;
574        u32::try_from(batch.num_rows()).map_err(|_| DataError::TooManyRows)?;
575        let new_keys = latest
576            .iter()
577            .filter(|&&(_, key)| self.keys.get(key).is_none())
578            .count();
579        if new_keys > self.free.len() + (MAX_ROWS - self.slots.len()) {
580            return Err(DataError::TooManyRows);
581        }
582        let mut plan = Vec::with_capacity(latest.len());
583        for &(row, key) in latest {
584            // Lossless: the batch has fewer than 2^32 rows (checked above).
585            let at = Slot {
586                batch: segment,
587                row: row as u32,
588            };
589            plan.push(match self.keys.get(key) {
590                Some(id) => (id, self.slots[id as usize], at, None),
591                None => (self.allocate_id(), None, at, Some(key)),
592            });
593        }
594        Ok(plan)
595    }
596
597    /// A free row id; [`plan`](Self::plan) checked that one exists.
598    fn allocate_id(&mut self) -> RowId {
599        if let Some(id) = self.free.pop() {
600            return id;
601        }
602        self.slots.push(None);
603        // Lossless: `plan` keeps the slot count at most 2^32.
604        (self.slots.len() - 1) as RowId
605    }
606
607    /// Copies the live rows of `chunk` into new batches at the end of `out`: one if Arrow can
608    /// build it, otherwise each half on its own, recursively, down to [`MIN_SPLIT_ROWS`]:
609    /// fewer rows use fewer distinct strings, so they fit a narrow dictionary key type.
610    fn gather_into(
611        &self,
612        chunk: &[(RowId, Slot)],
613        out: &mut Vec<RecordBatch>,
614    ) -> Result<(), DataError> {
615        match self.gather(chunk) {
616            Ok(batch) => {
617                out.push(batch);
618                Ok(())
619            }
620            Err(_) if chunk.len() > MIN_SPLIT_ROWS => {
621                let (first, second) = chunk.split_at(chunk.len() / 2);
622                self.gather_into(first, out)?;
623                self.gather_into(second, out)
624            }
625            Err(error) => Err(error),
626        }
627    }
628
629    /// The live rows of `chunk`, copied into one batch.
630    fn gather(&self, chunk: &[(RowId, Slot)]) -> Result<RecordBatch, DataError> {
631        // Only the batches the chunk reads go to Arrow, which may size a dictionary column by
632        // all of its inputs' values.
633        let mut position: FxHashMap<usize, usize> = FxHashMap::default();
634        let mut sources = Vec::new();
635        let indices: Vec<(usize, usize)> = chunk
636            .iter()
637            .map(|(_, slot)| {
638                let batch = slot.batch as usize;
639                let at = *position.entry(batch).or_insert_with(|| {
640                    sources.push(batch);
641                    sources.len() - 1
642                });
643                (at, slot.row as usize)
644            })
645            .collect();
646        let columns = (0..self.schema.fields().len())
647            .map(|column| {
648                let arrays: Vec<&dyn Array> = sources
649                    .iter()
650                    .map(|&batch| self.batches[batch].column(column).as_ref())
651                    .collect();
652                interleave_column(&arrays, &indices)
653            })
654            .collect::<Result<Vec<ArrayRef>, _>>()
655            .map_err(DataError::arrow)?;
656        RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(DataError::arrow)
657    }
658
659    /// Removes rows from every view and group, each at the place its values were when it
660    /// entered.
661    fn drop_from_views(&mut self, rows: impl Iterator<Item = (RowId, Slot)> + Clone) {
662        let source = Rows {
663            batches: &self.batches,
664            slots: &self.slots,
665        };
666        for (_, view) in self.views.iter_mut().flatten() {
667            for (id, slot) in rows.clone() {
668                view.remove(source, id, slot.row_ref());
669            }
670        }
671        for (_, groups) in self.groups.iter_mut().flatten() {
672            for (id, slot) in rows.clone() {
673                groups.remove(source, id, slot.row_ref());
674            }
675        }
676    }
677
678    /// Adds changed rows to every view and group, or rebuilds them.
679    fn add_to_views(&mut self, rebuild: bool, rows: impl Iterator<Item = (RowId, Slot)> + Clone) {
680        let source = Rows {
681            batches: &self.batches,
682            slots: &self.slots,
683        };
684        for (_, view) in self.views.iter_mut().flatten() {
685            if rebuild {
686                view.rebuild(source);
687            } else {
688                for (id, slot) in rows.clone() {
689                    view.insert(source, id, slot.row_ref());
690                }
691            }
692        }
693        for (_, groups) in self.groups.iter_mut().flatten() {
694            if rebuild {
695                groups.rebuild(source);
696            } else {
697                for (id, slot) in rows.clone() {
698                    groups.insert(source, id, slot.row_ref());
699                }
700            }
701        }
702    }
703
704    /// Compacts when due and enabled. A failure is recorded, not returned: the change that
705    /// triggered it is applied and the table stays correct, only larger.
706    fn compact_if_due(&mut self) {
707        if !self.options.auto_compact {
708            return;
709        }
710        let dead_majority = self.dead >= self.retry_at_dead && self.dead > self.len();
711        // Counting only the batches stored since the last compaction: one that had to write
712        // many small batches does not run again at every change.
713        let added = self.batches.len().saturating_sub(self.compacted_batches);
714        let too_many_batches = added > MAX_BATCHES && self.compaction_error.is_none();
715        if !(dead_majority || too_many_batches) {
716            return;
717        }
718        if let Err(error) = self.compact() {
719            self.compaction_error = Some(error);
720            self.retry_at_dead = self.dead.saturating_mul(2);
721        }
722    }
723}
724
725/// The entry at `index` if its serial is `serial`.
726fn find<T>(entries: &[Option<(u64, T)>], index: usize, serial: u64) -> Option<&T> {
727    match entries.get(index)? {
728        Some((s, entry)) if *s == serial => Some(entry),
729        _ => None,
730    }
731}
732
733/// Stores `entry` under a new serial in the first free place (removed entries leave one),
734/// so a model that replaces its view on every sort keeps one; returns the serial and index.
735fn store<T>(entries: &mut Vec<Option<(u64, T)>>, entry: T) -> (u64, usize) {
736    let serial = next_serial();
737    match entries.iter().position(Option::is_none) {
738        Some(index) => {
739            entries[index] = Some((serial, entry));
740            (serial, index)
741        }
742        None => {
743            entries.push(Some((serial, entry)));
744            (serial, entries.len() - 1)
745        }
746    }
747}
748
749/// Removes the entry at `index` if its serial is `serial`; whether it did.
750fn take<T>(entries: &mut [Option<(u64, T)>], index: usize, serial: u64) -> bool {
751    entries
752        .get_mut(index)
753        .and_then(|entry| entry.take_if(|(s, _)| *s == serial))
754        .is_some()
755}
756
757/// `schema` with every field nullable: superseded rows and compaction never need a column
758/// to be free of nulls, and batches may carry nulls the first batch's schema did not.
759fn nullable(schema: &Schema) -> SchemaRef {
760    let fields: Vec<Field> = schema
761        .fields()
762        .iter()
763        .map(|f| f.as_ref().clone().with_nullable(true))
764        .collect();
765    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
766}
767
768/// The rows of a view, in its order. Cheap to get; reading never allocates.
769#[derive(Debug, Clone, Copy)]
770pub struct View<'a> {
771    state: &'a ViewState,
772    batches: &'a [RecordBatch],
773    slots: &'a [Option<Slot>],
774    reuses: &'a [u32],
775}
776
777impl View<'_> {
778    /// Rows in the view.
779    pub fn len(&self) -> usize {
780        self.state.len()
781    }
782
783    /// Whether the view is empty.
784    pub fn is_empty(&self) -> bool {
785        self.len() == 0
786    }
787
788    /// Where the row at `index` (0 is first) is. `O(log n)`.
789    pub fn row(&self, index: usize) -> Option<RowRef> {
790        let id = self.state.get(index)?;
791        self.slots
792            .get(id as usize)
793            .copied()
794            .flatten()
795            .map(Slot::row_ref)
796    }
797
798    /// The record of the row at `index`. `O(log n)`.
799    pub fn record(&self, index: usize) -> Option<RecordId> {
800        let id = self.state.get(index)?;
801        Some(RecordId::new(id, reuse(self.reuses, id)))
802    }
803
804    /// Where `record`'s row is in the view, or `None` if it is not in it (filtered out, or
805    /// its key deleted). `O(log n)`: binary searches comparing the row's sort values.
806    pub fn position(&self, record: RecordId) -> Option<usize> {
807        let (id, _) = live_slot(self.slots, self.reuses, record)?;
808        let rows = Rows {
809            batches: self.batches,
810            slots: self.slots,
811        };
812        self.state.position(rows, id)
813    }
814}
815
816/// How many times row id `id` was freed.
817fn reuse(reuses: &[u32], id: RowId) -> u32 {
818    reuses.get(id as usize).copied().unwrap_or(0)
819}
820
821/// `record`'s row id and slot, if the record's key is still live.
822fn live_slot(slots: &[Option<Slot>], reuses: &[u32], record: RecordId) -> Option<(RowId, Slot)> {
823    let (id, reused) = record.parts();
824    let slot = (*slots.get(id as usize)?)?;
825    (reuse(reuses, id) == reused).then_some((id, slot))
826}
827
828#[cfg(test)]
829mod tests {
830    use std::sync::Arc;
831
832    use arrow_array::{ArrayRef, Int64Array};
833
834    use super::*;
835
836    fn ids(keys: Vec<i64>) -> RecordBatch {
837        RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(keys)) as ArrayRef)]).unwrap()
838    }
839
840    #[test]
841    fn replacing_a_view_again_and_again_keeps_one_place() {
842        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
843        let mut view = table.add_view(ViewSpec::new()).unwrap();
844        for _ in 0..100 {
845            let next = table.add_view(ViewSpec::new()).unwrap();
846            table.remove_view(view);
847            view = next;
848        }
849        assert!(table.views.len() <= 2, "{} places", table.views.len());
850    }
851
852    #[test]
853    fn ids_of_deleted_keys_are_reused() {
854        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
855        for round in 0..100 {
856            table
857                .upsert(ids((0..10).map(|k| round * 10 + k).collect()))
858                .unwrap();
859            table
860                .delete(&Int64Array::from(
861                    (0..10).map(|k| round * 10 + k).collect::<Vec<_>>(),
862                ))
863                .unwrap();
864        }
865        assert!(
866            table.slots.len() <= 10,
867            "{} ids for 10 live keys",
868            table.slots.len()
869        );
870    }
871
872    #[test]
873    fn a_failed_automatic_compaction_is_recorded_and_backs_off() {
874        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
875        table.dead = MIN_DEAD_ROWS;
876        // A batch the table's schema cannot hold makes the next compaction fail.
877        let other = RecordBatch::try_from_iter([(
878            "other",
879            Arc::new(arrow_array::StringArray::from(vec!["x"])) as ArrayRef,
880        )])
881        .unwrap();
882        table.batches.push(other);
883        table.slots.push(Some(Slot { batch: 0, row: 0 }));
884        table.compact_if_due();
885        assert!(table.compaction_error().is_some());
886        assert_eq!(table.retry_at_dead, 2 * MIN_DEAD_ROWS);
887    }
888
889    #[test]
890    fn automatic_compaction_counts_the_batches_stored_since_the_last_one() {
891        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
892        table.upsert(ids(vec![1])).unwrap();
893        // As if the last compaction had to write 10 batches.
894        table.compacted_batches = 10;
895        while table.batches.len() < 10 + MAX_BATCHES {
896            table.batches.push(ids(vec![]));
897        }
898        table.compact_if_due();
899        assert_eq!(table.generation(), 0, "4,096 batches added");
900        table.batches.push(ids(vec![]));
901        table.compact_if_due();
902        assert_eq!(table.generation(), 1, "4,097 batches added");
903        assert_eq!(table.compacted_batches, 1);
904    }
905}