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