Skip to content

refactor: Port CaseExpr proto serialization hooks #22838

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.

Sign up for GitHub

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

Merged
adriangb merged 2 commits into from
Jun 9, 2026
Merged

refactor: Port CaseExpr proto serialization hooks #22838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9e46942
refactor: Port CaseExpr proto serialization hooks
chakkk309 Jun 9, 2026
7162b2e
test: Move CaseExpr proto tests into proto module
chakkk309 Jun 9, 2026
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
254 changes: 254 additions & 0 deletions datafusion/physical-expr/src/expressions/case.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1410,6 +1410,86 @@ impl PhysicalExpr for CaseExpr {
}
write!(f, "END")
}

#[cfg(feature = "proto")]
fn try_to_proto(
&self,
ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
use datafusion_proto_models::protobuf;

Ok(Some(protobuf::PhysicalExprNode {
expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Case(Box::new(
protobuf::PhysicalCaseNode {
expr: self
.expr()
.map(|expr| ctx.encode_child(expr).map(Box::new))
.transpose()?,
when_then_expr: self
.when_then_expr()
.iter()
.map(|(when_expr, then_expr)| {
Ok(protobuf::PhysicalWhenThen {
when_expr: Some(ctx.encode_child(when_expr)?),
then_expr: Some(ctx.encode_child(then_expr)?),
})
})
.collect::<Result<Vec<_>>>()?,
else_expr: self
.else_expr()
.map(|expr| ctx.encode_child(expr).map(Box::new))
.transpose()?,
},
))),
}))
}
}

#[cfg(feature = "proto")]
impl CaseExpr {
/// Reconstruct a [`CaseExpr`] from its protobuf representation.
pub fn try_from_proto(
node: &datafusion_proto_models::protobuf::PhysicalExprNode,
ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
) -> Result<Arc<dyn PhysicalExpr>> {
use datafusion_physical_expr_common::expect_expr_variant;
use datafusion_proto_models::protobuf;

let case = expect_expr_variant!(
node,
protobuf::physical_expr_node::ExprType::Case,
"CaseExpr",
);

Ok(Arc::new(CaseExpr::try_new(
case.expr
.as_deref()
.map(|expr| ctx.decode(expr))
.transpose()?,
case.when_then_expr
.iter()
.map(|when_then| {
Ok((
ctx.decode_required_expression(
when_then.when_expr.as_ref(),
"CaseExpr",
"when_expr",
)?,
ctx.decode_required_expression(
when_then.then_expr.as_ref(),
"CaseExpr",
"then_expr",
)?,
))
})
.collect::<Result<Vec<_>>>()?,
case.else_expr
.as_deref()
.map(|expr| ctx.decode(expr))
.transpose()?,
)?))
}
}

/// Attempts to const evaluate the given `predicate`.
Expand Down Expand Up @@ -3193,3 +3273,177 @@ mod tests {
Ok(())
}
}

#[cfg(all(test, feature = "proto"))]
mod proto_tests {
use super::*;
use crate::expressions::col;
use crate::proto_test_util::{
StubDecoder, StubEncoder, UnreachableDecoder, column_node,
};
use arrow::datatypes::Field;
use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
use datafusion_proto_models::protobuf;
use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalWhenThen};

fn proto_case_fixture() -> CaseExpr {
let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]);
CaseExpr::try_new(
Some(col("a", &schema).unwrap()),
vec![(lit(true), lit(1_i32))],
Some(lit(0_i32)),
)
.unwrap()
}

fn proto_when_then(
when_expr: Option<PhysicalExprNode>,
then_expr: Option<PhysicalExprNode>,
) -> PhysicalWhenThen {
PhysicalWhenThen {
when_expr,
then_expr,
}
}

fn proto_case_node(
expr: Option<Box<PhysicalExprNode>>,
when_then_expr: Vec<PhysicalWhenThen>,
else_expr: Option<Box<PhysicalExprNode>>,
) -> PhysicalExprNode {
PhysicalExprNode {
expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Case(Box::new(
protobuf::PhysicalCaseNode {
expr,
when_then_expr,
else_expr,
},
))),
}
}

#[test]
fn try_to_proto_encodes_case_expr() {
let case = proto_case_fixture();
let encoder = StubEncoder::ok();
let ctx = PhysicalExprEncodeCtx::new(&encoder);

let node = case
.try_to_proto(&ctx)
.unwrap()
.expect("CaseExpr should encode to Some(node)");

assert!(node.expr_id.is_none());
let case_node = match node.expr_type {
Some(protobuf::physical_expr_node::ExprType::Case(boxed)) => *boxed,
other => panic!("expected a CaseExpr node, got {other:?}"),
};
assert!(case_node.expr.is_some());
assert_eq!(case_node.when_then_expr.len(), 1);
assert!(case_node.when_then_expr[0].when_expr.is_some());
assert!(case_node.when_then_expr[0].then_expr.is_some());
assert!(case_node.else_expr.is_some());
}

#[test]
fn try_to_proto_propagates_child_encode_error() {
let case = proto_case_fixture();
// Call 1 is the optional CASE expr, call 2 is the WHEN expr.
let encoder = StubEncoder::failing_on(2);
let ctx = PhysicalExprEncodeCtx::new(&encoder);

let err = case.try_to_proto(&ctx).unwrap_err();
assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2")));
}

#[test]
fn try_from_proto_decodes_case_expr() {
let node = proto_case_node(
Some(Box::new(column_node("case"))),
vec![proto_when_then(
Some(column_node("when")),
Some(column_node("then")),
)],
Some(Box::new(column_node("else"))),
);
let schema = Schema::empty();
let decoder = StubDecoder::ok();
let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);

