from __future__ import absolute_import, division, print_function
from itertools import product
from math import ceil
import numpy as np
import pandas as pd
import h5py
try:
from sparse import COO
except ImportError:
raise ImportError("The 'sparse' package is required to use dask")
from dask.base import tokenize
import dask.dataframe as dd
import dask.array as da
from ..core import CSRReader, query_rect
from ..util import parse_cooler_uri, partition
def _get_group_info(path, grouppath, keys):
with h5py.File(path, "r") as f:
grp = f[grouppath]
if keys is None:
keys = list(grp.keys())
nrows = len(grp[keys[0]])
categoricals = {}
for key in keys:
dt = h5py.check_dtype(enum=grp[key].dtype)
if dt is not None:
categoricals[key] = sorted(dt, key=dt.__getitem__)
# Meta is an empty dataframe that serves as a compound "dtype"
meta = pd.DataFrame(
{key: np.array([], dtype=grp[key].dtype) for key in keys}, columns=keys
)
for key in categoricals:
meta[key] = pd.Categorical([], categories=categoricals[key], ordered=True)
return nrows, keys, meta, categoricals
def _slice_dataset(filepath, grouppath, key, slc, lock=None):
try:
if lock is not None:
lock.acquire()
with h5py.File(filepath, "r") as f:
return f[grouppath][key][slc]
finally:
if lock is not None:
lock.release()
def _slice_group(filepath, grouppath, keys, slc, lock=None):
try:
if lock is not None:
lock.acquire()
with h5py.File(filepath, "r") as f:
return {key: f[grouppath][key][slc] for key in keys}
finally:
if lock is not None:
lock.release()
def _restore_categories(data, categorical_columns):
for key, category_dict in categorical_columns.items():
data[key] = pd.Categorical.from_codes(data[key], category_dict, ordered=True)
return data
def read_table(group_uri, keys=None, chunksize=int(10e6), index=None, lock=None):
"""
Create a dask dataframe around a column-oriented table in HDF5.
A table is a group containing equal-length 1D datasets.
Parameters
----------
group_uri : str
URI to the HDF5 group storing the table.
keys : list, optional
list of HDF5 Dataset keys, default is to use all keys in the group
chunksize : int, optional
Chunk size
index : str, optional
Sorted column to use as index
lock : multiprocessing.Lock, optional
Lock to serialize HDF5 read/write access. Default is no lock.
Returns
-------
:class:`dask.dataframe.DataFrame`
Notes
-----
Learn more about the `dask <https://docs.dask.org/en/latest/>`_ project.
"""
filepath, grouppath = parse_cooler_uri(group_uri)
nrows, keys, meta, categoricals = _get_group_info(filepath, grouppath, keys)
# Make a unique task name
token = tokenize(filepath, grouppath, chunksize, keys)
task_name = "daskify-h5py-table-" + token
# Partition the table
divisions = (0,) + tuple(range(-1, nrows, chunksize))[1:]
if divisions[-1] != nrows - 1:
divisions = divisions + (nrows - 1,)
# Build the task graph
dsk = {}
for i in range(0, int(ceil(nrows / chunksize))):
slc = slice(i * chunksize, (i + 1) * chunksize)
data_dict = (_slice_group, filepath, grouppath, keys, slc, lock)
if categoricals:
data_dict = (_restore_categories, data_dict, categoricals)
dsk[task_name, i] = (pd.DataFrame, data_dict, None, meta.columns)
# Generate ddf from dask graph
df = dd.DataFrame(dsk, task_name, meta, divisions)
if index is not None:
df = df.set_index(index, sorted=True, drop=False)
return df
def _array_select(clr, i0, i1, j0, j1, field, sparse_array):
is_upper = clr._is_symm_upper
with clr.open("r") as h5:
dtype = h5['pixels'][field].dtype
reader = CSRReader(h5, field, max_chunk=500000000)
if is_upper:
i, j, v = query_rect(reader.query, i0, i1, j0, j1, duplex=True)
else:
i, j, v = reader.query(i0, i1, j0, j1)
if not len(v):
v = v.astype(dtype)
arr = COO((i - i0, j - j0), v, shape=(i1 - i0, j1 - j0))
if not sparse_array:
arr = arr.todense()
return arr
def load_dask_array(
clr, i0, i1, j0, j1, field="count", sparse_array=False, chunksize=256
):
"""
Create a parallel Dask array around the matrix representation of a cooler.
Parameters
----------
clr : :class:`cooler.Cooler`
Cooler object
i0, i1 : int
Row query range
j0, j1 : int
Column query range
field : str
Value column to query
sparse_array : bool, optional
Create a dask array backed by :class:`sparse.COO` sparse arrays
instead of dense numpy arrays (default).
chunksize : int, optional
Length of the rowwise chunks to partition the underlying data into.
Returns
-------
:class:`dask.array.Array`
"""
token = tokenize(clr.uri, i0, i1, j0, i1, field, chunksize)
task_name = "cooler-array-slice-" + token
shape = (i1 - i0, j1 - j0)
meta = _array_select(clr, 0, 0, 0, 0, field, sparse_array)
slices = [(lo, hi, j0, j1) for lo, hi in partition(0, shape[0], chunksize)]
chunks = (tuple(s[1] - s[0] for s in slices), (shape[1],))
keys = list(product([task_name], *[range(len(dim)) for dim in chunks]))
values = [(_array_select, clr, *slc, field, sparse_array) for slc in slices]
dsk = dict(zip(keys, values))
return da.Array(dsk, task_name, chunks, meta=meta, shape=shape)
