Skip to content

Commit

Permalink
fix coalece bug
Browse files Browse the repository at this point in the history
  • Loading branch information
Lordworms committed Feb 5, 2025
1 parent 86b4afc commit 7ff3769
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 3 deletions.
25 changes: 23 additions & 2 deletions datafusion/expr-common/src/type_coercion/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -595,7 +595,7 @@ fn type_union_resolution_coercion(

/// Handle type union resolution including struct type and others.
pub fn try_type_union_resolution(data_types: &[DataType]) -> Result<Vec<DataType>> {
let err = match try_type_union_resolution_with_struct(data_types) {
let err = match try_type_union_resolution_with_struct(data_types, false) {
Ok(struct_types) => return Ok(struct_types),
Err(e) => Some(e),
};
Expand All @@ -611,6 +611,7 @@ pub fn try_type_union_resolution(data_types: &[DataType]) -> Result<Vec<DataType
// Since field name is the key of the struct, so it shouldn't be updated to the common column name like "c0" or "c1"
pub fn try_type_union_resolution_with_struct(
data_types: &[DataType],
is_unique: bool,
) -> Result<Vec<DataType>> {
let mut keys_string: Option<String> = None;
for data_type in data_types {
Expand All @@ -632,6 +633,12 @@ pub fn try_type_union_resolution_with_struct(
}
}

let first_fields = if let DataType::Struct(fields) = &data_types[0] {
fields.clone()
} else {
return internal_err!("Struct type is checked is the previous function, so this should be unreachable");
};

let mut struct_types_map: HashMap<String, DataType> = if let DataType::Struct(
fields,
) = &data_types[0]
Expand All @@ -650,7 +657,7 @@ pub fn try_type_union_resolution_with_struct(
let field_name = field.name();
if let Some(existing_type) = struct_types_map.get_mut(field_name) {
if let Some(coerced_type) =
type_union_resolution_coercion(&field.data_type(), existing_type)
type_union_resolution_coercion(field.data_type(), existing_type)
{
*existing_type = coerced_type;
} else {
Expand All @@ -669,6 +676,20 @@ pub fn try_type_union_resolution_with_struct(
}
}

if is_unique {
let new_fields =
first_fields
.iter()
.map(|f| {
Arc::new(Arc::unwrap_or_clone(Arc::clone(f)).with_data_type(
struct_types_map.get(f.name()).unwrap().to_owned(),
))
})
.collect();
let unified_struct = DataType::Struct(new_fields);
return Ok(vec![unified_struct; data_types.len()]);
}

let mut final_struct_types = vec![];
for s in data_types {
let mut new_fields = vec![];
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-nested/src/make_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ impl ScalarUDFImpl for MakeArray {

fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
let mut errors = vec![];
match try_type_union_resolution_with_struct(arg_types) {
match try_type_union_resolution_with_struct(arg_types, true) {
Ok(r) => return Ok(r),
Err(e) => {
errors.push(e);
Expand Down

0 comments on commit 7ff3769

Please sign in to comment.