let decoded = CaseExpr::try_from_proto(&node, &ctx).unwrap();
let case = decoded
.downcast_ref::<CaseExpr>()
.expect("decoded expr should be a CaseExpr");

assert!(case.expr().is_some());
assert_eq!(case.when_then_expr().len(), 1);
assert!(case.else_expr().is_some());
}

#[test]
fn try_from_proto_rejects_non_case_node() {
let node = column_node("a");
let schema = Schema::empty();
let decoder = UnreachableDecoder;
let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);

let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err();
assert!(
matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a CaseExpr"))
);
}

#[test]
fn try_from_proto_rejects_missing_when_expr() {
let node = proto_case_node(
None,
vec![proto_when_then(None, Some(column_node("then")))],
None,
);
let schema = Schema::empty();
let decoder = UnreachableDecoder;
let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);

let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err();
assert!(
matches!(err, DataFusionError::Internal(msg) if msg.contains("CaseExpr is missing required field 'when_expr'"))
);
}

#[test]
fn try_from_proto_rejects_missing_then_expr() {
let node = proto_case_node(
None,
vec![proto_when_then(Some(column_node("when")), None)],
None,
);
let schema = Schema::empty();
let decoder = StubDecoder::ok();
let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);

let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err();
assert!(
matches!(err, DataFusionError::Internal(msg) if msg.contains("CaseExpr is missing required field 'then_expr'"))
);
}

#[test]
fn try_from_proto_propagates_child_decode_error() {
let node = proto_case_node(
Some(Box::new(column_node("case"))),
vec![proto_when_then(
Some(column_node("when")),
Some(column_node("then")),
)],
Some(Box::new(column_node("else"))),
);
let schema = Schema::empty();
let decoder = StubDecoder::failing_on(2);
let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);

let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err();
assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2")));
}
}
36 changes: 1 addition & 35 deletions datafusion/proto/src/physical_plan/from_proto.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,41 +303,7 @@ pub fn parse_physical_expr_with_converter(
ExprType::NotExpr(_) => NotExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::Negative(_) => NegativeExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::InList(_) => InListExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::Case(e) => Arc::new(CaseExpr::try_new(
e.expr
.as_ref()
.map(|e| {
proto_converter.proto_to_physical_expr(e.as_ref(), input_schema, ctx)
})
.transpose()?,
e.when_then_expr
.iter()
.map(|e| {
Ok((
parse_required_physical_expr(
e.when_expr.as_ref(),
ctx,
"when_expr",
input_schema,
proto_converter,
)?,
parse_required_physical_expr(
e.then_expr.as_ref(),
ctx,
"then_expr",
input_schema,
proto_converter,
)?,
))
})
.collect::<Result<Vec<_>>>()?,
e.else_expr
.as_ref()
.map(|e| {
proto_converter.proto_to_physical_expr(e.as_ref(), input_schema, ctx)
})
.transpose()?,
)?),
ExprType::Case(_) => CaseExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::TryCast(_) => TryCastExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::ScalarUdf(e) => {
Expand Down
59 changes: 2 additions & 57 deletions datafusion/proto/src/physical_plan/to_proto.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use datafusion_physical_expr::ScalarFunctionExpr;
use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr;
use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr};
use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
use datafusion_physical_plan::expressions::{CaseExpr, DynamicFilterPhysicalExpr};
use datafusion_physical_plan::expressions::DynamicFilterPhysicalExpr;
use datafusion_physical_plan::udaf::AggregateFunctionExpr;
use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr};
use datafusion_physical_plan::{
Expand Down Expand Up @@ -300,50 +300,7 @@ pub fn serialize_physical_expr_with_converter(
return Ok(node);
}

if let Some(expr) = expr.downcast_ref::<CaseExpr>() {
Ok(protobuf::PhysicalExprNode {
expr_id,
expr_type: Some(
protobuf::physical_expr_node::ExprType::Case(
Box::new(
protobuf::PhysicalCaseNode {
expr: expr
.expr()
.map(|exp| {
proto_converter
.physical_expr_to_proto(exp, codec)
.map(Box::new)
})
.transpose()?,
when_then_expr: expr
.when_then_expr()
.iter()
.map(|(when_expr, then_expr)| {
serialize_when_then_expr(
when_expr,
then_expr,
codec,
proto_converter,
)
})
.collect::<Result<
Vec<protobuf::PhysicalWhenThen>,
DataFusionError,
>>()?,
else_expr: expr
.else_expr()
.map(|a| {
proto_converter
.physical_expr_to_proto(a, codec)
.map(Box::new)
})
.transpose()?,
},
),
),
),
})
} else if let Some(expr) = expr.downcast_ref::<ScalarFunctionExpr>() {
if let Some(expr) = expr.downcast_ref::<ScalarFunctionExpr>() {
let mut buf = Vec::new();
codec.try_encode_udf(expr.fun(), &mut buf)?;
Ok(protobuf::PhysicalExprNode {
Expand Down Expand Up @@ -500,18 +457,6 @@ fn serialize_range_split_point(
})
}

fn serialize_when_then_expr(
when_expr: &Arc<dyn PhysicalExpr>,
then_expr: &Arc<dyn PhysicalExpr>,
codec: &dyn PhysicalExtensionCodec,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<protobuf::PhysicalWhenThen> {
Ok(protobuf::PhysicalWhenThen {
when_expr: Some(proto_converter.physical_expr_to_proto(when_expr, codec)?),
then_expr: Some(proto_converter.physical_expr_to_proto(then_expr, codec)?),
})
}

impl TryFromProto<&PartitionedFile> for protobuf::PartitionedFile {
type Error = DataFusionError;

Expand Down
Loading