1use crate::sqlselect::{Expr, JoinKind};
58use anyhow::Result;
59use serde_json::Value;
60use std::collections::HashMap;
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum JoinExec {
70 Auto,
71 NestedLoop,
72 Hash,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct JoinChoice {
78 pub kind: JoinKind,
79 pub table: String,
80 pub strategy: Strategy,
81 pub keys: usize,
84 pub left_rows: usize,
85 pub right_rows: usize,
86 pub out_rows: usize,
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum Strategy {
91 NestedLoop,
92 Hash,
93}
94
95impl std::fmt::Display for Strategy {
96 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97 f.write_str(match self {
98 Strategy::NestedLoop => "Nested Loop",
99 Strategy::Hash => "Hash Join",
100 })
101 }
102}
103
104pub const AUTO_HASH_MIN_PAIRS: usize = 64;
113
114#[derive(Debug, Clone, PartialEq, Eq, Hash)]
120pub enum HKey {
121 Num(u64),
123 Text(String),
125}
126
127fn canon(f: f64) -> u64 {
130 let f = if f == 0.0 { 0.0 } else { f };
131 f.to_bits()
132}
133
134pub fn hkey(v: &Value) -> Option<HKey> {
152 match v {
153 Value::Null => None,
154 Value::Number(n) => Some(match n.as_f64() {
155 Some(f) => HKey::Num(canon(f)),
156 None => HKey::Text(n.to_string()),
159 }),
160 Value::String(s) => match s.parse::<f64>() {
161 Ok(f) => Some(HKey::Num(canon(f))),
162 Err(_) => Some(HKey::Text(s.clone())),
163 },
164 Value::Bool(b) => Some(HKey::Text(if *b { "t" } else { "f" }.to_string())),
167 other => Some(HKey::Text(other.to_string())),
171 }
172}
173
174const PURE_FUNCS: &[&str] = &[
189 "lower",
190 "upper",
191 "length",
192 "char_length",
193 "character_length",
194 "coalesce",
195 "nullif",
196 "int2",
197 "int4",
198 "int8",
199 "text",
200 "quote_ident",
201 "format_type",
202 "array_to_string",
203 "current_schema",
204 "current_database",
205 "current_catalog",
206 "current_user",
207 "session_user",
208 "user",
209 "version",
210 "pg_get_userbyid",
211 "pg_table_is_visible",
212 "pg_type_is_visible",
213 "pg_function_is_visible",
214 "pg_encoding_to_char",
215 "pg_get_expr",
216 "pg_get_indexdef",
217 "pg_get_constraintdef",
218];
219
220#[derive(Debug, Clone, Copy, PartialEq, Eq)]
222enum Side {
223 Left,
225 Right,
227 Const,
230 Unusable,
234}
235
236fn side_of(e: &Expr, left: &[String], right: &str) -> Side {
237 let mut saw_left = false;
238 let mut saw_right = false;
239 let mut usable = true;
240 walk(e, left, right, &mut saw_left, &mut saw_right, &mut usable);
241 if !usable || (saw_left && saw_right) {
242 return Side::Unusable;
243 }
244 match (saw_left, saw_right) {
245 (true, false) => Side::Left,
246 (false, true) => Side::Right,
247 (false, false) => Side::Const,
248 (true, true) => unreachable!("handled above"),
249 }
250}
251
252fn walk(
253 e: &Expr,
254 left: &[String],
255 right: &str,
256 saw_left: &mut bool,
257 saw_right: &mut bool,
258 usable: &mut bool,
259) {
260 match e {
261 Expr::Column { qual, .. } => match qual {
262 Some(q) => {
263 if q.eq_ignore_ascii_case(right) {
264 *saw_right = true;
265 } else if left.iter().any(|b| b.eq_ignore_ascii_case(q)) {
266 *saw_left = true;
267 } else {
268 *usable = false;
271 }
272 }
273 None => *usable = false,
278 },
279 Expr::Literal(_) => {}
280 Expr::Star | Expr::QualifiedStar(_) => *usable = false,
281 Expr::Func { name, args } => {
282 if !PURE_FUNCS.iter().any(|f| f.eq_ignore_ascii_case(name)) {
283 *usable = false;
284 }
285 for a in args {
286 walk(a, left, right, saw_left, saw_right, usable);
287 }
288 }
289 Expr::Case { operand, whens, else_ } => {
290 if let Some(o) = operand {
291 walk(o, left, right, saw_left, saw_right, usable);
292 }
293 for (w, t) in whens {
294 walk(w, left, right, saw_left, saw_right, usable);
295 walk(t, left, right, saw_left, saw_right, usable);
296 }
297 if let Some(x) = else_ {
298 walk(x, left, right, saw_left, saw_right, usable);
299 }
300 }
301 Expr::Binary { left: l, right: r, .. } => {
302 walk(l, left, right, saw_left, saw_right, usable);
303 walk(r, left, right, saw_left, saw_right, usable);
304 }
305 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
306 walk(expr, left, right, saw_left, saw_right, usable);
307 }
308 Expr::InList { expr, list, .. } => {
309 walk(expr, left, right, saw_left, saw_right, usable);
310 for i in list {
311 walk(i, left, right, saw_left, saw_right, usable);
312 }
313 }
314 Expr::Index { expr, index } => {
315 walk(expr, left, right, saw_left, saw_right, usable);
316 walk(index, left, right, saw_left, saw_right, usable);
317 }
318 Expr::ArrayLit(items) => {
319 for i in items {
320 walk(i, left, right, saw_left, saw_right, usable);
321 }
322 }
323 Expr::Subquery(_) | Expr::Exists { .. } | Expr::ArrayQuery(_) | Expr::InSubquery { .. } => {
326 *usable = false
327 }
328 Expr::Quantified { left: l, right: r, .. } => {
329 walk(l, left, right, saw_left, saw_right, usable);
330 walk(r, left, right, saw_left, saw_right, usable);
331 }
332 }
333}
334
335fn conjuncts<'a>(e: &'a Expr, out: &mut Vec<&'a Expr>) {
341 match e {
342 Expr::Binary { op, left, right } if op == "AND" => {
343 conjuncts(left, out);
344 conjuncts(right, out);
345 }
346 other => out.push(other),
347 }
348}
349
350pub fn hash_keys(on: Option<&Expr>, left: &[String], right: &str) -> Vec<(Expr, Expr)> {
357 let Some(on) = on else { return vec![] };
358 let mut parts = vec![];
359 conjuncts(on, &mut parts);
360 let mut keys = vec![];
361 for p in parts {
362 let Expr::Binary { op, left: l, right: r } = p else { continue };
363 if op != "=" {
366 continue;
367 }
368 match (side_of(l, left, right), side_of(r, left, right)) {
369 (Side::Left, Side::Right) => keys.push(((**l).clone(), (**r).clone())),
370 (Side::Right, Side::Left) => keys.push(((**r).clone(), (**l).clone())),
371 _ => {}
372 }
373 }
374 keys
375}
376
377pub fn choose(exec: JoinExec, keys: usize, left_rows: usize, right_rows: usize) -> Strategy {
379 if keys == 0 {
380 return Strategy::NestedLoop;
383 }
384 match exec {
385 JoinExec::NestedLoop => Strategy::NestedLoop,
386 JoinExec::Hash => Strategy::Hash,
387 JoinExec::Auto => {
388 if left_rows.saturating_mul(right_rows) > AUTO_HASH_MIN_PAIRS {
389 Strategy::Hash
390 } else {
391 Strategy::NestedLoop
392 }
393 }
394 }
395}
396
397pub struct HashSide {
403 buckets: HashMap<Vec<HKey>, Vec<usize>>,
404 pub null_keyed: Vec<usize>,
408}
409
410impl HashSide {
411 pub fn build(
414 n: usize,
415 mut key_of: impl FnMut(usize) -> Result<Option<Vec<HKey>>>,
416 ) -> Result<Self> {
417 let mut buckets: HashMap<Vec<HKey>, Vec<usize>> = HashMap::new();
418 let mut null_keyed = vec![];
419 for i in 0..n {
420 match key_of(i)? {
421 Some(k) => buckets.entry(k).or_default().push(i),
424 None => null_keyed.push(i),
425 }
426 }
427 Ok(Self { buckets, null_keyed })
428 }
429
430 pub fn probe(&self, key: &[HKey]) -> &[usize] {
432 self.buckets.get(key).map(|v| v.as_slice()).unwrap_or(&[])
433 }
434
435 pub fn distinct_keys(&self) -> usize {
436 self.buckets.len()
437 }
438}
439
440#[cfg(test)]
441mod tests {
442 use super::*;
443 use crate::sqlselect::parse;
444 use serde_json::json;
445
446 fn corpus() -> Vec<Value> {
451 vec![
452 Value::Null,
453 json!(0),
454 json!(-0.0),
455 json!(0.0),
456 json!(1),
457 json!(1.0),
458 json!(-1),
459 json!(1000),
460 json!(0.1),
461 json!(9007199254740993i64),
462 json!(9007199254740992i64),
463 json!("0"),
464 json!("1"),
465 json!("1.0"),
466 json!("1.00"),
467 json!("01"),
468 json!("1e3"),
469 json!(" 1"),
470 json!("1abc"),
471 json!(""),
472 json!("t"),
473 json!("f"),
474 json!("true"),
475 json!("nan"),
476 json!("inf"),
477 json!("-0"),
478 json!("abc"),
479 json!("ABC"),
480 json!(true),
481 json!(false),
482 json!([1, 2]),
483 json!("[1,2]"),
484 json!({"a": 1}),
485 json!(r#"{"a":1}"#),
486 ]
487 }
488
489 fn equals(a: &Value, b: &Value) -> bool {
492 let sel = parse("SELECT l.v = r.v AS eq FROM l JOIN r ON 1 = 1").expect("parses");
493 let (la, lb) = (a.clone(), b.clone());
494 let resolve = move |t: &str| -> Result<Option<Box<dyn crate::sqlselect::Relation>>> {
495 Ok(Some(crate::sqlselect::from_vec(match t {
496 "l" => vec![json!({"v": la})],
497 _ => vec![json!({"v": lb})],
498 })))
499 };
500 let (_, rows) = crate::sqlselect::execute(&sel, &resolve).expect("runs");
501 rows.first().and_then(|r| r.get("eq")).and_then(|v| v.as_bool()) == Some(true)
502 }
503
504 #[test]
505 fn equality_implies_same_bucket() {
506 let c = corpus();
507 let mut equal_pairs = 0;
508 for a in &c {
509 for b in &c {
510 if !equals(a, b) {
511 continue;
512 }
513 equal_pairs += 1;
514 let (ka, kb) = (hkey(a), hkey(b));
515 assert!(
516 ka.is_some() && kb.is_some(),
517 "{a:?} = {b:?} is TRUE but a key is unhashable"
518 );
519 assert_eq!(
520 ka, kb,
521 "{a:?} = {b:?} is TRUE but they bucket apart — the hash \
522 join would LOSE this match"
523 );
524 }
525 }
526 assert!(equal_pairs > 40, "corpus proved too little: {equal_pairs} equal pairs");
529 }
530
531 #[test]
532 fn null_never_hashes() {
533 assert_eq!(hkey(&Value::Null), None);
534 for v in corpus() {
536 assert!(!equals(&Value::Null, &v));
537 assert!(!equals(&v, &Value::Null));
538 }
539 }
540
541 #[test]
542 fn the_non_transitive_case_is_real_and_survives() {
543 assert!(equals(&json!(1), &json!("1")));
545 assert!(equals(&json!(1), &json!("1.0")));
546 assert!(!equals(&json!("1"), &json!("1.0")));
547 assert_eq!(hkey(&json!(1)), hkey(&json!("1")));
550 assert_eq!(hkey(&json!(1)), hkey(&json!("1.0")));
551 assert_eq!(hkey(&json!("1")), hkey(&json!("1.0")));
552 }
553
554 #[test]
555 fn signed_zero_shares_a_bucket() {
556 assert_eq!(hkey(&json!(0.0)), hkey(&json!(-0.0)));
557 assert_eq!(hkey(&json!(0)), hkey(&json!(-0.0)));
558 }
559
560 #[test]
561 fn bool_and_its_text_share_a_bucket() {
562 assert!(equals(&json!(true), &json!("t")));
563 assert_eq!(hkey(&json!(true)), hkey(&json!("t")));
564 assert_eq!(hkey(&json!(false)), hkey(&json!("f")));
565 }
566
567 #[test]
568 fn composite_and_its_json_text_share_a_bucket() {
569 assert_eq!(hkey(&json!([1, 2])), hkey(&json!("[1,2]")));
570 }
571
572 fn keys_for(sql: &str) -> Vec<(Expr, Expr)> {
575 let s = parse(sql).expect("parses");
576 let left = vec![s.from.as_ref().unwrap().binding()];
577 let j = &s.joins[0];
578 hash_keys(j.on.as_ref(), &left, &j.table.binding())
579 }
580
581 #[test]
582 fn simple_equijoin_yields_one_key() {
583 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.y").len(), 1);
584 }
585
586 #[test]
587 fn key_pairs_are_normalised_left_then_right() {
588 let k = keys_for("SELECT 1 FROM a JOIN b ON b.y = a.x");
591 assert_eq!(k.len(), 1);
592 assert_eq!(k[0].0, Expr::Column { qual: Some("a".into()), name: "x".into() });
593 assert_eq!(k[0].1, Expr::Column { qual: Some("b".into()), name: "y".into() });
594 }
595
596 #[test]
597 fn multiple_equality_conjuncts_all_become_keys() {
598 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x AND a.y = b.y").len(), 2);
599 }
600
601 #[test]
602 fn non_equality_conjuncts_are_left_to_the_evaluator() {
603 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x AND a.n > b.n").len(), 1);
606 }
607
608 #[test]
609 fn or_is_never_split() {
610 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x OR a.y = b.y").is_empty());
611 }
612
613 #[test]
614 fn a_constant_side_is_not_a_key() {
615 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = 5").is_empty());
616 assert!(keys_for("SELECT 1 FROM a JOIN b ON 1 = 1").is_empty());
617 }
618
619 #[test]
620 fn same_side_equality_is_not_a_key() {
621 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = a.y").is_empty());
622 }
623
624 #[test]
625 fn a_bare_column_is_refused() {
626 assert!(keys_for("SELECT 1 FROM a JOIN b ON x = b.y").is_empty());
629 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = y").is_empty());
630 }
631
632 #[test]
633 fn an_expression_key_is_allowed_when_it_reads_one_side() {
634 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON lower(a.x) = lower(b.y)").len(), 1);
635 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.y::text").len(), 1);
636 }
637
638 #[test]
639 fn a_key_spanning_both_sides_is_refused() {
640 assert!(keys_for("SELECT 1 FROM a JOIN b ON coalesce(a.x, b.y) = b.z").is_empty());
641 }
642
643 #[test]
644 fn an_unknown_function_is_refused() {
645 let s = parse("SELECT 1 FROM a JOIN b ON a.x = b.y").expect("parses");
647 let left = vec!["a".to_string()];
648 let on = Expr::Binary {
649 op: "=".into(),
650 left: Box::new(Expr::Column { qual: Some("a".into()), name: "x".into() }),
651 right: Box::new(Expr::Func {
652 name: "random".into(),
653 args: vec![Expr::Column { qual: Some("b".into()), name: "y".into() }],
654 }),
655 };
656 assert!(hash_keys(Some(&on), &left, &s.joins[0].table.binding()).is_empty());
657 }
658
659 #[test]
660 fn cross_join_has_no_keys() {
661 assert!(keys_for("SELECT 1 FROM a CROSS JOIN b").is_empty());
662 }
663
664 #[test]
665 fn a_second_join_may_key_off_either_earlier_relation() {
666 let s = parse("SELECT 1 FROM a JOIN b ON a.x = b.x JOIN c ON b.y = c.y").expect("parses");
667 let left = vec!["a".to_string(), "b".to_string()];
668 let j = &s.joins[1];
669 assert_eq!(hash_keys(j.on.as_ref(), &left, &j.table.binding()).len(), 1);
670 }
671
672 #[test]
675 fn no_keys_forces_the_nested_loop_even_when_hash_is_requested() {
676 assert_eq!(choose(JoinExec::Hash, 0, 1000, 1000), Strategy::NestedLoop);
677 }
678
679 #[test]
680 fn auto_stays_on_the_reference_path_for_small_inputs() {
681 assert_eq!(choose(JoinExec::Auto, 1, 4, 4), Strategy::NestedLoop);
682 assert_eq!(choose(JoinExec::Auto, 1, 8, 8), Strategy::NestedLoop);
683 assert_eq!(choose(JoinExec::Auto, 1, 8, 9), Strategy::Hash);
684 }
685
686 #[test]
687 fn forcing_is_honoured_so_differential_tests_mean_something() {
688 assert_eq!(choose(JoinExec::NestedLoop, 2, 10_000, 10_000), Strategy::NestedLoop);
689 assert_eq!(choose(JoinExec::Hash, 2, 1, 1), Strategy::Hash);
690 }
691
692 #[test]
695 fn build_preserves_ascending_row_order_within_a_bucket() {
696 let vals = vec![json!("a"), json!("b"), json!("a"), json!("a")];
697 let side = HashSide::build(vals.len(), |i| Ok(hkey(&vals[i]).map(|k| vec![k])))
698 .expect("builds");
699 let k = vec![hkey(&json!("a")).unwrap()];
700 assert_eq!(side.probe(&k), &[0, 2, 3]);
701 assert_eq!(side.distinct_keys(), 2);
702 }
703
704 #[test]
705 fn null_keyed_rows_are_set_aside_not_dropped() {
706 let vals = vec![json!("a"), Value::Null, json!("b")];
707 let side = HashSide::build(vals.len(), |i| Ok(hkey(&vals[i]).map(|k| vec![k])))
708 .expect("builds");
709 assert_eq!(side.null_keyed, vec![1]);
710 assert!(side.probe(&[HKey::Text("zzz".into())]).is_empty());
711 assert_eq!(side.distinct_keys(), 2);
713 }
714
715 #[test]
716 fn a_compound_key_matches_only_on_every_column() {
717 let rows = vec![(json!(1), json!("x")), (json!(1), json!("y"))];
718 let side = HashSide::build(rows.len(), |i| {
719 Ok(match (hkey(&rows[i].0), hkey(&rows[i].1)) {
720 (Some(a), Some(b)) => Some(vec![a, b]),
721 _ => None,
722 })
723 })
724 .expect("builds");
725 let want = vec![hkey(&json!(1)).unwrap(), hkey(&json!("x")).unwrap()];
726 assert_eq!(side.probe(&want), &[0]);
727 }
728}