Merge sorted sequences of sorted numpy arrays about dask HOT 3 CLOSED

import numpy as np
from cytoolz import merge_sorted, concat, map, partition_all

def shard(n, x):
    """

    >>> list(shard(3, list(range(10)))) 
    [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
    """
    for i in range(0, len(x), n):
        yield x[i: i + n]


def array_merge_sorted(seqs, out=None, out_chunksize=2**14):
    """ Merge step of external sort

    Merged sorted sequences of numpy arrays in to out result
    """
    assert out is not None

    seqs2 = [concat(x.tolist() for x in seq) for seq in seqs]

    seq = merge_sorted(*seqs2)
    chunks = map(np.array, partition_all(out_chunksize, seq))

    for i, chunk in enumerate(chunks):
        out[i*out_chunksize: min(len(out), (i+1)*out_chunksize)] = chunk

    return out

from dask.frame.esort import array_merge_sorted, shard
import numpy as np

def test_shard():
    result = list(shard(3, np.arange(10)))
    assert result[0].tolist() == [0, 1, 2]
    assert result[1].tolist() == [3, 4, 5]
    assert result[2].tolist() == [6, 7, 8]
    assert result[3].tolist() == [9]


def test_esort():
    seqs = [np.random.random(size=(np.random.randint(100))) for i in range(5)]
    sorteds = [np.sort(x) for x in seqs]
    chunks = [shard(5, x) for x in sorteds]

    out = np.empty(shape=(sum(len(x) for x in seqs),))

    array_merge_sorted(chunks, out=out)

    assert (out == np.sort(np.concatenate(seqs))).all()