Merge pull request #501 from uwescience/bmyerz/partition-optimize-groupbys

Optimize groupby and distinct according to partitioning

Author: senderista

raco/backends/myria/myria.py

@@ -1,23 +1,23 @@
 import itertools
 import logging
 from collections import defaultdict, deque
+from functools import reduce
 from operator import mul
 
 from sqlalchemy.dialects import postgresql
 
 from raco import algebra, expression, rules, scheme
-from raco.algebra import convertcondition
+from raco import types
 from raco.algebra import Shuffle
-from raco.catalog import Catalog
-from raco.representation import RepresentationProperties
+from raco.algebra import convertcondition
 from raco.backends import Language, Algebra
-from raco.backends.sql.catalog import SQLCatalog
-from raco.expression import WORKERID, COUNTALL
-from raco.expression import UnnamedAttributeRef
+from raco.backends.sql.catalog import SQLCatalog, PostgresSQLFunctionProvider
+from raco.catalog import Catalog
 from raco.datastructure.UnionFind import UnionFind
-from raco import types
-from raco.rules import distributed_group_by
-from functools import reduce
+from raco.expression import UnnamedAttributeRef
+from raco.expression import WORKERID, COUNTALL
+from raco.representation import RepresentationProperties
+from raco.rules import distributed_group_by, check_partition_equality
 
 LOGGER = logging.getLogger(__name__)
 
@@ -1133,17 +1133,6 @@ def fire(self, exp):
         return exp
 
 
-def check_partition_equality(op, representation):
-    """Check to see if the operator has the required hash partitioning.
-    @param op operator
-    @param representation list of columns hash partitioned by,
-                        in the unnamed perspective
-    @return true if the op has an equal hash partitioning to representation
-    """
-
-    return op.partitioning().hash_partitioned == frozenset(representation)
-
-
 class ShuffleBeforeSetop(rules.Rule):
 
     def fire(self, exp):
@@ -1481,14 +1470,16 @@ def fire(self, op):
 
 class PushIntoSQL(rules.Rule):
 
-    def __init__(self, dialect=None):
+    def __init__(self, dialect=None, push_grouping=False):
         self.dialect = dialect or postgresql.dialect()
+        self.push_grouping = push_grouping
         super(PushIntoSQL, self).__init__()
 
     def fire(self, expr):
         if isinstance(expr, (algebra.Scan, algebra.ScanTemp)):
             return expr
-        cat = SQLCatalog()
+        cat = SQLCatalog(provider=PostgresSQLFunctionProvider(),
+                         push_grouping=self.push_grouping)
         try:
             sql_plan = cat.get_sql(expr)
             sql_string = sql_plan.compile(dialect=self.dialect)
