Skip to content

Commit 8ef3402

Browse files
james-willisjames-willis
committed
test: check that the physical CASE field matches the logical one
Compare Expr::Case::to_field with the physical return_field, including a branch cast to the common type by type coercion. Use the built-in JSON extension type in the CASE metadata.slt cases, add a coercion case, and move the CASE WHEN false case next to the other CASE cases.
1 parent 9216db0 commit 8ef3402

2 files changed

Lines changed: 155 additions & 64 deletions

File tree

‎datafusion/physical-expr/src/planner.rs‎

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -735,7 +735,7 @@ mod tests {
735735
use datafusion_expr::physical_planning_context::{
736736
ScalarSubqueryResults, SubqueryIndex,
737737
};
738-
use datafusion_expr::{LogicalPlanBuilder, col, scalar_subquery};
738+
use datafusion_expr::{LogicalPlanBuilder, case, col, scalar_subquery, when};
739739

740740
use super::*;
741741

@@ -1017,6 +1017,90 @@ mod tests {
10171017
Ok(())
10181018
}
10191019

1020+
/// The physical CASE reports the same field metadata as the logical CASE
1021+
#[test]
1022+
fn test_case_lowering_preserves_field_metadata() -> Result<()> {
1023+
let json = Metadata::new()
1024+
.with(EXTENSION_TYPE_NAME_KEY, "arrow.json")
1025+
.with(EXTENSION_TYPE_METADATA_KEY, "");
1026+
let schema = Schema::new(vec![
1027+
Field::new("id", DataType::Int32, true),
1028+
Field::new("json", DataType::Utf8, true).with_metadata(json.clone()),
1029+
Field::new("name", DataType::Utf8, true)
1030+
.with_metadata(Metadata::new().with("source", "name")),
1031+
Field::new("plain", DataType::Utf8, true),
1032+
Field::new("large", DataType::LargeUtf8, true),
1033+
]);
1034+
let df_schema = DFSchema::try_from(schema.clone())?;
1035+
let null_with_metadata = Expr::Literal(
1036+
ScalarValue::Utf8(None),
1037+
Some(FieldMetadata::from(HashMap::from([(
1038+
"source".to_string(),
1039+
"null_else".to_string(),
1040+
)]))),
1041+
);
1042+
1043+
let cases = [
1044+
(
1045+
when(col("id").gt(lit(1)), col("json")).otherwise(col("plain"))?,
1046+
json.clone(),
1047+
),
1048+
(
1049+
when(col("id").gt(lit(1)), col("name")).otherwise(col("json"))?,
1050+
json.clone().with("source", "name"),
1051+
),
1052+
(
1053+
case(col("id")).when(lit(1), col("json")).end()?,
1054+
json.clone(),
1055+
),
1056+
(
1057+
when(col("id").gt(lit(1)), col("plain")).otherwise(null_with_metadata)?,
1058+
Metadata::new().with("source", "null_else"),
1059+
),
1060+
(
1061+
when(
1062+
col("id").gt(lit(1)),
1063+
when(col("id").gt(lit(2)), col("json")).end()?,
1064+
)
1065+
.otherwise(col("plain"))?,
1066+
json.clone(),
1067+
),
1068+
// Type coercion casts a branch to the common type, which drops its
1069+
// extension type and keeps its other metadata
1070+
(
1071+
when(
1072+
col("id").gt(lit(1)),
1073+
Expr::Cast(Cast::new(Box::new(col("json")), DataType::LargeUtf8)),
1074+
)
1075+
.otherwise(col("large"))?,
1076+
Metadata::new(),
1077+
),
1078+
(
1079+
when(
1080+
col("id").gt(lit(1)),
1081+
Expr::Cast(Cast::new(Box::new(col("name")), DataType::LargeUtf8)),
1082+
)
1083+
.otherwise(col("large"))?,
1084+
Metadata::new().with("source", "name"),
1085+
),
1086+
];
1087+
1088+
for (expr, expected) in cases {
1089+
let (_, logical) = expr.to_field(&df_schema)?;
1090+
let physical = create_physical_expr(
1091+
&expr,
1092+
&df_schema,
1093+
&ExecutionProps::new(),
1094+
&PhysicalPlanningContext::default(),
1095+
)?
1096+
.return_field(&schema)?;
1097+
assert_eq!(logical.metadata(), &expected, "logical field of {expr}");
1098+
assert_eq!(physical.metadata(), &expected, "physical field of {expr}");
1099+
}
1100+
1101+
Ok(())
1102+
}
1103+
10201104
/// Test that deeply nested expressions do not cause a stack overflow.
10211105
///
10221106
/// This test only runs when the `recursive_protection` feature is enabled,

‎datafusion/sqllogictest/test_files/metadata.slt‎

Lines changed: 70 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -616,18 +616,6 @@ DROP TABLE min_max_with_metadata;
616616
statement ok
617617
RESET datafusion.execution.parquet.skip_metadata;
618618

619-
# A CASE whose WHEN is always false simplifies to NULL, which keeps the
620-
# extension type of the CASE
621-
query T
622-
select arrow_metadata(
623-
case when 1 = 0 then with_metadata(name, 'ARROW:extension:name', 'a.type') end,
624-
'ARROW:extension:name')
625-
from table_with_metadata;
626-
----
627-
a.type
628-
a.type
629-
a.type
630-
631619
# coalesce, nvl, ifnull and nvl2: conflicting extension types are a planning error
632620
statement error Error during planning: Cannot merge field metadata with differing values of 'ARROW:extension:name': 'a\.type' and 'b\.type'
633621
select coalesce(
@@ -708,19 +696,17 @@ select
708696
arrow_metadata(case when id > 1 then name else l_name end, 'metadata_key'),
709697
arrow_metadata(case id when 1 then name else l_name end, 'metadata_key'),
710698
arrow_metadata(
711-
case when id > 1
712-
then with_metadata(name, 'ARROW:extension:name', 'a.type')
713-
else with_metadata(l_name, 'ARROW:extension:name', 'a.type') end,
699+
case when id > 1 then cast(name as json) else cast(l_name as json) end,
714700
'ARROW:extension:name')
715701
from table_with_metadata;
716702
----
717-
the name field the name field a.type
718-
the name field the name field a.type
719-
the name field the name field a.type
703+
the name field the name field arrow.json
704+
the name field the name field arrow.json
705+
the name field the name field arrow.json
720706

721-
statement error Error during planning: Cannot merge field metadata with differing values of 'ARROW:extension:name': 'a\.type' and 'b\.type'
707+
statement error Error during planning: Cannot merge field metadata with differing values of 'ARROW:extension:name': 'arrow\.json' and 'b\.type'
722708
select case when id > 1
723-
then with_metadata(name, 'ARROW:extension:name', 'a.type')
709+
then cast(name as json)
724710
else with_metadata(l_name, 'ARROW:extension:name', 'b.type') end
725711
from table_with_metadata;
726712

@@ -729,74 +715,95 @@ from table_with_metadata;
729715
query TTTT
730716
select
731717
arrow_metadata(
732-
case when id > 1 then upper(name)
733-
else with_metadata(l_name, 'ARROW:extension:name', 'a.type') end,
718+
case when id > 1 then upper(name) else cast(l_name as json) end,
734719
'ARROW:extension:name'),
720+
arrow_metadata(coalesce(upper(name), cast(l_name as json)), 'ARROW:extension:name'),
735721
arrow_metadata(
736-
coalesce(upper(name), with_metadata(l_name, 'ARROW:extension:name', 'a.type')),
722+
case when id > 1 then cast(name as json) else NULL end,
737723
'ARROW:extension:name'),
724+
arrow_metadata(case when id > 1 then cast(name as json) end, 'ARROW:extension:name')
725+
from table_with_metadata;
726+
----
727+
arrow.json arrow.json arrow.json arrow.json
728+
arrow.json arrow.json arrow.json arrow.json
729+
arrow.json arrow.json arrow.json arrow.json
730+
731+
# A typed NULL ELSE contributes its metadata, including when no row reaches a
732+
# THEN branch. Each query uses a CASE only once: a repeated expression is
733+
# computed once in a lower projection, whose field comes from the logical plan.
734+
query TT
735+
select
738736
arrow_metadata(
739-
case when id > 1 then with_metadata(name, 'ARROW:extension:name', 'a.type')
740-
else NULL end,
737+
case when id > 100 then name else cast(NULL as json) end,
741738
'ARROW:extension:name'),
742739
arrow_metadata(
743-
case when id > 1 then with_metadata(name, 'ARROW:extension:name', 'a.type') end,
740+
case when id > 100 then upper(name) else cast(NULL as json) end,
744741
'ARROW:extension:name')
745742
from table_with_metadata;
746743
----
747-
a.type a.type a.type a.type
748-
a.type a.type a.type a.type
749-
a.type a.type a.type a.type
744+
arrow.json arrow.json
745+
arrow.json arrow.json
746+
arrow.json arrow.json
750747

751-
# A typed NULL ELSE with metadata contributes its metadata, including when no
752-
# row reaches a THEN branch
753-
query TTTT
748+
query TT
754749
select
755-
case when id > 100 then name
756-
else with_metadata(arrow_cast(NULL, 'Utf8'), 'source', 'null_else') end,
750+
case when id > 100 then name else cast(NULL as json) end,
751+
case when id > 100 then upper(name) else cast(NULL as json) end
752+
from table_with_metadata;
753+
----
754+
NULL NULL
755+
NULL NULL
756+
NULL NULL
757+
758+
# Coercing the branches to a common type casts a branch to a plain storage type,
759+
# which drops its extension type and keeps its other metadata
760+
query TTT
761+
select
762+
arrow_typeof(
763+
case when id > 1 then cast(name as json) else arrow_cast(l_name, 'LargeUtf8') end),
757764
arrow_metadata(
758-
case when id > 100 then name
759-
else with_metadata(arrow_cast(NULL, 'Utf8'), 'source', 'null_else') end,
760-
'source'),
761-
case when id > 100 then upper(name)
762-
else with_metadata(arrow_cast(NULL, 'Utf8'), 'source', 'null_else') end,
765+
case when id > 1 then cast(name as json) else arrow_cast(l_name, 'LargeUtf8') end,
766+
'ARROW:extension:name'),
763767
arrow_metadata(
764-
case when id > 100 then upper(name)
765-
else with_metadata(arrow_cast(NULL, 'Utf8'), 'source', 'null_else') end,
766-
'source')
768+
case when id > 1 then name else arrow_cast(l_name, 'LargeUtf8') end,
769+
'metadata_key')
767770
from table_with_metadata;
768771
----
769-
NULL null_else NULL null_else
770-
NULL null_else NULL null_else
771-
NULL null_else NULL null_else
772+
LargeUtf8 NULL the name field
773+
LargeUtf8 NULL the name field
774+
LargeUtf8 NULL the name field
772775

