use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use spg_sql::ast::Expr;
use spg_storage::{Catalog, PartitionBound, PartitionRole, Value};
use crate::EngineError;
use crate::conversions::literal_expr_to_value;
pub(crate) fn has_children(catalog: &Catalog, table_name: &str) -> bool {
is_partition_parent(catalog, table_name) || !children_of_parent(catalog, table_name).is_empty()
}
pub(crate) fn has_inheritance_children(catalog: &Catalog, table_name: &str) -> bool {
children_of_parent(catalog, table_name).iter().any(|c| {
catalog.get(c).is_some_and(|t| {
matches!(
t.schema().partition_role,
Some(PartitionRole::Inherits { .. })
)
})
})
}
pub(crate) fn is_partition_parent(catalog: &Catalog, table_name: &str) -> bool {
catalog
.get(table_name)
.map(|t| {
matches!(
t.schema().partition_role,
Some(PartitionRole::Parent { .. })
)
})
.unwrap_or(false)
}
pub(crate) fn children_of_parent(catalog: &Catalog, parent: &str) -> Vec<String> {
let mut out = Vec::new();
for name in catalog.table_names() {
let Some(t) = catalog.get(&name) else {
continue;
};
match &t.schema().partition_role {
Some(PartitionRole::Range { parent_name, .. })
| Some(PartitionRole::Default { parent_name })
| Some(PartitionRole::List { parent_name, .. })
| Some(PartitionRole::Hash { parent_name, .. })
if parent_name == parent =>
{
out.push(name);
}
Some(PartitionRole::Inherits { parent_names })
if parent_names.iter().any(|p| p == parent) =>
{
out.push(name);
}
_ => {}
}
}
out
}
pub(crate) fn evaluate_partition_bound(expr: Expr) -> Result<PartitionBound, EngineError> {
if let Expr::FunctionCall { name, args } = &expr
&& args.is_empty()
{
let upper = name.to_ascii_uppercase();
if upper == "MINVALUE" {
return Ok(PartitionBound::MinValue);
}
if upper == "MAXVALUE" {
return Ok(PartitionBound::MaxValue);
}
}
let value = literal_expr_to_value(expr)?;
match value {
Value::Timestamp(micros) => Ok(PartitionBound::TimestampTz(micros)),
Value::Date(days) => Ok(PartitionBound::Date(days)),
Value::BigInt(n) => Ok(PartitionBound::BigInt(n)),
Value::Int(n) => Ok(PartitionBound::Int(n)),
Value::SmallInt(n) => Ok(PartitionBound::SmallInt(n)),
Value::Text(s) => {
match crate::eval::parse_timestamp_literal(&s) {
Some(micros) => Ok(PartitionBound::TimestampTz(micros)),
None => Ok(PartitionBound::Text(s.into_owned())),
}
}
other => Err(EngineError::Unsupported(format!(
"PARTITION OF: bound must be a typed literal or \
MINVALUE/MAXVALUE, got {other:?}"
))),
}
}
fn bound_i64(b: &PartitionBound) -> Option<i64> {
match b {
PartitionBound::TimestampTz(m) => Some(*m),
PartitionBound::BigInt(n) => Some(*n),
PartitionBound::Int(n) => Some(i64::from(*n)),
PartitionBound::SmallInt(n) => Some(i64::from(*n)),
PartitionBound::Date(d) => Some(i64::from(*d)),
PartitionBound::MinValue | PartitionBound::MaxValue | PartitionBound::Text(_) => None,
}
}
pub(crate) fn value_to_bound(v: &spg_storage::Value) -> Option<PartitionBound> {
use spg_storage::Value;
match v {
Value::Timestamp(m) => Some(PartitionBound::TimestampTz(*m)),
Value::BigInt(n) => Some(PartitionBound::BigInt(*n)),
Value::Int(n) => Some(PartitionBound::Int(*n)),
Value::SmallInt(n) => Some(PartitionBound::SmallInt(*n)),
Value::Date(d) => Some(PartitionBound::Date(*d)),
Value::Text(s) => match crate::eval::parse_timestamp_literal(s) {
Some(m) => Some(PartitionBound::TimestampTz(m)),
None => Some(PartitionBound::Text(s.clone().into_owned())),
},
_ => None,
}
}
fn bound_cmp(a: &PartitionBound, b: &PartitionBound) -> core::cmp::Ordering {
use PartitionBound::{MaxValue, MinValue, Text};
use core::cmp::Ordering;
match (a, b) {
(MinValue, MinValue) | (MaxValue, MaxValue) => Ordering::Equal,
(MinValue, _) => Ordering::Less,
(_, MinValue) => Ordering::Greater,
(MaxValue, _) => Ordering::Greater,
(_, MaxValue) => Ordering::Less,
(Text(x), Text(y)) => x.cmp(y),
_ => match (bound_i64(a), bound_i64(b)) {
(Some(x), Some(y)) => x.cmp(&y),
_ => Ordering::Equal,
},
}
}
pub(crate) fn ranges_overlap(
a_lo: &PartitionBound,
a_hi: &PartitionBound,
b_lo: &PartitionBound,
b_hi: &PartitionBound,
) -> bool {
use core::cmp::Ordering;
bound_cmp(a_lo, b_hi) == Ordering::Less && bound_cmp(b_lo, a_hi) == Ordering::Less
}
#[allow(dead_code)] pub(crate) fn value_in_range(
key: &PartitionBound,
lower: &PartitionBound,
upper: &PartitionBound,
) -> bool {
use core::cmp::Ordering;
bound_cmp(key, lower) != Ordering::Less && bound_cmp(key, upper) == Ordering::Less
}
#[allow(dead_code)] pub(crate) fn pg_compatible_hash(value: &Value<'_>) -> u64 {
const FNV_OFFSET: u64 = 0xcbf29ce484222325;
const FNV_PRIME: u64 = 0x100000001b3;
let mut h: u64 = FNV_OFFSET;
let mut feed = |bytes: &[u8]| {
for &b in bytes {
h ^= u64::from(b);
h = h.wrapping_mul(FNV_PRIME);
}
};
match value {
Value::Null => feed(&[0u8]),
Value::Bool(b) => {
feed(&[1u8]);
feed(&[u8::from(*b)]);
}
Value::SmallInt(n) => {
feed(&[2u8]);
feed(&n.to_le_bytes());
}
Value::Int(n) => {
feed(&[3u8]);
feed(&n.to_le_bytes());
}
Value::BigInt(n) => {
feed(&[4u8]);
feed(&n.to_le_bytes());
}
Value::Float(f) => {
feed(&[6u8]);
feed(&f.to_bits().to_le_bytes());
}
Value::Text(s) => {
feed(&[7u8]);
feed(s.as_bytes());
}
Value::Date(d) => {
feed(&[8u8]);
feed(&d.to_le_bytes());
}
Value::Timestamp(m) => {
feed(&[9u8]);
feed(&m.to_le_bytes());
}
other => {
feed(&[255u8]);
feed(format!("{other:?}").as_bytes());
}
}
h
}
pub(crate) fn bound_to_diag(b: &PartitionBound) -> String {
match b {
PartitionBound::MinValue => "MINVALUE".to_string(),
PartitionBound::MaxValue => "MAXVALUE".to_string(),
PartitionBound::TimestampTz(m) => format!("'{m}'::timestamptz"),
PartitionBound::BigInt(n) => format!("{n}::bigint"),
PartitionBound::Int(n) => format!("{n}::integer"),
PartitionBound::SmallInt(n) => format!("{n}::smallint"),
PartitionBound::Date(d) => format!("{d}::date"),
PartitionBound::Text(s) => format!("'{}'", s.replace('\'', "''")),
}
}
#[cfg(test)]
mod tests {
use super::*;
use PartitionBound::{MaxValue, MinValue, TimestampTz};
#[test]
fn bound_cmp_min_max_sentinels() {
use core::cmp::Ordering;
assert_eq!(bound_cmp(&MinValue, &MinValue), Ordering::Equal);
assert_eq!(bound_cmp(&MaxValue, &MaxValue), Ordering::Equal);
assert_eq!(bound_cmp(&MinValue, &TimestampTz(0)), Ordering::Less);
assert_eq!(bound_cmp(&TimestampTz(0), &MaxValue), Ordering::Less);
assert_eq!(
bound_cmp(&TimestampTz(100), &TimestampTz(99)),
Ordering::Greater
);
}
#[test]
fn ranges_overlap_half_open() {
let a_lo = TimestampTz(0);
let a_hi = TimestampTz(100);
let b_lo = TimestampTz(100);
let b_hi = TimestampTz(200);
let c_lo = TimestampTz(50);
let c_hi = TimestampTz(150);
assert!(!ranges_overlap(&a_lo, &a_hi, &b_lo, &b_hi));
assert!(ranges_overlap(&a_lo, &a_hi, &c_lo, &c_hi));
assert!(ranges_overlap(&b_lo, &b_hi, &c_lo, &c_hi));
assert!(ranges_overlap(&a_lo, &a_hi, &a_lo, &a_hi));
assert!(!ranges_overlap(&a_lo, &a_lo, &a_lo, &a_lo));
assert!(ranges_overlap(&MinValue, &MaxValue, &a_lo, &a_hi));
}
#[test]
fn value_in_range_inclusive_lower_exclusive_upper() {
use PartitionBound::{Date, Int};
assert!(value_in_range(
&TimestampTz(50),
&TimestampTz(0),
&TimestampTz(100)
));
assert!(value_in_range(
&TimestampTz(0),
&TimestampTz(0),
&TimestampTz(100)
));
assert!(!value_in_range(
&TimestampTz(100),
&TimestampTz(0),
&TimestampTz(100)
));
assert!(!value_in_range(
&TimestampTz(-1),
&TimestampTz(0),
&TimestampTz(100)
));
assert!(value_in_range(&TimestampTz(i64::MIN), &MinValue, &MaxValue));
assert!(value_in_range(
&TimestampTz(i64::MAX - 1),
&MinValue,
&MaxValue
));
assert!(!value_in_range(&TimestampTz(0), &MaxValue, &MaxValue));
assert!(!value_in_range(&TimestampTz(0), &MinValue, &MinValue));
assert!(value_in_range(&Int(5), &Int(0), &Int(10)));
assert!(!value_in_range(&Int(10), &Int(0), &Int(10))); assert!(value_in_range(&Int(0), &Int(0), &Int(10))); assert!(value_in_range(&Int(12), &Int(10), &MaxValue));
assert!(value_in_range(&Date(100), &Date(0), &Date(365)));
assert!(!value_in_range(&Date(400), &Date(0), &Date(365)));
}
}