poc.rs

use std::slice::Iter;
use std::ops::{Add, Mul, Neg};
use std::borrow::Cow;
use std::fmt::Debug;
use std::rc::Rc;

/// Our string type is copy-on-write string refs
type Str<'a> = Cow<'a, str>;

/// Marker trait for the values
trait Value: Clone + Debug {}

impl Value for f64 {}
impl<'a> Value for Str<'a> {}

/// A column is reference counted vector of one of supported
/// types - it erases the type for its consumers
#[derive(Debug, Clone)]
enum Column<'a> {
    F64(Rc<Vec<f64>>),
    STR(Rc<Vec<Str<'a>>>),
}

/// Convert vectors to columns - erases the concrete type
/// implement `std::convert::From` to get `std::convert::Into` for free
impl<'a> From<Vec<f64>> for Column<'a> {
    fn from(vec: Vec<f64>) -> Self {
        Column::F64(Rc::from(vec))
    }
}

impl<'a> From<Vec<Str<'a>>> for Column<'a> {
    fn from(vec: Vec<Str<'a>>) -> Self {
        Column::STR(Rc::from(vec))
    }
}

/// Recovers the elements type through iterators and slices
trait ColumnIter<'a>: Value + 'a {
    fn as_slice<'b>(col: &'b Column<'a>) -> &'b [Self];
    fn iter<'b>(col: &'b Column<'a>) -> Iter<'b, Self> {
        Self::as_slice(col).iter()
    }
}

impl<'a> ColumnIter<'a> for f64 {
    fn as_slice<'b>(col: &'b Column<'a>) -> &'b [f64] {
        if let Column::F64(ref vec) = *col {
            vec
        } else {
            panic!("Improper cast of {:?} to [f64]", col)
        }
    }
}

impl<'a> ColumnIter<'a> for Str<'a> {
    fn as_slice<'b>(col: &'b Column<'a>) -> &'b [Str<'a>] {
        if let Column::STR(ref vec) = *col {
            vec
        } else {
            panic!("Improper cast of {:?} to [Str]", col)
        }
    }
}

/// `ColumnType` is the type of the elements if the columns.
/// It composes all column traits and is used as a type bound
/// to bring all the dependencies at once
trait ColumnType<'a>: Value + 'a {
    fn to_column(Vec<Self>) -> Column<'a>;
    fn iter<'b>(&'b Column<'a>) -> Iter<'b, Self>;
    fn as_slice<'b>(&'b Column<'a>) -> &'b [Self];
}

/// Implement `ColumnType` for each type that implements
/// `ColumnIter<Self>` and `From<Vec<Self>>` for `Column`
impl<'a, T> ColumnType<'a> for T
where
    T: ColumnIter<'a> + Value + 'a,
    Column<'a>: From<Vec<T>>,
{
    fn to_column(vec: Vec<T>) -> Column<'a> {
        vec.into()
    }

    fn iter<'b>(col: &'b Column<'a>) -> Iter<'b, T> {
        T::iter(col)
    }

    fn as_slice<'b>(col: &'b Column<'a>) -> &'b [T] {
        T::as_slice(col)
    }
}

impl<'a> Column<'a> {
    /// Construct a column from a vector
    pub fn from<T: ColumnType<'a>>(vec: Vec<T>) -> Column<'a> {
        T::to_column(vec)
    }

    /// column.iter()
    pub fn iter<T: ColumnType<'a>>(&self) -> Iter<T> {
        T::iter(self)
    }

    /// column.as_slice()
    pub fn as_slice<T: ColumnType<'a>>(&self) -> &[T] {
        T::as_slice(self)
    }
}

/// The `Frame` holds `Column`-s

#[derive(Debug)]
struct Frame<'a> {
    columns: Vec<Column<'a>>,
    rows: usize,
}

/// Expr is just a boxed `&Frame -> Column` function

type Expr<'a> = Box<Fn(&Frame<'a>) -> Column<'a> + 'a>;

/// Select a `Column`
fn col<'a>(col_id: usize) -> Expr<'a> {
    Box::new(move |frame: &Frame<'a>| -> Column<'a> { frame.columns[col_id].clone() })
}

/// Literals
fn val<'a, T: ColumnType<'a>>(value: T) -> Expr<'a> {
    Box::new(move |frame: &Frame| Column::from(vec![value.clone(); frame.rows]))
}

