def xds_from_parquet(store, columns=None, chunks=None, **kwargs):
if isinstance(store, DaskMSStore):
pass
elif isinstance(store, (str, Path)):
store = DaskMSStore(f"{store}", **kwargs.pop("storage_options", {}))
else:
raise TypeError(f"store '{store}' must be "
f"Path, str or DaskMSStore")
# If any kwargs are added, they should be popped prior to this check.
if len(kwargs) > 0:
warnings.warn(
f"The following unsupported kwargs were ignored in "
f"xds_from_parquet: {kwargs}", UserWarning)
columns = promote_columns(columns)
if chunks is None:
pass
elif isinstance(chunks, (tuple, list)):
if len(chunks) == 0 or any(not isinstance(c, dict) for c in chunks):
raise TypeError("chunks must be None or dict or list of dict")
elif isinstance(chunks, dict):
chunks = [chunks]
else:
raise TypeError("chunks must be None or dict or list of dict")
table_path = "" if store.table else "MAIN"
fragments = list(map(Path, store.rglob("*.parquet")))
ds_cfg = defaultdict(list)
# Iterate over all parquet files in the directory tree
# and group them by partition
partition_schemas = set()
for fragment in fragments:
*partitions, _ = fragment.relative_to(Path(table_path)).parts
fragment = ParquetFileProxy(store, str(fragment))
fragment_meta = fragment.metadata
metadata = json.loads(fragment_meta.metadata[DASKMS_METADATA.encode()])
partition_meta = metadata[DASKMS_PARTITION_KEY]
partition_meta = tuple(tuple((f, v)) for f, v in partition_meta)
partitions = _partition_values(partitions, partition_meta)
partition_schemas.add(partition_meta)
ds_cfg[partitions].append(fragment)
# Sanity check partition schemas of all parquet files
if len(partition_schemas) == 0:
raise ValueError(f"No parquet files found in {store.path}")
elif len(partition_schemas) != 1:
raise ValueError(f"Multiple partitions discovered {partition_schemas}")
partition_schemas = partition_schemas.pop()
datasets = []
# Now create a dataset per partition
for p, (partition, fragments) in enumerate(sorted(ds_cfg.items())):
fragments = list(sorted(fragments))
column_arrays = defaultdict(list)
fragment_rows = [f.metadata.num_rows for f in fragments]
# Returns a dictionary of lists mapping fragments to partitions.
partition_chunks = partition_chunking(p, fragment_rows, chunks)
for pieces in partition_chunks.values():
chunk_fragments = [fragments[i] for i, _ in pieces]
chunk_ranges = [r for _, r in pieces]
chunk_metas = [f.metadata for f in chunk_fragments]
rows = sum(end - start for start, end in chunk_ranges)
# NOTE(JSKenyon): This assumes that the schema/fields are
# consistent between fragments. This should be ok.
exemplar_schema = chunk_metas[0].schema.to_arrow_schema()
exemplar_fields = {
n: exemplar_schema.field(n)
for n in exemplar_schema.names
}
for column, field in column_iterator(exemplar_fields, columns):
field_metadata = field.metadata[DASKMS_METADATA.encode()]
field_metadata = json.loads(field_metadata)
dims = tuple(field_metadata["dims"])
if isinstance(field.type, TensorType):
shape = (rows, ) + field.type.shape
else:
shape = (rows, )
assert len(shape) == len(dims)
dtype = field.type.to_pandas_dtype()
meta = np.empty((0, ) * len(dims), dtype)
new_axes = {d: s for d, s in zip(dims, shape)}
read = da.blockwise(fragment_reader,
dims,
chunk_fragments,
None,
chunk_ranges,
None,
column,
None,
shape,
None,
dtype,
None,
adjust_chunks={"row": rows},
new_axes=new_axes,
meta=meta)
column_arrays[column].append((read, dims))
data_vars = {}
for column, values in column_arrays.items():
arrays, array_dims = zip(*values)
array_dims = set(array_dims)
if not len(array_dims) == 1:
raise ValueError(f"Inconsistent array dimensions "
f"{array_dims} for {column}")
data_vars[column] = (array_dims.pop(), da.concatenate(arrays))
attrs = dict(partition)
attrs[DASKMS_PARTITION_KEY] = partition_schemas
datasets.append(Dataset(data_vars, attrs=attrs))
return datasets
def xds_from_zarr(store, columns=None, chunks=None, **kwargs):
"""
Reads the zarr data store in `store` and returns list of
Dataset's containing the data.
Parameters
----------
store : str or Path
Path containing the data
columns : list of str or str or None
Columns to read. `None` or `"ALL"` stores all columns on each dataset.
Otherwise, a list of columns should be supplied.
chunks: dict or list of dicts
chunking schema for each dataset
**kwargs: optional
Returns
-------
writes : Dataset or list of Datasets
Dataset(s) representing write operations
"""
if isinstance(store, DaskMSStore):
pass
elif isinstance(store, (Path, str)):
store = DaskMSStore(f"{store}", **kwargs.pop("storage_options", {}))
else:
raise TypeError(f"store '{store}' must be "
f"Path, str or DaskMSStore")
# If any kwargs are added, they should be popped prior to this check.
if len(kwargs) > 0:
warnings.warn(
f"The following unsupported kwargs were ignored in "
f"xds_from_zarr: {kwargs}", UserWarning)
columns = promote_columns(columns)
if chunks is None:
pass
elif isinstance(chunks, (tuple, list)):
if not all(isinstance(v, dict) for v in chunks):
raise TypeError("chunks must be None, a dict or a list of dicts")
elif isinstance(chunks, dict):
chunks = [chunks]
else:
raise TypeError("chunks must be None, a dict or a list of dicts")
datasets = []
numpy_vars = []
# NOTE(JSKenyon): Iterating over all the zarr groups/arrays is VERY
# expensive if the metadata has not been consolidated.
zc.consolidate_metadata(store.map)
table_path = store.table if store.table else "MAIN"
table_group = zarr.open_consolidated(store.map)[table_path]
for g, (group_name,
group) in enumerate(sorted(table_group.groups(),
key=group_sortkey)):
group_attrs = decode_attr(dict(group.attrs))
dask_ms_attrs = group_attrs.pop(DASKMS_ATTR_KEY)
natural_chunks = dask_ms_attrs["chunks"]
group_chunks = {d: tuple(dc) for d, dc in natural_chunks.items()}
if chunks:
# Defer to user-supplied chunking strategy
try:
group_chunks.update(chunks[g])
except IndexError:
group_chunks.update(chunks[-1]) # Reuse last chunking.
pass
data_vars = {}
coords = {}
for name, zarray in column_iterator(group, columns):
attrs = decode_attr(dict(zarray.attrs[DASKMS_ATTR_KEY]))
dims = attrs["dims"]
coordinate = attrs.get("coordinate", False)
array_chunks = tuple(
group_chunks.get(d, s) for d, s in zip(dims, zarray.shape))
array_chunks = da.core.normalize_chunks(array_chunks, zarray.shape)
ext_args = extent_args(dims, array_chunks)
token_name = f"read~{name}-{tokenize(zarray, *ext_args)}"
read = da.blockwise(zarr_getter,
dims,
zarray,
None,
*ext_args,
concatenate=False,
name=token_name,
meta=np.empty((0, ) * zarray.ndim,
zarray.dtype))
read = inlined_array(read, ext_args[::2])
var = Variable(dims, read, attrs)
(coords if coordinate else data_vars)[name] = var
# Save numpy arrays for reification
typ = decode_type(attrs["array_type"])
if typ is np.ndarray:
numpy_vars.append(var)
elif typ is da.Array:
pass
else:
raise TypeError(f"Unknown array_type '{attrs['array_type']}'")
datasets.append(Dataset(data_vars, coords=coords, attrs=group_attrs))
# Reify any numpy arrays directly into their variables
for v, a in zip(numpy_vars, dask.compute(v.data for v in numpy_vars)[0]):
v.data = a
return datasets
def xds_to_parquet(xds, path, columns=None):
path, table = store_path_split(path)
if not isinstance(path, Path):
path = Path(path)
columns = promote_columns(columns)
if isinstance(xds, Dataset):
xds = [xds]
elif isinstance(xds, (tuple, list)):
if not all(isinstance(ds, Dataset) for ds in xds):
raise TypeError("xds must be a Dataset or list of Datasets")
else:
raise TypeError("xds must be a Dataset or list of Datasets")
datasets = []
base_schema = ArrowSchema.from_datasets(xds)
for ds_id, ds in enumerate(xds):
arrow_schema = base_schema.with_attributes(ds)
fragment = ParquetFragment(path / table, arrow_schema, ds_id)
chunk_ids = da.arange(len(ds.chunks["row"]), chunks=1)
args = [chunk_ids, ("row", )]
data_var_it = column_iterator(ds.data_vars, columns)
coord_it = column_iterator(ds.coords, columns)
for column, variable in itertools.chain(data_var_it, coord_it):
if not isinstance(variable.data, da.Array):
raise ValueError(f"Column {column} does not "
f"contain a dask Array")
if len(variable.dims[0]) == 0 or variable.dims[0] != "row":
raise ValueError(f"Column {column} dimensions "
f"{variable.dims} don't start with 'row'")
args.extend((column, None, variable.data, variable.dims))
for dim, chunk in zip(variable.dims[1:], variable.data.chunks[1:]):
if len(chunk) != 1:
raise ValueError(f"Chunking in {dim} is not yet "
f"supported.")
writes = da.blockwise(fragment.write, ("row", ),
*args,
align_arrays=False,
adjust_chunks={"row": 1},
meta=np.empty((0, ), np.bool))
writes = inlined_array(writes, chunk_ids)
# Transfer any partition information over to the write dataset
partition = ds.attrs.get(DASKMS_PARTITION_KEY, False)
if not partition:
attrs = None
else:
attrs = {
DASKMS_PARTITION_KEY: partition,
**{k: getattr(ds, k)
for k, _ in partition}
}
datasets.append(Dataset({"WRITE": (("row", ), writes)}, attrs=attrs))
return datasets
def xds_to_zarr(xds, store, columns=None, rechunk=False, **kwargs):
"""
Stores a dataset of list of datasets defined by `xds` in
file location `store`.
Parameters
----------
xds : Dataset or list of Datasets
Data
store : str or Path
Path to store the data
columns : list of str or str or None
Columns to store. `None` or `"ALL"` stores all columns on each dataset.
Otherwise, a list of columns should be supplied. All coordinates
associated with a specified column will be written automatically.
rechunk : bool
Controls whether dask arrays should be automatically rechunked to be
consistent with existing on-disk zarr arrays while writing to disk.
**kwargs : optional
Returns
-------
writes : Dataset
A Dataset representing the write operations
"""
if isinstance(store, DaskMSStore):
pass
elif isinstance(store, (Path, str)):
store = DaskMSStore(f"{store}", **kwargs.pop("storage_options", {}))
else:
raise TypeError(f"store '{store}' must be "
f"Path, str or DaskMSStore")
# If any kwargs are added, they should be popped prior to this check.
if len(kwargs) > 0:
warnings.warn(
f"The following unsupported kwargs were ignored in "
f"xds_to_zarr: {kwargs}", UserWarning)
columns = promote_columns(columns)
if isinstance(xds, Dataset):
xds = [xds]
elif isinstance(xds, (tuple, list)):
if not all(isinstance(ds, Dataset) for ds in xds):
raise TypeError("xds must be a Dataset or list of Datasets")
else:
raise TypeError("xds must be a Dataset or list of Datasets")
write_datasets = []
for di, ds in enumerate(xds):
data_vars, coords = select_vars_and_coords(ds, columns)
# Create a new ds which is consistent with what we want to write.
ds = Dataset(data_vars, coords=coords, attrs=ds.attrs)
ds, group = prepare_zarr_group(di, ds, store, rechunk=rechunk)
data_vars = dict(_gen_writes(ds.data_vars, ds.chunks, group))
# Include coords in the write dataset so they're reified
data_vars.update(
dict(_gen_writes(ds.coords, ds.chunks, group, indirect_dims=True)))
# Transfer any partition information over to the write dataset
partition = ds.attrs.get(DASKMS_PARTITION_KEY, False)
if not partition:
attrs = None
else:
attrs = {
DASKMS_PARTITION_KEY: partition,
**{k: getattr(ds, k)
for k, _ in partition}
}
write_datasets.append(Dataset(data_vars, attrs=attrs))
return write_datasets
def xds_from_parquet(store, columns=None, chunks=None):
store, table = store_path_split(store)
store = store / table
if not isinstance(store, Path):
store = Path(store)
columns = promote_columns(columns)
if chunks is None:
pass
elif isinstance(chunks, (tuple, list)):
if len(chunks) == 0 or any(not isinstance(c, dict) for c in chunks):
raise TypeError("chunks must be None or dict or list of dict")
elif isinstance(chunks, dict):
chunks = [chunks]
else:
raise TypeError("chunks must be None or dict or list of dict")
fragments = store.rglob("*.parquet")
ds_cfg = defaultdict(list)
# Iterate over all parquet files in the directory tree
# and group them by partition
partition_schemas = set()
for fragment in fragments:
*partitions, parquet_file = fragment.relative_to(store).parts
fragment = ParquetFileProxy(fragment)
fragment_meta = fragment.metadata
metadata = json.loads(fragment_meta.metadata[DASKMS_METADATA.encode()])
partition_meta = metadata[DASKMS_PARTITION_KEY]
partition_meta = tuple(tuple((f, v)) for f, v in partition_meta)
partitions = _partition_values(partitions, partition_meta)
partition_schemas.add(partition_meta)
ds_cfg[partitions].append(fragment)
# Sanity check partition schemas of all parquet files
if len(partition_schemas) == 0:
raise ValueError(f"No parquet files found in {store}")
elif len(partition_schemas) != 1:
raise ValueError(f"Multiple partitions discovered {partition_schemas}")
partition_schemas = partition_schemas.pop()
datasets = []
# Now create a dataset per partition
for p, (partition, fragments) in enumerate(sorted(ds_cfg.items())):
fragments = list(sorted(fragments))
column_arrays = defaultdict(list)
fragment_rows = [f.metadata.num_rows for f in fragments]
for (f, (start, end)) in partition_chunking(p, fragment_rows, chunks):
fragment = fragments[f]
fragment_meta = fragment.metadata
rows = fragment_meta.num_rows
schema = fragment_meta.schema.to_arrow_schema()
fields = {n: schema.field(n) for n in schema.names}
for column, field in column_iterator(fields, columns):
field_metadata = field.metadata[DASKMS_METADATA.encode()]
field_metadata = json.loads(field_metadata)
dims = tuple(field_metadata["dims"])
if isinstance(field.type, TensorType):
shape = (rows,) + field.type.shape
else:
shape = (rows,)
assert len(shape) == len(dims)
meta = np.empty((0,)*len(dims), field.type.to_pandas_dtype())
new_axes = {d: s for d, s in zip(dims, shape)}
read = da.blockwise(fragment.read_column, dims,
column, None,
start, None,
end, None,
new_axes=new_axes,
meta=meta)
column_arrays[column].append((read, dims))
data_vars = {}
for column, values in column_arrays.items():
arrays, array_dims = zip(*values)
array_dims = set(array_dims)
if not len(array_dims) == 1:
raise ValueError(f"Inconsistent array dimensions "
f"{array_dims} for {column}")
data_vars[column] = (array_dims.pop(), da.concatenate(arrays))
attrs = dict(partition)
attrs[DASKMS_PARTITION_KEY] = partition_schemas
datasets.append(Dataset(data_vars, attrs=attrs))
return datasets
def xds_from_zarr(store, columns=None, chunks=None):
"""
Reads the zarr data store in `store` and returns list of
Dataset's containing the data.
Parameters
----------
store : str or Path
Path containing the data
columns : list of str or str or None
Columns to read. `None` or `"ALL"` stores all columns on each dataset.
Otherwise, a list of columns should be supplied.
chunks: dict or list of dicts
chunking schema for each dataset
Returns
-------
writes : Dataset or list of Datasets
Dataset(s) representing write operations
"""
store, table = store_path_split(store)
if isinstance(store, Path):
store = str(store)
if not isinstance(store, str):
raise TypeError("store must be a Path, str")
columns = promote_columns(columns)
if chunks is None:
pass
elif isinstance(chunks, (tuple, list)):
if not all(isinstance(v, dict) for v in chunks):
raise TypeError("chunks must be None, a dict or a list of dicts")
elif isinstance(chunks, dict):
chunks = [chunks]
else:
raise TypeError("chunks must be None, a dict or a list of dicts")
datasets = []
numpy_vars = []
table_group = zarr.open(store)[table]
for g, (group_name, group) in enumerate(sorted(table_group.groups())):
group_attrs = decode_attr(dict(group.attrs))
dask_ms_attrs = group_attrs.pop(DASKMS_ATTR_KEY)
natural_chunks = dask_ms_attrs["chunks"]
group_chunks = {d: tuple(dc) for d, dc in natural_chunks.items()}
if chunks:
# Defer to user-supplied chunking strategy
try:
group_chunks.update(chunks[g])
except IndexError:
pass
data_vars = {}
coords = {}
for name, zarray in column_iterator(group, columns):
attrs = decode_attr(dict(zarray.attrs[DASKMS_ATTR_KEY]))
dims = attrs["dims"]
coordinate = attrs.get("coordinate", False)
array_chunks = tuple(
group_chunks.get(d, s) for d, s in zip(dims, zarray.shape))
array_chunks = da.core.normalize_chunks(array_chunks, zarray.shape)
ext_args = extent_args(dims, array_chunks)
token_name = f"read~{name}-{tokenize(zarray, *ext_args)}"
read = da.blockwise(zarr_getter,
dims,
zarray,
None,
*ext_args,
concatenate=False,
name=token_name,
meta=np.empty((0, ) * zarray.ndim,
zarray.dtype))
read = inlined_array(read, ext_args[::2])
var = Variable(dims, read, attrs)
(coords if coordinate else data_vars)[name] = var
# Save numpy arrays for reification
typ = decode_type(attrs["array_type"])
if typ is np.ndarray:
numpy_vars.append(var)
elif typ is da.Array:
pass
else:
raise TypeError(f"Unknown array_type '{attrs['array_type']}'")
datasets.append(Dataset(data_vars, coords=coords, attrs=group_attrs))
# Reify any numpy arrays directly into their variables
for v, a in zip(numpy_vars, dask.compute(v.data for v in numpy_vars)[0]):
v.data = a
return datasets
def xds_to_zarr(xds, store, columns=None):
"""
Stores a dataset of list of datasets defined by `xds` in
file location `store`.
Parameters
----------
xds : Dataset or list of Datasets
Data
store : str or Path
Path to store the data
columns : list of str or str or None
Columns to store. `None` or `"ALL"` stores all columns on each dataset.
Otherwise, a list of columns should be supplied.
Returns
-------
writes : Dataset
A Dataset representing the write operations
"""
store, table = store_path_split(store)
if isinstance(store, Path):
store = str(store)
if not isinstance(store, str):
raise TypeError(f"store '{store}' must be Path or str")
columns = promote_columns(columns)
if isinstance(xds, Dataset):
xds = [xds]
elif isinstance(xds, (tuple, list)):
if not all(isinstance(ds, Dataset) for ds in xds):
raise TypeError("xds must be a Dataset or list of Datasets")
else:
raise TypeError("xds must be a Dataset or list of Datasets")
write_datasets = []
for di, ds in enumerate(xds):
group = prepare_zarr_group(di, ds, store, table)
write_args = (ds.chunks, columns, group)
data_vars = dict(_gen_writes(ds.data_vars, *write_args))
# Include coords in the write dataset so they're reified
data_vars.update(dict(_gen_writes(ds.coords, *write_args)))
# Transfer any partition information over to the write dataset
partition = ds.attrs.get(DASKMS_PARTITION_KEY, False)
if not partition:
attrs = None
else:
attrs = {
DASKMS_PARTITION_KEY: partition,
**{k: getattr(ds, k)
for k, _ in partition}
}
write_datasets.append(Dataset(data_vars, attrs=attrs))
return write_datasets
def xds_to_parquet(xds, store, columns=None, **kwargs):
if isinstance(store, DaskMSStore):
pass
elif isinstance(store, (str, Path)):
store = DaskMSStore(f"{store}", **kwargs.pop("storage_options", {}))
else:
raise TypeError(f"store '{store}' must be "
f"Path, str or DaskMSStore")
# If any kwargs are added, they should be popped prior to this check.
if len(kwargs) > 0:
warnings.warn(
f"The following unsupported kwargs were ignored in "
f"xds_to_parquet: {kwargs}", UserWarning)
columns = promote_columns(columns)
if isinstance(xds, Dataset):
xds = [xds]
elif isinstance(xds, (tuple, list)):
if not all(isinstance(ds, Dataset) for ds in xds):
raise TypeError("xds must be a Dataset or list of Datasets")
else:
raise TypeError("xds must be a Dataset or list of Datasets")
datasets = []
base_schema = ArrowSchema.from_datasets(xds)
for ds_id, ds in enumerate(xds):
arrow_schema = base_schema.with_attributes(ds)
fragment = ParquetFragment(store, store.table, arrow_schema, ds_id)
chunk_ids = da.arange(len(ds.chunks["row"]), chunks=1)
args = [chunk_ids, ("row", )]
data_var_it = column_iterator(ds.data_vars, columns)
coord_it = column_iterator(ds.coords, columns)
for column, variable in itertools.chain(data_var_it, coord_it):
if not isinstance(variable.data, da.Array):
raise ValueError(f"Column {column} does not "
f"contain a dask Array")
if len(variable.dims[0]) == 0 or variable.dims[0] != "row":
raise ValueError(f"Column {column} dimensions "
f"{variable.dims} don't start with 'row'")
args.extend((column, None, variable.data, variable.dims))
for dim, chunk in zip(variable.dims[1:], variable.data.chunks[1:]):
if len(chunk) != 1:
raise ValueError(f"Chunking in {dim} is not yet "
f"supported.")
writes = da.blockwise(fragment.write, ("row", ),
*args,
align_arrays=False,
adjust_chunks={"row": 1},
meta=np.empty((0, ), bool))
writes = inlined_array(writes, chunk_ids)
# Transfer any partition information over to the write dataset
partition = ds.attrs.get(DASKMS_PARTITION_KEY, False)
if not partition:
attrs = None
else:
attrs = {
DASKMS_PARTITION_KEY: partition,
**{k: getattr(ds, k)
for k, _ in partition}
}
datasets.append(Dataset({"WRITE": (("row", ), writes)}, attrs=attrs))
return datasets