quill_sql/plan/logical_plan/
util.rs

1use crate::catalog::{ColumnRef, Schema};
2use crate::error::QuillSQLResult;
3use crate::expression::{Expr, ExprTrait};
4use crate::plan::logical_plan::JoinType;
5use crate::plan::logical_plan::LogicalPlan;
6use std::sync::Arc;
7
8pub fn build_join_schema(
9    left: &Schema,
10    right: &Schema,
11    join_type: JoinType,
12) -> QuillSQLResult<Schema> {
13    fn nullify_columns(columns: &[ColumnRef]) -> Vec<ColumnRef> {
14        columns
15            .iter()
16            .map(|f| Arc::new(f.as_ref().clone().with_nullable(true)))
17            .collect()
18    }
19
20    let left_cols = &left.columns;
21    let right_cols = &right.columns;
22
23    let columns: Vec<ColumnRef> = match join_type {
24        JoinType::Inner | JoinType::Cross => {
25            left_cols.iter().chain(right_cols.iter()).cloned().collect()
26        }
27        JoinType::LeftOuter => left_cols
28            .iter()
29            .chain(&nullify_columns(right_cols))
30            .cloned()
31            .collect(),
32        JoinType::RightOuter => nullify_columns(left_cols)
33            .iter()
34            .chain(right_cols.iter())
35            .cloned()
36            .collect(),
37        JoinType::FullOuter => nullify_columns(left_cols)
38            .iter()
39            .chain(&nullify_columns(right_cols))
40            .cloned()
41            .collect(),
42    };
43    Ok(Schema { columns })
44}
45
46pub fn project_schema(input: &LogicalPlan, exprs: &[Expr]) -> QuillSQLResult<Schema> {
47    let input_schema = &input.schema();
48    let mut columns = vec![];
49    for expr in exprs {
50        columns.push(expr.to_column(input_schema)?)
51    }
52    Ok(Schema::new(columns))
53}