Skip to content

Commit e0753c3

Browse files
committed
Reuse shared CASE metadata merge and preserve custom nullability
1 parent 8656718 commit e0753c3

1 file changed

Lines changed: 52 additions & 49 deletions

File tree

‎datafusion/expr/src/expr_schema.rs‎

Lines changed: 52 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ use arrow::datatypes::FieldRef;
2222
use arrow::datatypes::{DataType, Field};
2323
use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
2424
use datafusion_common::datatype::FieldExt;
25+
use datafusion_common::metadata::coerced_fields_metadata;
2526
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
2627
use 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

Comments
 (0)