@@ -22,6 +22,7 @@ use arrow::datatypes::FieldRef;
2222use arrow:: datatypes:: { DataType , Field } ;
2323use arrow_schema:: extension:: { EXTENSION_TYPE_METADATA_KEY , EXTENSION_TYPE_NAME_KEY } ;
2424use datafusion_common:: datatype:: FieldExt ;
25+ use datafusion_common:: metadata:: coerced_fields_metadata;
2526use datafusion_common:: tree_node:: { TreeNode , TreeNodeRecursion } ;
2627use datafusion_common:: {
2728 Column , DataFusionError , ExprSchema , Result , ScalarValue , Spans , TableReference ,
@@ -282,11 +283,12 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result<FieldRef>
282283 return internal_err ! ( "Missing CASE result fields" ) ;
283284 }
284285 let start = fields. len ( ) - count;
286+ let metadata = coerced_fields_metadata (
287+ fields[ start..] . iter ( ) . map ( |branch| branch. field . as_ref ( ) ) ,
288+ ) ?;
285289 let mut then_type = DataType :: Null ;
286290 let mut else_type = DataType :: Null ;
287291 let mut branch_type = None ;
288- let mut metadata: Option < HashMap < String , String > > = None ;
289- let mut extension: Option < ( String , Option < String > ) > = None ;
290292 let mut type_conflict = false ;
291293 let mut certainly_null = true ;
292294 for ( index, branch) in fields. drain ( start..) . enumerate ( ) {
@@ -308,60 +310,22 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result<FieldRef>
308310 type_conflict = true ;
309311 }
310312 branch_type. get_or_insert_with ( || data_type. clone ( ) ) ;
311- let branch_metadata = branch. field . metadata ( ) ;
312- if let Some ( name) = branch_metadata. get ( EXTENSION_TYPE_NAME_KEY ) {
313- match & extension {
314- Some ( ( first_name, _) ) if first_name != name => {
315- return plan_err ! (
316- "Cannot merge field metadata with differing values of \
317- '{EXTENSION_TYPE_NAME_KEY}': '{first_name}' and '{name}'"
318- ) ;
319- }
320- Some ( _) => { }
321- None => {
322- extension = Some ( (
323- name. clone ( ) ,
324- branch_metadata
325- . get ( EXTENSION_TYPE_METADATA_KEY )
326- . cloned ( ) ,
327- ) ) ;
328- }
329- }
330- }
331- for ( key, value) in branch_metadata {
332- if key != EXTENSION_TYPE_NAME_KEY
333- && key != EXTENSION_TYPE_METADATA_KEY
334- {
335- metadata
336- . get_or_insert_with ( HashMap :: new)
337- . entry ( key. clone ( ) )
338- . or_insert_with ( || value. clone ( ) ) ;
339- }
340- }
341313 }
342314 let data_type = if then_type. is_null ( ) {
343315 else_type
344316 } else {
345317 then_type
346318 } ;
347- if type_conflict || branch_type. as_ref ( ) != Some ( & data_type) {
348- metadata = None ;
349- extension = None ;
350- }
351- let mut metadata = metadata. unwrap_or_default ( ) ;
352- if let Some ( ( name, extension_metadata) ) = extension {
353- metadata. insert ( EXTENSION_TYPE_NAME_KEY . to_owned ( ) , name) ;
354- if let Some ( extension_metadata) = extension_metadata {
355- metadata. insert (
356- EXTENSION_TYPE_METADATA_KEY . to_owned ( ) ,
357- extension_metadata,
358- ) ;
359- }
360- }
319+ let field = Field :: new ( "" , data_type, true ) ;
320+ let field = if !type_conflict
321+ && branch_type. as_ref ( ) == Some ( field. data_type ( ) )
322+ {
323+ metadata. add_to_field ( field)
324+ } else {
325+ field
326+ } ;
361327 fields. push ( BranchField {
362- field : Arc :: new (
363- Field :: new ( "" , data_type, true ) . with_metadata ( metadata) ,
364- ) ,
328+ field : Arc :: new ( field) ,
365329 certainly_null,
366330 } ) ;
367331 }
@@ -1981,6 +1945,45 @@ mod tests {
19811945 Ok ( ( ) )
19821946 }
19831947
1948+ #[ test]
1949+ fn test_plain_case_preserves_custom_nullability ( ) -> Result < ( ) > {
1950+ #[ derive( Debug ) ]
1951+ struct CustomNullableSchema {
1952+ field : FieldRef ,
1953+ nullable : bool ,
1954+ }
1955+
1956+ impl ExprSchema for CustomNullableSchema {
1957+ fn nullable ( & self , _column : & Column ) -> Result < bool > {
1958+ Ok ( self . nullable )
1959+ }
1960+
1961+ fn field_from_column ( & self , _column : & Column ) -> Result < & FieldRef > {
1962+ Ok ( & self . field )
1963+ }
1964+ }
1965+
1966+ for nullable in [ false , true ] {
1967+ let schema = CustomNullableSchema {
1968+ field : Arc :: new ( Field :: new ( "x" , DataType :: Int32 , !nullable) ) ,
1969+ nullable,
1970+ } ;
1971+ let ( field, column_nullable) =
1972+ schema. field_and_nullable ( & Column :: from_name ( "x" ) ) ?;
1973+ assert_eq ! ( field. is_nullable( ) , !nullable) ;
1974+ assert_eq ! ( column_nullable, nullable) ;
1975+
1976+ for expression in [
1977+ when ( lit ( true ) , col ( "x" ) ) . otherwise ( lit ( 0_i32 ) ) ?,
1978+ when ( lit ( true ) , lit ( 0_i32 ) ) . otherwise ( col ( "x" ) ) ?,
1979+ ] {
1980+ assert_eq ! ( expression. nullable( & schema) ?, nullable) ;
1981+ assert_eq ! ( expression. to_field( & schema) ?. 1 . is_nullable( ) , nullable) ;
1982+ }
1983+ }
1984+ Ok ( ( ) )
1985+ }
1986+
19841987 #[ test]
19851988 fn test_plain_case_preserves_custom_nullable_error ( ) -> Result < ( ) > {
19861989 let schema = MockExprSchema :: new ( )
0 commit comments