Skip to content

fix: fold cast null to substrait typed null #15854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions datafusion/substrait/src/logical_plan/producer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1590,6 +1590,24 @@ pub fn from_cast(
schema: &DFSchemaRef,
) -> Result<Expression> {
let Cast { expr, data_type } = cast;
// since substrait Null must be typed, so if we see a cast(null, dt), we make it a typed null
if let Expr::Literal(lit) = expr.as_ref() {
// only the untyped(a null scalar value) null literal need this special handling
// since all other kind of nulls are already typed and can be handled by substrait
// e.g. null::<Int32Type> or null::<Utf8Type>
if matches!(lit, ScalarValue::Null) {
let lit = Literal {
nullable: true,
type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
literal_type: Some(LiteralType::Null(to_substrait_type(
data_type, true,
)?)),
};
return Ok(Expression {
rex_type: Some(RexType::Literal(lit)),
});
}
}
Ok(Expression {
rex_type: Some(RexType::Cast(Box::new(
substrait::proto::expression::Cast {
Expand Down Expand Up @@ -2575,6 +2593,7 @@ mod test {
use datafusion::common::scalar::ScalarStructBuilder;
use datafusion::common::DFSchema;
use datafusion::execution::{SessionState, SessionStateBuilder};
use datafusion::logical_expr::ExprSchemable;
use datafusion::prelude::SessionContext;
use std::sync::LazyLock;

Expand Down Expand Up @@ -2912,4 +2931,70 @@ mod test {

assert!(matches!(err, Err(DataFusionError::SchemaError(_, _))));
}

#[tokio::test]
async fn fold_cast_null() {
let state = SessionStateBuilder::default().build();
let empty_schema = DFSchemaRef::new(DFSchema::empty());
let field = Field::new("out", DataType::Int32, false);

let expr = Expr::Literal(ScalarValue::Null)
.cast_to(&DataType::Int32, &empty_schema)
.unwrap();

let typed_null =
to_substrait_extended_expr(&[(&expr, &field)], &empty_schema, &state)
.unwrap();

if let ExprType::Expression(expr) =
typed_null.referred_expr[0].expr_type.as_ref().unwrap()
{
let lit = Literal {
nullable: true,
type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
literal_type: Some(LiteralType::Null(
to_substrait_type(&DataType::Int32, true).unwrap(),
)),
};
let expected = Expression {
rex_type: Some(RexType::Literal(lit)),
};
assert_eq!(*expr, expected);
} else {
panic!("Expected expression type");
}

// a typed null should not be folded
let expr = Expr::Literal(ScalarValue::Int64(None))
.cast_to(&DataType::Int32, &empty_schema)
.unwrap();

let typed_null =
to_substrait_extended_expr(&[(&expr, &field)], &empty_schema, &state)
.unwrap();

if let ExprType::Expression(expr) =
typed_null.referred_expr[0].expr_type.as_ref().unwrap()
{
let cast_expr = substrait::proto::expression::Cast {
r#type: Some(to_substrait_type(&DataType::Int32, true).unwrap()),
input: Some(Box::new(Expression {
rex_type: Some(RexType::Literal(Literal {
nullable: true,
type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
literal_type: Some(LiteralType::Null(
to_substrait_type(&DataType::Int64, true).unwrap(),
)),
})),
})),
failure_behavior: FailureBehavior::ThrowException as i32,
};
let expected = Expression {
rex_type: Some(RexType::Cast(Box::new(cast_expr))),
};
assert_eq!(*expr, expected);
} else {
panic!("Expected expression type");
}
}
}
Loading