use std::fmt;
use crate::error::{Error, Result};
use crate::types::{Field, LogicalType};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum Width {
#[default]
Exact,
Month,
Quarter,
Year,
Auto,
}
impl Width {
pub const DEFAULT: Self = Self::Quarter;
pub const TARGET: u64 = 200_000;
#[must_use]
pub fn for_span(rows: u64, days: u64) -> Self {
if rows == 0 || days == 0 {
return Self::DEFAULT;
}
let partitions = days / 30 + 1;
if rows / partitions >= Self::TARGET { Self::Month } else { Self::Quarter }
}
#[must_use]
pub fn tag(self) -> u8 {
match self {
Self::Exact => 0,
Self::Month => 1,
Self::Quarter => 2,
Self::Year => 3,
Self::Auto => 4,
}
}
#[must_use]
pub fn from_tag(tag: u8) -> Option<Self> {
match tag {
0 => Some(Self::Exact),
1 => Some(Self::Month),
2 => Some(Self::Quarter),
3 => Some(Self::Year),
4 => Some(Self::Auto),
_ => None,
}
}
}
impl fmt::Display for Width {
fn fmt(&self, out: &mut fmt::Formatter<'_>) -> fmt::Result {
out.write_str(match self {
Self::Exact => "EXACT",
Self::Month => "MONTH",
Self::Quarter => "QUARTER",
Self::Year => "YEAR",
Self::Auto => "AUTO",
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Clustering {
columns: Vec<u32>,
width: Width,
}
impl Clustering {
pub fn new(columns: Vec<u32>, width: Width, fields: &[Field]) -> Result<Self> {
if columns.is_empty() {
return Err(Error::invalid_input("a clustering declaration names no column"));
}
for (at, &column) in columns.iter().enumerate() {
if column as usize >= fields.len() {
return Err(Error::invalid_input(
"a clustering declaration names a column the table does not have",
));
}
if columns[..at].contains(&column) {
return Err(Error::invalid_input(
"a clustering declaration names the same column twice",
));
}
}
let leading = &fields[columns[0] as usize];
if width != Width::Exact
&& !matches!(leading.ty, LogicalType::Date | LogicalType::Timestamp)
{
return Err(Error::invalid_input(format!(
"a clustering declaration buckets {} by {width}, which only a date or a timestamp \
has",
leading.name
)));
}
Ok(Self { columns, width })
}
pub fn over(columns: Vec<u32>, fields: &[Field]) -> Result<Self> {
let leading = columns.first().and_then(|&at| fields.get(at as usize));
let width = match leading.map(|field| &field.ty) {
Some(LogicalType::Date | LogicalType::Timestamp) => Width::Auto,
_ => Width::Exact,
};
Self::new(columns, width, fields)
}
#[must_use]
pub fn fitted(&self, rows: u64, days: u64) -> Self {
match self.width {
Width::Auto => {
Self { columns: self.columns.clone(), width: Width::for_span(rows, days) }
}
_ => self.clone(),
}
}
#[must_use]
pub fn columns(&self) -> &[u32] {
&self.columns
}
#[must_use]
pub fn width(&self) -> Width {
self.width
}
#[must_use]
pub fn partition(&self) -> u32 {
self.columns[0]
}
#[must_use]
pub fn describe(&self, names: &[String]) -> String {
let named =
|at: u32| names.get(at as usize).cloned().unwrap_or_else(|| format!("column {at}"));
let leading = match self.width {
Width::Exact => named(self.partition()),
other => format!("{}({})", other.to_string().to_lowercase(), named(self.partition())),
};
let rest = self.columns[1..].iter().map(|&at| named(at)).collect::<Vec<_>>();
std::iter::once(leading).chain(rest).collect::<Vec<_>>().join(", ")
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Declared {
table: String,
width: Option<Width>,
columns: Vec<String>,
}
impl Declared {
#[must_use]
pub fn table(&self) -> &str {
&self.table
}
#[must_use]
pub fn width(&self) -> Option<Width> {
self.width
}
#[must_use]
pub fn columns(&self) -> &[String] {
&self.columns
}
}
#[must_use]
pub fn is_clustering_setting(name: &str) -> bool {
name.eq_ignore_ascii_case("cluster_by") || name.eq_ignore_ascii_case("rudb.cluster_by")
}
pub fn parse_clustering(setting: &str) -> Result<Vec<Declared>> {
let mut declared = Vec::new();
for entry in entries(setting) {
declared.push(parse_entry(&entry)?);
}
Ok(declared)
}
fn entries(setting: &str) -> Vec<String> {
let mut entries = Vec::new();
let mut current = String::new();
let mut depth = 0usize;
for character in setting.chars() {
match character {
'(' => depth += 1,
')' => depth = depth.saturating_sub(1),
',' if depth == 0 => {
entries.push(std::mem::take(&mut current));
continue;
}
_ => {}
}
current.push(character);
}
entries.push(current);
entries
.into_iter()
.map(|entry| entry.trim().to_owned())
.filter(|entry| !entry.is_empty())
.collect()
}
fn parse_entry(entry: &str) -> Result<Declared> {
let Some((table, rest)) = entry.split_once('(') else {
return Err(malformed(format!("expected `table(column, ...)` and found `{entry}`")));
};
let Some(inside) = rest.trim_end().strip_suffix(')') else {
return Err(malformed(format!("`{entry}` is missing its closing parenthesis")));
};
let table = table.trim();
if table.is_empty() {
return Err(malformed(format!("`{entry}` names no table")));
}
let mut columns = Vec::new();
for column in inside.split(',').map(str::trim).filter(|column| !column.is_empty()) {
columns.push(column.to_owned());
}
if columns.is_empty() {
return Err(malformed(format!("`{entry}` names no column")));
}
let (width, leading) = split_width(&columns[0])?;
for column in &columns[1..] {
if column.contains('(') {
return Err(malformed(format!(
"`{column}` is bucketed and only the leading column of `{table}` can be"
)));
}
}
columns[0] = leading;
Ok(Declared { table: table.to_owned(), width, columns })
}
fn split_width(leading: &str) -> Result<(Option<Width>, String)> {
let Some((word, rest)) = leading.split_once('(') else {
return Ok((None, leading.to_owned()));
};
let Some(column) = rest.trim_end().strip_suffix(')') else {
return Err(malformed(format!("`{leading}` is missing its closing parenthesis")));
};
let column = column.trim();
if column.is_empty() {
return Err(malformed(format!("`{leading}` names no column")));
}
let word = word.trim();
let width = [Width::Exact, Width::Month, Width::Quarter, Width::Year, Width::Auto]
.into_iter()
.find(|width| width.to_string().eq_ignore_ascii_case(word))
.ok_or_else(|| {
malformed(format!(
"`{word}` is not a partition width, which is one of exact, month, quarter, year or \
auto"
))
})?;
Ok((Some(width), column.to_owned()))
}
fn malformed(message: impl Into<String>) -> Error {
Error::invalid_input(format!("invalid rudb clustering: {}", message.into()))
}
#[cfg(test)]
mod tests {
use super::{Clustering, Width, parse_clustering};
use crate::types::{Field, LogicalType};
fn lineitem() -> Vec<Field> {
vec![
Field::new("l_orderkey", LogicalType::BigInt),
Field::new("l_linenumber", LogicalType::Integer),
Field::new("l_shipdate", LogicalType::Date),
]
}
#[test]
fn a_declaration_that_could_never_be_satisfied_is_refused() {
let fields = lineitem();
assert!(Clustering::new(Vec::new(), Width::Exact, &fields).is_err(), "no column at all");
assert!(Clustering::new(vec![3], Width::Exact, &fields).is_err(), "past the end");
assert!(Clustering::new(vec![0, 1, 0], Width::Exact, &fields).is_err(), "twice");
assert!(Clustering::new(vec![2, 0], Width::Month, &fields).is_ok());
}
#[test]
fn a_calendar_bucket_on_a_column_with_no_calendar_in_it_is_refused() {
let fields = lineitem();
let complaint = Clustering::new(vec![0, 2], Width::Month, &fields)
.expect_err("a bigint has no months")
.to_string();
assert!(complaint.contains("l_orderkey"), "{complaint}");
assert!(complaint.contains("MONTH"), "{complaint}");
assert!(
Clustering::new(vec![0, 2], Width::Exact, &fields).is_ok(),
"no bucket, no problem"
);
}
#[test]
fn a_declaration_with_no_width_takes_the_quarter_on_a_date_and_nothing_elsewhere() {
let fields = lineitem();
let dated = Clustering::over(vec![2, 0, 1], &fields).expect("a date leads");
assert_eq!(dated.width(), Width::Auto, "nobody said, so the rows will say");
assert_eq!(dated.columns(), [2, 0, 1], "the columns are the ones asked for, in order");
let keyed = Clustering::over(vec![0, 2], &fields).expect("a bigint leads");
assert_eq!(keyed.width(), Width::Exact);
assert!(Clustering::over(Vec::new(), &fields).is_err(), "no column at all");
assert!(Clustering::over(vec![7], &fields).is_err(), "past the end");
}
#[test]
fn every_width_survives_its_byte() {
for width in [Width::Exact, Width::Month, Width::Quarter, Width::Year, Width::Auto] {
assert_eq!(Width::from_tag(width.tag()), Some(width), "{width}");
}
assert_eq!(Width::from_tag(5), None, "a tag from a build that knows more than this one");
}
#[test]
fn the_same_table_at_two_scales_gets_two_widths() {
let span = 2525;
assert_eq!(Width::for_span(6_001_215, span), Width::Quarter, "SF1");
assert_eq!(Width::for_span(59_986_052, span), Width::Month, "SF10");
assert_eq!(Width::for_span(600_572, span), Width::Quarter, "SF0.1");
assert_eq!(Width::for_span(1_500_000, span), Width::Quarter, "SF1 orders");
assert_eq!(Width::for_span(0, span), Width::DEFAULT);
assert_eq!(Width::for_span(6_001_215, 0), Width::DEFAULT);
}
#[test]
fn fitting_a_declaration_resolves_the_automatic_width_and_only_that_one() {
let fields = lineitem();
let auto = Clustering::over(vec![2, 0, 1], &fields).expect("a date leads");
assert_eq!(auto.width(), Width::Auto);
let fitted = auto.fitted(6_001_215, 2525);
assert_eq!(fitted.width(), Width::Quarter, "the rows decided");
assert_eq!(fitted.columns(), auto.columns(), "and nothing else moved");
for width in [Width::Exact, Width::Month, Width::Quarter, Width::Year] {
let asked = Clustering::new(vec![2, 0, 1], width, &fields).expect("valid");
assert_eq!(asked.fitted(59_986_052, 2525), asked, "{width}");
}
}
#[test]
fn the_two_clustered_tpch_tables_parse_into_two_declarations() {
let setting = "lineitem(quarter(l_shipdate), l_orderkey, l_linenumber), \
orders(quarter(o_orderdate), o_orderkey)";
let declared = parse_clustering(setting).expect("parse");
assert_eq!(declared.len(), 2, "a comma inside a column list is not a separator");
assert_eq!(declared[0].table(), "lineitem");
assert_eq!(declared[0].width(), Some(Width::Quarter));
assert_eq!(declared[0].columns(), ["l_shipdate", "l_orderkey", "l_linenumber"]);
assert_eq!(declared[1].table(), "orders");
assert_eq!(declared[1].columns(), ["o_orderdate", "o_orderkey"]);
let spread = "lineitem(month(l_shipdate), l_orderkey),\n orders(o_orderkey),\n";
let declared = parse_clustering(spread).expect("parse");
assert_eq!(declared.len(), 2);
assert_eq!(declared[0].width(), Some(Width::Month));
assert_eq!(declared[1].width(), None);
assert_eq!(parse_clustering("t(exact(d))").expect("parse")[0].width(), Some(Width::Exact));
assert!(parse_clustering("").expect("parse").is_empty(), "a reset names no table");
}
#[test]
fn what_a_declaration_describes_itself_as_parses_back_into_the_same_declaration() {
let fields = lineitem();
let names = fields.iter().map(|field| field.name.clone()).collect::<Vec<_>>();
let at = |name: &String| names.iter().position(|it| it == name).expect("a column") as u32;
for width in [Width::Month, Width::Quarter, Width::Year, Width::Auto] {
let asked = Clustering::new(vec![2, 0, 1], width, &fields).expect("valid");
let written = format!("lineitem({})", asked.describe(&names));
let read = parse_clustering(&written).expect("parse");
let columns: Vec<u32> = read[0].columns().iter().map(at).collect();
let again = Clustering::new(columns, read[0].width().expect("a width"), &fields)
.expect("valid");
assert_eq!(again, asked, "{written}");
}
let exact = Clustering::new(vec![2, 0, 1], Width::Exact, &fields).expect("valid");
let written = format!("lineitem({})", exact.describe(&names));
assert_eq!(written, "lineitem(l_shipdate, l_orderkey, l_linenumber)");
let read = parse_clustering(&written).expect("parse");
assert_eq!(read[0].width(), None, "a bare date column names no width");
let columns: Vec<u32> = read[0].columns().iter().map(at).collect();
assert_eq!(
Clustering::over(columns, &fields).expect("valid").width(),
Width::Auto,
"and a silent date declaration leaves the width to the data"
);
}
#[test]
fn a_declaration_that_is_not_a_table_and_a_column_list_is_refused() {
for bad in [
"lineitem",
"lineitem(",
"(l_shipdate)",
"lineitem()",
"lineitem(day(l_shipdate))",
"lineitem(l_shipdate, month(l_orderkey))",
"lineitem(month())",
] {
let complaint = parse_clustering(bad).expect_err(bad).message().to_owned();
assert!(complaint.contains("clustering"), "{bad}: {complaint}");
}
let declared = parse_clustering("t(year)").expect("parse");
assert_eq!(declared[0].columns(), ["year"]);
assert_eq!(declared[0].width(), None);
}
#[test]
fn a_declaration_reads_back_the_way_it_was_written() {
let fields = lineitem();
let names = fields.iter().map(|field| field.name.clone()).collect::<Vec<_>>();
let stage_zero = Clustering::new(vec![2, 0, 1], Width::Month, &fields).expect("valid");
assert_eq!(stage_zero.describe(&names), "month(l_shipdate), l_orderkey, l_linenumber");
let plain = Clustering::new(vec![0], Width::Exact, &fields).expect("valid");
assert_eq!(plain.describe(&names), "l_orderkey");
}
}