Coerce expression

blaze/compute/numpy.py

@@ -12,7 +12,7 @@
 from ..expr import std, var, count, nunique, Summary, IsIn
 from ..expr import BinOp, UnaryOp, USub, Not, nelements, Repeat, Concat, Interp
 from ..expr import UTCFromTimestamp, DateTimeTruncate
-from ..expr import Transpose, TensorDot
+from ..expr import Transpose, TensorDot, Coerce
 from ..utils import keywords
 
 from .core import base, compute
@@ -364,3 +364,8 @@ def join_ndarray(expr, lhs, rhs, **kwargs):
     if isinstance(rhs, np.ndarray):
         rhs = DataFrame(rhs)
     return compute_up(expr, lhs, rhs, **kwargs)
+
+
+@dispatch(Coerce, np.ndarray)
+def compute_up(expr, data, **kwargs):
+    return data.astype(to_numpy_dtype(expr.schema))

blaze/compute/sql.py

@@ -42,15 +42,15 @@
 
 from multipledispatch import MDNotImplementedError
 
-from odo.backends.sql import metadata_of_engine
+from odo.backends.sql import metadata_of_engine, dshape_to_alchemy
 
 from ..dispatch import dispatch
 
 from .core import compute_up, compute, base
 
 from ..expr import Projection, Selection, Field, Broadcast, Expr, IsIn, Slice
-from ..expr import (BinOp, UnaryOp, USub, Join, mean, var, std, Reduction,
-                    count, FloorDiv, UnaryStringFunction, strlen, DateTime)
+from ..expr import (BinOp, UnaryOp, Join, mean, var, std, Reduction, count,
+                    FloorDiv, UnaryStringFunction, strlen, DateTime, Coerce)
 from ..expr import nunique, Distinct, By, Sort, Head, Label, ReLabel, Merge
 from ..expr import common_subexpression, Summary, Like, nelements
 
@@ -838,7 +838,7 @@ def compute_up(expr, data, **kwargs):
         stop = start + 1
     else:
         raise TypeError('type %r not supported for slicing wih SQL backend'
-                         % type(index).__name__)
+                        % type(index).__name__)
 
     warnings.warn('The order of the result set from a Slice expression '
                   'computed against the SQL backend is not deterministic.')
@@ -847,3 +847,8 @@ def compute_up(expr, data, **kwargs):
         return select(data).offset(start)
     else:
         return select(data).offset(start).limit(stop - start)
+
+
+@dispatch(Coerce, ColumnElement)
+def compute_up(expr, data, **kwargs):
+    return sa.cast(data, dshape_to_alchemy(expr.to)).label(expr._name)

blaze/compute/tests/test_numpy_compute.py

@@ -467,3 +467,10 @@ def test_timedelta_arith():
     sym = symbol('s', discover(dates))
     assert (compute(sym + delta, dates) == dates + delta).all()
     assert (compute(sym - delta, dates) == dates - delta).all()
+
+
+def test_coerce():
+    x = np.arange(1, 3)
+    s = symbol('s', discover(x))
+    np.testing.assert_array_equal(compute(s.coerce('float64'), x),
+                                  np.arange(1.0, 3.0))

blaze/compute/tests/test_sql_compute.py

@@ -1682,3 +1682,23 @@ def test_insert_from_subselect(sql_with_float):
         odo(sql_with_float, pd.DataFrame).iloc[2:].reset_index(drop=True),
         pd.DataFrame([{'c': 1.0}, {'c': 2.0}]),
     )
+
+
+def test_coerce(sql):
+    n = sql.name
+    t = symbol(n, discover(sql))
+    expr = t.B.coerce(to='float64')
+
+    # "B" because we're capitalized
+    expected = 'SELECT cast({t}."B" AS FLOAT(53)) AS "B" FROM {t}'.format(t=n)
+    result = compute(expr, sql)
+    assert normalize(str(result)) == normalize(expected)
+
+
+def test_coerce_bool_and_sum(sql):
+    n = sql.name
+    t = symbol(n, discover(sql))
+    expr = (t.B > 1.0).coerce(to='int32').sum()
+    result = compute(expr, sql).scalar()
+    expected = odo(compute(t.B, sql), pd.Series).gt(1).sum()
+    assert result == expected

blaze/expr/expressions.py

@@ -22,7 +22,7 @@
 __all__ = ['Expr', 'ElemWise', 'Field', 'Symbol', 'discover', 'Projection',
            'projection', 'Selection', 'selection', 'Label', 'label', 'Map',
            'ReLabel', 'relabel', 'Apply', 'Slice', 'shape', 'ndim', 'label',
-           'symbol']
+           'symbol', 'Coerce', 'coerce']
 
 
 _attr_cache = dict()
@@ -705,6 +705,21 @@ def dshape(self):
         return dshape(self._dshape)
 
 
+class Coerce(Expr):
+    __slots__ = '_hash', '_child', 'to'
+
+    @property
+    def schema(self):
+        return self.to
+
+    @property
+    def dshape(self):
+        return DataShape(*(self._child.shape + (self.schema,)))
+
+    def __str__(self):
+        return '%s.coerce(to=%r)' % (self._child, str(self.schema))
+
+
 def apply(expr, func, dshape, splittable=False):
     return Apply(expr, func, datashape.dshape(dshape), splittable)
 
@@ -757,13 +772,18 @@ def ndim(expr):
     return len(shape(expr))
 
 
+def coerce(expr, to):
+    return Coerce(expr, dshape(to) if isinstance(to, _strtypes) else to)
+
+
 dshape_method_list.extend([
     (lambda ds: True, set([apply])),
     (iscollection, set([shape, ndim])),
+    (lambda ds: iscollection(ds) and isscalar(ds.measure), set([coerce]))
 ])
 
 schema_method_list.extend([
-    (isscalar, set([label, relabel])),
+    (isscalar, set([label, relabel, coerce])),
     (isrecord, set([relabel])),
 ])
 

blaze/expr/tests/test_expr.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from datashape import dshape, var, datetime_
+from datashape import dshape, var, datetime_, float32
 
 from blaze.expr import symbol, label, Field
 
@@ -136,3 +136,15 @@ def test_hash_to_different_values():
     from blaze.expr.expressions import _attr_cache
     assert (expr, '_and') in _attr_cache
     assert (expr2, '_and') in _attr_cache
+
+
+def test_coerce():
+    s = symbol('s', var * float32)
+    assert str(s.coerce('int64')) == "s.coerce(to='int64')"
+
+
+@pytest.mark.xfail(raises=AttributeError, reason='Should this be valid?')
+def test_coerce_record():
+    s = symbol('s', 'var * {a: int64, b: float64}')
+    expr = s.coerce('{a: float64, b: float32}')
+    assert str(expr) == "s.coerce(to='{a: float64, b: float32}')"