Repository: ibis-project/ibis

Issue:
SQLite: add Bucket translation rule



Discussion:

Resolving patch:
diff --git a/ibis/impala/compiler.py b/ibis/impala/compiler.py
--- a/ibis/impala/compiler.py
+++ b/ibis/impala/compiler.py
@@ -866,54 +866,6 @@ def _searched_case(translator, expr):
     return formatter.get_result()
 
 
-def _bucket(translator, expr):
-    import operator
-
-    op = expr.op()
-    stmt = ibis.case()
-
-    if op.closed == 'left':
-        l_cmp = operator.le
-        r_cmp = operator.lt
-    else:
-        l_cmp = operator.lt
-        r_cmp = operator.le
-
-    user_num_buckets = len(op.buckets) - 1
-
-    bucket_id = 0
-    if op.include_under:
-        if user_num_buckets > 0:
-            cmp = operator.lt if op.close_extreme else r_cmp
-        else:
-            cmp = operator.le if op.closed == 'right' else operator.lt
-        stmt = stmt.when(cmp(op.arg, op.buckets[0]), bucket_id)
-        bucket_id += 1
-
-    for j, (lower, upper) in enumerate(zip(op.buckets, op.buckets[1:])):
-        if (op.close_extreme and
-            ((op.closed == 'right' and j == 0) or
-             (op.closed == 'left' and j == (user_num_buckets - 1)))):
-            stmt = stmt.when((lower <= op.arg) & (op.arg <= upper),
-                             bucket_id)
-        else:
-            stmt = stmt.when(l_cmp(lower, op.arg) & r_cmp(op.arg, upper),
-                             bucket_id)
-        bucket_id += 1
-
-    if op.include_over:
-        if user_num_buckets > 0:
-            cmp = operator.lt if op.close_extreme else l_cmp
-        else:
-            cmp = operator.lt if op.closed == 'right' else operator.le
-
-        stmt = stmt.when(cmp(op.buckets[-1], op.arg), bucket_id)
-        bucket_id += 1
-
-    case_expr = stmt.end().name(expr._name)
-    return _searched_case(translator, case_expr)
-
-
 def _category_label(translator, expr):
     op = expr.op()
 
@@ -1228,13 +1180,14 @@ def _value_list(translator, expr):
     ops.MinRank: _subtract_one,
 }
 
+_expr_rewrites = comp.ExprTranslator._rewrites.copy()
 
-_expr_rewrites = {
+_expr_rewrites.update({
     ops.Any: _any_expand,
     ops.All: _all_expand,
     ops.NotAny: _notany_expand,
     ops.NotAll: _notall_expand,
-}
+})
 
 
 _binary_infix_ops = {
@@ -1364,7 +1317,6 @@ def _value_list(translator, expr):
     ops.Contains: _binary_infix_op('IN'),
     ops.NotContains: _binary_infix_op('NOT IN'),
 
-    analytics.Bucket: _bucket,
     analytics.CategoryLabel: _category_label,
 
     ops.SimpleCase: _simple_case,
diff --git a/ibis/sql/alchemy.py b/ibis/sql/alchemy.py
--- a/ibis/sql/alchemy.py
+++ b/ibis/sql/alchemy.py
@@ -281,9 +281,8 @@ def unary(sa_func):
     return fixed_arity(sa_func, 1)
 
 
-_expr_rewrites = {
+_expr_rewrites = comp.ExprTranslator._rewrites.copy()
 
-}
 
 _operation_registry = {
     ops.And: fixed_arity(sql.and_, 2),
diff --git a/ibis/sql/compiler.py b/ibis/sql/compiler.py
--- a/ibis/sql/compiler.py
+++ b/ibis/sql/compiler.py
@@ -17,11 +17,13 @@
 from ibis.compat import lzip
 import ibis.common as com
 import ibis.expr.analysis as L
+import ibis.expr.analytics as analytics
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
 
 import ibis.sql.transforms as transforms
 import ibis.util as util
+import ibis
 
 
 # ---------------------------------------------------------------------
@@ -1025,8 +1027,59 @@ def _key_in(self, key, memo_attr, parent_contexts=False):
         return False
 
 
+def _bucket(expr):
+    import operator
+
+    op = expr.op()
+    stmt = ibis.case()
+
+    if op.closed == 'left':
+        l_cmp = operator.le
+        r_cmp = operator.lt
+    else:
+        l_cmp = operator.lt
+        r_cmp = operator.le
+
+    user_num_buckets = len(op.buckets) - 1
+
+    bucket_id = 0
+    if op.include_under:
+        if user_num_buckets > 0:
+            cmp = operator.lt if op.close_extreme else r_cmp
+        else:
+            cmp = operator.le if op.closed == 'right' else operator.lt
+        stmt = stmt.when(cmp(op.arg, op.buckets[0]), bucket_id)
+        bucket_id += 1
+
+    for j, (lower, upper) in enumerate(zip(op.buckets, op.buckets[1:])):
+        if (op.close_extreme and
+            ((op.closed == 'right' and j == 0) or
+             (op.closed == 'left' and j == (user_num_buckets - 1)))):
+            stmt = stmt.when((lower <= op.arg) & (op.arg <= upper),
+                             bucket_id)
+        else:
+            stmt = stmt.when(l_cmp(lower, op.arg) & r_cmp(op.arg, upper),
+                             bucket_id)
+        bucket_id += 1
+
+    if op.include_over:
+        if user_num_buckets > 0:
+            cmp = operator.lt if op.close_extreme else l_cmp
+        else:
+            cmp = operator.lt if op.closed == 'right' else operator.le
+
+        stmt = stmt.when(cmp(op.buckets[-1], op.arg), bucket_id)
+        bucket_id += 1
+
+    return stmt.end().name(expr._name)
+
+
 class ExprTranslator(object):
 
+    _rewrites = {
+        analytics.Bucket: _bucket,
+    }
+
     def __init__(self, expr, context=None, named=False, permit_subquery=False):
         self.expr = expr
         self.permit_subquery = permit_subquery


Test patch:
diff --git a/ibis/sql/sqlite/tests/test_functions.py b/ibis/sql/sqlite/tests/test_functions.py
--- a/ibis/sql/sqlite/tests/test_functions.py
+++ b/ibis/sql/sqlite/tests/test_functions.py
@@ -206,6 +206,12 @@ def test_misc_builtins_work(self):
         exprs = [
             (d > 20).ifelse(10, -20),
             (d > 20).ifelse(10, -20).abs(),
+
+            # tier and histogram
+            d.bucket([0, 10, 25, 50, 100]),
+            d.bucket([0, 10, 25, 50], include_over=True),
+            d.bucket([0, 10, 25, 50], include_over=True, close_extreme=False),
+            d.bucket([10, 25, 50, 100], include_under=True),
         ]
         self._execute_projection(t, exprs)