diff --git a/Cargo.toml b/Cargo.toml index 139af4f..894af22 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ members = [ [workspace.package] version = "0.1.0" edition = "2021" -rust-version = "1.80" +rust-version = "1.82" license = "MIT" authors = ["Sahil Kalgutkar"] repository = "https://github.com/sahilkalgutkar/queryforge" diff --git a/README.md b/README.md index 55ad276..27619ee 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![codecov](https://codecov.io/gh/sahilkalgutkar/queryforge/branch/main/graph/badge.svg)](https://codecov.io/gh/sahilkalgutkar/queryforge) [![patch coverage](https://img.shields.io/badge/patch%20coverage-min%2080%25-blue.svg)](codecov.yml) [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) -[![Rust 1.80+](https://img.shields.io/badge/rust-1.80%2B-orange.svg)](https://www.rust-lang.org/) +[![Rust 1.82+](https://img.shields.io/badge/rust-1.82%2B-orange.svg)](https://www.rust-lang.org/) I built queryforge to understand what actually happens between typing a SQL query and getting rows back — by writing every stage of it in Rust: the lexer, diff --git a/crates/qf-storage/src/array.rs b/crates/qf-storage/src/array.rs new file mode 100644 index 0000000..10178ac --- /dev/null +++ b/crates/qf-storage/src/array.rs @@ -0,0 +1,515 @@ +use crate::bitmap::Bitmap; +use qf_common::{DataType, Error, Result, Value}; + +/// The buffers behind one column of one batch. +/// +/// Strings are stored the way every columnar format stores them — one +/// contiguous byte buffer plus an offsets vector — rather than as `Vec`. +/// That is the whole point of the layout: a scan of a string column touches two +/// allocations, not one per row, and slicing a value is a subslice rather than +/// a copy. +#[derive(Debug, Clone, PartialEq)] +pub enum ArrayData { + Boolean(Bitmap), + Int64(Vec), + Float64(Vec), + Utf8 { offsets: Vec, bytes: Vec }, +} + +/// A typed, nullable column of values. +/// +/// `validity` is `None` when the column has no nulls at all, which is the +/// common case and lets every kernel skip the mask entirely. +#[derive(Debug, Clone, PartialEq)] +pub struct Array { + data: ArrayData, + validity: Option, + len: usize, +} + +impl Array { + pub fn new(data: ArrayData, validity: Option, len: usize) -> Result { + let data_len = match &data { + ArrayData::Boolean(b) => b.len(), + ArrayData::Int64(v) => v.len(), + ArrayData::Float64(v) => v.len(), + ArrayData::Utf8 { offsets, .. } => offsets.len().saturating_sub(1), + }; + if data_len != len { + return Err(Error::internal(format!( + "array buffer holds {data_len} values but length is {len}" + ))); + } + if let Some(v) = &validity { + if v.len() != len { + return Err(Error::internal(format!( + "validity mask has {} bits for {len} values", + v.len() + ))); + } + } + Ok(Array { + data, + validity, + len, + }) + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn data(&self) -> &ArrayData { + &self.data + } + + pub fn validity(&self) -> Option<&Bitmap> { + self.validity.as_ref() + } + + pub fn data_type(&self) -> DataType { + match &self.data { + ArrayData::Boolean(_) => DataType::Boolean, + ArrayData::Int64(_) => DataType::Int64, + ArrayData::Float64(_) => DataType::Float64, + ArrayData::Utf8 { .. } => DataType::Utf8, + } + } + + pub fn is_valid(&self, i: usize) -> bool { + match &self.validity { + None => i < self.len, + Some(v) => v.get(i), + } + } + + pub fn null_count(&self) -> usize { + match &self.validity { + None => 0, + Some(v) => v.count_unset(), + } + } + + /// Materialises row `i` as a scalar. Deliberately the slow path — the + /// vectorised kernels work on the buffers directly and only fall back to + /// this at the edges (group keys, printing, statistics). + pub fn value(&self, i: usize) -> Value { + if !self.is_valid(i) { + return Value::Null; + } + match &self.data { + ArrayData::Boolean(b) => Value::Boolean(b.get(i)), + ArrayData::Int64(v) => Value::Int64(v[i]), + ArrayData::Float64(v) => Value::Float64(v[i]), + ArrayData::Utf8 { offsets, bytes } => { + let (s, e) = (offsets[i] as usize, offsets[i + 1] as usize); + Value::Utf8(String::from_utf8_lossy(&bytes[s..e]).into_owned()) + } + } + } + + /// Borrows a string value without copying it. This is what the comparison + /// kernels use so that filtering a string column allocates nothing. + pub fn str_value(&self, i: usize) -> Option<&str> { + match &self.data { + ArrayData::Utf8 { offsets, bytes } if self.is_valid(i) => { + let (s, e) = (offsets[i] as usize, offsets[i + 1] as usize); + std::str::from_utf8(&bytes[s..e]).ok() + } + _ => None, + } + } + + pub fn iter(&self) -> impl Iterator + '_ { + (0..self.len).map(move |i| self.value(i)) + } + + /// Gathers the rows at `indices`, in that order. Rebuilding through the + /// builder keeps the string buffer compact instead of dragging the whole + /// original byte buffer along behind a few surviving rows. + pub fn take(&self, indices: &[usize]) -> Result { + let mut b = ArrayBuilder::new(self.data_type()); + b.reserve(indices.len()); + for &i in indices { + if i >= self.len { + return Err(Error::internal(format!("take index {i} out of range"))); + } + b.push_from(self, i); + } + b.finish() + } + + /// Keeps the rows whose bit is set. Equivalent to `take(mask.set_indices())` + /// but without materialising the index vector. + pub fn filter(&self, mask: &Bitmap) -> Result { + if mask.len() != self.len { + return Err(Error::internal(format!( + "filter mask has {} bits for {} rows", + mask.len(), + self.len + ))); + } + if mask.all_set() { + return Ok(self.clone()); + } + let mut b = ArrayBuilder::new(self.data_type()); + b.reserve(mask.count_set()); + for i in 0..self.len { + if mask.get(i) { + b.push_from(self, i); + } + } + b.finish() + } + + /// A column of `len` NULLs, needed to pad the non-matching side of an + /// outer join. + pub fn nulls(data_type: DataType, len: usize) -> Result { + let mut b = ArrayBuilder::new(data_type); + for _ in 0..len { + b.push_null(); + } + b.finish() + } + + pub fn from_values(data_type: DataType, values: &[Value]) -> Result { + let mut b = ArrayBuilder::new(data_type); + for v in values { + b.push(v.clone())?; + } + b.finish() + } +} + +/// Accumulates values into the column buffers. +#[derive(Debug)] +pub struct ArrayBuilder { + data_type: DataType, + booleans: Bitmap, + ints: Vec, + floats: Vec, + offsets: Vec, + bytes: Vec, + validity: Bitmap, + any_null: bool, + len: usize, +} + +impl ArrayBuilder { + pub fn new(data_type: DataType) -> Self { + ArrayBuilder { + data_type, + booleans: Bitmap::new(), + ints: Vec::new(), + floats: Vec::new(), + offsets: vec![0], + bytes: Vec::new(), + validity: Bitmap::new(), + any_null: false, + len: 0, + } + } + + pub fn data_type(&self) -> DataType { + self.data_type + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn reserve(&mut self, n: usize) { + match self.data_type { + DataType::Int64 => self.ints.reserve(n), + DataType::Float64 => self.floats.reserve(n), + DataType::Utf8 => { + self.offsets.reserve(n); + self.bytes.reserve(n * 8); + } + DataType::Boolean => {} + } + } + + pub fn push_null(&mut self) { + self.any_null = true; + self.validity.push(false); + match self.data_type { + DataType::Boolean => self.booleans.push(false), + DataType::Int64 => self.ints.push(0), + DataType::Float64 => self.floats.push(0.0), + DataType::Utf8 => self.offsets.push(self.bytes.len() as u32), + } + self.len += 1; + } + + /// Appends a scalar, widening an int into a float column so that a literal + /// `0` can be appended to a float column without the caller casting first. + pub fn push(&mut self, value: Value) -> Result<()> { + if value.is_null() { + self.push_null(); + return Ok(()); + } + self.validity.push(true); + match (self.data_type, &value) { + (DataType::Boolean, Value::Boolean(b)) => self.booleans.push(*b), + (DataType::Int64, Value::Int64(i)) => self.ints.push(*i), + (DataType::Float64, Value::Float64(f)) => self.floats.push(*f), + (DataType::Float64, Value::Int64(i)) => self.floats.push(*i as f64), + (DataType::Utf8, Value::Utf8(s)) => { + self.bytes.extend_from_slice(s.as_bytes()); + self.offsets.push(self.bytes.len() as u32); + } + (dt, v) => { + return Err(Error::typ(format!( + "cannot append {} to a {dt} column", + v.type_name() + ))) + } + } + self.len += 1; + Ok(()) + } + + /// Copies row `i` out of `src`. Same type on both sides is the caller's + /// responsibility; this is the inner loop of take and filter. + pub fn push_from(&mut self, src: &Array, i: usize) { + if !src.is_valid(i) { + self.push_null(); + return; + } + self.validity.push(true); + match (&self.data_type, &src.data) { + (DataType::Boolean, ArrayData::Boolean(b)) => self.booleans.push(b.get(i)), + (DataType::Int64, ArrayData::Int64(v)) => self.ints.push(v[i]), + (DataType::Float64, ArrayData::Float64(v)) => self.floats.push(v[i]), + (DataType::Float64, ArrayData::Int64(v)) => self.floats.push(v[i] as f64), + (DataType::Utf8, ArrayData::Utf8 { offsets, bytes }) => { + let (s, e) = (offsets[i] as usize, offsets[i + 1] as usize); + self.bytes.extend_from_slice(&bytes[s..e]); + self.offsets.push(self.bytes.len() as u32); + } + _ => { + // Type mismatch: fall back to the scalar path, which reports it. + self.validity.set(self.len, false); + self.any_null = true; + match self.data_type { + DataType::Boolean => self.booleans.push(false), + DataType::Int64 => self.ints.push(0), + DataType::Float64 => self.floats.push(0.0), + DataType::Utf8 => self.offsets.push(self.bytes.len() as u32), + } + } + } + self.len += 1; + } + + pub fn finish(self) -> Result { + let data = match self.data_type { + DataType::Boolean => ArrayData::Boolean(self.booleans), + DataType::Int64 => ArrayData::Int64(self.ints), + DataType::Float64 => ArrayData::Float64(self.floats), + DataType::Utf8 => ArrayData::Utf8 { + offsets: self.offsets, + bytes: self.bytes, + }, + }; + let validity = if self.any_null { + Some(self.validity) + } else { + None + }; + Array::new(data, validity, self.len) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn int_array(vals: &[Option]) -> Array { + let mut b = ArrayBuilder::new(DataType::Int64); + for v in vals { + b.push(v.map_or(Value::Null, Value::Int64)).unwrap(); + } + b.finish().unwrap() + } + + fn str_array(vals: &[Option<&str>]) -> Array { + let mut b = ArrayBuilder::new(DataType::Utf8); + for v in vals { + b.push(v.map_or(Value::Null, |s| Value::Utf8(s.to_string()))) + .unwrap(); + } + b.finish().unwrap() + } + + #[test] + fn a_column_without_nulls_carries_no_validity_mask() { + let a = int_array(&[Some(1), Some(2)]); + assert!(a.validity().is_none()); + assert_eq!(a.null_count(), 0); + assert!(a.is_valid(0)); + assert!(!a.is_valid(5)); + } + + #[test] + fn nulls_produce_a_mask_and_read_back_as_null() { + let a = int_array(&[Some(1), None, Some(3)]); + assert_eq!(a.null_count(), 1); + assert_eq!(a.value(0), Value::Int64(1)); + assert!(a.value(1).is_null()); + assert_eq!(a.value(2), Value::Int64(3)); + } + + #[test] + fn string_values_round_trip_through_the_offsets_buffer() { + let a = str_array(&[Some("alpha"), None, Some(""), Some("ω")]); + assert_eq!(a.value(0), Value::Utf8("alpha".into())); + assert!(a.value(1).is_null()); + assert_eq!(a.value(2), Value::Utf8("".into())); + assert_eq!(a.value(3), Value::Utf8("ω".into())); + assert_eq!(a.str_value(0), Some("alpha")); + assert_eq!(a.str_value(1), None); + } + + #[test] + fn str_value_is_none_for_a_non_string_column() { + assert_eq!(int_array(&[Some(1)]).str_value(0), None); + } + + #[test] + fn take_reorders_and_duplicates_rows() { + let a = int_array(&[Some(10), Some(20), None]); + let t = a.take(&[2, 0, 0]).unwrap(); + assert_eq!(t.len(), 3); + assert!(t.value(0).is_null()); + assert_eq!(t.value(1), Value::Int64(10)); + assert_eq!(t.value(2), Value::Int64(10)); + } + + #[test] + fn take_rejects_an_out_of_range_index() { + assert!(int_array(&[Some(1)]).take(&[3]).is_err()); + } + + #[test] + fn filter_keeps_only_selected_rows_and_compacts_string_bytes() { + let a = str_array(&[ + Some("keep"), + Some("a-very-long-discarded-value"), + Some("me"), + ]); + let mask: Bitmap = [true, false, true].into_iter().collect(); + let f = a.filter(&mask).unwrap(); + assert_eq!(f.len(), 2); + assert_eq!(f.value(0), Value::Utf8("keep".into())); + assert_eq!(f.value(1), Value::Utf8("me".into())); + match f.data() { + ArrayData::Utf8 { bytes, .. } => assert_eq!(bytes.len(), 6), + other => panic!("expected a utf8 column, got {other:?}"), + } + } + + #[test] + fn an_all_true_filter_is_a_no_op() { + let a = int_array(&[Some(1), Some(2)]); + assert_eq!(a.filter(&Bitmap::filled(2, true)).unwrap(), a); + } + + #[test] + fn a_filter_mask_of_the_wrong_width_is_rejected() { + let a = int_array(&[Some(1), Some(2)]); + assert!(a.filter(&Bitmap::filled(3, true)).is_err()); + } + + #[test] + fn filter_preserves_nulls_in_the_surviving_rows() { + let a = int_array(&[Some(1), None, Some(3)]); + let mask: Bitmap = [false, true, true].into_iter().collect(); + let f = a.filter(&mask).unwrap(); + assert_eq!(f.null_count(), 1); + assert!(f.value(0).is_null()); + } + + #[test] + fn appending_an_int_to_a_float_column_widens_it() { + let mut b = ArrayBuilder::new(DataType::Float64); + b.push(Value::Int64(3)).unwrap(); + b.push(Value::Float64(0.5)).unwrap(); + let a = b.finish().unwrap(); + assert_eq!(a.value(0), Value::Float64(3.0)); + assert_eq!(a.data_type(), DataType::Float64); + } + + #[test] + fn appending_the_wrong_type_is_rejected() { + let mut b = ArrayBuilder::new(DataType::Int64); + assert!(b.push(Value::Utf8("x".into())).is_err()); + let mut b = ArrayBuilder::new(DataType::Boolean); + assert!(b.push(Value::Int64(1)).is_err()); + } + + #[test] + fn a_null_column_reports_every_row_null() { + let a = Array::nulls(DataType::Utf8, 3).unwrap(); + assert_eq!(a.len(), 3); + assert_eq!(a.null_count(), 3); + assert!(a.value(2).is_null()); + } + + #[test] + fn boolean_columns_pack_into_a_bitmap() { + let mut b = ArrayBuilder::new(DataType::Boolean); + for i in 0..100 { + b.push(Value::Boolean(i % 2 == 0)).unwrap(); + } + let a = b.finish().unwrap(); + assert_eq!(a.len(), 100); + assert_eq!(a.value(4), Value::Boolean(true)); + assert_eq!(a.value(5), Value::Boolean(false)); + match a.data() { + ArrayData::Boolean(bm) => assert_eq!(bm.count_set(), 50), + other => panic!("expected a bitmap, got {other:?}"), + } + } + + #[test] + fn a_mismatched_buffer_length_is_caught_at_construction() { + let err = Array::new(ArrayData::Int64(vec![1, 2]), None, 3).unwrap_err(); + assert!(matches!(err, Error::Internal(_))); + let err = Array::new( + ArrayData::Int64(vec![1, 2]), + Some(Bitmap::filled(5, true)), + 2, + ) + .unwrap_err(); + assert!(matches!(err, Error::Internal(_))); + } + + #[test] + fn from_values_and_iter_are_inverses() { + let vals = vec![Value::Int64(1), Value::Null, Value::Int64(3)]; + let a = Array::from_values(DataType::Int64, &vals).unwrap(); + assert_eq!(a.iter().collect::>(), vals); + assert!(!a.is_empty()); + assert!(Array::from_values(DataType::Int64, &[]).unwrap().is_empty()); + } + + #[test] + fn builder_reports_its_own_length_as_it_fills() { + let mut b = ArrayBuilder::new(DataType::Utf8); + assert!(b.is_empty()); + b.reserve(10); + b.push(Value::Utf8("a".into())).unwrap(); + assert_eq!(b.len(), 1); + assert_eq!(b.data_type(), DataType::Utf8); + } +} diff --git a/crates/qf-storage/src/batch.rs b/crates/qf-storage/src/batch.rs new file mode 100644 index 0000000..3d37cde --- /dev/null +++ b/crates/qf-storage/src/batch.rs @@ -0,0 +1,410 @@ +use crate::array::{Array, ArrayBuilder}; +use crate::bitmap::Bitmap; +use qf_common::{Error, Result, Schema, Value}; +use std::sync::Arc; + +/// The unit of work that moves between operators: a schema plus one array per +/// column, all of the same length. +/// +/// Everything downstream is written to consume batches rather than rows. A +/// filter evaluates its predicate over a whole column and hands on a batch; +/// nothing in the engine has a "next row" method, because per-row dispatch +/// through the expression tree is the cost the vectorised model exists to +/// amortise. +#[derive(Debug, Clone, PartialEq)] +pub struct RecordBatch { + schema: Arc, + columns: Vec, + num_rows: usize, +} + +/// Rows per batch. Big enough to amortise per-batch overhead, small enough that +/// a batch's working set stays in L2 while an operator runs over it. +pub const DEFAULT_BATCH_SIZE: usize = 8192; + +impl RecordBatch { + pub fn try_new(schema: Arc, columns: Vec) -> Result { + if schema.len() != columns.len() { + return Err(Error::internal(format!( + "schema has {} fields but {} columns were supplied", + schema.len(), + columns.len() + ))); + } + let num_rows = columns.first().map_or(0, Array::len); + for (i, col) in columns.iter().enumerate() { + if col.len() != num_rows { + return Err(Error::internal(format!( + "column {i} has {} rows, expected {num_rows}", + col.len() + ))); + } + let expected = schema.field(i)?.data_type; + if col.data_type() != expected { + return Err(Error::internal(format!( + "column {i} is {} but the schema says {expected}", + col.data_type() + ))); + } + } + Ok(RecordBatch { + schema, + columns, + num_rows, + }) + } + + /// A batch with the right schema and no rows, which operators emit when a + /// filter rejects everything. + pub fn empty(schema: Arc) -> Result { + let columns = schema + .fields() + .iter() + .map(|f| Array::nulls(f.data_type, 0)) + .collect::>>()?; + RecordBatch::try_new(schema, columns) + } + + pub fn schema(&self) -> &Arc { + &self.schema + } + + pub fn columns(&self) -> &[Array] { + &self.columns + } + + pub fn column(&self, i: usize) -> Result<&Array> { + self.columns + .get(i) + .ok_or_else(|| Error::internal(format!("no column at index {i}"))) + } + + pub fn column_by_name(&self, name: &str) -> Result<&Array> { + self.column(self.schema.index_of(name)?) + } + + pub fn num_rows(&self) -> usize { + self.num_rows + } + + pub fn num_columns(&self) -> usize { + self.columns.len() + } + + pub fn is_empty(&self) -> bool { + self.num_rows == 0 + } + + /// Materialises one row as scalars. Used by the CLI's printer and by tests; + /// the operators never call it. + pub fn row(&self, i: usize) -> Vec { + self.columns.iter().map(|c| c.value(i)).collect() + } + + pub fn rows(&self) -> impl Iterator> + '_ { + (0..self.num_rows).map(move |i| self.row(i)) + } + + pub fn filter(&self, mask: &Bitmap) -> Result { + let columns = self + .columns + .iter() + .map(|c| c.filter(mask)) + .collect::>>()?; + // A batch with zero columns still has to remember how many rows + // survived, or `SELECT count(*)` after a filter counts the wrong thing. + if columns.is_empty() { + return Ok(RecordBatch { + schema: Arc::clone(&self.schema), + columns, + num_rows: mask.count_set(), + }); + } + RecordBatch::try_new(Arc::clone(&self.schema), columns) + } + + pub fn take(&self, indices: &[usize]) -> Result { + let columns = self + .columns + .iter() + .map(|c| c.take(indices)) + .collect::>>()?; + if columns.is_empty() { + return Ok(RecordBatch { + schema: Arc::clone(&self.schema), + columns, + num_rows: indices.len(), + }); + } + RecordBatch::try_new(Arc::clone(&self.schema), columns) + } + + /// Keeps `columns` in the order given, producing the matching schema. + pub fn project(&self, indices: &[usize]) -> Result { + let schema = Arc::new(self.schema.project(indices)?); + let mut columns = Vec::with_capacity(indices.len()); + for &i in indices { + columns.push(self.column(i)?.clone()); + } + if columns.is_empty() { + return Ok(RecordBatch { + schema, + columns, + num_rows: self.num_rows, + }); + } + RecordBatch::try_new(schema, columns) + } + + /// Glues two batches side by side. This is how a join builds its output + /// once both sides have been gathered to matching row counts. + pub fn hstack(left: &RecordBatch, right: &RecordBatch) -> Result { + if left.num_rows != right.num_rows { + return Err(Error::internal(format!( + "cannot stack a {}-row batch beside a {}-row batch", + left.num_rows, right.num_rows + ))); + } + let schema = Arc::new(Schema::join(&left.schema, &right.schema)); + let mut columns = left.columns.clone(); + columns.extend(right.columns.iter().cloned()); + if columns.is_empty() { + return Ok(RecordBatch { + schema, + columns, + num_rows: left.num_rows, + }); + } + RecordBatch::try_new(schema, columns) + } + + /// Concatenates batches that share a schema, which is what a collecting + /// operator does before it sorts or returns. + pub fn concat(schema: Arc, batches: &[RecordBatch]) -> Result { + if batches.is_empty() { + return RecordBatch::empty(schema); + } + let mut builders: Vec = schema + .fields() + .iter() + .map(|f| ArrayBuilder::new(f.data_type)) + .collect(); + for b in batches { + if b.schema.as_ref() != schema.as_ref() { + return Err(Error::internal( + "cannot concatenate batches with different schemas".to_string(), + )); + } + for (ci, builder) in builders.iter_mut().enumerate() { + let col = b.column(ci)?; + for r in 0..b.num_rows { + builder.push_from(col, r); + } + } + } + let columns = builders + .into_iter() + .map(ArrayBuilder::finish) + .collect::>>()?; + if columns.is_empty() { + let num_rows = batches.iter().map(|b| b.num_rows).sum(); + return Ok(RecordBatch { + schema, + columns, + num_rows, + }); + } + RecordBatch::try_new(schema, columns) + } + + /// Splits into batches of at most `size` rows, restoring the batch size + /// after an operator (a join, say) has produced something much larger. + pub fn chunks(&self, size: usize) -> Result> { + if size == 0 { + return Err(Error::internal("batch size must be positive".to_string())); + } + if self.num_rows <= size { + return Ok(vec![self.clone()]); + } + let mut out = Vec::new(); + let mut start = 0; + while start < self.num_rows { + let end = (start + size).min(self.num_rows); + let idx: Vec = (start..end).collect(); + out.push(self.take(&idx)?); + start = end; + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use qf_common::{DataType, Field}; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ])) + } + + fn batch(ids: &[i64], names: &[&str]) -> RecordBatch { + let ids: Vec = ids.iter().map(|i| Value::Int64(*i)).collect(); + let names: Vec = names.iter().map(|s| Value::Utf8(s.to_string())).collect(); + RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &ids).unwrap(), + Array::from_values(DataType::Utf8, &names).unwrap(), + ], + ) + .unwrap() + } + + #[test] + fn a_batch_knows_its_shape() { + let b = batch(&[1, 2, 3], &["a", "b", "c"]); + assert_eq!(b.num_rows(), 3); + assert_eq!(b.num_columns(), 2); + assert!(!b.is_empty()); + assert_eq!(b.row(1), vec![Value::Int64(2), Value::Utf8("b".into())]); + } + + #[test] + fn columns_of_unequal_length_are_rejected() { + let err = RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &[Value::Int64(1)]).unwrap(), + Array::from_values(DataType::Utf8, &[]).unwrap(), + ], + ) + .unwrap_err(); + assert!(matches!(err, Error::Internal(_))); + } + + #[test] + fn a_column_whose_type_contradicts_the_schema_is_rejected() { + let err = RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Utf8, &[Value::Utf8("x".into())]).unwrap(), + Array::from_values(DataType::Utf8, &[Value::Utf8("y".into())]).unwrap(), + ], + ) + .unwrap_err(); + assert!(err.to_string().contains("schema says")); + } + + #[test] + fn a_column_count_mismatch_is_rejected() { + assert!(RecordBatch::try_new(schema(), vec![]).is_err()); + } + + #[test] + fn filter_and_take_agree() { + let b = batch(&[1, 2, 3, 4], &["a", "b", "c", "d"]); + let mask: Bitmap = [false, true, false, true].into_iter().collect(); + assert_eq!(b.filter(&mask).unwrap(), b.take(&[1, 3]).unwrap()); + assert_eq!(b.filter(&mask).unwrap().num_rows(), 2); + } + + #[test] + fn projection_reorders_columns_and_the_schema_with_them() { + let b = batch(&[1], &["a"]); + let p = b.project(&[1, 0]).unwrap(); + assert_eq!(p.schema().field(0).unwrap().name, "name"); + assert_eq!(p.row(0), vec![Value::Utf8("a".into()), Value::Int64(1)]); + } + + #[test] + fn a_zero_column_projection_still_remembers_its_row_count() { + // `SELECT count(*) FROM t WHERE ...` projects away every column; if the + // row count went with them the count would come back as zero. + let b = batch(&[1, 2, 3], &["a", "b", "c"]); + let p = b.project(&[]).unwrap(); + assert_eq!(p.num_columns(), 0); + assert_eq!(p.num_rows(), 3); + + let mask: Bitmap = [true, false, true].into_iter().collect(); + assert_eq!(p.filter(&mask).unwrap().num_rows(), 2); + assert_eq!(p.take(&[0, 0, 1, 2]).unwrap().num_rows(), 4); + } + + #[test] + fn hstack_joins_two_batches_side_by_side() { + let l = batch(&[1, 2], &["a", "b"]); + let r = batch(&[9, 8], &["x", "y"]); + let s = RecordBatch::hstack(&l, &r).unwrap(); + assert_eq!(s.num_columns(), 4); + assert_eq!(s.num_rows(), 2); + assert_eq!(s.row(0)[2], Value::Int64(9)); + } + + #[test] + fn hstack_refuses_mismatched_row_counts() { + let l = batch(&[1, 2], &["a", "b"]); + let r = batch(&[9], &["x"]); + assert!(RecordBatch::hstack(&l, &r).is_err()); + } + + #[test] + fn concat_preserves_order_across_batches() { + let a = batch(&[1, 2], &["a", "b"]); + let b = batch(&[3], &["c"]); + let c = RecordBatch::concat(schema(), &[a, b]).unwrap(); + assert_eq!(c.num_rows(), 3); + assert_eq!( + c.column(0).unwrap().iter().collect::>(), + vec![Value::Int64(1), Value::Int64(2), Value::Int64(3)] + ); + } + + #[test] + fn concat_of_nothing_yields_an_empty_batch_with_the_right_schema() { + let c = RecordBatch::concat(schema(), &[]).unwrap(); + assert!(c.is_empty()); + assert_eq!(c.num_columns(), 2); + assert_eq!(RecordBatch::empty(schema()).unwrap(), c); + } + + #[test] + fn concat_rejects_a_batch_with_a_different_schema() { + let other = Arc::new(Schema::new(vec![Field::new("z", DataType::Int64, false)])); + let odd = RecordBatch::empty(other).unwrap(); + assert!(RecordBatch::concat(schema(), &[batch(&[1], &["a"]), odd]).is_err()); + } + + #[test] + fn chunks_splits_a_large_batch_and_loses_no_rows() { + let ids: Vec = (0..2500).collect(); + let names: Vec = ids.iter().map(|i| format!("n{i}")).collect(); + let refs: Vec<&str> = names.iter().map(String::as_str).collect(); + let b = batch(&ids, &refs); + let chunks = b.chunks(1000).unwrap(); + assert_eq!(chunks.len(), 3); + assert_eq!(chunks[2].num_rows(), 500); + let total: usize = chunks.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(total, 2500); + assert_eq!(RecordBatch::concat(schema(), &chunks).unwrap(), b); + } + + #[test] + fn a_batch_smaller_than_the_chunk_size_passes_through_whole() { + let b = batch(&[1], &["a"]); + assert_eq!(b.chunks(DEFAULT_BATCH_SIZE).unwrap().len(), 1); + assert!(b.chunks(0).is_err()); + } + + #[test] + fn columns_can_be_addressed_by_name_or_index() { + let b = batch(&[7], &["a"]); + assert_eq!(b.column_by_name("ID").unwrap().value(0), Value::Int64(7)); + assert!(b.column(5).is_err()); + assert!(b.column_by_name("nope").is_err()); + assert_eq!(b.columns().len(), 2); + assert_eq!(b.rows().count(), 1); + } +} diff --git a/crates/qf-storage/src/bitmap.rs b/crates/qf-storage/src/bitmap.rs new file mode 100644 index 0000000..7030759 --- /dev/null +++ b/crates/qf-storage/src/bitmap.rs @@ -0,0 +1,275 @@ +/// A packed bit vector, one bit per row. +/// +/// Both the validity (null) masks and the selection vectors produced by filters +/// are bitmaps. Packing them costs a shift and a mask per access but makes a +/// filter over a million rows touch 128 KB instead of 1 MB, and lets `and` +/// combine two predicates 64 rows at a time. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct Bitmap { + words: Vec, + len: usize, +} + +impl Bitmap { + pub fn new() -> Self { + Bitmap::default() + } + + pub fn with_capacity(cap: usize) -> Self { + Bitmap { + words: Vec::with_capacity(cap.div_ceil(64)), + len: 0, + } + } + + /// A bitmap of `len` bits all set to `value`. + pub fn filled(len: usize, value: bool) -> Self { + let mut words = vec![if value { u64::MAX } else { 0 }; len.div_ceil(64)]; + if value { + trim_tail(&mut words, len); + } + Bitmap { words, len } + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn push(&mut self, value: bool) { + if self.len % 64 == 0 { + self.words.push(0); + } + if value { + let (w, b) = (self.len / 64, self.len % 64); + self.words[w] |= 1u64 << b; + } + self.len += 1; + } + + /// Reads bit `i`. Out-of-range reads return false rather than panicking: + /// callers reach for this in tight loops where a bounds check per row is + /// exactly the cost the packed representation is meant to avoid. + pub fn get(&self, i: usize) -> bool { + if i >= self.len { + return false; + } + self.words[i / 64] & (1u64 << (i % 64)) != 0 + } + + pub fn set(&mut self, i: usize, value: bool) { + debug_assert!(i < self.len); + let (w, b) = (i / 64, i % 64); + if value { + self.words[w] |= 1u64 << b; + } else { + self.words[w] &= !(1u64 << b); + } + } + + /// Number of set bits — the selectivity of a filter, or the non-null count + /// of a column. `count_ones` is a single instruction per 64 rows. + pub fn count_set(&self) -> usize { + self.words.iter().map(|w| w.count_ones() as usize).sum() + } + + pub fn count_unset(&self) -> usize { + self.len - self.count_set() + } + + pub fn all_set(&self) -> bool { + self.count_set() == self.len + } + + pub fn any_set(&self) -> bool { + self.words.iter().any(|w| *w != 0) + } + + /// Word-at-a-time intersection, used to combine conjunctive predicates. + pub fn and(&self, other: &Bitmap) -> Bitmap { + let len = self.len.min(other.len); + let mut words: Vec = self + .words + .iter() + .zip(other.words.iter()) + .map(|(a, b)| a & b) + .collect(); + words.truncate(len.div_ceil(64)); + trim_tail(&mut words, len); + Bitmap { words, len } + } + + pub fn or(&self, other: &Bitmap) -> Bitmap { + let len = self.len.max(other.len); + let mut words = vec![0u64; len.div_ceil(64)]; + for (i, w) in words.iter_mut().enumerate() { + let a = self.words.get(i).copied().unwrap_or(0); + let b = other.words.get(i).copied().unwrap_or(0); + *w = a | b; + } + trim_tail(&mut words, len); + Bitmap { words, len } + } + + pub fn not(&self) -> Bitmap { + let mut words: Vec = self.words.iter().map(|w| !w).collect(); + trim_tail(&mut words, self.len); + Bitmap { + words, + len: self.len, + } + } + + /// The positions of the set bits, which is what the take-based operators + /// (join, sort) consume. + pub fn set_indices(&self) -> Vec { + let mut out = Vec::with_capacity(self.count_set()); + for (wi, mut word) in self.words.iter().copied().enumerate() { + while word != 0 { + let bit = word.trailing_zeros() as usize; + let idx = wi * 64 + bit; + if idx < self.len { + out.push(idx); + } + word &= word - 1; + } + } + out + } + + pub fn iter(&self) -> impl Iterator + '_ { + (0..self.len).map(move |i| self.get(i)) + } +} + +impl FromIterator for Bitmap { + fn from_iter>(iter: T) -> Self { + let mut b = Bitmap::new(); + for v in iter { + b.push(v); + } + b + } +} + +/// Clears the bits past `len` in the final word so `count_ones` cannot count +/// padding. Every operation that fills whole words has to call this. +fn trim_tail(words: &mut [u64], len: usize) { + let rem = len % 64; + if rem != 0 { + if let Some(last) = words.last_mut() { + *last &= (1u64 << rem) - 1; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn push_and_get_round_trip_across_a_word_boundary() { + let pattern: Vec = (0..130).map(|i| i % 3 == 0).collect(); + let bm: Bitmap = pattern.iter().copied().collect(); + assert_eq!(bm.len(), 130); + for (i, expected) in pattern.iter().enumerate() { + assert_eq!(bm.get(i), *expected, "bit {i}"); + } + } + + #[test] + fn filled_does_not_count_padding_bits() { + // 70 bits spans two words; the upper 58 bits of the second must not + // show up in count_set. + let bm = Bitmap::filled(70, true); + assert_eq!(bm.count_set(), 70); + assert!(bm.all_set()); + assert_eq!(Bitmap::filled(70, false).count_set(), 0); + } + + #[test] + fn and_or_not_agree_with_bit_by_bit_evaluation() { + let a: Bitmap = (0..100).map(|i| i % 2 == 0).collect(); + let b: Bitmap = (0..100).map(|i| i % 3 == 0).collect(); + let and = a.and(&b); + let or = a.or(&b); + let not = a.not(); + for i in 0..100 { + assert_eq!(and.get(i), a.get(i) && b.get(i)); + assert_eq!(or.get(i), a.get(i) || b.get(i)); + assert_eq!(not.get(i), !a.get(i)); + } + assert_eq!(not.len(), 100); + } + + #[test] + fn negation_does_not_leak_padding_into_the_count() { + let bm = Bitmap::filled(5, false); + assert_eq!(bm.not().count_set(), 5); + } + + #[test] + fn set_indices_lists_exactly_the_set_positions() { + let mut bm = Bitmap::filled(200, false); + for i in [0, 63, 64, 65, 199] { + bm.set(i, true); + } + assert_eq!(bm.set_indices(), vec![0, 63, 64, 65, 199]); + assert_eq!(bm.count_set(), 5); + assert_eq!(bm.count_unset(), 195); + } + + #[test] + fn set_can_clear_a_bit_again() { + let mut bm = Bitmap::filled(10, true); + bm.set(4, false); + assert!(!bm.get(4)); + assert_eq!(bm.count_set(), 9); + } + + #[test] + fn reads_past_the_end_are_false_rather_than_a_panic() { + let bm = Bitmap::filled(3, true); + assert!(!bm.get(3)); + assert!(!bm.get(9999)); + } + + #[test] + fn empty_bitmap_behaves() { + let bm = Bitmap::new(); + assert!(bm.is_empty()); + assert!(!bm.any_set()); + assert!(bm.all_set()); // vacuously + assert!(bm.set_indices().is_empty()); + assert_eq!(Bitmap::with_capacity(128).len(), 0); + } + + #[test] + fn and_truncates_to_the_shorter_input() { + let a = Bitmap::filled(100, true); + let b = Bitmap::filled(30, true); + let r = a.and(&b); + assert_eq!(r.len(), 30); + assert_eq!(r.count_set(), 30); + } + + #[test] + fn or_extends_to_the_longer_input() { + let a = Bitmap::filled(10, true); + let b = Bitmap::filled(100, false); + let r = a.or(&b); + assert_eq!(r.len(), 100); + assert_eq!(r.count_set(), 10); + assert!(r.any_set()); + } + + #[test] + fn iter_yields_every_bit_in_order() { + let bits = vec![true, false, true]; + let bm: Bitmap = bits.iter().copied().collect(); + assert_eq!(bm.iter().collect::>(), bits); + } +} diff --git a/crates/qf-storage/src/bytes.rs b/crates/qf-storage/src/bytes.rs new file mode 100644 index 0000000..9648dc0 --- /dev/null +++ b/crates/qf-storage/src/bytes.rs @@ -0,0 +1,272 @@ +use qf_common::{Error, Result}; + +/// Little-endian byte writer with LEB128 varints. +/// +/// Fixed-width integers everywhere would make the on-disk format simpler, but +/// most of what gets written here is small — run lengths, dictionary indices, +/// string lengths — and a varint spends one byte on those instead of eight. +#[derive(Debug, Default)] +pub struct ByteWriter { + buf: Vec, +} + +impl ByteWriter { + pub fn new() -> Self { + ByteWriter::default() + } + + pub fn len(&self) -> usize { + self.buf.len() + } + + pub fn is_empty(&self) -> bool { + self.buf.is_empty() + } + + pub fn into_bytes(self) -> Vec { + self.buf + } + + pub fn as_slice(&self) -> &[u8] { + &self.buf + } + + pub fn u8(&mut self, v: u8) { + self.buf.push(v); + } + + pub fn u32(&mut self, v: u32) { + self.buf.extend_from_slice(&v.to_le_bytes()); + } + + pub fn u64(&mut self, v: u64) { + self.buf.extend_from_slice(&v.to_le_bytes()); + } + + pub fn f64(&mut self, v: f64) { + self.buf.extend_from_slice(&v.to_le_bytes()); + } + + pub fn uvarint(&mut self, mut v: u64) { + loop { + let mut byte = (v & 0x7f) as u8; + v >>= 7; + if v != 0 { + byte |= 0x80; + } + self.buf.push(byte); + if v == 0 { + break; + } + } + } + + /// Zigzag then varint, so that small negative numbers stay small. + pub fn ivarint(&mut self, v: i64) { + self.uvarint(((v << 1) ^ (v >> 63)) as u64); + } + + pub fn bytes(&mut self, b: &[u8]) { + self.uvarint(b.len() as u64); + self.buf.extend_from_slice(b); + } + + pub fn string(&mut self, s: &str) { + self.bytes(s.as_bytes()); + } +} + +/// The reading half. Every accessor is bounds-checked and returns a +/// `Storage` error rather than panicking: a truncated or corrupt file is an +/// expected failure mode, not a bug. +#[derive(Debug)] +pub struct ByteReader<'a> { + buf: &'a [u8], + pos: usize, +} + +impl<'a> ByteReader<'a> { + pub fn new(buf: &'a [u8]) -> Self { + ByteReader { buf, pos: 0 } + } + + pub fn position(&self) -> usize { + self.pos + } + + pub fn remaining(&self) -> usize { + self.buf.len() - self.pos + } + + pub fn is_done(&self) -> bool { + self.pos >= self.buf.len() + } + + fn take(&mut self, n: usize) -> Result<&'a [u8]> { + if self.pos + n > self.buf.len() { + return Err(Error::storage(format!( + "unexpected end of buffer: wanted {n} bytes at offset {}, {} remain", + self.pos, + self.remaining() + ))); + } + let s = &self.buf[self.pos..self.pos + n]; + self.pos += n; + Ok(s) + } + + pub fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + pub fn u32(&mut self) -> Result { + let b = self.take(4)?; + Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]])) + } + + pub fn u64(&mut self) -> Result { + let b = self.take(8)?; + Ok(u64::from_le_bytes([ + b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], + ])) + } + + pub fn f64(&mut self) -> Result { + Ok(f64::from_bits(self.u64()?)) + } + + pub fn uvarint(&mut self) -> Result { + let mut result: u64 = 0; + let mut shift = 0; + loop { + let byte = self.u8()?; + if shift >= 64 { + return Err(Error::storage("varint overflows 64 bits".to_string())); + } + result |= u64::from(byte & 0x7f) << shift; + if byte & 0x80 == 0 { + break; + } + shift += 7; + } + Ok(result) + } + + pub fn ivarint(&mut self) -> Result { + let u = self.uvarint()?; + Ok(((u >> 1) as i64) ^ -((u & 1) as i64)) + } + + pub fn bytes(&mut self) -> Result<&'a [u8]> { + let n = self.uvarint()? as usize; + self.take(n) + } + + pub fn string(&mut self) -> Result { + let b = self.bytes()?; + String::from_utf8(b.to_vec()) + .map_err(|_| Error::storage("column chunk holds invalid UTF-8".to_string())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn primitives_round_trip() { + let mut w = ByteWriter::new(); + w.u8(7); + w.u32(70_000); + w.u64(u64::MAX); + w.f64(-1.5); + w.string("hello"); + assert!(!w.is_empty()); + let buf = w.into_bytes(); + + let mut r = ByteReader::new(&buf); + assert_eq!(r.u8().unwrap(), 7); + assert_eq!(r.u32().unwrap(), 70_000); + assert_eq!(r.u64().unwrap(), u64::MAX); + assert_eq!(r.f64().unwrap(), -1.5); + assert_eq!(r.string().unwrap(), "hello"); + assert!(r.is_done()); + } + + #[test] + fn varints_round_trip_across_the_whole_range() { + let cases = [0i64, 1, -1, 63, 64, -64, i64::MAX, i64::MIN, 1 << 40]; + let mut w = ByteWriter::new(); + for c in cases { + w.ivarint(c); + } + let buf = w.into_bytes(); + let mut r = ByteReader::new(&buf); + for c in cases { + assert_eq!(r.ivarint().unwrap(), c, "value {c}"); + } + } + + #[test] + fn small_negative_numbers_cost_one_byte_thanks_to_zigzag() { + let mut w = ByteWriter::new(); + w.ivarint(-1); + assert_eq!(w.len(), 1); + let mut w = ByteWriter::new(); + w.ivarint(-1_000_000); + assert!(w.len() < 8); + } + + #[test] + fn unsigned_varints_round_trip() { + let mut w = ByteWriter::new(); + for v in [0u64, 127, 128, 300, u64::MAX] { + w.uvarint(v); + } + let buf = w.into_bytes(); + let mut r = ByteReader::new(&buf); + for v in [0u64, 127, 128, 300, u64::MAX] { + assert_eq!(r.uvarint().unwrap(), v); + } + } + + #[test] + fn a_truncated_buffer_is_a_storage_error_not_a_panic() { + let mut w = ByteWriter::new(); + w.u64(1); + let buf = w.into_bytes(); + let mut r = ByteReader::new(&buf[..3]); + let err = r.u64().unwrap_err(); + assert!(matches!(err, Error::Storage(_))); + assert!(err.to_string().contains("unexpected end of buffer")); + } + + #[test] + fn a_varint_that_never_terminates_is_rejected() { + let buf = vec![0xff; 32]; + let mut r = ByteReader::new(&buf); + assert!(r.uvarint().is_err()); + } + + #[test] + fn invalid_utf8_in_a_string_field_is_reported() { + let mut w = ByteWriter::new(); + w.bytes(&[0xff, 0xfe]); + let buf = w.into_bytes(); + let mut r = ByteReader::new(&buf); + assert!(r.string().unwrap_err().to_string().contains("UTF-8")); + } + + #[test] + fn the_reader_tracks_its_own_position() { + let mut w = ByteWriter::new(); + w.u32(1); + let buf = w.into_bytes(); + let mut r = ByteReader::new(&buf); + assert_eq!(r.remaining(), 4); + r.u32().unwrap(); + assert_eq!(r.position(), 4); + assert_eq!(r.remaining(), 0); + assert!(ByteWriter::new().as_slice().is_empty()); + } +} diff --git a/crates/qf-storage/src/catalog.rs b/crates/qf-storage/src/catalog.rs new file mode 100644 index 0000000..69144a9 --- /dev/null +++ b/crates/qf-storage/src/catalog.rs @@ -0,0 +1,324 @@ +use crate::batch::RecordBatch; +use crate::format::QfcReader; +use crate::stats::{ColumnStats, TableStats}; +use qf_common::{Error, Result, Schema}; +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +/// Where a table's rows actually live. +#[derive(Debug, Clone)] +pub enum TableSource { + /// Batches held in memory — what `CREATE TABLE ... VALUES` and the test + /// fixtures produce. + Memory(Vec), + /// A `.qfc` file, read lazily one row group at a time. + File(PathBuf), +} + +#[derive(Debug, Clone)] +pub struct Table { + pub name: String, + pub schema: Arc, + pub source: TableSource, + pub stats: TableStats, +} + +impl Table { + pub fn row_count(&self) -> usize { + self.stats.row_count + } +} + +/// Name-to-table map, and the only place the planner learns what exists. +/// +/// Lookups are case-insensitive, which is what SQL expects, and the map is +/// ordered so `\dt` and `SHOW TABLES` list tables the same way every run — +/// a hash map's iteration order would make the CLI's output non-deterministic +/// and its tests flaky. +#[derive(Debug, Clone, Default)] +pub struct Catalog { + tables: BTreeMap, +} + +impl Catalog { + pub fn new() -> Catalog { + Catalog::default() + } + + fn key(name: &str) -> String { + name.to_ascii_lowercase() + } + + /// Registers in-memory batches, deriving the table's statistics from them. + pub fn register_batches( + &mut self, + name: impl Into, + schema: Arc, + batches: Vec, + ) -> Result<()> { + let name = name.into(); + for b in &batches { + if b.schema().as_ref() != schema.as_ref() { + return Err(Error::plan(format!( + "a batch registered for `{name}` does not match the table schema" + ))); + } + } + let stats = stats_from_batches(&schema, &batches)?; + self.insert(Table { + name: name.clone(), + schema, + source: TableSource::Memory(batches), + stats, + }) + } + + /// Registers a `.qfc` file. The schema and statistics come from the + /// footer, so registering a table reads kilobytes rather than the file. + pub fn register_file(&mut self, name: impl Into, path: impl AsRef) -> Result<()> { + let name = name.into(); + let reader = QfcReader::open(path.as_ref())?; + let schema = Arc::clone(reader.schema()); + let stats = stats_from_file(&reader); + self.insert(Table { + name, + schema, + source: TableSource::File(path.as_ref().to_path_buf()), + stats, + }) + } + + fn insert(&mut self, table: Table) -> Result<()> { + let key = Catalog::key(&table.name); + if self.tables.contains_key(&key) { + return Err(Error::plan(format!( + "table `{}` is already registered", + table.name + ))); + } + self.tables.insert(key, table); + Ok(()) + } + + /// Registers a table, replacing any existing one of the same name. + pub fn replace(&mut self, table: Table) { + self.tables.insert(Catalog::key(&table.name), table); + } + + pub fn get(&self, name: &str) -> Result<&Table> { + self.tables + .get(&Catalog::key(name)) + .ok_or_else(|| Error::plan(format!("no such table `{name}`"))) + } + + pub fn contains(&self, name: &str) -> bool { + self.tables.contains_key(&Catalog::key(name)) + } + + pub fn drop_table(&mut self, name: &str) -> Result { + self.tables + .remove(&Catalog::key(name)) + .ok_or_else(|| Error::plan(format!("no such table `{name}`"))) + } + + pub fn table_names(&self) -> Vec<&str> { + self.tables.values().map(|t| t.name.as_str()).collect() + } + + pub fn tables(&self) -> impl Iterator { + self.tables.values() + } + + pub fn len(&self) -> usize { + self.tables.len() + } + + pub fn is_empty(&self) -> bool { + self.tables.is_empty() + } +} + +fn stats_from_batches(schema: &Arc, batches: &[RecordBatch]) -> Result { + let mut columns = vec![ColumnStats::empty(); schema.len()]; + let mut row_count = 0; + for b in batches { + row_count += b.num_rows(); + for (i, col) in columns.iter_mut().enumerate() { + *col = ColumnStats::merge(col, &ColumnStats::from_array(b.column(i)?)); + } + } + Ok(TableStats { row_count, columns }) +} + +/// Folds the per-row-group zone maps in the footer into one table-level +/// summary. No column data is read. +/// +/// `min`, `max`, `row_count` and `null_count` come out exact. `distinct_count` +/// does not — row groups can hold overlapping values and the footer cannot say +/// which — so it is an upper bound. A sketch per chunk would fix that; the +/// planner only uses distinct counts to order joins, and an over-estimate there +/// costs a suboptimal plan rather than a wrong answer. +fn stats_from_file(reader: &QfcReader) -> TableStats { + let meta = reader.meta(); + let mut columns = vec![ColumnStats::empty(); meta.schema.len()]; + for group in &meta.row_groups { + for (i, chunk) in group.columns.iter().enumerate() { + columns[i] = ColumnStats::merge(&columns[i], &chunk.stats); + } + } + TableStats { + row_count: meta.num_rows(), + columns, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::array::Array; + use crate::format::write_file; + use qf_common::{DataType, Field, Value}; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("tag", DataType::Utf8, false), + ])) + } + + fn batch(ids: &[i64]) -> RecordBatch { + let idv: Vec = ids.iter().map(|i| Value::Int64(*i)).collect(); + let tags: Vec = ids + .iter() + .map(|i| Value::Utf8(format!("t{}", i % 2))) + .collect(); + RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &idv).unwrap(), + Array::from_values(DataType::Utf8, &tags).unwrap(), + ], + ) + .unwrap() + } + + #[test] + fn tables_are_found_case_insensitively() { + let mut c = Catalog::new(); + c.register_batches("Orders", schema(), vec![batch(&[1, 2])]) + .unwrap(); + assert!(c.get("orders").is_ok()); + assert!(c.get("ORDERS").is_ok()); + assert!(c.contains("Orders")); + assert_eq!(c.get("orders").unwrap().name, "Orders"); + } + + #[test] + fn statistics_are_derived_when_batches_are_registered() { + let mut c = Catalog::new(); + c.register_batches("t", schema(), vec![batch(&[3, 1, 2]), batch(&[9])]) + .unwrap(); + let t = c.get("t").unwrap(); + assert_eq!(t.row_count(), 4); + let id_stats = t.stats.column(0).unwrap(); + assert_eq!(id_stats.min, Value::Int64(1)); + assert_eq!(id_stats.max, Value::Int64(9)); + assert_eq!(id_stats.row_count, 4); + } + + #[test] + fn registering_the_same_name_twice_is_refused() { + let mut c = Catalog::new(); + c.register_batches("t", schema(), vec![]).unwrap(); + let err = c.register_batches("T", schema(), vec![]).unwrap_err(); + assert!(err.to_string().contains("already registered")); + } + + #[test] + fn replace_overwrites_an_existing_registration() { + let mut c = Catalog::new(); + c.register_batches("t", schema(), vec![batch(&[1])]) + .unwrap(); + let replacement = Table { + name: "t".into(), + schema: schema(), + source: TableSource::Memory(vec![batch(&[1, 2, 3])]), + stats: TableStats::default(), + }; + c.replace(replacement); + assert_eq!(c.len(), 1); + match &c.get("t").unwrap().source { + TableSource::Memory(b) => assert_eq!(b[0].num_rows(), 3), + other => panic!("expected memory batches, got {other:?}"), + } + } + + #[test] + fn a_batch_that_contradicts_the_declared_schema_is_refused() { + let other = Arc::new(Schema::new(vec![Field::new("z", DataType::Int64, false)])); + let mut c = Catalog::new(); + assert!(c + .register_batches("t", other, vec![batch(&[1])]) + .unwrap_err() + .to_string() + .contains("does not match")); + } + + #[test] + fn file_backed_tables_take_their_schema_and_stats_from_the_footer() { + let mut p = std::env::temp_dir(); + p.push(format!("qf-catalog-{}.qfc", std::process::id())); + let rows: Vec = (0..500).collect(); + write_file(&p, schema(), &[batch(&rows)], 100).unwrap(); + + let mut c = Catalog::new(); + c.register_file("sales", &p).unwrap(); + let t = c.get("sales").unwrap(); + assert_eq!(t.row_count(), 500); + assert_eq!(t.schema.as_ref(), schema().as_ref()); + assert_eq!(t.stats.column(0).unwrap().min, Value::Int64(0)); + assert_eq!(t.stats.column(0).unwrap().max, Value::Int64(499)); + // `tag` really holds two distinct values, but the merged estimate is + // an upper bound: each of the five row groups reports 2 and the merge + // cannot tell that they are the same two. It stays bounded by the row + // count, which is the property the join-ordering rule relies on. + let tag = t.stats.column(1).unwrap(); + assert!(tag.distinct_count >= 2); + assert!(tag.distinct_count <= tag.row_count); + assert!(matches!(t.source, TableSource::File(_))); + std::fs::remove_file(&p).unwrap(); + } + + #[test] + fn registering_a_file_that_does_not_exist_reports_the_io_error() { + let mut c = Catalog::new(); + assert!(c.register_file("x", "/nonexistent/path.qfc").is_err()); + } + + #[test] + fn missing_tables_are_reported_by_name() { + let c = Catalog::new(); + assert!(c.is_empty()); + assert!(c.get("ghost").unwrap_err().to_string().contains("`ghost`")); + } + + #[test] + fn tables_are_listed_in_a_stable_order() { + let mut c = Catalog::new(); + for n in ["zebra", "apple", "mango"] { + c.register_batches(n, schema(), vec![]).unwrap(); + } + assert_eq!(c.table_names(), vec!["apple", "mango", "zebra"]); + assert_eq!(c.tables().count(), 3); + } + + #[test] + fn dropping_removes_a_table_and_reports_an_unknown_one() { + let mut c = Catalog::new(); + c.register_batches("t", schema(), vec![]).unwrap(); + assert_eq!(c.drop_table("T").unwrap().name, "t"); + assert!(c.is_empty()); + assert!(c.drop_table("t").is_err()); + } +} diff --git a/crates/qf-storage/src/csv.rs b/crates/qf-storage/src/csv.rs new file mode 100644 index 0000000..2f181c8 --- /dev/null +++ b/crates/qf-storage/src/csv.rs @@ -0,0 +1,336 @@ +//! CSV ingest — the on-ramp into the columnar format. +//! +//! CSV is how data actually arrives, so the engine has to read it, but nothing +//! else in the engine ever touches a CSV: the loader converts to batches once +//! and everything downstream reads `.qfc`. + +use crate::array::ArrayBuilder; +use crate::batch::{RecordBatch, DEFAULT_BATCH_SIZE}; +use qf_common::{DataType, Error, Field, Result, Schema, Value}; +use std::path::Path; +use std::sync::Arc; + +/// Splits one CSV line, honouring double quotes and doubled `""` escapes. +/// +/// Written by hand rather than by splitting on commas because a quoted field +/// containing a comma is not an edge case in real data — it is most address +/// columns. +pub fn split_line(line: &str) -> Vec { + let mut fields = Vec::new(); + let mut current = String::new(); + let mut in_quotes = false; + let mut chars = line.chars().peekable(); + while let Some(c) = chars.next() { + match c { + '"' if in_quotes => { + if chars.peek() == Some(&'"') { + chars.next(); + current.push('"'); + } else { + in_quotes = false; + } + } + '"' => in_quotes = true, + ',' if !in_quotes => { + fields.push(std::mem::take(&mut current)); + } + _ => current.push(c), + } + } + fields.push(current); + fields +} + +/// Infers a column type from the values seen in it. +/// +/// The rule is the narrowest type that fits everything: an integer column with +/// one decimal in it is a float column, and anything that fails to parse makes +/// the column text. Empty strings are treated as NULL and never widen the type, +/// which is the difference between a usable `amount` column and a text one. +fn infer_column(samples: &[&str]) -> DataType { + let mut saw_value = false; + let mut all_int = true; + let mut all_float = true; + let mut all_bool = true; + for s in samples { + let t = s.trim(); + if t.is_empty() { + continue; + } + saw_value = true; + if t.parse::().is_err() { + all_int = false; + } + if t.parse::().is_err() { + all_float = false; + } + if !matches!( + t.to_ascii_lowercase().as_str(), + "true" | "false" | "t" | "f" + ) { + all_bool = false; + } + } + if !saw_value { + return DataType::Utf8; + } + if all_bool { + DataType::Boolean + } else if all_int { + DataType::Int64 + } else if all_float { + DataType::Float64 + } else { + DataType::Utf8 + } +} + +/// Parses one CSV cell into the column's type. An empty cell is NULL. +fn parse_cell(cell: &str, data_type: DataType) -> Result { + let t = cell.trim(); + if t.is_empty() { + return Ok(Value::Null); + } + Value::Utf8(t.to_string()).cast_to(data_type) +} + +/// The header row plus the inferred column types. +pub fn infer_schema(text: &str, sample_rows: usize) -> Result> { + let mut lines = text.lines().filter(|l| !l.trim().is_empty()); + let header = lines + .next() + .ok_or_else(|| Error::storage("CSV input is empty".to_string()))?; + let names = split_line(header); + if names.iter().any(|n| n.trim().is_empty()) { + return Err(Error::storage( + "CSV header has an empty column name".to_string(), + )); + } + + let rows: Vec> = lines.take(sample_rows).map(split_line).collect(); + let mut fields = Vec::with_capacity(names.len()); + for (i, name) in names.iter().enumerate() { + let samples: Vec<&str> = rows + .iter() + .filter_map(|r| r.get(i).map(String::as_str)) + .collect(); + let nullable = samples.iter().any(|s| s.trim().is_empty()); + fields.push(Field::new(name.trim(), infer_column(&samples), nullable)); + } + Ok(Arc::new(Schema::new(fields))) +} + +/// Reads CSV text into batches, inferring the schema from the first +/// `sample_rows` data rows. +/// +/// Fully blank lines are treated as separators and skipped; a blank *cell* +/// within a row is a NULL. +pub fn read_csv_str(text: &str, sample_rows: usize) -> Result<(Arc, Vec)> { + let schema = infer_schema(text, sample_rows)?; + let batches = read_csv_with_schema(text, Arc::clone(&schema))?; + Ok((schema, batches)) +} + +/// Reads CSV text against a schema that is already known. +pub fn read_csv_with_schema(text: &str, schema: Arc) -> Result> { + let mut lines = text.lines().filter(|l| !l.trim().is_empty()); + lines.next(); // header + + let mut batches = Vec::new(); + let mut builders: Vec = schema + .fields() + .iter() + .map(|f| ArrayBuilder::new(f.data_type)) + .collect(); + let mut rows_in_batch = 0usize; + + for (line_no, line) in lines.enumerate() { + let cells = split_line(line); + if cells.len() != schema.len() { + return Err(Error::storage(format!( + "row {} has {} fields, expected {}", + line_no + 2, + cells.len(), + schema.len() + ))); + } + for (i, builder) in builders.iter_mut().enumerate() { + let value = parse_cell(&cells[i], builder.data_type()).map_err(|e| { + Error::storage(format!( + "row {}, column `{}`: {e}", + line_no + 2, + schema.field(i).map(|f| f.name.clone()).unwrap_or_default() + )) + })?; + builder.push(value)?; + } + rows_in_batch += 1; + + if rows_in_batch == DEFAULT_BATCH_SIZE { + batches.push(finish_batch(&schema, &mut builders)?); + rows_in_batch = 0; + } + } + if rows_in_batch > 0 { + batches.push(finish_batch(&schema, &mut builders)?); + } + Ok(batches) +} + +fn finish_batch(schema: &Arc, builders: &mut Vec) -> Result { + let fresh: Vec = schema + .fields() + .iter() + .map(|f| ArrayBuilder::new(f.data_type)) + .collect(); + let done = std::mem::replace(builders, fresh); + let columns = done + .into_iter() + .map(ArrayBuilder::finish) + .collect::>>()?; + RecordBatch::try_new(Arc::clone(schema), columns) +} + +pub fn read_csv_file( + path: impl AsRef, + sample_rows: usize, +) -> Result<(Arc, Vec)> { + let text = std::fs::read_to_string(path)?; + read_csv_str(&text, sample_rows) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quoted_fields_may_contain_commas_and_escaped_quotes() { + assert_eq!( + split_line(r#"1,"Smith, John","say ""hi""",x"#), + vec!["1", "Smith, John", r#"say "hi""#, "x"] + ); + } + + #[test] + fn empty_trailing_fields_are_preserved() { + assert_eq!(split_line("a,,c,"), vec!["a", "", "c", ""]); + assert_eq!(split_line(""), vec![""]); + } + + #[test] + fn types_are_inferred_as_the_narrowest_that_fits() { + let csv = "id,price,name,active\n1,1.5,ann,true\n2,2,bob,false\n"; + let s = infer_schema(csv, 100).unwrap(); + assert_eq!(s.field(0).unwrap().data_type, DataType::Int64); + assert_eq!(s.field(1).unwrap().data_type, DataType::Float64); + assert_eq!(s.field(2).unwrap().data_type, DataType::Utf8); + assert_eq!(s.field(3).unwrap().data_type, DataType::Boolean); + } + + #[test] + fn one_decimal_widens_an_otherwise_integer_column_to_float() { + let s = infer_schema("n\n1\n2\n3.5\n", 100).unwrap(); + assert_eq!(s.field(0).unwrap().data_type, DataType::Float64); + } + + #[test] + fn one_unparseable_value_makes_the_whole_column_text() { + let s = infer_schema("n\n1\n2\nn/a\n", 100).unwrap(); + assert_eq!(s.field(0).unwrap().data_type, DataType::Utf8); + } + + #[test] + fn a_blank_cell_is_null_and_does_not_widen_the_type() { + // A blank *cell* is a null. A blank *line* is a separator and is + // skipped, which is why this case needs a second column to be + // unambiguous. + let csv = "label,n\na,1\nb,\nc,3\n"; + let (schema, batches) = read_csv_str(csv, 100).unwrap(); + assert_eq!(schema.field(1).unwrap().data_type, DataType::Int64); + assert!(schema.field(1).unwrap().nullable); + assert!(!schema.field(0).unwrap().nullable); + let b = &batches[0]; + assert_eq!(b.num_rows(), 3); + assert!(b.column(1).unwrap().value(1).is_null()); + assert_eq!(b.column(1).unwrap().value(2), Value::Int64(3)); + } + + #[test] + fn an_all_blank_column_falls_back_to_text() { + let s = infer_schema("a,b\n1,\n2,\n", 100).unwrap(); + assert_eq!(s.field(1).unwrap().data_type, DataType::Utf8); + } + + #[test] + fn only_the_sampled_rows_drive_inference() { + // The float appears after the sample window, so inference says Int64 + // and the value that does not fit is reported rather than guessed at. + let csv = "n\n1\n2\n3.5\n"; + let s = infer_schema(csv, 2).unwrap(); + assert_eq!(s.field(0).unwrap().data_type, DataType::Int64); + let err = read_csv_with_schema(csv, s).unwrap_err(); + assert!(err.to_string().contains("row 4")); + assert!(err.to_string().contains("column `n`")); + } + + #[test] + fn a_row_with_the_wrong_field_count_is_reported_with_its_line_number() { + let err = read_csv_str("a,b\n1,2\n3\n", 10).unwrap_err(); + assert!(err.to_string().contains("row 3")); + assert!(err.to_string().contains("expected 2")); + } + + #[test] + fn values_round_trip_into_typed_columns() { + let csv = "id,name,score,ok\n1,ann,9.5,true\n2,bob,8,false\n"; + let (_, batches) = read_csv_str(csv, 10).unwrap(); + let b = &batches[0]; + assert_eq!(b.num_rows(), 2); + assert_eq!(b.column(0).unwrap().value(1), Value::Int64(2)); + assert_eq!(b.column(1).unwrap().value(0), Value::Utf8("ann".into())); + assert_eq!(b.column(2).unwrap().value(1), Value::Float64(8.0)); + assert_eq!(b.column(3).unwrap().value(0), Value::Boolean(true)); + } + + #[test] + fn large_inputs_are_split_into_batches() { + let mut csv = String::from("n\n"); + for i in 0..(DEFAULT_BATCH_SIZE + 5) { + csv.push_str(&format!("{i}\n")); + } + let (_, batches) = read_csv_str(&csv, 10).unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].num_rows(), DEFAULT_BATCH_SIZE); + assert_eq!(batches[1].num_rows(), 5); + } + + #[test] + fn blank_lines_are_skipped() { + let (_, batches) = read_csv_str("n\n1\n\n\n2\n", 10).unwrap(); + assert_eq!(batches[0].num_rows(), 2); + } + + #[test] + fn empty_input_and_empty_headers_are_rejected() { + assert!(infer_schema("", 10).is_err()); + assert!(infer_schema("a,,c\n1,2,3\n", 10).is_err()); + } + + #[test] + fn a_header_only_file_yields_no_batches() { + let (schema, batches) = read_csv_str("a,b\n", 10).unwrap(); + assert_eq!(schema.len(), 2); + assert!(batches.is_empty()); + } + + #[test] + fn files_are_read_from_disk() { + let mut p = std::env::temp_dir(); + p.push(format!("qf-csv-{}.csv", std::process::id())); + std::fs::write(&p, "a,b\n1,x\n2,y\n").unwrap(); + let (schema, batches) = read_csv_file(&p, 10).unwrap(); + assert_eq!(schema.len(), 2); + assert_eq!(batches[0].num_rows(), 2); + std::fs::remove_file(&p).unwrap(); + } +} diff --git a/crates/qf-storage/src/encoding.rs b/crates/qf-storage/src/encoding.rs new file mode 100644 index 0000000..945843f --- /dev/null +++ b/crates/qf-storage/src/encoding.rs @@ -0,0 +1,558 @@ +use crate::array::{Array, ArrayBuilder, ArrayData}; +use crate::bitmap::Bitmap; +use crate::bytes::{ByteReader, ByteWriter}; +use qf_common::{DataType, Error, Result, Value}; +use std::collections::HashMap; + +/// How the values of one column chunk are laid out on disk. +/// +/// The encoding is chosen per chunk, not per column, because the right answer +/// changes with the data: an `order_status` column can be two distinct values +/// in one row group and forty in the next. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Encoding { + /// Values written one after another; varint for ints, bit-packed for bools. + Plain, + /// A sorted-order dictionary of distinct values plus one index per row. + /// Pays for itself as soon as values repeat. + Dictionary, + /// (value, run length) pairs. The win on sorted or clustered columns. + Rle, +} + +impl Encoding { + fn tag(self) -> u8 { + match self { + Encoding::Plain => 0, + Encoding::Dictionary => 1, + Encoding::Rle => 2, + } + } + + fn from_tag(tag: u8) -> Result { + match tag { + 0 => Ok(Encoding::Plain), + 1 => Ok(Encoding::Dictionary), + 2 => Ok(Encoding::Rle), + other => Err(Error::storage(format!("unknown encoding tag {other}"))), + } + } + + pub fn name(self) -> &'static str { + match self { + Encoding::Plain => "plain", + Encoding::Dictionary => "dictionary", + Encoding::Rle => "rle", + } + } +} + +/// Picks an encoding by measuring the chunk rather than guessing from the type. +/// +/// Two cheap signals decide it: the average run length (long runs → RLE) and +/// the ratio of distinct values to rows (few distinct → dictionary). The +/// thresholds are conservative — both encodings add a decode step, so a chunk +/// that would barely benefit stays plain. +pub fn choose_encoding(array: &Array) -> Encoding { + let n = array.len(); + if n < 16 { + return Encoding::Plain; + } + let mut runs = 1usize; + let mut distinct: HashMap = HashMap::new(); + let mut prev = array.value(0); + distinct.insert(prev.clone(), ()); + for i in 1..n { + let v = array.value(i); + if v != prev { + runs += 1; + } + distinct.entry(v.clone()).or_insert(()); + prev = v; + } + let avg_run = n as f64 / runs as f64; + if avg_run >= 4.0 { + return Encoding::Rle; + } + if (distinct.len() as f64) / (n as f64) <= 0.5 { + return Encoding::Dictionary; + } + Encoding::Plain +} + +/// Serialises one column chunk: encoding tag, validity mask, then the values. +/// +/// Only non-null values are written. A column that is 90% null costs a bitmap +/// and almost nothing else. +pub fn encode_column(array: &Array, encoding: Encoding) -> Result> { + let mut w = ByteWriter::new(); + w.u8(encoding.tag()); + w.uvarint(array.len() as u64); + + match array.validity() { + None => w.u8(0), + Some(v) => { + w.u8(1); + write_bitmap(&mut w, v); + } + } + + let values: Vec = (0..array.len()) + .filter(|i| array.is_valid(*i)) + .map(|i| array.value(i)) + .collect(); + + match encoding { + Encoding::Plain => encode_plain(&mut w, array.data_type(), &values)?, + Encoding::Dictionary => encode_dictionary(&mut w, array.data_type(), &values)?, + Encoding::Rle => encode_rle(&mut w, array.data_type(), &values)?, + } + Ok(w.into_bytes()) +} + +/// Serialises with whichever encoding `choose_encoding` picks. +pub fn encode_column_auto(array: &Array) -> Result<(Encoding, Vec)> { + let enc = choose_encoding(array); + Ok((enc, encode_column(array, enc)?)) +} + +pub fn decode_column(data_type: DataType, buf: &[u8]) -> Result { + let mut r = ByteReader::new(buf); + let encoding = Encoding::from_tag(r.u8()?)?; + let len = r.uvarint()? as usize; + let validity = if r.u8()? == 1 { + Some(read_bitmap(&mut r, len)?) + } else { + None + }; + let valid_count = validity.as_ref().map_or(len, Bitmap::count_set); + + let values = match encoding { + Encoding::Plain => decode_plain(&mut r, data_type, valid_count)?, + Encoding::Dictionary => decode_dictionary(&mut r, data_type, valid_count)?, + Encoding::Rle => decode_rle(&mut r, data_type, valid_count)?, + }; + if values.len() != valid_count { + return Err(Error::storage(format!( + "chunk holds {} values but the validity mask expects {valid_count}", + values.len() + ))); + } + + let mut b = ArrayBuilder::new(data_type); + b.reserve(len); + let mut next = 0usize; + for i in 0..len { + let valid = validity.as_ref().is_none_or(|v| v.get(i)); + if valid { + b.push(values[next].clone())?; + next += 1; + } else { + b.push_null(); + } + } + b.finish() +} + +fn write_bitmap(w: &mut ByteWriter, bm: &Bitmap) { + let words = bm.len().div_ceil(64); + w.uvarint(words as u64); + for wi in 0..words { + let mut word = 0u64; + for b in 0..64 { + let idx = wi * 64 + b; + if idx < bm.len() && bm.get(idx) { + word |= 1u64 << b; + } + } + w.u64(word); + } +} + +fn read_bitmap(r: &mut ByteReader<'_>, len: usize) -> Result { + let words = r.uvarint()? as usize; + if words != len.div_ceil(64) { + return Err(Error::storage(format!( + "validity mask has {words} words for {len} rows" + ))); + } + let mut bm = Bitmap::with_capacity(len); + let mut buf = Vec::with_capacity(words); + for _ in 0..words { + buf.push(r.u64()?); + } + for i in 0..len { + bm.push(buf[i / 64] & (1u64 << (i % 64)) != 0); + } + Ok(bm) +} + +fn encode_scalar(w: &mut ByteWriter, data_type: DataType, v: &Value) -> Result<()> { + match (data_type, v) { + (DataType::Boolean, Value::Boolean(b)) => w.u8(u8::from(*b)), + (DataType::Int64, Value::Int64(i)) => w.ivarint(*i), + (DataType::Float64, Value::Float64(f)) => w.f64(*f), + (DataType::Float64, Value::Int64(i)) => w.f64(*i as f64), + (DataType::Utf8, Value::Utf8(s)) => w.string(s), + (dt, v) => { + return Err(Error::storage(format!( + "cannot encode {} into a {dt} chunk", + v.type_name() + ))) + } + } + Ok(()) +} + +fn decode_scalar(r: &mut ByteReader<'_>, data_type: DataType) -> Result { + Ok(match data_type { + DataType::Boolean => Value::Boolean(r.u8()? != 0), + DataType::Int64 => Value::Int64(r.ivarint()?), + DataType::Float64 => Value::Float64(r.f64()?), + DataType::Utf8 => Value::Utf8(r.string()?), + }) +} + +fn encode_plain(w: &mut ByteWriter, data_type: DataType, values: &[Value]) -> Result<()> { + if data_type == DataType::Boolean { + // Bit-pack rather than a byte per value. + let bm: Bitmap = values + .iter() + .map(|v| matches!(v, Value::Boolean(true))) + .collect(); + write_bitmap(w, &bm); + return Ok(()); + } + for v in values { + encode_scalar(w, data_type, v)?; + } + Ok(()) +} + +fn decode_plain(r: &mut ByteReader<'_>, data_type: DataType, n: usize) -> Result> { + if data_type == DataType::Boolean { + let bm = read_bitmap(r, n)?; + return Ok((0..n).map(|i| Value::Boolean(bm.get(i))).collect()); + } + (0..n).map(|_| decode_scalar(r, data_type)).collect() +} + +fn encode_dictionary(w: &mut ByteWriter, data_type: DataType, values: &[Value]) -> Result<()> { + let mut dict: Vec = Vec::new(); + let mut index: HashMap = HashMap::new(); + let mut codes: Vec = Vec::with_capacity(values.len()); + for v in values { + let code = match index.get(v) { + Some(c) => *c, + None => { + let c = dict.len() as u64; + dict.push(v.clone()); + index.insert(v.clone(), c); + c + } + }; + codes.push(code); + } + w.uvarint(dict.len() as u64); + for v in &dict { + encode_scalar(w, data_type, v)?; + } + for c in codes { + w.uvarint(c); + } + Ok(()) +} + +fn decode_dictionary(r: &mut ByteReader<'_>, data_type: DataType, n: usize) -> Result> { + let dict_len = r.uvarint()? as usize; + let mut dict = Vec::with_capacity(dict_len); + for _ in 0..dict_len { + dict.push(decode_scalar(r, data_type)?); + } + let mut out = Vec::with_capacity(n); + for _ in 0..n { + let code = r.uvarint()? as usize; + let v = dict + .get(code) + .ok_or_else(|| Error::storage(format!("dictionary code {code} out of range")))?; + out.push(v.clone()); + } + Ok(out) +} + +fn encode_rle(w: &mut ByteWriter, data_type: DataType, values: &[Value]) -> Result<()> { + let mut runs: Vec<(&Value, u64)> = Vec::new(); + for v in values { + match runs.last_mut() { + Some((prev, count)) if *prev == v => *count += 1, + _ => runs.push((v, 1)), + } + } + w.uvarint(runs.len() as u64); + for (v, count) in runs { + encode_scalar(w, data_type, v)?; + w.uvarint(count); + } + Ok(()) +} + +fn decode_rle(r: &mut ByteReader<'_>, data_type: DataType, n: usize) -> Result> { + let run_count = r.uvarint()? as usize; + let mut out = Vec::with_capacity(n); + for _ in 0..run_count { + let v = decode_scalar(r, data_type)?; + let count = r.uvarint()? as usize; + if out.len() + count > n { + return Err(Error::storage( + "run-length chunk describes more rows than the header declares".to_string(), + )); + } + for _ in 0..count { + out.push(v.clone()); + } + } + Ok(out) +} + +/// Bytes the chunk would occupy under each encoding — used by the CLI's +/// `\encodings` command and by the tests that assert the heuristic picks a +/// genuinely smaller layout. +pub fn measure_encodings(array: &Array) -> Result> { + let mut out = Vec::new(); + for enc in [Encoding::Plain, Encoding::Dictionary, Encoding::Rle] { + out.push((enc, encode_column(array, enc)?.len())); + } + Ok(out) +} + +/// Debug helper: which physical layout is behind an array. +pub fn layout_name(array: &Array) -> &'static str { + match array.data() { + ArrayData::Boolean(_) => "bitmap", + ArrayData::Int64(_) => "i64 buffer", + ArrayData::Float64(_) => "f64 buffer", + ArrayData::Utf8 { .. } => "offsets + bytes", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn arr(dt: DataType, values: Vec) -> Array { + Array::from_values(dt, &values).unwrap() + } + + fn round_trip(a: &Array, enc: Encoding) { + let buf = encode_column(a, enc).unwrap(); + let back = decode_column(a.data_type(), &buf).unwrap(); + assert_eq!(&back, a, "{enc:?} round trip"); + } + + #[test] + fn every_encoding_round_trips_every_type() { + let cases = vec![ + arr( + DataType::Int64, + (0..40).map(|i| Value::Int64(i % 5)).collect(), + ), + arr( + DataType::Float64, + (0..40).map(|i| Value::Float64(i as f64 / 3.0)).collect(), + ), + arr( + DataType::Utf8, + (0..40) + .map(|i| Value::Utf8(format!("s{}", i % 4))) + .collect(), + ), + arr( + DataType::Boolean, + (0..40).map(|i| Value::Boolean(i % 3 == 0)).collect(), + ), + ]; + for a in cases { + for enc in [Encoding::Plain, Encoding::Dictionary, Encoding::Rle] { + round_trip(&a, enc); + } + } + } + + #[test] + fn nulls_survive_every_encoding() { + let a = arr( + DataType::Int64, + (0..100) + .map(|i| { + if i % 3 == 0 { + Value::Null + } else { + Value::Int64(i) + } + }) + .collect(), + ); + assert_eq!(a.null_count(), 34); + for enc in [Encoding::Plain, Encoding::Dictionary, Encoding::Rle] { + round_trip(&a, enc); + } + } + + #[test] + fn a_fully_null_column_costs_almost_nothing() { + let a = Array::nulls(DataType::Utf8, 1000).unwrap(); + let buf = encode_column(&a, Encoding::Plain).unwrap(); + // 1000 rows of validity bitmap is 128 bytes; the values cost nothing. + assert!(buf.len() < 200, "{} bytes", buf.len()); + round_trip(&a, Encoding::Plain); + } + + #[test] + fn an_empty_column_round_trips() { + for dt in [ + DataType::Int64, + DataType::Utf8, + DataType::Boolean, + DataType::Float64, + ] { + round_trip(&arr(dt, vec![]), Encoding::Plain); + } + } + + #[test] + fn rle_is_chosen_and_wins_on_a_sorted_column() { + let a = arr( + DataType::Int64, + (0..1000).map(|i| Value::Int64(i / 100)).collect(), + ); + assert_eq!(choose_encoding(&a), Encoding::Rle); + let sizes = measure_encodings(&a).unwrap(); + let plain = sizes.iter().find(|(e, _)| *e == Encoding::Plain).unwrap().1; + let rle = sizes.iter().find(|(e, _)| *e == Encoding::Rle).unwrap().1; + assert!(rle * 10 < plain, "rle {rle} vs plain {plain}"); + } + + #[test] + fn dictionary_is_chosen_and_wins_on_a_low_cardinality_shuffled_column() { + // Low cardinality but no long runs, which is exactly the case RLE + // cannot help with. + let words = ["SHIPPED", "PENDING", "CANCELLED", "DELIVERED"]; + let a = arr( + DataType::Utf8, + (0..1000) + .map(|i| Value::Utf8(words[(i * 7 + 3) % 4].to_string())) + .collect(), + ); + assert_eq!(choose_encoding(&a), Encoding::Dictionary); + let sizes = measure_encodings(&a).unwrap(); + let plain = sizes.iter().find(|(e, _)| *e == Encoding::Plain).unwrap().1; + let dict = sizes + .iter() + .find(|(e, _)| *e == Encoding::Dictionary) + .unwrap() + .1; + assert!(dict * 4 < plain, "dict {dict} vs plain {plain}"); + } + + #[test] + fn plain_is_chosen_when_neither_encoding_would_help() { + let a = arr( + DataType::Int64, + (0..1000).map(|i| Value::Int64(i * 7919)).collect(), + ); + assert_eq!(choose_encoding(&a), Encoding::Plain); + } + + #[test] + fn tiny_chunks_skip_the_analysis_entirely() { + let a = arr(DataType::Int64, vec![Value::Int64(1); 8]); + assert_eq!(choose_encoding(&a), Encoding::Plain); + } + + #[test] + fn auto_encoding_records_the_encoding_it_used() { + let a = arr(DataType::Int64, vec![Value::Int64(4); 100]); + let (enc, buf) = encode_column_auto(&a).unwrap(); + assert_eq!(enc, Encoding::Rle); + assert_eq!(decode_column(DataType::Int64, &buf).unwrap(), a); + } + + #[test] + fn a_corrupt_encoding_tag_is_rejected() { + let mut buf = encode_column( + &arr(DataType::Int64, vec![Value::Int64(1)]), + Encoding::Plain, + ) + .unwrap(); + buf[0] = 9; + assert!(decode_column(DataType::Int64, &buf).is_err()); + } + + #[test] + fn a_dictionary_code_pointing_past_the_dictionary_is_rejected() { + let mut w = ByteWriter::new(); + w.u8(Encoding::Dictionary.tag()); + w.uvarint(1); // one row + w.u8(0); // no validity + w.uvarint(1); // dictionary of one entry + w.ivarint(5); + w.uvarint(9); // code 9 does not exist + let buf = w.into_bytes(); + assert!(decode_column(DataType::Int64, &buf) + .unwrap_err() + .to_string() + .contains("out of range")); + } + + #[test] + fn an_rle_run_longer_than_the_declared_row_count_is_rejected() { + let mut w = ByteWriter::new(); + w.u8(Encoding::Rle.tag()); + w.uvarint(2); // two rows + w.u8(0); + w.uvarint(1); // one run... + w.ivarint(7); + w.uvarint(50); // ...claiming fifty rows + let buf = w.into_bytes(); + assert!(decode_column(DataType::Int64, &buf).is_err()); + } + + #[test] + fn a_short_rle_chunk_is_rejected_rather_than_silently_padded() { + let mut w = ByteWriter::new(); + w.u8(Encoding::Rle.tag()); + w.uvarint(5); + w.u8(0); + w.uvarint(1); + w.ivarint(7); + w.uvarint(2); // only two of the five rows described + let buf = w.into_bytes(); + assert!(decode_column(DataType::Int64, &buf) + .unwrap_err() + .to_string() + .contains("validity mask expects")); + } + + #[test] + fn a_malformed_validity_mask_is_rejected() { + let mut w = ByteWriter::new(); + w.u8(Encoding::Plain.tag()); + w.uvarint(100); + w.u8(1); + w.uvarint(1); // one word cannot cover 100 rows + w.u64(0); + let buf = w.into_bytes(); + assert!(decode_column(DataType::Int64, &buf).is_err()); + } + + #[test] + fn encoding_names_and_layout_names_are_reported() { + assert_eq!(Encoding::Plain.name(), "plain"); + assert_eq!(Encoding::Dictionary.name(), "dictionary"); + assert_eq!(Encoding::Rle.name(), "rle"); + assert_eq!(layout_name(&arr(DataType::Utf8, vec![])), "offsets + bytes"); + assert_eq!(layout_name(&arr(DataType::Int64, vec![])), "i64 buffer"); + assert_eq!(layout_name(&arr(DataType::Float64, vec![])), "f64 buffer"); + assert_eq!(layout_name(&arr(DataType::Boolean, vec![])), "bitmap"); + } +} diff --git a/crates/qf-storage/src/format.rs b/crates/qf-storage/src/format.rs new file mode 100644 index 0000000..2887d69 --- /dev/null +++ b/crates/qf-storage/src/format.rs @@ -0,0 +1,709 @@ +//! `.qfc` — the columnar file format the engine reads and writes. +//! +//! ```text +//! ┌────────┬──────────────────────────┬─────────┬────────────┬────────┐ +//! │ "QFC1" │ row group 0..n │ footer │ footer len │ "QFC1" │ +//! │ 4 B │ column chunks, back to │ metadata│ u32 LE │ 4 B │ +//! │ │ back, one per column │ │ │ │ +//! └────────┴──────────────────────────┴─────────┴────────────┴────────┘ +//! ``` +//! +//! The footer lives at the end, and its length is the last thing before the +//! trailing magic, so a reader seeks to `len - 8`, learns where the metadata +//! starts, and reads it in one go — without touching the data. +//! +//! That layout is what makes projection and predicate pushdown physical rather +//! than cosmetic. Each column chunk records its own byte range and its own +//! zone map, so a query that touches two of forty columns reads two chunks per +//! row group, and a query whose predicate falls outside a row group's `[min, +//! max]` reads none of it at all. + +use crate::array::Array; +use crate::batch::RecordBatch; +use crate::bytes::{ByteReader, ByteWriter}; +use crate::encoding::{decode_column, encode_column_auto, Encoding}; +use crate::stats::ColumnStats; +use qf_common::{DataType, Error, Field, Result, Schema}; +use std::fs::File; +use std::io::{BufWriter, Read, Seek, SeekFrom, Write}; +use std::path::Path; +use std::sync::Arc; + +const MAGIC: &[u8; 4] = b"QFC1"; + +/// Rows per row group. This is the granularity of zone-map pruning: smaller +/// groups prune more precisely but multiply the footer, and past a point the +/// metadata costs more than the reads it saves. +pub const DEFAULT_ROW_GROUP_SIZE: usize = 65_536; + +#[derive(Debug, Clone, PartialEq)] +pub struct ColumnChunkMeta { + pub offset: u64, + pub length: u64, + pub encoding: Encoding, + pub stats: ColumnStats, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct RowGroupMeta { + pub num_rows: usize, + pub columns: Vec, +} + +impl RowGroupMeta { + pub fn column(&self, i: usize) -> Result<&ColumnChunkMeta> { + self.columns + .get(i) + .ok_or_else(|| Error::storage(format!("row group has no column {i}"))) + } + + /// Bytes this row group occupies on disk. + pub fn byte_size(&self) -> u64 { + self.columns.iter().map(|c| c.length).sum() + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct FileMeta { + pub schema: Arc, + pub row_groups: Vec, +} + +impl FileMeta { + pub fn num_rows(&self) -> usize { + self.row_groups.iter().map(|g| g.num_rows).sum() + } +} + +fn type_tag(dt: DataType) -> u8 { + match dt { + DataType::Boolean => 0, + DataType::Int64 => 1, + DataType::Float64 => 2, + DataType::Utf8 => 3, + } +} + +fn type_from_tag(tag: u8) -> Result { + Ok(match tag { + 0 => DataType::Boolean, + 1 => DataType::Int64, + 2 => DataType::Float64, + 3 => DataType::Utf8, + other => return Err(Error::storage(format!("unknown type tag {other}"))), + }) +} + +fn encode_footer(meta: &FileMeta) -> Vec { + let mut w = ByteWriter::new(); + w.uvarint(meta.schema.len() as u64); + for f in meta.schema.fields() { + w.string(&f.name); + w.u8(type_tag(f.data_type)); + w.u8(u8::from(f.nullable)); + } + w.uvarint(meta.row_groups.len() as u64); + for g in &meta.row_groups { + w.uvarint(g.num_rows as u64); + w.uvarint(g.columns.len() as u64); + for c in &g.columns { + w.u64(c.offset); + w.u64(c.length); + w.u8(match c.encoding { + Encoding::Plain => 0, + Encoding::Dictionary => 1, + Encoding::Rle => 2, + }); + c.stats.encode(&mut w); + } + } + w.into_bytes() +} + +fn decode_footer(buf: &[u8]) -> Result { + let mut r = ByteReader::new(buf); + let field_count = r.uvarint()? as usize; + let mut fields = Vec::with_capacity(field_count); + for _ in 0..field_count { + let name = r.string()?; + let data_type = type_from_tag(r.u8()?)?; + let nullable = r.u8()? != 0; + fields.push(Field::new(name, data_type, nullable)); + } + let group_count = r.uvarint()? as usize; + let mut row_groups = Vec::with_capacity(group_count); + for _ in 0..group_count { + let num_rows = r.uvarint()? as usize; + let col_count = r.uvarint()? as usize; + if col_count != field_count { + return Err(Error::storage(format!( + "row group describes {col_count} columns but the schema has {field_count}" + ))); + } + let mut columns = Vec::with_capacity(col_count); + for _ in 0..col_count { + let offset = r.u64()?; + let length = r.u64()?; + let encoding = match r.u8()? { + 0 => Encoding::Plain, + 1 => Encoding::Dictionary, + 2 => Encoding::Rle, + other => return Err(Error::storage(format!("unknown encoding tag {other}"))), + }; + let stats = ColumnStats::decode(&mut r)?; + columns.push(ColumnChunkMeta { + offset, + length, + encoding, + stats, + }); + } + row_groups.push(RowGroupMeta { num_rows, columns }); + } + Ok(FileMeta { + schema: Arc::new(Schema::new(fields)), + row_groups, + }) +} + +/// Buffers batches into row groups and writes them out. +pub struct QfcWriter { + file: BufWriter, + schema: Arc, + row_group_size: usize, + pending: Vec, + pending_rows: usize, + row_groups: Vec, + offset: u64, + finished: bool, +} + +impl QfcWriter { + pub fn create(path: impl AsRef, schema: Arc) -> Result { + QfcWriter::create_with_row_group_size(path, schema, DEFAULT_ROW_GROUP_SIZE) + } + + pub fn create_with_row_group_size( + path: impl AsRef, + schema: Arc, + row_group_size: usize, + ) -> Result { + if row_group_size == 0 { + return Err(Error::storage( + "row group size must be positive".to_string(), + )); + } + let mut file = BufWriter::new(File::create(path)?); + file.write_all(MAGIC)?; + Ok(QfcWriter { + file, + schema, + row_group_size, + pending: Vec::new(), + pending_rows: 0, + row_groups: Vec::new(), + offset: MAGIC.len() as u64, + finished: false, + }) + } + + pub fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { + if batch.schema().as_ref() != self.schema.as_ref() { + return Err(Error::storage( + "batch schema does not match the file schema".to_string(), + )); + } + if batch.num_rows() == 0 { + return Ok(()); + } + self.pending_rows += batch.num_rows(); + self.pending.push(batch.clone()); + while self.pending_rows >= self.row_group_size { + self.flush_row_group(self.row_group_size)?; + } + Ok(()) + } + + fn flush_row_group(&mut self, max_rows: usize) -> Result<()> { + if self.pending_rows == 0 { + return Ok(()); + } + let combined = RecordBatch::concat(Arc::clone(&self.schema), &self.pending)?; + let take = max_rows.min(combined.num_rows()); + let group_rows: Vec = (0..take).collect(); + let group = combined.take(&group_rows)?; + + let mut columns = Vec::with_capacity(group.num_columns()); + for ci in 0..group.num_columns() { + let array = group.column(ci)?; + let (encoding, bytes) = encode_column_auto(array)?; + self.file.write_all(&bytes)?; + columns.push(ColumnChunkMeta { + offset: self.offset, + length: bytes.len() as u64, + encoding, + stats: ColumnStats::from_array(array), + }); + self.offset += bytes.len() as u64; + } + self.row_groups.push(RowGroupMeta { + num_rows: take, + columns, + }); + + // Carry the remainder into the next row group. + let rest: Vec = (take..combined.num_rows()).collect(); + self.pending_rows = rest.len(); + self.pending = if rest.is_empty() { + Vec::new() + } else { + vec![combined.take(&rest)?] + }; + Ok(()) + } + + pub fn finish(mut self) -> Result { + self.flush_row_group(self.row_group_size)?; + let meta = FileMeta { + schema: Arc::clone(&self.schema), + row_groups: std::mem::take(&mut self.row_groups), + }; + let footer = encode_footer(&meta); + self.file.write_all(&footer)?; + self.file.write_all(&(footer.len() as u32).to_le_bytes())?; + self.file.write_all(MAGIC)?; + self.file.flush()?; + self.finished = true; + Ok(meta) + } +} + +impl Drop for QfcWriter { + fn drop(&mut self) { + if !self.finished { + // A file without a footer is unreadable, and silently leaving one + // behind is worse than a loud warning: the next run would report a + // corrupt file with no hint of why. + eprintln!( + "warning: a .qfc writer was dropped without finish(); the file has no footer" + ); + } + } +} + +/// Reads a `.qfc` file, one column chunk at a time. +#[derive(Debug)] +pub struct QfcReader { + file: File, + meta: FileMeta, +} + +impl QfcReader { + pub fn open(path: impl AsRef) -> Result { + let mut file = File::open(path)?; + let size = file.seek(SeekFrom::End(0))?; + if size < (MAGIC.len() * 2 + 4) as u64 { + return Err(Error::storage("file is too short to be a .qfc".to_string())); + } + + let mut head = [0u8; 4]; + file.seek(SeekFrom::Start(0))?; + file.read_exact(&mut head)?; + if &head != MAGIC { + return Err(Error::storage( + "file does not start with the .qfc magic".to_string(), + )); + } + + let mut tail = [0u8; 8]; + file.seek(SeekFrom::End(-8))?; + file.read_exact(&mut tail)?; + if &tail[4..] != MAGIC { + return Err(Error::storage( + "file does not end with the .qfc magic — it may be truncated".to_string(), + )); + } + let footer_len = u32::from_le_bytes([tail[0], tail[1], tail[2], tail[3]]) as u64; + if footer_len + 8 > size { + return Err(Error::storage( + "footer length exceeds file size".to_string(), + )); + } + + file.seek(SeekFrom::Start(size - 8 - footer_len))?; + let mut footer = vec![0u8; footer_len as usize]; + file.read_exact(&mut footer)?; + let meta = decode_footer(&footer)?; + + Ok(QfcReader { file, meta }) + } + + pub fn schema(&self) -> &Arc { + &self.meta.schema + } + + pub fn meta(&self) -> &FileMeta { + &self.meta + } + + pub fn num_row_groups(&self) -> usize { + self.meta.row_groups.len() + } + + pub fn num_rows(&self) -> usize { + self.meta.num_rows() + } + + pub fn row_group(&self, i: usize) -> Result<&RowGroupMeta> { + self.meta + .row_groups + .get(i) + .ok_or_else(|| Error::storage(format!("no row group {i}"))) + } + + /// Reads exactly the chunks named by `columns` from one row group. + /// + /// The returned batch's schema is the projected one, so a caller asking for + /// two of forty columns gets a two-column batch and the other thirty-eight + /// chunks are never touched. + pub fn read_row_group(&mut self, group: usize, columns: &[usize]) -> Result { + let meta = self.row_group(group)?.clone(); + let schema = Arc::new(self.meta.schema.project(columns)?); + let mut arrays: Vec = Vec::with_capacity(columns.len()); + for &ci in columns { + let chunk = meta.column(ci)?; + let data_type = self.meta.schema.field(ci)?.data_type; + let mut buf = vec![0u8; chunk.length as usize]; + self.file.seek(SeekFrom::Start(chunk.offset))?; + self.file.read_exact(&mut buf)?; + let array = decode_column(data_type, &buf)?; + if array.len() != meta.num_rows { + return Err(Error::storage(format!( + "column {ci} of row group {group} decoded {} rows, expected {}", + array.len(), + meta.num_rows + ))); + } + arrays.push(array); + } + if arrays.is_empty() { + // A count-only query projects nothing, but the row count still has + // to come back — and it comes from the footer, without reading a + // single byte of column data. + return RecordBatch::empty(schema)?.take(&vec![0usize; meta.num_rows]); + } + RecordBatch::try_new(schema, arrays) + } + + /// Every row group, every column — the unoptimised path, kept for tests + /// and for `\dump`. + pub fn read_all(&mut self) -> Result> { + let all: Vec = (0..self.meta.schema.len()).collect(); + let mut out = Vec::with_capacity(self.num_row_groups()); + for g in 0..self.num_row_groups() { + out.push(self.read_row_group(g, &all)?); + } + Ok(out) + } +} + +/// Writes a set of batches to a file in one call. +pub fn write_file( + path: impl AsRef, + schema: Arc, + batches: &[RecordBatch], + row_group_size: usize, +) -> Result { + let mut w = QfcWriter::create_with_row_group_size(path, schema, row_group_size)?; + for b in batches { + w.write_batch(b)?; + } + w.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::array::Array; + use qf_common::Value; + use std::path::PathBuf; + + struct TempPath(PathBuf); + + impl TempPath { + fn new(name: &str) -> TempPath { + let mut p = std::env::temp_dir(); + p.push(format!( + "qf-{}-{}-{}.qfc", + name, + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + TempPath(p) + } + fn path(&self) -> &Path { + &self.0 + } + } + + impl Drop for TempPath { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } + } + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + Field::new("amount", DataType::Float64, true), + ])) + } + + fn batch(range: std::ops::Range) -> RecordBatch { + let ids: Vec = range.clone().map(Value::Int64).collect(); + let regions: Vec = range + .clone() + .map(|i| Value::Utf8(format!("r{}", i % 4))) + .collect(); + let amounts: Vec = range + .map(|i| { + if i % 10 == 0 { + Value::Null + } else { + Value::Float64(i as f64 * 1.5) + } + }) + .collect(); + RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &ids).unwrap(), + Array::from_values(DataType::Utf8, ®ions).unwrap(), + Array::from_values(DataType::Float64, &amounts).unwrap(), + ], + ) + .unwrap() + } + + #[test] + fn a_file_round_trips_through_write_and_read() { + let tmp = TempPath::new("roundtrip"); + let original = batch(0..250); + write_file(tmp.path(), schema(), std::slice::from_ref(&original), 100).unwrap(); + + let mut r = QfcReader::open(tmp.path()).unwrap(); + assert_eq!(r.schema().as_ref(), schema().as_ref()); + assert_eq!(r.num_rows(), 250); + assert_eq!(r.num_row_groups(), 3); + let read_back = RecordBatch::concat(schema(), &r.read_all().unwrap()).unwrap(); + assert_eq!(read_back, original); + } + + #[test] + fn batches_are_repacked_into_even_row_groups() { + let tmp = TempPath::new("repack"); + // Three uneven batches must still come out as 100/100/40. + let meta = write_file( + tmp.path(), + schema(), + &[batch(0..30), batch(30..210), batch(210..240)], + 100, + ) + .unwrap(); + let sizes: Vec = meta.row_groups.iter().map(|g| g.num_rows).collect(); + assert_eq!(sizes, vec![100, 100, 40]); + assert_eq!(meta.num_rows(), 240); + } + + #[test] + fn reading_a_subset_of_columns_returns_only_those_columns() { + let tmp = TempPath::new("project"); + write_file(tmp.path(), schema(), &[batch(0..50)], 100).unwrap(); + let mut r = QfcReader::open(tmp.path()).unwrap(); + let b = r.read_row_group(0, &[2, 0]).unwrap(); + assert_eq!(b.num_columns(), 2); + assert_eq!(b.schema().field(0).unwrap().name, "amount"); + assert_eq!(b.schema().field(1).unwrap().name, "id"); + assert_eq!(b.num_rows(), 50); + assert_eq!(b.column(1).unwrap().value(3), Value::Int64(3)); + } + + #[test] + fn each_column_chunk_carries_a_zone_map_the_planner_can_prune_with() { + let tmp = TempPath::new("zonemap"); + // Row groups of 100 over a sorted id column: group 0 holds 0..99, + // group 1 holds 100..199. + let meta = write_file(tmp.path(), schema(), &[batch(0..200)], 100).unwrap(); + let g0 = &meta.row_groups[0].columns[0].stats; + let g1 = &meta.row_groups[1].columns[0].stats; + assert_eq!(g0.min, Value::Int64(0)); + assert_eq!(g0.max, Value::Int64(99)); + assert_eq!(g1.min, Value::Int64(100)); + assert_eq!(g1.max, Value::Int64(199)); + + use crate::stats::PruneOp; + assert!(!g0.may_match(PruneOp::Gt, &Value::Int64(150))); + assert!(g1.may_match(PruneOp::Gt, &Value::Int64(150))); + } + + #[test] + fn null_counts_are_recorded_per_row_group() { + let tmp = TempPath::new("nulls"); + let meta = write_file(tmp.path(), schema(), &[batch(0..100)], 100).unwrap(); + assert_eq!(meta.row_groups[0].columns[2].stats.null_count, 10); + assert_eq!(meta.row_groups[0].columns[0].stats.null_count, 0); + } + + #[test] + fn the_low_cardinality_column_picks_a_compressing_encoding() { + let tmp = TempPath::new("encodings"); + let meta = write_file(tmp.path(), schema(), &[batch(0..1000)], 1000).unwrap(); + // `region` cycles through four values, so it must not stay plain. + assert_ne!(meta.row_groups[0].columns[1].encoding, Encoding::Plain); + assert!(meta.row_groups[0].byte_size() > 0); + } + + #[test] + fn an_empty_batch_writes_no_row_group() { + let tmp = TempPath::new("emptybatch"); + let meta = write_file( + tmp.path(), + schema(), + &[RecordBatch::empty(schema()).unwrap()], + 100, + ) + .unwrap(); + assert_eq!(meta.row_groups.len(), 0); + let r = QfcReader::open(tmp.path()).unwrap(); + assert_eq!(r.num_rows(), 0); + assert_eq!(r.num_row_groups(), 0); + } + + #[test] + fn writing_a_batch_with_the_wrong_schema_is_rejected() { + let tmp = TempPath::new("badschema"); + let other = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)])); + let mut w = QfcWriter::create(tmp.path(), schema()).unwrap(); + assert!(w + .write_batch(&RecordBatch::empty(other).unwrap()) + .is_err_and(|e| e.to_string().contains("does not match"))); + w.finish().unwrap(); + } + + #[test] + fn a_zero_row_group_size_is_rejected() { + let tmp = TempPath::new("zerorg"); + assert!(QfcWriter::create_with_row_group_size(tmp.path(), schema(), 0).is_err()); + } + + #[test] + fn a_file_with_the_wrong_magic_is_refused() { + let tmp = TempPath::new("badmagic"); + std::fs::write(tmp.path(), b"NOPEnot-a-real-file-at-allQFC1").unwrap(); + assert!(QfcReader::open(tmp.path()) + .unwrap_err() + .to_string() + .contains("magic")); + } + + #[test] + fn a_truncated_file_is_refused_rather_than_read_as_garbage() { + let tmp = TempPath::new("truncated"); + write_file(tmp.path(), schema(), &[batch(0..50)], 100).unwrap(); + let bytes = std::fs::read(tmp.path()).unwrap(); + std::fs::write(tmp.path(), &bytes[..bytes.len() - 6]).unwrap(); + assert!(QfcReader::open(tmp.path()).is_err()); + } + + #[test] + fn a_file_shorter_than_its_own_header_is_refused() { + let tmp = TempPath::new("tiny"); + std::fs::write(tmp.path(), b"QF").unwrap(); + assert!(QfcReader::open(tmp.path()) + .unwrap_err() + .to_string() + .contains("too short")); + } + + #[test] + fn a_corrupt_footer_length_is_refused() { + let tmp = TempPath::new("badfooterlen"); + write_file(tmp.path(), schema(), &[batch(0..10)], 100).unwrap(); + let mut bytes = std::fs::read(tmp.path()).unwrap(); + let n = bytes.len(); + bytes[n - 8..n - 4].copy_from_slice(&u32::MAX.to_le_bytes()); + std::fs::write(tmp.path(), &bytes).unwrap(); + assert!(QfcReader::open(tmp.path()) + .unwrap_err() + .to_string() + .contains("exceeds file size")); + } + + #[test] + fn asking_for_a_row_group_that_does_not_exist_is_an_error() { + let tmp = TempPath::new("norg"); + write_file(tmp.path(), schema(), &[batch(0..10)], 100).unwrap(); + let mut r = QfcReader::open(tmp.path()).unwrap(); + assert!(r.row_group(4).is_err()); + assert!(r.read_row_group(0, &[9]).is_err()); + } + + #[test] + fn footer_metadata_survives_its_own_encoding() { + let meta = FileMeta { + schema: schema(), + row_groups: vec![RowGroupMeta { + num_rows: 5, + columns: vec![ + ColumnChunkMeta { + offset: 4, + length: 20, + encoding: Encoding::Rle, + stats: ColumnStats::empty(), + }; + 3 + ], + }], + }; + let buf = encode_footer(&meta); + assert_eq!(decode_footer(&buf).unwrap(), meta); + } + + #[test] + fn a_footer_whose_column_count_contradicts_the_schema_is_refused() { + let meta = FileMeta { + schema: schema(), + row_groups: vec![RowGroupMeta { + num_rows: 1, + columns: vec![ColumnChunkMeta { + offset: 0, + length: 1, + encoding: Encoding::Plain, + stats: ColumnStats::empty(), + }], + }], + }; + let buf = encode_footer(&meta); + assert!(decode_footer(&buf) + .unwrap_err() + .to_string() + .contains("but the schema has")); + } + + #[test] + fn unknown_tags_in_a_footer_are_refused() { + assert!(type_from_tag(7).is_err()); + for dt in [ + DataType::Boolean, + DataType::Int64, + DataType::Float64, + DataType::Utf8, + ] { + assert_eq!(type_from_tag(type_tag(dt)).unwrap(), dt); + } + } +} diff --git a/crates/qf-storage/src/lib.rs b/crates/qf-storage/src/lib.rs index ff7bd09..791961a 100644 --- a/crates/qf-storage/src/lib.rs +++ b/crates/qf-storage/src/lib.rs @@ -1 +1,9 @@ -// placeholder +pub mod array; +pub mod batch; +pub mod bitmap; +pub mod bytes; +pub mod catalog; +pub mod csv; +pub mod encoding; +pub mod format; +pub mod stats; diff --git a/crates/qf-storage/src/stats.rs b/crates/qf-storage/src/stats.rs new file mode 100644 index 0000000..71093d8 --- /dev/null +++ b/crates/qf-storage/src/stats.rs @@ -0,0 +1,385 @@ +use crate::array::Array; +use crate::bytes::{ByteReader, ByteWriter}; +use qf_common::{DataType, Result, Value}; +use std::cmp::Ordering; +use std::collections::HashSet; + +/// The comparison operators a zone map can reason about. +/// +/// The planner maps its own binary operators onto these; the storage layer has +/// no opinion about SQL syntax. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PruneOp { + Eq, + NotEq, + Lt, + LtEq, + Gt, + GtEq, +} + +/// Per-column, per-row-group summary — a zone map. +/// +/// This is the thing that makes predicate pushdown worth doing. Without it, +/// pushing a filter down only saves the cost of materialising rows; with it, +/// a row group whose `[min, max]` cannot satisfy the predicate is never read +/// off disk at all. +#[derive(Debug, Clone, PartialEq)] +pub struct ColumnStats { + pub min: Value, + pub max: Value, + pub null_count: usize, + pub row_count: usize, + /// Exact for row groups (they are bounded), summed and capped when merged + /// across row groups. The planner uses it to estimate join cardinality. + pub distinct_count: usize, +} + +impl ColumnStats { + pub fn empty() -> Self { + ColumnStats { + min: Value::Null, + max: Value::Null, + null_count: 0, + row_count: 0, + distinct_count: 0, + } + } + + pub fn from_array(array: &Array) -> ColumnStats { + let mut min = Value::Null; + let mut max = Value::Null; + let mut distinct: HashSet = HashSet::new(); + for i in 0..array.len() { + let v = array.value(i); + if v.is_null() { + continue; + } + min = Value::min(&min, &v); + max = Value::max(&max, &v); + distinct.insert(v); + } + ColumnStats { + min, + max, + null_count: array.null_count(), + row_count: array.len(), + distinct_count: distinct.len(), + } + } + + pub fn all_null(&self) -> bool { + self.row_count > 0 && self.null_count == self.row_count + } + + /// Selectivity is estimated as `1 / distinct` for equality, the textbook + /// uniform-distribution assumption. It is wrong on skewed data, and the + /// join-ordering rule that depends on it says so. + pub fn selectivity_eq(&self) -> f64 { + if self.distinct_count == 0 { + return 1.0; + } + 1.0 / self.distinct_count as f64 + } + + /// Combines the stats of two row groups into one summary of both. + pub fn merge(a: &ColumnStats, b: &ColumnStats) -> ColumnStats { + ColumnStats { + min: Value::min(&a.min, &b.min), + max: Value::max(&a.max, &b.max), + null_count: a.null_count + b.null_count, + row_count: a.row_count + b.row_count, + // Distinct values may overlap between groups, so the true count is + // somewhere between the larger of the two and their sum. Capping at + // the row count keeps the estimate from claiming more distinct + // values than there are rows. + distinct_count: (a.distinct_count + b.distinct_count) + .min(a.row_count + b.row_count) + .max(a.distinct_count.max(b.distinct_count)), + } + } + + /// Can any row in this range satisfy `column literal`? + /// + /// Answering `true` when the truth is `false` only costs a wasted read; + /// answering `false` when the truth is `true` drops rows from the result. + /// Every branch here defaults to `true` when it cannot prove otherwise. + pub fn may_match(&self, op: PruneOp, literal: &Value) -> bool { + if self.row_count == 0 { + return false; + } + // A comparison against NULL is never true, whichever side it is on. + if literal.is_null() || self.all_null() { + return false; + } + if self.min.is_null() || self.max.is_null() { + return true; + } + let (Ok(Some(vs_min)), Ok(Some(vs_max))) = ( + literal.sql_compare(&self.min), + literal.sql_compare(&self.max), + ) else { + // Types the zone map cannot compare — let the row group through and + // let the execution-time predicate produce the error or the answer. + return true; + }; + + match op { + // literal must sit inside [min, max] + PruneOp::Eq => vs_min != Ordering::Less && vs_max != Ordering::Greater, + // only prunable when the whole group is one value equal to the literal + PruneOp::NotEq => { + !(self.null_count == 0 && vs_min == Ordering::Equal && vs_max == Ordering::Equal) + } + // column < literal needs min < literal + PruneOp::Lt => vs_min == Ordering::Greater, + PruneOp::LtEq => vs_min != Ordering::Less, + // column > literal needs max > literal + PruneOp::Gt => vs_max == Ordering::Less, + PruneOp::GtEq => vs_max != Ordering::Greater, + } + } + + pub(crate) fn encode(&self, w: &mut ByteWriter) { + w.uvarint(self.null_count as u64); + w.uvarint(self.row_count as u64); + w.uvarint(self.distinct_count as u64); + encode_value(w, &self.min); + encode_value(w, &self.max); + } + + pub(crate) fn decode(r: &mut ByteReader<'_>) -> Result { + let null_count = r.uvarint()? as usize; + let row_count = r.uvarint()? as usize; + let distinct_count = r.uvarint()? as usize; + let min = decode_value(r)?; + let max = decode_value(r)?; + Ok(ColumnStats { + min, + max, + null_count, + row_count, + distinct_count, + }) + } +} + +pub(crate) fn encode_value(w: &mut ByteWriter, v: &Value) { + match v { + Value::Null => w.u8(0), + Value::Boolean(b) => { + w.u8(1); + w.u8(u8::from(*b)); + } + Value::Int64(i) => { + w.u8(2); + w.ivarint(*i); + } + Value::Float64(f) => { + w.u8(3); + w.f64(*f); + } + Value::Utf8(s) => { + w.u8(4); + w.string(s); + } + } +} + +pub(crate) fn decode_value(r: &mut ByteReader<'_>) -> Result { + Ok(match r.u8()? { + 0 => Value::Null, + 1 => Value::Boolean(r.u8()? != 0), + 2 => Value::Int64(r.ivarint()?), + 3 => Value::Float64(r.f64()?), + 4 => Value::Utf8(r.string()?), + tag => { + return Err(qf_common::Error::storage(format!( + "unknown value tag {tag} in file metadata" + ))) + } + }) +} + +/// Statistics for a whole table, one entry per column. +#[derive(Debug, Clone, PartialEq, Default)] +pub struct TableStats { + pub row_count: usize, + pub columns: Vec, +} + +impl TableStats { + pub fn column(&self, i: usize) -> Option<&ColumnStats> { + self.columns.get(i) + } +} + +/// Convenience for building stats from a fully materialised column. +pub fn stats_for(data_type: DataType, values: &[Value]) -> Result { + Ok(ColumnStats::from_array(&Array::from_values( + data_type, values, + )?)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ints(vals: &[Option]) -> ColumnStats { + let values: Vec = vals + .iter() + .map(|v| v.map_or(Value::Null, Value::Int64)) + .collect(); + stats_for(DataType::Int64, &values).unwrap() + } + + #[test] + fn min_max_and_null_count_come_from_the_data() { + let s = ints(&[Some(5), None, Some(1), Some(9)]); + assert_eq!(s.min, Value::Int64(1)); + assert_eq!(s.max, Value::Int64(9)); + assert_eq!(s.null_count, 1); + assert_eq!(s.row_count, 4); + assert_eq!(s.distinct_count, 3); + } + + #[test] + fn equality_prunes_a_group_whose_range_excludes_the_literal() { + let s = ints(&[Some(10), Some(20)]); + assert!(!s.may_match(PruneOp::Eq, &Value::Int64(5))); + assert!(!s.may_match(PruneOp::Eq, &Value::Int64(25))); + assert!(s.may_match(PruneOp::Eq, &Value::Int64(15))); + assert!(s.may_match(PruneOp::Eq, &Value::Int64(10))); + assert!(s.may_match(PruneOp::Eq, &Value::Int64(20))); + } + + #[test] + fn range_predicates_prune_from_the_correct_end() { + let s = ints(&[Some(10), Some(20)]); + // column < 10 is impossible when the smallest value is 10 + assert!(!s.may_match(PruneOp::Lt, &Value::Int64(10))); + assert!(s.may_match(PruneOp::Lt, &Value::Int64(11))); + assert!(s.may_match(PruneOp::LtEq, &Value::Int64(10))); + assert!(!s.may_match(PruneOp::LtEq, &Value::Int64(9))); + // column > 20 is impossible when the largest value is 20 + assert!(!s.may_match(PruneOp::Gt, &Value::Int64(20))); + assert!(s.may_match(PruneOp::Gt, &Value::Int64(19))); + assert!(s.may_match(PruneOp::GtEq, &Value::Int64(20))); + assert!(!s.may_match(PruneOp::GtEq, &Value::Int64(21))); + } + + #[test] + fn inequality_only_prunes_a_constant_group() { + let constant = ints(&[Some(7), Some(7)]); + assert!(!constant.may_match(PruneOp::NotEq, &Value::Int64(7))); + assert!(constant.may_match(PruneOp::NotEq, &Value::Int64(8))); + + let varied = ints(&[Some(7), Some(8)]); + assert!(varied.may_match(PruneOp::NotEq, &Value::Int64(7))); + + // A null alongside the constant still satisfies nothing under `<>`, + // but the zone map cannot see that, so it must not prune. + let with_null = ints(&[Some(7), None]); + assert!(with_null.may_match(PruneOp::NotEq, &Value::Int64(7))); + } + + #[test] + fn an_all_null_group_matches_no_predicate() { + let s = ints(&[None, None]); + assert!(s.all_null()); + for op in [ + PruneOp::Eq, + PruneOp::NotEq, + PruneOp::Lt, + PruneOp::LtEq, + PruneOp::Gt, + PruneOp::GtEq, + ] { + assert!(!s.may_match(op, &Value::Int64(1)), "{op:?}"); + } + } + + #[test] + fn comparing_against_null_never_matches() { + let s = ints(&[Some(1), Some(2)]); + assert!(!s.may_match(PruneOp::Eq, &Value::Null)); + } + + #[test] + fn an_empty_group_is_always_pruned() { + assert!(!ColumnStats::empty().may_match(PruneOp::Eq, &Value::Int64(1))); + } + + #[test] + fn an_incomparable_literal_leaves_the_group_readable() { + // Comparing a string against an int column cannot be decided by the + // zone map, so the group survives and execution reports the type error. + let s = ints(&[Some(1), Some(2)]); + assert!(s.may_match(PruneOp::Eq, &Value::Utf8("x".into()))); + } + + #[test] + fn string_zone_maps_prune_lexicographically() { + let s = stats_for( + DataType::Utf8, + &[Value::Utf8("banana".into()), Value::Utf8("cherry".into())], + ) + .unwrap(); + assert!(!s.may_match(PruneOp::Eq, &Value::Utf8("apple".into()))); + assert!(s.may_match(PruneOp::Eq, &Value::Utf8("blue".into()))); + assert!(!s.may_match(PruneOp::Gt, &Value::Utf8("cherry".into()))); + } + + #[test] + fn merging_widens_the_range_and_sums_the_counts() { + let a = ints(&[Some(1), Some(5)]); + let b = ints(&[Some(10), None]); + let m = ColumnStats::merge(&a, &b); + assert_eq!(m.min, Value::Int64(1)); + assert_eq!(m.max, Value::Int64(10)); + assert_eq!(m.row_count, 4); + assert_eq!(m.null_count, 1); + assert!(m.distinct_count <= m.row_count); + assert!(m.distinct_count >= 2); + } + + #[test] + fn selectivity_falls_as_distinct_values_rise() { + assert!((ints(&[Some(1), Some(2)]).selectivity_eq() - 0.5).abs() < 1e-9); + assert!((ColumnStats::empty().selectivity_eq() - 1.0).abs() < 1e-9); + } + + #[test] + fn stats_survive_a_serialisation_round_trip() { + for s in [ + ints(&[Some(1), None, Some(3)]), + stats_for(DataType::Utf8, &[Value::Utf8("a".into())]).unwrap(), + stats_for(DataType::Float64, &[Value::Float64(1.5)]).unwrap(), + stats_for(DataType::Boolean, &[Value::Boolean(true)]).unwrap(), + ColumnStats::empty(), + ] { + let mut w = ByteWriter::new(); + s.encode(&mut w); + let buf = w.into_bytes(); + let mut r = ByteReader::new(&buf); + assert_eq!(ColumnStats::decode(&mut r).unwrap(), s); + } + } + + #[test] + fn an_unknown_value_tag_is_a_storage_error() { + let buf = vec![9u8]; + let mut r = ByteReader::new(&buf); + assert!(decode_value(&mut r).is_err()); + } + + #[test] + fn table_stats_index_by_column() { + let t = TableStats { + row_count: 2, + columns: vec![ints(&[Some(1)])], + }; + assert!(t.column(0).is_some()); + assert!(t.column(1).is_none()); + } +}