initial impl of lazy evaluation

ref #11

Author: nevi-me

Cargo.toml

@@ -14,4 +14,7 @@ byteorder = "1"
 flatbuffers = "0.5"
 array_tool = "1"
 postgres = {version = "0.16.0-rc.1", features = ["default", "with-chrono-0_4", "with-uuid-0_7"]}
-chrono = "0.4"
+chrono = "0.4"
+# for lazy evaluation
+serde = { version = "1.0", features = ["derive"] }
+serde_json = { version = "1.0", features = ["preserve_order"] }

src/evaluation.rs

@@ -0,0 +1,95 @@
+//!
+//! Lazy evaluation for DataFrames
+//!
+//! This is an experimental initial implementation
+
+use crate::dataframe::DataFrame;
+use crate::expression::*;
+use crate::functions::scalar::ScalarFunctions as Scalar;
+use crate::table;
+use ::std::sync::Arc;
+use arrow::array::*;
+use arrow::datatypes::DataType;
+
+pub trait Evaluate {
+    /// Evaluate an operation against a data source
+    fn evaluate(self, operation: &Operation) -> Self;
+}
+
+impl Evaluate for DataFrame {
+    fn evaluate(self, operation: &Operation) -> Self {
+        // get the input columns from the dataframe
+        let columns: Vec<&table::Column> = operation
+            .inputs
+            .clone()
+            .into_iter()
+            .map(|col: Column| self.column_by_name(&col.name))
+            .collect();
+        match &operation.expression {
+            Expression::Scalar(expr) => match expr {
+                ScalarExpression::Add => {
+                    // we are adding 2 columns together to create a third
+                    let column: Vec<ArrayRef> =
+                        if let ColumnType::Scalar(dtype) = &operation.output.column_type {
+                            match dtype {
+                                DataType::Int16 => Scalar::add(
+                                    table::column_to_arrays_i16(columns.get(0).unwrap()),
+                                    table::column_to_arrays_i16(columns.get(1).unwrap()),
+                                )
+                                .unwrap()
+                                .into_iter()
+                                .map(|arr| Arc::new(arr) as ArrayRef)
+                                .collect(),
+                                DataType::Float64 => Scalar::add(
+                                    table::column_to_arrays_f64(columns.get(0).unwrap()),
+                                    table::column_to_arrays_f64(columns.get(1).unwrap()),
+                                )
+                                .unwrap()
+                                .into_iter()
+                                .map(|arr| Arc::new(arr) as ArrayRef)
+                                .collect(),
+                                _ => panic!("Unsupported operation"),
+                            }
+                        } else {
+                            unreachable!()
+                        };
+                    self.with_column(
+                        &operation.output.name,
+                        table::Column::from_arrays(column, operation.output.clone().into()),
+                    )
+                }
+                _ => panic!("Scalar Expression {:?} not supported", expr),
+            },
+            // expr @ _ => panic!("Expression {:?} not supported", expr),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use crate::operation::{AddOperation, ScalarOperation};
+
+    #[test]
+    fn test_evaluation() {
+        let dataframe = DataFrame::from_csv("./test/data/uk_cities_with_headers.csv", None);
+        let a = Column {
+            name: "lat".to_owned(),
+            column_type: ColumnType::Scalar(DataType::Float64),
+        };
+        let b = Column {
+            name: "lng".to_owned(),
+            column_type: ColumnType::Scalar(DataType::Float64),
+        };
+
+        let add = AddOperation::transform(vec![a, b], Some("lat_lng".to_owned())).unwrap();
+
+        let out_dataframe = dataframe.evaluate(&add);
+
+        assert_eq!(
+            "Schema { fields: [Field { name: \"city\", data_type: Utf8, nullable: false }, Field { name: \"lat\", data_type: Float64, nullable: false }, Field { name: \"lng\", data_type: Float64, nullable: false }, Field { name: \"lat_lng\", data_type: Float64, nullable: true }] }",
+            format!("{:?}", out_dataframe.schema())
+        );
+    }
+}

src/expression.rs

@@ -0,0 +1,113 @@
+//! Expressions that generate operations and computations
+
+use arrow::datatypes::DataType;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum ColumnType {
+    Array(DataType),
+    Scalar(DataType),
+}
+
+impl From<DataType> for ColumnType {
+    fn from(dtype: DataType) -> Self {
+        match dtype {
+            DataType::Struct(_) => panic!("struct array conversion not yet supported"),
+            DataType::List(inner) => ColumnType::Array(*inner.clone()),
+            _ => ColumnType::Scalar(dtype.clone()),
+        }
+    }
+}
+
+impl From<ColumnType> for DataType {
+    fn from(from: ColumnType) -> Self {
+        match from {
+            ColumnType::Array(dtype) => DataType::List(Box::new(dtype)),
+            ColumnType::Scalar(dtype) => dtype.clone(),
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Column {
+    pub(crate) name: String,
+    pub(crate) column_type: ColumnType,
+}
+
+impl From<arrow::datatypes::Field> for Column {
+    fn from(field: arrow::datatypes::Field) -> Self {
+        Column {
+            name: field.name().clone(),
+            column_type: field.data_type().clone().into(),
+        }
+    }
+}
+
+impl From<Column> for arrow::datatypes::Field {
+    fn from(column: Column) -> Self {
+        arrow::datatypes::Field::new(column.name.as_str(), column.column_type.into(), true)
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Frame {
+    pub(crate) name: String,
+    pub(crate) columns: Vec<Column>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum Transformation {
+    Aggregate,
+    Calculate(Operation),
+    Join,
+    Group,
+    // can add other transform types here
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Operation {
+    pub(crate) name: String,
+    pub(crate) inputs: Vec<Column>,
+    pub(crate) output: Column,
+    pub(crate) expression: Expression,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Computation {
+    pub(crate) input: Vec<Frame>,
+    pub(crate) operations: Vec<Operation>,
+    pub(crate) output: Frame,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum Expression {
+    Scalar(ScalarExpression),
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum ScalarExpression {
+    Add,
+    Subtract,
+    Divide,
+    Multiply,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn debug() {
+        let frame = Frame {
+            name: "Input Table 1".to_owned(),
+            columns: vec![Column {
+                name: "id".to_owned(),
+                column_type: ColumnType::Scalar(DataType::Int64),
+            }],
+        };
+
+        assert_eq!("Frame { name: \"Input Table 1\", columns: [Column { name: \"id\", column_type: Scalar(Int64) }] }", format!("{:?}", frame));
+        let as_json = serde_json::to_string(&frame).unwrap();
+        assert_eq!("{\"name\":\"Input Table 1\",\"columns\":[{\"name\":\"id\",\"column_type\":{\"Scalar\":\"Int64\"}}]}", as_json);
+    }
+}

src/lib.rs

@@ -6,7 +6,10 @@
 
 pub mod dataframe;
 pub mod error;
+pub mod evaluation;
+pub mod expression;
 pub mod functions;
 pub mod io;
+pub mod operation;
 pub mod table;
 pub mod utils;

src/operation.rs

@@ -0,0 +1,86 @@
+//! Operations
+
+use crate::evaluation::*;
+use crate::expression::*;
+use arrow::datatypes::DataType;
+use arrow::error::ArrowError;
+
+pub trait ScalarOperation {
+    fn name() -> &'static str;
+    fn transform(inputs: Vec<Column>, name: Option<String>) -> Result<Operation, ArrowError>;
+}
+
+/// Operation to add two numeric columns together
+pub struct AddOperation;
+
+impl ScalarOperation for AddOperation {
+    fn name() -> &'static str {
+        "add"
+    }
+
+    fn transform(inputs: Vec<Column>, name: Option<String>) -> Result<Operation, ArrowError> {
+        // add n columns together provided that they are of the same data type
+        // for now we support 2 inputs at a time
+        if inputs.len() != 2 {
+            Err(ArrowError::ComputeError(
+                "Add operation expects 2 inputs".to_string(),
+            ))
+        } else {
+            let a = &inputs[0];
+            let b = &inputs[1];
+            match (&a.column_type, &b.column_type) {
+                (ColumnType::Array(_), _) | (_, ColumnType::Array(_)) => {
+                    Err(ArrowError::ComputeError(
+                        "Add operation only works on scalar columns".to_string(),
+                    ))
+                }
+                (ColumnType::Scalar(from_type), ColumnType::Scalar(to_type)) => {
+                    if from_type != to_type {
+                        Err(ArrowError::ComputeError(
+                            "Add operation currently only works on the same data types".to_string(),
+                        ))
+                    } else {
+                        Ok(Operation {
+                            name: Self::name().to_string(),
+                            inputs: inputs.clone(),
+                            output: Column {
+                                name: name.unwrap_or(format!(
+                                    "{}({}, {})",
+                                    Self::name(),
+                                    &a.name,
+                                    &b.name
+                                )),
+                                column_type: ColumnType::Scalar(from_type.clone()),
+                            },
+                            expression: Expression::Scalar(ScalarExpression::Add),
+                        })
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn scalar_operations() {
+        let a = Column {
+            name: "a".to_owned(),
+            column_type: ColumnType::Scalar(DataType::Int64),
+        };
+        let b = Column {
+            name: "b".to_owned(),
+            column_type: ColumnType::Scalar(DataType::Int64),
+        };
+
+        let add = AddOperation::transform(vec![a, b], None).unwrap();
+
+        assert_eq!(
+            "Operation { name: \"add\", inputs: [Column { name: \"a\", column_type: Scalar(Int64) }, Column { name: \"b\", column_type: Scalar(Int64) }], output: Column { name: \"add(a, b)\", column_type: Scalar(Int64) }, expression: Scalar(Add) }",
+            format!("{:?}", add)
+        );
+    }
+}