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