def test_join_option_types_outer():
a = symbol('a', 'var * {x: ?int}')
b = symbol('b', 'var * {x: int}')
assert (join(a, b, 'x', how='outer').dshape ==
join(b, a, 'x', how='outer').dshape ==
dshape('var * {x: ?int}'))
def test_clean_complex_join():
metadata = sa.MetaData()
lhs = sa.Table('amounts', metadata,
sa.Column('name', sa.String),
sa.Column('amount', sa.Integer))
rhs = sa.Table('ids', metadata,
sa.Column('name', sa.String),
sa.Column('id', sa.Integer))
L = symbol('L', 'var * {name: string, amount: int}')
R = symbol('R', 'var * {name: string, id: int}')
joined = join(L[L.amount > 0], R, 'name')
result = compute(joined, {L: lhs, R: rhs})
expected1 = """
SELECT amounts.name, amounts.amount, ids.id
FROM amounts JOIN ids ON amounts.name = ids.name
WHERE amounts.amount > :amount_1"""
expected2 = """
SELECT alias.name, alias.amount, ids.id
FROM (SELECT amounts.name AS name, amounts.amount AS amount
FROM amounts
WHERE amounts.amount > :amount_1) AS alias
JOIN ids ON alias.name = ids.name"""
assert (normalize(str(result)) == normalize(expected1) or
normalize(str(result)) == normalize(expected2))
def test_join_by_arcs():
df_idx = DataFrame([['A', 1],
['B', 2],
['C', 3]],
columns=['name', 'node_id'])
df_arc = DataFrame([[1, 3],
[2, 3],
[3, 1]],
columns=['node_out', 'node_id'])
t_idx = symbol('t_idx', 'var * {name: string, node_id: int32}')
t_arc = symbol('t_arc', 'var * {node_out: int32, node_id: int32}')
joined = join(t_arc, t_idx, "node_id")
want = by(joined['name'], count=joined['node_id'].count())
result = compute(want, {t_arc: df_arc, t_idx: df_idx})
result_pandas = pd.merge(df_arc, df_idx, on='node_id')
gb = result_pandas.groupby('name')
expected = gb.node_id.count().reset_index().rename(columns={
'node_id': 'count'
})
tm.assert_frame_equal(result, expected)
assert list(result.columns) == ['name', 'count']
def test_join_complex_clean():
metadata = sa.MetaData()
name = sa.Table('name', metadata,
sa.Column('id', sa.Integer),
sa.Column('name', sa.String),
)
city = sa.Table('place', metadata,
sa.Column('id', sa.Integer),
sa.Column('city', sa.String),
sa.Column('country', sa.String),
)
tname = symbol('name', discover(name))
tcity = symbol('city', discover(city))
ns = {tname: name, tcity: city}
expr = join(tname[tname.id > 0], tcity, 'id')
result = compute(expr, ns)
expected1 = """
SELECT name.id, name.name, place.city, place.country
FROM name JOIN place ON name.id = place.id
WHERE name.id > :id_1"""
expected2 = """
SELECT alias.id, alias.name, place.city, place.country
FROM (SELECT name.id as id, name.name AS name
FROM name
WHERE name.id > :id_1) AS alias
JOIN place ON alias.id = place.id"""
assert (normalize(str(result)) == normalize(expected1) or
normalize(str(result)) == normalize(expected2))
def test_multi_column_join():
metadata = sa.MetaData()
lhs = sa.Table('aaa', metadata,
sa.Column('x', sa.Integer),
sa.Column('y', sa.Integer),
sa.Column('z', sa.Integer))
rhs = sa.Table('bbb', metadata,
sa.Column('w', sa.Integer),
sa.Column('x', sa.Integer),
sa.Column('y', sa.Integer))
L = symbol('L', 'var * {x: int, y: int, z: int}')
R = symbol('R', 'var * {w: int, x: int, y: int}')
joined = join(L, R, ['x', 'y'])
expected = lhs.join(rhs, (lhs.c.x == rhs.c.x)
& (lhs.c.y == rhs.c.y))
expected = select(list(unique(expected.columns, key=lambda c:
c.name))).select_from(expected)
result = compute(joined, {L: lhs, R: rhs})
assert str(result) == str(expected)
assert str(select(result)) == str(select(expected))
# Schemas match
print(result.c.keys())
print(joined.fields)
assert list(result.c.keys()) == list(joined.fields)
def test_coalesce():
data = pd.Series([0, None, 1, None, 2, None], dtype=object)
s = symbol('s', 'var * ?int')
t = symbol('t', 'int')
u = symbol('u', '?int')
v = symbol('v', 'var * int')
w = symbol('w', 'var * ?int')
# array to scalar
tm.assert_series_equal(
compute(coalesce(s, t), {s: data, t: -1}),
pd.Series([0, -1, 1, -1, 2, -1], dtype=object),
)
# array to scalar with NULL
tm.assert_series_equal(
compute(coalesce(s, u), {s: data, u: None}),
pd.Series([0, None, 1, None, 2, None], dtype=object),
)
# array to array
tm.assert_series_equal(
compute(coalesce(s, v), {
s: data, v: np.array([-1, -2, -3, -4, -5, -6]),
}),
pd.Series([0, -2, 1, -4, 2, -6], dtype=object),
)
# array to array with NULL
tm.assert_series_equal(
compute(coalesce(s, w), {
s: data, w: np.array([-1, None, -3, -4, -5, -6]),
}),
pd.Series([0, None, 1, -4, 2, -6], dtype=object),
)
def test_selection_of_join():
metadata = sa.MetaData()
name = sa.Table('name', metadata,
sa.Column('id', sa.Integer),
sa.Column('name', sa.String),
)
city = sa.Table('place', metadata,
sa.Column('id', sa.Integer),
sa.Column('city', sa.String),
sa.Column('country', sa.String),
)
tname = symbol('name', discover(name))
tcity = symbol('city', discover(city))
ns = {tname: name, tcity: city}
j = join(tname, tcity, 'id')
expr = j[j.city == 'NYC'].name
result = compute(expr, ns)
assert normalize(str(result)) == normalize("""
SELECT name.name
FROM name JOIN place ON name.id = place.id
WHERE place.city = :city_1""")
def test_multi_column_join():
left = [(1, 2, 3),
(2, 3, 4),
(1, 3, 5)]
left = DataFrame(left, columns=['x', 'y', 'z'])
right = [(1, 2, 30),
(1, 3, 50),
(1, 3, 150)]
right = DataFrame(right, columns=['x', 'y', 'w'])
L = symbol('L', 'var * {x: int, y: int, z: int}')
R = symbol('R', 'var * {x: int, y: int, w: int}')
j = join(L, R, ['x', 'y'])
expected = [(1, 2, 3, 30),
(1, 3, 5, 50),
(1, 3, 5, 150)]
expected = DataFrame(expected, columns=['x', 'y', 'z', 'w'])
result = compute(j, {L: left, R: right})
print(result)
assert str(result) == str(expected)
assert list(result.columns) == list(j.fields)
def test_outer_join():
left = [(1, 'Alice', 100),
(2, 'Bob', 200),
(4, 'Dennis', 400)]
right = [('NYC', 1),
('Boston', 1),
('LA', 3),
('Moscow', 4)]
L = symbol('L', 'var * {id: int, name: string, amount: real}')
R = symbol('R', 'var * {city: string, id: int}')
assert set(compute(join(L, R), {L: left, R: right})) == set(
[(1, 'Alice', 100, 'NYC'),
(1, 'Alice', 100, 'Boston'),
(4, 'Dennis', 400, 'Moscow')])
assert set(compute(join(L, R, how='left'), {L: left, R: right})) == set(
[(1, 'Alice', 100, 'NYC'),
(1, 'Alice', 100, 'Boston'),
(2, 'Bob', 200, None),
(4, 'Dennis', 400, 'Moscow')])
assert set(compute(join(L, R, how='right'), {L: left, R: right})) == set(
[(1, 'Alice', 100, 'NYC'),
(1, 'Alice', 100, 'Boston'),
(3, None, None, 'LA'),
(4, 'Dennis', 400, 'Moscow')])
assert set(compute(join(L, R, how='outer'), {L: left, R: right})) == set(
[(1, 'Alice', 100, 'NYC'),
(1, 'Alice', 100, 'Boston'),
(2, 'Bob', 200, None),
(3, None, None, 'LA'),
(4, 'Dennis', 400, 'Moscow')])
def test_basic():
expr = (x + y) * 3
assert eval(str(expr)).isidentical(expr)
assert expr.isidentical(
Mult(Add(symbol('x', 'real'), symbol('y', 'real')), 3),
)
def test_join_on_single_column():
a = symbol('a', 'var * {x: int, y: int, z: int}')
b = symbol('b', 'var * {x: int, y: int, w: int}')
expr = join(a, b.x)
assert expr.on_right == 'x'
def test_concat_arr():
a = symbol('a', '3 * int32')
b = symbol('b', '5 * int32')
v = symbol('v', 'var * int32')
assert concat(a, b).dshape == dshape('8 * int32')
assert concat(a, v).dshape == dshape('var * int32')
def test_concat_table():
a = symbol('a', '3 * {a: int32, b: int32}')
b = symbol('a', '5 * {a: int32, b: int32}')
v = symbol('v', 'var * {a: int32, b: int32}')
assert concat(a, b).dshape == dshape('8 * {a: int32, b: int32}')
assert concat(a, v).dshape == dshape('var * {a: int32, b: int32}')
def test_symbol_name():
t = symbol('t', '10 * {people: string, amount: int}')
r = symbol('r', 'var * int64')
with pytest.raises(AttributeError):
t.name
with pytest.raises(AttributeError):
r.name
def test_eq():
assert symbol('t', 'var * {a: string, b: int}').isidentical(
symbol('t', 'var * {a: string, b: int}'),
)
assert not symbol('t', 'var * {b: string, a: int}').isidentical(
symbol('t', 'var * {a: string, b: int}'),
)
def test_bottom_up_until_type_break():
s = symbol('s', 'var * {name: string, amount: int}')
data = np.array([('Alice', 100), ('Bob', 200), ('Charlie', 300)],
dtype=[('name', 'S7'), ('amount', 'i4')])
e = (s.amount + 1).distinct()
expr, scope = bottom_up_until_type_break(e, {s: data})
amount = symbol('amount', 'var * real', token=1)
assert expr.isidentical(amount)
assert len(scope) == 1
assert amount in scope
assert (scope[amount] == np.array([101, 201, 301], dtype='i4')).all()
# This computation has a type change midstream, so we stop and get the
# unfinished computation.
e = s.amount.sum() + 1
expr, scope = bottom_up_until_type_break(e, {s: data})
amount_sum = symbol('amount_sum', 'int')
assert expr.isidentical(amount_sum + 1)
assert len(scope) == 1
assert amount_sum in scope
assert scope[amount_sum] == 600
# ensure that we work on binops with one child
x = symbol('x', 'real')
expr, scope = bottom_up_until_type_break(x + x, {x: 1})
assert len(scope) == 1
x2 = list(scope.keys())[0]
assert isinstance(x2, Symbol)
assert isinstance(expr, Symbol)
assert scope[x2] == 2
def test_multi_column_join():
left = [(1, 2, 3),
(2, 3, 4),
(1, 3, 5)]
left = DataFrame(left, columns=['x', 'y', 'z'])
right = [(1, 2, 30),
(1, 3, 50),
(1, 3, 150)]
right = DataFrame(right, columns=['x', 'y', 'w'])
lsym = symbol('lsym', 'var * {x: int, y: int, z: int}')
rsym = symbol('rsym', 'var * {x: int, y: int, w: int}')
j = join(lsym, rsym, ['x', 'y'])
expected = [(1, 2, 3, 30),
(1, 3, 5, 50),
(1, 3, 5, 150)]
expected = DataFrame(expected, columns=['x', 'y', 'z', 'w'])
result = compute(j, {lsym: left, rsym: right})
print(result)
tm.assert_frame_equal(result, expected)
assert list(result.columns) == list(j.fields)
def test_concat_arr():
s_data = np.arange(15)
t_data = np.arange(15, 30)
s = symbol("s", discover(s_data))
t = symbol("t", discover(t_data))
assert (compute(concat(s, t), {s: s_data, t: t_data}) == np.arange(30)).all()
def test_selection_inner_inputs():
s_data = np.arange(5).reshape(5, 1)
t_data = np.arange(5).reshape(5, 1)
s = symbol("s", "var * {a: int64}")
t = symbol("t", "var * {a: int64}")
assert (compute(s[s.a == t.a], {s: s_data, t: t_data}) == s_data).all()
def test_concat_axis_too_great():
a = symbol('a', '3 * 5 * int32')
b = symbol('b', '3 * 5 * int32')
with pytest.raises(ValueError) as excinfo:
concat(a, b, axis=2)
assert "must be in range: [0, 2)" in str(excinfo.value)
def test_broadcast_compute_against_numbers_and_arrays():
A = symbol('A', '5 * float32')
a = symbol('a', 'float32')
b = symbol('b', 'float32')
x = np.arange(5, dtype='f4')
expr = Broadcast((A, b), (a, b), a + b)
result = compute(expr, {A: x, b: 10})
assert eq(result, x + 10)
def test_broadcast_compute_against_numbers_and_arrays():
A = symbol("A", "5 * float32")
a = symbol("a", "float32")
b = symbol("b", "float32")
x = np.arange(5, dtype="f4")
expr = Broadcast((A, b), (a, b), a + b)
result = compute(expr, {A: x, b: 10})
assert eq(result, x + 10)
def test_multi_dataset_broadcast_with_Record_types():
x = symbol('x', '3 * {p: int, q: int}')
y = symbol('y', '3 * int')
a = [(1, 1), (2, 2), (3, 3)]
b = [10, 20, 30]
assert list(compute(x.p + x.q + y, {x: iter(a), y: iter(b)})) == [12, 24, 36]
def test_agg_shape_in_tabular_case_with_explicit_chunk():
t = symbol('t', '1000 * {name: string, amount: int, id: int}')
c = symbol('chunk', 100 * t.schema)
expr = by(t.name, total=t.amount.sum())
(chunk, chunk_expr), (agg, agg_expr) = split(t, expr, chunk=c)
assert agg.dshape == dshape('var * {name: string, total: int64}')
def test_element():
x = symbol('x', '5 * 3 * float32')
assert isscalar(x[1, 2].dshape)
assert x[1, 2].dshape == dshape('float32')
assert str(x[1, 2]) == 'x[1, 2]'
x = symbol('x', '5 * float32')
assert isscalar(x[3].dshape)
def test_multi_dataset_broadcast():
x = symbol('x', '3 * int')
y = symbol('y', '3 * int')
a = [1, 2, 3]
b = [10, 20, 30]
assert list(compute(x + y, {x: a, y: b})) == [11, 22, 33]
assert list(compute(2*x + (y + 1), {x: a, y: b})) == [13, 25, 37]
def test_truncate_datetime():
s = symbol('x', 'datetime')
assert compute(s.truncate(2, 'days'), datetime(2002, 1, 3, 12, 30)) ==\
date(2002, 1, 2)
s = symbol('x', 'var * datetime')
assert list(compute(s.truncate(2, 'days'),
[datetime(2002, 1, 3, 12, 30)])) ==\
[date(2002, 1, 2)]
def test_greatest_mixed(dtype):
s_data = np.array([2, 1], dtype=dtype)
t_data = np.array([1, 2], dtype=dtype)
s = symbol('s', discover(s_data))
t = symbol('t', discover(t_data))
expr = greatest(s, t)
result = compute(expr, {s: s_data, t: t_data})
expected = np.maximum(s_data, t_data)
assert np.all(result == expected)
def test_greatest(dtype):
s_data = np.arange(15, dtype=dtype).reshape(5, 3)
t_data = np.arange(15, 30, dtype=dtype).reshape(5, 3)
s = symbol('s', discover(s_data))
t = symbol('t', discover(t_data))
expr = greatest(s, t)
result = compute(expr, {s: s_data, t: t_data})
expected = np.maximum(s_data, t_data)
assert np.all(result == expected)
def test_floating_binary_math(func, kwargs):
s_data = np.arange(15).reshape(5, 3)
t_data = np.arange(15, 30).reshape(5, 3)
s = symbol('s', discover(s_data))
t = symbol('t', discover(t_data))
scope = {s: s_data, t: t_data}
result = compute(getattr(blaze, func)(s, t), scope, **kwargs)
expected = getattr(np, binary_name_map.get(func, func))(s_data, t_data)
np.testing.assert_allclose(result, expected)
def test_mixed(recdata):
s = symbol('s', discover(recdata))
expr = (s.x + 1).sum(axis=1)
assert eq(compute(expr, recdata), compute(expr, rec))
def test_pre_compute_on_small_csv_gives_dataframe():
csv = CSV(example('iris.csv'))
s = symbol('s', discover(csv))
assert isinstance(pre_compute(s.species, csv), (Series, DataFrame))
def test_nrows_3D_records(recdata):
s = symbol('s', discover(recdata))
assert not hasattr(s, 'nrows')
a = symbol('a', 'int')
expr = tdata.name[tdata.balance > a]
assert expr_repr(expr) == 'data[data.balance > a].name'
def test_isidentical_regr():
# regression test for #1387
tdata = np.array([(np.nan, ), (np.nan, )], dtype=[('a', 'float64')])
ds = data(tdata)
assert ds.a.isidentical(ds.a)
@pytest.mark.parametrize(
'data,dshape,exp_type',
[
(1, symbol('x', 'int').dshape, int),
# test 1-d to series
(into(da.core.Array, [1, 2],
chunks=(10, )), dshape('2 * int'), pd.Series),
# test 2-d tabular to dataframe
(into(da.core.Array, [{
'a': 1,
'b': 2
}, {
'a': 3,
'b': 4
}],
chunks=(10, 10)), dshape('2 * {a: int, b: int}'), pd.DataFrame),
# test 2-d non tabular to ndarray
(into(da.core.Array, [[1, 2], [3, 4]],
chunks=(10, 10)), dshape('2 * 2 * int'), np.ndarray)
def test_concretehead_failure():
t = symbol('t', 'var * {x:int, y:int}')
d = t[t['x'] > 100]
with pytest.raises(ValueError):
concrete_head(d)
def test_arithmetic_and_then_slicing(data):
s = symbol('s', discover(data))
assert eq(compute((2 * s + 1)[0], data, pre_compute=False), 2 * x[0] + 1)
def test_optimize_slicing(data):
a = symbol('a', discover(data))
b = symbol('b', discover(data))
assert optimize((a + 1)[:3], data).isidentical(a[:3] + 1)
assert optimize((a + b)[:3], data).isidentical(a[:3] + b[:3])
def test_series_columnwise():
s = Series([1, 2, 3], name='a')
t = symbol('t', 'var * {a: int64}')
result = compute(t.a + 1, s)
assert_series_equal(s + 1, result)
def test_field_on_series():
expr = symbol('s', 'var * int')
data = Series([1, 2, 3, 4], name='s')
assert_series_equal(compute(expr.s, data), data)
def test_compute_client_with_multiple_datasets():
c = bz_data('blaze://localhost:6363')
s = symbol('s', discover(c))
assert compute(s.accounts.amount.sum() + s.accounts2.amount.sum(),
{s: c}) == 600
def test_resources_fail():
t = symbol('t', 'var * {x: int, y: int}')
d = t[t['x'] > 100]
with pytest.raises(ValueError):
compute(d)
def test_pre_compute_on_large_csv_gives_chunked_reader():
csv = CSV(example('iris.csv'))
s = symbol('s', discover(csv))
assert isinstance(pre_compute(s.species, csv, comfortable_memory=10),
(chunks(pd.DataFrame), pd.io.parsers.TextFileReader))
def test_pre_compute_with_head_on_large_csv_yields_iterator():
csv = CSV(example('iris.csv'))
s = symbol('s', discover(csv))
assert isinstance(
pre_compute(s.species.head(), csv, comfortable_memory=10), Iterator)
def test_optimize_slicing_on_file(file):
f = symbol('f', discover(file))
assert optimize((f.x + 1)[:5], file).isidentical(f.x[:5] + 1)
def test_arithmetic_on_small_array_from_file(file):
""" Want to make sure that we call pre_compute on Dataset
Even when it's not the leaf data input. """
s = symbol('s', discover(file))
assert eq(compute(s.x + 1, file), x + 1)
def test_repr_html_on_no_resources_symbol():
t = symbol('t', '5 * {id: int, name: string, balance: int}')
assert to_html(t) == 't'
def test_compute_chunks_on_single_csv():
csv = CSV(example('iris.csv'))
s = symbol('s', discover(csv))
expr = s.sepal_length.max()
assert compute(expr, {s: csv}, comfortable_memory=10, chunksize=50) == 7.9
def test_field_access_on_group(file):
s = symbol('s', '{x: 20 * 24 * float32}')
d = compute(s.x, file['/'])
# assert isinstance(d, h5py.Dataset)
assert eq(d[:], x)
def test_summary_with_multiple_children():
t = symbol('t', 'var * {x: int, y: int, z: int}')
assert summary(a=t.x.sum() + t.y.sum())._child.isidentical(t)
def test_custom_expressions():
ec = Client('localhost:6363')
t = symbol('t', discover(ec))
assert list(map(tuple, compute(CustomExpr(t.accounts),
ec))) == into(list, df)
def test_pre_compute_doesnt_collapse_slices(data):
s = symbol('s', discover(data))
assert pre_compute(s[:5], data) is data
def test_axis_kwarg_is_normalized_to_tuple():
x = symbol('x', '5 * 3 * float32')
exprs = [x.sum(), x.sum(axis=1), x.sum(axis=[1]), x.std(), x.mean(axis=1)]
for expr in exprs:
assert isinstance(expr.axis, tuple)
def test_arithmetic_on_small_array(data):
s = symbol('s', discover(data))
assert eq(compute(s + 1, data), compute(s + 1, x))
def test_reduction_dshape():
x = symbol('x', '5 * 3 * float32')
assert x.sum().dshape == dshape('float64')
assert x.sum(axis=0).dshape == dshape('3 * float64')
assert x.sum(axis=1).dshape == dshape('5 * float64')
assert x.sum(axis=(0, 1)).dshape == dshape('float64')
def test_compute_on_file(file):
s = symbol('s', discover(file))
assert eq(compute(s.x.sum(axis=1), file), x.sum(axis=1))
assert eq(compute(s.x.sum(), file, chunksize=(4, 6)), x.sum())
def test_summary_str():
x = symbol('x', '5 * 3 * float32')
assert 'keepdims' not in str(summary(a=x.min(), b=x.max()))
def test_nelements_records(recdata):
s = symbol('s', discover(recdata))
assert compute(s.nelements(), recdata) == np.prod(recdata.shape)
np.testing.assert_array_equal(
compute(s.nelements(axis=0), recdata),
np.zeros(recdata.shape[1]) + recdata.shape[0])
def test_summary_keepdims():
x = symbol('x', '5 * 3 * float32')
assert summary(a=x.min(), b=x.max()).dshape == \
dshape('{a: float32, b: float32}')
assert summary(a=x.min(), b=x.max(), keepdims=True).dshape == \
dshape('1 * 1 * {a: float32, b: float32}')
from datetime import datetime, timedelta
import pandas as pd
import pandas.util.testing as tm
import numpy as np
from pandas import DataFrame, Series
from string import ascii_lowercase
from blaze.compute.core import compute
from blaze import dshape, discover, transform
from blaze.expr import symbol, join, by, summary, distinct, shape
from blaze.expr import (merge, exp, mean, count, nunique, sum, min, max, any,
var, std, concat)
from blaze.compatibility import builtins, xfail, assert_series_equal
t = symbol('t', 'var * {name: string, amount: int, id: int}')
df = DataFrame([['Alice', 100, 1], ['Bob', 200, 2], ['Alice', 50, 3]],
columns=['name', 'amount', 'id'])
tbig = symbol('tbig',
'var * {name: string, sex: string[1], amount: int, id: int}')
dfbig = DataFrame(
[['Alice', 'F', 100, 1], ['Alice', 'F', 100, 3], ['Drew', 'F', 100, 4],
['Drew', 'M', 100, 5], ['Drew', 'M', 200, 5]],
columns=['name', 'sex', 'amount', 'id'])
def test_series_columnwise():
s = Series([1, 2, 3], name='a')
def test_dir():
t = symbol('t', '10 * int')
assert 'mean' in dir(t)
t = symbol('t', 'int')
assert 'mean' not in dir(t)
