use crate::sqlselect::{Expr, JoinKind};
use anyhow::Result;
use serde_json::Value;
use std::collections::HashMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JoinExec {
Auto,
NestedLoop,
Hash,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JoinChoice {
pub kind: JoinKind,
pub table: String,
pub strategy: Strategy,
pub keys: usize,
pub left_rows: usize,
pub right_rows: usize,
pub out_rows: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Strategy {
NestedLoop,
Hash,
}
impl std::fmt::Display for Strategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
Strategy::NestedLoop => "Nested Loop",
Strategy::Hash => "Hash Join",
})
}
}
pub const AUTO_HASH_MIN_PAIRS: usize = 64;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum HKey {
Num(u64),
Text(String),
}
fn canon(f: f64) -> u64 {
let f = if f == 0.0 { 0.0 } else { f };
f.to_bits()
}
pub fn hkey(v: &Value) -> Option<HKey> {
match v {
Value::Null => None,
Value::Number(n) => Some(match n.as_f64() {
Some(f) => HKey::Num(canon(f)),
None => HKey::Text(n.to_string()),
}),
Value::String(s) => match s.parse::<f64>() {
Ok(f) => Some(HKey::Num(canon(f))),
Err(_) => Some(HKey::Text(s.clone())),
},
Value::Bool(b) => Some(HKey::Text(if *b { "t" } else { "f" }.to_string())),
other => Some(HKey::Text(other.to_string())),
}
}
const PURE_FUNCS: &[&str] = &[
"lower",
"upper",
"length",
"char_length",
"character_length",
"coalesce",
"nullif",
"int2",
"int4",
"int8",
"text",
"quote_ident",
"format_type",
"array_to_string",
"current_schema",
"current_database",
"current_catalog",
"current_user",
"session_user",
"user",
"version",
"pg_get_userbyid",
"pg_table_is_visible",
"pg_type_is_visible",
"pg_function_is_visible",
"pg_encoding_to_char",
"pg_get_expr",
"pg_get_indexdef",
"pg_get_constraintdef",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Side {
Left,
Right,
Const,
Unusable,
}
fn side_of(e: &Expr, left: &[String], right: &str) -> Side {
let mut saw_left = false;
let mut saw_right = false;
let mut usable = true;
walk(e, left, right, &mut saw_left, &mut saw_right, &mut usable);
if !usable || (saw_left && saw_right) {
return Side::Unusable;
}
match (saw_left, saw_right) {
(true, false) => Side::Left,
(false, true) => Side::Right,
(false, false) => Side::Const,
(true, true) => unreachable!("handled above"),
}
}
fn walk(
e: &Expr,
left: &[String],
right: &str,
saw_left: &mut bool,
saw_right: &mut bool,
usable: &mut bool,
) {
match e {
Expr::Column { qual, .. } => match qual {
Some(q) => {
if q.eq_ignore_ascii_case(right) {
*saw_right = true;
} else if left.iter().any(|b| b.eq_ignore_ascii_case(q)) {
*saw_left = true;
} else {
*usable = false;
}
}
None => *usable = false,
},
Expr::Literal(_) => {}
Expr::Star | Expr::QualifiedStar(_) => *usable = false,
Expr::Func { name, args } => {
if !PURE_FUNCS.iter().any(|f| f.eq_ignore_ascii_case(name)) {
*usable = false;
}
for a in args {
walk(a, left, right, saw_left, saw_right, usable);
}
}
Expr::Agg { .. } => *usable = false,
Expr::Case { operand, whens, else_ } => {
if let Some(o) = operand {
walk(o, left, right, saw_left, saw_right, usable);
}
for (w, t) in whens {
walk(w, left, right, saw_left, saw_right, usable);
walk(t, left, right, saw_left, saw_right, usable);
}
if let Some(x) = else_ {
walk(x, left, right, saw_left, saw_right, usable);
}
}
Expr::Binary { left: l, right: r, .. } => {
walk(l, left, right, saw_left, saw_right, usable);
walk(r, left, right, saw_left, saw_right, usable);
}
Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
walk(expr, left, right, saw_left, saw_right, usable);
}
Expr::InList { expr, list, .. } => {
walk(expr, left, right, saw_left, saw_right, usable);
for i in list {
walk(i, left, right, saw_left, saw_right, usable);
}
}
Expr::Index { expr, index } => {
walk(expr, left, right, saw_left, saw_right, usable);
walk(index, left, right, saw_left, saw_right, usable);
}
Expr::ArrayLit(items) => {
for i in items {
walk(i, left, right, saw_left, saw_right, usable);
}
}
Expr::Subquery(_) | Expr::Exists { .. } | Expr::ArrayQuery(_) | Expr::InSubquery { .. } => {
*usable = false
}
Expr::Quantified { left: l, right: r, .. } => {
walk(l, left, right, saw_left, saw_right, usable);
walk(r, left, right, saw_left, saw_right, usable);
}
}
}
fn conjuncts<'a>(e: &'a Expr, out: &mut Vec<&'a Expr>) {
match e {
Expr::Binary { op, left, right } if op == "AND" => {
conjuncts(left, out);
conjuncts(right, out);
}
other => out.push(other),
}
}
pub fn hash_keys(on: Option<&Expr>, left: &[String], right: &str) -> Vec<(Expr, Expr)> {
let Some(on) = on else { return vec![] };
let mut parts = vec![];
conjuncts(on, &mut parts);
let mut keys = vec![];
for p in parts {
let Expr::Binary { op, left: l, right: r } = p else { continue };
if op != "=" {
continue;
}
match (side_of(l, left, right), side_of(r, left, right)) {
(Side::Left, Side::Right) => keys.push(((**l).clone(), (**r).clone())),
(Side::Right, Side::Left) => keys.push(((**r).clone(), (**l).clone())),
_ => {}
}
}
keys
}
pub fn choose(exec: JoinExec, keys: usize, left_rows: usize, right_rows: usize) -> Strategy {
if keys == 0 {
return Strategy::NestedLoop;
}
match exec {
JoinExec::NestedLoop => Strategy::NestedLoop,
JoinExec::Hash => Strategy::Hash,
JoinExec::Auto => {
if left_rows.saturating_mul(right_rows) > AUTO_HASH_MIN_PAIRS {
Strategy::Hash
} else {
Strategy::NestedLoop
}
}
}
}
pub struct HashSide {
buckets: HashMap<Vec<HKey>, Vec<usize>>,
pub null_keyed: Vec<usize>,
}
impl HashSide {
pub fn build(
n: usize,
mut key_of: impl FnMut(usize) -> Result<Option<Vec<HKey>>>,
) -> Result<Self> {
let mut buckets: HashMap<Vec<HKey>, Vec<usize>> = HashMap::new();
let mut null_keyed = vec![];
for i in 0..n {
match key_of(i)? {
Some(k) => buckets.entry(k).or_default().push(i),
None => null_keyed.push(i),
}
}
Ok(Self { buckets, null_keyed })
}
pub fn probe(&self, key: &[HKey]) -> &[usize] {
self.buckets.get(key).map(|v| v.as_slice()).unwrap_or(&[])
}
pub fn distinct_keys(&self) -> usize {
self.buckets.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sqlselect::parse;
use serde_json::json;
fn corpus() -> Vec<Value> {
vec![
Value::Null,
json!(0),
json!(-0.0),
json!(0.0),
json!(1),
json!(1.0),
json!(-1),
json!(1000),
json!(0.1),
json!(9007199254740993i64),
json!(9007199254740992i64),
json!("0"),
json!("1"),
json!("1.0"),
json!("1.00"),
json!("01"),
json!("1e3"),
json!(" 1"),
json!("1abc"),
json!(""),
json!("t"),
json!("f"),
json!("true"),
json!("nan"),
json!("inf"),
json!("-0"),
json!("abc"),
json!("ABC"),
json!(true),
json!(false),
json!([1, 2]),
json!("[1,2]"),
json!({"a": 1}),
json!(r#"{"a":1}"#),
]
}
fn equals(a: &Value, b: &Value) -> bool {
let sel = parse("SELECT l.v = r.v AS eq FROM l JOIN r ON 1 = 1").expect("parses");
let (la, lb) = (a.clone(), b.clone());
let resolve = move |t: &str| -> Result<Option<Box<dyn crate::sqlselect::Relation>>> {
Ok(Some(crate::sqlselect::from_vec(match t {
"l" => vec![json!({"v": la})],
_ => vec![json!({"v": lb})],
})))
};
let (_, rows) = crate::sqlselect::execute(&sel, &resolve).expect("runs");
rows.first().and_then(|r| r.get("eq")).and_then(|v| v.as_bool()) == Some(true)
}
#[test]
fn equality_implies_same_bucket() {
let c = corpus();
let mut equal_pairs = 0;
for a in &c {
for b in &c {
if !equals(a, b) {
continue;
}
equal_pairs += 1;
let (ka, kb) = (hkey(a), hkey(b));
assert!(
ka.is_some() && kb.is_some(),
"{a:?} = {b:?} is TRUE but a key is unhashable"
);
assert_eq!(
ka, kb,
"{a:?} = {b:?} is TRUE but they bucket apart — the hash \
join would LOSE this match"
);
}
}
assert!(equal_pairs > 40, "corpus proved too little: {equal_pairs} equal pairs");
}
#[test]
fn null_never_hashes() {
assert_eq!(hkey(&Value::Null), None);
for v in corpus() {
assert!(!equals(&Value::Null, &v));
assert!(!equals(&v, &Value::Null));
}
}
#[test]
fn the_non_transitive_case_is_real_and_survives() {
assert!(equals(&json!(1), &json!("1")));
assert!(equals(&json!(1), &json!("1.0")));
assert!(!equals(&json!("1"), &json!("1.0")));
assert_eq!(hkey(&json!(1)), hkey(&json!("1")));
assert_eq!(hkey(&json!(1)), hkey(&json!("1.0")));
assert_eq!(hkey(&json!("1")), hkey(&json!("1.0")));
}
#[test]
fn signed_zero_shares_a_bucket() {
assert_eq!(hkey(&json!(0.0)), hkey(&json!(-0.0)));
assert_eq!(hkey(&json!(0)), hkey(&json!(-0.0)));
}
#[test]
fn bool_and_its_text_share_a_bucket() {
assert!(equals(&json!(true), &json!("t")));
assert_eq!(hkey(&json!(true)), hkey(&json!("t")));
assert_eq!(hkey(&json!(false)), hkey(&json!("f")));
}
#[test]
fn composite_and_its_json_text_share_a_bucket() {
assert_eq!(hkey(&json!([1, 2])), hkey(&json!("[1,2]")));
}
fn keys_for(sql: &str) -> Vec<(Expr, Expr)> {
let s = parse(sql).expect("parses");
let left = vec![s.from.as_ref().unwrap().binding()];
let j = &s.joins[0];
hash_keys(j.on.as_ref(), &left, &j.table.binding())
}
#[test]
fn simple_equijoin_yields_one_key() {
assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.y").len(), 1);
}
#[test]
fn key_pairs_are_normalised_left_then_right() {
let k = keys_for("SELECT 1 FROM a JOIN b ON b.y = a.x");
assert_eq!(k.len(), 1);
assert_eq!(k[0].0, Expr::Column { qual: Some("a".into()), name: "x".into() });
assert_eq!(k[0].1, Expr::Column { qual: Some("b".into()), name: "y".into() });
}
#[test]
fn multiple_equality_conjuncts_all_become_keys() {
assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x AND a.y = b.y").len(), 2);
}
#[test]
fn non_equality_conjuncts_are_left_to_the_evaluator() {
assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x AND a.n > b.n").len(), 1);
}
#[test]
fn or_is_never_split() {
assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x OR a.y = b.y").is_empty());
}
#[test]
fn a_constant_side_is_not_a_key() {
assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = 5").is_empty());
assert!(keys_for("SELECT 1 FROM a JOIN b ON 1 = 1").is_empty());
}
#[test]
fn same_side_equality_is_not_a_key() {
assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = a.y").is_empty());
}
#[test]
fn a_bare_column_is_refused() {
assert!(keys_for("SELECT 1 FROM a JOIN b ON x = b.y").is_empty());
assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = y").is_empty());
}
#[test]
fn an_expression_key_is_allowed_when_it_reads_one_side() {
assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON lower(a.x) = lower(b.y)").len(), 1);
assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.y::text").len(), 1);
}
#[test]
fn a_key_spanning_both_sides_is_refused() {
assert!(keys_for("SELECT 1 FROM a JOIN b ON coalesce(a.x, b.y) = b.z").is_empty());
}
#[test]
fn an_unknown_function_is_refused() {
let s = parse("SELECT 1 FROM a JOIN b ON a.x = b.y").expect("parses");
let left = vec!["a".to_string()];
let on = Expr::Binary {
op: "=".into(),
left: Box::new(Expr::Column { qual: Some("a".into()), name: "x".into() }),
right: Box::new(Expr::Func {
name: "random".into(),
args: vec![Expr::Column { qual: Some("b".into()), name: "y".into() }],
}),
};
assert!(hash_keys(Some(&on), &left, &s.joins[0].table.binding()).is_empty());
}
#[test]
fn cross_join_has_no_keys() {
assert!(keys_for("SELECT 1 FROM a CROSS JOIN b").is_empty());
}
#[test]
fn a_second_join_may_key_off_either_earlier_relation() {
let s = parse("SELECT 1 FROM a JOIN b ON a.x = b.x JOIN c ON b.y = c.y").expect("parses");
let left = vec!["a".to_string(), "b".to_string()];
let j = &s.joins[1];
assert_eq!(hash_keys(j.on.as_ref(), &left, &j.table.binding()).len(), 1);
}
#[test]
fn no_keys_forces_the_nested_loop_even_when_hash_is_requested() {
assert_eq!(choose(JoinExec::Hash, 0, 1000, 1000), Strategy::NestedLoop);
}
#[test]
fn auto_stays_on_the_reference_path_for_small_inputs() {
assert_eq!(choose(JoinExec::Auto, 1, 4, 4), Strategy::NestedLoop);
assert_eq!(choose(JoinExec::Auto, 1, 8, 8), Strategy::NestedLoop);
assert_eq!(choose(JoinExec::Auto, 1, 8, 9), Strategy::Hash);
}
#[test]
fn forcing_is_honoured_so_differential_tests_mean_something() {
assert_eq!(choose(JoinExec::NestedLoop, 2, 10_000, 10_000), Strategy::NestedLoop);
assert_eq!(choose(JoinExec::Hash, 2, 1, 1), Strategy::Hash);
}
#[test]
fn build_preserves_ascending_row_order_within_a_bucket() {
let vals = vec![json!("a"), json!("b"), json!("a"), json!("a")];
let side = HashSide::build(vals.len(), |i| Ok(hkey(&vals[i]).map(|k| vec![k])))
.expect("builds");
let k = vec![hkey(&json!("a")).unwrap()];
assert_eq!(side.probe(&k), &[0, 2, 3]);
assert_eq!(side.distinct_keys(), 2);
}
#[test]
fn null_keyed_rows_are_set_aside_not_dropped() {
let vals = vec![json!("a"), Value::Null, json!("b")];
let side = HashSide::build(vals.len(), |i| Ok(hkey(&vals[i]).map(|k| vec![k])))
.expect("builds");
assert_eq!(side.null_keyed, vec![1]);
assert!(side.probe(&[HKey::Text("zzz".into())]).is_empty());
assert_eq!(side.distinct_keys(), 2);
}
#[test]
fn a_compound_key_matches_only_on_every_column() {
let rows = vec![(json!(1), json!("x")), (json!(1), json!("y"))];
let side = HashSide::build(rows.len(), |i| {
Ok(match (hkey(&rows[i].0), hkey(&rows[i].1)) {
(Some(a), Some(b)) => Some(vec![a, b]),
_ => None,
})
})
.expect("builds");
let want = vec![hkey(&json!(1)).unwrap(), hkey(&json!("x")).unwrap()];
assert_eq!(side.probe(&want), &[0]);
}
}