def test_filter(N=223):
C = Context()
dfm = C.iterates(N)
for n in [1, 2, 3, 7, 20]:
ans = dfm.filter(lambda x: x % n == n - 1)
assert ans.len() == N // n
def test_group(N=0, M=1):
C = Context()
def groups(e, out):
key = e % M
if key not in out:
out[key] = []
out[key].append(e)
dfm = C \
. iterates(N) \
. group( groups, lambda x: x, M )
assert dfm.len() == min(N, M)
i0 = C.rank * (M // C.procs) + min(C.rank, M % C.procs)
i1 = (C.rank + 1) * (M // C.procs) + min(C.rank + 1, M % C.procs)
for eg in dfm.E:
for e in eg:
loc = e % M
assert i0 <= loc
assert loc < i1
# use flatMap and check sizes again
fdf = dfm.flatMap(lambda x: x).len()
assert fdf == N
def test_head(N=107, M=3):
import numpy as np
C = Context()
h = C \
. iterates(N) \
. head(M)
assert len(h) == min(N, M)
assert h[0] == 0
assert h[-1] == len(h) - 1
def test_flatmap():
C = Context()
dfm = C . iterates(100) \
. flatMap(lambda x: [a for a in str(x)])
ans = dfm.collect()
if C.rank == 0:
N = 2 * 100 - 10
assert len(ans) == N
assert ans[0] == '0'
assert ans[-2] == '9'
assert ans[-1] == '9'
def test_reduce(N=101):
C = Context()
dfm = C.iterates(N).map(lambda x: [x])
s = [0]
def add(a, b):
a[0] += b[0]
return a
ans = dfm.reduce(add, s)
assert ans[0] == N * (N - 1) // 2
s = []
def append(a, b):
a.extend(b)
return a
ans = dfm.reduce(append, s)
assert len(ans) == N
if N >= 1:
assert ans[0] == 0 and ans[-1] == N - 1
def test_nodemap(N=100):
C = Context()
def test_fn(rank, elems):
assert rank == C.rank
return [(rank, len(elems))]
dfm = C . iterates(N) \
. nodeMap( test_fn )
assert dfm.len() == C.procs
v = dfm.head()
if N > 0:
assert v[0][0] == 0 and v[0][1] == (N // C.procs) + (N % C.procs != 0)
def test_repartition(N=0, M=1):
C = Context()
start = C \
. iterates(N) \
. map( lambda x: np.ones((x,4)) )
K = start.map(len).reduce(lambda a, b: a + b, 0)
if C.rank == 0:
assert K == N * (N - 1) // 2
dfm = start \
. repartition(len,
lambda df,rng: [df[r0:r1] for r0,r1 in rng],
np.vstack, M)
K = dfm.map(len).reduce(lambda a, b: a + b, 0)
if C.rank == 0:
print(N, K, [len(e) for e in dfm.E])
assert K == N * (N - 1) // 2
for e in dfm.E:
assert len(e) >= N // M
def test_dfm(N=97):
import numpy as np
C = Context()
dfm = C \
. iterates(N) \
. map( lambda x: np.ones(x) ) \
assert dfm.len() == N
assert len(dfm.E) >= N // C.procs
# Note: What follows requires correct ordering of iterates:
ans = dfm.collect()
if C.rank == 0:
for i in range(N):
assert len(ans[i]) == i
dfm = dfm.map(lambda x: x.sum())
ans = dfm.collect()
if C.rank == 0:
for i in range(N):
assert ans[i] == i