Repository: ibis-project/ibis

Issue:
t['a', 'b'] and t.select(['a', 'b']) behave differently in the presence of a join



Discussion:

Resolving patch:
diff --git a/ibis/expr/analysis.py b/ibis/expr/analysis.py
--- a/ibis/expr/analysis.py
+++ b/ibis/expr/analysis.py
@@ -748,72 +748,69 @@ def __init__(self, parent, proj_exprs):
         self.parent = parent
         self.input_exprs = proj_exprs
         self.resolved_exprs = [parent._ensure_expr(e) for e in proj_exprs]
-
-        node = self.parent.op()
-
-        if isinstance(node, ops.Selection):
-            roots = [node]
-        else:
-            roots = node.root_tables()
-
-        self.parent_roots = roots
-
-        clean_exprs = []
-
-        for expr in self.resolved_exprs:
-            # Perform substitution only if we share common roots
-            expr = windowize_function(expr)
-            clean_exprs.append(expr)
-
-        self.clean_exprs = clean_exprs
+        node = parent.op()
+        self.parent_roots = (
+            [node] if isinstance(node, ops.Selection) else node.root_tables()
+        )
+        self.clean_exprs = list(map(windowize_function, self.resolved_exprs))
 
     def get_result(self):
         roots = self.parent_roots
         first_root = roots[0]
 
         if len(roots) == 1 and isinstance(first_root, ops.Selection):
-            fused_op = self._check_fusion(first_root)
+            fused_op = self.try_fusion(first_root)
             if fused_op is not None:
                 return fused_op
 
         return ops.Selection(self.parent, self.clean_exprs)
 
-    def _check_fusion(self, root):
-        roots = root.table._root_tables()
-        validator = ExprValidator([root.table])
+    def try_fusion(self, root):
+        root_table = root.table
+        roots = root_table._root_tables()
+        validator = ExprValidator([root_table])
         fused_exprs = []
         can_fuse = False
 
-        resolved = _maybe_resolve_exprs(root.table, self.input_exprs)
+        if not isinstance(root_table.op(), ops.Join):
+            resolved = _maybe_resolve_exprs(root_table, self.input_exprs)
+        else:
+            # joins cannot be used to resolve expressions, but we still may be
+            # able to fuse columns from a projection off of a join. In that
+            # case, use the projection's input expressions as the columns with
+            # which to attempt fusion
+            resolved = self.clean_exprs
+
         if not resolved:
             return None
 
+        root_selections = root.selections
+        parent_op = self.parent.op()
         for val in resolved:
             # XXX
             lifted_val = substitute_parents(val)
 
             # a * projection
             if isinstance(val, ir.TableExpr) and (
-                self.parent.op().compatible_with(val.op())
+                parent_op.compatible_with(val.op())
                 # gross we share the same table root. Better way to
                 # detect?
                 or len(roots) == 1
                 and val._root_tables()[0] is roots[0]
             ):
                 can_fuse = True
-
                 have_root = False
-                for y in root.selections:
+                for root_sel in root_selections:
                     # Don't add the * projection twice
-                    if y.equals(root.table):
-                        fused_exprs.append(root.table)
+                    if root_sel.equals(root_table):
+                        fused_exprs.append(root_table)
                         have_root = True
                         continue
-                    fused_exprs.append(y)
+                    fused_exprs.append(root_sel)
 
                 # This was a filter, so implicitly a select *
-                if not have_root and len(root.selections) == 0:
-                    fused_exprs = [root.table] + fused_exprs
+                if not have_root and not root_selections:
+                    fused_exprs = [root_table] + fused_exprs
             elif validator.validate(lifted_val):
                 can_fuse = True
                 fused_exprs.append(lifted_val)
@@ -825,13 +822,12 @@ def _check_fusion(self, root):
 
         if can_fuse:
             return ops.Selection(
-                root.table,
+                root_table,
                 fused_exprs,
                 predicates=root.predicates,
                 sort_keys=root.sort_keys,
             )
-        else:
-            return None
+        return None
 
 
 def _maybe_resolve_exprs(table, exprs):


Test patch:
diff --git a/ibis/expr/tests/test_value_exprs.py b/ibis/expr/tests/test_value_exprs.py
--- a/ibis/expr/tests/test_value_exprs.py
+++ b/ibis/expr/tests/test_value_exprs.py
@@ -1490,3 +1490,53 @@ def test_nullif_fail(left, right):
         left.nullif(right)
     with pytest.raises(com.IbisTypeError):
         right.nullif(left)