773776
# Nested CASE, and the CASE rewrites of coalesce, nvl, ifnull and nvl2, keep
774-
# the extension type at execution
777+
# the extension type at execution. The rewrites take different arguments so
778+
# that they are not one repeated expression.
775779
query TTTTT
776780
select
777781
arrow_metadata(
778782
case when id > 1
779-
then case when id > 2
780-
then with_metadata(name, 'ARROW:extension:name', 'a.type') end
783+
then case when id > 2 then cast(name as json) end
781784
else l_name end,
782785
'ARROW:extension:name'),
783-
arrow_metadata(
784-
coalesce(with_metadata(name, 'ARROW:extension:name', 'a.type'), l_name),
785-
'ARROW:extension:name'),
786-
arrow_metadata(
787-
nvl(with_metadata(name, 'ARROW:extension:name', 'a.type'), l_name),
788-
'ARROW:extension:name'),
789-
arrow_metadata(
790-
ifnull(with_metadata(name, 'ARROW:extension:name', 'a.type'), l_name),
791-
'ARROW:extension:name'),
792-
arrow_metadata(
793-
nvl2(id, with_metadata(name, 'ARROW:extension:name', 'a.type'), l_name),
794-
'ARROW:extension:name')
786+
arrow_metadata(coalesce(cast(name as json), l_name), 'ARROW:extension:name'),
787+
arrow_metadata(nvl(cast(name as json), upper(l_name)), 'ARROW:extension:name'),
788+
arrow_metadata(ifnull(cast(name as json), lower(l_name)), 'ARROW:extension:name'),
789+
arrow_metadata(nvl2(id, cast(name as json), l_name), 'ARROW:extension:name')
790+
from table_with_metadata;
791+
----
792+
arrow.json arrow.json arrow.json arrow.json arrow.json
793+
arrow.json arrow.json arrow.json arrow.json arrow.json
794+
arrow.json arrow.json arrow.json arrow.json arrow.json
795+
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')
795802
from table_with_metadata;
796803
----
797-
a.type a.type a.type a.type a.type
798-
a.type a.type a.type a.type a.type
799-
a.type a.type a.type a.type a.type
804+
arrow.json
805+
arrow.json
806+
arrow.json
800807

801808
statement ok
802809
drop table table_with_metadata;

0 commit comments

Comments
 (0)