From 10edb65a065b59b5d9f9e5423d2f53479b4f41bf Mon Sep 17 00:00:00 2001 From: BMK Date: Mon, 16 Mar 2026 02:19:24 -0400 Subject: [PATCH] feat: add Zarr file format support with zarrReader --- aidrin/file_handling/file_parser.py | 5 +- aidrin/file_handling/readers/zarr_reader.py | 133 ++++++++++++++++++++ 2 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 aidrin/file_handling/readers/zarr_reader.py diff --git a/aidrin/file_handling/file_parser.py b/aidrin/file_handling/file_parser.py index 88009ddf..0a0b8468 100644 --- a/aidrin/file_handling/file_parser.py +++ b/aidrin/file_handling/file_parser.py @@ -6,7 +6,7 @@ from aidrin.file_handling.readers.hdf5_reader import hdf5Reader from aidrin.file_handling.readers.json_reader import jsonReader from aidrin.file_handling.readers.npz_reader import npzReader - +from aidrin.file_handling.readers.zarr_reader import zarrReader # Notes: # To add support for new file types: # - Add a new subclass of BaseFileReader with a .read() method @@ -21,7 +21,9 @@ ".xls, .xlsb, .xlsx, .xlsm": excelReader, ".json": jsonReader, ".h5": hdf5Reader, + ".zarr": zarrReader, # Add additional file types here + } # Supported file types. Read on front end to create select features. @@ -31,6 +33,7 @@ (".json", "JSON"), (".npz", "NumPy"), (".h5", "HDF5"), + (".zarr", "Zarr"), # Add additional file types here using the format: # (file_type,file_type_name) ] diff --git a/aidrin/file_handling/readers/zarr_reader.py b/aidrin/file_handling/readers/zarr_reader.py new file mode 100644 index 00000000..fd41447d --- /dev/null +++ b/aidrin/file_handling/readers/zarr_reader.py @@ -0,0 +1,133 @@ +import zarr +import numpy as np +import pandas as pd +from aidrin.file_handling.readers.base_reader import BaseFileReader + + +class zarrReader(BaseFileReader): + # zarr is similar to hdf5 - both store data in groups and arrays + # i learned this when i was reading about scientific data formats + # zarr is newer and works better for cloud storage compared to hdf5 + + def read(self): + try: + rows = [] + + # open the zarr store - can be a folder or zip file + store = zarr.open(self.file_path, mode='r') + + def recurse(group, prefix=''): + # going through all arrays inside zarr store + # zarr stores data in groups like folders in a computer + for key in group.keys(): + full_key = f"{prefix}/{key}".strip('/') + item = group[key] + + if isinstance(item, zarr.Array): + try: + # convert zarr array to numpy then to dataframe + data = item[:] + + # handle fill values similar to hdf5 reader + if hasattr(data, 'dtype') and data.dtype.kind in ('f', 'i', 'u'): + fill_val = item.fill_value + if fill_val is not None and fill_val != 0: + data = data.astype(np.float64) + data[data == fill_val] = np.nan + + # flatten to 2d if needed + if data.ndim > 2: + data = data.reshape(-1, data.shape[-1]) + elif data.ndim == 1: + data = data.reshape(-1, 1) + + df = pd.DataFrame(data) + df.columns = [f"{full_key}_{c}" + for c in df.columns] + + for _, row in df.iterrows(): + rows.append(row.to_dict()) + + except Exception as e: + self.logger.warning( + f"could not read array {full_key}: {e}" + ) + + elif isinstance(item, zarr.Group): + # if its a group go deeper into it + recurse(item, full_key) + + recurse(store) + + if not rows: + self.logger.warning("no data found in zarr file") + return None + + df = pd.DataFrame(rows) + self.logger.info(f"zarr file read successfully, shape: {df.shape}") + return df + + except Exception as e: + self.logger.error(f"error reading zarr file: {e}") + return None + + def parse(self): + # returns all group names inside the zarr store + # similar to parse() in hdf5_reader + try: + store = zarr.open(self.file_path, mode='r') + group_names = [] + + def collect_groups(group, prefix=''): + for key in group.keys(): + full_key = f"{prefix}/{key}".strip('/') + item = group[key] + if isinstance(item, zarr.Group): + group_names.append(full_key) + collect_groups(item, full_key) + + collect_groups(store) + self.logger.info(f"zarr groups found: {group_names}") + return group_names + + except Exception as e: + self.logger.error(f"error parsing zarr file: {e}") + return None + + def filter(self, kept_keys): + # filters zarr store to keep only selected groups + # useful when zarr file has many groups and user wants specific ones + try: + if isinstance(kept_keys, str): + kept_keys = [k.strip() for k in kept_keys.split(',')] + + kept_keys = set(kept_keys) + src = zarr.open(self.file_path, mode='r') + + # create a new in memory store with only kept groups + filtered = zarr.group() + + def copy_filtered(src_group, tgt_group, prefix=''): + for key in src_group.keys(): + full_key = f"{prefix}/{key}".strip('/') + item = src_group[key] + + if isinstance(item, zarr.Array): + if prefix in kept_keys or full_key in kept_keys: + tgt_group.create_dataset( + key, data=item[:], + chunks=item.chunks + ) + elif isinstance(item, zarr.Group): + if full_key in kept_keys: + sub = tgt_group.require_group(key) + copy_filtered(item, sub, full_key) + else: + copy_filtered(item, tgt_group, full_key) + + copy_filtered(src, filtered) + return filtered + + except Exception as e: + self.logger.error(f"error filtering zarr file: {e}") + return None \ No newline at end of file