use std::cmp::Ordering;
use arrow::datatypes::*;
#[derive(Debug, Clone, PartialEq)]
pub enum ScalarValue {
Boolean(bool),
Float32(f32),
Float64(f64),
Int8(i8),
Int16(i16),
Int32(i32),
Int64(i64),
UInt8(u8),
UInt16(u16),
UInt32(u32),
UInt64(u64),
Utf8(String),
Struct(Vec<ScalarValue>),
}
impl PartialOrd for ScalarValue {
fn partial_cmp(&self, other: &ScalarValue) -> Option<Ordering> {
match self {
&ScalarValue::Float64(l) => match other {
&ScalarValue::Float64(r) => l.partial_cmp(&r),
&ScalarValue::Int64(r) => l.partial_cmp(&(r as f64)),
_ => unimplemented!("type coercion rules missing"),
},
&ScalarValue::Int64(l) => match other {
&ScalarValue::Float64(r) => (l as f64).partial_cmp(&r),
&ScalarValue::Int64(r) => l.partial_cmp(&r),
_ => unimplemented!("type coercion rules missing"),
},
&ScalarValue::Utf8(ref l) => match other {
&ScalarValue::Utf8(ref r) => l.partial_cmp(r),
_ => unimplemented!("type coercion rules missing"),
},
&ScalarValue::Struct(_) => None,
_ => unimplemented!("type coercion rules missing"),
}
}
}
impl ScalarValue {
pub fn to_string(&self) -> String {
match self {
&ScalarValue::Boolean(b) => b.to_string(),
&ScalarValue::Int8(l) => l.to_string(),
&ScalarValue::Int16(l) => l.to_string(),
&ScalarValue::Int32(l) => l.to_string(),
&ScalarValue::Int64(l) => l.to_string(),
&ScalarValue::UInt8(l) => l.to_string(),
&ScalarValue::UInt16(l) => l.to_string(),
&ScalarValue::UInt32(l) => l.to_string(),
&ScalarValue::UInt64(l) => l.to_string(),
&ScalarValue::Float32(d) => d.to_string(),
&ScalarValue::Float64(d) => d.to_string(),
&ScalarValue::Utf8(ref s) => s.clone(),
&ScalarValue::Struct(ref v) => {
let s: Vec<String> = v.iter().map(|v| v.to_string()).collect();
s.join(",")
}
}
}
}
#[derive(Debug, Clone)]
pub struct FunctionMeta {
pub name: String,
pub args: Vec<Field>,
pub return_type: DataType,
}
pub trait Row {
fn get(&self, index: usize) -> &ScalarValue;
fn to_string(&self) -> String;
}
impl Row for Vec<ScalarValue> {
fn get(&self, index: usize) -> &ScalarValue {
&self[index]
}
fn to_string(&self) -> String {
let value_strings: Vec<String> = self.iter().map(|v| v.to_string()).collect();
value_strings.join(",")
}
}
#[derive(Debug, Clone)]
pub enum Operator {
Eq,
NotEq,
Lt,
LtEq,
Gt,
GtEq,
Plus,
Minus,
Multiply,
Divide,
Modulus,
}
#[derive(Debug, Clone)]
pub enum Expr {
Column(usize),
Literal(ScalarValue),
BinaryExpr {
left: Box<Expr>,
op: Operator,
right: Box<Expr>,
},
Sort { expr: Box<Expr>, asc: bool },
ScalarFunction { name: String, args: Vec<Expr> },
}
impl Expr {
pub fn eq(&self, other: &Expr) -> Expr {
Expr::BinaryExpr {
left: Box::new(self.clone()),
op: Operator::Eq,
right: Box::new(other.clone()),
}
}
pub fn gt(&self, other: &Expr) -> Expr {
Expr::BinaryExpr {
left: Box::new(self.clone()),
op: Operator::Gt,
right: Box::new(other.clone()),
}
}
pub fn lt(&self, other: &Expr) -> Expr {
Expr::BinaryExpr {
left: Box::new(self.clone()),
op: Operator::Lt,
right: Box::new(other.clone()),
}
}
}
#[derive(Debug, Clone)]
pub enum LogicalPlan {
Limit {
limit: usize,
input: Box<LogicalPlan>,
schema: Schema,
},
Projection {
expr: Vec<Expr>,
input: Box<LogicalPlan>,
schema: Schema,
},
Selection {
expr: Expr,
input: Box<LogicalPlan>,
schema: Schema,
},
Sort {
expr: Vec<Expr>,
input: Box<LogicalPlan>,
schema: Schema,
},
TableScan {
schema_name: String,
table_name: String,
schema: Schema,
},
CsvFile {
filename: String,
schema: Schema,
},
EmptyRelation,
}
impl LogicalPlan {
pub fn schema(&self) -> Schema {
match self {
&LogicalPlan::EmptyRelation => Schema::empty(),
&LogicalPlan::TableScan { ref schema, .. } => schema.clone(),
&LogicalPlan::CsvFile { ref schema, .. } => schema.clone(),
&LogicalPlan::Projection { ref schema, .. } => schema.clone(),
&LogicalPlan::Selection { ref schema, .. } => schema.clone(),
&LogicalPlan::Sort { ref schema, .. } => schema.clone(),
&LogicalPlan::Limit { ref schema, .. } => schema.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::Expr::*;
use super::LogicalPlan::*;
use super::ScalarValue::*;
use super::*;
#[test]
fn serde() {
let schema = Schema {
columns: vec![
Field {
name: "id".to_string(),
data_type: DataType::Int32,
nullable: false,
},
Field {
name: "name".to_string(),
data_type: DataType::Utf8,
nullable: false,
},
],
};
let csv = CsvFile {
filename: "test/data/people.csv".to_string(),
schema: schema.clone(),
};
let filter_expr = BinaryExpr {
left: Box::new(Column(0)),
op: Operator::Eq,
right: Box::new(Literal(Int64(2))),
};
let plan = Selection {
expr: filter_expr,
input: Box::new(csv),
schema: schema.clone(),
};
}
}