Skip to content

Commit 39697af

Browse files
james-willisjames-willis
committed
fix: keep CASE field metadata through the other CASE simplifications
A CASE with a constant WHEN now keeps the CASE when dropping a branch would change its field metadata, and a literal result takes the CASE's metadata. A Boolean CASE whose branch has metadata is no longer rewritten to AND/OR, which cannot carry it.
1 parent 8ef3402 commit 39697af

3 files changed

Lines changed: 220 additions & 50 deletions

File tree

‎datafusion/expr/src/expr_schema.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,8 @@ impl ExprSchemable for Expr {
507507
/// - **Aliases**: Merge underlying expr metadata with alias-specific metadata, preferring the alias metadata
508508
/// - **List unnesting**: Preserve the element field metadata
509509
/// - **Binary expressions**: field metadata is empty
510+
/// - **CASE expressions**: Merges the metadata of the THEN and ELSE branches; the operand
511+
/// and WHEN conditions do not contribute
510512
/// - **Boolean expressions**: field metadata is empty
511513
/// - **Cast expressions**: Type-only casts pass through source metadata (stripping extension
512514
/// type keys); casts with explicit target fields use target metadata exactly

‎datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs‎

Lines changed: 160 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1482,22 +1482,15 @@ impl TreeNodeRewriter for Simplifier<'_> {
14821482
.iter()
14831483
.any(|(when, _)| is_true(when.as_ref()) || is_false(when.as_ref())) =>
14841484
{
1485-
// CASE WHEN false THEN A END --> NULL
1486-
// The NULL keeps the CASE's field metadata, such as an extension type
1487-
if else_expr.is_none()
1488-
&& when_then_expr
1489-
.iter()
1490-
.all(|(when, _)| is_false(when.as_ref()))
1491-
{
1492-
let (_, field) = Expr::Case(Case::new(None, when_then_expr, None))
1493-
.to_field(info.schema().as_ref())?;
1494-
let metadata = FieldMetadata::new_from_field(&field);
1495-
let null = Expr::Literal(
1496-
ScalarValue::try_new_null(field.data_type())?,
1497-
(!metadata.is_empty()).then_some(metadata),
1498-
);
1499-
return Ok(Transformed::yes(null));
1500-
}
1485+
// Dropping a branch must not change the field metadata of the CASE,
1486+
// which merges the metadata of all its value branches
1487+
let original = Expr::Case(Case::new(
1488+
None,
1489+
when_then_expr.clone(),
1490+
else_expr.clone(),
1491+
));
1492+
let schema = info.schema().as_ref();
1493+
let (_, field) = original.to_field(schema)?;
15011494

15021495
let mut new_when_then_expr = Vec::with_capacity(when_then_expr.len());
15031496

@@ -1514,18 +1507,37 @@ impl TreeNodeRewriter for Simplifier<'_> {
15141507
}
15151508

15161509
// Exclude CASE statement altogether if there are no when-then expressions left
1517-
// CASE WHEN false THEN A ELSE B END --> B
1518-
if new_when_then_expr.is_empty()
1519-
&& let Some(else_expr) = else_expr
1520-
{
1521-
return Ok(Transformed::yes(*else_expr));
1522-
}
1510+
let simplified = match (new_when_then_expr.is_empty(), else_expr) {
1511+
// CASE WHEN false THEN A END --> NULL
1512+
(true, None) => {
1513+
Expr::Literal(ScalarValue::try_new_null(field.data_type())?, None)
1514+
}
1515+
// CASE WHEN false THEN A ELSE B END --> B
1516+
(true, Some(else_expr)) => *else_expr,
1517+
(false, else_expr) => Expr::Case(Case {
1518+
expr: None,
1519+
when_then_expr: new_when_then_expr,
1520+
else_expr,
1521+
}),
1522+
};
15231523

1524-
Transformed::yes(Expr::Case(Case {
1525-
expr: None,
1526-
when_then_expr: new_when_then_expr,
1527-
else_expr,
1528-
}))
1524+
match simplified {
1525+
// A literal takes the CASE's field metadata, such as an extension type
1526+
Expr::Literal(value, _) => {
1527+
let metadata = FieldMetadata::new_from_field(&field);
1528+
Transformed::yes(Expr::Literal(
1529+
value,
1530+
(!metadata.is_empty()).then_some(metadata),
1531+
))
1532+
}
1533+
simplified
1534+
if simplified.to_field(schema)?.1.metadata()
1535+
== field.metadata() =>
1536+
{
1537+
Transformed::yes(simplified)
1538+
}
1539+
_ => Transformed::no(original),
1540+
}
15291541
}
15301542

