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
15const REBUILD_RATIO: usize = 8;
17const MIN_DEAD_ROWS: usize = 4_096;
20const MAX_BATCHES: usize = 4_096;
22const COMPACT_ROWS: usize = 65_536;
24const MIN_SPLIT_ROWS: usize = 128;
27const MAX_ROWS: usize = 1 << 32;
29
30static NEXT_SERIAL: AtomicU64 = AtomicU64::new(1);
32
33fn next_serial() -> u64 {
34 NEXT_SERIAL.fetch_add(1, AtomicOrdering::Relaxed)
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
42pub struct RowRef {
43 pub batch: usize,
45 pub row: usize,
47}
48
49#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
67#[non_exhaustive]
68pub struct Changes {
69 pub inserted: usize,
71 pub updated: usize,
73 pub deleted: usize,
75}
76
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79#[non_exhaustive]
80pub struct TableOptions {
81 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 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
107pub struct ViewId {
108 serial: u64,
110 index: usize,
111}
112
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
116pub struct GroupsId {
117 serial: u64,
118 index: usize,
119}
120
121#[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 #[must_use]
139 pub const fn to_bits(self) -> u64 {
140 self.0
141 }
142
143 #[must_use]
146 pub const fn from_bits(bits: u64) -> Self {
147 Self(bits)
148 }
149
150 fn parts(self) -> (RowId, u32) {
152 (self.0 as u32, (self.0 >> 32) as u32)
154 }
155}
156
157type Change<'a> = (RowId, Option<Slot>, Slot, Option<KeyRef<'a>>);
160
161#[derive(Debug)]
179pub struct LiveTable {
180 instance: u64,
182 options: TableOptions,
183 schema: SchemaRef,
184 key_column: usize,
185 key_read: KeyFn,
186 keys: KeyIndex,
187 batches: Vec<RecordBatch>,
188 slots: Vec<Option<Slot>>,
190 free: Vec<RowId>,
191 reuses: Vec<u32>,
194 dead: usize,
196 views: Vec<Option<(u64, ViewState)>>,
198 groups: Vec<Option<(u64, GroupState)>>,
199 generation: u64,
200 compaction_error: Option<DataError>,
202 retry_at_dead: usize,
204 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 pub fn new(schema: SchemaRef, key: &str) -> Result<Self, DataError> {
241 Self::with_options(schema, key, TableOptions::default())
242 }
243
244 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 pub fn schema(&self) -> &SchemaRef {
283 &self.schema
284 }
285
286 pub fn instance(&self) -> u64 {
290 self.instance
291 }
292
293 pub fn len(&self) -> usize {
295 self.keys.len()
296 }
297
298 pub fn is_empty(&self) -> bool {
300 self.len() == 0
301 }
302
303 pub fn batches(&self) -> &[RecordBatch] {
307 &self.batches
308 }
309
310 pub(crate) fn live_rows(&self) -> impl Iterator<Item = RowRef> + '_ {
312 self.slots.iter().flatten().map(|slot| slot.row_ref())
313 }
314
315 pub fn generation(&self) -> u64 {
318 self.generation
319 }
320
321 pub fn compaction_error(&self) -> Option<&DataError> {
325 self.compaction_error.as_ref()
326 }
327
328 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 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 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 pub(crate) fn id_of(&self, key: KeyRef<'_>) -> Option<RowId> {
347 self.keys.get(key)
348 }
349
350 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 pub(crate) fn live_ids(&self) -> impl Iterator<Item = (RowId, RowRef)> + '_ {
358 self.slots
360 .iter()
361 .enumerate()
362 .filter_map(|(id, slot)| Some((id as RowId, (*slot)?.row_ref())))
363 }
364
365 pub(crate) fn key_column(&self) -> usize {
367 self.key_column
368 }
369
370 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 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 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 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 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 let places = compacted.iter().enumerate().flat_map(|(batch, rows)| {
489 (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 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 pub fn window(&self) -> Option<&Window> {
535 self.window.as_ref().map(WindowState::spec)
536 }
537
538 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 pub fn remove_view(&mut self, id: ViewId) -> bool {
556 take(&mut self.views, id.index, id.serial)
557 }
558
559 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 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 pub fn remove_groups(&mut self, id: GroupsId) -> bool {
591 take(&mut self.groups, id.index, id.serial)
592 }
593
594 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 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 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 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 fn allocate_id(&mut self) -> RowId {
675 if let Some(id) = self.free.pop() {
676 return id;
677 }
678 self.slots.push(None);
679 (self.slots.len() - 1) as RowId
681 }
682
683 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 fn gather(&self, chunk: &[(RowId, Slot)]) -> Result<RecordBatch, DataError> {
707 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 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 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 self.reuses[id as usize] = self.reuses[id as usize].wrapping_add(1);
780 self.dead += 1;
781 }
782
783 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 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 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 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
863fn 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
871fn 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
887fn 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
895fn 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#[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 pub fn len(&self) -> usize {
948 self.state.len()
949 }
950
951 pub fn is_empty(&self) -> bool {
953 self.len() == 0
954 }
955
956 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 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 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
984fn reuse(reuses: &[u32], id: RowId) -> u32 {
986 reuses.get(id as usize).copied().unwrap_or(0)
987}
988
989fn 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 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 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}