You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ag...@apache.org on 2023/03/21 18:38:16 UTC
[arrow-datafusion] branch main updated: Add cast expression with bool, integers and decimal128 support (#5137)
This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 63acd5773 Add cast expression with bool, integers and decimal128 support (#5137)
63acd5773 is described below
commit 63acd57736524728370617e20af0b538a20b6864
Author: Nuttiiya Seekhao <37...@users.noreply.github.com>
AuthorDate: Tue Mar 21 14:38:08 2023 -0400
Add cast expression with bool, integers and decimal128 support (#5137)
Address comment on unwrap
---
datafusion/substrait/src/logical_plan/consumer.rs | 42 ++++++++++++-
datafusion/substrait/src/logical_plan/producer.rs | 69 +++++++++++++++++++++-
.../substrait/tests/roundtrip_logical_plan.rs | 10 ++++
3 files changed, 119 insertions(+), 2 deletions(-)
diff --git a/datafusion/substrait/src/logical_plan/consumer.rs b/datafusion/substrait/src/logical_plan/consumer.rs
index 27f8093f7..b7cacf131 100644
--- a/datafusion/substrait/src/logical_plan/consumer.rs
+++ b/datafusion/substrait/src/logical_plan/consumer.rs
@@ -16,12 +16,13 @@
// under the License.
use async_recursion::async_recursion;
+use datafusion::arrow::datatypes::DataType;
use datafusion::common::{DFField, DFSchema, DFSchemaRef};
-use datafusion::logical_expr::expr;
use datafusion::logical_expr::{
aggregate_function, BinaryExpr, Case, Expr, LogicalPlan, Operator,
};
use datafusion::logical_expr::{build_join_schema, LogicalPlanBuilder};
+use datafusion::logical_expr::{expr, Cast};
use datafusion::prelude::JoinType;
use datafusion::sql::TableReference;
use datafusion::{
@@ -721,12 +722,51 @@ pub async fn from_substrait_rex(
))),
}
}
+ Some(RexType::Cast(cast)) => match cast.as_ref().r#type.as_ref() {
+ Some(output_type) => Ok(Arc::new(Expr::Cast(Cast::new(
+ Box::new(
+ from_substrait_rex(
+ cast.as_ref().input.as_ref().unwrap().as_ref(),
+ input_schema,
+ extensions,
+ )
+ .await?
+ .as_ref()
+ .clone(),
+ ),
+ from_substrait_type(output_type)?,
+ )))),
+ None => Err(DataFusionError::Substrait(
+ "Cast experssion without output type is not allowed".to_string(),
+ )),
+ },
_ => Err(DataFusionError::NotImplemented(
"unsupported rex_type".to_string(),
)),
}
}
+fn from_substrait_type(dt: &substrait::proto::Type) -> Result<DataType> {
+ match &dt.kind {
+ Some(s_kind) => match s_kind {
+ r#type::Kind::Bool(_) => Ok(DataType::Boolean),
+ r#type::Kind::I8(_) => Ok(DataType::Int8),
+ r#type::Kind::I16(_) => Ok(DataType::Int16),
+ r#type::Kind::I32(_) => Ok(DataType::Int32),
+ r#type::Kind::I64(_) => Ok(DataType::Int64),
+ r#type::Kind::Decimal(d) => {
+ Ok(DataType::Decimal128(d.precision as u8, d.scale as i8))
+ }
+ _ => Err(DataFusionError::NotImplemented(format!(
+ "Unsupported Substrait type: {s_kind:?}"
+ ))),
+ },
+ _ => Err(DataFusionError::NotImplemented(
+ "`None` Substrait kind is not supported".to_string(),
+ )),
+ }
+}
+
fn from_substrait_null(null_type: &Type) -> Result<ScalarValue> {
if let Some(kind) = &null_type.kind {
match kind {
diff --git a/datafusion/substrait/src/logical_plan/producer.rs b/datafusion/substrait/src/logical_plan/producer.rs
index f891cf6cc..1de26a343 100644
--- a/datafusion/substrait/src/logical_plan/producer.rs
+++ b/datafusion/substrait/src/logical_plan/producer.rs
@@ -18,6 +18,7 @@
use std::{collections::HashMap, mem, sync::Arc};
use datafusion::{
+ arrow::datatypes::DataType,
error::{DataFusionError, Result},
prelude::JoinType,
scalar::ScalarValue,
@@ -26,7 +27,7 @@ use datafusion::{
use datafusion::common::DFSchemaRef;
#[allow(unused_imports)]
use datafusion::logical_expr::aggregate_function;
-use datafusion::logical_expr::expr::{BinaryExpr, Case, Sort};
+use datafusion::logical_expr::expr::{BinaryExpr, Case, Cast, Sort};
use datafusion::logical_expr::{expr, Between, JoinConstraint, LogicalPlan, Operator};
use datafusion::prelude::{binary_expr, Expr};
use substrait::proto::{
@@ -577,6 +578,21 @@ pub fn to_substrait_rex(
rex_type: Some(RexType::IfThen(Box::new(IfThen { ifs, r#else }))),
})
}
+ Expr::Cast(Cast { expr, data_type }) => {
+ Ok(Expression {
+ rex_type: Some(RexType::Cast(Box::new(
+ substrait::proto::expression::Cast {
+ r#type: Some(to_substrait_type(data_type)?),
+ input: Some(Box::new(to_substrait_rex(
+ expr,
+ schema,
+ extension_info,
+ )?)),
+ failure_behavior: 0, // FAILURE_BEHAVIOR_UNSPECIFIED
+ },
+ ))),
+ })
+ }
Expr::Literal(value) => {
let literal_type = match value {
ScalarValue::Int8(Some(n)) => Some(LiteralType::I8(*n as i32)),
@@ -626,6 +642,57 @@ pub fn to_substrait_rex(
}
}
+fn to_substrait_type(dt: &DataType) -> Result<substrait::proto::Type> {
+ let default_type_ref = 0;
+ let default_nullability = r#type::Nullability::Required as i32;
+ match dt {
+ DataType::Null => Err(DataFusionError::Internal(
+ "Null cast is not valid".to_string(),
+ )),
+ DataType::Boolean => Ok(substrait::proto::Type {
+ kind: Some(r#type::Kind::Bool(r#type::Boolean {
+ type_variation_reference: default_type_ref,
+ nullability: default_nullability,
+ })),
+ }),
+ DataType::Int8 => Ok(substrait::proto::Type {
+ kind: Some(r#type::Kind::I8(r#type::I8 {
+ type_variation_reference: default_type_ref,
+ nullability: default_nullability,
+ })),
+ }),
+ DataType::Int16 => Ok(substrait::proto::Type {
+ kind: Some(r#type::Kind::I16(r#type::I16 {
+ type_variation_reference: default_type_ref,
+ nullability: default_nullability,
+ })),
+ }),
+ DataType::Int32 => Ok(substrait::proto::Type {
+ kind: Some(r#type::Kind::I32(r#type::I32 {
+ type_variation_reference: default_type_ref,
+ nullability: default_nullability,
+ })),
+ }),
+ DataType::Int64 => Ok(substrait::proto::Type {
+ kind: Some(r#type::Kind::I64(r#type::I64 {
+ type_variation_reference: default_type_ref,
+ nullability: default_nullability,
+ })),
+ }),
+ DataType::Decimal128(p, s) => Ok(substrait::proto::Type {
+ kind: Some(r#type::Kind::Decimal(r#type::Decimal {
+ type_variation_reference: default_type_ref,
+ nullability: default_nullability,
+ scale: *s as i32,
+ precision: *p as i32,
+ })),
+ }),
+ _ => Err(DataFusionError::NotImplemented(format!(
+ "Unsupported cast type: {dt:?}"
+ ))),
+ }
+}
+
fn try_to_substrait_null(v: &ScalarValue) -> Result<LiteralType> {
let default_type_ref = 0;
let default_nullability = r#type::Nullability::Nullable as i32;
diff --git a/datafusion/substrait/tests/roundtrip_logical_plan.rs b/datafusion/substrait/tests/roundtrip_logical_plan.rs
index 85b3ef19f..9aa430bb0 100644
--- a/datafusion/substrait/tests/roundtrip_logical_plan.rs
+++ b/datafusion/substrait/tests/roundtrip_logical_plan.rs
@@ -179,6 +179,16 @@ mod tests {
.await
}
+ #[tokio::test]
+ async fn cast_decimal_to_int() -> Result<()> {
+ roundtrip("SELECT * FROM data WHERE a = CAST(2.5 AS int)").await
+ }
+
+ #[tokio::test]
+ async fn implicit_cast() -> Result<()> {
+ roundtrip("SELECT * FROM data WHERE a = b").await
+ }
+
#[tokio::test]
async fn aggregate_case() -> Result<()> {
assert_expected_plan(