Skip to content

Commit

Permalink
Enable non-literal k in MAP[k], in get_field UDF.
Browse files Browse the repository at this point in the history
Requiring k to be a literal should only apply to STRUCT.k
With MAP[k] it is an excessive typechecking restriction,
compared to what the operation is supposed to do.
  • Loading branch information
vgapeyev committed Aug 23, 2024
1 parent e466803 commit 008153b
Showing 1 changed file with 8 additions and 21 deletions.
29 changes: 8 additions & 21 deletions datafusion/functions/src/core/getfield.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,17 +63,11 @@ impl ScalarUDFImpl for GetFieldFunc {
args.len()
);
}

let name = match &args[1] {
Expr::Literal(name) => name,
_ => {
return exec_err!(
"get_field function requires the argument field_name to be a string"
);
}
};

Ok(format!("{}[{}]", args[0].display_name()?, name))
Ok(format!(
"{}[{}]",
args[0].display_name()?,
args[1].display_name()?
))
}

fn signature(&self) -> &Signature {
Expand All @@ -97,16 +91,9 @@ impl ScalarUDFImpl for GetFieldFunc {
);
}

let name = match &args[1] {
Expr::Literal(name) => name,
_ => {
return exec_err!(
"get_field function requires the argument field_name to be a string"
);
}
};
let key = &args[1];
let data_type = args[0].get_type(schema)?;
match (data_type, name) {
match (data_type, key) {
(DataType::Map(fields, _), _) => {
match fields.data_type() {
DataType::Struct(fields) if fields.len() == 2 => {
Expand All @@ -120,7 +107,7 @@ impl ScalarUDFImpl for GetFieldFunc {
_ => plan_err!("Map fields must contain a Struct with exactly 2 fields"),
}
}
(DataType::Struct(fields), ScalarValue::Utf8(Some(s))) => {
(DataType::Struct(fields), Expr::Literal(ScalarValue::Utf8(Some(s)))) => {
if s.is_empty() {
plan_err!(
"Struct based indexed access requires a non empty string"
Expand Down

0 comments on commit 008153b

Please sign in to comment.