LiaCastaneda commented on code in PR #27099:
URL: https://github.com/apache/datafusion/pull/27099#discussion_r4523618486
##########
datafusion/substrait/src/logical_plan/consumer/expr/literal.rs:
##########
@@ -573,6 +578,161 @@ pub(crate) fn from_substrait_literal(
Ok(scalar_value)
}
+/// Gives struct-typed literals among `exprs` the field names of an expression
of the
+/// same shape next to them.
+///
+/// Substrait types carry no struct field names, so a struct literal is
consumed with
+/// generated names (`c0`, `c1`, ...). Next to a column of the same shape, as
in
+/// `coalesce(s, [])` with `s: List<Struct{a, b}>`, the literal would be
+/// `List<Struct{c0, c1}>`, and DataFusion matches struct fields by name, so
the two types
+/// can't be unified. Substrait matches struct fields by position, so taking
the column's
+/// names keeps the meaning of the plan.
+pub(crate) fn name_struct_literals_like_siblings(
+ exprs: &mut [&mut Expr],
+ input_schema: &DFSchema,
+) {
+ // the names come from the first expression that isn't a literal and
contains a struct
+ let Some(named_type) = exprs
+ .iter()
+ .filter(|e| !is_literal(e) && !matches!(***e, Expr::Lambda(_)))
+ .filter_map(|e| e.get_type(input_schema).ok())
+ .find(contains_struct)
+ else {
+ return;
+ };
+
+ for expr in exprs.iter_mut() {
+ name_struct_literal_like(expr, &named_type);
+ }
+}
+
+/// If `expr` is a struct-typed literal, possibly aliased, gives it the field
names of
+/// `named_type` when the two have the same shape. Anything else is left
unchanged.
+pub(crate) fn name_struct_literal_like(expr: &mut Expr, named_type: &DataType)
{
+ match expr {
+ Expr::Alias(alias) => name_struct_literal_like(&mut alias.expr,
named_type),
+ Expr::Literal(value, _) => {
+ let value_type = value.data_type();
+ if !contains_struct(&value_type) {
+ return;
+ }
+ let Some(renamed_type) = with_field_names_of(&value_type,
named_type) else {
+ return;
+ };
+ // if renaming fails, keep the literal: coercion reports the
mismatch as before
+ if let Ok(renamed) = with_data_type(value, &renamed_type) {
+ *value = renamed;
+ }
+ }
+ _ => {}
+ }
+}
+
+/// Returns true if `expr` is a literal, possibly aliased.
+pub(crate) fn is_literal(expr: &Expr) -> bool {
+ match expr {
+ Expr::Alias(alias) => is_literal(&alias.expr),
+ Expr::Literal(..) => true,
+ _ => false,
+ }
+}
+
+pub(crate) fn contains_struct(data_type: &DataType) -> bool {
+ match data_type {
+ DataType::Struct(_) => true,
+ DataType::List(f)
+ | DataType::LargeList(f)
+ | DataType::FixedSizeList(f, _)
+ | DataType::Map(f, _) => contains_struct(f.data_type()),
+ _ => false,
+ }
+}
+
+/// Returns `data_type` with the field names (and field metadata) of `named`,
at every
+/// level, or `None` if the two don't have the same shape. Everything else,
such as leaf
+/// types and nullability, stays as in `data_type`.
+fn with_field_names_of(data_type: &DataType, named: &DataType) ->
Option<DataType> {
+ let renamed_field = |field: &FieldRef, named: &FieldRef| ->
Option<FieldRef> {
+ let data_type = with_field_names_of(field.data_type(),
named.data_type())?;
+ Some(Arc::new(
+ Field::new(named.name(), data_type, field.is_nullable())
+ .with_metadata(named.metadata().clone()),
+ ))
+ };
+ match (data_type, named) {
Review Comment:
yes makes sense, added `Dictionary` as an extra handled type
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
