Skip to content

Commit 3467011

Browse files
authored
[bugfix] ScalarFunctionExpr does not preserve the nullable flag on roundtrip (#13830)
* [test] coalesce round trip schema mismatch * [proto] added the nullable flag in PhysicalScalarUdfNode * [bugfix] propagate the nullable flag for serialized scalar UDFS
1 parent 2439979 commit 3467011

File tree

6 files changed

+72
-7
lines changed

6 files changed

+72
-7
lines changed

datafusion/proto/proto/datafusion.proto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -838,6 +838,7 @@ message PhysicalScalarUdfNode {
838838
repeated PhysicalExprNode args = 2;
839839
optional bytes fun_definition = 3;
840840
datafusion_common.ArrowType return_type = 4;
841+
bool nullable = 5;
841842
}
842843

843844
message PhysicalAggregateExprNode {

datafusion/proto/src/generated/pbjson.rs

Lines changed: 17 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

datafusion/proto/src/generated/prost.rs

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

datafusion/proto/src/physical_plan/from_proto.rs

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -358,12 +358,15 @@ pub fn parse_physical_expr(
358358

359359
let args = parse_physical_exprs(&e.args, registry, input_schema, codec)?;
360360

361-
Arc::new(ScalarFunctionExpr::new(
362-
e.name.as_str(),
363-
scalar_fun_def,
364-
args,
365-
convert_required!(e.return_type)?,
366-
))
361+
Arc::new(
362+
ScalarFunctionExpr::new(
363+
e.name.as_str(),
364+
scalar_fun_def,
365+
args,
366+
convert_required!(e.return_type)?,
367+
)
368+
.with_nullable(e.nullable),
369+
)
367370
}
368371
ExprType::LikeExpr(like_expr) => Arc::new(LikeExpr::new(
369372
like_expr.negated,

datafusion/proto/src/physical_plan/to_proto.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,7 @@ pub fn serialize_physical_expr(
347347
args: serialize_physical_exprs(expr.args(), codec)?,
348348
fun_definition: (!buf.is_empty()).then_some(buf),
349349
return_type: Some(expr.return_type().try_into()?),
350+
nullable: expr.nullable(),
350351
},
351352
)),
352353
})

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use crate::cases::{
3939
use datafusion::arrow::array::ArrayRef;
4040
use datafusion::arrow::compute::kernels::sort::SortOptions;
4141
use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema};
42+
use datafusion::datasource::empty::EmptyTable;
4243
use datafusion::datasource::file_format::csv::CsvSink;
4344
use datafusion::datasource::file_format::json::JsonSink;
4445
use datafusion::datasource::file_format::parquet::ParquetSink;
@@ -83,7 +84,7 @@ use datafusion::physical_plan::windows::{
8384
WindowAggExec,
8485
};
8586
use datafusion::physical_plan::{
86-
ExecutionPlan, InputOrderMode, Partitioning, PhysicalExpr, Statistics,
87+
displayable, ExecutionPlan, InputOrderMode, Partitioning, PhysicalExpr, Statistics,
8788
};
8889
use datafusion::prelude::SessionContext;
8990
use datafusion::scalar::ScalarValue;
@@ -106,6 +107,7 @@ use datafusion_proto::physical_plan::{
106107
AsExecutionPlan, DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
107108
};
108109
use datafusion_proto::protobuf;
110+
use datafusion_proto::protobuf::PhysicalPlanNode;
109111

110112
/// Perform a serde roundtrip and assert that the string representation of the before and after plans
111113
/// are identical. Note that this often isn't sufficient to guarantee that no information is
@@ -1525,3 +1527,42 @@ fn roundtrip_unnest() -> Result<()> {
15251527
);
15261528
roundtrip_test(Arc::new(unnest))
15271529
}
1530+
1531+
#[tokio::test]
1532+
async fn roundtrip_coalesce() -> Result<()> {
1533+
let ctx = SessionContext::new();
1534+
ctx.register_table(
1535+
"t",
1536+
Arc::new(EmptyTable::new(Arc::new(Schema::new(Fields::from([
1537+
Arc::new(Field::new("f", DataType::Int64, false)),
1538+
]))))),
1539+
)?;
1540+
let df = ctx.sql("select coalesce(f) as f from t").await?;
1541+
let plan = df.create_physical_plan().await?;
1542+
1543+
let node = PhysicalPlanNode::try_from_physical_plan(
1544+
plan.clone(),
1545+
&DefaultPhysicalExtensionCodec {},
1546+
)?;
1547+
let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice())
1548+
.map_err(|e| DataFusionError::External(Box::new(e)))?;
1549+
let restored = node.try_into_physical_plan(
1550+
&ctx,
1551+
ctx.runtime_env().as_ref(),
1552+
&DefaultPhysicalExtensionCodec {},
1553+
)?;
1554+
1555+
assert_eq!(
1556+
plan.schema(),
1557+
restored.schema(),
1558+
"Schema mismatch for plans:\n>> initial:\n{}>> final: \n{}",
1559+
displayable(plan.as_ref())
1560+
.set_show_schema(true)
1561+
.indent(true),
1562+
displayable(restored.as_ref())
1563+
.set_show_schema(true)
1564+
.indent(true),
1565+
);
1566+
1567+
Ok(())
1568+
}

0 commit comments

Comments
 (0)