[GitHub] [arrow-datafusion] liukun4515 commented on a diff in pull request #3727: Consolidate and better tests for expression re-rewriting / aliasing

[GitBox <gi...@apache.org>]

liukun4515 commented on code in PR #3727:
URL: https://github.com/apache/arrow-datafusion/pull/3727#discussion_r988526071


##########
datafusion/optimizer/src/utils.rs:
##########
@@ -315,13 +316,63 @@ pub fn alias_cols(cols: &[Column]) -> Vec<Expr> {
         .collect()
 }
 
+/// Rewrites `expr` using `rewriter`, ensuring that the output has the
+/// same name as `expr` prior to rewrite, adding an alias if necessary.
+///
+/// This is important when optimzing plans to ensure the the output
+/// schema of plan nodes don't change after optimization
+pub fn rewrite_preserving_name<R>(expr: Expr, rewriter: &mut R) -> Result<Expr>
+where
+    R: ExprRewriter<Expr>,
+{
+    let original_name = name_for_alias(&expr)?;
+    let expr = expr.rewrite(rewriter)?;
+    add_alias_if_changed(original_name, expr)
+}
+
+/// Return the name to use for the specific Expr, recursing into
+/// `Expr::Sort` as appropriate
+fn name_for_alias(expr: &Expr) -> Result<String> {
+    match expr {
+        Expr::Sort { expr, .. } => name_for_alias(expr),

Review Comment:
   I missed this issue https://github.com/apache/arrow-datafusion/pull/3710
   
   But I want to know why we need do the special branch for `Sort` Expr?



##########
datafusion/optimizer/src/utils.rs:
##########
@@ -315,13 +316,63 @@ pub fn alias_cols(cols: &[Column]) -> Vec<Expr> {
         .collect()
 }
 
+/// Rewrites `expr` using `rewriter`, ensuring that the output has the
+/// same name as `expr` prior to rewrite, adding an alias if necessary.
+///
+/// This is important when optimzing plans to ensure the the output
+/// schema of plan nodes don't change after optimization
+pub fn rewrite_preserving_name<R>(expr: Expr, rewriter: &mut R) -> Result<Expr>
+where
+    R: ExprRewriter<Expr>,
+{
+    let original_name = name_for_alias(&expr)?;
+    let expr = expr.rewrite(rewriter)?;
+    add_alias_if_changed(original_name, expr)
+}
+
+/// Return the name to use for the specific Expr, recursing into
+/// `Expr::Sort` as appropriate
+fn name_for_alias(expr: &Expr) -> Result<String> {
+    match expr {
+        Expr::Sort { expr, .. } => name_for_alias(expr),

Review Comment:
   I missed this issue https://github.com/apache/arrow-datafusion/pull/3710
   
   But I want to know why we need to do the special branch for `Sort` Expr?



-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org