fix bug in predicate pd after renaming node (#2860)

polars/polars-core/src/vector_hasher.rs

@@ -231,31 +231,36 @@ pub(crate) trait AsU64 {
 }
 
 impl AsU64 for u32 {
+    #[inline]
     fn as_u64(self) -> u64 {
         self as u64
     }
 }
 
 impl AsU64 for u64 {
+    #[inline]
     fn as_u64(self) -> u64 {
         self
     }
 }
 
 impl AsU64 for i32 {
+    #[inline]
     fn as_u64(self) -> u64 {
         let asu32: u32 = unsafe { std::mem::transmute(self) };
         asu32 as u64
     }
 }
 
 impl AsU64 for i64 {
+    #[inline]
     fn as_u64(self) -> u64 {
         unsafe { std::mem::transmute(self) }
     }
 }
 
 impl AsU64 for Option<u32> {
+    #[inline]
     fn as_u64(self) -> u64 {
         match self {
             Some(v) => v as u64,
@@ -266,12 +271,14 @@ impl AsU64 for Option<u32> {
 }
 
 impl AsU64 for Option<u64> {
+    #[inline]
     fn as_u64(self) -> u64 {
         self.unwrap_or(u64::MAX >> 2)
     }
 }
 
 impl AsU64 for [u8; 9] {
+    #[inline]
     fn as_u64(self) -> u64 {
         // the last byte includes the null information.
         // that one is skipped. Worst thing that could happen is unbalanced partition.
@@ -280,12 +287,14 @@ impl AsU64 for [u8; 9] {
 }
 const BUILD_HASHER: RandomState = RandomState::with_seeds(0, 0, 0, 0);
 impl AsU64 for [u8; 17] {
+    #[inline]
     fn as_u64(self) -> u64 {
         <[u8]>::get_hash(&self, &BUILD_HASHER)
     }
 }
 
 impl AsU64 for [u8; 13] {
+    #[inline]
     fn as_u64(self) -> u64 {
         <[u8]>::get_hash(&self, &BUILD_HASHER)
     }

polars/polars-lazy/src/frame/mod.rs

@@ -332,6 +332,12 @@ impl LazyFrame {
     }
 
     fn rename_imp(self, existing: Vec<String>, new: Vec<String>) -> Self {
+        let mut schema = (*self.schema()).clone();
+
+        for (old, new) in existing.iter().zip(&new) {
+            let _ = schema.rename(old, new.clone());
+        }
+
         self.with_columns(
             existing
                 .iter()
@@ -354,7 +360,7 @@ impl LazyFrame {
                 Ok(df)
             },
             None,
-            None,
+            Some(schema),
             Some("RENAME"),
         )
     }

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs

@@ -81,17 +81,16 @@ impl PredicatePushDown {
             let new_inputs = inputs
                 .iter()
                 .map(|&node| {
-                    // first we check if we are able to push down the predicate pass this node
+                    // first we check if we are able to push down the predicate passed this node
                     // it could be that this node just added the column where we base the predicate on
                     let input_schema = lp_arena.get(node).schema(lp_arena);
                     let mut pushdown_predicates = optimizer::init_hashmap();
-                    for &predicate in acc_predicates.values() {
+                    for (name, &predicate) in acc_predicates.iter() {
                         // we can pushdown the predicate
                         if check_input_node(predicate, input_schema, expr_arena) {
-                            let name = get_insertion_name(expr_arena, predicate, input_schema);
                             insert_and_combine_predicate(
                                 &mut pushdown_predicates,
-                                name,
+                                name.clone(),
                                 predicate,
                                 expr_arena,
                             )

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/utils.rs

@@ -188,6 +188,7 @@ where
             if let AExpr::Alias(_, name) = projection_aexpr {
                 // if this alias refers to one of the predicates in the upper nodes
                 // we rename the column of the predicate before we push it downwards.
+
                 if let Some(predicate) = acc_predicates.remove(&*name) {
                     if projection_is_boundary {
                         local_predicates.push(predicate);

polars/tests/it/lazy/mod.rs

@@ -1,6 +1,7 @@
 mod explodes;
 mod expressions;
 mod groupby_dynamic;
+mod predicate_queries;
 mod projection_queries;
 
 use polars::prelude::*;

polars/tests/it/lazy/predicate_queries.rs

@@ -0,0 +1,21 @@
+use super::*;
+
+#[test]
+fn test_predicate_after_renaming() -> Result<()> {
+    let df = df![
+        "foo" => [1, 2, 3],
+        "bar" => [3, 2, 1]
+    ]?
+    .lazy()
+    .rename(["foo", "bar"], ["foo2", "bar2"])
+    .filter(col("foo2").eq(col("bar2")))
+    .collect()?;
+
+    let expected = df![
+        "foo2" => [2],
+        "bar2" => [2],
+    ]?;
+    assert!(df.frame_equal(&expected));
+
+    Ok(())
+}