03 Dask源码剖析-Dask的数据模型-Array

Dask源码剖析是一个专栏，更多章节请点击文章列表查看。后续我会更新更多内容上来。

Collection：Array

前面我们了解了Bag数据模型。通过阅读Bag的源码我们大概熟悉了Dask的数据模型的套路。操作有加载，有计算，都是先转成Delayed，并按照逻辑构建成Graph。真正的计算得到结果需要执行compute提交到集群或在本地执行，并且dask很智能的会graph进行优化。

本节我们看下Array数据结构。
从源码目录里，我们可以很快的定位到Array的包：

就在dask源码包的下面，有个array的包。而其核心代码都在core.py中。

还是老规矩，先构建一些简单的demo，然后根据demo通过debug的形式看下是如何实现的。(例子来自dask-tutorial)

Array的创建

这回dask-tutorial给的示例是构建了一套demo的数据集，先运行以下，生成random.hdf5

%run prep.py -d random
# 通过三方库（可pip install h5py获取） h5py 加载数据集
# 下述操作只是获取了一个h5py文件指针，并不会把所有数据加载到内存。
import h5py
import os
f = h5py.File(os.path.join('data', 'random.hdf5'), mode='r')
# 数据没实际加载到内存，
dset = f['/x']

from_array加载

import dask.array as da
x = da.from_array(dset, chunks=(1000000,))
x

通过dask array提供的from_array函数可以构建Array对象。但hdf5里面的数据是和bag的from_sequence函数一样加载到内存里了，还是只是构建了一个delayed对象呢？
我们看下from_array的源码：

# dask/array/core.py
def from_array(
    x,
    chunks="auto",
    name=None,
    lock=False,
    asarray=None,
    fancy=True,
    getitem=None,
    meta=None,
):
    """ 从一个长得像array的对象中，创建一个dask array对象。
    什么叫长得像array的对象？其实就是像python原生的数组、np.ndarray，这些具备数组行为的对象。具体行为（方法）必须有``.shape``, ``.ndim``, ``.dtype``,然后还要支持numpy风格的切片（slicing）
    参数：
    ----------
    x : 具备数组行为的对象
    chunks : int, tuple
        如何对数组进行分块. 必须是下述形式之一:

        - 传类似1000这样的整数，即分解成每个维度长度为1000的块，比如2250*2750，会分解成（3*3）个(1000,1000)的块.
        - 按shape进行分块，例如 (1000, 1000).
        - 像((1000, 1000, 500), (400, 400))这种，在每个维度上给出每块分块大小.
        - 指定每块的大小, 比如 "100 MiB" 
        - 传入 "auto" 会使用 ``array.chunk-size`` 作为分块大小.

        - 传入-1 或 None 则作为1整块处理，不进行分块.
    name : str, optional
        array的名称. 默认是字符 ``x``的hash.
        默认情况下，使用Python标准的sha1算法. 当安装了 cityhash, xxhash 或 murmurhash.
       会变成相应的算法，这些算法在大规模的特征标记（创建token化的name时）可以获得更好的性能。
        
        当使用 ``name=False`` 会使用随机名称代替hashing
        .. 注意::
		   由于name会被用于graph的key，所以必须保证key的唯一性
    lock : bool or Lock, optional
        如果 ``x`` 不支持并行读，那么传入True（Dask会帮你创建）或者指定的Lock对象。
    asarray : bool, optional
        asarray在numpy中，会将类似array的数据结构转换成ndarray（是直接值引用，也就是说不会
        新开辟内存）。这种操作可以很方便的让任意类array的数据结构使用numpy的特性。
        dask的asarray与之类似，这个参数默认为None，当__array_function__未定义时，效果类似于
        将这个参数设置为True，即将各个分块转换成numpy array。
        而当asarray为False的时候，分块不转换成ndarray。
    fancy : bool, optional
        fancy indexing是numpy的概念，很简单：即指传递索引数组以便一次得到多个数组元素。
        使用fancy indexing时要特别注意的一点是返回数组的shape反映的是索引数组的shape而不是被索引的原数组的shape。
        如果 ``x`` 不支持 fancy indexing (例如 用list或array做索引) 那么此参数为 False. 默认是 True.
    meta : Array-like, optional
        dask array 结果的元数据.  一般由输入数组的slicing（分片操作）得到
        默认是输入数组.
    """
    if isinstance(x, Array):
        raise ValueError(
            "Array is already a dask array. Use 'asarray' or " "'rechunk' instead."
        )
    elif is_dask_collection(x):
        warnings.warn(
            "Passing an object to dask.array.from_array which is already a "
            "Dask collection. This can lead to unexpected behavior."
        )
	# 像例子中hdf5对象是不会在此转numpy的，否则相当于在client端加载数据了。
    if isinstance(x, (list, tuple, memoryview) + np.ScalarType):
        x = np.array(x)

    if asarray is None:
        # 例子中的hdf5是不具备__array_function__方法的，后续会借助numpy的asarray转ndarray
        asarray = not hasattr(x, "__array_function__")

    previous_chunks = getattr(x, "chunks", None)
    # 这里chunk是一个tuple，存储了各chunks的分块大小
    chunks = normalize_chunks(
        chunks, x.shape, dtype=x.dtype, previous_chunks=previous_chunks
    )
    # 命名
    if name in (None, True):
        token = tokenize(x, chunks)
        original_name = "array-original-" + token
        name = name or "array-" + token
    elif name is False:
        original_name = name = "array-" + str(uuid.uuid1())
    else:
        original_name = name
    # 是否加锁，例子里是支持并行读的，所以没有加锁，lock为False
    if lock is True:
        lock = SerializableLock()

    # Always use the getter for h5py etc. Not using isinstance(x, np.ndarray)
    # because np.matrix is a subclass of np.ndarray.
    if type(x) is np.ndarray and all(len(c) == 1 for c in chunks):
        # No slicing needed
        dsk = {(name,) + (0,) * x.ndim: x}
    else:
        # getitem实际是函数的参数，但是文档里并没有暴露，应该是暂时不希望用户传入此参数。
        # 实际上getitem是定义了如何从x里加载数据到内存np.ndarray的方法。
        # 一般来说不传的话，会有三种gettter方法：
        if getitem is None:
            if type(x) is np.ndarray and not lock:
                # simpler and cleaner, but missing all the nuances of getter
                getitem = operator.getitem
            elif fancy:
                # hdf5文件会通过此getter来获取数据，getter的代码我们在下面解读一下。
                getitem = getter
            else:
                # 如果fancy参数为false，会用此逻辑获取
                getitem = getter_nofancy
        # 这里生成了graph，getem是此处的关键，代码我们在下面解读一下。
        dsk = getem(
            original_name,
            chunks,
            getitem=getitem,
            shape=x.shape,
            out_name=name,
            lock=lock,
            asarray=asarray,
            dtype=x.dtype,
        )
        dsk[original_name] = x

    # TileDB（一种分布式格点数据存储数据库），这里dask.array提供了对其的支持，可以看到和其他数据源的区别在于：
    # TileDB并没有把x作为meta，这里暂不深究。
    if x.__class__.__module__.split(".")[0] == "tiledb" and hasattr(x, "_ctx_"):
        return Array(dsk, name, chunks, dtype=x.dtype)
    # 和文档中提到的一样，如果meta元数据为空，则把它本身赋值给meta
    if meta is None:
        meta = x

    return Array(dsk, name, chunks, meta=meta, dtype=getattr(x, "dtype", None))