15311543
// CASE
@@ -1552,7 +1564,8 @@ impl TreeNodeRewriter for Simplifier<'_> {
15521564
|| (when_then_expr.iter().all(|(_, then)| is_bool_lit(then))
15531565
&& when_then_expr.iter().filter(|(_, then)| is_true(then)).count() < 3))
15541566
&& info.is_boolean_type(&when_then_expr[0].1)?
1555-
&& can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref()) =>
1567+
&& can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref())
1568+
&& !case_branch_has_metadata(&when_then_expr, else_expr.as_deref(), info)? =>
15561569
{
15571570
// String disjunction of all the when predicates encountered so far. Not nullable.
15581571
let mut filter_expr = lit(false);
@@ -1610,7 +1623,8 @@ impl TreeNodeRewriter for Simplifier<'_> {
16101623
.count()
16111624
< 3
16121625
&& else_expr.as_deref().is_none_or(is_bool_lit)
1613-
&& can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref()) =>
1626+
&& can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref())
1627+
&& !case_branch_has_metadata(&when_then_expr, else_expr.as_deref(), info)? =>
16141628
{
16151629
Transformed::yes(
16161630
Expr::Case(Case {
@@ -2408,6 +2422,25 @@ fn can_lower_case_to_boolean(
24082422
&& else_expr.is_none_or(is_leaf)
24092423
}
24102424

2425+
/// Whether a value branch of a CASE has field metadata, which AND/OR cannot keep
2426+
fn case_branch_has_metadata(
2427+
when_then_expr: &[(Box<Expr>, Box<Expr>)],
2428+
else_expr: Option<&Expr>,
2429+
info: &SimplifyContext,
2430+
) -> Result<bool> {
2431+
let branches = when_then_expr
2432+
.iter()
2433+
.map(|(_, then)| then.as_ref())
2434+
.chain(else_expr);
2435+
for expr in branches {
2436+
let (_, field) = expr.to_field(info.schema().as_ref())?;
2437+
if !field.metadata().is_empty() {
2438+
return Ok(true);
2439+
}
2440+
}
2441+
Ok(false)
2442+
}
2443+
24112444
/// Returns expression testing a boolean `expr` for being exactly `true` (not `false` or NULL).
24122445
fn is_exactly_true(expr: Expr, info: &SimplifyContext) -> Result<Expr> {
24132446
if !info.nullable(&expr)? {
@@ -4601,28 +4634,114 @@ mod tests {
46014634
}
46024635

46034636
#[test]
4604-
fn simplify_expr_case_when_false_keeps_field_metadata() -> Result<()> {
4637+
fn simplify_expr_case_keeps_field_metadata() -> Result<()> {
46054638
let metadata = FieldMetadata::from(HashMap::from([(
46064639
"ARROW:extension:name".to_string(),
46074640
"a.type".to_string(),
46084641
)]));
46094642
let schema = Arc::new(DFSchema::from_unqualified_fields(
4610-
vec![metadata.add_to_field(Field::new("ext", DataType::Utf8, true))].into(),
4643+
vec![
4644+
metadata.add_to_field(Field::new("ext", DataType::Utf8, true)),
4645+
Field::new("plain", DataType::Utf8, true),
4646+
metadata.add_to_field(Field::new("flag", DataType::Boolean, true)),
4647+
Field::new("b1", DataType::Boolean, true),
4648+
Field::new("b2", DataType::Boolean, true),
4649+
Field::new("b3", DataType::Boolean, true),
4650+
]
4651+
.into(),
46114652
HashMap::new(),
46124653
)?);
46134654
let simplifier =
46144655
ExprSimplifier::new(SimplifyContext::builder().with_schema(schema).build());
46154656

4616-
// CASE WHEN false THEN ext END --> NULL with ext's extension type
4617-
let expr = Expr::Case(Case::new(
4618-
None,
4619-
vec![(Box::new(lit(false)), Box::new(col("ext")))],
4620-
None,
4621-
));
4622-
assert_eq!(
4623-
simplifier.simplify(expr)?,
4624-
Expr::Literal(ScalarValue::Utf8(None), Some(metadata))
4625-
);
4657+
// A CASE that simplifies to a literal gives it the CASE's metadata
4658+
let to_literal = [
4659+
// CASE WHEN false THEN ext END --> NULL
4660+
(
4661+
Case::new(
4662+
None,
4663+
vec![(Box::new(lit(false)), Box::new(col("ext")))],
4664+
None,
4665+
),
4666+
ScalarValue::Utf8(None),
4667+
),
4668+
// CASE WHEN false THEN ext ELSE NULL END --> NULL
4669+
(
4670+
Case::new(
4671+
None,
4672+
vec![(Box::new(lit(false)), Box::new(col("ext")))],
4673+
Some(Box::new(Expr::Literal(ScalarValue::Utf8(None), None))),
4674+
),
4675+
ScalarValue::Utf8(None),
4676+
),
4677+
// CASE WHEN true THEN 'x' ELSE ext END --> 'x'
4678+
(
4679+
Case::new(
4680+
None,
4681+
vec![(Box::new(lit(true)), Box::new(lit("x")))],
4682+
Some(Box::new(col("ext"))),
4683+
),
4684+
ScalarValue::from("x"),
4685+
),
4686+
];
4687+
for (case, expected) in to_literal {
4688+
assert_eq!(
4689+
simplifier.simplify(Expr::Case(case))?,
4690+
Expr::Literal(expected, Some(metadata.clone()))
4691+
);
4692+
}
4693+
4694+
// Rewrites that would drop the metadata of a branch leave the CASE as is
4695+
let unchanged = [
4696+
// CASE WHEN false THEN ext ELSE plain END
4697+
Case::new(
4698+
None,
4699+
vec![(Box::new(lit(false)), Box::new(col("ext")))],
4700+
Some(Box::new(col("plain"))),
4701+
),
4702+
// CASE WHEN true THEN plain ELSE ext END
4703+
Case::new(
4704+
None,
4705+
vec![(Box::new(lit(true)), Box::new(col("plain")))],
4706+
Some(Box::new(col("ext"))),
4707+
),
4708+
// CASE WHEN b1 THEN plain WHEN false THEN ext END
4709+
Case::new(
4710+
None,
4711+
vec![
4712+
(Box::new(col("b1")), Box::new(col("plain"))),
4713+
(Box::new(lit(false)), Box::new(col("ext"))),
4714+
],
4715+
None,
4716+
),
4717+
// CASE WHEN b1 THEN flag END, which AND/OR cannot express with metadata
4718+
Case::new(
4719+
None,
4720+
vec![(Box::new(col("b1")), Box::new(col("flag")))],
4721+
None,
4722+
),
4723+
// CASE WHEN b1 THEN true WHEN b2 THEN true WHEN b3 THEN true ELSE false END,
4724+
// with metadata on the first THEN, which NOT(CASE ...) cannot keep
4725+
Case::new(
4726+
None,
4727+
vec![
4728+
(
4729+
Box::new(col("b1")),
4730+
Box::new(Expr::Literal(
4731+
ScalarValue::Boolean(Some(true)),
4732+
Some(metadata.clone()),
4733+
)),
4734+
),
4735+
(Box::new(col("b2")), Box::new(lit(true))),
4736+
(Box::new(col("b3")), Box::new(lit(true))),
4737+
],
4738+
Some(Box::new(lit(false))),
4739+
),
4740+
];
4741+
for case in unchanged {
4742+
let expr = Expr::Case(case);
4743+
assert_eq!(simplifier.simplify(expr.clone())?, expr);
4744+
}
46264745
Ok(())
46274746
}
46284747

‎datafusion/sqllogictest/test_files/metadata.slt‎

Lines changed: 58 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -793,17 +793,66 @@ arrow.json arrow.json arrow.json arrow.json arrow.json
793793
arrow.json arrow.json arrow.json arrow.json arrow.json
794794
arrow.json arrow.json arrow.json arrow.json arrow.json
795795

796-
# A CASE whose WHEN is always false simplifies to NULL, which keeps the
797-
# extension type of the CASE
798-
query T
799-
select arrow_metadata(
800-
case when 1 = 0 then cast(name as json) end,
801-
'ARROW:extension:name')
796+
# A CASE with a constant WHEN keeps its extension type when it is simplified:
797+
# - no ELSE, or a NULL ELSE: the CASE becomes a NULL with the extension type
798+
# - a branch that is dropped has the extension type: the CASE is kept
799+
query TTTTT
800+
select
801+
arrow_metadata(
802+
case when 1 = 0 then cast(name as json) end,
803+
'ARROW:extension:name'),
804+
arrow_metadata(
805+
case when 1 = 0 then cast(name as json) else NULL end,
806+
'ARROW:extension:name'),
807+
arrow_metadata(
808+
case when 1 = 0 then cast(name as json) else upper(l_name) end,
809+
'ARROW:extension:name'),
810+
arrow_metadata(
811+
case when 1 = 1 then upper(l_name) else cast(name as json) end,
812+
'ARROW:extension:name'),
813+
arrow_metadata(
814+
case when id > 1 then upper(l_name) when 1 = 0 then cast(name as json) end,
815+
'ARROW:extension:name')
802816
from table_with_metadata;
803817
----
804-
arrow.json
805-
arrow.json
806-
arrow.json
818+
arrow.json arrow.json arrow.json arrow.json arrow.json
819+
arrow.json arrow.json arrow.json arrow.json arrow.json
820+
arrow.json arrow.json arrow.json arrow.json arrow.json
821+
822+
# The same applies to coalesce and nvl2, whose first argument here is never
823+
# NULL, so their CASE always takes its first branch
824+
query TT
825+
select
826+
arrow_metadata(
827+
coalesce(nonnull_name, cast(name as json)),
828+
'ARROW:extension:name'),
829+
arrow_metadata(
830+
nvl2(nonnull_name, upper(l_name), cast(name as json)),
831+
'ARROW:extension:name')
832+
from table_with_metadata;
833+
----
834+
arrow.json arrow.json
835+
arrow.json arrow.json
836+
arrow.json arrow.json
837+
838+
# A Boolean CASE is not rewritten to AND/OR when a branch has metadata
839+
statement ok
840+
create table bool_with_metadata as
841+
select
842+
id > 1 as p,
843+
with_metadata(id < 3, 'ARROW:extension:name', 'a.type') as f
844+
from table_with_metadata;
845+
846+
query T
847+
select arrow_metadata(case when p then f end, 'ARROW:extension:name')
848+
from bool_with_metadata;
849+
----
850+
a.type
851+
a.type
852+
a.type
853+
854+
statement ok
855+
drop table bool_with_metadata;
807856

808857
statement ok
809858
drop table table_with_metadata;

0 commit comments

Comments
 (0)