fix(backends): ensure select after filter works

Author: cpcloud

ibis/backends/pandas/execution/generic.py

@@ -969,7 +969,10 @@ def execute_series_notnnull(op, data, **kwargs):
 
 @execute_node.register(ops.IsNan, (pd.Series, floating_types))
 def execute_isnan(op, data, **kwargs):
-    return np.isnan(data)
+    try:
+        return np.isnan(data)
+    except (TypeError, ValueError):
+        return data != data
 
 
 @execute_node.register(ops.IsInf, (pd.Series, floating_types))

ibis/backends/tests/test_generic.py

@@ -662,3 +662,22 @@ def test_where_column(backend, alltypes, df):
     )
 
     backend.assert_series_equal(result, expected)
+
+
+def test_select_filter(backend, alltypes, df):
+    t = alltypes
+
+    expr = t.select("int_col").filter(t.string_col == "4")
+    result = expr.execute()
+
+    expected = df.loc[df.string_col == "4", ["int_col"]].reset_index(drop=True)
+    backend.assert_frame_equal(result, expected)
+
+
+def test_select_filter_select(backend, alltypes, df):
+    t = alltypes
+    expr = t.select("int_col").filter(t.string_col == "4").int_col
+    result = expr.execute().rename("int_col")
+
+    expected = df.loc[df.string_col == "4", "int_col"].reset_index(drop=True)
+    backend.assert_series_equal(result, expected)

ibis/expr/analysis.py

@@ -12,7 +12,11 @@
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
 from ibis import util
-from ibis.common.exceptions import ExpressionError, IbisTypeError
+from ibis.common.exceptions import (
+    ExpressionError,
+    IbisTypeError,
+    IntegrityError,
+)
 from ibis.expr.window import window
 
 # ---------------------------------------------------------------------
@@ -343,21 +347,31 @@ def _filter_selection(expr, predicates):
     # the parent tables in the join being projected
 
     op = expr.op()
-    if not op.blocks():
-        # Potential fusion opportunity. The predicates may need to be
-        # rewritten in terms of the child table. This prevents the broken
-        # ref issue (described in more detail in #59)
+    # Potential fusion opportunity. The predicates may need to be
+    # rewritten in terms of the child table. This prevents the broken
+    # ref issue (described in more detail in #59)
+    try:
         simplified_predicates = tuple(
             sub_for(predicate, [(expr, op.table)])
             if not is_reduction(predicate)
             else predicate
             for predicate in predicates
         )
-
-        if shares_all_roots(simplified_predicates, op.table):
+    except IntegrityError:
+        pass
+    else:
+        if shares_all_roots(simplified_predicates, op.table) and not any(
+            # we can't push down filters on unnest because unnest changes the
+            # shape and potential values of the data: unnest can potentially
+            # produce NULLs
+            #
+            # the getattr shenanigans is to handle Alias
+            isinstance(getattr(sel.op(), "arg", sel).op(), ops.Unnest)
+            for sel in op.selections
+        ):
             result = ops.Selection(
                 op.table,
-                [],
+                selections=op.selections,
                 predicates=op.predicates + simplified_predicates,
                 sort_keys=op.sort_keys,
             )

ibis/tests/expr/test_analysis.py

@@ -284,7 +284,7 @@ def test_select_filter_mutate_fusion():
     #
     # eventually we will bring this back, but we're trading off the ability
     # to remove materialize for some performance in the short term
-    assert len(first_selection.op().selections) == 0
+    assert len(first_selection.op().selections) == 1
     assert len(first_selection.op().predicates) == 1
 
 

ibis/tests/sql/test_compiler.py

@@ -229,19 +229,15 @@ def test_agg_filter():
   FROM my_table
 ),
 t1 AS (
-  SELECT `a`, `b2`
+  SELECT t0.`a`, t0.`b2`
   FROM t0
+  WHERE t0.`a` < 100
 )
-SELECT t2.*
-FROM (
-  SELECT t1.*
-  FROM t1
-  WHERE t1.`a` < 100
-) t2
-WHERE t2.`a` = (
-  SELECT max(t1.`a`) AS `blah`
+SELECT t1.*
+FROM t1
+WHERE t1.`a` = (
+  SELECT max(`a`) AS `blah`
   FROM t1
-  WHERE t1.`a` < 100
 )"""
     assert result == expected
 
@@ -259,19 +255,15 @@ def test_agg_filter_with_alias():
   FROM my_table
 ),
 t1 AS (
-  SELECT `a`, `b2`
+  SELECT t0.`a`, t0.`b2`
   FROM t0
+  WHERE t0.`a` < 100
 )
-SELECT t2.*
-FROM (
-  SELECT t1.*
-  FROM t1
-  WHERE t1.`a` < 100
-) t2
-WHERE t2.`a` = (
-  SELECT max(t1.`a`) AS `blah`
+SELECT t1.*
+FROM t1
+WHERE t1.`a` = (
+  SELECT max(`a`) AS `blah`
   FROM t1
-  WHERE t1.`a` < 100
 )"""
     assert result == expected
 
@@ -358,8 +350,8 @@ def test_subquery_where_location():
   FROM (
     SELECT `float_col`, `timestamp_col`, `int_col`, `string_col`
     FROM alltypes
+    WHERE `timestamp_col` < '20140101'
   ) t1
-  WHERE `timestamp_col` < '20140101'
   GROUP BY 1
 ) t0"""
     assert result == expected

ibis/tests/sql/test_select_sql.py

@@ -1488,16 +1488,17 @@ def test_filter_predicates():
         expr = projected
 
     expected = """\
-SELECT t0.*
+SELECT *
 FROM (
   SELECT *
   FROM (
     SELECT *
     FROM t
     WHERE (lower(`color`) LIKE '%de%') AND
           (locate('de', lower(`color`)) - 1 >= 0)
-  ) t2
-) t0
-WHERE regexp_like(lower(t0.`color`), '.*ge.*')"""
+  ) t1
+  WHERE regexp_like(lower(`color`), '.*ge.*')
+) t0"""
 
-    assert Compiler.to_sql(expr) == expected
+    result = Compiler.to_sql(expr)
+    assert result == expected