quill_sql/catalog/
schema.rs

1use super::column::{Column, ColumnRef};
2use crate::catalog::DataType;
3use crate::error::QuillSQLError;
4use crate::error::QuillSQLResult;
5use crate::utils::table_ref::TableReference;
6use std::sync::{Arc, LazyLock};
7
8pub type SchemaRef = Arc<Schema>;
9
10pub static EMPTY_SCHEMA_REF: LazyLock<SchemaRef> = LazyLock::new(|| Arc::new(Schema::empty()));
11pub static INSERT_OUTPUT_SCHEMA_REF: LazyLock<SchemaRef> = LazyLock::new(|| {
12    Arc::new(Schema::new(vec![Column::new(
13        "insert_rows",
14        DataType::Int32,
15        false,
16    )]))
17});
18pub static UPDATE_OUTPUT_SCHEMA_REF: LazyLock<SchemaRef> = LazyLock::new(|| {
19    Arc::new(Schema::new(vec![Column::new(
20        "update_rows",
21        DataType::Int32,
22        false,
23    )]))
24});
25pub static DELETE_OUTPUT_SCHEMA_REF: LazyLock<SchemaRef> = LazyLock::new(|| {
26    Arc::new(Schema::new(vec![Column::new(
27        "delete_rows",
28        DataType::Int32,
29        false,
30    )]))
31});
32
33#[derive(Debug, Clone, Eq, PartialEq)]
34pub struct Schema {
35    pub columns: Vec<ColumnRef>,
36}
37
38impl Schema {
39    pub fn new(columns: Vec<Column>) -> Self {
40        Self::new_with_check(columns.into_iter().map(Arc::new).collect())
41    }
42
43    fn new_with_check(columns: Vec<ColumnRef>) -> Self {
44        for (idx1, col1) in columns.iter().enumerate() {
45            for col2 in columns.iter().skip(idx1 + 1) {
46                match (&col1.relation, &col2.relation) {
47                    (Some(rel1), Some(rel2)) => {
48                        assert!(!(rel1.resolved_eq(rel2) && col1.name == col2.name));
49                    }
50                    (None, None) => {
51                        if col1.name == col2.name {
52                            assert_eq!(col1.relation, col2.relation);
53                        }
54                    }
55                    (Some(_), None) | (None, Some(_)) => {}
56                }
57            }
58        }
59        Self { columns }
60    }
61
62    pub fn empty() -> Self {
63        Self { columns: vec![] }
64    }
65
66    pub fn try_merge(schemas: impl IntoIterator<Item = Self>) -> QuillSQLResult<Self> {
67        let mut columns = Vec::new();
68        for schema in schemas {
69            columns.extend(schema.columns);
70        }
71        Ok(Self::new_with_check(columns))
72    }
73
74    pub fn project(&self, indices: &[usize]) -> QuillSQLResult<Schema> {
75        let new_columns = indices
76            .iter()
77            .map(|i| self.column_with_index(*i))
78            .collect::<QuillSQLResult<Vec<ColumnRef>>>()?;
79        Ok(Schema::new_with_check(new_columns))
80    }
81
82    pub fn column_with_name(
83        &self,
84        relation: Option<&TableReference>,
85        name: &str,
86    ) -> QuillSQLResult<ColumnRef> {
87        let index = self.index_of(relation, name)?;
88        Ok(self.columns[index].clone())
89    }
90
91    pub fn column_with_index(&self, index: usize) -> QuillSQLResult<ColumnRef> {
92        self.columns
93            .get(index)
94            .cloned()
95            .ok_or_else(|| QuillSQLError::Plan(format!("Unable to get column with index {index}")))
96    }
97
98    /// Find the index of the column with the given name.
99    pub fn index_of(&self, relation: Option<&TableReference>, name: &str) -> QuillSQLResult<usize> {
100        let (idx, _) = self
101            .columns
102            .iter()
103            .enumerate()
104            .find(|(_, col)| {
105                let name_matches = col.name.eq_ignore_ascii_case(name);
106                match (relation, &col.relation) {
107                    (Some(rel), Some(col_rel)) => name_matches && rel.resolved_eq(col_rel),
108                    (Some(_), None) => false,
109                    (None, Some(_)) | (None, None) => name_matches,
110                }
111            })
112            .ok_or_else(|| QuillSQLError::Plan(format!("Unable to get column named \"{name}\"")))?;
113        Ok(idx)
114    }
115
116    pub fn column_count(&self) -> usize {
117        self.columns.len()
118    }
119}
quill_sql/catalog/schema.rs

quill_sql/catalog/
schema.rs
