use std::time::Instant;
use rudb_common::{Error, LogicalType, Result, Value};
use rudb_vector::vector::VECTOR_SIZE;
use rudb_vector::{Chunk, Vector};
use crate::zone::{Probe, Zone};
#[derive(Debug, Clone)]
pub struct MemoryTable {
types: Vec<LogicalType>,
chunks: Vec<Chunk>,
zones: Vec<Zone>,
rows: usize,
stats_ns: u64,
}
impl MemoryTable {
#[must_use]
pub fn new(types: Vec<LogicalType>) -> Self {
Self { types, chunks: Vec::new(), zones: Vec::new(), rows: 0, stats_ns: 0 }
}
#[must_use]
pub fn types(&self) -> &[LogicalType] {
&self.types
}
#[must_use]
pub fn width(&self) -> usize {
self.types.len()
}
#[must_use]
pub fn len(&self) -> usize {
self.rows
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.rows == 0
}
#[must_use]
pub fn chunk_count(&self) -> usize {
self.chunks.len()
}
pub fn append(&mut self, chunk: Chunk) -> Result<()> {
if chunk.width() != self.types.len() {
return Err(Error::internal(format!(
"a chunk of {} columns appended to a table of {}",
chunk.width(),
self.types.len()
)));
}
for (index, (held, wanted)) in chunk.types().iter().zip(&self.types).enumerate() {
if held != wanted {
return Err(Error::internal(format!(
"column {index} of the chunk is {held} and the table's is {wanted}"
)));
}
}
if chunk.is_empty() {
return Ok(());
}
let started = Instant::now();
let zone = Zone::of(&chunk);
self.stats_ns += started.elapsed().as_nanos() as u64;
self.rows += chunk.len();
self.zones.push(zone);
self.chunks.push(chunk);
Ok(())
}
#[must_use]
pub fn stats_ns(&self) -> u64 {
self.stats_ns
}
#[must_use]
pub fn zone(&self, index: usize) -> Option<&Zone> {
self.zones.get(index)
}
#[must_use]
pub fn skips(&self, index: usize, probes: &[Probe]) -> bool {
self.zones.get(index).is_some_and(|zone| zone.skips(probes))
}
pub fn append_rows(&mut self, rows: &[Vec<Value>]) -> Result<()> {
for (index, row) in rows.iter().enumerate() {
if row.len() != self.types.len() {
return Err(Error::internal(format!(
"row {index} has {} values and the table has {} columns",
row.len(),
self.types.len()
)));
}
}
for batch in rows.chunks(VECTOR_SIZE) {
let mut columns = Vec::with_capacity(self.types.len());
for (position, ty) in self.types.iter().enumerate() {
let down: Vec<Value> = batch.iter().map(|row| row[position].clone()).collect();
columns.push(Vector::from_values(ty.clone(), &down)?);
}
self.append(Chunk::with_rows(columns, batch.len())?)?;
}
Ok(())
}
pub fn read(&self, chunk: usize, columns: &[usize]) -> Result<Chunk> {
let held = self.chunks.get(chunk).ok_or_else(|| {
Error::internal(format!(
"chunk {chunk} of a table that has {} chunks",
self.chunks.len()
))
})?;
let mut picked = Vec::with_capacity(columns.len());
for &column in columns {
picked.push(held.column(column)?.clone());
}
Chunk::with_rows(picked, held.len())
}
#[must_use]
pub fn chunk(&self, index: usize) -> Option<&Chunk> {
self.chunks.get(index)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn people() -> MemoryTable {
let mut table = MemoryTable::new(vec![LogicalType::Integer, LogicalType::Varchar]);
table
.append_rows(&[
vec![Value::Integer(1), Value::Varchar("ada".to_string())],
vec![Value::Integer(2), Value::Null],
vec![Value::Integer(3), Value::Varchar("grace".to_string())],
])
.expect("three rows of the table's own types");
table
}
#[test]
fn rows_go_in_and_come_back_out() {
let table = people();
assert_eq!(table.len(), 3);
assert_eq!(table.chunk_count(), 1);
let chunk = table.read(0, &[0, 1]).expect("both columns of the only chunk");
assert_eq!(chunk.value_at(0, 1), Value::Varchar("ada".to_string()));
assert_eq!(chunk.value_at(1, 1), Value::Null);
assert_eq!(chunk.value_at(2, 0), Value::Integer(3));
}
#[test]
fn a_read_gives_back_only_the_columns_it_was_asked_for() {
let table = people();
let chunk = table.read(0, &[1]).expect("the second column");
assert_eq!(chunk.width(), 1);
assert_eq!(chunk.len(), 3);
assert_eq!(chunk.value_at(2, 0), Value::Varchar("grace".to_string()));
}
#[test]
fn a_read_of_no_columns_still_says_how_many_rows() {
let table = people();
let chunk = table.read(0, &[]).expect("no columns");
assert_eq!(chunk.width(), 0);
assert_eq!(chunk.len(), 3);
}
#[test]
fn more_rows_than_a_vector_become_more_than_one_chunk() {
let mut table = MemoryTable::new(vec![LogicalType::BigInt]);
let rows: Vec<Vec<Value>> =
(0..VECTOR_SIZE + 5).map(|n| vec![Value::BigInt(n as i64)]).collect();
table.append_rows(&rows).expect("bigints");
assert_eq!(table.len(), VECTOR_SIZE + 5);
assert_eq!(table.chunk_count(), 2);
let last = table.read(1, &[0]).expect("the second chunk");
assert_eq!(last.len(), 5);
assert_eq!(last.value_at(4, 0), Value::BigInt((VECTOR_SIZE + 4) as i64));
}
#[test]
fn a_chunk_of_the_wrong_types_is_caught() {
let mut table = MemoryTable::new(vec![LogicalType::Integer]);
let wrong = Chunk::new(vec![
Vector::from_values(LogicalType::Varchar, &[Value::Varchar("x".to_string())])
.expect("a string column"),
])
.expect("one column");
let error = table.append(wrong).expect_err("a varchar is not an integer");
assert!(error.message().contains("column 0"), "{error}");
}
#[test]
fn a_row_of_the_wrong_width_is_caught() {
let mut table = MemoryTable::new(vec![LogicalType::Integer, LogicalType::Integer]);
let error =
table.append_rows(&[vec![Value::Integer(1)]]).expect_err("a row of one is not a row");
assert!(error.message().contains("row 0"), "{error}");
}
#[test]
fn an_empty_chunk_is_not_stored() {
let mut table = MemoryTable::new(vec![LogicalType::Integer]);
table.append(Chunk::empty(&[LogicalType::Integer])).expect("an empty chunk is allowed");
assert_eq!(table.chunk_count(), 0);
assert!(table.is_empty());
}
}