/// Unary operator consturctor
/// takes an expression and a function over the elements of the
/// expression results and creates a new expression
///
/// Expr<Arg> -> (Arg -> Res) -> Expr<Res>
///
fn expr1<'a, Arg, Res, F>(arg: Expr<'a>, f: F) -> Expr<'a>
where
    Arg: ColumnType<'a>,
    Res: ColumnType<'a>,
    F: Fn(&Arg) -> Res + 'a,
{
    Box::new(move |frame: &Frame<'a>| -> Column<'a> {
        Column::from(arg(frame).iter().map(&f).collect::<Vec<Res>>())
    })
}

/// Binary operator consturctor
/// takes two expressions and a function over the elements (left,right)
/// of the expressions results and creates a new expression
///
/// Expr<Left> -> Expr<Right> -> (Left -> Right -> Result) -> Expr<Result>
///
fn expr2<'a, Left, Right, Res, F>(left: Expr<'a>, right: Expr<'a>, f: F) -> Expr<'a>
where
    Left: ColumnType<'a>,
    Right: ColumnType<'a>,
    Res: ColumnType<'a>,
    F: Fn(&Left, &Right) -> Res + 'a,
{
    Box::new(move |frame: &Frame<'a>| -> Column<'a> {
        Column::from(
            Iterator::zip(left(frame).iter(), right(frame).iter())
                .map(|(l, r)| f(l, r))
                .collect::<Vec<Res>>(),
        )
    })
}

/// Some example unary operators

/// numeric negation
fn neg<'a, T>(arg: Expr<'a>) -> Expr<'a>
where
    T: ColumnType<'a>,
    for<'x> &'x T: Neg<Output = T>,
{
    expr1(arg, move |num: &T| -num)
}

/// string length
fn length<'a>(arg: Expr<'a>) -> Expr<'a> {
    expr1(arg, move |in_str: &Str<'a>| -> f64 { in_str.len() as f64 })
}

/// uppercase a string (unicode aware)
fn upper<'a>(arg: Expr<'a>) -> Expr<'a> {
    expr1(arg, move |in_str: &Str<'a>| -> Str<'a> {
        in_str.chars().flat_map(|c| c.to_uppercase()).collect()
    })
}

/// Some example binary operators

/// numeric addition
fn add<'a, T>(left: Expr<'a>, right: Expr<'a>) -> Expr<'a>
where
    T: ColumnType<'a>,
    for<'x> &'x T: Add<Output = T>,
{
    expr2(left, right, move |l: &T, r: &T| l + r)
}

/// numeric multiplication
fn mul<'a, T>(left: Expr<'a>, right: Expr<'a>) -> Expr<'a>
where
    T: ColumnType<'a>,
    for<'x> &'x T: Mul<Output = T>,
{
    expr2(left, right, move |l: &T, r: &T| l * r)
}

fn main() {
    // Example frame
    let num_col = Column::from(vec![1.0, 2.0, 3.0]);
    let str_col = Column::from(
        vec!["foo", "bar", "baz"]
            .iter()
            .map(|&str| Cow::from(str))
            .collect(),
    );
    let frame = Frame {
        columns: vec![num_col.clone(), num_col, str_col],
        rows: 3,
    };
    // Projection params
    let select: Vec<Expr> = vec![
        // a
        col(0),
        // -a
        neg::<f64>(col(0)),
        // upper(c)
        upper(col(2)),
        // c
        col(2),
        // length(c)
        length(col(2)),
        // a + b
        add::<f64>(col(0), col(1)),
        // a * c
        mul::<f64>(col(0), col(1)),
        // a * 3
        mul::<f64>(col(0), val(3.0)),
        // (a*b) + (a+b)
        add::<f64>(mul::<f64>(col(0), col(1)), add::<f64>(col(0), col(1))),
    ];

    // Execute the projection
    let mut result = Frame {
        columns: vec![],
        rows: 3,
    };
    for _ in 1..1000000 {
        result = Frame {
            columns: vec![],
            rows: 3,
        };
        select.iter().for_each(|expr| {
            result.columns.push(expr(&frame));
        });
    }
    // Print the result and also the input - it should not be touched
    println!("Input:  {:?}", frame);
    println!("Output: {:?}", result);
}