Skip to main content

gluesql_core/plan/statement/query/join/
hash.rs

1use {
2    super::{InnerJoinPlan, LeftOuterJoinPlan},
3    crate::plan::{ExprPlan, SourcePlan},
4    serde::{Deserialize, Serialize},
5};
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
8pub enum HashJoinInputPlan {
9    Source(SourcePlan),
10    InnerJoin(Box<InnerJoinPlan>),
11    LeftOuterJoin(Box<LeftOuterJoinPlan>),
12}
13
14#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
15pub struct HashJoinPlan {
16    pub input: HashJoinInputPlan,
17    pub right: SourcePlan,
18    pub input_key: ExprPlan,
19    pub right_key: ExprPlan,
20    pub right_filter: Option<ExprPlan>,
21}
22
23impl HashJoinPlan {
24    pub(crate) fn base_source(&self) -> &SourcePlan {
25        match &self.input {
26            HashJoinInputPlan::Source(source) => source,
27            HashJoinInputPlan::InnerJoin(join) => join.base_source(),
28            HashJoinInputPlan::LeftOuterJoin(join) => join.base_source(),
29        }
30    }
31
32    pub(crate) fn base_source_mut(&mut self) -> &mut SourcePlan {
33        match &mut self.input {
34            HashJoinInputPlan::Source(source) => source,
35            HashJoinInputPlan::InnerJoin(join) => join.base_source_mut(),
36            HashJoinInputPlan::LeftOuterJoin(join) => join.base_source_mut(),
37        }
38    }
39
40    pub(crate) fn joined_sources(&self) -> Vec<&SourcePlan> {
41        let mut sources = match &self.input {
42            HashJoinInputPlan::Source(_) => Vec::new(),
43            HashJoinInputPlan::InnerJoin(join) => join.joined_sources(),
44            HashJoinInputPlan::LeftOuterJoin(join) => join.joined_sources(),
45        };
46        sources.push(&self.right);
47
48        sources
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use {
55        super::{HashJoinInputPlan, HashJoinPlan},
56        crate::{
57            data::Value,
58            plan::{
59                ExprPlan, InnerJoinInputPlan, InnerJoinPlan, LeftOuterJoinInputPlan,
60                LeftOuterJoinPlan, NestedLoopJoinInputPlan, NestedLoopJoinPlan, SourcePlan,
61                TableAccessPlan, TableSourcePlan,
62            },
63        },
64        pretty_assertions::assert_eq,
65    };
66
67    fn table(name: &str) -> SourcePlan {
68        SourcePlan::Table(TableSourcePlan {
69            name: name.to_owned(),
70            alias: None,
71            access: TableAccessPlan::FullScan,
72        })
73    }
74
75    fn expr() -> ExprPlan {
76        ExprPlan::Value(Value::Bool(true))
77    }
78
79    fn inner_join() -> InnerJoinPlan {
80        InnerJoinPlan {
81            input: InnerJoinInputPlan::NestedLoop(NestedLoopJoinPlan {
82                input: NestedLoopJoinInputPlan::Source(table("A")),
83                right: table("B"),
84            }),
85        }
86    }
87
88    fn left_outer_join() -> LeftOuterJoinPlan {
89        LeftOuterJoinPlan {
90            input: LeftOuterJoinInputPlan::NestedLoop(NestedLoopJoinPlan {
91                input: NestedLoopJoinInputPlan::Source(table("A")),
92                right: table("B"),
93            }),
94        }
95    }
96
97    #[test]
98    fn accepts_each_input() {
99        let mut actual = HashJoinPlan {
100            input: HashJoinInputPlan::Source(table("A")),
101            right: table("B"),
102            input_key: expr(),
103            right_key: expr(),
104            right_filter: None,
105        };
106        let expected = HashJoinInputPlan::Source(table("A"));
107        assert_eq!(actual.input, expected);
108        assert_eq!(actual.base_source(), &table("A"));
109        let expected = [table("B")];
110        assert_eq!(actual.joined_sources(), expected.iter().collect::<Vec<_>>());
111        *actual.base_source_mut() = table("source");
112        assert_eq!(actual.base_source(), &table("source"));
113
114        let mut actual = HashJoinPlan {
115            input: HashJoinInputPlan::InnerJoin(Box::new(inner_join())),
116            right: table("C"),
117            input_key: expr(),
118            right_key: expr(),
119            right_filter: None,
120        };
121        let expected = HashJoinInputPlan::InnerJoin(Box::new(inner_join()));
122        assert_eq!(actual.input, expected);
123        assert_eq!(actual.base_source(), &table("A"));
124        let expected = [table("B"), table("C")];
125        assert_eq!(actual.joined_sources(), expected.iter().collect::<Vec<_>>());
126        *actual.base_source_mut() = table("inner");
127        assert_eq!(actual.base_source(), &table("inner"));
128
129        let mut actual = HashJoinPlan {
130            input: HashJoinInputPlan::LeftOuterJoin(Box::new(left_outer_join())),
131            right: table("C"),
132            input_key: expr(),
133            right_key: expr(),
134            right_filter: None,
135        };
136        let expected = HashJoinInputPlan::LeftOuterJoin(Box::new(left_outer_join()));
137        assert_eq!(actual.input, expected);
138        assert_eq!(actual.base_source(), &table("A"));
139        let expected = [table("B"), table("C")];
140        assert_eq!(actual.joined_sources(), expected.iter().collect::<Vec<_>>());
141        *actual.base_source_mut() = table("left-outer");
142        assert_eq!(actual.base_source(), &table("left-outer"));
143    }
144}