This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new d7c0e420c remove type coercion in the binary physical expr (#3396)
d7c0e420c is described below
commit d7c0e420ce6362a253e1f55a873a0732052e0e2f
Author: Kun Liu <[email protected]>
AuthorDate: Sat Sep 24 19:51:12 2022 +0800
remove type coercion in the binary physical expr (#3396)
* remove type coercion binary from phy
* fix test case
* revert the fix for #3387
* type coercion before simplify expression
* complete remove the type coercion in the physical plan
* refactor
* merge master
* refactor
* do type coercion in the simplify expression
* Add comments
* fix: fmt
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/core/src/execution/context.rs | 4 +
.../src/physical_optimizer/aggregate_statistics.rs | 5 +-
.../core/src/physical_plan/file_format/parquet.rs | 22 +-
datafusion/core/src/physical_plan/planner.rs | 4 +-
datafusion/core/tests/sql/aggregates.rs | 10 +-
datafusion/core/tests/sql/decimal.rs | 114 +++++-----
datafusion/core/tests/sql/predicates.rs | 13 +-
datafusion/core/tests/sql/select.rs | 4 +-
datafusion/expr/src/binary_rule.rs | 229 +++++++++++++++++++++
datafusion/optimizer/src/simplify_expressions.rs | 96 ++++++---
datafusion/optimizer/src/type_coercion.rs | 57 +++--
datafusion/optimizer/tests/integration-test.rs | 12 +-
datafusion/physical-expr/src/expressions/binary.rs | 211 +++++++++----------
13 files changed, 529 insertions(+), 252 deletions(-)
diff --git a/datafusion/core/src/execution/context.rs
b/datafusion/core/src/execution/context.rs
index ec6374d7a..27476313c 100644
--- a/datafusion/core/src/execution/context.rs
+++ b/datafusion/core/src/execution/context.rs
@@ -1454,7 +1454,11 @@ impl SessionState {
rules.push(Arc::new(FilterNullJoinKeys::default()));
}
rules.push(Arc::new(ReduceOuterJoin::new()));
+ // TODO: https://github.com/apache/arrow-datafusion/issues/3557
+ // remove this, after the issue fixed.
rules.push(Arc::new(TypeCoercion::new()));
+ // after the type coercion, can do simplify expression again
+ rules.push(Arc::new(SimplifyExpressions::new()));
rules.push(Arc::new(FilterPushDown::new()));
rules.push(Arc::new(LimitPushDown::new()));
rules.push(Arc::new(SingleDistinctToGroupBy::new()));
diff --git a/datafusion/core/src/physical_optimizer/aggregate_statistics.rs
b/datafusion/core/src/physical_optimizer/aggregate_statistics.rs
index 4a941ec4b..bb1e49cf3 100644
--- a/datafusion/core/src/physical_optimizer/aggregate_statistics.rs
+++ b/datafusion/core/src/physical_optimizer/aggregate_statistics.rs
@@ -261,6 +261,7 @@ mod tests {
use arrow::array::{Int32Array, Int64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
+ use datafusion_physical_expr::expressions::cast;
use datafusion_physical_expr::PhysicalExpr;
use crate::error::Result;
@@ -525,7 +526,7 @@ mod tests {
expressions::binary(
expressions::col("a", &schema)?,
Operator::Gt,
- expressions::lit(1u32),
+ cast(expressions::lit(1u32), &schema, DataType::Int32)?,
&schema,
)?,
source,
@@ -568,7 +569,7 @@ mod tests {
expressions::binary(
expressions::col("a", &schema)?,
Operator::Gt,
- expressions::lit(1u32),
+ cast(expressions::lit(1u32), &schema, DataType::Int32)?,
&schema,
)?,
source,
diff --git a/datafusion/core/src/physical_plan/file_format/parquet.rs
b/datafusion/core/src/physical_plan/file_format/parquet.rs
index f4c52e5dc..ff6507c7f 100644
--- a/datafusion/core/src/physical_plan/file_format/parquet.rs
+++ b/datafusion/core/src/physical_plan/file_format/parquet.rs
@@ -871,13 +871,14 @@ mod tests {
physical_plan::collect,
};
use arrow::array::Float32Array;
+ use arrow::datatypes::DataType::Decimal128;
use arrow::record_batch::RecordBatch;
use arrow::{
array::{Int64Array, Int8Array, StringArray},
datatypes::{DataType, Field},
};
use chrono::{TimeZone, Utc};
- use datafusion_expr::{col, lit};
+ use datafusion_expr::{cast, col, lit};
use futures::StreamExt;
use object_store::local::LocalFileSystem;
use object_store::path::Path;
@@ -1768,6 +1769,7 @@ mod tests {
// In this case, construct four types of statistics to filtered with
the decimal predication.
// INT32: c1 > 5, the c1 is decimal(9,2)
+ // The type of scalar value if decimal(9,2), don't need to do cast
let expr = col("c1").gt(lit(ScalarValue::Decimal128(Some(500), 9, 2)));
let schema =
Schema::new(vec![Field::new("c1", DataType::Decimal128(9, 2),
false)]);
@@ -1809,11 +1811,15 @@ mod tests {
);
// INT32: c1 > 5, but parquet decimal type has different precision or
scale to arrow decimal
+ // The c1 type is decimal(9,0) in the parquet file, and the type of
scalar is decimal(5,2).
+ // We should convert all type to the coercion type, which is
decimal(11,2)
// The decimal of arrow is decimal(5,2), the decimal of parquet is
decimal(9,0)
- let expr = col("c1").gt(lit(ScalarValue::Decimal128(Some(500), 5, 2)));
+ let expr = cast(col("c1"), DataType::Decimal128(11, 2)).gt(cast(
+ lit(ScalarValue::Decimal128(Some(500), 5, 2)),
+ Decimal128(11, 2),
+ ));
let schema =
- Schema::new(vec![Field::new("c1", DataType::Decimal128(5, 2),
false)]);
- // The decimal of parquet is decimal(9,0)
+ Schema::new(vec![Field::new("c1", DataType::Decimal128(9, 0),
false)]);
let schema_descr = get_test_schema_descr(vec![(
"c1",
PhysicalType::INT32,
@@ -1901,11 +1907,13 @@ mod tests {
vec![1]
);
- // FIXED_LENGTH_BYTE_ARRAY: c1 = 100, the c1 is decimal(28,2)
+ // FIXED_LENGTH_BYTE_ARRAY: c1 = decimal128(100000, 28, 3), the c1 is
decimal(18,2)
// the type of parquet is decimal(18,2)
- let expr = col("c1").eq(lit(ScalarValue::Decimal128(Some(100000), 28,
3)));
let schema =
- Schema::new(vec![Field::new("c1", DataType::Decimal128(18, 3),
false)]);
+ Schema::new(vec![Field::new("c1", DataType::Decimal128(18, 2),
false)]);
+ // cast the type of c1 to decimal(28,3)
+ let left = cast(col("c1"), DataType::Decimal128(28, 3));
+ let expr = left.eq(lit(ScalarValue::Decimal128(Some(100000), 28, 3)));
let schema_descr = get_test_schema_descr(vec![(
"c1",
PhysicalType::FIXED_LEN_BYTE_ARRAY,
diff --git a/datafusion/core/src/physical_plan/planner.rs
b/datafusion/core/src/physical_plan/planner.rs
index 48b78722f..8b247a452 100644
--- a/datafusion/core/src/physical_plan/planner.rs
+++ b/datafusion/core/src/physical_plan/planner.rs
@@ -1685,7 +1685,7 @@ mod tests {
use crate::execution::runtime_env::RuntimeEnv;
use crate::logical_plan::plan::Extension;
use crate::physical_plan::{
- expressions, DisplayFormatType, Partitioning, Statistics,
+ expressions, DisplayFormatType, Partitioning, PhysicalPlanner,
Statistics,
};
use crate::prelude::{SessionConfig, SessionContext};
use crate::scalar::ScalarValue;
@@ -1736,10 +1736,10 @@ mod tests {
let exec_plan = plan(&logical_plan).await?;
// verify that the plan correctly casts u8 to i64
+ // the cast from u8 to i64 for literal will be simplified, and get
lit(int64(5))
// the cast here is implicit so has CastOptions with safe=true
let expected = "BinaryExpr { left: Column { name: \"c7\", index: 2 },
op: Lt, right: Literal { value: Int64(5) } }";
assert!(format!("{:?}", exec_plan).contains(expected));
-
Ok(())
}
diff --git a/datafusion/core/tests/sql/aggregates.rs
b/datafusion/core/tests/sql/aggregates.rs
index b7f24992c..357addbc0 100644
--- a/datafusion/core/tests/sql/aggregates.rs
+++ b/datafusion/core/tests/sql/aggregates.rs
@@ -1834,11 +1834,11 @@ async fn aggregate_avg_add() -> Result<()> {
assert_eq!(results.len(), 1);
let expected = vec![
-
"+--------------+---------------------------+---------------------------+---------------------------+",
- "| AVG(test.c1) | AVG(test.c1) + Float64(1) | AVG(test.c1) +
Float64(2) | Float64(1) + AVG(test.c1) |",
-
"+--------------+---------------------------+---------------------------+---------------------------+",
- "| 1.5 | 2.5 | 3.5
| 2.5 |",
-
"+--------------+---------------------------+---------------------------+---------------------------+",
+
"+--------------+-------------------------+-------------------------+-------------------------+",
+ "| AVG(test.c1) | AVG(test.c1) + Int64(1) | AVG(test.c1) + Int64(2) |
Int64(1) + AVG(test.c1) |",
+
"+--------------+-------------------------+-------------------------+-------------------------+",
+ "| 1.5 | 2.5 | 3.5 |
2.5 |",
+
"+--------------+-------------------------+-------------------------+-------------------------+",
];
assert_batches_sorted_eq!(expected, &results);
diff --git a/datafusion/core/tests/sql/decimal.rs
b/datafusion/core/tests/sql/decimal.rs
index 0898be62c..9d32f1c31 100644
--- a/datafusion/core/tests/sql/decimal.rs
+++ b/datafusion/core/tests/sql/decimal.rs
@@ -376,25 +376,25 @@ async fn decimal_arithmetic_op() -> Result<()> {
actual[0].schema().field(0).data_type()
);
let expected = vec![
- "+----------------------------------------------------+",
- "| decimal_simple.c1 + Decimal128(Some(1000000),27,6) |",
- "+----------------------------------------------------+",
- "| 1.000010 |",
- "| 1.000020 |",
- "| 1.000020 |",
- "| 1.000030 |",
- "| 1.000030 |",
- "| 1.000030 |",
- "| 1.000040 |",
- "| 1.000040 |",
- "| 1.000040 |",
- "| 1.000040 |",
- "| 1.000050 |",
- "| 1.000050 |",
- "| 1.000050 |",
- "| 1.000050 |",
- "| 1.000050 |",
- "+----------------------------------------------------+",
+ "+------------------------------+",
+ "| decimal_simple.c1 + Int64(1) |",
+ "+------------------------------+",
+ "| 1.000010 |",
+ "| 1.000020 |",
+ "| 1.000020 |",
+ "| 1.000030 |",
+ "| 1.000030 |",
+ "| 1.000030 |",
+ "| 1.000040 |",
+ "| 1.000040 |",
+ "| 1.000040 |",
+ "| 1.000040 |",
+ "| 1.000050 |",
+ "| 1.000050 |",
+ "| 1.000050 |",
+ "| 1.000050 |",
+ "| 1.000050 |",
+ "+------------------------------+",
];
assert_batches_eq!(expected, &actual);
// array decimal(10,6) + array decimal(12,7) => decimal(13,7)
@@ -434,25 +434,25 @@ async fn decimal_arithmetic_op() -> Result<()> {
actual[0].schema().field(0).data_type()
);
let expected = vec![
- "+----------------------------------------------------+",
- "| decimal_simple.c1 - Decimal128(Some(1000000),27,6) |",
- "+----------------------------------------------------+",
- "| -0.999990 |",
- "| -0.999980 |",
- "| -0.999980 |",
- "| -0.999970 |",
- "| -0.999970 |",
- "| -0.999970 |",
- "| -0.999960 |",
- "| -0.999960 |",
- "| -0.999960 |",
- "| -0.999960 |",
- "| -0.999950 |",
- "| -0.999950 |",
- "| -0.999950 |",
- "| -0.999950 |",
- "| -0.999950 |",
- "+----------------------------------------------------+",
+ "+------------------------------+",
+ "| decimal_simple.c1 - Int64(1) |",
+ "+------------------------------+",
+ "| -0.999990 |",
+ "| -0.999980 |",
+ "| -0.999980 |",
+ "| -0.999970 |",
+ "| -0.999970 |",
+ "| -0.999970 |",
+ "| -0.999960 |",
+ "| -0.999960 |",
+ "| -0.999960 |",
+ "| -0.999960 |",
+ "| -0.999950 |",
+ "| -0.999950 |",
+ "| -0.999950 |",
+ "| -0.999950 |",
+ "| -0.999950 |",
+ "+------------------------------+",
];
assert_batches_eq!(expected, &actual);
@@ -492,25 +492,25 @@ async fn decimal_arithmetic_op() -> Result<()> {
actual[0].schema().field(0).data_type()
);
let expected = vec![
- "+-----------------------------------------------------+",
- "| decimal_simple.c1 * Decimal128(Some(20000000),31,6) |",
- "+-----------------------------------------------------+",
- "| 0.000200 |",
- "| 0.000400 |",
- "| 0.000400 |",
- "| 0.000600 |",
- "| 0.000600 |",
- "| 0.000600 |",
- "| 0.000800 |",
- "| 0.000800 |",
- "| 0.000800 |",
- "| 0.000800 |",
- "| 0.001000 |",
- "| 0.001000 |",
- "| 0.001000 |",
- "| 0.001000 |",
- "| 0.001000 |",
- "+-----------------------------------------------------+",
+ "+-------------------------------+",
+ "| decimal_simple.c1 * Int64(20) |",
+ "+-------------------------------+",
+ "| 0.000200 |",
+ "| 0.000400 |",
+ "| 0.000400 |",
+ "| 0.000600 |",
+ "| 0.000600 |",
+ "| 0.000600 |",
+ "| 0.000800 |",
+ "| 0.000800 |",
+ "| 0.000800 |",
+ "| 0.000800 |",
+ "| 0.001000 |",
+ "| 0.001000 |",
+ "| 0.001000 |",
+ "| 0.001000 |",
+ "| 0.001000 |",
+ "+-------------------------------+",
];
assert_batches_eq!(expected, &actual);
diff --git a/datafusion/core/tests/sql/predicates.rs
b/datafusion/core/tests/sql/predicates.rs
index 5b1fa92e5..67432b08f 100644
--- a/datafusion/core/tests/sql/predicates.rs
+++ b/datafusion/core/tests/sql/predicates.rs
@@ -389,6 +389,7 @@ async fn csv_in_set_test() -> Result<()> {
#[tokio::test]
async fn multiple_or_predicates() -> Result<()> {
+ // TODO https://github.com/apache/arrow-datafusion/issues/3587
let ctx = SessionContext::new();
register_tpch_csv(&ctx, "lineitem").await?;
register_tpch_csv(&ctx, "part").await?;
@@ -424,15 +425,13 @@ async fn multiple_or_predicates() -> Result<()> {
let plan = state.optimize(&plan)?;
// Note that we expect `#part.p_partkey = #lineitem.l_partkey` to have been
// factored out and appear only once in the following plan
- let expected =vec![
+ let expected = vec![
"Explain [plan_type:Utf8, plan:Utf8]",
" Projection: #lineitem.l_partkey [l_partkey:Int64]",
- " Projection: #part.p_size >= Int32(1) AS #part.p_size >=
Int32(1)Int32(1)#part.p_size, #lineitem.l_partkey, #lineitem.l_quantity,
#part.p_brand, #part.p_size [#part.p_size >=
Int32(1)Int32(1)#part.p_size:Boolean;N, l_partkey:Int64,
l_quantity:Decimal128(15, 2), p_brand:Utf8, p_size:Int32]",
- " Filter: #part.p_brand = Utf8(\"Brand#12\") AND
#lineitem.l_quantity >= Decimal128(Some(100),15,2) AND #lineitem.l_quantity <=
Decimal128(Some(1100),15,2) AND #part.p_size <= Int32(5) OR #part.p_brand =
Utf8(\"Brand#23\") AND #lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND
#lineitem.l_quantity <= Decimal128(Some(2000),15,2) AND #part.p_size <=
Int32(10) OR #part.p_brand = Utf8(\"Brand#34\") AND #lineitem.l_quantity >=
Decimal128(Some(2000),15,2) AND #lineitem.l_quan [...]
- " Inner Join: #lineitem.l_partkey = #part.p_partkey
[l_partkey:Int64, l_quantity:Decimal128(15, 2), p_partkey:Int64, p_brand:Utf8,
p_size:Int32]",
- " TableScan: lineitem projection=[l_partkey, l_quantity]
[l_partkey:Int64, l_quantity:Decimal128(15, 2)]",
- " Filter: #part.p_size >= Int32(1) [p_partkey:Int64,
p_brand:Utf8, p_size:Int32]",
- " TableScan: part projection=[p_partkey, p_brand, p_size],
partial_filters=[#part.p_size >= Int32(1)] [p_partkey:Int64, p_brand:Utf8,
p_size:Int32]",
+ " Filter: #part.p_brand = Utf8(\"Brand#12\") AND
#lineitem.l_quantity >= Decimal128(Some(100),15,2) AND #lineitem.l_quantity <=
Decimal128(Some(1100),15,2) AND CAST(#part.p_size AS Int64) BETWEEN Int64(1)
AND Int64(5) OR #part.p_brand = Utf8(\"Brand#23\") AND #lineitem.l_quantity >=
Decimal128(Some(1000),15,2) AND #lineitem.l_quantity <=
Decimal128(Some(2000),15,2) AND CAST(#part.p_size AS Int64) BETWEEN Int64(1)
AND Int64(10) OR #part.p_brand = Utf8(\"Brand#34\") AND #lineite [...]
+ " Inner Join: #lineitem.l_partkey = #part.p_partkey
[l_partkey:Int64, l_quantity:Decimal128(15, 2), p_partkey:Int64, p_brand:Utf8,
p_size:Int32]",
+ " TableScan: lineitem projection=[l_partkey, l_quantity]
[l_partkey:Int64, l_quantity:Decimal128(15, 2)]",
+ " TableScan: part projection=[p_partkey, p_brand, p_size]
[p_partkey:Int64, p_brand:Utf8, p_size:Int32]",
];
let formatted = plan.display_indent_schema().to_string();
let actual: Vec<&str> = formatted.trim().lines().collect();
diff --git a/datafusion/core/tests/sql/select.rs
b/datafusion/core/tests/sql/select.rs
index 461764e1e..5823f4c0f 100644
--- a/datafusion/core/tests/sql/select.rs
+++ b/datafusion/core/tests/sql/select.rs
@@ -523,12 +523,12 @@ async fn use_between_expression_in_select_query() ->
Result<()> {
.unwrap()
.to_string();
+ // TODO https://github.com/apache/arrow-datafusion/issues/3587
// Only test that the projection exprs are correct, rather than entire
output
let needle = "ProjectionExec: expr=[c1@0 >= 2 AND c1@0 <= 3 as test.c1
BETWEEN Int64(2) AND Int64(3)]";
assert_contains!(&formatted, needle);
- let needle = "Projection: #test.c1 >= Int64(2) AND #test.c1 <= Int64(3)";
+ let needle = "Projection: #test.c1 BETWEEN Int64(2) AND Int64(3)";
assert_contains!(&formatted, needle);
-
Ok(())
}
diff --git a/datafusion/expr/src/binary_rule.rs
b/datafusion/expr/src/binary_rule.rs
index c3c1c4290..77e02cf3a 100644
--- a/datafusion/expr/src/binary_rule.rs
+++ b/datafusion/expr/src/binary_rule.rs
@@ -298,6 +298,8 @@ fn mathematics_numerical_coercion(
};
// same type => all good
+ // TODO: remove this
+ // bug: https://github.com/apache/arrow-datafusion/issues/3387
if lhs_type == rhs_type {
return Some(lhs_type.clone());
}
@@ -632,6 +634,7 @@ fn null_coercion(lhs_type: &DataType, rhs_type: &DataType)
-> Option<DataType> {
_ => None,
}
}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -799,4 +802,230 @@ mod tests {
Some(rhs_type.clone())
);
}
+
+ macro_rules! test_coercion_binary_rule {
+ ($A_TYPE:expr, $B_TYPE:expr, $OP:expr, $C_TYPE:expr) => {{
+ let result = coerce_types(&$A_TYPE, &$OP, &$B_TYPE)?;
+ assert_eq!(result, $C_TYPE);
+ }};
+ }
+
+ #[test]
+ fn test_type_coercion() -> Result<()> {
+ test_coercion_binary_rule!(
+ DataType::Utf8,
+ DataType::Utf8,
+ Operator::Like,
+ DataType::Utf8
+ );
+ test_coercion_binary_rule!(
+ DataType::Utf8,
+ DataType::Date32,
+ Operator::Eq,
+ DataType::Date32
+ );
+ test_coercion_binary_rule!(
+ DataType::Utf8,
+ DataType::Date64,
+ Operator::Lt,
+ DataType::Date64
+ );
+ test_coercion_binary_rule!(
+ DataType::Utf8,
+ DataType::Utf8,
+ Operator::RegexMatch,
+ DataType::Utf8
+ );
+ test_coercion_binary_rule!(
+ DataType::Utf8,
+ DataType::Utf8,
+ Operator::RegexNotMatch,
+ DataType::Utf8
+ );
+ test_coercion_binary_rule!(
+ DataType::Utf8,
+ DataType::Utf8,
+ Operator::RegexNotIMatch,
+ DataType::Utf8
+ );
+ test_coercion_binary_rule!(
+ DataType::Int16,
+ DataType::Int64,
+ Operator::BitwiseAnd,
+ DataType::Int64
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn test_type_coercion_arithmetic() -> Result<()> {
+ // integer
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::UInt32,
+ Operator::Plus,
+ DataType::Int32
+ );
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::UInt16,
+ Operator::Minus,
+ DataType::Int32
+ );
+ test_coercion_binary_rule!(
+ DataType::Int8,
+ DataType::Int64,
+ Operator::Multiply,
+ DataType::Int64
+ );
+ // float
+ test_coercion_binary_rule!(
+ DataType::Float32,
+ DataType::Int32,
+ Operator::Plus,
+ DataType::Float32
+ );
+ test_coercion_binary_rule!(
+ DataType::Float32,
+ DataType::Float64,
+ Operator::Multiply,
+ DataType::Float64
+ );
+ // decimal
+ // bug: https://github.com/apache/arrow-datafusion/issues/3387 will be
fixed in the next pr
+ // test_coercion_binary_rule!(
+ // DataType::Decimal128(10, 2),
+ // DataType::Decimal128(10, 2),
+ // Operator::Plus,
+ // DataType::Decimal128(11, 2)
+ // );
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::Decimal128(10, 2),
+ Operator::Plus,
+ DataType::Decimal128(13, 2)
+ );
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::Decimal128(10, 2),
+ Operator::Minus,
+ DataType::Decimal128(13, 2)
+ );
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::Decimal128(10, 2),
+ Operator::Multiply,
+ DataType::Decimal128(21, 2)
+ );
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::Decimal128(10, 2),
+ Operator::Divide,
+ DataType::Decimal128(23, 11)
+ );
+ test_coercion_binary_rule!(
+ DataType::Int32,
+ DataType::Decimal128(10, 2),
+ Operator::Modulo,
+ DataType::Decimal128(10, 2)
+ );
+ // TODO add other data type
+ Ok(())
+ }
+
+ #[test]
+ fn test_type_coercion_compare() -> Result<()> {
+ // boolean
+ test_coercion_binary_rule!(
+ DataType::Boolean,
+ DataType::Boolean,
+ Operator::Eq,
+ DataType::Boolean
+ );
+ // float
+ test_coercion_binary_rule!(
+ DataType::Float32,
+ DataType::Int64,
+ Operator::Eq,
+ DataType::Float32
+ );
+ test_coercion_binary_rule!(
+ DataType::Float32,
+ DataType::Float64,
+ Operator::GtEq,
+ DataType::Float64
+ );
+ // signed integer
+ test_coercion_binary_rule!(
+ DataType::Int8,
+ DataType::Int32,
+ Operator::LtEq,
+ DataType::Int32
+ );
+ test_coercion_binary_rule!(
+ DataType::Int64,
+ DataType::Int32,
+ Operator::LtEq,
+ DataType::Int64
+ );
+ // unsigned integer
+ test_coercion_binary_rule!(
+ DataType::UInt32,
+ DataType::UInt8,
+ Operator::Gt,
+ DataType::UInt32
+ );
+ // numeric/decimal
+ test_coercion_binary_rule!(
+ DataType::Int64,
+ DataType::Decimal128(10, 0),
+ Operator::Eq,
+ DataType::Decimal128(20, 0)
+ );
+ test_coercion_binary_rule!(
+ DataType::Int64,
+ DataType::Decimal128(10, 2),
+ Operator::Lt,
+ DataType::Decimal128(22, 2)
+ );
+ test_coercion_binary_rule!(
+ DataType::Float64,
+ DataType::Decimal128(10, 3),
+ Operator::Gt,
+ DataType::Decimal128(30, 15)
+ );
+ test_coercion_binary_rule!(
+ DataType::Int64,
+ DataType::Decimal128(10, 0),
+ Operator::Eq,
+ DataType::Decimal128(20, 0)
+ );
+ test_coercion_binary_rule!(
+ DataType::Decimal128(14, 2),
+ DataType::Decimal128(10, 3),
+ Operator::GtEq,
+ DataType::Decimal128(15, 3)
+ );
+
+ // TODO add other data type
+ Ok(())
+ }
+
+ #[test]
+ fn test_type_coercion_logical_op() -> Result<()> {
+ test_coercion_binary_rule!(
+ DataType::Boolean,
+ DataType::Boolean,
+ Operator::And,
+ DataType::Boolean
+ );
+
+ test_coercion_binary_rule!(
+ DataType::Boolean,
+ DataType::Boolean,
+ Operator::Or,
+ DataType::Boolean
+ );
+ Ok(())
+ }
}
diff --git a/datafusion/optimizer/src/simplify_expressions.rs
b/datafusion/optimizer/src/simplify_expressions.rs
index a9c0d4b1f..6ab0eb87c 100644
--- a/datafusion/optimizer/src/simplify_expressions.rs
+++ b/datafusion/optimizer/src/simplify_expressions.rs
@@ -18,6 +18,7 @@
//! Simplify expressions optimizer rule
use crate::expr_simplifier::ExprSimplifiable;
+use crate::type_coercion::TypeCoercionRewriter;
use crate::{expr_simplifier::SimplifyInfo, OptimizerConfig, OptimizerRule};
use arrow::array::new_null_array;
use arrow::datatypes::{DataType, Field, Schema};
@@ -33,6 +34,7 @@ use datafusion_expr::{
ColumnarValue, Expr, ExprSchemable, Operator, Volatility,
};
use datafusion_physical_expr::{create_physical_expr,
execution_props::ExecutionProps};
+use std::sync::Arc;
/// Provides simplification information based on schema and properties
pub(crate) struct SimplifyContext<'a, 'b> {
@@ -360,6 +362,9 @@ pub struct ConstEvaluator<'a> {
execution_props: &'a ExecutionProps,
input_schema: DFSchema,
input_batch: RecordBatch,
+ // Needed until we ensure type coercion is done before any optimizations
+ // https://github.com/apache/arrow-datafusion/issues/3557
+ type_coercion_helper: TypeCoercionRewriter,
}
impl<'a> ExprRewriter for ConstEvaluator<'a> {
@@ -411,16 +416,17 @@ impl<'a> ConstEvaluator<'a> {
static DUMMY_COL_NAME: &str = ".";
let schema = Schema::new(vec![Field::new(DUMMY_COL_NAME,
DataType::Null, true)]);
let input_schema = DFSchema::try_from(schema.clone())?;
-
// Need a single "input" row to produce a single output row
let col = new_null_array(&DataType::Null, 1);
let input_batch = RecordBatch::try_new(std::sync::Arc::new(schema),
vec![col])?;
+ let type_coercion =
TypeCoercionRewriter::new(Arc::new(input_schema.clone()));
Ok(Self {
can_evaluate: vec![],
execution_props,
input_schema,
input_batch,
+ type_coercion_helper: type_coercion,
})
}
@@ -484,11 +490,20 @@ impl<'a> ConstEvaluator<'a> {
}
/// Internal helper to evaluates an Expr
- pub(crate) fn evaluate_to_scalar(&self, expr: Expr) -> Result<ScalarValue>
{
+ pub(crate) fn evaluate_to_scalar(&mut self, expr: Expr) ->
Result<ScalarValue> {
if let Expr::Literal(s) = expr {
return Ok(s);
}
+ // TODO: https://github.com/apache/arrow-datafusion/issues/3582
+ // TODO: https://github.com/apache/arrow-datafusion/issues/3556
+ // Do the type coercion in the simplify expression
+ // this is just a work around for removing the type coercion in the
physical phase
+ // we need to support eval the result without the physical expr.
+ // If we don't do the type coercion, we will meet the
+ // https://github.com/apache/arrow-datafusion/issues/3556 when create
the physical expr
+ // to try evaluate the result.
+ let expr = expr.rewrite(&mut self.type_coercion_helper)?;
let phys_expr = create_physical_expr(
&expr,
&self.input_schema,
@@ -804,24 +819,28 @@ impl<'a, S: SimplifyInfo> ExprRewriter for Simplifier<'a,
S> {
//
// Rules for Between
//
+ // TODO https://github.com/apache/arrow-datafusion/issues/3587
+ // we remove between optimization temporarily, and will recover it
after above issue fixed.
+ // We should check compatibility for `expr` `low` and `high` expr
first.
+ // The rule only can work, when these three exprs can be casted to
a same data type.
// a between 3 and 5 --> a >= 3 AND a <=5
// a not between 3 and 5 --> a < 3 OR a > 5
- Between {
- expr,
- low,
- high,
- negated,
- } => {
- if negated {
- let l = *expr.clone();
- let r = *expr;
- or(l.lt(*low), r.gt(*high))
- } else {
- and(expr.clone().gt_eq(*low), expr.lt_eq(*high))
- }
- }
+ // Between {
+ // expr,
+ // low,
+ // high,
+ // negated,
+ // } => {
+ // if negated {
+ // let l = *expr.clone();
+ // let r = *expr;
+ // or(l.lt(*low), r.gt(*high))
+ // } else {
+ // and(expr.clone().gt_eq(*low), expr.lt_eq(*high))
+ // }
+ // }
expr => {
// no additional rewrites possible
expr
@@ -1107,6 +1126,12 @@ mod tests {
assert_eq!(simplify(expr_eq), lit(true));
}
+ #[test]
+ fn test_simplify_with_type_coercion() {
+ let expr_plus = binary_expr(lit(1_i32), Operator::Plus, lit(1_i64));
+ assert_eq!(simplify(expr_plus), lit(2_i64));
+ }
+
// ------------------------------
// --- ConstEvaluator tests -----
// ------------------------------
@@ -1186,7 +1211,6 @@ mod tests {
let ts_nanos = 1599566400000000000i64;
let time = chrono::Utc.timestamp_nanos(ts_nanos);
let ts_string = "2020-09-08T12:05:00+00:00";
-
// now() --> ts
test_evaluate_with_start_time(now_expr(),
lit_timestamp_nano(ts_nanos), &time);
@@ -1194,7 +1218,7 @@ mod tests {
let expr = cast_to_int64_expr(now_expr()) + lit(100);
test_evaluate_with_start_time(expr, lit(ts_nanos + 100), &time);
- // now() < cast(to_timestamp(...) as int) + 50000 ---> true
+ // CAST(now() as int64) < cast(to_timestamp(...) as int64) + 50000
---> true
let expr = cast_to_int64_expr(now_expr())
.lt(cast_to_int64_expr(to_timestamp_expr(ts_string)) + lit(50000));
test_evaluate_with_start_time(expr, lit(true), &time);
@@ -1517,7 +1541,7 @@ mod tests {
expr: None,
when_then_expr: vec![
(Box::new(col("c1")), Box::new(lit(true)),),
- (Box::new(col("c2")), Box::new(lit(false)),)
+ (Box::new(col("c2")), Box::new(lit(false)),),
],
else_expr: Some(Box::new(lit(true))),
})),
@@ -1536,7 +1560,7 @@ mod tests {
expr: None,
when_then_expr: vec![
(Box::new(col("c1")), Box::new(lit(true)),),
- (Box::new(col("c2")), Box::new(lit(false)),)
+ (Box::new(col("c2")), Box::new(lit(false)),),
],
else_expr: Some(Box::new(lit(true))),
})),
@@ -1564,6 +1588,8 @@ mod tests {
// null || false is always null
assert_eq!(simplify(lit_bool_null().or(lit(false))), lit_bool_null(),);
+ // TODO change the result
+ // https://github.com/apache/arrow-datafusion/issues/3587
// ( c1 BETWEEN Int32(0) AND Int32(10) ) OR Boolean(NULL)
// it can be either NULL or TRUE depending on the value of `c1
BETWEEN Int32(0) AND Int32(10)`
// and should not be rewritten
@@ -1573,13 +1599,15 @@ mod tests {
low: Box::new(lit(0)),
high: Box::new(lit(10)),
};
+ let between_expr = expr.clone();
let expr = expr.or(lit_bool_null());
let result = simplify(expr);
- let expected_expr = or(
- and(col("c1").gt_eq(lit(0)), col("c1").lt_eq(lit(10))),
- lit_bool_null(),
- );
+ let expected_expr = or(between_expr, lit_bool_null());
+ // let expected_expr = or(
+ // and(col("c1").gt_eq(lit(0)), col("c1").lt_eq(lit(10))),
+ // lit_bool_null(),
+ //);
assert_eq!(expected_expr, result);
}
@@ -1602,6 +1630,8 @@ mod tests {
// null && false is always false
assert_eq!(simplify(lit_bool_null().and(lit(false))), lit(false),);
+ // TODO change the result
+ // https://github.com/apache/arrow-datafusion/issues/3587
// c1 BETWEEN Int32(0) AND Int32(10) AND Boolean(NULL)
// it can be either NULL or FALSE depending on the value of `c1
BETWEEN Int32(0) AND Int32(10)`
// and the Boolean(NULL) should remain
@@ -1611,17 +1641,21 @@ mod tests {
low: Box::new(lit(0)),
high: Box::new(lit(10)),
};
+ let between_expr = expr.clone();
let expr = expr.and(lit_bool_null());
let result = simplify(expr);
- let expected_expr = and(
- and(col("c1").gt_eq(lit(0)), col("c1").lt_eq(lit(10))),
- lit_bool_null(),
- );
+ let expected_expr = and(between_expr, lit_bool_null());
+ // let expected_expr = and(
+ // and(col("c1").gt_eq(lit(0)), col("c1").lt_eq(lit(10))),
+ // lit_bool_null(),
+ // );
assert_eq!(expected_expr, result);
}
#[test]
+ #[ignore]
+ // https://github.com/apache/arrow-datafusion/issues/3587
fn simplify_expr_between() {
// c2 between 3 and 4 is c2 >= 3 and c2 <= 4
let expr = Expr::Between {
@@ -2037,7 +2071,7 @@ mod tests {
let ts_string = "2020-09-08T12:05:00+00:00";
let time = chrono::Utc.timestamp_nanos(1599566400000000000i64);
- // now() < cast(to_timestamp(...) as int) + 5000000000
+ // cast(now() as int) < cast(to_timestamp(...) as int) + 5000000000
let plan = LogicalPlanBuilder::from(table_scan)
.filter(
cast_to_int64_expr(now_expr())
@@ -2209,6 +2243,8 @@ mod tests {
}
#[test]
+ #[ignore]
+ // https://github.com/apache/arrow-datafusion/issues/3587
fn simplify_not_between() {
let table_scan = test_table_scan();
let qual = Expr::Between {
@@ -2230,6 +2266,8 @@ mod tests {
}
#[test]
+ #[ignore]
+ // https://github.com/apache/arrow-datafusion/issues/3587
fn simplify_not_not_between() {
let table_scan = test_table_scan();
let qual = Expr::Between {
diff --git a/datafusion/optimizer/src/type_coercion.rs
b/datafusion/optimizer/src/type_coercion.rs
index 0f22c01d4..bf99d61d9 100644
--- a/datafusion/optimizer/src/type_coercion.rs
+++ b/datafusion/optimizer/src/type_coercion.rs
@@ -17,7 +17,6 @@
//! Optimizer rule for type validation and coercion
-use crate::simplify_expressions::ConstEvaluator;
use crate::{OptimizerConfig, OptimizerRule};
use arrow::datatypes::DataType;
use datafusion_common::{DFSchema, DFSchemaRef, DataFusionError, Result};
@@ -30,7 +29,6 @@ use datafusion_expr::{
LogicalPlan, Operator,
};
use datafusion_expr::{ExprSchemable, Signature};
-use datafusion_physical_expr::execution_props::ExecutionProps;
use std::sync::Arc;
#[derive(Default)]
@@ -69,14 +67,8 @@ impl OptimizerRule for TypeCoercion {
},
);
- let mut execution_props = ExecutionProps::new();
- execution_props.query_execution_start_time =
- optimizer_config.query_execution_start_time();
- let const_evaluator = ConstEvaluator::try_new(&execution_props)?;
-
let mut expr_rewrite = TypeCoercionRewriter {
schema: Arc::new(schema),
- const_evaluator,
};
let original_expr_names: Vec<Option<String>> = plan
@@ -110,12 +102,17 @@ impl OptimizerRule for TypeCoercion {
}
}
-struct TypeCoercionRewriter<'a> {
- schema: DFSchemaRef,
- const_evaluator: ConstEvaluator<'a>,
+pub(crate) struct TypeCoercionRewriter {
+ pub(crate) schema: DFSchemaRef,
+}
+
+impl TypeCoercionRewriter {
+ pub(crate) fn new(schema: DFSchemaRef) -> TypeCoercionRewriter {
+ TypeCoercionRewriter { schema }
+ }
}
-impl ExprRewriter for TypeCoercionRewriter<'_> {
+impl ExprRewriter for TypeCoercionRewriter {
fn pre_visit(&mut self, _expr: &Expr) -> Result<RewriteRecursion> {
Ok(RewriteRecursion::Continue)
}
@@ -124,20 +121,20 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
match expr {
Expr::IsTrue(expr) => {
let expr = is_true(get_casted_expr_for_bool_op(&expr,
&self.schema)?);
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::IsNotTrue(expr) => {
let expr = is_not_true(get_casted_expr_for_bool_op(&expr,
&self.schema)?);
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::IsFalse(expr) => {
let expr = is_false(get_casted_expr_for_bool_op(&expr,
&self.schema)?);
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::IsNotFalse(expr) => {
let expr =
is_not_false(get_casted_expr_for_bool_op(&expr,
&self.schema)?);
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::Like {
negated,
@@ -157,7 +154,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
pattern,
escape_char,
};
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::ILike {
negated,
@@ -177,7 +174,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
pattern,
escape_char,
};
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::IsUnknown(expr) => {
// will convert the
binary(expr,IsNotDistinctFrom,lit(Boolean(None));
@@ -186,7 +183,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
let coerced_type =
coerce_types(&left_type, &Operator::IsNotDistinctFrom,
&right_type)?;
let expr = is_unknown(expr.cast_to(&coerced_type,
&self.schema)?);
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::IsNotUnknown(expr) => {
// will convert the
binary(expr,IsDistinctFrom,lit(Boolean(None));
@@ -195,7 +192,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
let coerced_type =
coerce_types(&left_type, &Operator::IsDistinctFrom,
&right_type)?;
let expr = is_not_unknown(expr.cast_to(&coerced_type,
&self.schema)?);
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::BinaryExpr {
ref left,
@@ -223,7 +220,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
right.clone().cast_to(&coerced_type,
&self.schema)?,
),
};
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
}
}
@@ -264,7 +261,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
low: Box::new(low.cast_to(&coercion_type, &self.schema)?),
high: Box::new(high.cast_to(&coercion_type,
&self.schema)?),
};
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::ScalarUDF { fun, args } => {
let new_expr = coerce_arguments_for_signature(
@@ -276,7 +273,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
fun,
args: new_expr,
};
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
Expr::InList {
expr,
@@ -309,7 +306,7 @@ impl ExprRewriter for TypeCoercionRewriter<'_> {
list: cast_list_expr,
negated,
};
- expr.rewrite(&mut self.const_evaluator)
+ Ok(expr)
}
}
}
@@ -402,7 +399,7 @@ mod test {
let mut config = OptimizerConfig::default();
let plan = rule.optimize(&plan, &mut config)?;
assert_eq!(
- "Projection: #a < Float64(2)\n EmptyRelation",
+ "Projection: #a < CAST(UInt32(2) AS Float64)\n EmptyRelation",
&format!("{:?}", plan)
);
Ok(())
@@ -430,7 +427,7 @@ mod test {
let mut config = OptimizerConfig::default();
let plan = rule.optimize(&plan, &mut config)?;
assert_eq!(
- "Projection: #a < Float64(2) OR #a < Float64(2)\
+ "Projection: #a < CAST(UInt32(2) AS Float64) OR #a <
CAST(UInt32(2) AS Float64)\
\n EmptyRelation",
&format!("{:?}", plan)
);
@@ -461,7 +458,7 @@ mod test {
let mut config = OptimizerConfig::default();
let plan = rule.optimize(&plan, &mut config)?;
assert_eq!(
- "Projection: Utf8(\"a\")\n EmptyRelation",
+ "Projection: TestScalarUDF(CAST(Int32(123) AS Float32))\n
EmptyRelation",
&format!("{:?}", plan)
);
Ok(())
@@ -540,7 +537,7 @@ mod test {
let mut config = OptimizerConfig::default();
let plan = rule.optimize(&plan, &mut config)?;
assert_eq!(
- "Projection: #a IN ([Int64(1), Int64(4), Int64(8)])\n
EmptyRelation",
+ "Projection: #a IN ([CAST(Int32(1) AS Int64), CAST(Int8(4) AS
Int64), Int64(8)])\n EmptyRelation",
&format!("{:?}", plan)
);
// a in (1,4,8), a is decimal
@@ -558,7 +555,7 @@ mod test {
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr],
empty, None)?);
let plan = rule.optimize(&plan, &mut config)?;
assert_eq!(
- "Projection: CAST(#a AS Decimal128(24, 4)) IN
([Decimal128(Some(10000),24,4), Decimal128(Some(40000),24,4),
Decimal128(Some(80000),24,4)])\n EmptyRelation",
+ "Projection: CAST(#a AS Decimal128(24, 4)) IN ([CAST(Int32(1) AS
Decimal128(24, 4)), CAST(Int8(4) AS Decimal128(24, 4)), CAST(Int64(8) AS
Decimal128(24, 4))])\n EmptyRelation",
&format!("{:?}", plan)
);
Ok(())
@@ -656,7 +653,7 @@ mod test {
let mut config = OptimizerConfig::default();
let plan = rule.optimize(&plan, &mut config).unwrap();
assert_eq!(
- "Projection: #a LIKE Utf8(NULL)\n EmptyRelation",
+ "Projection: #a LIKE CAST(NULL AS Utf8)\n EmptyRelation",
&format!("{:?}", plan)
);
diff --git a/datafusion/optimizer/tests/integration-test.rs
b/datafusion/optimizer/tests/integration-test.rs
index cb171780c..554e3cceb 100644
--- a/datafusion/optimizer/tests/integration-test.rs
+++ b/datafusion/optimizer/tests/integration-test.rs
@@ -79,12 +79,13 @@ fn intersect() -> Result<()> {
#[test]
fn between_date32_plus_interval() -> Result<()> {
+ // TODO: https://github.com/apache/arrow-datafusion/issues/3587
let sql = "SELECT count(1) FROM test \
WHERE col_date32 between '1998-03-18' AND cast('1998-03-18' as date) +
INTERVAL '90 days'";
let plan = test_sql(sql)?;
let expected =
"Projection: #COUNT(UInt8(1))\n Aggregate: groupBy=[[]],
aggr=[[COUNT(UInt8(1))]]\
- \n Filter: #test.col_date32 >= Date32(\"10303\") AND
#test.col_date32 <= Date32(\"10393\")\
+ \n Filter: #test.col_date32 BETWEEN Date32(\"10303\") AND
Date32(\"10393\")\
\n TableScan: test projection=[col_date32]";
assert_eq!(expected, format!("{:?}", plan));
Ok(())
@@ -92,18 +93,21 @@ fn between_date32_plus_interval() -> Result<()> {
#[test]
fn between_date64_plus_interval() -> Result<()> {
+ // TODO: https://github.com/apache/arrow-datafusion/issues/3587
let sql = "SELECT count(1) FROM test \
WHERE col_date64 between '1998-03-18T00:00:00' AND cast('1998-03-18' as
date) + INTERVAL '90 days'";
let plan = test_sql(sql)?;
let expected =
"Projection: #COUNT(UInt8(1))\n Aggregate: groupBy=[[]],
aggr=[[COUNT(UInt8(1))]]\
- \n Filter: #test.col_date64 >= Date64(\"890179200000\") AND
#test.col_date64 <= Date64(\"897955200000\")\
+ \n Filter: #test.col_date64 BETWEEN Date64(\"890179200000\") AND
Date64(\"897955200000\")\
\n TableScan: test projection=[col_date64]";
assert_eq!(expected, format!("{:?}", plan));
Ok(())
}
fn test_sql(sql: &str) -> Result<LogicalPlan> {
+ // TODO should make align with rules in the context
+ // https://github.com/apache/arrow-datafusion/issues/3524
let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
// Simplify expressions first to maximize the chance
// of applying other optimizations
@@ -121,8 +125,10 @@ fn test_sql(sql: &str) -> Result<LogicalPlan> {
Arc::new(RewriteDisjunctivePredicate::new()),
Arc::new(FilterNullJoinKeys::default()),
Arc::new(ReduceOuterJoin::new()),
- Arc::new(FilterPushDown::new()),
Arc::new(TypeCoercion::new()),
+ // after the type coercion, can do simplify expression again
+ Arc::new(SimplifyExpressions::new()),
+ Arc::new(FilterPushDown::new()),
Arc::new(LimitPushDown::new()),
Arc::new(SingleDistinctToGroupBy::new()),
];
diff --git a/datafusion/physical-expr/src/expressions/binary.rs
b/datafusion/physical-expr/src/expressions/binary.rs
index b3704dc70..02bf0e5bd 100644
--- a/datafusion/physical-expr/src/expressions/binary.rs
+++ b/datafusion/physical-expr/src/expressions/binary.rs
@@ -73,12 +73,11 @@ use kernels_arrow::{
use arrow::datatypes::{DataType, Schema, TimeUnit};
use arrow::record_batch::RecordBatch;
-use crate::expressions::try_cast;
use crate::PhysicalExpr;
use datafusion_common::ScalarValue;
use datafusion_common::{DataFusionError, Result};
use datafusion_expr::binary_rule::binary_operator_data_type;
-use datafusion_expr::{binary_rule::coerce_types, ColumnarValue, Operator};
+use datafusion_expr::{ColumnarValue, Operator};
/// Binary expression
#[derive(Debug)]
@@ -912,25 +911,6 @@ impl BinaryExpr {
}
}
-/// return two physical expressions that are optionally coerced to a
-/// common type that the binary operator supports.
-fn binary_cast(
- lhs: Arc<dyn PhysicalExpr>,
- op: &Operator,
- rhs: Arc<dyn PhysicalExpr>,
- input_schema: &Schema,
-) -> Result<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> {
- let lhs_type = &lhs.data_type(input_schema)?;
- let rhs_type = &rhs.data_type(input_schema)?;
-
- let result_type = coerce_types(lhs_type, op, rhs_type)?;
-
- Ok((
- try_cast(lhs, input_schema, result_type.clone())?,
- try_cast(rhs, input_schema, result_type)?,
- ))
-}
-
/// Create a binary expression whose arguments are correctly coerced.
/// This function errors if it is not possible to coerce the arguments
/// to computational types supported by the operator.
@@ -940,17 +920,25 @@ pub fn binary(
rhs: Arc<dyn PhysicalExpr>,
input_schema: &Schema,
) -> Result<Arc<dyn PhysicalExpr>> {
- let (l, r) = binary_cast(lhs, &op, rhs, input_schema)?;
- Ok(Arc::new(BinaryExpr::new(l, op, r)))
+ let lhs_type = &lhs.data_type(input_schema)?;
+ let rhs_type = &rhs.data_type(input_schema)?;
+ if !lhs_type.eq(rhs_type) {
+ return Err(DataFusionError::Internal(format!(
+ "The type of {} {} {} of binary physical should be same",
+ lhs_type, op, rhs_type
+ )));
+ }
+ Ok(Arc::new(BinaryExpr::new(lhs, op, rhs)))
}
#[cfg(test)]
mod tests {
use super::*;
+ use crate::expressions::try_cast;
use crate::expressions::{col, lit};
use arrow::datatypes::{ArrowNumericType, Field, Int32Type, SchemaRef};
- use arrow::util::display::array_value_to_string;
use datafusion_common::Result;
+ use datafusion_expr::binary_rule::coerce_types;
// Create a binary expression without coercion. Used here when we do not
want to coerce the expressions
// to valid types. Usage can result in an execution (after plan) error.
@@ -1051,17 +1039,20 @@ mod tests {
// 4. verify that the resulting expression is of type C
// 5. verify that the results of evaluation are $VEC
macro_rules! test_coercion {
- ($A_ARRAY:ident, $A_TYPE:expr, $A_VEC:expr, $B_ARRAY:ident,
$B_TYPE:expr, $B_VEC:expr, $OP:expr, $C_ARRAY:ident, $C_TYPE:expr, $VEC:expr)
=> {{
+ ($A_ARRAY:ident, $A_TYPE:expr, $A_VEC:expr, $B_ARRAY:ident,
$B_TYPE:expr, $B_VEC:expr, $OP:expr, $C_ARRAY:ident, $C_TYPE:expr, $VEC:expr,)
=> {{
let schema = Schema::new(vec![
Field::new("a", $A_TYPE, false),
Field::new("b", $B_TYPE, false),
]);
let a = $A_ARRAY::from($A_VEC);
let b = $B_ARRAY::from($B_VEC);
+ let result_type = coerce_types(&$A_TYPE, &$OP, &$B_TYPE)?;
+
+ let left = try_cast(col("a", &schema)?, &schema,
result_type.clone())?;
+ let right = try_cast(col("b", &schema)?, &schema, result_type)?;
// verify that we can construct the expression
- let expression =
- binary(col("a", &schema)?, $OP, col("b", &schema)?, &schema)?;
+ let expression = binary(left, $OP, right, &schema)?;
let batch = RecordBatch::try_new(
Arc::new(schema.clone()),
vec![Arc::new(a), Arc::new(b)],
@@ -1100,7 +1091,7 @@ mod tests {
Operator::Plus,
Int32Array,
DataType::Int32,
- vec![2i32, 4i32]
+ vec![2i32, 4i32],
);
test_coercion!(
Int32Array,
@@ -1112,7 +1103,7 @@ mod tests {
Operator::Plus,
Int32Array,
DataType::Int32,
- vec![2i32]
+ vec![2i32],
);
test_coercion!(
Float32Array,
@@ -1124,7 +1115,7 @@ mod tests {
Operator::Plus,
Float32Array,
DataType::Float32,
- vec![2f32]
+ vec![2f32],
);
test_coercion!(
Float32Array,
@@ -1136,7 +1127,7 @@ mod tests {
Operator::Multiply,
Float32Array,
DataType::Float32,
- vec![2f32]
+ vec![2f32],
);
test_coercion!(
StringArray,
@@ -1148,7 +1139,7 @@ mod tests {
Operator::Like,
BooleanArray,
DataType::Boolean,
- vec![true, false]
+ vec![true, false],
);
test_coercion!(
StringArray,
@@ -1160,7 +1151,7 @@ mod tests {
Operator::Eq,
BooleanArray,
DataType::Boolean,
- vec![true, true]
+ vec![true, true],
);
test_coercion!(
StringArray,
@@ -1172,7 +1163,7 @@ mod tests {
Operator::Lt,
BooleanArray,
DataType::Boolean,
- vec![true, false]
+ vec![true, false],
);
test_coercion!(
StringArray,
@@ -1184,7 +1175,7 @@ mod tests {
Operator::Eq,
BooleanArray,
DataType::Boolean,
- vec![true, true]
+ vec![true, true],
);
test_coercion!(
StringArray,
@@ -1196,7 +1187,7 @@ mod tests {
Operator::Lt,
BooleanArray,
DataType::Boolean,
- vec![true, false]
+ vec![true, false],
);
test_coercion!(
StringArray,
@@ -1208,7 +1199,7 @@ mod tests {
Operator::RegexMatch,
BooleanArray,
DataType::Boolean,
- vec![true, false, true, false, false]
+ vec![true, false, true, false, false],
);
test_coercion!(
StringArray,
@@ -1220,7 +1211,7 @@ mod tests {
Operator::RegexIMatch,
BooleanArray,
DataType::Boolean,
- vec![true, true, true, true, false]
+ vec![true, true, true, true, false],
);
test_coercion!(
StringArray,
@@ -1232,7 +1223,7 @@ mod tests {
Operator::RegexNotMatch,
BooleanArray,
DataType::Boolean,
- vec![false, true, false, true, true]
+ vec![false, true, false, true, true],
);
test_coercion!(
StringArray,
@@ -1244,7 +1235,7 @@ mod tests {
Operator::RegexNotIMatch,
BooleanArray,
DataType::Boolean,
- vec![false, false, false, false, true]
+ vec![false, false, false, false, true],
);
test_coercion!(
LargeStringArray,
@@ -1256,7 +1247,7 @@ mod tests {
Operator::RegexMatch,
BooleanArray,
DataType::Boolean,
- vec![true, false, true, false, false]
+ vec![true, false, true, false, false],
);
test_coercion!(
LargeStringArray,
@@ -1268,7 +1259,7 @@ mod tests {
Operator::RegexIMatch,
BooleanArray,
DataType::Boolean,
- vec![true, true, true, true, false]
+ vec![true, true, true, true, false],
);
test_coercion!(
LargeStringArray,
@@ -1280,7 +1271,7 @@ mod tests {
Operator::RegexNotMatch,
BooleanArray,
DataType::Boolean,
- vec![false, true, false, true, true]
+ vec![false, true, false, true, true],
);
test_coercion!(
LargeStringArray,
@@ -1292,7 +1283,7 @@ mod tests {
Operator::RegexNotIMatch,
BooleanArray,
DataType::Boolean,
- vec![false, false, false, false, true]
+ vec![false, false, false, false, true],
);
test_coercion!(
Int16Array,
@@ -1304,7 +1295,7 @@ mod tests {
Operator::BitwiseAnd,
Int64Array,
DataType::Int64,
- vec![0i64, 0i64, 1i64]
+ vec![0i64, 0i64, 1i64],
);
test_coercion!(
Int16Array,
@@ -1316,7 +1307,7 @@ mod tests {
Operator::BitwiseOr,
Int64Array,
DataType::Int64,
- vec![11i64, 6i64, 7i64]
+ vec![11i64, 6i64, 7i64],
);
test_coercion!(
Int16Array,
@@ -1328,7 +1319,7 @@ mod tests {
Operator::BitwiseXor,
Int64Array,
DataType::Int64,
- vec![9i64, 4i64, 6i64]
+ vec![9i64, 4i64, 6i64],
);
Ok(())
}
@@ -1352,72 +1343,36 @@ mod tests {
dict_builder.append_null();
dict_builder.append("three")?;
dict_builder.append("four")?;
- let dict_array = dict_builder.finish();
+ let dict_array = Arc::new(dict_builder.finish()) as ArrayRef;
- let str_array =
- StringArray::from(vec![Some("not one"), Some("two"), None,
Some("four")]);
+ let str_array = Arc::new(StringArray::from(vec![
+ Some("not one"),
+ Some("two"),
+ None,
+ Some("four"),
+ ])) as ArrayRef;
let schema = Arc::new(Schema::new(vec![
- Field::new("dict", dict_type, true),
- Field::new("str", string_type, true),
+ Field::new("a", dict_type.clone(), true),
+ Field::new("b", string_type.clone(), true),
]));
- let batch = RecordBatch::try_new(
- schema.clone(),
- vec![Arc::new(dict_array), Arc::new(str_array)],
- )?;
-
- let expected = "false\n\n\ntrue";
-
- // Test 1: dict = str
-
- // verify that we can construct the expression
- let expression = binary(
- col("dict", &schema)?,
- Operator::Eq,
- col("str", &schema)?,
- &schema,
- )?;
- assert_eq!(expression.data_type(&schema)?, DataType::Boolean);
-
- // evaluate and verify the result type matched
- let result = expression.evaluate(&batch)?.into_array(batch.num_rows());
- assert_eq!(result.data_type(), &DataType::Boolean);
-
- // verify that the result itself is correct
- assert_eq!(expected, array_to_string(&result)?);
+ // Test 1: a = b
+ let result = BooleanArray::from(vec![Some(false), None, None,
Some(true)]);
+ apply_logic_op(&schema, &dict_array, &str_array, Operator::Eq,
result)?;
// Test 2: now test the other direction
- // str = dict
-
- // verify that we can construct the expression
- let expression = binary(
- col("str", &schema)?,
- Operator::Eq,
- col("dict", &schema)?,
- &schema,
- )?;
- assert_eq!(expression.data_type(&schema)?, DataType::Boolean);
-
- // evaluate and verify the result type matched
- let result = expression.evaluate(&batch)?.into_array(batch.num_rows());
- assert_eq!(result.data_type(), &DataType::Boolean);
-
- // verify that the result itself is correct
- assert_eq!(expected, array_to_string(&result)?);
+ // b = a
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", string_type, true),
+ Field::new("b", dict_type, true),
+ ]));
+ let result = BooleanArray::from(vec![Some(false), None, None,
Some(true)]);
+ apply_logic_op(&schema, &str_array, &dict_array, Operator::Eq,
result)?;
Ok(())
}
- // Convert the array to a newline delimited string of pretty printed values
- fn array_to_string(array: &ArrayRef) -> Result<String> {
- let s = (0..array.len())
- .map(|i| array_value_to_string(array, i))
- .collect::<std::result::Result<Vec<_>,
arrow::error::ArrowError>>()?
- .join("\n");
- Ok(s)
- }
-
#[test]
fn plus_op() -> Result<()> {
let schema = Schema::new(vec![
@@ -1543,8 +1498,13 @@ mod tests {
op: Operator,
expected: BooleanArray,
) -> Result<()> {
- let arithmetic_op =
- binary_simple(col("a", schema)?, op, col("b", schema)?, schema);
+ let left_type = left.data_type();
+ let right_type = right.data_type();
+ let result_type = coerce_types(left_type, &op, right_type)?;
+
+ let left_expr = try_cast(col("a", schema)?, schema,
result_type.clone())?;
+ let right_expr = try_cast(col("b", schema)?, schema, result_type)?;
+ let arithmetic_op = binary_simple(left_expr, op, right_expr, schema);
let data: Vec<ArrayRef> = vec![left.clone(), right.clone()];
let batch = RecordBatch::try_new(schema.clone(), data)?;
let result =
arithmetic_op.evaluate(&batch)?.into_array(batch.num_rows());
@@ -1562,8 +1522,19 @@ mod tests {
expected: &BooleanArray,
) -> Result<()> {
let scalar = lit(scalar.clone());
+ let op_type = coerce_types(&scalar.data_type(schema)?, &op,
arr.data_type())?;
+ let left_expr = if op_type.eq(&scalar.data_type(schema)?) {
+ scalar
+ } else {
+ try_cast(scalar, schema, op_type.clone())?
+ };
+ let right_expr = if op_type.eq(arr.data_type()) {
+ col("a", schema)?
+ } else {
+ try_cast(col("a", schema)?, schema, op_type)?
+ };
- let arithmetic_op = binary_simple(scalar, op, col("a", schema)?,
schema);
+ let arithmetic_op = binary_simple(left_expr, op, right_expr, schema);
let batch = RecordBatch::try_new(Arc::clone(schema),
vec![Arc::clone(arr)])?;
let result =
arithmetic_op.evaluate(&batch)?.into_array(batch.num_rows());
assert_eq!(result.as_ref(), expected);
@@ -1580,8 +1551,19 @@ mod tests {
expected: &BooleanArray,
) -> Result<()> {
let scalar = lit(scalar.clone());
+ let op_type = coerce_types(arr.data_type(), &op,
&scalar.data_type(schema)?)?;
+ let right_expr = if op_type.eq(&scalar.data_type(schema)?) {
+ scalar
+ } else {
+ try_cast(scalar, schema, op_type.clone())?
+ };
+ let left_expr = if op_type.eq(arr.data_type()) {
+ col("a", schema)?
+ } else {
+ try_cast(col("a", schema)?, schema, op_type)?
+ };
- let arithmetic_op = binary_simple(col("a", schema)?, op, scalar,
schema);
+ let arithmetic_op = binary_simple(left_expr, op, right_expr, schema);
let batch = RecordBatch::try_new(Arc::clone(schema),
vec![Arc::clone(arr)])?;
let result =
arithmetic_op.evaluate(&batch)?.into_array(batch.num_rows());
assert_eq!(result.as_ref(), expected);
@@ -2405,8 +2387,19 @@ mod tests {
op: Operator,
expected: ArrayRef,
) -> Result<()> {
- let arithmetic_op =
- binary_simple(col("a", schema)?, op, col("b", schema)?, schema);
+ let op_type = coerce_types(left.data_type(), &op, right.data_type())?;
+ let left_expr = if left.data_type().eq(&op_type) {
+ col("a", schema)?
+ } else {
+ try_cast(col("a", schema)?, schema, op_type.clone())?
+ };
+
+ let right_expr = if right.data_type().eq(&op_type) {
+ col("b", schema)?
+ } else {
+ try_cast(col("b", schema)?, schema, op_type)?
+ };
+ let arithmetic_op = binary_simple(left_expr, op, right_expr, schema);
let data: Vec<ArrayRef> = vec![left.clone(), right.clone()];
let batch = RecordBatch::try_new(schema.clone(), data)?;
let result =
arithmetic_op.evaluate(&batch)?.into_array(batch.num_rows());
@@ -2487,6 +2480,7 @@ mod tests {
expect,
)
.unwrap();
+
// divide: int32 array divide decimal array
let schema = Arc::new(Schema::new(vec