@@ -1713,7 +1704,9 @@ def opt_rules(self, **kwargs):
 
         if kwargs.get('push_sql', False):
             opt_grps_sequence.append([
-                PushIntoSQL(dialect=kwargs.get('dialect'))])
+                PushIntoSQL(dialect=kwargs.get('dialect'),
+                            push_grouping=kwargs.get(
+                                'push_sql_grouping', False))])
 
         compile_grps_sequence = [
             myriafy,

raco/backends/sql/catalog.py

@@ -12,6 +12,7 @@
 import raco.scheme as scheme
 import raco.types as types
 from raco.representation import RepresentationProperties
+import abc
 
 
 type_to_raco = {Integer: types.LONG_TYPE,
@@ -30,9 +31,46 @@
                 types.DATETIME_TYPE: DateTime}
 
 
+class SQLFunctionProvider(object):
+    """Interface for translating function names. For Raco functions
+    not understood by SQLAlchemy, like stdev, we cannot rely
+    on SQLAlchemy's compiler to translate function
+    names to the given dialect.
+    For functions not understood by SQLAlchemy, the SQLAlchemy compiler
+    just emits them verbatim."""
+
+    __metaclass__ = abc.ABCMeta
+
+    @abc.abstractmethod
+    def convert_unary_expr(self, expr, input):
+        pass
+
+
+class _DefaultSQLFunctionProvider(SQLFunctionProvider):
+    def convert_unary_expr(self, expr, input):
+        # just use the function name without complaining
+        fname = expr.__class__.__name__.lower()
+        return getattr(func, fname)(input)
+
+
+class PostgresSQLFunctionProvider(SQLFunctionProvider):
+    def convert_unary_expr(self, expr, input):
+        fname = expr.__class__.__name__.lower()
+
+        # replacements
+        if fname == "stdev":
+            return func.stddev_samp(input)
+
+        # Warning: may create some functions not available in Postgres
+        return getattr(func, fname)(input)
+
+
 class SQLCatalog(Catalog):
-    def __init__(self, engine=None):
+    def __init__(self, engine=None, push_grouping=False,
+                 provider=_DefaultSQLFunctionProvider()):
         self.engine = engine
+        self.push_grouping = push_grouping
+        self.provider = provider
         self.metadata = MetaData()
 
     @staticmethod
@@ -106,10 +144,11 @@ def _convert_zeroary_expr(self, cols, expr, input_scheme):
 
     def _convert_unary_expr(self, cols, expr, input_scheme):
         input = self._convert_expr(cols, expr.input, input_scheme)
-        if isinstance(expr, expression.MAX):
-            return func.max(input)
-        if isinstance(expr, expression.MIN):
-            return func.min(input)
+
+        c = self.provider.convert_unary_expr(expr, input)
+        if c is not None:
+            return c
+
         raise NotImplementedError("expression {} to sql".format(type(expr)))
 
     def _convert_binary_expr(self, cols, expr, input_scheme):
@@ -157,9 +196,11 @@ def _get_unary_sql(self, plan):
             return select(clause, from_obj=input)
 
         elif isinstance(plan, algebra.GroupBy):
-            if len(plan.grouping_list) > 0:
+            if (not self.push_grouping) and len(plan.grouping_list) > 0:
                 raise NotImplementedError(
-                    "convert aggregate with grouping to sql -- Myria faster")
+                    """convert aggregate with grouping to sql
+                    -- Myria is faster. If you want to push group by into
+                    SQL use the flag push_sql_grouping""")
             a = [self._convert_expr(cols, e, input_sch)
                  for e in plan.aggregate_list]
             g = [self._convert_expr(cols, e, input_sch)
@@ -169,6 +210,9 @@ def _get_unary_sql(self, plan):
                 return sel
             return sel.group_by(*g)
 
+        elif isinstance(plan, algebra.Distinct):
+            return select(['*'], from_obj=input, distinct=True)
+
         raise NotImplementedError("convert {op} to sql".format(op=type(plan)))
 
     def _get_binary_sql(self, plan):

raco/myrial/optimizer_tests.py

@@ -1,14 +1,18 @@
 import collections
 import random
+import sys
+import re
 
 from raco.algebra import *
 from raco.expression import NamedAttributeRef as AttRef
 from raco.expression import UnnamedAttributeRef as AttIndex
 from raco.expression import StateVar
+from raco.expression import aggregate
 
 from raco.backends.myria import (
     MyriaShuffleConsumer, MyriaShuffleProducer, MyriaHyperShuffleProducer,
-    MyriaBroadcastConsumer, MyriaQueryScan, MyriaSplitConsumer)
+    MyriaBroadcastConsumer, MyriaQueryScan, MyriaSplitConsumer, MyriaDupElim,
+    MyriaGroupBy)
 from raco.backends.myria import (MyriaLeftDeepTreeAlgebra,
                                  MyriaHyperCubeAlgebra)
 from raco.compile import optimize
@@ -1003,3 +1007,170 @@ def test_projecting_join_maintains_partitioning(self):
         # (in general, info could be h($0) && h($2)
         self.assertEquals(pp.partitioning().hash_partitioned,
                           frozenset([AttIndex(0)]))
