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