use crate::epoch::Epoch;
use crate::memtable::Row;
use crate::pma::Pma;
use crate::rowid::RowId;
use std::collections::BTreeMap;
use std::sync::Arc;
type VersionKey = (RowId, Epoch);
#[derive(Clone)]
struct MutableRunSegment {
pma: Pma<VersionKey, Row>,
byte_size: u64,
}
#[derive(Clone)]
pub struct MutableRun {
frozen: Arc<Vec<Arc<MutableRunSegment>>>,
active: MutableRunSegment,
byte_size: u64,
}
impl Default for MutableRun {
fn default() -> Self {
Self::new()
}
}
impl MutableRun {
pub fn new() -> Self {
Self {
frozen: Arc::new(Vec::new()),
active: MutableRunSegment {
pma: Pma::new(),
byte_size: 0,
},
byte_size: 0,
}
}
pub fn insert_many(&mut self, rows: Vec<Row>) {
let batch: Vec<(VersionKey, Row)> = rows
.into_iter()
.map(|r| {
let bytes = r.estimated_bytes();
self.byte_size = self.byte_size.saturating_add(bytes);
self.active.byte_size = self.active.byte_size.saturating_add(bytes);
((r.row_id, r.committed_epoch), r)
})
.collect();
self.active.pma.extend_sorted(batch);
}
pub fn len(&self) -> usize {
self.active.pma.len()
+ self
.frozen
.iter()
.map(|segment| segment.pma.len())
.sum::<usize>()
}
pub fn is_empty(&self) -> bool {
self.active.pma.is_empty() && self.frozen.is_empty()
}
pub fn approx_bytes(&self) -> u64 {
self.byte_size
}
pub fn get_version(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<(Epoch, Row)> {
let mut best: Option<(Epoch, Row)> = None;
for pma in self
.frozen
.iter()
.map(|segment| &segment.pma)
.chain(std::iter::once(&self.active.pma))
{
for ((rid, epoch), row) in pma.iter_from(&(row_id, Epoch::ZERO)) {
if *rid != row_id {
break;
}
if *epoch <= snapshot_epoch
&& best.as_ref().map_or(true, |(current, _)| *epoch > *current)
{
best = Some((*epoch, row.clone()));
}
}
}
best
}
pub fn visible_versions(&self, snapshot_epoch: Epoch) -> Vec<Row> {
let mut by_row: BTreeMap<RowId, Row> = BTreeMap::new();
for pma in self
.frozen
.iter()
.map(|segment| &segment.pma)
.chain(std::iter::once(&self.active.pma))
{
for ((rid, epoch), row) in pma.iter() {
if *epoch <= snapshot_epoch
&& by_row
.get(rid)
.map_or(true, |current| *epoch > current.committed_epoch)
{
by_row.insert(*rid, row.clone());
}
}
}
by_row.into_values().collect()
}
pub(crate) fn seal(&mut self) {
if self.active.pma.is_empty() {
return;
}
let active = std::mem::replace(
&mut self.active,
MutableRunSegment {
pma: Pma::new(),
byte_size: 0,
},
);
Arc::make_mut(&mut self.frozen).push(Arc::new(active));
if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
self.consolidate();
}
}
fn consolidate(&mut self) {
let mut rows = self
.frozen
.iter()
.flat_map(|segment| segment.pma.iter().map(|(_, row)| row.clone()))
.collect::<Vec<_>>();
rows.sort_by_key(|row| (row.row_id, row.committed_epoch));
let mut pma = Pma::new();
pma.extend_sorted(
rows.into_iter()
.map(|row| ((row.row_id, row.committed_epoch), row))
.collect(),
);
self.frozen = Arc::new(vec![Arc::new(MutableRunSegment {
pma,
byte_size: self.byte_size,
})]);
}
#[cfg(test)]
pub(crate) fn frozen_layer_count(&self) -> usize {
self.frozen.len()
}
pub fn drain_sorted(&mut self) -> Vec<Row> {
let mut out = self
.frozen
.iter()
.flat_map(|segment| segment.pma.iter().map(|(_, row)| row.clone()))
.chain(self.active.pma.iter().map(|(_, row)| row.clone()))
.collect::<Vec<_>>();
out.sort_by_key(|row| (row.row_id, row.committed_epoch));
self.frozen = Arc::new(Vec::new());
self.active = MutableRunSegment {
pma: Pma::new(),
byte_size: 0,
};
self.byte_size = 0;
out
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memtable::Value;
fn row(id: u64, epoch: u64, v: i64) -> Row {
Row::new(RowId(id), Epoch(epoch)).with_column(1, Value::Int64(v))
}
fn tomb(id: u64, epoch: u64) -> Row {
Row {
row_id: RowId(id),
committed_epoch: Epoch(epoch),
columns: std::collections::HashMap::new(),
deleted: true,
}
}
fn int_of(r: &Row) -> i64 {
match r.columns.get(&1) {
Some(Value::Int64(x)) => *x,
_ => panic!("expected Int64 column"),
}
}
#[test]
fn get_version_returns_newest_visible() {
let mut mr = MutableRun::new();
mr.insert_many(vec![row(1, 1, 10), row(1, 3, 30), row(1, 9, 90)]);
assert_eq!(int_of(&mr.get_version(RowId(1), Epoch(5)).unwrap().1), 30);
assert_eq!(int_of(&mr.get_version(RowId(1), Epoch(9)).unwrap().1), 90);
assert!(mr.get_version(RowId(1), Epoch(0)).is_none());
assert!(mr.get_version(RowId(2), Epoch(100)).is_none());
}
#[test]
fn tombstone_is_returned_as_a_version() {
let mut mr = MutableRun::new();
mr.insert_many(vec![row(1, 1, 10), tomb(1, 2)]);
let v = mr.get_version(RowId(1), Epoch(5)).unwrap().1;
assert!(v.deleted);
let v0 = mr.get_version(RowId(1), Epoch(1)).unwrap().1;
assert!(!v0.deleted);
}
#[test]
fn sealed_generations_share_rows_and_consolidate() {
let mut writer = MutableRun::new();
for id in 0..crate::MAX_READ_GENERATION_LAYERS as u64 + 2 {
writer.insert_many(vec![row(id, id + 1, id as i64)]);
writer.seal();
}
assert!(writer.frozen_layer_count() < crate::MAX_READ_GENERATION_LAYERS);
let generation = writer.clone();
writer.insert_many(vec![row(99, 99, 99)]);
assert!(generation.get_version(RowId(99), Epoch(99)).is_none());
assert!(writer.get_version(RowId(99), Epoch(99)).is_some());
}
#[test]
fn visible_versions_dedups_to_newest_ascending() {
let mut mr = MutableRun::new();
mr.insert_many(vec![
row(3, 1, 30),
row(1, 1, 10),
row(2, 9, 20), row(1, 3, 11), row(3, 2, 31),
]);
let out = mr.visible_versions(Epoch(5));
let got: Vec<(u64, i64)> = out.iter().map(|r| (r.row_id.0, int_of(r))).collect();
assert_eq!(got, vec![(1, 11), (3, 31)], "row 2 hidden, newest wins");
}
#[test]
fn drain_sorted_is_ascending_version_order_and_empties() {
let mut mr = MutableRun::new();
mr.insert_many(vec![row(3, 1, 0), row(1, 2, 0), row(1, 1, 0), row(2, 1, 0)]);
let out = mr.drain_sorted();
let keys: Vec<(u64, u64)> = out
.iter()
.map(|r| (r.row_id.0, r.committed_epoch.0))
.collect();
assert_eq!(keys, vec![(1, 1), (1, 2), (2, 1), (3, 1)]);
assert!(mr.is_empty());
assert_eq!(mr.approx_bytes(), 0);
}
#[test]
fn many_inserts_stay_queryable() {
let mut mr = MutableRun::new();
let mut rows = Vec::new();
for i in 0..500u64 {
rows.push(row(i, 1, i as i64));
}
mr.insert_many(rows);
assert_eq!(mr.len(), 500);
for i in 0..500u64 {
assert_eq!(
int_of(&mr.get_version(RowId(i), Epoch(1)).unwrap().1),
i as i64
);
}
assert!(mr.approx_bytes() > 0);
}
}