Repository: ibis-project/ibis

Issue:
SQLAlchemy: implement materialized join

SQLAlchemy: implement materialized join



Discussion:

Resolving patch:
diff --git a/ibis/sql/alchemy.py b/ibis/sql/alchemy.py
--- a/ibis/sql/alchemy.py
+++ b/ibis/sql/alchemy.py
@@ -74,14 +74,19 @@ def schema_from_table(table):
     for c in table.columns.values():
         type_class = type(c.type)
 
-        if c.type in _sqla_type_to_ibis:
-            ibis_class = _sqla_type_to_ibis[c.type]
-        elif type_class in _sqla_type_to_ibis:
-            ibis_class = _sqla_type_to_ibis[type_class]
+        if isinstance(c.type, sa.types.NUMERIC):
+            t = dt.Decimal(c.type.precision,
+                           c.type.scale,
+                           nullable=c.nullable)
         else:
-            raise NotImplementedError(c.type)
+            if c.type in _sqla_type_to_ibis:
+                ibis_class = _sqla_type_to_ibis[c.type]
+            elif type_class in _sqla_type_to_ibis:
+                ibis_class = _sqla_type_to_ibis[type_class]
+            else:
+                raise NotImplementedError(c.type)
+            t = ibis_class(c.nullable)
 
-        t = ibis_class(c.nullable)
         types.append(t)
 
     return dt.Schema(names, types)
@@ -92,7 +97,7 @@ def table_from_schema(name, meta, schema):
     sqla_cols = []
 
     for cname, itype in zip(schema.names, schema.types):
-        ctype = _ibis_type_to_sqla[type(itype)]
+        ctype = _to_sqla_type(itype)
 
         col = sa.Column(cname, ctype, nullable=itype.nullable)
         sqla_cols.append(col)
@@ -100,6 +105,13 @@ def table_from_schema(name, meta, schema):
     return sa.Table(name, meta, *sqla_cols)
 
 
+def _to_sqla_type(itype):
+    if isinstance(itype, dt.Decimal):
+        return sa.types.NUMERIC(itype.precision, itype.scale)
+    else:
+        return _ibis_type_to_sqla[type(itype)]
+
+
 def fixed_arity(sa_func, arity):
     if isinstance(sa_func, six.string_types):
         sa_func = getattr(sa.func, sa_func)
@@ -587,8 +599,12 @@ def _add_select(self, table_set):
                 arg = self._translate(expr, named=True)
             elif isinstance(expr, ir.TableExpr):
                 if expr.equals(self.table_set):
-                    # the select * case
-                    arg = table_set
+                    cached_table = self.context.get_table(expr)
+                    if cached_table is None:
+                        # the select * case from materialized join
+                        arg = '*'
+                    else:
+                        arg = table_set
                 else:
                     arg = self.context.get_table(expr)
                     if arg is None:


Test patch:
diff --git a/ibis/sql/tests/test_sqlalchemy.py b/ibis/sql/tests/test_sqlalchemy.py
--- a/ibis/sql/tests/test_sqlalchemy.py
+++ b/ibis/sql/tests/test_sqlalchemy.py
@@ -220,7 +220,7 @@ def test_joins(self):
         ipred = region.r_regionkey == nation.n_regionkey
         spred = rt.c.r_regionkey == nt.c.n_regionkey
 
-        joins = [
+        fully_mat_joins = [
             (region.inner_join(nation, ipred),
              rt.join(nt, spred)),
 
@@ -230,10 +230,47 @@ def test_joins(self):
             (region.outer_join(nation, ipred),
              rt.outerjoin(nt, spred)),
         ]
-        for ibis_joined, joined_sqla in joins:
-            expected = sa.select([rt, nt]).select_from(joined_sqla)
+        for ibis_joined, joined_sqla in fully_mat_joins:
+            expected = sa.select(['*']).select_from(joined_sqla)
             self._compare_sqla(ibis_joined, expected)
 
+        subselect_joins = [
+            (region.inner_join(nation, ipred).projection(nation),
+             rt.join(nt, spred)),
+
+            (region.left_join(nation, ipred).projection(nation),
+             rt.join(nt, spred, isouter=True)),
+
+            (region.outer_join(nation, ipred).projection(nation),
+             rt.outerjoin(nt, spred)),
+        ]
+        for ibis_joined, joined_sqla in subselect_joins:
+            expected = sa.select([nt]).select_from(joined_sqla)
+            self._compare_sqla(ibis_joined, expected)
+
+    def test_join_just_materialized(self):
+        joined = self._case_join_just_materialized()
+
+        rt, nt, ct = self._sqla_tables(['tpch_region', 'tpch_nation',
+                                        'tpch_customer'])
+        nt = nt.alias('t0')
+        rt = rt.alias('t1')
+        ct = ct.alias('t2')
+
+        sqla_joined = (nt.join(rt, nt.c.n_regionkey == rt.c.r_regionkey)
+                       .join(ct, nt.c.n_nationkey == ct.c.c_nationkey))
+
+        expected = sa.select(['*']).select_from(sqla_joined)
+
+        self._compare_sqla(joined, expected)
+
+    def _sqla_tables(self, tables):
+        result = []
+        for t in tables:
+            ibis_table = self.con.table(t)
+            result.append(self._to_sqla(ibis_table))
+        return result
+
     def test_simple_case(self):
         self.con.table('alltypes')
         st = self.con.meta.tables['alltypes']