+
+
+@pytest.mark.parametrize(
+    "join_method",
+    [
+        "left_join",
+        pytest.param(
+            "right_join",
+            marks=pytest.mark.xfail(
+                raises=AttributeError, reason="right_join is not an ibis API"
+            ),
+        ),
+        "inner_join",
+        "outer_join",
+        "asof_join",
+        pytest.param(
+            "semi_join",
+            marks=pytest.mark.xfail(
+                raises=com.IbisTypeError,
+                reason=(
+                    "semi_join only gives access to the left table's "
+                    "columns"
+                ),
+            ),
+        ),
+    ],
+)
+@pytest.mark.xfail(
+    raises=(com.IbisError, AttributeError),
+    reason="Select from unambiguous joins not implemented",
+)
+def test_select_on_unambiguous_join(join_method):
+    t = ibis.table([("a0", dt.int64), ("b1", dt.string)], name="t")
+    s = ibis.table([("a1", dt.int64), ("b2", dt.string)], name="s")
+    method = getattr(t, join_method)
+    join = method(s, t.b1 == s.b2)
+    expr1 = join["a0", "a1"]
+    expr2 = join[["a0", "a1"]]
+    expr3 = join.select(["a0", "a1"])
+    assert expr1.equals(expr2)
+    assert expr1.equals(expr3)
+
+
+def test_chained_select_on_join():
+    t = ibis.table([("a", dt.int64)], name="t")
+    s = ibis.table([("a", dt.int64), ("b", dt.string)], name="s")
+    join = t.join(s)[t.a, s.b]
+    expr1 = join["a", "b"]
+    expr2 = join.select(["a", "b"])
+    assert expr1.equals(expr2)
diff --git a/ibis/pandas/execution/tests/test_join.py b/ibis/pandas/execution/tests/test_join.py
--- a/ibis/pandas/execution/tests/test_join.py
+++ b/ibis/pandas/execution/tests/test_join.py
@@ -4,6 +4,7 @@
 from pytest import param
 
 import ibis
+import ibis.common as com
 
 pytestmark = pytest.mark.pandas
 
@@ -338,3 +339,79 @@ def test_keyed_asof_join_with_tolerance(
         tolerance=pd.Timedelta('2D'),
     )
     tm.assert_frame_equal(result[expected.columns], expected)
+
+
+@pytest.mark.parametrize(
+    "how",
+    [
+        "left",
+        pytest.param(
+            "right",
+            marks=pytest.mark.xfail(
+                raises=AttributeError, reason="right_join is not an ibis API"
+            ),
+        ),
+        "inner",
+        "outer",
+    ],
+)
+@pytest.mark.parametrize(
+    "func",
+    [
+        pytest.param(lambda join: join["a0", "a1"], id="tuple"),
+        pytest.param(lambda join: join[["a0", "a1"]], id="list"),
+        pytest.param(lambda join: join.select(["a0", "a1"]), id="select"),
+    ],
+)
+@pytest.mark.xfail(
+    raises=(com.IbisError, AttributeError),
+    reason="Select from unambiguous joins not implemented",
+)
+def test_select_on_unambiguous_join(how, func):
+    df_t = pd.DataFrame(dict(a0=[1, 2, 3], b1=list("aab")))
+    df_s = pd.DataFrame(dict(a1=[2, 3, 4], b2=list("abc")))
+    con = ibis.pandas.connect({"t": df_t, "s": df_s})
+    t = con.table("t")
+    s = con.table("s")
+    method = getattr(t, "{}_join".format(how))
+    join = method(s, t.b1 == s.b2)
+    expected = pd.merge(df_t, df_s, left_on=["b1"], right_on=["b2"], how=how)[
+        ["a0", "a1"]
+    ]
+    assert not expected.empty
+    expr = func(join)
+    result = expr.execute()
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "func",
+    [
+        pytest.param(lambda join: join["a0", "a1"], id="tuple"),
+        pytest.param(lambda join: join[["a0", "a1"]], id="list"),
+        pytest.param(lambda join: join.select(["a0", "a1"]), id="select"),
+    ],
+)
+@pytest.mark.xfail(
+    raises=(com.IbisError, AttributeError),
+    reason="Select from unambiguous joins not implemented",
+)
+@merge_asof_minversion
+def test_select_on_unambiguous_asof_join(func):
+    df_t = pd.DataFrame(
+        dict(a0=[1, 2, 3], b1=pd.date_range("20180101", periods=3))
+    )
+    df_s = pd.DataFrame(
+        dict(a1=[2, 3, 4], b2=pd.date_range("20171230", periods=3))
+    )
+    con = ibis.pandas.connect({"t": df_t, "s": df_s})
+    t = con.table("t")
+    s = con.table("s")
+    join = t.asof_join(s, t.b1 == s.b2)
+    expected = pd.merge_asof(df_t, df_s, left_on=["b1"], right_on=["b2"])[
+        ["a0", "a1"]
+    ]
+    assert not expected.empty
+    expr = func(join)
+    result = expr.execute()
+    tm.assert_frame_equal(result, expected)