fix(rust, python): ensure that no-projection is seen as select all in… (#5356)

Author: ritchie46

polars/polars-core/src/datatypes/aliases.rs

@@ -69,3 +69,15 @@ impl<K> InitHashMaps for PlIndexSet<K> {
         Self::with_capacity_and_hasher(capacity, Default::default())
     }
 }
+
+impl<K, V> InitHashMaps for PlIndexMap<K, V> {
+    type HashMap = Self;
+
+    fn new() -> Self::HashMap {
+        Self::with_capacity_and_hasher(0, Default::default())
+    }
+
+    fn with_capacity(capacity: usize) -> Self::HashMap {
+        Self::with_capacity_and_hasher(capacity, Default::default())
+    }
+}

polars/polars-lazy/polars-plan/src/logical_plan/optimizer/cache_states.rs

@@ -15,7 +15,7 @@ fn get_upper_projections(
     use ALogicalPlan::*;
     // during projection pushdown all accumulated p
     match parent {
-        Projection { expr, .. } => {
+        Projection { expr, .. } | HStack { exprs: expr, .. } => {
             let mut out = Vec::with_capacity(expr.len());
             for node in expr {
                 out.extend(aexpr_to_leaf_names_iter(*node, expr_arena));
@@ -92,12 +92,9 @@ pub(super) fn set_cache_states(
                     if let Some(names) = get_upper_projections(parent_node, lp_arena, expr_arena) {
                         entry.1.extend(names);
                     }
-                    // if there is no projection above, it maybe that the
-                    // cache is underneath another cache and projection pushdown never reached it.
-                    // other trails may take care of that cache
-                    // if there is no other cache above, then there was no projection and we must take
+                    // There was no projection and we must take
                     // all columns
-                    else if previous_cache.is_none() {
+                    else {
                         let schema = lp.schema(lp_arena);
                         entry
                             .1

polars/polars-lazy/polars-plan/src/logical_plan/optimizer/cse.rs

@@ -239,12 +239,14 @@ fn lp_node_equal(a: &ALogicalPlan, b: &ALogicalPlan, expr_arena: &Arena<AExpr>)
 }
 
 /// Iterate from two leaf location upwards and find the latest matching node.
+///
+/// Returns the matching nodes
 fn longest_subgraph(
     trail_a: &Trail,
     trail_b: &Trail,
     lp_arena: &Arena<ALogicalPlan>,
     expr_arena: &Arena<AExpr>,
-) -> Option<(usize, Node, Node)> {
+) -> Option<(Node, Node)> {
     if trail_a.is_empty() || trail_b.is_empty() {
         return None;
     }
@@ -276,7 +278,7 @@ fn longest_subgraph(
     }
     // previous node was equal
     if i > 0 {
-        Some((i - 1, prev_node_a, prev_node_b))
+        Some((prev_node_a, prev_node_b))
     } else {
         None
     }
@@ -317,14 +319,14 @@ pub(crate) fn elim_cmn_subplans(
     let mut changed = false;
 
     let mut cache_mapping = BTreeMap::new();
+    let mut cache_counts = PlHashMap::with_capacity(trail_ends.len());
 
-    // insert cache nodes
     for combination in trail_ends.iter() {
         // both are the same, but only point to a different location
         // in our arena so we hash one and store the hash for both locations
         // this will ensure all matches have the same hash.
-        let node1 = combination.1 .0;
-        let node2 = combination.2 .0;
+        let node1 = combination.0 .0;
+        let node2 = combination.1 .0;
 
         let cache_id = match (cache_mapping.get(&node1), cache_mapping.get(&node2)) {
             (Some(h), _) => *h,
@@ -345,12 +347,31 @@ pub(crate) fn elim_cmn_subplans(
                 cache_id
             }
         };
+        *cache_counts.entry(cache_id).or_insert(0usize) += 1;
+    }
+
+    // insert cache nodes
+    for combination in trail_ends.iter() {
+        // both are the same, but only point to a different location
+        // in our arena so we hash one and store the hash for both locations
+        // this will ensure all matches have the same hash.
+        let node1 = combination.0 .0;
+        let node2 = combination.1 .0;
+
+        let cache_id = match (cache_mapping.get(&node1), cache_mapping.get(&node2)) {
+            (Some(h), _) => *h,
+            (_, Some(h)) => *h,
+            _ => {
+                unreachable!()
+            }
+        };
+        let cache_count = *cache_counts.get(&cache_id).unwrap();
 
         // reassign old nodes to another location as we are going to replace
         // them with a cache node
-        for inp_node in [combination.1, combination.2] {
+        for inp_node in [combination.0, combination.1] {
             if let ALogicalPlan::Cache { count, .. } = lp_arena.get_mut(inp_node) {
-                *count += 1;
+                *count = cache_count;
             } else {
                 let lp = lp_arena.get(inp_node).clone();
 
@@ -360,7 +381,7 @@ pub(crate) fn elim_cmn_subplans(
                     input: node,
                     id: cache_id,
                     // remove after one cache hit.
-                    count: 1,
+                    count: cache_count,
                 };
                 lp_arena.replace(inp_node, cache_lp.clone());
             };

polars/polars-lazy/src/tests/cse.rs

@@ -198,3 +198,59 @@ fn test_cse_joins_4954() -> PolarsResult<()> {
 
     Ok(())
 }
+#[test]
+#[cfg(feature = "semi_anti_join")]
+fn test_cache_with_partial_projection() -> PolarsResult<()> {
+    let lf1 = df![
+        "id" => ["a"],
+        "x" => [1],
+        "freq" => [2]
+    ]?
+    .lazy();
+
+    let lf2 = df![
+        "id" => ["a"]
+    ]?
+    .lazy();
+
+    let q = lf2
+        .join(
+            lf1.clone().select([col("id"), col("freq")]),
+            [col("id")],
+            [col("id")],
+            JoinType::Semi,
+        )
+        .join(
+            lf1.clone().filter(col("x").neq(lit(8))),
+            [col("id")],
+            [col("id")],
+            JoinType::Semi,
+        )
+        .join(
+            lf1.clone().filter(col("x").neq(lit(8))),
+            [col("id")],
+            [col("id")],
+            JoinType::Semi,
+        );
+
+    let q = q.with_common_subplan_elimination(true);
+
+    let (mut expr_arena, mut lp_arena) = get_arenas();
+    let lp = q.optimize(&mut lp_arena, &mut expr_arena).unwrap();
+
+    // ensure we get two different caches
+    // and ensure that every cache only has 1 hit.
+    let cache_ids = (&lp_arena)
+        .iter(lp)
+        .flat_map(|(_, lp)| {
+            use ALogicalPlan::*;
+            match lp {
+                Cache { id, .. } => Some(*id),
+                _ => None,
+            }
+        })
+        .collect::<BTreeSet<_>>();
+    assert_eq!(cache_ids.len(), 2);
+
+    Ok(())
+}