Skip to content
Snippets Groups Projects
Commit bcc62421 authored by Nils-Arne Dreier's avatar Nils-Arne Dreier
Browse files

fix: avoid reopen the zarr again and again

parent 6a558e1b
No related branches found
No related tags found
No related merge requests found
Pipeline #101571 waiting for manual action
......@@ -2,16 +2,17 @@ import zarr
def get_var_group(v):
z = zarr.open(v.store)
parent_group_path = "/".join(v.path.split("/")[:-1])
if parent_group_path == "":
return z
if not hasattr(v, "root"):
v.root = zarr.open(v.store)
group_path = "/".join(v.path.split("/")[:-1])
if group_path == "":
return v.root
else:
return z[parent_group_path]
return v.root[group_path]
def get_var_parent_group(v):
var_group = get_var_group(v)
parent_var_path = var_group.attrs["hiopy::parent"]
parent_group = zarr.open(v.store)[parent_var_path]
parent_group = v.root[parent_var_path]
return parent_group
......@@ -79,14 +79,15 @@ def main():
)
# find all variables considered to be written in the input datasets:
def collect_data_vars(group):
def collect_data_vars(group, root):
for _name, item in group.arrays():
if "hiopy::enable" in item.attrs and item.attrs["hiopy::enable"]:
item.root = root
yield item
for _name, item in group.groups():
yield from collect_data_vars(item)
yield from collect_data_vars(item, root)
all_data_vars = list(chain(*[collect_data_vars(z) for z in args.datasets]))
all_data_vars = list(chain(*[collect_data_vars(z, z) for z in args.datasets]))
logging.info(f"Found {len(all_data_vars)} variables")
if len(all_data_vars) == 0:
raise RuntimeError("No variables found by the hiopy worker.")
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment