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 of key `key`, if it is live.
346    fn key_id(&self, key: &Scalar) -> Option<RowId> {
347        let key = match (key, &self.keys) {
348            (Scalar::Int(k), KeyIndex::Int(_)) => KeyRef::Int(*k),
349            (Scalar::Int(k), KeyIndex::UInt(_)) => KeyRef::UInt(u64::try_from(*k).ok()?),
350            (Scalar::UInt(k), KeyIndex::UInt(_)) => KeyRef::UInt(*k),
351            (Scalar::UInt(k), KeyIndex::Int(_)) => KeyRef::Int(i64::try_from(*k).ok()?),
352            (Scalar::Str(k), KeyIndex::Str(_)) => KeyRef::Str(k),
353            _ => return None,
354        };
355        self.keys.get(key)
356    }
357
358    /// Inserts `batch`'s rows, replacing the rows of keys already present; within the batch
359    /// the last row for a key wins. Then, with a [`Window`], deletes the rows it leaves
360    /// behind (counted in [`Changes::deleted`]; a row already past it when it came counts as
361    /// inserted and deleted). Views follow.
362    ///
363    /// # Errors
364    /// A column error if `batch` has other column names or data types,
365    /// [`DataError::NullKey`] if a row has no key, [`DataError::TooManyRows`] past 2^32 live
366    /// keys or rows in a batch. Nothing changes on error.
367    pub fn upsert(&mut self, batch: RecordBatch) -> Result<Changes, DataError> {
368        self.check(&batch)?;
369        let latest = self.latest_rows(&batch)?;
370        if latest.is_empty() {
371            return Ok(Changes::default());
372        }
373        let plan = self.plan(&batch, &latest)?;
374        let updated = plan.iter().filter(|c| c.1.is_some()).count();
375        let mut changes = Changes {
376            inserted: plan.len() - updated,
377            updated,
378            deleted: 0,
379        };
380        let rebuild = plan.len().saturating_mul(REBUILD_RATIO) > self.len() + changes.inserted;
381
382        self.batches.push(batch.clone());
383        if !rebuild {
384            self.drop_from_views(plan.iter().filter_map(|c| Some((c.0, c.1?))));
385        }
386        for &(id, old, at, key) in &plan {
387            self.slots[id as usize] = Some(at);
388            self.dead += usize::from(old.is_some());
389            if let Some(key) = key {
390                self.keys.insert(key, id);
391            }
392        }
393        // Earlier rows of keys repeated within the batch were never live.
394        self.dead += batch.num_rows() - plan.len();
395        self.add_to_views(rebuild, plan.iter().map(|c| (c.0, c.2)));
396        changes.deleted = self.expire();
397        self.compact_if_due();
398        Ok(changes)
399    }
400
401    /// Deletes the rows of `keys` (an array of the key column's kind; nulls and unknown keys
402    /// are skipped). Views follow.
403    ///
404    /// ```
405    /// # use std::sync::Arc;
406    /// # use densa_data::arrow_array::{ArrayRef, Int64Array, RecordBatch};
407    /// # use densa_data::LiveTable;
408    /// # let ids: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));
409    /// # let rows = RecordBatch::try_from_iter([("id", ids)])?;
410    /// # let mut table = LiveTable::new(rows.schema(), "id")?;
411    /// # table.upsert(rows)?;
412    /// // Keys go in an Arrow array of the key column's type: here an Int64 key.
413    /// table.delete(&Int64Array::from(vec![1, 3]))?;
414    /// assert_eq!(table.len(), 1);
415    /// # Ok::<(), Box<dyn std::error::Error>>(())
416    /// ```
417    ///
418    /// # Errors
419    /// [`DataError::KeyType`] if `keys` is not of the key column's kind: nothing changes.
420    pub fn delete(&mut self, keys: &dyn Array) -> Result<Changes, DataError> {
421        let read = key_reader(keys.data_type())
422            .filter(|_| self.keys.same_kind(keys.data_type()))
423            .ok_or(DataError::KeyType)?;
424        let mut deleted = 0;
425        for row in 0..keys.len() {
426            let Some(id) = read(keys, row).and_then(|key| self.keys.remove(key)) else {
427                continue;
428            };
429            if let Some(slot) = self.slots[id as usize] {
430                self.drop_from_views(std::iter::once((id, slot)));
431            }
432            self.release(id);
433            deleted += 1;
434        }
435        self.compact_if_due();
436        Ok(Changes {
437            deleted,
438            ..Changes::default()
439        })
440    }
441
442    /// Rewrites the stored batches with only the live rows, in batches of up to 65,536 rows
443    /// (fewer where one batch cannot hold them: more distinct strings than a dictionary
444    /// column's key type indexes, over 2 GiB of strings). Dictionary columns of strings are
445    /// re-encoded with the strings their rows use. Views keep their order;
446    /// [`generation`](Self::generation) increases.
447    ///
448    /// # Errors
449    /// [`DataError::Arrow`] if Arrow cannot build a batch of even 128 of the rows; nothing
450    /// changes.
451    pub fn compact(&mut self) -> Result<(), DataError> {
452        let live: Vec<(RowId, Slot)> = self
453            .slots
454            .iter()
455            .enumerate()
456            .filter_map(|(id, slot)| Some((RowId::try_from(id).ok()?, (*slot)?)))
457            .collect();
458        let mut compacted = Vec::new();
459        for chunk in live.chunks(COMPACT_ROWS) {
460            self.gather_into(chunk, &mut compacted)?;
461        }
462        // The live rows, in order, fill the new batches in order.
463        let places = compacted.iter().enumerate().flat_map(|(batch, rows)| {
464            // Lossless: at most 2^32 live rows, so as many batches or rows in one.
465            (0..rows.num_rows()).map(move |row| Slot {
466                batch: batch as u32,
467                row: row as u32,
468            })
469        });
470        for (&(id, _), slot) in live.iter().zip(places) {
471            self.slots[id as usize] = Some(slot);
472        }
473        self.compacted_batches = compacted.len();
474        self.batches = compacted;
475        self.dead = 0;
476        self.generation += 1;
477        self.compaction_error = None;
478        self.retry_at_dead = MIN_DEAD_ROWS;
479        Ok(())
480    }
481
482    /// Keeps the table to the rows within `window` from now on, deleting the rows already
483    /// past it; `None` keeps every row again (rows already deleted stay deleted). Views
484    /// follow.
485    ///
486    /// # Errors
487    /// [`DataError::UnknownColumn`], [`DataError::UnsupportedColumn`], or
488    /// [`DataError::WindowSpan`] if the span does not fit the column: nothing changes.
489    pub fn set_window(&mut self, window: Option<Window>) -> Result<Changes, DataError> {
490        let Some(window) = window else {
491            self.window = None;
492            return Ok(Changes::default());
493        };
494        let mut state = WindowState::compile(&window, &self.schema)?;
495        state.rows.rebuild(Rows {
496            batches: &self.batches,
497            slots: &self.slots,
498        });
499        self.window = Some(state);
500        let deleted = self.expire();
501        self.compact_if_due();
502        Ok(Changes {
503            deleted,
504            ..Changes::default()
505        })
506    }
507
508    /// The table's window, if it has one.
509    pub fn window(&self) -> Option<&Window> {
510        self.window.as_ref().map(WindowState::spec)
511    }
512
513    /// Adds a view of the rows matching `spec`, filled from the current rows.
514    ///
515    /// # Errors
516    /// [`DataError::UnknownColumn`], [`DataError::UnsupportedColumn`] or
517    /// [`DataError::LiteralMismatch`] if the spec does not fit the schema.
518    pub fn add_view(&mut self, spec: ViewSpec) -> Result<ViewId, DataError> {
519        let mut view = ViewState::compile(&spec, &self.schema)?;
520        view.rebuild(Rows {
521            batches: &self.batches,
522            slots: &self.slots,
523        });
524        let (serial, index) = store(&mut self.views, view);
525        Ok(ViewId { serial, index })
526    }
527
528    /// Removes a view; `false` if it was already gone or belongs to another table. Its id
529    /// never names another view (its storage may be reused).
530    pub fn remove_view(&mut self, id: ViewId) -> bool {
531        take(&mut self.views, id.index, id.serial)
532    }
533
534    /// The rows of view `id`, in order; `None` if it was removed or belongs to another
535    /// table.
536    pub fn view(&self, id: ViewId) -> Option<View<'_>> {
537        let state = find(&self.views, id.index, id.serial)?;
538        Some(View {
539            state,
540            batches: &self.batches,
541            slots: &self.slots,
542            reuses: &self.reuses,
543        })
544    }
545
546    /// Adds grouped aggregates of the rows matching `spec`, filled from the current rows.
547    ///
548    /// # Errors
549    /// [`DataError::UnknownColumn`] or [`DataError::UnsupportedColumn`] for a key or an
550    /// aggregated column, [`DataError::NotNumeric`] for a sum or mean of a column that is
551    /// not numeric, [`DataError::DuplicateName`] if two result columns share a name,
552    /// [`DataError::LiteralMismatch`] for a filter that does not fit.
553    pub fn add_groups(&mut self, spec: GroupSpec) -> Result<GroupsId, DataError> {
554        let mut groups = GroupState::compile(&spec, &self.schema)?;
555        groups.rebuild(Rows {
556            batches: &self.batches,
557            slots: &self.slots,
558        });
559        let (serial, index) = store(&mut self.groups, groups);
560        Ok(GroupsId { serial, index })
561    }
562
563    /// Removes grouped aggregates; `false` if already gone or another table's. Their id
564    /// never names other aggregates (their storage may be reused).
565    pub fn remove_groups(&mut self, id: GroupsId) -> bool {
566        take(&mut self.groups, id.index, id.serial)
567    }
568
569    /// The groups of `id`, in key order; `None` if removed or another table's.
570    pub fn groups(&self, id: GroupsId) -> Option<Groups<'_>> {
571        let state = find(&self.groups, id.index, id.serial)?;
572        Some(Groups {
573            state,
574            rows: Rows {
575                batches: &self.batches,
576                slots: &self.slots,
577            },
578        })
579    }
580
581    fn check(&self, batch: &RecordBatch) -> Result<(), DataError> {
582        let ours = self.schema.fields();
583        let theirs = batch.schema_ref().fields();
584        if ours.len() != theirs.len() {
585            return Err(DataError::ColumnCount {
586                expected: ours.len(),
587                found: theirs.len(),
588            });
589        }
590        match ours
591            .iter()
592            .zip(theirs.iter())
593            .position(|(a, b)| a.name() != b.name() || a.data_type() != b.data_type())
594        {
595            Some(column) => Err(DataError::ColumnMismatch { column }),
596            None => Ok(()),
597        }
598    }
599
600    /// The last row of each distinct key in `batch`, in row order.
601    fn latest_rows<'b>(
602        &self,
603        batch: &'b RecordBatch,
604    ) -> Result<Vec<(usize, KeyRef<'b>)>, DataError> {
605        let keys = batch.column(self.key_column).as_ref();
606        let mut last: FxHashMap<KeyRef<'b>, usize> =
607            FxHashMap::with_capacity_and_hasher(batch.num_rows(), Default::default());
608        for row in 0..batch.num_rows() {
609            let key = (self.key_read)(keys, row).ok_or(DataError::NullKey { row })?;
610            last.insert(key, row);
611        }
612        let mut rows: Vec<(usize, KeyRef<'b>)> = last.into_iter().map(|(k, r)| (r, k)).collect();
613        rows.sort_unstable_by_key(|&(row, _)| row);
614        Ok(rows)
615    }
616
617    /// Each changed row's id, old and new place and new key. Checks every limit before
618    /// taking any id, so an error leaves the table unchanged.
619    fn plan<'b>(
620        &mut self,
621        batch: &RecordBatch,
622        latest: &[(usize, KeyRef<'b>)],
623    ) -> Result<Vec<Change<'b>>, DataError> {
624        let segment = u32::try_from(self.batches.len()).map_err(|_| DataError::TooManyRows)?;
625        u32::try_from(batch.num_rows()).map_err(|_| DataError::TooManyRows)?;
626        let new_keys = latest
627            .iter()
628            .filter(|&&(_, key)| self.keys.get(key).is_none())
629            .count();
630        if new_keys > self.free.len() + (MAX_ROWS - self.slots.len()) {
631            return Err(DataError::TooManyRows);
632        }
633        let mut plan = Vec::with_capacity(latest.len());
634        for &(row, key) in latest {
635            // Lossless: the batch has fewer than 2^32 rows (checked above).
636            let at = Slot {
637                batch: segment,
638                row: row as u32,
639            };
640            plan.push(match self.keys.get(key) {
641                Some(id) => (id, self.slots[id as usize], at, None),
642                None => (self.allocate_id(), None, at, Some(key)),
643            });
644        }
645        Ok(plan)
646    }
647
648    /// A free row id; [`plan`](Self::plan) checked that one exists.
649    fn allocate_id(&mut self) -> RowId {
650        if let Some(id) = self.free.pop() {
651            return id;
652        }
653        self.slots.push(None);
654        // Lossless: `plan` keeps the slot count at most 2^32.
655        (self.slots.len() - 1) as RowId
656    }
657
658    /// Copies the live rows of `chunk` into new batches at the end of `out`: one if Arrow can
659    /// build it, otherwise each half on its own, recursively, down to [`MIN_SPLIT_ROWS`]:
660    /// fewer rows use fewer distinct strings, so they fit a narrow dictionary key type.
661    fn gather_into(
662        &self,
663        chunk: &[(RowId, Slot)],
664        out: &mut Vec<RecordBatch>,
665    ) -> Result<(), DataError> {
666        match self.gather(chunk) {
667            Ok(batch) => {
668                out.push(batch);
669                Ok(())
670            }
671            Err(_) if chunk.len() > MIN_SPLIT_ROWS => {
672                let (first, second) = chunk.split_at(chunk.len() / 2);
673                self.gather_into(first, out)?;
674                self.gather_into(second, out)
675            }
676            Err(error) => Err(error),
677        }
678    }
679
680    /// The live rows of `chunk`, copied into one batch.
681    fn gather(&self, chunk: &[(RowId, Slot)]) -> Result<RecordBatch, DataError> {
682        // Only the batches the chunk reads go to Arrow, which may size a dictionary column by
683        // all of its inputs' values.
684        let mut position: FxHashMap<usize, usize> = FxHashMap::default();
685        let mut sources = Vec::new();
686        let indices: Vec<(usize, usize)> = chunk
687            .iter()
688            .map(|(_, slot)| {
689                let batch = slot.batch as usize;
690                let at = *position.entry(batch).or_insert_with(|| {
691                    sources.push(batch);
692                    sources.len() - 1
693                });
694                (at, slot.row as usize)
695            })
696            .collect();
697        let columns = (0..self.schema.fields().len())
698            .map(|column| {
699                let arrays: Vec<&dyn Array> = sources
700                    .iter()
701                    .map(|&batch| self.batches[batch].column(column).as_ref())
702                    .collect();
703                interleave_column(&arrays, &indices)
704            })
705            .collect::<Result<Vec<ArrayRef>, _>>()
706            .map_err(DataError::arrow)?;
707        RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(DataError::arrow)
708    }
709
710    /// Deletes the rows past the window; how many.
711    fn expire(&mut self) -> usize {
712        let Some(window) = &self.window else {
713            return 0;
714        };
715        let count = window.expired(Rows {
716            batches: &self.batches,
717            slots: &self.slots,
718        });
719        if count == 0 {
720            return 0;
721        }
722        let gone: Vec<(RowId, Slot)> = (0..count)
723            .filter_map(|index| {
724                let id = window.rows.get(index)?;
725                Some((id, self.slots[id as usize]?))
726            })
727            .collect();
728        let rebuild = count.saturating_mul(REBUILD_RATIO) > self.len();
729        if !rebuild {
730            self.drop_from_views(gone.iter().copied());
731        }
732        for &(id, slot) in &gone {
733            let at = slot.row_ref();
734            let keys = self.batches[at.batch].column(self.key_column).as_ref();
735            if let Some(key) = (self.key_read)(keys, at.row) {
736                self.keys.remove(key);
737            }
738            self.release(id);
739        }
740        if rebuild {
741            self.add_to_views(true, std::iter::empty());
742        }
743        gone.len()
744    }
745
746    /// Frees row `id`, already out of every view and group: its key's record ends.
747    fn release(&mut self, id: RowId) {
748        self.slots[id as usize] = None;
749        self.free.push(id);
750        if self.reuses.len() <= id as usize {
751            self.reuses.resize(id as usize + 1, 0);
752        }
753        // Wrapping: an id freed 2^32 times could give an old record back; harmless odds.
754        self.reuses[id as usize] = self.reuses[id as usize].wrapping_add(1);
755        self.dead += 1;
756    }
757
758    /// Removes rows from every view, group and the window, each at the place its values were
759    /// when it entered.
760    fn drop_from_views(&mut self, rows: impl Iterator<Item = (RowId, Slot)> + Clone) {
761        let source = Rows {
762            batches: &self.batches,
763            slots: &self.slots,
764        };
765        for (_, view) in self.views.iter_mut().flatten() {
766            for (id, slot) in rows.clone() {
767                view.remove(source, id, slot.row_ref());
768            }
769        }
770        for (_, groups) in self.groups.iter_mut().flatten() {
771            for (id, slot) in rows.clone() {
772                groups.remove(source, id, slot.row_ref());
773            }
774        }
775        if let Some(window) = &mut self.window {
776            for (id, slot) in rows {
777                window.rows.remove(source, id, slot.row_ref());
778            }
779        }
780    }
781
782    /// Adds changed rows to every view, group and the window, or rebuilds them.
783    fn add_to_views(&mut self, rebuild: bool, rows: impl Iterator<Item = (RowId, Slot)> + Clone) {
784        let source = Rows {
785            batches: &self.batches,
786            slots: &self.slots,
787        };
788        for (_, view) in self.views.iter_mut().flatten() {
789            if rebuild {
790                view.rebuild(source);
791            } else {
792                for (id, slot) in rows.clone() {
793                    view.insert(source, id, slot.row_ref());
794                }
795            }
796        }
797        for (_, groups) in self.groups.iter_mut().flatten() {
798            if rebuild {
799                groups.rebuild(source);
800            } else {
801                for (id, slot) in rows.clone() {
802                    groups.insert(source, id, slot.row_ref());
803                }
804            }
805        }
806        if let Some(window) = &mut self.window {
807            if rebuild {
808                window.rows.rebuild(source);
809            } else {
810                for (id, slot) in rows {
811                    window.rows.insert(source, id, slot.row_ref());
812                }
813            }
814        }
815    }
816
817    /// Compacts when due and enabled. A failure is recorded, not returned: the change that
818    /// triggered it is applied and the table stays correct, only larger.
819    fn compact_if_due(&mut self) {
820        if !self.options.auto_compact {
821            return;
822        }
823        let dead_majority = self.dead >= self.retry_at_dead && self.dead > self.len();
824        // Counting only the batches stored since the last compaction: one that had to write
825        // many small batches does not run again at every change.
826        let added = self.batches.len().saturating_sub(self.compacted_batches);
827        let too_many_batches = added > MAX_BATCHES && self.compaction_error.is_none();
828        if !(dead_majority || too_many_batches) {
829            return;
830        }
831        if let Err(error) = self.compact() {
832            self.compaction_error = Some(error);
833            self.retry_at_dead = self.dead.saturating_mul(2);
834        }
835    }
836}
837
838/// The entry at `index` if its serial is `serial`.
839fn find<T>(entries: &[Option<(u64, T)>], index: usize, serial: u64) -> Option<&T> {
840    match entries.get(index)? {
841        Some((s, entry)) if *s == serial => Some(entry),
842        _ => None,
843    }
844}
845
846/// Stores `entry` under a new serial in the first free place (removed entries leave one),
847/// so a model that replaces its view on every sort keeps one; returns the serial and index.
848fn store<T>(entries: &mut Vec<Option<(u64, T)>>, entry: T) -> (u64, usize) {
849    let serial = next_serial();
850    match entries.iter().position(Option::is_none) {
851        Some(index) => {
852            entries[index] = Some((serial, entry));
853            (serial, index)
854        }
855        None => {
856            entries.push(Some((serial, entry)));
857            (serial, entries.len() - 1)
858        }
859    }
860}
861
862/// Removes the entry at `index` if its serial is `serial`; whether it did.
863fn take<T>(entries: &mut [Option<(u64, T)>], index: usize, serial: u64) -> bool {
864    entries
865        .get_mut(index)
866        .and_then(|entry| entry.take_if(|(s, _)| *s == serial))
867        .is_some()
868}
869
870/// `schema` with every field nullable: superseded rows and compaction never need a column
871/// to be free of nulls, and batches may carry nulls the first batch's schema did not.
872fn nullable(schema: &Schema) -> SchemaRef {
873    let fields: Vec<Field> = schema
874        .fields()
875        .iter()
876        .map(|f| f.as_ref().clone().with_nullable(true))
877        .collect();
878    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
879}
880
881/// The rows of a view, in its order. Cheap to get; reading never allocates.
882///
883/// A row is a [`RowRef`]: which of [`LiveTable::batches`] holds it, and where. Read its
884/// values from that batch's columns:
885///
886/// ```
887/// use std::sync::Arc;
888/// use densa_data::arrow_array::cast::AsArray;
889/// use densa_data::arrow_array::{ArrayRef, Float64Array, RecordBatch, StringArray};
890/// use densa_data::{LiveTable, Order, ViewSpec};
891///
892/// let rows = RecordBatch::try_from_iter([
893///     ("symbol", Arc::new(StringArray::from(vec!["AAA", "BBB"])) as ArrayRef),
894///     ("price", Arc::new(Float64Array::from(vec![150.0, 90.0])) as ArrayRef),
895/// ])?;
896/// let mut table = LiveTable::new(rows.schema(), "symbol")?;
897/// let cheapest = table.add_view(ViewSpec::new().sort_by("price", Order::Ascending))?;
898/// table.upsert(rows)?;
899///
900/// let view = table.view(cheapest).unwrap();
901/// let symbols: Vec<&str> = (0..view.len())
902///     .filter_map(|i| view.row(i))
903///     .map(|at| {
904///         let batch = &table.batches()[at.batch];
905///         // The column's Arrow type, as the table's schema has it (Utf8 here).
906///         batch.column_by_name("symbol").unwrap().as_string::<i32>().value(at.row)
907///     })
908///     .collect();
909/// assert_eq!(symbols, ["BBB", "AAA"]);
910/// # Ok::<(), Box<dyn std::error::Error>>(())
911/// ```
912#[derive(Debug, Clone, Copy)]
913pub struct View<'a> {
914    state: &'a ViewState,
915    batches: &'a [RecordBatch],
916    slots: &'a [Option<Slot>],
917    reuses: &'a [u32],
918}
919
920impl View<'_> {
921    /// Rows in the view.
922    pub fn len(&self) -> usize {
923        self.state.len()
924    }
925
926    /// Whether the view is empty.
927    pub fn is_empty(&self) -> bool {
928        self.len() == 0
929    }
930
931    /// Where the row at `index` (0 is first) is. `O(log n)`.
932    pub fn row(&self, index: usize) -> Option<RowRef> {
933        let id = self.state.get(index)?;
934        self.slots
935            .get(id as usize)
936            .copied()
937            .flatten()
938            .map(Slot::row_ref)
939    }
940
941    /// The record of the row at `index`. `O(log n)`.
942    pub fn record(&self, index: usize) -> Option<RecordId> {
943        let id = self.state.get(index)?;
944        Some(RecordId::new(id, reuse(self.reuses, id)))
945    }
946
947    /// Where `record`'s row is in the view, or `None` if it is not in it (filtered out, or
948    /// its key deleted). `O(log n)`: binary searches comparing the row's sort values.
949    pub fn position(&self, record: RecordId) -> Option<usize> {
950        let (id, _) = live_slot(self.slots, self.reuses, record)?;
951        let rows = Rows {
952            batches: self.batches,
953            slots: self.slots,
954        };
955        self.state.position(rows, id)
956    }
957}
958
959/// How many times row id `id` was freed.
960fn reuse(reuses: &[u32], id: RowId) -> u32 {
961    reuses.get(id as usize).copied().unwrap_or(0)
962}
963
964/// `record`'s row id and slot, if the record's key is still live.
965fn live_slot(slots: &[Option<Slot>], reuses: &[u32], record: RecordId) -> Option<(RowId, Slot)> {
966    let (id, reused) = record.parts();
967    let slot = (*slots.get(id as usize)?)?;
968    (reuse(reuses, id) == reused).then_some((id, slot))
969}
970
971#[cfg(test)]
972mod tests {
973    use std::sync::Arc;
974
975    use arrow_array::{ArrayRef, Int64Array};
976
977    use super::*;
978
979    fn ids(keys: Vec<i64>) -> RecordBatch {
980        RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(keys)) as ArrayRef)]).unwrap()
981    }
982
983    #[test]
984    fn replacing_a_view_again_and_again_keeps_one_place() {
985        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
986        let mut view = table.add_view(ViewSpec::new()).unwrap();
987        for _ in 0..100 {
988            let next = table.add_view(ViewSpec::new()).unwrap();
989            table.remove_view(view);
990            view = next;
991        }
992        assert!(table.views.len() <= 2, "{} places", table.views.len());
993    }
994
995    #[test]
996    fn ids_of_deleted_keys_are_reused() {
997        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
998        for round in 0..100 {
999            table
1000                .upsert(ids((0..10).map(|k| round * 10 + k).collect()))
1001                .unwrap();
1002            table
1003                .delete(&Int64Array::from(
1004                    (0..10).map(|k| round * 10 + k).collect::<Vec<_>>(),
1005                ))
1006                .unwrap();
1007        }
1008        assert!(
1009            table.slots.len() <= 10,
1010            "{} ids for 10 live keys",
1011            table.slots.len()
1012        );
1013    }
1014
1015    #[test]
1016    fn a_failed_automatic_compaction_is_recorded_and_backs_off() {
1017        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
1018        table.dead = MIN_DEAD_ROWS;
1019        // A batch the table's schema cannot hold makes the next compaction fail.
1020        let other = RecordBatch::try_from_iter([(
1021            "other",
1022            Arc::new(arrow_array::StringArray::from(vec!["x"])) as ArrayRef,
1023        )])
1024        .unwrap();
1025        table.batches.push(other);
1026        table.slots.push(Some(Slot { batch: 0, row: 0 }));
1027        table.compact_if_due();
1028        assert!(table.compaction_error().is_some());
1029        assert_eq!(table.retry_at_dead, 2 * MIN_DEAD_ROWS);
1030    }
1031
1032    #[test]
1033    fn automatic_compaction_counts_the_batches_stored_since_the_last_one() {
1034        let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
1035        table.upsert(ids(vec![1])).unwrap();
1036        // As if the last compaction had to write 10 batches.
1037        table.compacted_batches = 10;
1038        while table.batches.len() < 10 + MAX_BATCHES {
1039            table.batches.push(ids(vec![]));
1040        }
1041        table.compact_if_due();
1042        assert_eq!(table.generation(), 0, "4,096 batches added");
1043        table.batches.push(ids(vec![]));
1044        table.compact_if_due();
1045        assert_eq!(table.generation(), 1, "4,097 batches added");
1046        assert_eq!(table.compacted_batches, 1);
1047    }
1048}