blaginin commented on code in PR #14684:
URL: https://github.com/apache/datafusion/pull/14684#discussion_r1958693053
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -1779,37 +1779,82 @@ impl DataFrame {
.config_options()
.sql_parser
.enable_ident_normalization;
+
let old_column: Column = if ident_opts {
Column::from_qualified_name(old_name)
} else {
Column::from_qualified_name_ignore_case(old_name)
};
- let (qualifier_rename, field_rename) =
- match self.plan.schema().qualified_field_from_column(&old_column) {
- Ok(qualifier_and_field) => qualifier_and_field,
- // no-op if field not found
- Err(DataFusionError::SchemaError(
- SchemaError::FieldNotFound { .. },
- _,
- )) => return Ok(self),
- Err(err) => return Err(err),
- };
- let projection = self
- .plan
- .schema()
- .iter()
- .map(|(qualifier, field)| {
- if qualifier.eq(&qualifier_rename) && field.as_ref() ==
field_rename {
- col(Column::from((qualifier, field))).alias(new_name)
- } else {
- col(Column::from((qualifier, field)))
- }
- })
- .collect::<Vec<_>>();
- let project_plan = LogicalPlanBuilder::from(self.plan)
- .project(projection)?
- .build()?;
+ let project_plan = if let LogicalPlan::Projection(Projection {
+ expr,
+ input,
+ schema,
+ ..
+ }) = self.plan
+ {
+ // special case: we already have a projection on top, so we can
reuse it rather than creating a new one
+ let (qualifier_rename, field_rename) =
+ match schema.qualified_field_from_column(&old_column) {
+ Ok(qualifier_and_field) => qualifier_and_field,
+ // no-op if field not found
+ Err(DataFusionError::SchemaError(
+ SchemaError::FieldNotFound { .. },
+ _,
+ )) => {
+ self.plan = LogicalPlan::Projection(
+ Projection::try_new_with_schema(expr, input,
schema)?,
+ );
+ return Ok(self);
+ }
+ Err(err) => return Err(err),
+ };
+
+ let expr: Vec<_> = expr
+ .into_iter()
+ .map(|e| {
+ let (qualifier, field) = e.qualified_name();
+
+ if qualifier.as_ref().eq(&qualifier_rename)
+ && field.as_str() == field_rename.name()
+ {
+ e.alias_qualified(qualifier, new_name.to_string())
+ } else {
+ e
+ }
+ })
+ .collect();
+ LogicalPlan::Projection(Projection::try_new(expr, input)?)
+ } else {
Review Comment:
That's very fair! I also dislike repeating `qualifier_rename, field_rename`
twice, but it's happening because they reference `plan` and must come after its
destruction (hence needed twice)
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
