Repository: ibis-project/ibis

Issue:
KeyError: 'timestamp' in aggregation with 'where'
Hello, 

Table 'ta' has a column 'ts' which is a timestamp.  This works as expected: 

    ta.group_by('uuid').aggregate(min_date=ta.ts.min()).compile()

But the following produces keyError: 'timestamp'.

    ta.group_by('uuid').aggregate(min_date=ta.ts.min(where=ta.search_level==1)).compile()

Casting to a different data type as below works fine, so it seems to be a problem with timestamps?
   
    (ta.group_by('uuid')
        .aggregate(min_date=ta.ts.cast('string').min(where=ta.search_level==1)).compile())

Here's the full traceback:

```python
ta.group_by('uuid').aggregate(min_date=ta.ts.min(where=ta.search_level==1)).compile()

    KeyError                                  Traceback (most recent call last)
<ipython-input-70-ec191fa92eab> in <module>()
----> 1 ta.group_by('uuid').aggregate(min_date=ta.ts.min(where=ta.search_level==1)).compile()

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/types.py in compile(self, limit)
    155         """
    156         from ibis.client import compile
--> 157         return compile(self, limit=limit)
    158 
    159     def verify(self):

/home/pteehan/.local/lib/python3.5/site-packages/ibis/client.py in compile(expr, limit)
    317 def compile(expr, limit=None):
    318     backend = find_backend(expr)
--> 319     return backend.compile(expr, limit=limit)
    320 
    321 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/client.py in compile(self, expr, params, limit)
    243         """
    244         ast = self._build_ast_ensure_limit(expr, limit)
--> 245         queries = [query.compile() for query in ast.queries]
    246         return queries[0] if len(queries) == 1 else queries
    247 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/client.py in <listcomp>(.0)
    243         """
    244         ast = self._build_ast_ensure_limit(expr, limit)
--> 245         queries = [query.compile() for query in ast.queries]
    246         return queries[0] if len(queries) == 1 else queries
    247 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/impala/compiler.py in compile(self)
     99 
    100         # SELECT
--> 101         select_frag = self.format_select_set()
    102 
    103         # FROM, JOIN, UNION

/home/pteehan/.local/lib/python3.5/site-packages/ibis/impala/compiler.py in format_select_set(self)
    143         for expr in self.select_set:
    144             if isinstance(expr, ir.ValueExpr):
--> 145                 expr_str = self._translate(expr, named=True)
    146             elif isinstance(expr, ir.TableExpr):
    147                 # A * selection, possibly prefixed

/home/pteehan/.local/lib/python3.5/site-packages/ibis/sql/compiler.py in _translate(self, expr, context, named, permit_subquery)
   1302                                      named=named,
   1303                                      permit_subquery=permit_subquery)
-> 1304         return translator.get_result()
   1305 
   1306     def equals(self, other):

/home/pteehan/.local/lib/python3.5/site-packages/ibis/sql/compiler.py in get_result(self)
   1090         Build compiled SQL expression from the bottom up and return as a string
   1091         """
-> 1092         translated = self.translate(self.expr)
   1093         if self._needs_name(self.expr):
   1094             # TODO: this could fail in various ways

/home/pteehan/.local/lib/python3.5/site-packages/ibis/sql/compiler.py in translate(self, expr)
   1129         elif type(op) in self._registry:
   1130             formatter = self._registry[type(op)]
-> 1131             return formatter(self, expr)
   1132         else:
   1133             raise com.TranslationError('No translator rule for {0}'

/home/pteehan/.local/lib/python3.5/site-packages/ibis/impala/compiler.py in formatter(translator, expr)
    654         arg, where = op.args[:2]
    655 
--> 656         return _reduction_format(translator, func_name, arg, where)
    657     return formatter
    658 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/impala/compiler.py in _reduction_format(translator, func_name, arg, where)
    639 def _reduction_format(translator, func_name, arg, where):
    640     if where is not None:
--> 641         case = where.ifelse(arg, ibis.NA)
    642         arg = translator.translate(case)
    643     else:

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/api.py in ifelse(arg, true_expr, false_expr)
   1166     # must be used.
   1167     case = _ops.SearchedCaseBuilder()
-> 1168     return case.when(arg, true_expr).else_(false_expr).end()
   1169 
   1170 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/operations.py in end(self)
   1272 
   1273         op = SearchedCase(self.cases, self.results, default)
-> 1274         return op.to_expr()
   1275 
   1276 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/types.py in to_expr(self)
    272 
    273     def to_expr(self):
--> 274         klass = self.output_type()
    275         return klass(self)
    276 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/operations.py in output_type(self)
   1320         cases, results, default = self.args
   1321         out_exprs = results + [default]
-> 1322         typename = rules.highest_precedence_type(out_exprs)
   1323         return rules.shape_like_args(cases, typename)
   1324 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/rules.py in highest_precedence_type(exprs)
    114     # the selected highest precedence type
    115     selector = _TypePrecedence(exprs)
--> 116     return selector.get_result()
    117 
    118 

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/rules.py in get_result(self)
    153 
    154     def get_result(self):
--> 155         highest_type = self._get_highest_type()
    156         self._check_casts(highest_type)
    157         return highest_type

/home/pteehan/.local/lib/python3.5/site-packages/ibis/expr/rules.py in _get_highest_type(self)
    166             if not v:
    167                 continue
--> 168             score = self._precedence[k.name()]
    169 
    170             scores.append((score, k))

KeyError: 'timestamp'
````


Discussion:

Resolving patch:
diff --git a/ibis/expr/rules.py b/ibis/expr/rules.py
--- a/ibis/expr/rules.py
+++ b/ibis/expr/rules.py
@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import defaultdict
+from collections import Counter
+
 import operator
 
+from toolz import first
+
 from ibis.common import IbisTypeError
 from ibis.compat import py_string
 import ibis.expr.datatypes as dt
@@ -108,76 +111,55 @@ def _get_type(self):
             raise NotImplementedError
 
 
+# Impala type precedence (more complex in other database implementations)
+# timestamp
+# double
+# float
+# decimal
+# bigint
+# int
+# smallint
+# tinyint
+# boolean
+# string
+
+
+_TYPE_PRECEDENCE = {
+    'timestamp': 10,
+    'double': 9,
+    'float': 8,
+    'decimal': 7,
+    'int64': 6,
+    'int32': 5,
+    'int16': 4,
+    'int8': 3,
+    'boolean': 2,
+    'string': 1,
+    'null': 0
+}
+
+
 def highest_precedence_type(exprs):
     # Return the highest precedence type from the passed expressions. Also
     # verifies that there are valid implicit casts between any of the types and
     # the selected highest precedence type
-    selector = _TypePrecedence(exprs)
-    return selector.get_result()
-
-
-class _TypePrecedence(object):
-    # Impala type precedence (more complex in other database implementations)
-    # timestamp
-    # double
-    # float
-    # decimal
-    # bigint
-    # int
-    # smallint
-    # tinyint
-    # boolean
-    # string
-
-    _precedence = {
-        'double': 9,
-        'float': 8,
-        'decimal': 7,
-        'int64': 6,
-        'int32': 5,
-        'int16': 4,
-        'int8': 3,
-        'boolean': 2,
-        'string': 1,
-        'null': 0
-    }
-
-    def __init__(self, exprs):
-        self.exprs = exprs
-
-        if len(exprs) == 0:
-            raise ValueError('Must pass at least one expression')
-
-        self.type_counts = defaultdict(lambda: 0)
-        self._count_types()
-
-    def get_result(self):
-        highest_type = self._get_highest_type()
-        self._check_casts(highest_type)
-        return highest_type
-
-    def _count_types(self):
-        for expr in self.exprs:
-            self.type_counts[expr.type()] += 1
-
-    def _get_highest_type(self):
-        scores = []
-        for k, v in self.type_counts.items():
-            if not v:
-                continue
-            score = self._precedence[k.name.lower()]
-
-            scores.append((score, k))
-
-        scores.sort()
-        return scores[-1][1]
-
-    def _check_casts(self, typename):
-        for expr in self.exprs:
-            if not expr._can_cast_implicit(typename):
-                raise ValueError('Expression with type {0} cannot be '
-                                 'implicitly casted to {1}'
-                                 .format(expr.type(), typename))
+    if not exprs:
+        raise ValueError('Must pass at least one expression')
+
+    type_counts = Counter(expr.type() for expr in exprs)
+    scores = (
+        (_TYPE_PRECEDENCE[k.name.lower()], k) for k, v in type_counts.items()
+    )
+    _, highest_type = max(scores, key=first)
+
+    for expr in exprs:
+        if not expr._can_cast_implicit(highest_type):
+            raise TypeError(
+                'Expression with type {0} cannot be implicitly casted to {1}'
+                .format(expr.type(), highest_type)
+            )
+
+    return highest_type
 
 
 def _int_bounds_promotion(ltype, rtype, op):


Test patch:
diff --git a/ibis/expr/tests/test_datatypes.py b/ibis/expr/tests/test_datatypes.py
--- a/ibis/expr/tests/test_datatypes.py
+++ b/ibis/expr/tests/test_datatypes.py
@@ -1,6 +1,7 @@
 import pytest
 import ibis
 from ibis.expr import datatypes as dt
+from ibis.expr.rules import highest_precedence_type
 
 
 def test_array():
@@ -185,3 +186,9 @@ def test_whole_schema():
         ],
     )
     assert customers.schema() == expected
+
+
+def test_precedence_with_no_arguments():
+    with pytest.raises(ValueError) as e:
+        highest_precedence_type([])
+    assert str(e.value) == 'Must pass at least one expression'
diff --git a/ibis/expr/tests/test_timestamp.py b/ibis/expr/tests/test_timestamp.py
--- a/ibis/expr/tests/test_timestamp.py
+++ b/ibis/expr/tests/test_timestamp.py
@@ -13,11 +13,13 @@
 # limitations under the License.
 
 import pandas as pd
+from datetime import datetime
 
 import ibis
 import ibis.expr.api as api
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
+from ibis.expr.rules import highest_precedence_type
 
 from ibis.expr.tests.mocks import MockConnection
 from ibis.compat import unittest
@@ -107,3 +109,10 @@ def test_comparisons_pandas_timestamp(self):
         # op = expr2.op()
         # assert isinstance(op, ops.Greater)
         # assert isinstance(op.right, ir.TimestampScalar)
+
+
+def test_timestamp_precedence():
+    ts = ibis.literal(datetime.now())
+    null_ts = ibis.NA
+    highest_type = highest_precedence_type([ts, null_ts])
+    assert highest_type == 'timestamp'
diff --git a/ibis/impala/tests/test_exprs.py b/ibis/impala/tests/test_exprs.py
--- a/ibis/impala/tests/test_exprs.py
+++ b/ibis/impala/tests/test_exprs.py
@@ -1564,3 +1564,24 @@ def test_unions_with_ctes(self):
 
         expr = join1.union(join2)
         self.con.explain(expr)
+
+
+def test_where_with_timestamp():
+    t = ibis.table(
+        [
+            ('uuid', 'string'),
+            ('ts', 'timestamp'),
+            ('search_level', 'int64'),
+        ],
+        name='t'
+    )
+    expr = t.group_by(t.uuid).aggregate(
+        min_date=t.ts.min(where=t.search_level == 1)
+    )
+    result = ibis.impala.compile(expr)
+    expected = """\
+SELECT `uuid`,
+       min(CASE WHEN `search_level` = 1 THEN `ts` ELSE NULL END) AS `min_date`
+FROM t
+GROUP BY 1"""
+    assert result == expected