otter_sql/
table.rs

1//! Tables and rows.
2
3use sqlparser::ast::{ColumnOption, ColumnOptionDef, DataType};
4
5use crate::{column::Column, identifier::ColumnRef, value::Value, vm::RuntimeError, BoundedString};
6
7pub(super) const TABLE_UNIQUE_KEY_NAME: &str = "__otter_unique_key";
8
9pub(super) const TABLE_TEMPORARY_NAME: &str = "__otter_temporary_table";
10
11#[derive(Debug, Clone)]
12/// A table in a database.
13///
14/// Contains both the metadata and the actual data.
15pub struct Table {
16    name: BoundedString,
17    pub(super) raw_columns: Vec<Column>,
18    /// The table's data.
19    // TODO: provide methods that verify the data while adding
20    pub(super) raw_data: Vec<RawRow>,
21    row_id: u64,
22}
23
24impl Table {
25    pub(super) fn new(name: BoundedString, mut columns: Vec<Column>) -> Self {
26        // every table has a default unique key
27        columns.insert(
28            0,
29            Column::new(
30                TABLE_UNIQUE_KEY_NAME.into(),
31                DataType::UnsignedInt(None),
32                vec![ColumnOptionDef {
33                    name: None,
34                    option: ColumnOption::Unique { is_primary: false },
35                }],
36                true,
37            ),
38        );
39
40        Self {
41            name,
42            raw_columns: columns,
43            raw_data: Vec::new(),
44            row_id: 0,
45        }
46    }
47
48    pub(super) fn new_temp(num: usize) -> Self {
49        Self::new(
50            format!("{}_{}", TABLE_TEMPORARY_NAME, num).as_str().into(),
51            Vec::new(),
52        )
53    }
54
55    pub(super) fn new_from(table: &Self) -> Self {
56        Self {
57            name: table.name,
58            raw_columns: table.raw_columns.clone(),
59            raw_data: Vec::new(),
60            row_id: 0,
61        }
62    }
63
64    /// Add a new row of data to the table.
65    ///
66    /// **Note**: validation is not implemented yet and the return type is subject to change.
67    pub fn new_row(&mut self, mut data: Vec<Value>) -> &mut Self {
68        data.insert(0, Value::Int64(self.row_id as i64));
69        self.raw_data.push(RawRow { raw_data: data });
70        self.row_id += 1;
71        self
72    }
73
74    /// Retrieve a copy of all of the table's non-internal data.
75    pub fn all_data(&self) -> Vec<Row> {
76        self.raw_data
77            .iter()
78            .cloned()
79            .map(|row| Row::from_raw(row, self))
80            .collect()
81    }
82
83    /// The table's name.
84    pub fn name(&self) -> &BoundedString {
85        &self.name
86    }
87
88    /// The table's (non-internal) columns.
89    pub fn columns(&self) -> impl Iterator<Item = &Column> {
90        self.raw_columns.iter().filter(|c| !c.is_internal())
91    }
92
93    /// Number of non-internal columns.
94    pub fn num_columns(&self) -> usize {
95        // TODO: keep track of this count instead of calculating every time
96        self.columns().count()
97    }
98
99    /// Add a new column to the table.
100    ///
101    /// **Note**: this does not yet modify any of the rows. They must be kept consistent
102    /// externally using [`add_column_data`](`Self::add_column_data`).
103    // TODO: does not add the column data to the rows.
104    pub fn add_column(&mut self, column: Column) -> &mut Self {
105        self.raw_columns.push(column);
106        self
107    }
108
109    /// Add data for a new column to all rows.
110    pub fn add_column_data(
111        &mut self,
112        col_name: &BoundedString,
113        data: Vec<Value>,
114    ) -> Result<&mut Self, RuntimeError> {
115        let (col_index, _) = self.get_column(col_name)?;
116
117        if !self.is_empty() && self.raw_data.len() != data.len() {
118            return Err(RuntimeError::TableNewColumnSizeMismatch {
119                table_name: *self.name(),
120                table_len: self.raw_data.len(),
121                col_name: *col_name,
122                col_len: data.len(),
123            });
124        }
125
126        if !self.is_empty() {
127            let first_row_size = self.raw_data[0].raw_data.len();
128            if first_row_size == col_index {
129                // column is the last one. just push it at the end.
130                for (row, new_data) in self.raw_data.iter_mut().zip(data.into_iter()) {
131                    row.raw_data.push(new_data);
132                }
133            } else if first_row_size == self.raw_columns.len() {
134                // column data is already added. we replace it.
135                for (row, new_data) in self.raw_data.iter_mut().zip(data.into_iter()) {
136                    row.raw_data[col_index] = new_data;
137                }
138            } else {
139                // when the column is somewhere in the middle or beginning.
140                // perhaps an expensive operation!
141                for (row, new_data) in self.raw_data.iter_mut().zip(data.into_iter()) {
142                    row.raw_data.insert(col_index, new_data)
143                }
144            }
145        } else {
146            for value in data {
147                self.new_row(vec![value]);
148            }
149        }
150
151        Ok(self)
152    }
153
154    /// Map column name to its index and definition.
155    pub(super) fn get_column(
156        &self,
157        col_name: &BoundedString,
158    ) -> Result<(usize, &Column), RuntimeError> {
159        let idx = self.raw_columns.iter().position(|c| c.name() == col_name);
160        if let Some(idx) = idx {
161            Ok((idx, &self.raw_columns[idx]))
162        } else {
163            return Err(RuntimeError::ColumnNotFound(ColumnRef {
164                schema_name: None,
165                table_name: Some(*self.name()),
166                col_name: *col_name,
167            }));
168        }
169    }
170
171    /// Retrieve all data of a column.
172    pub fn get_column_data(&self, col_name: &BoundedString) -> Result<Vec<Value>, RuntimeError> {
173        let (col_index, _) = self.get_column(col_name)?;
174
175        Ok(self
176            .raw_data
177            .iter()
178            .map(|row| row.raw_data[col_index].clone())
179            .collect())
180    }
181
182    /// Rename the table.
183    pub fn rename(&mut self, new_name: BoundedString) {
184        self.name = new_name;
185    }
186
187    /// Whether the table has no rows.
188    pub fn is_empty(&self) -> bool {
189        self.raw_data.is_empty()
190    }
191
192    /// Whether the table has no defined columns.
193    pub fn has_no_columns(&self) -> bool {
194        self.columns().next().is_none()
195    }
196
197    /// Create a new row filled with sentinel values for the data type.
198    ///
199    /// Note: does not add the row to the table.
200    pub(super) fn sentinel_row(&self) -> Result<Row, RuntimeError> {
201        let data = self
202            .raw_columns
203            .iter()
204            .map(|c| Value::sentinel_value(c.data_type()))
205            .collect::<Result<Vec<_>, _>>()?;
206        Ok(Row { data })
207    }
208}
209
210#[cfg(feature = "terminal-output")]
211#[cfg_attr(docsrs, doc(cfg(feature = "terminal-output")))]
212impl std::fmt::Display for Table {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        use tabled::{builder::Builder, Style};
215
216        let mut builder = Builder::default();
217
218        for row in self.all_data() {
219            builder.add_record(row.data().into_iter().map(|v| v.to_string()));
220        }
221
222        builder.set_columns(self.columns().map(|c| c.name().to_string()));
223
224        let mut table = builder.build();
225
226        table.with(Style::rounded());
227
228        write!(f, "{}", table)
229    }
230}
231
232/// Trait to retrieve data from something that looks like a row in a table.
233pub trait RowLike {
234    /// Copy or move of the data contained in the row.
235    ///
236    /// Whether the operation is a copy or move depends on whether the row is a reference to a row
237    /// in a table or is an actual row in a table.
238    fn data(self) -> Vec<Value>;
239
240    /// References to all values contained in the row.
241    fn data_shared(&self) -> Vec<&Value>;
242}
243
244/// A row stored in a table. Represents a relation in relational algebra terms.
245#[derive(Debug, Clone, PartialEq)]
246pub struct Row {
247    /// Values for each column in the row.
248    data: Vec<Value>,
249}
250
251impl RowLike for Row {
252    fn data(self) -> Vec<Value> {
253        self.data
254    }
255
256    fn data_shared(&self) -> Vec<&Value> {
257        self.data.iter().collect()
258    }
259}
260
261impl Row {
262    pub fn new(data: Vec<Value>) -> Self {
263        Self { data }
264    }
265
266    pub(super) fn from_raw(raw: RawRow, table: &Table) -> Row {
267        Row {
268            data: raw
269                .raw_data
270                .into_iter()
271                .enumerate()
272                .filter_map(|(i, value)| {
273                    if table.raw_columns[i].is_internal() {
274                        None
275                    } else {
276                        Some(value)
277                    }
278                })
279                .collect(),
280        }
281    }
282
283    pub fn to_shared(&self) -> RowShared {
284        RowShared::from_row(self)
285    }
286}
287
288/// A reference to a row in a table. Does not contain internal columns.
289#[derive(Debug, Clone, PartialEq)]
290pub struct RowShared<'a> {
291    data: Vec<&'a Value>,
292}
293
294impl<'a> RowLike for RowShared<'a> {
295    fn data_shared(&self) -> Vec<&Value> {
296        self.data.clone()
297    }
298
299    fn data(self) -> Vec<Value> {
300        self.data.into_iter().cloned().collect()
301    }
302}
303
304impl<'a> RowShared<'a> {
305    pub(super) fn from_raw<'b>(raw: &'a RawRow, table: &'b Table) -> Self {
306        Self {
307            data: raw
308                .raw_data
309                .iter()
310                .enumerate()
311                .filter_map(|(i, value)| {
312                    if table.raw_columns[i].is_internal() {
313                        None
314                    } else {
315                        Some(value)
316                    }
317                })
318                .collect(),
319        }
320    }
321
322    pub fn from_row(row: &'a Row) -> Self {
323        Self {
324            data: row.data_shared(),
325        }
326    }
327}
328
329impl<'a> From<&'a Row> for RowShared<'a> {
330    fn from(row: &'a Row) -> Self {
331        Self::from_row(row)
332    }
333}
334
335/// A row in a table, including internal columns.
336#[derive(Debug, Clone, PartialEq)]
337pub(super) struct RawRow {
338    /// Values for each column in the row.
339    pub(crate) raw_data: Vec<Value>,
340}
341
342impl RowLike for RawRow {
343    fn data(self) -> Vec<Value> {
344        self.raw_data
345    }
346
347    fn data_shared(&self) -> Vec<&Value> {
348        self.raw_data.iter().collect()
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use sqlparser::ast::DataType;
355
356    use super::Table;
357    use crate::{column::Column, table::Row, value::Value};
358
359    #[test]
360    fn create_table() {
361        let mut table = Table::new("test".into(), vec![]);
362        assert_eq!(table.name(), "test");
363
364        table.add_column(Column::new(
365            "col1".into(),
366            DataType::Int(None),
367            vec![],
368            false,
369        ));
370
371        table.add_column(Column::new(
372            "col2".into(),
373            DataType::Int(None),
374            vec![],
375            true,
376        ));
377
378        assert_eq!(table.columns().collect::<Vec<_>>().len(), 1);
379        assert_eq!(table.columns().next().unwrap().name(), "col1");
380
381        table.new_row(vec![Value::Int64(1), Value::Int64(2)]);
382
383        assert_eq!(
384            table.all_data(),
385            vec![Row {
386                data: vec![Value::Int64(1)],
387            }]
388        );
389    }
390}
otter_sql/table.rs

otter_sql/
table.rs
