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 }
315}
316
317fn conjuncts<'a>(e: &'a Expr, out: &mut Vec<&'a Expr>) {
323 match e {
324 Expr::Binary { op, left, right } if op == "AND" => {
325 conjuncts(left, out);
326 conjuncts(right, out);
327 }
328 other => out.push(other),
329 }
330}
331
332pub fn hash_keys(on: Option<&Expr>, left: &[String], right: &str) -> Vec<(Expr, Expr)> {
339 let Some(on) = on else { return vec![] };
340 let mut parts = vec![];
341 conjuncts(on, &mut parts);
342 let mut keys = vec![];
343 for p in parts {
344 let Expr::Binary { op, left: l, right: r } = p else { continue };
345 if op != "=" {
348 continue;
349 }
350 match (side_of(l, left, right), side_of(r, left, right)) {
351 (Side::Left, Side::Right) => keys.push(((**l).clone(), (**r).clone())),
352 (Side::Right, Side::Left) => keys.push(((**r).clone(), (**l).clone())),
353 _ => {}
354 }
355 }
356 keys
357}
358
359pub fn choose(exec: JoinExec, keys: usize, left_rows: usize, right_rows: usize) -> Strategy {
361 if keys == 0 {
362 return Strategy::NestedLoop;
365 }
366 match exec {
367 JoinExec::NestedLoop => Strategy::NestedLoop,
368 JoinExec::Hash => Strategy::Hash,
369 JoinExec::Auto => {
370 if left_rows.saturating_mul(right_rows) > AUTO_HASH_MIN_PAIRS {
371 Strategy::Hash
372 } else {
373 Strategy::NestedLoop
374 }
375 }
376 }
377}
378
379pub struct HashSide {
385 buckets: HashMap<Vec<HKey>, Vec<usize>>,
386 pub null_keyed: Vec<usize>,
390}
391
392impl HashSide {
393 pub fn build(
396 n: usize,
397 mut key_of: impl FnMut(usize) -> Result<Option<Vec<HKey>>>,
398 ) -> Result<Self> {
399 let mut buckets: HashMap<Vec<HKey>, Vec<usize>> = HashMap::new();
400 let mut null_keyed = vec![];
401 for i in 0..n {
402 match key_of(i)? {
403 Some(k) => buckets.entry(k).or_default().push(i),
406 None => null_keyed.push(i),
407 }
408 }
409 Ok(Self { buckets, null_keyed })
410 }
411
412 pub fn probe(&self, key: &[HKey]) -> &[usize] {
414 self.buckets.get(key).map(|v| v.as_slice()).unwrap_or(&[])
415 }
416
417 pub fn distinct_keys(&self) -> usize {
418 self.buckets.len()
419 }
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use crate::sqlselect::parse;
426 use serde_json::json;
427
428 fn corpus() -> Vec<Value> {
433 vec![
434 Value::Null,
435 json!(0),
436 json!(-0.0),
437 json!(0.0),
438 json!(1),
439 json!(1.0),
440 json!(-1),
441 json!(1000),
442 json!(0.1),
443 json!(9007199254740993i64),
444 json!(9007199254740992i64),
445 json!("0"),
446 json!("1"),
447 json!("1.0"),
448 json!("1.00"),
449 json!("01"),
450 json!("1e3"),
451 json!(" 1"),
452 json!("1abc"),
453 json!(""),
454 json!("t"),
455 json!("f"),
456 json!("true"),
457 json!("nan"),
458 json!("inf"),
459 json!("-0"),
460 json!("abc"),
461 json!("ABC"),
462 json!(true),
463 json!(false),
464 json!([1, 2]),
465 json!("[1,2]"),
466 json!({"a": 1}),
467 json!(r#"{"a":1}"#),
468 ]
469 }
470
471 fn equals(a: &Value, b: &Value) -> bool {
474 let sel = parse("SELECT l.v = r.v AS eq FROM l JOIN r ON 1 = 1").expect("parses");
475 let (la, lb) = (a.clone(), b.clone());
476 let resolve = move |t: &str| -> Result<Option<Box<dyn crate::sqlselect::Relation>>> {
477 Ok(Some(crate::sqlselect::from_vec(match t {
478 "l" => vec![json!({"v": la})],
479 _ => vec![json!({"v": lb})],
480 })))
481 };
482 let (_, rows) = crate::sqlselect::execute(&sel, &resolve).expect("runs");
483 rows.first().and_then(|r| r.get("eq")).and_then(|v| v.as_bool()) == Some(true)
484 }
485
486 #[test]
487 fn equality_implies_same_bucket() {
488 let c = corpus();
489 let mut equal_pairs = 0;
490 for a in &c {
491 for b in &c {
492 if !equals(a, b) {
493 continue;
494 }
495 equal_pairs += 1;
496 let (ka, kb) = (hkey(a), hkey(b));
497 assert!(
498 ka.is_some() && kb.is_some(),
499 "{a:?} = {b:?} is TRUE but a key is unhashable"
500 );
501 assert_eq!(
502 ka, kb,
503 "{a:?} = {b:?} is TRUE but they bucket apart — the hash \
504 join would LOSE this match"
505 );
506 }
507 }
508 assert!(equal_pairs > 40, "corpus proved too little: {equal_pairs} equal pairs");
511 }
512
513 #[test]
514 fn null_never_hashes() {
515 assert_eq!(hkey(&Value::Null), None);
516 for v in corpus() {
518 assert!(!equals(&Value::Null, &v));
519 assert!(!equals(&v, &Value::Null));
520 }
521 }
522
523 #[test]
524 fn the_non_transitive_case_is_real_and_survives() {
525 assert!(equals(&json!(1), &json!("1")));
527 assert!(equals(&json!(1), &json!("1.0")));
528 assert!(!equals(&json!("1"), &json!("1.0")));
529 assert_eq!(hkey(&json!(1)), hkey(&json!("1")));
532 assert_eq!(hkey(&json!(1)), hkey(&json!("1.0")));
533 assert_eq!(hkey(&json!("1")), hkey(&json!("1.0")));
534 }
535
536 #[test]
537 fn signed_zero_shares_a_bucket() {
538 assert_eq!(hkey(&json!(0.0)), hkey(&json!(-0.0)));
539 assert_eq!(hkey(&json!(0)), hkey(&json!(-0.0)));
540 }
541
542 #[test]
543 fn bool_and_its_text_share_a_bucket() {
544 assert!(equals(&json!(true), &json!("t")));
545 assert_eq!(hkey(&json!(true)), hkey(&json!("t")));
546 assert_eq!(hkey(&json!(false)), hkey(&json!("f")));
547 }
548
549 #[test]
550 fn composite_and_its_json_text_share_a_bucket() {
551 assert_eq!(hkey(&json!([1, 2])), hkey(&json!("[1,2]")));
552 }
553
554 fn keys_for(sql: &str) -> Vec<(Expr, Expr)> {
557 let s = parse(sql).expect("parses");
558 let left = vec![s.from.as_ref().unwrap().binding()];
559 let j = &s.joins[0];
560 hash_keys(j.on.as_ref(), &left, &j.table.binding())
561 }
562
563 #[test]
564 fn simple_equijoin_yields_one_key() {
565 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.y").len(), 1);
566 }
567
568 #[test]
569 fn key_pairs_are_normalised_left_then_right() {
570 let k = keys_for("SELECT 1 FROM a JOIN b ON b.y = a.x");
573 assert_eq!(k.len(), 1);
574 assert_eq!(k[0].0, Expr::Column { qual: Some("a".into()), name: "x".into() });
575 assert_eq!(k[0].1, Expr::Column { qual: Some("b".into()), name: "y".into() });
576 }
577
578 #[test]
579 fn multiple_equality_conjuncts_all_become_keys() {
580 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x AND a.y = b.y").len(), 2);
581 }
582
583 #[test]
584 fn non_equality_conjuncts_are_left_to_the_evaluator() {
585 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x AND a.n > b.n").len(), 1);
588 }
589
590 #[test]
591 fn or_is_never_split() {
592 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.x OR a.y = b.y").is_empty());
593 }
594
595 #[test]
596 fn a_constant_side_is_not_a_key() {
597 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = 5").is_empty());
598 assert!(keys_for("SELECT 1 FROM a JOIN b ON 1 = 1").is_empty());
599 }
600
601 #[test]
602 fn same_side_equality_is_not_a_key() {
603 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = a.y").is_empty());
604 }
605
606 #[test]
607 fn a_bare_column_is_refused() {
608 assert!(keys_for("SELECT 1 FROM a JOIN b ON x = b.y").is_empty());
611 assert!(keys_for("SELECT 1 FROM a JOIN b ON a.x = y").is_empty());
612 }
613
614 #[test]
615 fn an_expression_key_is_allowed_when_it_reads_one_side() {
616 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON lower(a.x) = lower(b.y)").len(), 1);
617 assert_eq!(keys_for("SELECT 1 FROM a JOIN b ON a.x = b.y::text").len(), 1);
618 }
619
620 #[test]
621 fn a_key_spanning_both_sides_is_refused() {
622 assert!(keys_for("SELECT 1 FROM a JOIN b ON coalesce(a.x, b.y) = b.z").is_empty());
623 }
624
625 #[test]
626 fn an_unknown_function_is_refused() {
627 let s = parse("SELECT 1 FROM a JOIN b ON a.x = b.y").expect("parses");
629 let left = vec!["a".to_string()];
630 let on = Expr::Binary {
631 op: "=".into(),
632 left: Box::new(Expr::Column { qual: Some("a".into()), name: "x".into() }),
633 right: Box::new(Expr::Func {
634 name: "random".into(),
635 args: vec![Expr::Column { qual: Some("b".into()), name: "y".into() }],
636 }),
637 };
638 assert!(hash_keys(Some(&on), &left, &s.joins[0].table.binding()).is_empty());
639 }
640
641 #[test]
642 fn cross_join_has_no_keys() {
643 assert!(keys_for("SELECT 1 FROM a CROSS JOIN b").is_empty());
644 }
645
646 #[test]
647 fn a_second_join_may_key_off_either_earlier_relation() {
648 let s = parse("SELECT 1 FROM a JOIN b ON a.x = b.x JOIN c ON b.y = c.y").expect("parses");
649 let left = vec!["a".to_string(), "b".to_string()];
650 let j = &s.joins[1];
651 assert_eq!(hash_keys(j.on.as_ref(), &left, &j.table.binding()).len(), 1);
652 }
653
654 #[test]
657 fn no_keys_forces_the_nested_loop_even_when_hash_is_requested() {
658 assert_eq!(choose(JoinExec::Hash, 0, 1000, 1000), Strategy::NestedLoop);
659 }
660
661 #[test]
662 fn auto_stays_on_the_reference_path_for_small_inputs() {
663 assert_eq!(choose(JoinExec::Auto, 1, 4, 4), Strategy::NestedLoop);
664 assert_eq!(choose(JoinExec::Auto, 1, 8, 8), Strategy::NestedLoop);
665 assert_eq!(choose(JoinExec::Auto, 1, 8, 9), Strategy::Hash);
666 }
667
668 #[test]
669 fn forcing_is_honoured_so_differential_tests_mean_something() {
670 assert_eq!(choose(JoinExec::NestedLoop, 2, 10_000, 10_000), Strategy::NestedLoop);
671 assert_eq!(choose(JoinExec::Hash, 2, 1, 1), Strategy::Hash);
672 }
673
674 #[test]
677 fn build_preserves_ascending_row_order_within_a_bucket() {
678 let vals = vec![json!("a"), json!("b"), json!("a"), json!("a")];
679 let side = HashSide::build(vals.len(), |i| Ok(hkey(&vals[i]).map(|k| vec![k])))
680 .expect("builds");
681 let k = vec![hkey(&json!("a")).unwrap()];
682 assert_eq!(side.probe(&k), &[0, 2, 3]);
683 assert_eq!(side.distinct_keys(), 2);
684 }
685
686 #[test]
687 fn null_keyed_rows_are_set_aside_not_dropped() {
688 let vals = vec![json!("a"), Value::Null, json!("b")];
689 let side = HashSide::build(vals.len(), |i| Ok(hkey(&vals[i]).map(|k| vec![k])))
690 .expect("builds");
691 assert_eq!(side.null_keyed, vec![1]);
692 assert!(side.probe(&[HKey::Text("zzz".into())]).is_empty());
693 assert_eq!(side.distinct_keys(), 2);
695 }
696
697 #[test]
698 fn a_compound_key_matches_only_on_every_column() {
699 let rows = vec![(json!(1), json!("x")), (json!(1), json!("y"))];
700 let side = HashSide::build(rows.len(), |i| {
701 Ok(match (hkey(&rows[i].0), hkey(&rows[i].1)) {
702 (Some(a), Some(b)) => Some(vec![a, b]),
703 _ => None,
704 })
705 })
706 .expect("builds");
707 let want = vec![hkey(&json!(1)).unwrap(), hkey(&json!("x")).unwrap()];
708 assert_eq!(side.probe(&want), &[0]);
709 }
710}