use std::fs::{File, OpenOptions};
use std::io::{BufWriter, Read, Seek, SeekFrom, Write};
use std::path::PathBuf;
use rudb_common::{Error, LogicalType, Result, Value};
use rudb_vector::{Buffer, Chunk, Data, StringColumn, VECTOR_SIZE, Validity, Vector};
const ALL_VALID: u8 = 0;
const ALL_NULL: u8 = 1;
const MASK: u8 = 2;
const EMPTY: u8 = 0;
const VARLEN: u8 = 1;
#[derive(Debug)]
pub(crate) struct Runs {
path: PathBuf,
types: Vec<LogicalType>,
writer: Option<BufWriter<File>>,
reader: Option<File>,
begun: bool,
rows: usize,
blocks: usize,
starts: Vec<u64>,
spans: Vec<Vec<u64>>,
column: usize,
block: usize,
read: usize,
cursors: Vec<u64>,
bytes: u64,
}
impl Runs {
pub(crate) fn new(tag: &str, types: Vec<LogicalType>) -> Result<Self> {
let unique = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |since| since.as_nanos());
let path = std::env::temp_dir().join(format!("rudb-{tag}-{}-{unique}", std::process::id()));
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(true)
.open(&path)
.map_err(|e| {
Error::io(format!("could not open a run file at {}: {e}", path.display()))
})?;
Ok(Self {
path,
types,
writer: Some(BufWriter::with_capacity(1 << 20, file)),
reader: None,
begun: false,
rows: 0,
blocks: 0,
starts: Vec::new(),
spans: Vec::new(),
column: 0,
block: 0,
read: 0,
cursors: Vec::new(),
bytes: 0,
})
}
pub(crate) fn chunks(&self) -> u64 {
self.blocks as u64
}
pub(crate) fn rows(&self) -> u64 {
self.rows as u64
}
pub(crate) fn bytes(&self) -> u64 {
self.bytes
}
pub(crate) fn begin(&mut self, rows: usize) -> Result<()> {
if self.begun {
return Err(Error::internal("a run file was begun twice"));
}
self.begun = true;
self.rows = rows;
self.blocks = rows.div_ceil(VECTOR_SIZE);
Ok(())
}
pub(crate) fn column(&mut self, whole: &Vector) -> Result<()> {
if whole.len() != self.rows {
return Err(Error::internal(format!(
"a run file of {} rows was given a column of {}",
self.rows,
whole.len()
)));
}
for block in 0..self.blocks {
let at = block * VECTOR_SIZE;
self.part(&whole.slice(at, (self.rows - at).min(VECTOR_SIZE))?)?;
}
Ok(())
}
pub(crate) fn part(&mut self, values: &Vector) -> Result<()> {
if !self.begun {
return Err(Error::internal("a run file was written to before it was begun"));
}
let Some(ty) = self.types.get(self.column) else {
return Err(Error::internal(format!(
"a run file of {} columns was given another one",
self.types.len()
)));
};
let at = self.block * VECTOR_SIZE;
let rows = (self.rows - at).min(VECTOR_SIZE);
if values.len() != rows {
return Err(Error::internal(format!(
"a run file block of {rows} rows was given {}",
values.len()
)));
}
if values.logical_type() != ty {
return Err(Error::internal(format!(
"a run file column of {ty} was given {}",
values.logical_type()
)));
}
let flat = values.flatten()?;
let Some(writer) = self.writer.as_mut() else {
return Err(Error::internal("a run file was written to after it was read"));
};
let mut out = Sink { writer, written: 0 };
put_validity(&mut out, &flat, rows)?;
put_payload(&mut out, &flat, rows)?;
let written = out.written;
if self.block == 0 {
self.starts.push(self.bytes);
self.spans.push(Vec::with_capacity(self.blocks));
}
if let Some(spans) = self.spans.get_mut(self.column) {
spans.push(written);
}
self.bytes += written;
self.block += 1;
if self.block == self.blocks {
self.column += 1;
self.block = 0;
}
Ok(())
}
pub(crate) fn next_chunk(&mut self) -> Result<Option<Chunk>> {
if let Some(mut writer) = self.writer.take() {
writer.flush().map_err(|e| Error::io(format!("could not finish a run file: {e}")))?;
let file = writer
.into_inner()
.map_err(|e| Error::io(format!("could not finish a run file: {e}")))?;
if self.blocks > 0 && self.starts.len() != self.types.len() {
return Err(Error::internal(format!(
"a run file of {} columns was read with {} of them written",
self.types.len(),
self.starts.len()
)));
}
self.cursors.clone_from(&self.starts);
self.reader = Some(file);
}
if self.read >= self.blocks {
return Ok(None);
}
let block = self.read;
let at = block * VECTOR_SIZE;
let rows = (self.rows - at).min(VECTOR_SIZE);
let mut columns = Vec::with_capacity(self.types.len());
let mut bytes = Vec::new();
for (position, ty) in self.types.iter().enumerate() {
let (Some(file), Some(span), Some(cursor)) = (
self.reader.as_mut(),
self.spans.get(position).and_then(|spans| spans.get(block)),
self.cursors.get_mut(position),
) else {
return Err(Error::internal("a run file was read before it was written"));
};
file.seek(SeekFrom::Start(*cursor))
.map_err(|e| Error::io(format!("could not seek in a run file: {e}")))?;
bytes.clear();
bytes.resize(usize::try_from(*span).unwrap_or(usize::MAX), 0);
file.read_exact(&mut bytes)
.map_err(|e| Error::io(format!("could not read a run file: {e}")))?;
*cursor += *span;
let mut src: &[u8] = &bytes;
let validity = take_validity(&mut src, rows)?;
let column = take_payload(&mut src, ty, rows)?;
columns.push(column.with_validity(validity));
}
self.read += 1;
Ok(Some(Chunk::with_rows(columns, rows)?))
}
}
impl Drop for Runs {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.path);
}
}
fn put_validity(out: &mut Sink<'_>, column: &Vector, rows: usize) -> Result<()> {
let validity = column.validity();
match validity {
Validity::AllValid => out.put(&[ALL_VALID]),
Validity::AllInvalid => out.put(&[ALL_NULL]),
Validity::Mask(_) => {
out.put(&[MASK])?;
let mut bytes = Vec::with_capacity(rows);
for row in 0..rows {
bytes.push(u8::from(validity.is_valid(row)));
}
out.put(&bytes)
}
}
}
fn take_validity(src: &mut &[u8], rows: usize) -> Result<Validity> {
let mut tag = [0u8; 1];
fill(src, &mut tag)?;
match tag[0] {
ALL_VALID => Ok(Validity::AllValid),
ALL_NULL => Ok(Validity::AllInvalid),
MASK => {
let mut bytes = vec![0u8; rows];
fill(src, &mut bytes)?;
Ok(Validity::from_iter(rows, |row| bytes[row] != 0))
}
other => Err(Error::internal(format!("a run file has {other} where a validity byte goes"))),
}
}
macro_rules! layouts {
($(($tag:literal, $variant:ident, $native:ty)),+ $(,)?) => {
fn put_payload(out: &mut Sink<'_>, column: &Vector, rows: usize) -> Result<()> {
let Some(data) = column.data() else {
return Err(Error::internal(format!(
"a run file was given a {} column that did not flatten",
column.logical_type()
)));
};
match data {
$(Data::$variant(values) => {
out.put(&[$tag])?;
let mut bytes = Vec::with_capacity(values.len() * <$native as Plain>::WIDTH);
for value in values.as_slice() {
value.put(&mut bytes);
}
out.put(&bytes)
})+
Data::Varlen(strings) => {
out.put(&[VARLEN])?;
put_strings(out, strings, rows)
}
Data::Empty => out.put(&[EMPTY]),
_ => Err(Error::internal(format!(
"a run file has no layout for a {} column",
column.logical_type()
))),
}
}
fn take_payload(src: &mut &[u8], ty: &LogicalType, rows: usize) -> Result<Vector> {
let mut tag = [0u8; 1];
fill(src, &mut tag)?;
let data = match tag[0] {
$($tag => {
let mut bytes = vec![0u8; rows * <$native as Plain>::WIDTH];
fill(src, &mut bytes)?;
let mut values = Vec::with_capacity(rows);
for at in 0..rows {
values.push(<$native as Plain>::get(&bytes[at * <$native as Plain>::WIDTH..]));
}
Data::$variant(Buffer::from_vec(values))
})+
VARLEN => Data::Varlen(take_strings(src, rows)?),
EMPTY => return Ok(Vector::constant(ty.clone(), Value::Null, rows)),
other => {
return Err(Error::internal(format!(
"a run file has {other} where the layout of a {ty} column goes"
)));
}
};
Vector::flat(ty.clone(), data)
}
};
}
layouts!(
(2, Bool, bool),
(3, Int8, i8),
(4, Int16, i16),
(5, Int32, i32),
(6, Int64, i64),
(7, Int128, i128),
(8, UInt8, u8),
(9, UInt16, u16),
(10, UInt32, u32),
(11, UInt64, u64),
(12, UInt128, u128),
(13, Float32, f32),
(14, Float64, f64),
(15, Interval, (i32, i32, i64)),
);
fn put_strings(out: &mut Sink<'_>, strings: &StringColumn, rows: usize) -> Result<()> {
let mut bytes = Vec::with_capacity(rows * 16);
for row in 0..rows {
let text = strings.bytes(row).unwrap_or_default();
bytes.extend_from_slice(&u32::try_from(text.len()).unwrap_or(u32::MAX).to_le_bytes());
bytes.extend_from_slice(text);
}
out.put(&bytes)
}
fn take_strings(src: &mut &[u8], rows: usize) -> Result<StringColumn> {
let mut strings = StringColumn::with_capacity(rows);
let mut text = Vec::new();
for _ in 0..rows {
let len = take_u32(src)? as usize;
text.clear();
text.resize(len, 0);
fill(src, &mut text)?;
strings.push_bytes(&text);
}
Ok(strings)
}
trait Plain: Sized {
const WIDTH: usize;
fn put(&self, out: &mut Vec<u8>);
fn get(bytes: &[u8]) -> Self;
}
macro_rules! plain_numbers {
($($native:ty),+ $(,)?) => {
$(impl Plain for $native {
const WIDTH: usize = std::mem::size_of::<$native>();
fn put(&self, out: &mut Vec<u8>) {
out.extend_from_slice(&self.to_le_bytes());
}
fn get(bytes: &[u8]) -> Self {
let mut at = [0u8; Self::WIDTH];
at.copy_from_slice(&bytes[..Self::WIDTH]);
Self::from_le_bytes(at)
}
})+
};
}
plain_numbers!(i8, i16, i32, i64, i128, u8, u16, u32, u64, u128, f32, f64);
impl Plain for bool {
const WIDTH: usize = 1;
fn put(&self, out: &mut Vec<u8>) {
out.push(u8::from(*self));
}
fn get(bytes: &[u8]) -> Self {
bytes[0] != 0
}
}
impl Plain for (i32, i32, i64) {
const WIDTH: usize = 16;
fn put(&self, out: &mut Vec<u8>) {
self.0.put(out);
self.1.put(out);
self.2.put(out);
}
fn get(bytes: &[u8]) -> Self {
(i32::get(bytes), i32::get(&bytes[4..]), i64::get(&bytes[8..]))
}
}
fn take_u32(src: &mut &[u8]) -> Result<u32> {
let mut bytes = [0u8; 4];
fill(src, &mut bytes)?;
Ok(u32::from_le_bytes(bytes))
}
fn fill(src: &mut &[u8], into: &mut [u8]) -> Result<()> {
if src.len() < into.len() {
return Err(Error::internal("a run file block ended in the middle of a value"));
}
let (take, left) = src.split_at(into.len());
into.copy_from_slice(take);
*src = left;
Ok(())
}
struct Sink<'a> {
writer: &'a mut BufWriter<File>,
written: u64,
}
impl Sink<'_> {
fn put(&mut self, bytes: &[u8]) -> Result<()> {
self.writer
.write_all(bytes)
.map_err(|e| Error::io(format!("could not write to a run file: {e}")))?;
self.written += bytes.len() as u64;
Ok(())
}
}
#[cfg(test)]
mod tests {
use rudb_common::{LogicalType, Value};
use rudb_vector::VECTOR_SIZE;
use super::{Buffer, Chunk, Data, Runs, StringColumn, Validity, Vector};
fn ints(values: &[i32]) -> Vector {
Vector::flat(LogicalType::Integer, Data::Int32(Buffer::from_vec(values.to_vec())))
.expect("integers are an i32 layout")
}
fn text(values: &[&str]) -> Vector {
let mut strings = StringColumn::with_capacity(values.len());
for value in values {
strings.push(value);
}
Vector::flat(LogicalType::Varchar, Data::Varlen(strings)).expect("strings are a varlen")
}
fn wrote(tag: &str, columns: &[Vector]) -> Runs {
let types = columns.iter().map(|column| column.logical_type().clone()).collect();
let mut runs = Runs::new(tag, types).expect("a temporary file");
runs.begin(columns.first().map_or(0, Vector::len)).expect("the row count");
for column in columns {
runs.column(column).expect("a column");
}
runs
}
fn rows_of(runs: &mut Runs) -> Vec<Vec<Value>> {
let mut out = Vec::new();
while let Some(chunk) = runs.next_chunk().expect("readable") {
for row in 0..chunk.len() {
out.push(chunk.row(row).collect());
}
}
out
}
fn rows_in(chunk: &Chunk) -> Vec<Vec<Value>> {
(0..chunk.len()).map(|row| chunk.row(row).collect()).collect()
}
#[test]
fn rows_come_back_in_the_order_they_went_in() {
let chunk =
Chunk::new(vec![ints(&[1, 2, 3]), text(&["one", "two", "three"])]).expect("three rows");
let mut runs = wrote("test-order", chunk.columns());
assert_eq!(runs.chunks(), 1);
assert_eq!(runs.rows(), 3);
assert!(runs.bytes() > 0, "something was written");
assert_eq!(rows_of(&mut runs), rows_in(&chunk));
}
#[test]
fn a_run_longer_than_a_block_comes_back_as_several() {
let rows = VECTOR_SIZE * 2 + 7;
let numbers: Vec<i32> = (0..rows as i32).collect();
let words: Vec<String> = numbers.iter().map(|value| format!("row {value}")).collect();
let borrowed: Vec<&str> = words.iter().map(String::as_str).collect();
let mut runs = wrote("test-blocks", &[ints(&numbers), text(&borrowed)]);
assert_eq!(runs.chunks(), 3);
assert_eq!(runs.rows(), rows as u64);
let back = rows_of(&mut runs);
assert_eq!(back.len(), rows);
assert_eq!(back[0], vec![Value::Integer(0), Value::Varchar("row 0".into())]);
assert_eq!(
back[VECTOR_SIZE],
vec![Value::Integer(VECTOR_SIZE as i32), Value::Varchar(format!("row {VECTOR_SIZE}"))],
"the first row of the second block"
);
assert_eq!(
back[rows - 1],
vec![Value::Integer(rows as i32 - 1), Value::Varchar(format!("row {}", rows - 1))],
"and the last row of the short one"
);
}
#[test]
fn long_strings_and_empty_ones_survive() {
let long = "x".repeat(400);
let chunk = Chunk::new(vec![text(&[&long, "", "short"])]).expect("three rows");
let mut runs = wrote("test-strings", chunk.columns());
assert_eq!(rows_of(&mut runs), rows_in(&chunk));
}
#[test]
fn the_three_validity_cases_all_come_back() {
let some = ints(&[7, 8, 9]).with_validity(Validity::from_iter(3, |row| row != 1));
let none = ints(&[1, 2, 3]).with_validity(Validity::AllInvalid);
let mut runs = wrote("test-validity", &[ints(&[4, 5, 6]), some, none]);
let back = runs.next_chunk().expect("readable").expect("a chunk");
assert_eq!(back.value_at(1, 0), Value::Integer(5), "all valid");
assert_eq!(back.value_at(0, 1), Value::Integer(7), "either side of the null");
assert_eq!(back.value_at(1, 1), Value::Null, "the masked null");
assert_eq!(back.value_at(2, 1), Value::Integer(9));
assert_eq!(back.value_at(0, 2), Value::Null, "all null");
}
#[test]
fn a_column_that_is_not_flat_is_flattened_on_the_way_out() {
let coded = Vector::dictionary(vec![1, 0, 1], text(&["no", "yes"])).expect("a dictionary");
let same = Vector::constant(LogicalType::Integer, Value::Integer(42), 3);
let chunk = Chunk::new(vec![coded, same]).expect("three rows");
let mut runs = wrote("test-flat", chunk.columns());
assert_eq!(rows_of(&mut runs), rows_in(&chunk));
}
#[test]
fn a_column_of_the_wrong_length_is_refused() {
let mut runs = Runs::new("test-length", vec![LogicalType::Integer]).expect("a file");
runs.begin(2).expect("the row count");
let why = runs.column(&ints(&[1])).expect_err("one row into a file of two");
assert!(why.to_string().contains("was given a column of 1"), "{why}");
}
#[test]
fn a_column_of_the_wrong_type_is_refused() {
let mut runs = Runs::new("test-type", vec![LogicalType::Varchar]).expect("a file");
runs.begin(1).expect("the row count");
let why = runs.column(&ints(&[1])).expect_err("an integer into a varchar column");
assert!(why.to_string().contains("was given INTEGER"), "{why}");
}
#[test]
fn a_column_past_the_last_one_is_refused() {
let mut runs = wrote("test-width", &[ints(&[1])]);
let why = runs.column(&ints(&[2])).expect_err("a second column into a file of one");
assert!(why.to_string().contains("was given another one"), "{why}");
}
#[test]
fn a_run_missing_a_column_is_refused_rather_than_read() {
let types = vec![LogicalType::Integer, LogicalType::Integer];
let mut runs = Runs::new("test-missing", types).expect("a file");
runs.begin(1).expect("the row count");
runs.column(&ints(&[1])).expect("the first column");
let why = runs.next_chunk().expect_err("one column of two");
assert!(why.to_string().contains("with 1 of them written"), "{why}");
}
#[test]
fn a_file_nobody_wrote_to_reads_back_as_nothing() {
let mut runs = Runs::new("test-nothing", vec![LogicalType::Integer]).expect("a file");
assert!(runs.next_chunk().expect("readable").is_none());
}
#[test]
fn a_run_of_no_rows_reads_back_as_nothing() {
let mut runs = wrote("test-norows", &[ints(&[])]);
assert_eq!(runs.chunks(), 0);
assert!(runs.next_chunk().expect("readable").is_none());
}
#[test]
fn a_run_cannot_be_begun_twice() {
let mut runs = Runs::new("test-begin", vec![LogicalType::Integer]).expect("a file");
runs.begin(1).expect("the row count");
let why = runs.begin(2).expect_err("a second row count");
assert!(why.to_string().contains("begun twice"), "{why}");
}
#[test]
fn the_file_is_gone_when_the_runs_are_dropped() {
let runs = Runs::new("test-drop", vec![LogicalType::Integer]).expect("a file");
let path = runs.path.clone();
assert!(path.exists(), "it is there while the runs are");
drop(runs);
assert!(!path.exists(), "and gone after");
}
}