check alias in whole expr on opt (#3032)

pola-rs · Apr 1, 2022 · 3e7631f · 3e7631f
1 parent fb2c57e
commit 3e7631f
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 8 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/optimizer/projection_pushdown.rs b/polars/polars-lazy/src/logical_plan/optimizer/projection_pushdown.rs
@@ -1,4 +1,5 @@
 use crate::logical_plan::Context;
+use crate::prelude::iterator::ArenaExprIter;
 use crate::prelude::*;
 use crate::utils::{
     aexpr_assign_renamed_root, aexpr_to_root_names, aexpr_to_root_nodes, check_input_node,
@@ -255,16 +256,30 @@ impl ProjectionPushDown {
                     //
                     // In this query, bar cannot pass this projection, as it would not exist in DF.
                     if !acc_projections.is_empty() {
-                        if let AExpr::Alias(_, name) = expr_arena.get(*e) {
-                            if projected_names.remove(name) {
-                                acc_projections = acc_projections
-                                    .into_iter()
-                                    .filter(|expr| {
-                                        !aexpr_to_root_names(*expr, expr_arena).contains(name)
-                                    })
-                                    .collect();
+                        for (_, ae) in (&*expr_arena).iter(*e) {
+                            if let AExpr::Alias(_, name) = ae {
+                                if projected_names.remove(name) {
+                                    acc_projections = acc_projections
+                                        .into_iter()
+                                        .filter(|expr| {
+                                            !aexpr_to_root_names(*expr, expr_arena).contains(name)
+                                        })
+                                        .collect();
+                                }
                             }
                         }
+                        // if has_aexpr(*e, expr_arena, |ae| matches!(ae, AExpr::Alias(_, _))) {}
+                        //
+                        // if let AExpr::Alias(_, name) = expr_arena.get(*e) {
+                        //     if projected_names.remove(name) {
+                        //         acc_projections = acc_projections
+                        //             .into_iter()
+                        //             .filter(|expr| {
+                        //                 !aexpr_to_root_names(*expr, expr_arena).contains(name)
+                        //             })
+                        //             .collect();
+                        //     }
+                        // }
                     }
 
                     add_expr_to_accumulated(

diff --git a/polars/tests/it/lazy/queries.rs b/polars/tests/it/lazy/queries.rs
@@ -108,3 +108,19 @@ fn max_on_empty_df_3027() -> Result<()> {
     assert_eq!(out.shape(), (0, 3));
     Ok(())
 }
+
+#[test]
+fn test_alias_before_cast() -> Result<()> {
+    let out = df![
+        "a" => [1, 2, 3],
+    ]?
+    .lazy()
+    .select([col("a").alias("d").cast(DataType::Int32)])
+    .select([all()])
+    .collect()?;
+    assert_eq!(
+        Vec::from(out.column("d")?.i32()?),
+        &[Some(1), Some(2), Some(3)]
+    );
+    Ok(())
+}