# dask/array/core.py
def getter(a, b, asarray=True, lock=None):
    # a：数据源，b：slice
    if isinstance(b, tuple) and any(x is None for x in b):
        b2 = tuple(x for x in b if x is not None)
        b3 = tuple(
            None if x is None else slice(None, None)
            for x in b
            if not isinstance(x, Integral)
        )
        # 通过递归，解决多维度获取分块的问题
        return getter(a, b2, asarray=asarray, lock=lock)[b3]

    if lock:
        lock.acquire()
    try:
        c = a[b]
        if asarray:
            # 转numpy的ndarray
            c = np.asarray(c)
    finally:
        if lock:
            lock.release()
    return c
def getem(
    arr,
    chunks,
    getitem=getter,
    shape=None,
    out_name=None,
    lock=False,
    asarray=True,
    dtype=None,
):
    """dask从各种array-like的数据中获取分块数据的方法，其中getter即上面提到的3类getter之一。

    >>> getem('X', chunks=(2, 3), shape=(4, 6))  # doctest: +SKIP
    {('X', 0, 0): (getter, 'X', (slice(0, 2), slice(0, 3))),
     ('X', 1, 0): (getter, 'X', (slice(2, 4), slice(0, 3))),
     ('X', 1, 1): (getter, 'X', (slice(2, 4), slice(3, 6))),
     ('X', 0, 1): (getter, 'X', (slice(0, 2), slice(3, 6)))}

    >>> getem('X', chunks=((2, 2), (3, 3)))  # doctest: +SKIP
    {('X', 0, 0): (getter, 'X', (slice(0, 2), slice(0, 3))),
     ('X', 1, 0): (getter, 'X', (slice(2, 4), slice(0, 3))),
     ('X', 1, 1): (getter, 'X', (slice(2, 4), slice(3, 6))),
     ('X', 0, 1): (getter, 'X', (slice(0, 2), slice(3, 6)))}
    """
    out_name = out_name or arr
    chunks = normalize_chunks(chunks, shape, dtype=dtype)
    # 这里会生成一个迭代器，输出(out_name,0、1、2...)作为key
    keys = product([out_name], *(range(len(bds)) for bds in chunks))
    # 这里会构建slices的列表，标记start/end，便于分块
    slices = slices_from_chunks(chunks)

    if (
        has_keyword(getitem, "asarray")
        and has_keyword(getitem, "lock")
        and (not asarray or lock)
    ):
        values = [(getitem, arr, x, asarray, lock) for x in slices]
    else:
        # Common case, drop extra parameters
        values = [(getitem, arr, x) for x in slices]
    # ok，这里会将(name,分块数作为key)，对应的(getitem,arr,x)作为value，getitem暂时不会调用，
    # 会作为delayed对象后续执行按需调用。
    # arr是数据源，x在这里是slice，包含了重要的start/end信息。回过头来可以再看下上面getitem的源码，便于更深入理解。
    return dict(zip(keys, values)

上述整个从h5py对象到da.array的过程，会生成一个graph，用思维导图表示即：

这样我们就通过from_array成功构建了一个dask-array对象。其在jupyter里的展示还是很友好的：

小结一下，dask array数据结构与bag没有太大差异，也是将数据转换成了graph。如果传入的x支持异步加载（如hdf5），那么实际数据仍然是懒加载的。
其他的构建方式还有

• from_zarr
• from_tiledb
• from_npy_stack
• from_delayed
  这里就不一一介绍了，大致原理一致