This is an automated email from the ASF dual-hosted git repository.
paddyhoran pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 37676f7 ARROW-4628: [Rust] [DataFusion] Implement type coercion query
optimizer rule
37676f7 is described below
commit 37676f704cbc1eb622c2de7ad4b4827848e8e1a7
Author: Andy Grove <[email protected]>
AuthorDate: Thu Mar 21 15:10:56 2019 -0400
ARROW-4628: [Rust] [DataFusion] Implement type coercion query optimizer rule
This PR refactors the existing type coercion logic, to remove it from the
SQL query planner and into an optimizer rule and also makes it more complete,
with improved unit tests.
It also converts more functions to return Result instead of using unwrap
and removes some dead code and removes some duplicated code by introducing a
new `utils` file in `optimizer` module.
Author: Andy Grove <[email protected]>
Closes #3939 from andygrove/type_coercion and squashes the following
commits:
d94d8fdb <Andy Grove> use correct error type
f8b25e9e <Andy Grove> manual merge with latest from master
30ef3396 <Andy Grove> rebase
98447920 <Andy Grove> Roll back some changes to reduce scope of PR
1a76ca3b <Andy Grove> Implement type coercion optimizer rule
---
rust/datafusion/src/execution/context.rs | 12 +-
rust/datafusion/src/logicalplan.rs | 25 +--
rust/datafusion/src/optimizer/mod.rs | 1 +
rust/datafusion/src/optimizer/optimizer.rs | 2 +-
.../src/optimizer/projection_push_down.rs | 42 ++--
rust/datafusion/src/optimizer/type_coercion.rs | 211 +++++++++++++++++++++
rust/datafusion/src/optimizer/utils.rs | 59 ++++++
rust/datafusion/src/sql/planner.rs | 43 ++---
8 files changed, 321 insertions(+), 74 deletions(-)
diff --git a/rust/datafusion/src/execution/context.rs
b/rust/datafusion/src/execution/context.rs
index 22c0263..8aee99f 100644
--- a/rust/datafusion/src/execution/context.rs
+++ b/rust/datafusion/src/execution/context.rs
@@ -37,6 +37,7 @@ use crate::execution::relation::{DataSourceRelation,
Relation};
use crate::logicalplan::*;
use crate::optimizer::optimizer::OptimizerRule;
use crate::optimizer::projection_push_down::ProjectionPushDown;
+use crate::optimizer::type_coercion::TypeCoercionRule;
use crate::optimizer::utils;
use crate::sql::parser::{DFASTNode, DFParser};
use crate::sql::planner::{SchemaProvider, SqlToRel};
@@ -107,8 +108,15 @@ impl ExecutionContext {
/// Optimize the logical plan by applying optimizer rules
fn optimize(&self, plan: &LogicalPlan) -> Result<Arc<LogicalPlan>> {
- let mut rule = ProjectionPushDown::new();
- Ok(rule.optimize(plan)?)
+ let rules: Vec<Box<OptimizerRule>> = vec![
+ Box::new(ProjectionPushDown::new()),
+ Box::new(TypeCoercionRule::new()),
+ ];
+ let mut plan = Arc::new(plan.clone());
+ for mut rule in rules {
+ plan = rule.optimize(&plan)?;
+ }
+ Ok(plan)
}
/// Execute a logical plan and produce a Relation (a schema-aware iterator
over a series
diff --git a/rust/datafusion/src/logicalplan.rs
b/rust/datafusion/src/logicalplan.rs
index 432165f..2e2bee2 100644
--- a/rust/datafusion/src/logicalplan.rs
+++ b/rust/datafusion/src/logicalplan.rs
@@ -450,53 +450,54 @@ impl fmt::Debug for LogicalPlan {
}
}
-pub fn can_coerce_from(left: &DataType, other: &DataType) -> bool {
+pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
use self::DataType::*;
- match left {
- Int8 => match other {
+ match type_into {
+ Int8 => match type_from {
Int8 => true,
_ => false,
},
- Int16 => match other {
+ Int16 => match type_from {
Int8 | Int16 => true,
_ => false,
},
- Int32 => match other {
+ Int32 => match type_from {
Int8 | Int16 | Int32 => true,
_ => false,
},
- Int64 => match other {
+ Int64 => match type_from {
Int8 | Int16 | Int32 | Int64 => true,
_ => false,
},
- UInt8 => match other {
+ UInt8 => match type_from {
UInt8 => true,
_ => false,
},
- UInt16 => match other {
+ UInt16 => match type_from {
UInt8 | UInt16 => true,
_ => false,
},
- UInt32 => match other {
+ UInt32 => match type_from {
UInt8 | UInt16 | UInt32 => true,
_ => false,
},
- UInt64 => match other {
+ UInt64 => match type_from {
UInt8 | UInt16 | UInt32 | UInt64 => true,
_ => false,
},
- Float32 => match other {
+ Float32 => match type_from {
Int8 | Int16 | Int32 | Int64 => true,
UInt8 | UInt16 | UInt32 | UInt64 => true,
Float32 => true,
_ => false,
},
- Float64 => match other {
+ Float64 => match type_from {
Int8 | Int16 | Int32 | Int64 => true,
UInt8 | UInt16 | UInt32 | UInt64 => true,
Float32 | Float64 => true,
_ => false,
},
+ Utf8 => true,
_ => false,
}
}
diff --git a/rust/datafusion/src/optimizer/mod.rs
b/rust/datafusion/src/optimizer/mod.rs
index f507f5a..b43fba4 100644
--- a/rust/datafusion/src/optimizer/mod.rs
+++ b/rust/datafusion/src/optimizer/mod.rs
@@ -19,4 +19,5 @@
pub mod optimizer;
pub mod projection_push_down;
+pub mod type_coercion;
pub mod utils;
diff --git a/rust/datafusion/src/optimizer/optimizer.rs
b/rust/datafusion/src/optimizer/optimizer.rs
index f4e368c..4fc184b 100644
--- a/rust/datafusion/src/optimizer/optimizer.rs
+++ b/rust/datafusion/src/optimizer/optimizer.rs
@@ -17,8 +17,8 @@
//! Query optimizer traits
+use crate::error::Result;
use crate::logicalplan::LogicalPlan;
-use arrow::error::Result;
use std::sync::Arc;
/// An optimizer rules performs a transformation on a logical plan to produce
an optimized logical plan.
diff --git a/rust/datafusion/src/optimizer/projection_push_down.rs
b/rust/datafusion/src/optimizer/projection_push_down.rs
index 2ffdfc9..9ffa54b 100644
--- a/rust/datafusion/src/optimizer/projection_push_down.rs
+++ b/rust/datafusion/src/optimizer/projection_push_down.rs
@@ -18,11 +18,12 @@
//! Projection Push Down optimizer rule ensures that only referenced columns
are
//! loaded into memory
+use crate::error::{ExecutionError, Result};
use crate::logicalplan::Expr;
use crate::logicalplan::LogicalPlan;
use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
use arrow::datatypes::{Field, Schema};
-use arrow::error::{ArrowError, Result};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
@@ -56,7 +57,7 @@ impl ProjectionPushDown {
schema,
} => {
// collect all columns referenced by projection expressions
- self.collect_exprs(&expr, accum);
+ utils::exprlist_to_column_indices(&expr, accum);
// push projection down
let input = self.optimize_plan(&input, accum, mapping)?;
@@ -72,7 +73,7 @@ impl ProjectionPushDown {
}
LogicalPlan::Selection { expr, input } => {
// collect all columns referenced by filter expression
- self.collect_expr(expr, accum);
+ utils::expr_to_column_indices(expr, accum);
// push projection down
let input = self.optimize_plan(&input, accum, mapping)?;
@@ -92,8 +93,8 @@ impl ProjectionPushDown {
schema,
} => {
// collect all columns referenced by grouping and aggregate
expressions
- self.collect_exprs(&group_expr, accum);
- self.collect_exprs(&aggr_expr, accum);
+ utils::exprlist_to_column_indices(&group_expr, accum);
+ utils::exprlist_to_column_indices(&aggr_expr, accum);
// push projection down
let input = self.optimize_plan(&input, accum, mapping)?;
@@ -115,7 +116,7 @@ impl ProjectionPushDown {
schema,
} => {
// collect all columns referenced by sort expressions
- self.collect_exprs(&expr, accum);
+ utils::exprlist_to_column_indices(&expr, accum);
// push projection down
let input = self.optimize_plan(&input, accum, mapping)?;
@@ -161,7 +162,9 @@ impl ProjectionPushDown {
// can rewrite expressions as we walk back up the tree
if mapping.len() != 0 {
- return Err(ArrowError::ComputeError("illegal
state".to_string()));
+ return Err(ExecutionError::InternalError(
+ "illegal state".to_string(),
+ ));
}
for i in 0..schema.fields().len() {
@@ -190,29 +193,6 @@ impl ProjectionPushDown {
}
}
- fn collect_exprs(&self, expr: &Vec<Expr>, accum: &mut HashSet<usize>) {
- expr.iter().for_each(|e| self.collect_expr(e, accum));
- }
-
- fn collect_expr(&self, expr: &Expr, accum: &mut HashSet<usize>) {
- match expr {
- Expr::Column(i) => {
- accum.insert(*i);
- }
- Expr::Literal(_) => { /* not needed */ }
- Expr::IsNull(e) => self.collect_expr(e, accum),
- Expr::IsNotNull(e) => self.collect_expr(e, accum),
- Expr::BinaryExpr { left, right, .. } => {
- self.collect_expr(left, accum);
- self.collect_expr(right, accum);
- }
- Expr::Cast { expr, .. } => self.collect_expr(expr, accum),
- Expr::Sort { expr, .. } => self.collect_expr(expr, accum),
- Expr::AggregateFunction { args, .. } => self.collect_exprs(args,
accum),
- Expr::ScalarFunction { args, .. } => self.collect_exprs(args,
accum),
- }
- }
-
fn rewrite_exprs(
&self,
expr: &Vec<Expr>,
@@ -269,7 +249,7 @@ impl ProjectionPushDown {
fn new_index(&self, mapping: &HashMap<usize, usize>, i: &usize) ->
Result<usize> {
match mapping.get(i) {
Some(j) => Ok(*j),
- _ => Err(ArrowError::ComputeError(
+ _ => Err(ExecutionError::InternalError(
"Internal error computing new column index".to_string(),
)),
}
diff --git a/rust/datafusion/src/optimizer/type_coercion.rs
b/rust/datafusion/src/optimizer/type_coercion.rs
new file mode 100644
index 0000000..b2ba386
--- /dev/null
+++ b/rust/datafusion/src/optimizer/type_coercion.rs
@@ -0,0 +1,211 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The type_coercion optimizer rule ensures that all binary operators are
operating on
+//! compatible types by adding explicit cast operations to expressions. For
example,
+//! the operation `c_float + c_int` would be rewritten as `c_float +
CAST(c_int AS float)`.
+//! This keeps the runtime query execution code much simpler.
+
+use std::sync::Arc;
+
+use arrow::datatypes::Schema;
+
+use crate::error::{ExecutionError, Result};
+use crate::logicalplan::Expr;
+use crate::logicalplan::LogicalPlan;
+use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
+
+/// Implementation of type coercion optimizer rule
+pub struct TypeCoercionRule {}
+
+impl OptimizerRule for TypeCoercionRule {
+ fn optimize(&mut self, plan: &LogicalPlan) -> Result<Arc<LogicalPlan>> {
+ match plan {
+ LogicalPlan::Projection {
+ expr,
+ input,
+ schema,
+ } => Ok(Arc::new(LogicalPlan::Projection {
+ expr: expr
+ .iter()
+ .map(|e| rewrite_expr(e, &schema))
+ .collect::<Result<Vec<_>>>()?,
+ input: self.optimize(input)?,
+ schema: schema.clone(),
+ })),
+ LogicalPlan::Selection { expr, input } => {
+ Ok(Arc::new(LogicalPlan::Selection {
+ expr: rewrite_expr(expr, input.schema())?,
+ input: self.optimize(input)?,
+ }))
+ }
+ LogicalPlan::Aggregate {
+ input,
+ group_expr,
+ aggr_expr,
+ schema,
+ } => Ok(Arc::new(LogicalPlan::Aggregate {
+ group_expr: group_expr
+ .iter()
+ .map(|e| rewrite_expr(e, &schema))
+ .collect::<Result<Vec<_>>>()?,
+ aggr_expr: aggr_expr
+ .iter()
+ .map(|e| rewrite_expr(e, &schema))
+ .collect::<Result<Vec<_>>>()?,
+ input: self.optimize(input)?,
+ schema: schema.clone(),
+ })),
+ LogicalPlan::TableScan { .. } => Ok(Arc::new(plan.clone())),
+ LogicalPlan::EmptyRelation { .. } => Ok(Arc::new(plan.clone())),
+ LogicalPlan::Limit { .. } => Ok(Arc::new(plan.clone())),
+ other => Err(ExecutionError::NotImplemented(format!(
+ "Type coercion optimizer rule does not support relation: {:?}",
+ other
+ ))),
+ }
+ }
+}
+
+impl TypeCoercionRule {
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+/// Rewrite an expression to include explicit CAST operations when required
+fn rewrite_expr(expr: &Expr, schema: &Schema) -> Result<Expr> {
+ match expr {
+ Expr::BinaryExpr { left, op, right } => {
+ let left = rewrite_expr(left, schema)?;
+ let right = rewrite_expr(right, schema)?;
+ let left_type = left.get_type(schema);
+ let right_type = right.get_type(schema);
+ if left_type == right_type {
+ Ok(expr.clone())
+ } else {
+ let super_type = utils::get_supertype(&left_type,
&right_type)?;
+ Ok(Expr::BinaryExpr {
+ left: Arc::new(left.cast_to(&super_type, schema)?),
+ op: op.clone(),
+ right: Arc::new(right.cast_to(&super_type, schema)?),
+ })
+ }
+ }
+ Expr::IsNull(e) => Ok(Expr::IsNull(Arc::new(rewrite_expr(e,
schema)?))),
+ Expr::IsNotNull(e) => Ok(Expr::IsNotNull(Arc::new(rewrite_expr(e,
schema)?))),
+ Expr::ScalarFunction {
+ name,
+ args,
+ return_type,
+ } => Ok(Expr::ScalarFunction {
+ name: name.clone(),
+ args: args
+ .iter()
+ .map(|a| rewrite_expr(a, schema))
+ .collect::<Result<Vec<_>>>()?,
+ return_type: return_type.clone(),
+ }),
+ Expr::AggregateFunction {
+ name,
+ args,
+ return_type,
+ } => Ok(Expr::AggregateFunction {
+ name: name.clone(),
+ args: args
+ .iter()
+ .map(|a| rewrite_expr(a, schema))
+ .collect::<Result<Vec<_>>>()?,
+ return_type: return_type.clone(),
+ }),
+ Expr::Cast { .. } => Ok(expr.clone()),
+ Expr::Column(_) => Ok(expr.clone()),
+ Expr::Literal(_) => Ok(expr.clone()),
+ other => Err(ExecutionError::NotImplemented(format!(
+ "Type coercion optimizer rule does not support expression: {:?}",
+ other
+ ))),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::logicalplan::Expr::*;
+ use crate::logicalplan::Operator;
+ use arrow::datatypes::{DataType, Field, Schema};
+
+ #[test]
+ fn test_add_i32_i64() {
+ binary_cast_test(
+ DataType::Int32,
+ DataType::Int64,
+ "CAST(#0 AS Int64) Plus #1",
+ );
+ binary_cast_test(
+ DataType::Int64,
+ DataType::Int32,
+ "#0 Plus CAST(#1 AS Int64)",
+ );
+ }
+
+ #[test]
+ fn test_add_f32_f64() {
+ binary_cast_test(
+ DataType::Float32,
+ DataType::Float64,
+ "CAST(#0 AS Float64) Plus #1",
+ );
+ binary_cast_test(
+ DataType::Float64,
+ DataType::Float32,
+ "#0 Plus CAST(#1 AS Float64)",
+ );
+ }
+
+ #[test]
+ fn test_add_i32_f32() {
+ binary_cast_test(
+ DataType::Int32,
+ DataType::Float32,
+ "CAST(#0 AS Float32) Plus #1",
+ );
+ binary_cast_test(
+ DataType::Float32,
+ DataType::Int32,
+ "#0 Plus CAST(#1 AS Float32)",
+ );
+ }
+
+ fn binary_cast_test(left_type: DataType, right_type: DataType, expected:
&str) {
+ let schema = Schema::new(vec![
+ Field::new("c0", left_type, true),
+ Field::new("c1", right_type, true),
+ ]);
+
+ let expr = Expr::BinaryExpr {
+ left: Arc::new(Column(0)),
+ op: Operator::Plus,
+ right: Arc::new(Column(1)),
+ };
+
+ let expr2 = rewrite_expr(&expr, &schema).unwrap();
+
+ assert_eq!(expected, format!("{:?}", expr2));
+ }
+}
diff --git a/rust/datafusion/src/optimizer/utils.rs
b/rust/datafusion/src/optimizer/utils.rs
index be97859..10f6a44 100644
--- a/rust/datafusion/src/optimizer/utils.rs
+++ b/rust/datafusion/src/optimizer/utils.rs
@@ -17,11 +17,40 @@
//! Collection of utility functions that are leveraged by the query optimizer
rules
+use std::collections::HashSet;
+
use arrow::datatypes::{DataType, Field, Schema};
use crate::error::{ExecutionError, Result};
use crate::logicalplan::Expr;
+/// Recursively walk a list of expression trees, collecting the unique set of
column indexes
+/// referenced in the expression
+pub fn exprlist_to_column_indices(expr: &Vec<Expr>, accum: &mut
HashSet<usize>) {
+ expr.iter().for_each(|e| expr_to_column_indices(e, accum));
+}
+
+/// Recursively walk an expression tree, collecting the unique set of column
indexes
+/// referenced in the expression
+pub fn expr_to_column_indices(expr: &Expr, accum: &mut HashSet<usize>) {
+ match expr {
+ Expr::Column(i) => {
+ accum.insert(*i);
+ }
+ Expr::Literal(_) => { /* not needed */ }
+ Expr::IsNull(e) => expr_to_column_indices(e, accum),
+ Expr::IsNotNull(e) => expr_to_column_indices(e, accum),
+ Expr::BinaryExpr { left, right, .. } => {
+ expr_to_column_indices(left, accum);
+ expr_to_column_indices(right, accum);
+ }
+ Expr::Cast { expr, .. } => expr_to_column_indices(expr, accum),
+ Expr::Sort { expr, .. } => expr_to_column_indices(expr, accum),
+ Expr::AggregateFunction { args, .. } =>
exprlist_to_column_indices(args, accum),
+ Expr::ScalarFunction { args, .. } => exprlist_to_column_indices(args,
accum),
+ }
+}
+
/// Create field meta-data from an expression, for use in a result set schema
pub fn expr_to_field(e: &Expr, input_schema: &Schema) -> Result<Field> {
match e {
@@ -182,3 +211,33 @@ fn _get_supertype(l: &DataType, r: &DataType) ->
Option<DataType> {
_ => None,
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::logicalplan::Expr;
+ use arrow::datatypes::DataType;
+ use std::collections::HashSet;
+ use std::sync::Arc;
+
+ #[test]
+ fn test_collect_expr() {
+ let mut accum: HashSet<usize> = HashSet::new();
+ expr_to_column_indices(
+ &Expr::Cast {
+ expr: Arc::new(Expr::Column(3)),
+ data_type: DataType::Float64,
+ },
+ &mut accum,
+ );
+ expr_to_column_indices(
+ &Expr::Cast {
+ expr: Arc::new(Expr::Column(3)),
+ data_type: DataType::Float64,
+ },
+ &mut accum,
+ );
+ assert_eq!(1, accum.len());
+ assert!(accum.contains(&3));
+ }
+}
diff --git a/rust/datafusion/src/sql/planner.rs
b/rust/datafusion/src/sql/planner.rs
index 8c84bb3..655694a 100644
--- a/rust/datafusion/src/sql/planner.rs
+++ b/rust/datafusion/src/sql/planner.rs
@@ -20,8 +20,8 @@
use std::string::String;
use std::sync::Arc;
-use crate::error::*;
-use crate::logicalplan::*;
+use crate::error::{ExecutionError, Result};
+use crate::logicalplan::{Expr, FunctionMeta, LogicalPlan, Operator,
ScalarValue};
use crate::optimizer::utils;
use arrow::datatypes::*;
@@ -279,25 +279,11 @@ impl SqlToRel {
&SQLOperator::NotLike => Operator::NotLike,
};
- let left_expr = self.sql_to_rex(&left, &schema)?;
- let right_expr = self.sql_to_rex(&right, &schema)?;
- let left_type = left_expr.get_type(schema);
- let right_type = right_expr.get_type(schema);
-
- match utils::get_supertype(&left_type, &right_type) {
- Ok(supertype) => Ok(Expr::BinaryExpr {
- left: Arc::new(left_expr.cast_to(&supertype, schema)?),
- op: operator,
- right: Arc::new(right_expr.cast_to(&supertype,
schema)?),
- }),
- Err(_) => {
- return Err(ExecutionError::General(format!(
- "No common supertype found for binary operator
{:?} \
- with input types {:?} and {:?}",
- operator, left_type, right_type
- )));
- }
- }
+ Ok(Expr::BinaryExpr {
+ left: Arc::new(self.sql_to_rex(&left, &schema)?),
+ op: operator,
+ right: Arc::new(self.sql_to_rex(&right, &schema)?),
+ })
}
// &ASTNode::SQLOrderBy { ref expr, asc } =>
Ok(Expr::Sort {
@@ -401,6 +387,7 @@ pub fn convert_data_type(sql: &SQLType) -> Result<DataType>
{
mod tests {
use super::*;
+ use crate::logicalplan::FunctionType;
use sqlparser::sqlparser::*;
#[test]
@@ -437,7 +424,7 @@ mod tests {
FROM person WHERE state = 'CO' AND age >= 21 AND age <= 65";
let expected =
"Projection: #0, #1, #2\
- \n Selection: #4 Eq Utf8(\"CO\") And CAST(#3 AS Int64) GtEq
Int64(21) And CAST(#3 AS Int64) LtEq Int64(65)\
+ \n Selection: #4 Eq Utf8(\"CO\") And #3 GtEq Int64(21) And #3
LtEq Int64(65)\
\n TableScan: person projection=None";
quick_test(sql, expected);
}
@@ -453,12 +440,12 @@ mod tests {
AND age < 65 \
AND age <= 65";
let expected = "Projection: #3, #1, #2\
- \n Selection: CAST(#3 AS Int64) Eq Int64(21) \
- And CAST(#3 AS Int64) NotEq Int64(21) \
- And CAST(#3 AS Int64) Gt Int64(21) \
- And CAST(#3 AS Int64) GtEq Int64(21) \
- And CAST(#3 AS Int64) Lt Int64(65) \
- And CAST(#3 AS Int64) LtEq Int64(65)\
+ \n Selection: #3 Eq Int64(21) \
+ And #3 NotEq Int64(21) \
+ And #3 Gt Int64(21) \
+ And #3 GtEq Int64(21) \
+ And #3 Lt Int64(65) \
+ And #3 LtEq Int64(65)\
\n TableScan: person projection=None";
quick_test(sql, expected);
}
