diff --git a/src/abstract_tree/select.rs b/src/abstract_tree/select.rs index ebe4377..bb96abc 100644 --- a/src/abstract_tree/select.rs +++ b/src/abstract_tree/select.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use crate::{AbstractTree, Expression, Identifier, Item, Value}; +use crate::{AbstractTree, Expression, Identifier, Item, Result, Table, Value, VariableMap}; #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, PartialOrd, Ord)] pub struct Select { @@ -52,9 +52,9 @@ impl AbstractTree for Select { }) } - fn run(&self, source: &str, context: &mut crate::VariableMap) -> crate::Result { + fn run(&self, source: &str, context: &mut VariableMap) -> Result { let value = self.expression.run(source, context)?; - let table = value.as_table()?; + let old_table = value.as_table()?; let column_names = if self.identifiers.len() > 0 { self.identifiers .iter() @@ -62,9 +62,51 @@ impl AbstractTree for Select { .map(|identifierier| identifierier.take_inner()) .collect() } else { - table.column_names().clone() + old_table.headers().clone() }; - let new_table = table.select(&column_names); + + let mut new_table = Table::new(column_names.to_vec()); + + for row in old_table.rows() { + let mut new_row = Vec::new(); + let mut row_context = VariableMap::new(); + + for (i, value) in row.iter().enumerate() { + let column_name = old_table.headers().get(i).unwrap(); + + row_context.set_value(column_name.clone(), value.clone())?; + + let new_table_column_index = + new_table + .headers() + .iter() + .enumerate() + .find_map(|(index, new_column_name)| { + if new_column_name == column_name { + Some(index) + } else { + None + } + }); + + if let Some(index) = new_table_column_index { + while new_row.len() < index + 1 { + new_row.push(Value::Empty); + } + new_row[index] = value.clone(); + } + } + + if let Some(where_clause) = &self.item { + let should_include = where_clause.run(source, &mut row_context)?.as_boolean()?; + + if should_include { + new_table.insert(new_row)?; + } + } else { + new_table.insert(new_row)?; + } + } Ok(Value::Table(new_table)) } diff --git a/src/value/table.rs b/src/value/table.rs index 331d887..2cafaee 100644 --- a/src/value/table.rs +++ b/src/value/table.rs @@ -1,4 +1,4 @@ -use crate::{Error, Result, Value, VariableMap}; +use crate::{Error, Item, Result, Value, VariableMap}; use comfy_table::{Cell, Color, ContentArrangement, Table as ComfyTable}; use serde::{Deserialize, Serialize}; use std::{ @@ -42,7 +42,7 @@ impl Table { self.rows.reserve(additional); } - pub fn column_names(&self) -> &Vec { + pub fn headers(&self) -> &Vec { &self.headers } @@ -86,46 +86,11 @@ impl Table { } pub fn get(&self, value: &Value) -> Option<&Vec> { - let primary_key = self.column_names().get(self.primary_key_index)?; + let primary_key = self.headers().get(self.primary_key_index)?; self.get_where(primary_key, value) } - pub fn select(&self, column_names: &[String]) -> Table { - let mut new_table = Table::new(column_names.to_vec()); - - for row in &self.rows { - let mut new_row = Vec::new(); - - for (i, value) in row.iter().enumerate() { - let column_name = self.headers.get(i).unwrap(); - let new_table_column_index = - new_table - .headers - .iter() - .enumerate() - .find_map(|(index, new_column_name)| { - if new_column_name == column_name { - Some(index) - } else { - None - } - }); - - if let Some(index) = new_table_column_index { - while new_row.len() < index + 1 { - new_row.push(Value::Empty); - } - new_row[index] = value.clone(); - } - } - - new_table.insert(new_row).unwrap(); - } - - new_table - } - pub fn get_where(&self, column_name: &str, expected: &Value) -> Option<&Vec> { let column_index = self.get_column_index(column_name)?;