use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
use std::sync::Arc;
use arrow_array::{Array, ArrayRef, RecordBatch};
use arrow_schema::{Field, Schema, SchemaRef};
use rustc_hash::FxHashMap;
use crate::dictionary::interleave_column;
use crate::group::{GroupState, Groups};
use crate::key::{key_reader, KeyFn, KeyIndex, KeyRef};
use crate::view::{Rows, ViewState};
use crate::window::WindowState;
use crate::{DataError, GroupSpec, RowId, Scalar, ViewSpec, Window};
const REBUILD_RATIO: usize = 8;
const MIN_DEAD_ROWS: usize = 4_096;
const MAX_BATCHES: usize = 4_096;
const COMPACT_ROWS: usize = 65_536;
const MIN_SPLIT_ROWS: usize = 128;
const MAX_ROWS: usize = 1 << 32;
static NEXT_SERIAL: AtomicU64 = AtomicU64::new(1);
fn next_serial() -> u64 {
NEXT_SERIAL.fetch_add(1, AtomicOrdering::Relaxed)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct RowRef {
pub batch: usize,
pub row: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Slot {
batch: u32,
row: u32,
}
impl Slot {
pub(crate) fn row_ref(self) -> RowRef {
RowRef {
batch: self.batch as usize,
row: self.row as usize,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct Changes {
pub inserted: usize,
pub updated: usize,
pub deleted: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct TableOptions {
pub auto_compact: bool,
}
impl Default for TableOptions {
fn default() -> Self {
Self { auto_compact: true }
}
}
impl TableOptions {
#[must_use]
pub fn with_auto_compact(mut self, auto_compact: bool) -> Self {
self.auto_compact = auto_compact;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ViewId {
serial: u64,
index: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct GroupsId {
serial: u64,
index: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct RecordId(u64);
impl RecordId {
fn new(id: RowId, reuse: u32) -> Self {
Self(u64::from(reuse) << 32 | u64::from(id))
}
#[must_use]
pub const fn to_bits(self) -> u64 {
self.0
}
#[must_use]
pub const fn from_bits(bits: u64) -> Self {
Self(bits)
}
fn parts(self) -> (RowId, u32) {
(self.0 as u32, (self.0 >> 32) as u32)
}
}
type Change<'a> = (RowId, Option<Slot>, Slot, Option<KeyRef<'a>>);
#[derive(Debug)]
pub struct LiveTable {
instance: u64,
options: TableOptions,
schema: SchemaRef,
key_column: usize,
key_read: KeyFn,
keys: KeyIndex,
batches: Vec<RecordBatch>,
slots: Vec<Option<Slot>>,
free: Vec<RowId>,
reuses: Vec<u32>,
dead: usize,
views: Vec<Option<(u64, ViewState)>>,
groups: Vec<Option<(u64, GroupState)>>,
generation: u64,
compaction_error: Option<DataError>,
retry_at_dead: usize,
compacted_batches: usize,
window: Option<WindowState>,
}
impl Clone for LiveTable {
fn clone(&self) -> Self {
Self {
instance: next_serial(),
options: self.options,
schema: Arc::clone(&self.schema),
key_column: self.key_column,
key_read: self.key_read,
keys: self.keys.clone(),
batches: self.batches.clone(),
slots: self.slots.clone(),
free: self.free.clone(),
reuses: self.reuses.clone(),
dead: self.dead,
views: self.views.clone(),
groups: self.groups.clone(),
generation: self.generation,
compaction_error: self.compaction_error.clone(),
retry_at_dead: self.retry_at_dead,
compacted_batches: self.compacted_batches,
window: self.window.clone(),
}
}
}
impl LiveTable {
pub fn new(schema: SchemaRef, key: &str) -> Result<Self, DataError> {
Self::with_options(schema, key, TableOptions::default())
}
pub fn with_options(
schema: SchemaRef,
key: &str,
options: TableOptions,
) -> Result<Self, DataError> {
let (key_column, field) = schema
.column_with_name(key)
.ok_or_else(|| DataError::UnknownColumn(key.to_owned()))?;
let unsupported = || DataError::UnsupportedKey(key.to_owned());
let key_read = key_reader(field.data_type()).ok_or_else(unsupported)?;
let keys = KeyIndex::for_type(field.data_type()).ok_or_else(unsupported)?;
Ok(Self {
instance: next_serial(),
options,
schema: nullable(&schema),
key_column,
key_read,
keys,
batches: Vec::new(),
slots: Vec::new(),
free: Vec::new(),
reuses: Vec::new(),
dead: 0,
views: Vec::new(),
groups: Vec::new(),
generation: 0,
compaction_error: None,
retry_at_dead: MIN_DEAD_ROWS,
compacted_batches: 0,
window: None,
})
}
pub fn schema(&self) -> &SchemaRef {
&self.schema
}
pub fn instance(&self) -> u64 {
self.instance
}
pub fn len(&self) -> usize {
self.keys.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn batches(&self) -> &[RecordBatch] {
&self.batches
}
pub(crate) fn live_rows(&self) -> impl Iterator<Item = RowRef> + '_ {
self.slots.iter().flatten().map(|slot| slot.row_ref())
}
pub fn generation(&self) -> u64 {
self.generation
}
pub fn compaction_error(&self) -> Option<&DataError> {
self.compaction_error.as_ref()
}
pub fn find(&self, key: &Scalar) -> Option<RowRef> {
let id = self.key_id(key)?;
self.slots[id as usize].map(Slot::row_ref)
}
pub fn record_of_key(&self, key: &Scalar) -> Option<RecordId> {
let id = self.key_id(key)?;
Some(RecordId::new(id, reuse(&self.reuses, id)))
}
pub fn row_of(&self, record: RecordId) -> Option<RowRef> {
live_slot(&self.slots, &self.reuses, record).map(|(_, slot)| slot.row_ref())
}
fn key_id(&self, key: &Scalar) -> Option<RowId> {
let key = match (key, &self.keys) {
(Scalar::Int(k), KeyIndex::Int(_)) => KeyRef::Int(*k),
(Scalar::Int(k), KeyIndex::UInt(_)) => KeyRef::UInt(u64::try_from(*k).ok()?),
(Scalar::UInt(k), KeyIndex::UInt(_)) => KeyRef::UInt(*k),
(Scalar::UInt(k), KeyIndex::Int(_)) => KeyRef::Int(i64::try_from(*k).ok()?),
(Scalar::Str(k), KeyIndex::Str(_)) => KeyRef::Str(k),
_ => return None,
};
self.keys.get(key)
}
pub fn upsert(&mut self, batch: RecordBatch) -> Result<Changes, DataError> {
self.check(&batch)?;
let latest = self.latest_rows(&batch)?;
if latest.is_empty() {
return Ok(Changes::default());
}
let plan = self.plan(&batch, &latest)?;
let updated = plan.iter().filter(|c| c.1.is_some()).count();
let mut changes = Changes {
inserted: plan.len() - updated,
updated,
deleted: 0,
};
let rebuild = plan.len().saturating_mul(REBUILD_RATIO) > self.len() + changes.inserted;
self.batches.push(batch.clone());
if !rebuild {
self.drop_from_views(plan.iter().filter_map(|c| Some((c.0, c.1?))));
}
for &(id, old, at, key) in &plan {
self.slots[id as usize] = Some(at);
self.dead += usize::from(old.is_some());
if let Some(key) = key {
self.keys.insert(key, id);
}
}
self.dead += batch.num_rows() - plan.len();
self.add_to_views(rebuild, plan.iter().map(|c| (c.0, c.2)));
changes.deleted = self.expire();
self.compact_if_due();
Ok(changes)
}
pub fn delete(&mut self, keys: &dyn Array) -> Result<Changes, DataError> {
let read = key_reader(keys.data_type())
.filter(|_| self.keys.same_kind(keys.data_type()))
.ok_or(DataError::KeyType)?;
let mut deleted = 0;
for row in 0..keys.len() {
let Some(id) = read(keys, row).and_then(|key| self.keys.remove(key)) else {
continue;
};
if let Some(slot) = self.slots[id as usize] {
self.drop_from_views(std::iter::once((id, slot)));
}
self.release(id);
deleted += 1;
}
self.compact_if_due();
Ok(Changes {
deleted,
..Changes::default()
})
}
pub fn compact(&mut self) -> Result<(), DataError> {
let live: Vec<(RowId, Slot)> = self
.slots
.iter()
.enumerate()
.filter_map(|(id, slot)| Some((RowId::try_from(id).ok()?, (*slot)?)))
.collect();
let mut compacted = Vec::new();
for chunk in live.chunks(COMPACT_ROWS) {
self.gather_into(chunk, &mut compacted)?;
}
let places = compacted.iter().enumerate().flat_map(|(batch, rows)| {
(0..rows.num_rows()).map(move |row| Slot {
batch: batch as u32,
row: row as u32,
})
});
for (&(id, _), slot) in live.iter().zip(places) {
self.slots[id as usize] = Some(slot);
}
self.compacted_batches = compacted.len();
self.batches = compacted;
self.dead = 0;
self.generation += 1;
self.compaction_error = None;
self.retry_at_dead = MIN_DEAD_ROWS;
Ok(())
}
pub fn set_window(&mut self, window: Option<Window>) -> Result<Changes, DataError> {
let Some(window) = window else {
self.window = None;
return Ok(Changes::default());
};
let mut state = WindowState::compile(&window, &self.schema)?;
state.rows.rebuild(Rows {
batches: &self.batches,
slots: &self.slots,
});
self.window = Some(state);
let deleted = self.expire();
self.compact_if_due();
Ok(Changes {
deleted,
..Changes::default()
})
}
pub fn window(&self) -> Option<&Window> {
self.window.as_ref().map(WindowState::spec)
}
pub fn add_view(&mut self, spec: ViewSpec) -> Result<ViewId, DataError> {
let mut view = ViewState::compile(&spec, &self.schema)?;
view.rebuild(Rows {
batches: &self.batches,
slots: &self.slots,
});
let (serial, index) = store(&mut self.views, view);
Ok(ViewId { serial, index })
}
pub fn remove_view(&mut self, id: ViewId) -> bool {
take(&mut self.views, id.index, id.serial)
}
pub fn view(&self, id: ViewId) -> Option<View<'_>> {
let state = find(&self.views, id.index, id.serial)?;
Some(View {
state,
batches: &self.batches,
slots: &self.slots,
reuses: &self.reuses,
})
}
pub fn add_groups(&mut self, spec: GroupSpec) -> Result<GroupsId, DataError> {
let mut groups = GroupState::compile(&spec, &self.schema)?;
groups.rebuild(Rows {
batches: &self.batches,
slots: &self.slots,
});
let (serial, index) = store(&mut self.groups, groups);
Ok(GroupsId { serial, index })
}
pub fn remove_groups(&mut self, id: GroupsId) -> bool {
take(&mut self.groups, id.index, id.serial)
}
pub fn groups(&self, id: GroupsId) -> Option<Groups<'_>> {
let state = find(&self.groups, id.index, id.serial)?;
Some(Groups {
state,
rows: Rows {
batches: &self.batches,
slots: &self.slots,
},
})
}
fn check(&self, batch: &RecordBatch) -> Result<(), DataError> {
let ours = self.schema.fields();
let theirs = batch.schema_ref().fields();
if ours.len() != theirs.len() {
return Err(DataError::ColumnCount {
expected: ours.len(),
found: theirs.len(),
});
}
match ours
.iter()
.zip(theirs.iter())
.position(|(a, b)| a.name() != b.name() || a.data_type() != b.data_type())
{
Some(column) => Err(DataError::ColumnMismatch { column }),
None => Ok(()),
}
}
fn latest_rows<'b>(
&self,
batch: &'b RecordBatch,
) -> Result<Vec<(usize, KeyRef<'b>)>, DataError> {
let keys = batch.column(self.key_column).as_ref();
let mut last: FxHashMap<KeyRef<'b>, usize> =
FxHashMap::with_capacity_and_hasher(batch.num_rows(), Default::default());
for row in 0..batch.num_rows() {
let key = (self.key_read)(keys, row).ok_or(DataError::NullKey { row })?;
last.insert(key, row);
}
let mut rows: Vec<(usize, KeyRef<'b>)> = last.into_iter().map(|(k, r)| (r, k)).collect();
rows.sort_unstable_by_key(|&(row, _)| row);
Ok(rows)
}
fn plan<'b>(
&mut self,
batch: &RecordBatch,
latest: &[(usize, KeyRef<'b>)],
) -> Result<Vec<Change<'b>>, DataError> {
let segment = u32::try_from(self.batches.len()).map_err(|_| DataError::TooManyRows)?;
u32::try_from(batch.num_rows()).map_err(|_| DataError::TooManyRows)?;
let new_keys = latest
.iter()
.filter(|&&(_, key)| self.keys.get(key).is_none())
.count();
if new_keys > self.free.len() + (MAX_ROWS - self.slots.len()) {
return Err(DataError::TooManyRows);
}
let mut plan = Vec::with_capacity(latest.len());
for &(row, key) in latest {
let at = Slot {
batch: segment,
row: row as u32,
};
plan.push(match self.keys.get(key) {
Some(id) => (id, self.slots[id as usize], at, None),
None => (self.allocate_id(), None, at, Some(key)),
});
}
Ok(plan)
}
fn allocate_id(&mut self) -> RowId {
if let Some(id) = self.free.pop() {
return id;
}
self.slots.push(None);
(self.slots.len() - 1) as RowId
}
fn gather_into(
&self,
chunk: &[(RowId, Slot)],
out: &mut Vec<RecordBatch>,
) -> Result<(), DataError> {
match self.gather(chunk) {
Ok(batch) => {
out.push(batch);
Ok(())
}
Err(_) if chunk.len() > MIN_SPLIT_ROWS => {
let (first, second) = chunk.split_at(chunk.len() / 2);
self.gather_into(first, out)?;
self.gather_into(second, out)
}
Err(error) => Err(error),
}
}
fn gather(&self, chunk: &[(RowId, Slot)]) -> Result<RecordBatch, DataError> {
let mut position: FxHashMap<usize, usize> = FxHashMap::default();
let mut sources = Vec::new();
let indices: Vec<(usize, usize)> = chunk
.iter()
.map(|(_, slot)| {
let batch = slot.batch as usize;
let at = *position.entry(batch).or_insert_with(|| {
sources.push(batch);
sources.len() - 1
});
(at, slot.row as usize)
})
.collect();
let columns = (0..self.schema.fields().len())
.map(|column| {
let arrays: Vec<&dyn Array> = sources
.iter()
.map(|&batch| self.batches[batch].column(column).as_ref())
.collect();
interleave_column(&arrays, &indices)
})
.collect::<Result<Vec<ArrayRef>, _>>()
.map_err(DataError::arrow)?;
RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(DataError::arrow)
}
fn expire(&mut self) -> usize {
let Some(window) = &self.window else {
return 0;
};
let count = window.expired(Rows {
batches: &self.batches,
slots: &self.slots,
});
if count == 0 {
return 0;
}
let gone: Vec<(RowId, Slot)> = (0..count)
.filter_map(|index| {
let id = window.rows.get(index)?;
Some((id, self.slots[id as usize]?))
})
.collect();
let rebuild = count.saturating_mul(REBUILD_RATIO) > self.len();
if !rebuild {
self.drop_from_views(gone.iter().copied());
}
for &(id, slot) in &gone {
let at = slot.row_ref();
let keys = self.batches[at.batch].column(self.key_column).as_ref();
if let Some(key) = (self.key_read)(keys, at.row) {
self.keys.remove(key);
}
self.release(id);
}
if rebuild {
self.add_to_views(true, std::iter::empty());
}
gone.len()
}
fn release(&mut self, id: RowId) {
self.slots[id as usize] = None;
self.free.push(id);
if self.reuses.len() <= id as usize {
self.reuses.resize(id as usize + 1, 0);
}
self.reuses[id as usize] = self.reuses[id as usize].wrapping_add(1);
self.dead += 1;
}
fn drop_from_views(&mut self, rows: impl Iterator<Item = (RowId, Slot)> + Clone) {
let source = Rows {
batches: &self.batches,
slots: &self.slots,
};
for (_, view) in self.views.iter_mut().flatten() {
for (id, slot) in rows.clone() {
view.remove(source, id, slot.row_ref());
}
}
for (_, groups) in self.groups.iter_mut().flatten() {
for (id, slot) in rows.clone() {
groups.remove(source, id, slot.row_ref());
}
}
if let Some(window) = &mut self.window {
for (id, slot) in rows {
window.rows.remove(source, id, slot.row_ref());
}
}
}
fn add_to_views(&mut self, rebuild: bool, rows: impl Iterator<Item = (RowId, Slot)> + Clone) {
let source = Rows {
batches: &self.batches,
slots: &self.slots,
};
for (_, view) in self.views.iter_mut().flatten() {
if rebuild {
view.rebuild(source);
} else {
for (id, slot) in rows.clone() {
view.insert(source, id, slot.row_ref());
}
}
}
for (_, groups) in self.groups.iter_mut().flatten() {
if rebuild {
groups.rebuild(source);
} else {
for (id, slot) in rows.clone() {
groups.insert(source, id, slot.row_ref());
}
}
}
if let Some(window) = &mut self.window {
if rebuild {
window.rows.rebuild(source);
} else {
for (id, slot) in rows {
window.rows.insert(source, id, slot.row_ref());
}
}
}
}
fn compact_if_due(&mut self) {
if !self.options.auto_compact {
return;
}
let dead_majority = self.dead >= self.retry_at_dead && self.dead > self.len();
let added = self.batches.len().saturating_sub(self.compacted_batches);
let too_many_batches = added > MAX_BATCHES && self.compaction_error.is_none();
if !(dead_majority || too_many_batches) {
return;
}
if let Err(error) = self.compact() {
self.compaction_error = Some(error);
self.retry_at_dead = self.dead.saturating_mul(2);
}
}
}
fn find<T>(entries: &[Option<(u64, T)>], index: usize, serial: u64) -> Option<&T> {
match entries.get(index)? {
Some((s, entry)) if *s == serial => Some(entry),
_ => None,
}
}
fn store<T>(entries: &mut Vec<Option<(u64, T)>>, entry: T) -> (u64, usize) {
let serial = next_serial();
match entries.iter().position(Option::is_none) {
Some(index) => {
entries[index] = Some((serial, entry));
(serial, index)
}
None => {
entries.push(Some((serial, entry)));
(serial, entries.len() - 1)
}
}
}
fn take<T>(entries: &mut [Option<(u64, T)>], index: usize, serial: u64) -> bool {
entries
.get_mut(index)
.and_then(|entry| entry.take_if(|(s, _)| *s == serial))
.is_some()
}
fn nullable(schema: &Schema) -> SchemaRef {
let fields: Vec<Field> = schema
.fields()
.iter()
.map(|f| f.as_ref().clone().with_nullable(true))
.collect();
Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
}
#[derive(Debug, Clone, Copy)]
pub struct View<'a> {
state: &'a ViewState,
batches: &'a [RecordBatch],
slots: &'a [Option<Slot>],
reuses: &'a [u32],
}
impl View<'_> {
pub fn len(&self) -> usize {
self.state.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn row(&self, index: usize) -> Option<RowRef> {
let id = self.state.get(index)?;
self.slots
.get(id as usize)
.copied()
.flatten()
.map(Slot::row_ref)
}
pub fn record(&self, index: usize) -> Option<RecordId> {
let id = self.state.get(index)?;
Some(RecordId::new(id, reuse(self.reuses, id)))
}
pub fn position(&self, record: RecordId) -> Option<usize> {
let (id, _) = live_slot(self.slots, self.reuses, record)?;
let rows = Rows {
batches: self.batches,
slots: self.slots,
};
self.state.position(rows, id)
}
}
fn reuse(reuses: &[u32], id: RowId) -> u32 {
reuses.get(id as usize).copied().unwrap_or(0)
}
fn live_slot(slots: &[Option<Slot>], reuses: &[u32], record: RecordId) -> Option<(RowId, Slot)> {
let (id, reused) = record.parts();
let slot = (*slots.get(id as usize)?)?;
(reuse(reuses, id) == reused).then_some((id, slot))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_array::{ArrayRef, Int64Array};
use super::*;
fn ids(keys: Vec<i64>) -> RecordBatch {
RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(keys)) as ArrayRef)]).unwrap()
}
#[test]
fn replacing_a_view_again_and_again_keeps_one_place() {
let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
let mut view = table.add_view(ViewSpec::new()).unwrap();
for _ in 0..100 {
let next = table.add_view(ViewSpec::new()).unwrap();
table.remove_view(view);
view = next;
}
assert!(table.views.len() <= 2, "{} places", table.views.len());
}
#[test]
fn ids_of_deleted_keys_are_reused() {
let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
for round in 0..100 {
table
.upsert(ids((0..10).map(|k| round * 10 + k).collect()))
.unwrap();
table
.delete(&Int64Array::from(
(0..10).map(|k| round * 10 + k).collect::<Vec<_>>(),
))
.unwrap();
}
assert!(
table.slots.len() <= 10,
"{} ids for 10 live keys",
table.slots.len()
);
}
#[test]
fn a_failed_automatic_compaction_is_recorded_and_backs_off() {
let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
table.dead = MIN_DEAD_ROWS;
let other = RecordBatch::try_from_iter([(
"other",
Arc::new(arrow_array::StringArray::from(vec!["x"])) as ArrayRef,
)])
.unwrap();
table.batches.push(other);
table.slots.push(Some(Slot { batch: 0, row: 0 }));
table.compact_if_due();
assert!(table.compaction_error().is_some());
assert_eq!(table.retry_at_dead, 2 * MIN_DEAD_ROWS);
}
#[test]
fn automatic_compaction_counts_the_batches_stored_since_the_last_one() {
let mut table = LiveTable::new(ids(vec![]).schema(), "id").unwrap();
table.upsert(ids(vec![1])).unwrap();
table.compacted_batches = 10;
while table.batches.len() < 10 + MAX_BATCHES {
table.batches.push(ids(vec![]));
}
table.compact_if_due();
assert_eq!(table.generation(), 0, "4,096 batches added");
table.batches.push(ids(vec![]));
table.compact_if_due();
assert_eq!(table.generation(), 1, "4,097 batches added");
assert_eq!(table.compacted_batches, 1);
}
}