def test_read_array_names(ms):
_, short_name, _ = table_path_split(ms)
datasets = xds_from_ms(ms)
for ds in datasets:
for k, v in ds.data_vars.items():
product = ("~[" + str(ds.FIELD_ID) +
"," + str(ds.DATA_DESC_ID) + "]")
prefix = "".join(("read~", k, product))
assert key_split(v.data.name) == prefix
def test_write_array_names(ms, tmp_path):
_, short_name, _ = table_path_split(ms)
datasets = xds_from_ms(ms)
out_table = str(tmp_path / short_name)
writes = xds_to_table(datasets, out_table, "ALL")
for ds in writes:
for k, v in ds.data_vars.items():
prefix = "".join(("write~", k))
assert key_split(v.data.name) == prefix
def visualize(self, x, filename="simple_computation_graph", format=None):
if hasattr(x, "dask"):
dsk = x.__dask_optimize__(x.dask, x.__dask_keys__())
else:
dsk = x
deps = {k: get_dependencies(dsk, k) for k in dsk}
g = graphviz.Digraph(araph_attr={"rankdir": "LR"})
nodes = set()
edges = set()
for k in dsk:
key = self._node_key(k)
if key not in nodes:
g.node(key, label=key_split(k), shape="rectangle")
nodes.add(key)
for dep in deps[k]:
dep_key = self._node_key(dep)
if dep_key not in nodes:
g.node(dep_key, label=key_split(dep), shape="rectangle")
nodes.add(dep_key)
# Avoid circular references
if dep_key != key and (dep_key, key) not in edges:
g.edge(dep_key, key)
edges.add((dep_key, key))
data = g.pipe(format=self.format)
display_cls = _get_display_cls(self.format)
if self.filename is None:
return display_cls(data=data)
full_filename = ".".join([filename, self.format])
with open(full_filename, "wb") as f:
f.write(data)
return display_cls(filename=full_filename)
def __str__(self):
return "<%s: %s, npartitions=%d>" % (type(self).__name__,
key_split(self._name), self.npartitions)