+
+    def test_no_shuffle_for_partitioned_distinct(self):
+        """Do not shuffle for Distinct if already partitioned"""
+
+        query = """
+        r = scan({part});
+        t = select distinct r.h from r;
+        store(t, OUTPUT);""".format(part=self.part_key)
+
+        lp = self.get_logical_plan(query)
+        pp = self.logical_to_physical(lp)
+
+        # shuffles should be removed and distinct not decomposed into two
+        self.assertEquals(self.get_count(pp, MyriaShuffleConsumer), 0)
+        self.assertEquals(self.get_count(pp, MyriaShuffleProducer), 0)
+        self.assertEquals(self.get_count(pp, MyriaDupElim), 1)
+
+        self.db.evaluate(pp)
+        result = self.db.get_table('OUTPUT')
+        expected = dict([((h,), 1) for _, h, _ in self.part_data])
+        self.assertEquals(result, expected)
+
+    def test_no_shuffle_for_partitioned_groupby(self):
+        """Do not shuffle for groupby if already partitioned"""
+
+        query = """
+        r = scan({part});
+        t = select r.h, MIN(r.i) from r;
+        store(t, OUTPUT);""".format(part=self.part_key)
+
+        lp = self.get_logical_plan(query)
+        pp = self.logical_to_physical(lp)
+
+        # shuffles should be removed and the groupby not decomposed into two
+        self.assertEquals(self.get_count(pp, MyriaShuffleConsumer), 0)
+        self.assertEquals(self.get_count(pp, MyriaShuffleProducer), 0)
+        self.assertEquals(self.get_count(pp, MyriaGroupBy), 1)
+
+    def test_partition_aware_groupby_into_sql(self):
+        """No shuffle for groupby also causes it to be pushed into sql"""
+
+        query = """
+        r = scan({part});
+        t = select r.h, MIN(r.i) from r;
+        store(t, OUTPUT);""".format(part=self.part_key)
+
+        lp = self.get_logical_plan(query)
+        pp = self.logical_to_physical(lp, push_sql=True,
+                                      push_sql_grouping=True)
+
+        # shuffles should be removed and the groupby not decomposed into two
+        self.assertEquals(self.get_count(pp, MyriaShuffleConsumer), 0)
+        self.assertEquals(self.get_count(pp, MyriaShuffleProducer), 0)
+
+        # should be pushed
+        self.assertEquals(self.get_count(pp, MyriaGroupBy), 0)
+        self.assertEquals(self.get_count(pp, MyriaQueryScan), 1)
+
+        self.db.evaluate(pp)
+        result = self.db.get_table('OUTPUT')
+        temp = dict([(h, sys.maxsize) for _, h, _ in self.part_data])
+        for _, h, i in self.part_data:
+            temp[h] = min(temp[h], i)
+        expected = dict(((h, i), 1) for h, i in temp.items())
+
+        self.assertEquals(result, expected)
+
+    def test_partition_aware_distinct_into_sql(self):
+        """No shuffle for distinct also causes it to be pushed into sql"""
+
+        query = """
+        r = scan({part});
+        t = select distinct r.h from r;
+        store(t, OUTPUT);""".format(part=self.part_key)
+
+        lp = self.get_logical_plan(query)
+        pp = self.logical_to_physical(lp, push_sql=True)
+
+        # shuffles should be removed and the groupby not decomposed into two
+        self.assertEquals(self.get_count(pp, MyriaShuffleConsumer), 0)
+        self.assertEquals(self.get_count(pp, MyriaShuffleProducer), 0)
+
+        # should be pushed
+        self.assertEquals(self.get_count(pp, MyriaGroupBy), 0)  # sanity
+        self.assertEquals(self.get_count(pp, MyriaDupElim), 0)
+        self.assertEquals(self.get_count(pp, MyriaQueryScan), 1)
+
+        self.db.evaluate(pp)
+        result = self.db.get_table('OUTPUT')
+        expected = dict([((h,), 1) for _, h, _ in self.part_data])
+        self.assertEquals(result, expected)
+
+    def test_push_half_groupby_into_sql(self):
+        """Push the first group by of decomposed group by into sql"""
+
+        query = """
+        r = scan({part});
+        t = select r.i, MIN(r.h) from r;
+        store(t, OUTPUT);""".format(part=self.part_key)
+
+        lp = self.get_logical_plan(query)
+        pp = self.logical_to_physical(lp, push_sql=True,
+                                      push_sql_grouping=True)
+
+        # wrong partition, so still has shuffle
+        self.assertEquals(self.get_count(pp, MyriaShuffleConsumer), 1)
+        self.assertEquals(self.get_count(pp, MyriaShuffleProducer), 1)
+
+        # one group by should be pushed
+        self.assertEquals(self.get_count(pp, MyriaGroupBy), 1)
+        self.assertEquals(self.get_count(pp, MyriaQueryScan), 1)
+
+        self.db.evaluate(pp)
+        result = self.db.get_table('OUTPUT')
+        temp = dict([(i, sys.maxsize) for _, _, i in self.part_data])
+        for _, h, i in self.part_data:
+            temp[i] = min(temp[i], h)
+        expected = dict(((k, v), 1) for k, v in temp.items())
+
+        self.assertEquals(result, expected)
+
+    def _check_aggregate_functions_pushed(
+            self,
+            func,
+            expected,
+            override=False):
+        if override:
+            agg = func
+        else:
+            agg = "{func}(r.i)".format(func=func)
+
+        query = """
+        r = scan({part});
+        t = select r.h, {agg} from r;
+        store(t, OUTPUT);""".format(part=self.part_key, agg=agg)
+
+        lp = self.get_logical_plan(query)
+        pp = self.logical_to_physical(lp, push_sql=True,
+                                      push_sql_grouping=True)
+
+        self.assertEquals(self.get_count(pp, MyriaQueryScan), 1)
+
+        for op in pp.walk():
+            if isinstance(op, MyriaQueryScan):
+                self.assertTrue(re.search(expected, op.sql))
+
+    def test_aggregate_AVG_pushed(self):
+        """AVG is translated properly for postgresql. This is
+        a function not in SQLAlchemy"""
+        self._check_aggregate_functions_pushed(
+            aggregate.AVG.__name__, 'avg')
+
+    def test_aggregate_STDDEV_pushed(self):
+        """STDEV is translated properly for postgresql. This is
+        a function that is named differently in Raco and postgresql"""
+        self._check_aggregate_functions_pushed(
+            aggregate.STDEV.__name__, 'stddev_samp')
+
+    def test_aggregate_COUNTALL_pushed(self):
+        """COUNTALL is translated properly for postgresql. This is
+        a function that is expressed differently in Raco and postgresql"""
+
+        # MyriaL parses count(*) to Raco COUNTALL. And COUNTALL
+        # should currently (under the no nulls semantics of Raco/Myria)
+        # translate to COUNT(something)
+        self._check_aggregate_functions_pushed(
+            'count(*)', r'count[(][a-zA-Z.]+[)]', True)